Examples with DisableOnCloseInputStream org.apache.nifi.controller.repository.io.DisableOnCloseInputStream used on opensource projects

Example 1 with DisableOnCloseInputStream

use of org.apache.nifi.controller.repository.io.DisableOnCloseInputStream in project nifi by apache.

the class StandardProcessSession method read.

@Override
public void read(FlowFile source, boolean allowSessionStreamManagement, InputStreamCallback reader) {
    verifyTaskActive();
    source = validateRecordState(source, true);
    final StandardRepositoryRecord record = records.get(source);
    try {
        ensureNotAppending(record.getCurrentClaim());
        claimCache.flush(record.getCurrentClaim());
    } catch (final IOException e) {
        throw new FlowFileAccessException("Failed to access ContentClaim for " + source.toString(), e);
    }
    try (final InputStream rawIn = getInputStream(source, record.getCurrentClaim(), record.getCurrentClaimOffset(), true);
        final InputStream limitedIn = new LimitedInputStream(rawIn, source.getSize());
        final InputStream disableOnCloseIn = new DisableOnCloseInputStream(limitedIn);
        final ByteCountingInputStream countingStream = new ByteCountingInputStream(disableOnCloseIn, this.bytesRead)) {
        // We want to differentiate between IOExceptions thrown by the repository and IOExceptions thrown from
        // Processor code. As a result, as have the FlowFileAccessInputStream that catches IOException from the repository
        // and translates into either FlowFileAccessException or ContentNotFoundException. We keep track of any
        // ContentNotFoundException because if it is thrown, the Processor code may catch it and do something else with it
        // but in reality, if it is thrown, we want to know about it and handle it, even if the Processor code catches it.
        final FlowFileAccessInputStream ffais = new FlowFileAccessInputStream(countingStream, source, record.getCurrentClaim());
        boolean cnfeThrown = false;
        try {
            incrementReadCount(source);
            reader.process(createTaskTerminationStream(ffais));
            // Allow processors to close the file after reading to avoid too many files open or do smart session stream management.
            if (this.currentReadClaimStream != null && !allowSessionStreamManagement) {
                currentReadClaimStream.close();
                currentReadClaimStream = null;
            }
        } catch (final ContentNotFoundException cnfe) {
            cnfeThrown = true;
            throw cnfe;
        } finally {
            decrementReadCount(source);
            bytesRead += countingStream.getBytesRead();
            // if cnfeThrown is true, we don't need to re-thrown the Exception; it will propagate.
            if (!cnfeThrown && ffais.getContentNotFoundException() != null) {
                throw ffais.getContentNotFoundException();
            }
        }
    } catch (final ContentNotFoundException nfe) {
        handleContentNotFound(nfe, record);
    } catch (final IOException ex) {
        throw new ProcessException("IOException thrown from " + connectableDescription + ": " + ex.toString(), ex);
    }
}

Example 2 with DisableOnCloseInputStream

use of org.apache.nifi.controller.repository.io.DisableOnCloseInputStream in project nifi by apache.

the class StandardProcessSession method getInputStream.

private InputStream getInputStream(final FlowFile flowFile, final ContentClaim claim, final long offset, final boolean allowCachingOfStream) throws ContentNotFoundException {
    // that there is no actual content.
    if (flowFile.getSize() == 0L) {
        return new ByteArrayInputStream(new byte[0]);
    }
    try {
        // callback for reading FlowFile 1 and if we used the same stream we'd be destroying the ability to read from FlowFile 1.
        if (allowCachingOfStream && readRecursionSet.isEmpty() && writeRecursionSet.isEmpty()) {
            if (currentReadClaim == claim) {
                if (currentReadClaimStream != null && currentReadClaimStream.getBytesConsumed() <= offset) {
                    final long bytesToSkip = offset - currentReadClaimStream.getBytesConsumed();
                    if (bytesToSkip > 0) {
                        StreamUtils.skip(currentReadClaimStream, bytesToSkip);
                    }
                    return new DisableOnCloseInputStream(currentReadClaimStream);
                }
            }
            claimCache.flush(claim);
            final InputStream rawInStream = context.getContentRepository().read(claim);
            if (currentReadClaimStream != null) {
                currentReadClaimStream.close();
            }
            currentReadClaim = claim;
            currentReadClaimStream = new ByteCountingInputStream(rawInStream);
            StreamUtils.skip(currentReadClaimStream, offset);
            // reuse the same InputStream for the next FlowFile
            return new DisableOnCloseInputStream(currentReadClaimStream);
        } else {
            claimCache.flush(claim);
            final InputStream rawInStream = context.getContentRepository().read(claim);
            try {
                StreamUtils.skip(rawInStream, offset);
            } catch (IOException ioe) {
                try {
                    rawInStream.close();
                } catch (final Exception e) {
                    ioe.addSuppressed(ioe);
                }
                throw ioe;
            }
            return rawInStream;
        }
    } catch (final ContentNotFoundException cnfe) {
        throw cnfe;
    } catch (final EOFException eof) {
        throw new ContentNotFoundException(claim, eof);
    } catch (final IOException ioe) {
        throw new FlowFileAccessException("Failed to read content of " + flowFile, ioe);
    }
}

Example 3 with DisableOnCloseInputStream

use of org.apache.nifi.controller.repository.io.DisableOnCloseInputStream in project nifi by apache.

the class StandardProcessSession method write.

