Examples with StandardDataPacket org.apache.nifi.remote.util.StandardDataPacket used on opensource projects

Example 1 with StandardDataPacket

use of org.apache.nifi.remote.util.StandardDataPacket in project nifi by apache.

the class TestStandardRemoteGroupPort method testReceiveRaw.

@Test
public void testReceiveRaw() throws Exception {
    setupMock(SiteToSiteTransportProtocol.RAW, TransferDirection.RECEIVE);
    setupMockProcessSession();
    final String peerUrl = "nifi://node1.example.com:9090";
    final PeerDescription peerDescription = new PeerDescription("node1.example.com", 9090, true);
    try (final SocketChannel socketChannel = SocketChannel.open()) {
        final CommunicationsSession commsSession = new SocketChannelCommunicationsSession(socketChannel);
        commsSession.setUserDn("nifi.node1.example.com");
        final Peer peer = new Peer(peerDescription, commsSession, peerUrl, REMOTE_CLUSTER_URL);
        doReturn(peer).when(transaction).getCommunicant();
        final String sourceFlowFileUuid = "flowfile-uuid";
        final Map<String, String> attributes = new HashMap<>();
        attributes.put(CoreAttributes.UUID.key(), sourceFlowFileUuid);
        final byte[] dataPacketContents = "DataPacket Contents".getBytes();
        final ByteArrayInputStream dataPacketInputStream = new ByteArrayInputStream(dataPacketContents);
        final DataPacket dataPacket = new StandardDataPacket(attributes, dataPacketInputStream, dataPacketContents.length);
        // Return null when it gets called second time.
        doReturn(dataPacket).doReturn(null).when(this.transaction).receive();
        port.onTrigger(processContext, processSession);
        // Assert provenance.
        final List<ProvenanceEventRecord> provenanceEvents = sessionState.getProvenanceEvents();
        assertEquals(1, provenanceEvents.size());
        final ProvenanceEventRecord provenanceEvent = provenanceEvents.get(0);
        assertEquals(ProvenanceEventType.RECEIVE, provenanceEvent.getEventType());
        assertEquals(peerUrl + "/" + sourceFlowFileUuid, provenanceEvent.getTransitUri());
        assertEquals("Remote DN=nifi.node1.example.com", provenanceEvent.getDetails());
        // Assert received flow files.
        processSession.assertAllFlowFilesTransferred(Relationship.ANONYMOUS);
        final List<MockFlowFile> flowFiles = processSession.getFlowFilesForRelationship(Relationship.ANONYMOUS);
        assertEquals(1, flowFiles.size());
        final MockFlowFile flowFile = flowFiles.get(0);
        flowFile.assertAttributeEquals(SiteToSiteAttributes.S2S_HOST.key(), peer.getHost());
        flowFile.assertAttributeEquals(SiteToSiteAttributes.S2S_ADDRESS.key(), peer.getHost() + ":" + peer.getPort());
    }
}

Example 2 with StandardDataPacket

use of org.apache.nifi.remote.util.StandardDataPacket in project nifi by apache.

the class TestHttpFlowFileServerProtocol method createClientDataPacket.

private DataPacket createClientDataPacket() {
    final String contents = "Content from client.";
    final byte[] bytes = contents.getBytes();
    final InputStream in = new ByteArrayInputStream(bytes);
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(CoreAttributes.UUID.key(), "client-flow-file-uuid");
    attributes.put("client-attr-1", "client-attr-1-value");
    attributes.put("client-attr-2", "client-attr-2-value");
    return new StandardDataPacket(attributes, in, bytes.length);
}

Example 3 with StandardDataPacket

use of org.apache.nifi.remote.util.StandardDataPacket in project nifi by apache.

the class StandardFlowFileCodec method decode.

@Override
public DataPacket decode(final InputStream stream) throws IOException, ProtocolException {
    final DataInputStream in = new DataInputStream(stream);
    final int numAttributes;
    try {
        numAttributes = in.readInt();
    } catch (final EOFException e) {
        // we're out of data.
        return null;
    }
    // generally result in an OutOfMemoryError.
    if (numAttributes > MAX_NUM_ATTRIBUTES) {
        throw new ProtocolException("FlowFile exceeds maximum number of attributes with a total of " + numAttributes);
    }
    final Map<String, String> attributes = new HashMap<>(numAttributes);
    for (int i = 0; i < numAttributes; i++) {
        final String attrName = readString(in);
        final String attrValue = readString(in);
        attributes.put(attrName, attrValue);
    }
    final long numBytes = in.readLong();
    return new StandardDataPacket(attributes, stream, numBytes);
}

