Examples with ByteCountingInputStream org.apache.nifi.stream.io.ByteCountingInputStream used on opensource projects

Example 6 with ByteCountingInputStream

use of org.apache.nifi.stream.io.ByteCountingInputStream 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 7 with ByteCountingInputStream

use of org.apache.nifi.stream.io.ByteCountingInputStream 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 8 with ByteCountingInputStream

use of org.apache.nifi.stream.io.ByteCountingInputStream 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 9 with ByteCountingInputStream

use of org.apache.nifi.stream.io.ByteCountingInputStream in project nifi by apache.

the class StandardProcessSession method read.

@Override
public InputStream read(FlowFile source) {
    verifyTaskActive();
    source = validateRecordState(source, true);
    final StandardRepositoryRecord record = records.get(source);
    try {
        ensureNotAppending(record.getCurrentClaim());
    } catch (final IOException e) {
        throw new FlowFileAccessException("Failed to access ContentClaim for " + source.toString(), e);
    }
    final InputStream rawIn = getInputStream(source, record.getCurrentClaim(), record.getCurrentClaimOffset(), true);
    final InputStream limitedIn = new LimitedInputStream(rawIn, source.getSize());
    final ByteCountingInputStream countingStream = new ByteCountingInputStream(limitedIn);
    final FlowFileAccessInputStream ffais = new FlowFileAccessInputStream(countingStream, source, record.getCurrentClaim());
    final FlowFile sourceFlowFile = source;
    final InputStream errorHandlingStream = new InputStream() {

        private boolean closed = false;

        @Override
        public int read() throws IOException {
            try {
                return ffais.read();
            } catch (final ContentNotFoundException cnfe) {
                handleContentNotFound(cnfe, record);
                close();
                throw cnfe;
            } catch (final FlowFileAccessException ffae) {
                LOG.error("Failed to read content from " + sourceFlowFile + "; rolling back session", ffae);
                rollback(true);
                close();
                throw ffae;
            }
        }

        @Override
        public int read(final byte[] b) throws IOException {
            return read(b, 0, b.length);
        }

        @Override
        public int read(final byte[] b, final int off, final int len) throws IOException {
            try {
                return ffais.read(b, off, len);
            } catch (final ContentNotFoundException cnfe) {
                handleContentNotFound(cnfe, record);
                close();
                throw cnfe;
            } catch (final FlowFileAccessException ffae) {
                LOG.error("Failed to read content from " + sourceFlowFile + "; rolling back session", ffae);
                rollback(true);
                close();
                throw ffae;
            }
        }

        @Override
        public void close() throws IOException {
            decrementReadCount(sourceFlowFile);
            if (!closed) {
                StandardProcessSession.this.bytesRead += countingStream.getBytesRead();
                closed = true;
            }
            ffais.close();
            openInputStreams.remove(sourceFlowFile);
        }

        @Override
        public int available() throws IOException {
            return ffais.available();
        }

        @Override
        public long skip(long n) throws IOException {
            return ffais.skip(n);
        }

        @Override
        public boolean markSupported() {
            return ffais.markSupported();
        }

        @Override
        public synchronized void mark(int readlimit) {
            ffais.mark(readlimit);
        }

        @Override
        public synchronized void reset() throws IOException {
            ffais.reset();
        }

        @Override
        public String toString() {
            return "ErrorHandlingInputStream[FlowFile=" + sourceFlowFile + "]";
        }
    };
    incrementReadCount(sourceFlowFile);
    openInputStreams.put(sourceFlowFile, errorHandlingStream);
    return createTaskTerminationStream(errorHandlingStream);
}

Example 10 with ByteCountingInputStream

use of org.apache.nifi.stream.io.ByteCountingInputStream 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);
    }
}