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

Example 1 with LimitedInputStream

use of org.apache.nifi.controller.repository.io.LimitedInputStream 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 LimitedInputStream

use of org.apache.nifi.controller.repository.io.LimitedInputStream 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 3 with LimitedInputStream

use of org.apache.nifi.controller.repository.io.LimitedInputStream 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 4 with LimitedInputStream

use of org.apache.nifi.controller.repository.io.LimitedInputStream 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);
    }
}

Example 5 with LimitedInputStream

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

the class FileSystemRepository method read.

@Override
public InputStream read(final ContentClaim claim) throws IOException {
    if (claim == null) {
        return new ByteArrayInputStream(new byte[0]);
    }
    final Path path = getPath(claim, true);
    final FileInputStream fis = new FileInputStream(path.toFile());
    if (claim.getOffset() > 0L) {
        try {
            StreamUtils.skip(fis, claim.getOffset());
        } catch (IOException ioe) {
            IOUtils.closeQuietly(fis);
            throw ioe;
        }
    }
    // see javadocs for claim.getLength() as to why we do this.
    if (claim.getLength() >= 0) {
        return new LimitedInputStream(fis, claim.getLength());
    } else {
        return fis;
    }
}