Examples with DataPacket org.apache.nifi.remote.protocol.DataPacket used on opensource projects

Example 21 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestStandardRemoteGroupPort method testSendBatch.

/**
 * Generate flow files to be sent, and execute port's onTrigger method.
 * Finally, this method verifies whether packets are sent as expected.
 * @param expectedNumberOfPackets Specify how many packets should be sent by each transaction.
 *                                E.g. passing {2, 2, 1}, would generate 5 flow files in total.
 *                                Based on the siteToSiteClientConfig batch parameters,
 *                                it's expected to be sent via 3 transactions,
 *                                transaction 0 will send flow file 0 and 1,
 *                                transaction 1 will send flow file 2 and 3,
 *                                and transaction 2 will send flow file 4.
 *                                Each flow file has different content size generated automatically.
 *                                The content size starts with 10, and increases as more flow files are generated.
 *                                E.g. flow file 1 will have 10 bytes, flow file 2 has 11 bytes, f3 has 12 and so on.
 */
private void testSendBatch(final int[] expectedNumberOfPackets) throws Exception {
    setupMockProcessSession();
    final String peerUrl = "http://node1.example.com:8080/nifi";
    final PeerDescription peerDescription = new PeerDescription("node1.example.com", 8080, false);
    final HttpCommunicationsSession commsSession = new HttpCommunicationsSession();
    final Peer peer = new Peer(peerDescription, commsSession, peerUrl, REMOTE_CLUSTER_URL);
    final String flowFileEndpointUri = "http://node1.example.com:8080/nifi-api/output-ports/port-id/transactions/transaction-id/flow-files";
    doReturn(peer).when(transaction).getCommunicant();
    commsSession.setDataTransferUrl(flowFileEndpointUri);
    // Capture packets being sent to the remote peer
    final AtomicInteger totalPacketsSent = new AtomicInteger(0);
    final List<List<DataPacket>> sentPackets = new ArrayList<>(expectedNumberOfPackets.length);
    final List<DataPacket> sentPacketsPerTransaction = new ArrayList<>();
    doAnswer(invocation -> {
        sentPacketsPerTransaction.add((DataPacket) invocation.getArguments()[0]);
        totalPacketsSent.incrementAndGet();
        return null;
    }).when(transaction).send(any(DataPacket.class));
    doAnswer(invocation -> {
        sentPackets.add(new ArrayList<>(sentPacketsPerTransaction));
        sentPacketsPerTransaction.clear();
        return null;
    }).when(transaction).confirm();
    // Execute onTrigger while offering new flow files.
    final List<MockFlowFile> flowFiles = new ArrayList<>();
    for (int i = 0; i < expectedNumberOfPackets.length; i++) {
        int numOfPackets = expectedNumberOfPackets[i];
        int startF = flowFiles.size();
        int endF = startF + numOfPackets;
        IntStream.range(startF, endF).forEach(f -> {
            final StringBuilder flowFileContents = new StringBuilder("0123456789");
            for (int c = 0; c < f; c++) {
                flowFileContents.append(c);
            }
            final byte[] bytes = flowFileContents.toString().getBytes();
            final MockFlowFile flowFile = spy(processSession.createFlowFile(bytes));
            when(flowFile.getSize()).then(invocation -> {
                // For testSendBatchByDuration
                Thread.sleep(1);
                return bytes.length;
            });
            sessionState.getFlowFileQueue().offer(flowFile);
            flowFiles.add(flowFile);
        });
        port.onTrigger(processContext, processSession);
    }
    // Verify transactions, sent packets, and provenance events.
    assertEquals(flowFiles.size(), totalPacketsSent.get());
    assertEquals("The number of transactions should match as expected.", expectedNumberOfPackets.length, sentPackets.size());
    final List<ProvenanceEventRecord> provenanceEvents = sessionState.getProvenanceEvents();
    assertEquals(flowFiles.size(), provenanceEvents.size());
    int f = 0;
    for (int i = 0; i < expectedNumberOfPackets.length; i++) {
        final List<DataPacket> dataPackets = sentPackets.get(i);
        assertEquals(expectedNumberOfPackets[i], dataPackets.size());
        for (int p = 0; p < dataPackets.size(); p++) {
            final FlowFile flowFile = flowFiles.get(f);
            // Assert sent packet
            final DataPacket dataPacket = dataPackets.get(p);
            assertEquals(flowFile.getSize(), dataPacket.getSize());
            // Assert provenance event
            final ProvenanceEventRecord provenanceEvent = provenanceEvents.get(f);
            assertEquals(ProvenanceEventType.SEND, provenanceEvent.getEventType());
            assertEquals(flowFileEndpointUri, provenanceEvent.getTransitUri());
            f++;
        }
    }
}

Example 22 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestHttpFlowFileServerProtocol method testReceiveOneFile.