@Override
public FlowFile write(FlowFile source, final StreamCallback writer) {
    verifyTaskActive();
    source = validateRecordState(source);
    final StandardRepositoryRecord record = records.get(source);
    final ContentClaim currClaim = record.getCurrentClaim();
    long writtenToFlowFile = 0L;
    ContentClaim newClaim = null;
    try {
        newClaim = claimCache.getContentClaim();
        claimLog.debug("Creating ContentClaim {} for 'write' for {}", newClaim, source);
        ensureNotAppending(newClaim);
        if (currClaim != null) {
            claimCache.flush(currClaim.getResourceClaim());
        }
        try (final InputStream is = getInputStream(source, currClaim, record.getCurrentClaimOffset(), true);
            final InputStream limitedIn = new LimitedInputStream(is, source.getSize());
            final InputStream disableOnCloseIn = new DisableOnCloseInputStream(limitedIn);
            final ByteCountingInputStream countingIn = new ByteCountingInputStream(disableOnCloseIn, bytesRead);
            final OutputStream os = claimCache.write(newClaim);
            final OutputStream disableOnCloseOut = new DisableOnCloseOutputStream(os);
            final ByteCountingOutputStream countingOut = new ByteCountingOutputStream(disableOnCloseOut)) {
            writeRecursionSet.add(source);
            // We want to differentiate between IOExceptions thrown by the repository and IOExceptions thrown from
            // Processor code. As a result, as have the FlowFileAccessInputStream that catches IOException from the repository
            // and translates into either FlowFileAccessException or ContentNotFoundException. We keep track of any
            // ContentNotFoundException because if it is thrown, the Processor code may catch it and do something else with it
            // but in reality, if it is thrown, we want to know about it and handle it, even if the Processor code catches it.
            final FlowFileAccessInputStream ffais = new FlowFileAccessInputStream(countingIn, source, currClaim);
            final FlowFileAccessOutputStream ffaos = new FlowFileAccessOutputStream(countingOut, source);
            boolean cnfeThrown = false;
            try {
                writer.process(createTaskTerminationStream(ffais), createTaskTerminationStream(ffaos));
            } catch (final ContentNotFoundException cnfe) {
                cnfeThrown = true;
                throw cnfe;
            } finally {
                writtenToFlowFile = countingOut.getBytesWritten();
                this.bytesWritten += writtenToFlowFile;
                this.bytesRead += countingIn.getBytesRead();
                writeRecursionSet.remove(source);
                // if cnfeThrown is true, we don't need to re-thrown the Exception; it will propagate.
                if (!cnfeThrown && ffais.getContentNotFoundException() != null) {
                    throw ffais.getContentNotFoundException();
                }
            }
        }
    } catch (final ContentNotFoundException nfe) {
        destroyContent(newClaim);
        handleContentNotFound(nfe, record);
    } catch (final IOException ioe) {
        destroyContent(newClaim);
        throw new ProcessException("IOException thrown from " + connectableDescription + ": " + ioe.toString(), ioe);
    } catch (final FlowFileAccessException ffae) {
        destroyContent(newClaim);
        throw ffae;
    } catch (final Throwable t) {
        destroyContent(newClaim);
        throw t;
    }
    removeTemporaryClaim(record);
    final FlowFileRecord newFile = new StandardFlowFileRecord.Builder().fromFlowFile(record.getCurrent()).contentClaim(newClaim).contentClaimOffset(Math.max(0L, newClaim.getLength() - writtenToFlowFile)).size(writtenToFlowFile).build();
    record.setWorking(newFile);
    return newFile;
}

Example 4 with DisableOnCloseInputStream

use of org.apache.nifi.controller.repository.io.DisableOnCloseInputStream in project nifi by apache.

the class StandardProcessSession method exportTo.

@Override
public void exportTo(FlowFile source, final OutputStream destination) {
    verifyTaskActive();
    source = validateRecordState(source);
    final StandardRepositoryRecord record = records.get(source);
    if (record.getCurrentClaim() == null) {
        return;
    }
    try {
        ensureNotAppending(record.getCurrentClaim());
        claimCache.flush(record.getCurrentClaim());
    } catch (final IOException e) {
        throw new FlowFileAccessException("Failed to access ContentClaim for " + source.toString(), e);
    }
    try (final InputStream rawIn = getInputStream(source, record.getCurrentClaim(), record.getCurrentClaimOffset(), true);
        final InputStream limitedIn = new LimitedInputStream(rawIn, source.getSize());
        final InputStream disableOnCloseIn = new DisableOnCloseInputStream(limitedIn);
        final ByteCountingInputStream countingStream = new ByteCountingInputStream(disableOnCloseIn, this.bytesRead)) {
        // but in reality, if it is thrown, we want to know about it and handle it, even if the Processor code catches it.
        try (final FlowFileAccessInputStream ffais = new FlowFileAccessInputStream(countingStream, source, record.getCurrentClaim())) {
            boolean cnfeThrown = false;
            try {
                incrementReadCount(source);
                StreamUtils.copy(ffais, createTaskTerminationStream(destination), source.getSize());
            } catch (final ContentNotFoundException cnfe) {
                cnfeThrown = true;
                throw cnfe;
            } finally {
                decrementReadCount(source);
                // if cnfeThrown is true, we don't need to re-throw the Exception; it will propagate.
                if (!cnfeThrown && ffais.getContentNotFoundException() != null) {
                    throw ffais.getContentNotFoundException();
                }
            }
        }
    } catch (final ContentNotFoundException nfe) {
        handleContentNotFound(nfe, record);
    } catch (final IOException ex) {
        throw new ProcessException("IOException thrown from " + connectableDescription + ": " + ex.toString(), ex);
    }
}