Example 4 with StandardDataPacket

use of org.apache.nifi.remote.util.StandardDataPacket in project nifi by apache.

the class StandardRemoteGroupPort method transferFlowFiles.

private int transferFlowFiles(final Transaction transaction, final ProcessContext context, final ProcessSession session, final FlowFile firstFlowFile) throws IOException, ProtocolException {
    FlowFile flowFile = firstFlowFile;
    try {
        final String userDn = transaction.getCommunicant().getDistinguishedName();
        final long startSendingNanos = System.nanoTime();
        final StopWatch stopWatch = new StopWatch(true);
        long bytesSent = 0L;
        final SiteToSiteClientConfig siteToSiteClientConfig = getSiteToSiteClient().getConfig();
        final long maxBatchBytes = siteToSiteClientConfig.getPreferredBatchSize();
        final int maxBatchCount = siteToSiteClientConfig.getPreferredBatchCount();
        final long preferredBatchDuration = siteToSiteClientConfig.getPreferredBatchDuration(TimeUnit.NANOSECONDS);
        final long maxBatchDuration = preferredBatchDuration > 0 ? preferredBatchDuration : BATCH_SEND_NANOS;
        final Set<FlowFile> flowFilesSent = new HashSet<>();
        boolean continueTransaction = true;
        while (continueTransaction) {
            final long startNanos = System.nanoTime();
            // call codec.encode within a session callback so that we have the InputStream to read the FlowFile
            final FlowFile toWrap = flowFile;
            session.read(flowFile, new InputStreamCallback() {

                @Override
                public void process(final InputStream in) throws IOException {
                    final DataPacket dataPacket = new StandardDataPacket(toWrap.getAttributes(), in, toWrap.getSize());
                    transaction.send(dataPacket);
                }
            });
            final long transferNanos = System.nanoTime() - startNanos;
            final long transferMillis = TimeUnit.MILLISECONDS.convert(transferNanos, TimeUnit.NANOSECONDS);
            flowFilesSent.add(flowFile);
            bytesSent += flowFile.getSize();
            logger.debug("{} Sent {} to {}", this, flowFile, transaction.getCommunicant().getUrl());
            final String transitUri = transaction.getCommunicant().createTransitUri(flowFile.getAttribute(CoreAttributes.UUID.key()));
            session.getProvenanceReporter().send(flowFile, transitUri, "Remote DN=" + userDn, transferMillis, false);
            session.remove(flowFile);
            final long sendingNanos = System.nanoTime() - startSendingNanos;
            if (maxBatchCount > 0 && flowFilesSent.size() >= maxBatchCount) {
                flowFile = null;
            } else if (maxBatchBytes > 0 && bytesSent >= maxBatchBytes) {
                flowFile = null;
            } else if (sendingNanos >= maxBatchDuration) {
                flowFile = null;
            } else {
                flowFile = session.get();
            }
            continueTransaction = (flowFile != null);
        }
        transaction.confirm();
        // consume input stream entirely, ignoring its contents. If we
        // don't do this, the Connection will not be returned to the pool
        stopWatch.stop();
        final String uploadDataRate = stopWatch.calculateDataRate(bytesSent);
        final long uploadMillis = stopWatch.getDuration(TimeUnit.MILLISECONDS);
        final String dataSize = FormatUtils.formatDataSize(bytesSent);
        transaction.complete();
        session.commit();
        final String flowFileDescription = (flowFilesSent.size() < 20) ? flowFilesSent.toString() : flowFilesSent.size() + " FlowFiles";
        logger.info("{} Successfully sent {} ({}) to {} in {} milliseconds at a rate of {}", new Object[] { this, flowFileDescription, dataSize, transaction.getCommunicant().getUrl(), uploadMillis, uploadDataRate });
        return flowFilesSent.size();
    } catch (final Exception e) {
        session.rollback();
        throw e;
    }
}

Example 5 with StandardDataPacket

use of org.apache.nifi.remote.util.StandardDataPacket in project nifi by apache.

the class AbstractFlowFileServerProtocol method transferFlowFiles.

