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

Example 11 with ByteCountingOutputStream

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

the class StandardProcessSession method write.

@Override
public FlowFile write(FlowFile source, final OutputStreamCallback writer) {
    verifyTaskActive();
    source = validateRecordState(source);
    final StandardRepositoryRecord record = records.get(source);
    long writtenToFlowFile = 0L;
    ContentClaim newClaim = null;
    try {
        newClaim = claimCache.getContentClaim();
        claimLog.debug("Creating ContentClaim {} for 'write' for {}", newClaim, source);
        ensureNotAppending(newClaim);
        try (final OutputStream stream = claimCache.write(newClaim);
            final OutputStream disableOnClose = new DisableOnCloseOutputStream(stream);
            final ByteCountingOutputStream countingOut = new ByteCountingOutputStream(disableOnClose)) {
            try {
                writeRecursionSet.add(source);
                final OutputStream ffaos = new FlowFileAccessOutputStream(countingOut, source);
                writer.process(createTaskTerminationStream(ffaos));
            } finally {
                writtenToFlowFile = countingOut.getBytesWritten();
                bytesWritten += countingOut.getBytesWritten();
            }
        } finally {
            writeRecursionSet.remove(source);
        }
    } catch (final ContentNotFoundException nfe) {
        // need to reset write claim before we can remove the claim
        resetWriteClaims();
        destroyContent(newClaim);
        handleContentNotFound(nfe, record);
    } catch (final FlowFileAccessException ffae) {
        // need to reset write claim before we can remove the claim
        resetWriteClaims();
        destroyContent(newClaim);
        throw ffae;
    } catch (final IOException ioe) {
        // need to reset write claim before we can remove the claim
        resetWriteClaims();
        destroyContent(newClaim);
        throw new ProcessException("IOException thrown from " + connectableDescription + ": " + ioe.toString(), ioe);
    } catch (final Throwable t) {
        // need to reset write claim before we can remove the claim
        resetWriteClaims();
        destroyContent(newClaim);
        throw t;
    }
    removeTemporaryClaim(record);
    final FlowFileRecord newFile = new StandardFlowFileRecord.Builder().fromFlowFile(record.getCurrent()).contentClaim(newClaim).contentClaimOffset(Math.max(0, newClaim.getLength() - writtenToFlowFile)).size(writtenToFlowFile).build();
    record.setWorking(newFile);
    return newFile;
}

Example 12 with ByteCountingOutputStream

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

the class FileSystemRepository method write.

private OutputStream write(final ContentClaim claim, final boolean append) throws IOException {
    if (claim == null) {
        throw new NullPointerException("ContentClaim cannot be null");
    }
    if (!(claim instanceof StandardContentClaim)) {
        // else, just throw an Exception because it is not valid for this Repository
        throw new IllegalArgumentException("Cannot write to " + claim + " because that Content Claim does belong to this Content Repository");
    }
    final StandardContentClaim scc = (StandardContentClaim) claim;
    if (claim.getLength() > 0) {
        throw new IllegalArgumentException("Cannot write to " + claim + " because it has already been written to.");
    }
    ByteCountingOutputStream claimStream = writableClaimStreams.get(scc.getResourceClaim());
    final int initialLength = append ? (int) Math.max(0, scc.getLength()) : 0;
    final ByteCountingOutputStream bcos = claimStream;
    final OutputStream out = new OutputStream() {

        private long bytesWritten = 0L;

        private boolean recycle = true;

        private boolean closed = false;

        @Override
        public String toString() {
            return "FileSystemRepository Stream [" + scc + "]";
        }

        @Override
        public synchronized void write(final int b) throws IOException {
            if (closed) {
                throw new IOException("Stream is closed");
            }
            try {
                bcos.write(b);
            } catch (final IOException ioe) {
                recycle = false;
                throw new IOException("Failed to write to " + this, ioe);
            }
            bytesWritten++;
            scc.setLength(bytesWritten + initialLength);
        }

        @Override
        public synchronized void write(final byte[] b) throws IOException {
            if (closed) {
                throw new IOException("Stream is closed");
            }
            try {
                bcos.write(b);
            } catch (final IOException ioe) {
                recycle = false;
                throw new IOException("Failed to write to " + this, ioe);
            }
            bytesWritten += b.length;
            scc.setLength(bytesWritten + initialLength);
        }

        @Override
        public synchronized void write(final byte[] b, final int off, final int len) throws IOException {
            if (closed) {
                throw new IOException("Stream is closed");
            }
            try {
                bcos.write(b, off, len);
            } catch (final IOException ioe) {
                recycle = false;
                throw new IOException("Failed to write to " + this, ioe);
            }
            bytesWritten += len;
            scc.setLength(bytesWritten + initialLength);
        }

        @Override
        public synchronized void flush() throws IOException {
            if (closed) {
                throw new IOException("Stream is closed");
            }
            bcos.flush();
        }

        @Override
        public synchronized void close() throws IOException {
            closed = true;
            if (alwaysSync) {
                ((FileOutputStream) bcos.getWrappedStream()).getFD().sync();
            }
            if (scc.getLength() < 0) {
                // If claim was not written to, set length to 0
                scc.setLength(0L);
            }
            // if we've not yet hit the threshold for appending to a resource claim, add the claim
            // to the writableClaimQueue so that the Resource Claim can be used again when create()
            // is called. In this case, we don't have to actually close the file stream. Instead, we
            // can just add it onto the queue and continue to use it for the next content claim.
            final long resourceClaimLength = scc.getOffset() + scc.getLength();
            if (recycle && resourceClaimLength < maxAppendableClaimLength) {
                final ClaimLengthPair pair = new ClaimLengthPair(scc.getResourceClaim(), resourceClaimLength);
                // We are checking that writableClaimStreams contains the resource claim as a key, as a sanity check.
                // It should always be there. However, we have encountered a bug before where we archived content before
                // we should have. As a result, the Resource Claim and the associated OutputStream were removed from the
                // writableClaimStreams map, and this caused a NullPointerException. Worse, the call here to
                // writableClaimQueue.offer() means that the ResourceClaim was then reused, which resulted in an endless
                // loop of NullPointerException's being thrown. As a result, we simply ensure that the Resource Claim does
                // in fact have an OutputStream associated with it before adding it back to the writableClaimQueue.
                final boolean enqueued = writableClaimStreams.get(scc.getResourceClaim()) != null && writableClaimQueue.offer(pair);
                if (enqueued) {
                    LOG.debug("Claim length less than max; Adding {} back to Writable Claim Queue", this);
                } else {
                    writableClaimStreams.remove(scc.getResourceClaim());
                    resourceClaimManager.freeze(scc.getResourceClaim());
                    bcos.close();
                    LOG.debug("Claim length less than max; Closing {} because could not add back to queue", this);
                    if (LOG.isTraceEnabled()) {
                        LOG.trace("Stack trace: ", new RuntimeException("Stack Trace for closing " + this));
                    }
                }
            } else {
                // we've reached the limit for this claim. Don't add it back to our queue.
                // Instead, just remove it and move on.
                // Mark the claim as no longer being able to be written to
                resourceClaimManager.freeze(scc.getResourceClaim());
                // ensure that the claim is no longer on the queue
                writableClaimQueue.remove(new ClaimLengthPair(scc.getResourceClaim(), resourceClaimLength));
                bcos.close();
                LOG.debug("Claim lenth >= max; Closing {}", this);
                if (LOG.isTraceEnabled()) {
                    LOG.trace("Stack trace: ", new RuntimeException("Stack Trace for closing " + this));
                }
            }
        }
    };
    LOG.debug("Writing to {}", out);
    if (LOG.isTraceEnabled()) {
        LOG.trace("Stack trace: ", new RuntimeException("Stack Trace for writing to " + out));
    }
    return out;
}

