Examples with SocketChannelRecordReader org.apache.nifi.record.listen.SocketChannelRecordReader used on opensource projects

Example 1 with SocketChannelRecordReader

use of org.apache.nifi.record.listen.SocketChannelRecordReader in project nifi by apache.

the class ListenTCPRecord method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final SocketChannelRecordReader socketRecordReader = pollForSocketRecordReader();
    if (socketRecordReader == null) {
        return;
    }
    if (socketRecordReader.isClosed()) {
        getLogger().warn("Unable to read records from {}, socket already closed", new Object[] { getRemoteAddress(socketRecordReader) });
        // still need to call close so the overall count is decremented
        IOUtils.closeQuietly(socketRecordReader);
        return;
    }
    final int recordBatchSize = context.getProperty(RECORD_BATCH_SIZE).asInteger();
    final String readerErrorHandling = context.getProperty(READER_ERROR_HANDLING_STRATEGY).getValue();
    final RecordSetWriterFactory recordSetWriterFactory = context.getProperty(RECORD_WRITER).asControllerService(RecordSetWriterFactory.class);
    // synchronize to ensure there are no stale values in the underlying SocketChannel
    synchronized (socketRecordReader) {
        FlowFile flowFile = session.create();
        try {
            // lazily creating the record reader here b/c we need a flow file, eventually shouldn't have to do this
            RecordReader recordReader = socketRecordReader.getRecordReader();
            if (recordReader == null) {
                recordReader = socketRecordReader.createRecordReader(flowFile, getLogger());
            }
            Record record;
            try {
                record = recordReader.nextRecord();
            } catch (final Exception e) {
                boolean timeout = false;
                // some of the underlying record libraries wrap the real exception in RuntimeException, so check each
                // throwable (starting with the current one) to see if its a SocketTimeoutException
                Throwable cause = e;
                while (cause != null) {
                    if (cause instanceof SocketTimeoutException) {
                        timeout = true;
                        break;
                    }
                    cause = cause.getCause();
                }
                if (timeout) {
                    getLogger().debug("Timeout reading records, will try again later", e);
                    socketReaders.offer(socketRecordReader);
                    session.remove(flowFile);
                    return;
                } else {
                    throw e;
                }
            }
            if (record == null) {
                getLogger().debug("No records available from {}, closing connection", new Object[] { getRemoteAddress(socketRecordReader) });
                IOUtils.closeQuietly(socketRecordReader);
                session.remove(flowFile);
                return;
            }
            String mimeType = null;
            WriteResult writeResult = null;
            final RecordSchema recordSchema = recordSetWriterFactory.getSchema(Collections.EMPTY_MAP, record.getSchema());
            try (final OutputStream out = session.write(flowFile);
                final RecordSetWriter recordWriter = recordSetWriterFactory.createWriter(getLogger(), recordSchema, out)) {
                // start the record set and write the first record from above
                recordWriter.beginRecordSet();
                writeResult = recordWriter.write(record);
                while (record != null && writeResult.getRecordCount() < recordBatchSize) {
                    // if keeping then null out the record to break out of the loop, which will transfer what we have and close the connection
                    try {
                        record = recordReader.nextRecord();
                    } catch (final SocketTimeoutException ste) {
                        getLogger().debug("Timeout reading records, will try again later", ste);
                        break;
                    } catch (final Exception e) {
                        if (ERROR_HANDLING_DISCARD.getValue().equals(readerErrorHandling)) {
                            throw e;
                        } else {
                            record = null;
                        }
                    }
                    if (record != null) {
                        writeResult = recordWriter.write(record);
                    }
                }
                writeResult = recordWriter.finishRecordSet();
                recordWriter.flush();
                mimeType = recordWriter.getMimeType();
            }
            // if we didn't write any records then we need to remove the flow file
            if (writeResult.getRecordCount() <= 0) {
                getLogger().debug("Removing flow file, no records were written");
                session.remove(flowFile);
            } else {
                final String sender = getRemoteAddress(socketRecordReader);
                final Map<String, String> attributes = new HashMap<>(writeResult.getAttributes());
                attributes.put(CoreAttributes.MIME_TYPE.key(), mimeType);
                attributes.put("tcp.sender", sender);
                attributes.put("tcp.port", String.valueOf(port));
                attributes.put("record.count", String.valueOf(writeResult.getRecordCount()));
                flowFile = session.putAllAttributes(flowFile, attributes);
                final String senderHost = sender.startsWith("/") && sender.length() > 1 ? sender.substring(1) : sender;
                final String transitUri = new StringBuilder().append("tcp").append("://").append(senderHost).append(":").append(port).toString();
                session.getProvenanceReporter().receive(flowFile, transitUri);
                session.transfer(flowFile, REL_SUCCESS);
            }
            getLogger().debug("Re-queuing connection for further processing...");
            socketReaders.offer(socketRecordReader);
        } catch (Exception e) {
            getLogger().error("Error processing records: " + e.getMessage(), e);
            IOUtils.closeQuietly(socketRecordReader);
            session.remove(flowFile);
            return;
        }
    }
}

Example 2 with SocketChannelRecordReader

use of org.apache.nifi.record.listen.SocketChannelRecordReader in project nifi by apache.

the class ListenTCPRecord method onStopped.

@OnStopped
public void onStopped() {
    if (dispatcher != null) {
        dispatcher.close();
        dispatcher = null;
    }
    SocketChannelRecordReader socketRecordReader;
    while ((socketRecordReader = socketReaders.poll()) != null) {
        IOUtils.closeQuietly(socketRecordReader.getRecordReader());
    }
}