@Override
public int transferFlowFiles(final Peer peer, final ProcessContext context, final ProcessSession session, final FlowFileCodec codec) throws IOException, ProtocolException {
    if (!handshakeCompleted) {
        throw new IllegalStateException("Handshake has not been completed");
    }
    if (shutdown) {
        throw new IllegalStateException("Protocol is shutdown");
    }
    logger.debug("{} Sending FlowFiles to {}", this, peer);
    final CommunicationsSession commsSession = peer.getCommunicationsSession();
    String remoteDn = commsSession.getUserDn();
    if (remoteDn == null) {
        remoteDn = "none";
    }
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        // we have no data to send. Notify the peer.
        logger.debug("{} No data to send to {}", this, peer);
        writeTransactionResponse(true, ResponseCode.NO_MORE_DATA, commsSession);
        return 0;
    }
    // we have data to send.
    logger.debug("{} Data is available to send to {}", this, peer);
    writeTransactionResponse(true, ResponseCode.MORE_DATA, commsSession);
    final StopWatch stopWatch = new StopWatch(true);
    long bytesSent = 0L;
    final Set<FlowFile> flowFilesSent = new HashSet<>();
    final CRC32 crc = new CRC32();
    // send data until we reach some batch size
    boolean continueTransaction = true;
    final long startNanos = System.nanoTime();
    String calculatedCRC = "";
    OutputStream os = new DataOutputStream(commsSession.getOutput().getOutputStream());
    while (continueTransaction) {
        final boolean useGzip = handshakeProperties.isUseGzip();
        final OutputStream flowFileOutputStream = useGzip ? new CompressionOutputStream(os) : os;
        logger.debug("{} Sending {} to {}", new Object[] { this, flowFile, peer });
        final CheckedOutputStream checkedOutputStream = new CheckedOutputStream(flowFileOutputStream, crc);
        final StopWatch transferWatch = new StopWatch(true);
        final FlowFile toSend = flowFile;
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(final InputStream in) throws IOException {
                final DataPacket dataPacket = new StandardDataPacket(toSend.getAttributes(), in, toSend.getSize());
                codec.encode(dataPacket, checkedOutputStream);
            }
        });
        final long transmissionMillis = transferWatch.getElapsed(TimeUnit.MILLISECONDS);
        // (CompressionOutputStream will not close the underlying stream when it's closed)
        if (useGzip) {
            checkedOutputStream.close();
        }
        flowFilesSent.add(flowFile);
        bytesSent += flowFile.getSize();
        final String transitUri = createTransitUri(peer, flowFile.getAttribute(CoreAttributes.UUID.key()));
        session.getProvenanceReporter().send(flowFile, transitUri, "Remote Host=" + peer.getHost() + ", Remote DN=" + remoteDn, transmissionMillis, false);
        session.remove(flowFile);
        // determine if we should check for more data on queue.
        final long sendingNanos = System.nanoTime() - startNanos;
        boolean poll = true;
        double batchDurationNanos = handshakeProperties.getBatchDurationNanos();
        if (sendingNanos >= batchDurationNanos && batchDurationNanos > 0L) {
            poll = false;
        }
        double batchBytes = handshakeProperties.getBatchBytes();
        if (bytesSent >= batchBytes && batchBytes > 0L) {
            poll = false;
        }
        double batchCount = handshakeProperties.getBatchCount();
        if (flowFilesSent.size() >= batchCount && batchCount > 0) {
            poll = false;
        }
        if (batchDurationNanos == 0 && batchBytes == 0 && batchCount == 0) {
            poll = (sendingNanos < DEFAULT_BATCH_NANOS);
        }
        if (poll) {
            // we've not elapsed the requested sending duration, so get more data.
            flowFile = session.get();
        } else {
            flowFile = null;
        }
        continueTransaction = (flowFile != null);
        if (continueTransaction) {
            logger.debug("{} Sending ContinueTransaction indicator to {}", this, peer);
            writeTransactionResponse(true, ResponseCode.CONTINUE_TRANSACTION, commsSession);
        } else {
            logger.debug("{} Sending FinishTransaction indicator to {}", this, peer);
            writeTransactionResponse(true, ResponseCode.FINISH_TRANSACTION, commsSession);
            calculatedCRC = String.valueOf(checkedOutputStream.getChecksum().getValue());
        }
    }
    FlowFileTransaction transaction = new FlowFileTransaction(session, context, stopWatch, bytesSent, flowFilesSent, calculatedCRC);
    return commitTransferTransaction(peer, transaction);
}