Example 13 with ByteCountingOutputStream

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

the class FileSystemRepository method remove.

private boolean remove(final ResourceClaim claim) {
    if (claim == null) {
        return false;
    }
    // If the claim is still in use, we won't remove it.
    if (claim.isInUse()) {
        return false;
    }
    Path path = null;
    try {
        path = getPath(claim);
    } catch (final ContentNotFoundException cnfe) {
    }
    // Ensure that we have no writable claim streams for this resource claim
    final ByteCountingOutputStream bcos = writableClaimStreams.remove(claim);
    if (bcos != null) {
        try {
            bcos.close();
        } catch (final IOException e) {
            LOG.warn("Failed to close Output Stream for {} due to {}", claim, e);
        }
    }
    final File file = path.toFile();
    if (!file.delete() && file.exists()) {
        LOG.warn("Unable to delete {} at path {}", new Object[] { claim, path });
        return false;
    }
    return true;
}

Example 14 with ByteCountingOutputStream

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

the class FileSystemRepository method create.

@Override
public ContentClaim create(final boolean lossTolerant) throws IOException {
    ResourceClaim resourceClaim;
    final long resourceOffset;
    final ClaimLengthPair pair = writableClaimQueue.poll();
    if (pair == null) {
        final long currentIndex = index.incrementAndGet();
        String containerName = null;
        boolean waitRequired = true;
        ContainerState containerState = null;
        for (long containerIndex = currentIndex; containerIndex < currentIndex + containers.size(); containerIndex++) {
            final long modulatedContainerIndex = containerIndex % containers.size();
            containerName = containerNames.get((int) modulatedContainerIndex);
            containerState = containerStateMap.get(containerName);
            if (!containerState.isWaitRequired()) {
                waitRequired = false;
                break;
            }
        }
        if (waitRequired) {
            containerState.waitForArchiveExpiration();
        }
        final long modulatedSectionIndex = currentIndex % SECTIONS_PER_CONTAINER;
        final String section = String.valueOf(modulatedSectionIndex);
        final String claimId = System.currentTimeMillis() + "-" + currentIndex;
        resourceClaim = resourceClaimManager.newResourceClaim(containerName, section, claimId, lossTolerant, true);
        resourceOffset = 0L;
        LOG.debug("Creating new Resource Claim {}", resourceClaim);
        // we always append because there may be another ContentClaim using the same resource claim.
        // However, we know that we will never write to the same claim from two different threads
        // at the same time because we will call create() to get the claim before we write to it,
        // and when we call create(), it will remove it from the Queue, which means that no other
        // thread will get the same Claim until we've finished writing to it.
        final File file = getPath(resourceClaim).toFile();
        ByteCountingOutputStream claimStream = new SynchronizedByteCountingOutputStream(new FileOutputStream(file, true), file.length());
        writableClaimStreams.put(resourceClaim, claimStream);
        incrementClaimantCount(resourceClaim, true);
    } else {
        resourceClaim = pair.getClaim();
        resourceOffset = pair.getLength();
        LOG.debug("Reusing Resource Claim {}", resourceClaim);
        incrementClaimantCount(resourceClaim, false);
    }
    final StandardContentClaim scc = new StandardContentClaim(resourceClaim, resourceOffset);
    return scc;
}