@Test
public void testReceiveOneFile() throws Exception {
    final HttpFlowFileServerProtocol serverProtocol = getDefaultHttpFlowFileServerProtocol();
    final String transactionId = "testReceiveOneFile";
    final String endpointUri = "https://remote-host:8443/nifi-api/input-ports/port-id/transactions/" + transactionId + "/flow-files";
    final Peer peer = getDefaultPeer(transactionId);
    final HttpServerCommunicationsSession commsSession = (HttpServerCommunicationsSession) peer.getCommunicationsSession();
    commsSession.putHandshakeParam(HandshakeProperty.BATCH_COUNT, "1");
    commsSession.setUserDn("unit-test");
    commsSession.setDataTransferUrl(endpointUri);
    final DataPacket dataPacket = createClientDataPacket();
    receiveFlowFiles(serverProtocol, transactionId, peer, dataPacket);
    // Commit transaction
    commsSession.setResponseCode(ResponseCode.CONFIRM_TRANSACTION);
    final int flowFileReceived = serverProtocol.commitReceiveTransaction(peer);
    assertEquals(1, flowFileReceived);
    // Assert provenance.
    final List<ProvenanceEventRecord> provenanceEvents = sessionState.getProvenanceEvents();
    assertEquals(1, provenanceEvents.size());
    final ProvenanceEventRecord provenanceEvent = provenanceEvents.get(0);
    assertEquals(ProvenanceEventType.RECEIVE, provenanceEvent.getEventType());
    assertEquals(endpointUri, provenanceEvent.getTransitUri());
    assertEquals("Remote Host=peer-host, Remote DN=unit-test", 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());
    flowFile.assertAttributeEquals("client-attr-1", "client-attr-1-value");
    flowFile.assertAttributeEquals("client-attr-2", "client-attr-2-value");
}

Example 23 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class SocketClient method createTransaction.

@Override
public Transaction createTransaction(final TransferDirection direction) throws IOException {
    if (closed) {
        throw new IllegalStateException("Client is closed");
    }
    final String portId = getPortIdentifier(direction);
    if (portId == null) {
        throw new IOException("Could not find Port with name '" + portName + "' for remote NiFi instance");
    }
    final EndpointConnection connectionState = pool.getEndpointConnection(direction, getConfig());
    if (connectionState == null) {
        return null;
    }
    final Transaction transaction;
    try {
        transaction = connectionState.getSocketClientProtocol().startTransaction(connectionState.getPeer(), connectionState.getCodec(), direction);
    } catch (final Throwable t) {
        pool.terminate(connectionState);
        throw new IOException("Unable to create Transaction to communicate with " + connectionState.getPeer(), t);
    }
    // Wrap the transaction in a new one that will return the EndpointConnectionState back to the pool whenever
    // the transaction is either completed or canceled.
    final AtomicReference<EndpointConnection> connectionStateRef = new AtomicReference<>(connectionState);
    return new Transaction() {

        @Override
        public void confirm() throws IOException {
            transaction.confirm();
        }

        @Override
        public TransactionCompletion complete() throws IOException {
            try {
                return transaction.complete();
            } finally {
                final EndpointConnection state = connectionStateRef.get();
                if (state != null) {
                    pool.offer(connectionState);
                    connectionStateRef.set(null);
                }
            }
        }

        @Override
        public void cancel(final String explanation) throws IOException {
            try {
                transaction.cancel(explanation);
            } finally {
                final EndpointConnection state = connectionStateRef.get();
                if (state != null) {
                    pool.terminate(connectionState);
                    connectionStateRef.set(null);
                }
            }
        }

        @Override
        public void error() {
            try {
                transaction.error();
            } finally {
                final EndpointConnection state = connectionStateRef.get();
                if (state != null) {
                    pool.terminate(connectionState);
                    connectionStateRef.set(null);
                }
            }
        }

        @Override
        public void send(final DataPacket dataPacket) throws IOException {
            transaction.send(dataPacket);
        }

        @Override
        public void send(final byte[] content, final Map<String, String> attributes) throws IOException {
            transaction.send(content, attributes);
        }

        @Override
        public DataPacket receive() throws IOException {
            return transaction.receive();
        }

        @Override
        public TransactionState getState() throws IOException {
            return transaction.getState();
        }

        @Override
        public Communicant getCommunicant() {
            return transaction.getCommunicant();
        }
    };
}

Example 24 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestHttpClient method testReceiveTimeoutAfterDataExchange.

@Test
public void testReceiveTimeoutAfterDataExchange() throws Exception {
    try (SiteToSiteClient client = getDefaultBuilder().timeout(1, TimeUnit.SECONDS).portName("output-timeout-data-ex").build()) {
        final Transaction transaction = client.createTransaction(TransferDirection.RECEIVE);
        assertNotNull(transaction);
        DataPacket packet = transaction.receive();
        assertNotNull(packet);
        consumeDataPacket(packet);
        try {
            transaction.receive();
            fail();
        } catch (IOException e) {
            logger.info("An exception was thrown as expected.", e);
            assertTrue(e.getCause() instanceof SocketTimeoutException);
        }
        confirmShouldFail(transaction);
        completeShouldFail(transaction);
    }
}

Example 25 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestHttpClient method testSendSlowClientSuccess.

@Test
public void testSendSlowClientSuccess() throws Exception {
    try (SiteToSiteClient client = getDefaultBuilder().idleExpiration(1000, TimeUnit.MILLISECONDS).portName("input-running").build()) {
        final Transaction transaction = client.createTransaction(TransferDirection.SEND);
        assertNotNull(transaction);
        serverChecksum = "3882825556";
        for (int i = 0; i < 3; i++) {
            DataPacket packet = new DataPacketBuilder().contents("Example contents from client.").attr("Client attr 1", "Client attr 1 value").attr("Client attr 2", "Client attr 2 value").build();
            transaction.send(packet);
            long written = ((Peer) transaction.getCommunicant()).getCommunicationsSession().getBytesWritten();
            logger.info("{} bytes have been written.", written);
            Thread.sleep(50);
        }
        transaction.confirm();
        transaction.complete();
    }
}