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

Example 1 with NonThreadSafeCircularBuffer

use of org.apache.nifi.stream.io.util.NonThreadSafeCircularBuffer in project nifi by apache.

the class PutSplunk method processDelimitedMessages.

/**
 * Read delimited messages from the FlowFile tracking which messages are sent successfully.
 */
private void processDelimitedMessages(final ProcessContext context, final ProcessSession session, final FlowFile flowFile, final ChannelSender sender, final String delimiter) {
    final String protocol = context.getProperty(PROTOCOL).getValue();
    final byte[] delimiterBytes = delimiter.getBytes(StandardCharsets.UTF_8);
    // The NonThreadSafeCircularBuffer allows us to add a byte from the stream one at a time and see if it matches
    // some pattern. We can use this to search for the delimiter as we read through the stream of bytes in the FlowFile
    final NonThreadSafeCircularBuffer buffer = new NonThreadSafeCircularBuffer(delimiterBytes);
    final AtomicLong messagesSent = new AtomicLong(0L);
    final FlowFileMessageBatch messageBatch = new FlowFileMessageBatch(session, flowFile);
    activeBatches.add(messageBatch);
    try (final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(final InputStream rawIn) throws IOException {
                // contents of a single message
                byte[] data = null;
                boolean streamFinished = false;
                int nextByte;
                try (final InputStream bufferedIn = new BufferedInputStream(rawIn);
                    final ByteCountingInputStream in = new ByteCountingInputStream(bufferedIn)) {
                    long messageStartOffset = in.getBytesConsumed();
                    // read until we're out of data.
                    while (!streamFinished) {
                        nextByte = in.read();
                        if (nextByte > -1) {
                            baos.write(nextByte);
                        }
                        if (nextByte == -1) {
                            // we ran out of data. This message is complete.
                            data = getMessage(baos, baos.size(), protocol);
                            streamFinished = true;
                        } else if (buffer.addAndCompare((byte) nextByte)) {
                            // we matched our delimiter. This message is complete. We want all of the bytes from the
                            // underlying BAOS except for the last 'delimiterBytes.length' bytes because we don't want
                            // the delimiter itself to be sent.
                            data = getMessage(baos, baos.size() - delimiterBytes.length, protocol);
                        }
                        if (data != null) {
                            final long messageEndOffset = in.getBytesConsumed();
                            // If the message has no data, ignore it.
                            if (data.length != 0) {
                                final long rangeStart = messageStartOffset;
                                try {
                                    sender.send(data);
                                    messageBatch.addSuccessfulRange(rangeStart, messageEndOffset);
                                    messagesSent.incrementAndGet();
                                } catch (final IOException e) {
                                    messageBatch.addFailedRange(rangeStart, messageEndOffset, e);
                                }
                            }
                            // reset BAOS so that we can start a new message.
                            baos.reset();
                            data = null;
                            messageStartOffset = in.getBytesConsumed();
                        }
                    }
                }
            }
        });
        messageBatch.setNumMessages(messagesSent.get());
    } catch (final IOException ioe) {
    // Since this can be thrown only from closing the ByteArrayOutputStream(), we have already
    // completed everything that we need to do, so there's nothing really to be done here
    }
}

Example 2 with NonThreadSafeCircularBuffer

use of org.apache.nifi.stream.io.util.NonThreadSafeCircularBuffer in project nifi by apache.

the class StreamUtils method copyExclusive.

/**
 * Copies data from in to out until either we are out of data (returns null) or we hit one of the byte patterns identified by the <code>stoppers</code> parameter (returns the byte pattern
 * matched). The byte pattern matched will NOT be copied to the output and will be un-read from the input.
 *
 * @param in the source to read bytes from
 * @param out the destination to write bytes to
 * @param maxBytes the maximum number of bytes to copy
 * @param stoppers byte patterns which will cause the copy to stop if found
 * @return the byte array matched, or null if end of stream was reached
 * @throws IOException for issues reading or writing to underlying streams
 */
public static byte[] copyExclusive(final InputStream in, final OutputStream out, final int maxBytes, final byte[]... stoppers) throws IOException {
    if (stoppers.length == 0) {
        return null;
    }
    int longest = 0;
    NonThreadSafeCircularBuffer longestBuffer = null;
    final List<NonThreadSafeCircularBuffer> circularBuffers = new ArrayList<>();
    for (final byte[] stopper : stoppers) {
        final NonThreadSafeCircularBuffer circularBuffer = new NonThreadSafeCircularBuffer(stopper);
        if (stopper.length > longest) {
            longest = stopper.length;
            longestBuffer = circularBuffer;
            circularBuffers.add(0, circularBuffer);
        } else {
            circularBuffers.add(circularBuffer);
        }
    }
    long bytesRead = 0;
    while (true) {
        final int next = in.read();
        if (next == -1) {
            return null;
        } else if (maxBytes > 0 && bytesRead++ > maxBytes) {
            throw new BytePatternNotFoundException("Did not encounter any byte pattern that was expected; data does not appear to be in the expected format");
        }
        for (final NonThreadSafeCircularBuffer circ : circularBuffers) {
            if (circ.addAndCompare((byte) next)) {
                // The longest buffer has some data that may not have been written out yet; we need to make sure
                // that we copy out those bytes.
                final int bytesToCopy = longest - circ.getByteArray().length;
                for (int i = 0; i < bytesToCopy; i++) {
                    final int oldestByte = longestBuffer.getOldestByte();
                    if (oldestByte != -1) {
                        out.write(oldestByte);
                        longestBuffer.addAndCompare((byte) 0);
                    }
                }
                return circ.getByteArray();
            }
        }
        if (longestBuffer.isFilled()) {
            out.write(longestBuffer.getOldestByte());
        }
    }
}

Example 3 with NonThreadSafeCircularBuffer

use of org.apache.nifi.stream.io.util.NonThreadSafeCircularBuffer in project nifi by apache.

the class StreamUtils method copyInclusive.

/**
 * Copies data from in to out until either we are out of data (returns null) or we hit one of the byte patterns identified by the <code>stoppers</code> parameter (returns the byte pattern
 * matched). The bytes in the stopper will be copied.
 *
 * @param in the source to read bytes from
 * @param out the destination to write bytes to
 * @param maxBytes the max bytes to copy
 * @param stoppers patterns of bytes which if seen will cause the copy to stop
 * @return the byte array matched, or null if end of stream was reached
 * @throws IOException if issues occur reading or writing bytes to the underlying streams
 */
public static byte[] copyInclusive(final InputStream in, final OutputStream out, final int maxBytes, final byte[]... stoppers) throws IOException {
    if (stoppers.length == 0) {
        return null;
    }
    final List<NonThreadSafeCircularBuffer> circularBuffers = new ArrayList<>();
    for (final byte[] stopper : stoppers) {
        circularBuffers.add(new NonThreadSafeCircularBuffer(stopper));
    }
    long bytesRead = 0;
    while (true) {
        final int next = in.read();
        if (next == -1) {
            return null;
        } else if (maxBytes > 0 && ++bytesRead >= maxBytes) {
            throw new BytePatternNotFoundException("Did not encounter any byte pattern that was expected; data does not appear to be in the expected format");
        }
        out.write(next);
        for (final NonThreadSafeCircularBuffer circ : circularBuffers) {
            if (circ.addAndCompare((byte) next)) {
                return circ.getByteArray();
            }
        }
    }
}