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

Example 1 with ByteCountingOutputStream

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

the class CompressableRecordWriter method resetWriteStream.

/**
 * Resets the streams to prepare for a new block
 *
 * @param eventId the first id that will be written to the new block
 * @throws IOException if unable to flush/close the current streams properly
 */
protected void resetWriteStream(final Long eventId) throws IOException {
    try {
        if (out != null) {
            out.flush();
        }
        final long byteOffset = (byteCountingOut == null) ? rawOutStream.getBytesWritten() : byteCountingOut.getBytesWritten();
        final TocWriter tocWriter = getTocWriter();
        if (compressed) {
            // We don't have to check if the writer is dirty because we will have already checked before calling this method.
            if (out != null) {
                out.close();
            }
            if (tocWriter != null && eventId != null) {
                tocWriter.addBlockOffset(rawOutStream.getBytesWritten(), eventId);
            }
            final OutputStream writableStream = new BufferedOutputStream(new GZIPOutputStream(new NonCloseableOutputStream(rawOutStream), 1), 65536);
            this.byteCountingOut = new ByteCountingOutputStream(writableStream, byteOffset);
        } else {
            if (tocWriter != null && eventId != null) {
                tocWriter.addBlockOffset(rawOutStream.getBytesWritten(), eventId);
            }
            this.byteCountingOut = rawOutStream;
        }
        this.out = new DataOutputStream(byteCountingOut);
        resetDirtyFlag();
    } catch (final IOException ioe) {
        markDirty();
        throw ioe;
    }
}

Example 2 with ByteCountingOutputStream

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

the class EventFileCompressor method compress.

public static void compress(final File input, final TocReader tocReader, final File output, final TocWriter tocWriter) throws IOException {
    try (final InputStream fis = new FileInputStream(input);
        final OutputStream fos = new FileOutputStream(output);
        final ByteCountingOutputStream byteCountingOut = new ByteCountingOutputStream(fos)) {
        int blockIndex = 0;
        while (true) {
            // Determine the min and max byte ranges for the current block.
            final long blockStart = tocReader.getBlockOffset(blockIndex);
            if (blockStart == -1) {
                break;
            }
            long blockEnd = tocReader.getBlockOffset(blockIndex + 1);
            if (blockEnd < 0) {
                blockEnd = input.length();
            }
            final long firstEventId = tocReader.getFirstEventIdForBlock(blockIndex);
            final long blockStartOffset = byteCountingOut.getBytesWritten();
            try (final OutputStream ncos = new NonCloseableOutputStream(byteCountingOut);
                final OutputStream gzipOut = new GZIPOutputStream(ncos, 1)) {
                StreamUtils.copy(fis, gzipOut, blockEnd - blockStart);
            }
            tocWriter.addBlockOffset(blockStartOffset, firstEventId);
            blockIndex++;
        }
    }
    // Close the TOC Reader and TOC Writer
    CloseableUtil.closeQuietly(tocReader, tocWriter);
}

Example 3 with ByteCountingOutputStream

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

the class RecordBin method offer.

public boolean offer(final FlowFile flowFile, final RecordReader recordReader, final ProcessSession flowFileSession, final boolean block) throws IOException, MalformedRecordException, SchemaNotFoundException {
    if (isComplete()) {
        logger.debug("RecordBin.offer for id={} returning false because {} is complete", new Object[] { flowFile.getId(), this });
        return false;
    }
    final boolean locked;
    if (block) {
        writeLock.lock();
        locked = true;
    } else {
        locked = writeLock.tryLock();
    }
    if (!locked) {
        logger.debug("RecordBin.offer for id={} returning false because failed to get lock for {}", new Object[] { flowFile.getId(), this });
        return false;
    }
    boolean flowFileMigrated = false;
    try {
        if (isComplete()) {
            logger.debug("RecordBin.offer for id={} returning false because {} is complete", new Object[] { flowFile.getId(), this });
            return false;
        }
        logger.debug("Migrating id={} to {}", new Object[] { flowFile.getId(), this });
        Record record;
        while ((record = recordReader.nextRecord()) != null) {
            if (recordWriter == null) {
                final OutputStream rawOut = session.write(merged);
                logger.debug("Created OutputStream using session {} for {}", new Object[] { session, this });
                this.out = new ByteCountingOutputStream(rawOut);
                recordWriter = writerFactory.createWriter(logger, record.getSchema(), out);
                recordWriter.beginRecordSet();
            }
            recordWriter.write(record);
            recordCount++;
        }
        // This will be closed by the MergeRecord class anyway but we have to close it
        // here because it needs to be closed before we are able to migrate the FlowFile
        // to a new Session.
        recordReader.close();
        flowFileSession.migrate(this.session, Collections.singleton(flowFile));
        flowFileMigrated = true;
        this.flowFiles.add(flowFile);
        if (isFull()) {
            logger.debug(this + " is now full. Completing bin.");
            complete("Bin is full");
        } else if (isOlderThan(thresholds.getMaxBinMillis(), TimeUnit.MILLISECONDS)) {
            logger.debug(this + " is now expired. Completing bin.");
            complete("Bin is older than " + thresholds.getMaxBinAge());
        }
        return true;
    } catch (final Exception e) {
        logger.error("Failed to create merged FlowFile from " + (flowFiles.size() + 1) + " input FlowFiles; routing originals to failure", e);
        try {
            // This will be closed by the MergeRecord class anyway but we have to close it
            // here because it needs to be closed before we are able to migrate the FlowFile
            // to a new Session.
            recordReader.close();
            if (recordWriter != null) {
                recordWriter.close();
            }
            if (this.out != null) {
                this.out.close();
            }
            if (!flowFileMigrated) {
                flowFileSession.migrate(this.session, Collections.singleton(flowFile));
                this.flowFiles.add(flowFile);
            }
        } finally {
            complete = true;
            session.remove(merged);
            session.transfer(flowFiles, MergeRecord.REL_FAILURE);
            session.commit();
        }
        return true;
    } finally {
        writeLock.unlock();
    }
}

Example 4 with ByteCountingOutputStream

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

the class HttpOutput method setOutputStream.

public void setOutputStream(OutputStream outputStream) {
    interruptableOut = new InterruptableOutputStream(outputStream);
    this.countingOut = new ByteCountingOutputStream(interruptableOut);
}

Example 5 with ByteCountingOutputStream

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

the class OrcFlowFileWriter method getStream.

@VisibleForTesting
public OutputStream getStream() throws IOException {
    if (rawWriter == null) {
        rawWriter = new ByteCountingOutputStream(flowFileOutputStream);
        rawWriter.write(OrcFile.MAGIC.getBytes());
        headerLength = rawWriter.getBytesWritten();
        writer = new OutStream("metadata", bufferSize, codec, new DirectStream(rawWriter));
        protobufWriter = CodedOutputStream.newInstance(writer);
    }
    return rawWriter;
}