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

Example 16 with ByteArrayInputStream

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

the class TestLeakyBucketThrottler method testInputStreamInterface.

@Test(timeout = 10000)
public void testInputStreamInterface() throws IOException {
    final byte[] data = new byte[1024 * 1024 * 4];
    // throttle rate at 1 MB/sec
    try (final LeakyBucketStreamThrottler throttler = new LeakyBucketStreamThrottler(1024 * 1024);
        final ByteArrayInputStream bais = new ByteArrayInputStream(data);
        final InputStream throttledIn = throttler.newThrottledInputStream(bais);
        final ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        final byte[] buffer = new byte[4096];
        final long start = System.currentTimeMillis();
        int len;
        while ((len = throttledIn.read(buffer)) > 0) {
            baos.write(buffer, 0, len);
        }
        final long millis = System.currentTimeMillis() - start;
        // should take 4 sec give or take
        assertTrue(millis > 3000);
        assertTrue(millis < 6000);
    }
}

Example 17 with ByteArrayInputStream

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

the class SnippetManager method parseBytes.

public static List<StandardSnippet> parseBytes(final byte[] bytes) {
    final List<StandardSnippet> snippets = new ArrayList<>();
    try (final InputStream rawIn = new ByteArrayInputStream(bytes);
        final DataInputStream in = new DataInputStream(rawIn)) {
        final int length = in.readInt();
        final byte[] buffer = new byte[length];
        StreamUtils.fillBuffer(in, buffer, true);
        final StandardSnippet snippet = StandardSnippetDeserializer.deserialize(new ByteArrayInputStream(buffer));
        snippets.add(snippet);
    } catch (final IOException e) {
        // should never happen because of streams being used
        throw new RuntimeException("Failed to parse bytes", e);
    }
    return snippets;
}

Example 18 with ByteArrayInputStream

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

the class HttpClientTransaction method readTransactionResponse.

@Override
protected Response readTransactionResponse() throws IOException {
    HttpCommunicationsSession commSession = (HttpCommunicationsSession) peer.getCommunicationsSession();
    ByteArrayOutputStream bos = new ByteArrayOutputStream();
    DataOutputStream dos = new DataOutputStream(bos);
    if (TransferDirection.RECEIVE.equals(direction)) {
        switch(state) {
            case TRANSACTION_STARTED:
            case DATA_EXCHANGED:
                logger.debug("{} {} readTransactionResponse. checksum={}", this, peer, commSession.getChecksum());
                if (StringUtils.isEmpty(commSession.getChecksum())) {
                    // We don't know if there's more data to receive, so just continue it.
                    ResponseCode.CONTINUE_TRANSACTION.writeResponse(dos);
                } else {
                    // We got a checksum to send to server.
                    if (TransactionState.TRANSACTION_STARTED.equals(state)) {
                        logger.debug("{} {} There's no transaction to confirm.", this, peer);
                        ResponseCode.CONFIRM_TRANSACTION.writeResponse(dos, "");
                    } else {
                        TransactionResultEntity transactionResult = apiClient.commitReceivingFlowFiles(transactionUrl, ResponseCode.CONFIRM_TRANSACTION, commSession.getChecksum());
                        ResponseCode responseCode = ResponseCode.fromCode(transactionResult.getResponseCode());
                        if (responseCode.containsMessage()) {
                            String message = transactionResult.getMessage();
                            responseCode.writeResponse(dos, message == null ? "" : message);
                        } else {
                            responseCode.writeResponse(dos);
                        }
                    }
                }
                break;
        }
    } else {
        switch(state) {
            case DATA_EXCHANGED:
                // Some flow files have been sent via stream, finish transferring.
                apiClient.finishTransferFlowFiles(commSession);
                ResponseCode.CONFIRM_TRANSACTION.writeResponse(dos, commSession.getChecksum());
                break;
            case TRANSACTION_CONFIRMED:
                TransactionResultEntity resultEntity = apiClient.commitTransferFlowFiles(transactionUrl, ResponseCode.CONFIRM_TRANSACTION);
                ResponseCode responseCode = ResponseCode.fromCode(resultEntity.getResponseCode());
                if (responseCode.containsMessage()) {
                    responseCode.writeResponse(dos, resultEntity.getMessage());
                } else {
                    responseCode.writeResponse(dos);
                }
                break;
        }
    }
    ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray());
    return Response.read(new DataInputStream(bis));
}

Example 19 with ByteArrayInputStream

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

the class SiteToSiteTestUtils method createDataPacket.

public static DataPacket createDataPacket(String contents) {
    try {
        byte[] bytes = contents.getBytes("UTF-8");
        ByteArrayInputStream is = new ByteArrayInputStream(bytes);
        return new StandardDataPacket(new HashMap<>(), is, bytes.length);
    } catch (UnsupportedEncodingException e) {
        throw new RuntimeException(e);
    }
}

Example 20 with ByteArrayInputStream

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

the class TestHttpClientTransaction method testReceiveOneFlowFile.

@Test
public void testReceiveOneFlowFile() throws IOException {
    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doReturn(true).when(apiClient).openConnectionForReceive(eq(transactionUrl), any(Peer.class));
    TransactionResultEntity resultEntity = new TransactionResultEntity();
    resultEntity.setResponseCode(CONFIRM_TRANSACTION.getCode());
    doReturn(resultEntity).when(apiClient).commitReceivingFlowFiles(eq(transactionUrl), eq(CONFIRM_TRANSACTION), eq("3680976076"));
    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    codec.encode(createDataPacket("contents on server 1"), serverResponseBos);
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.RECEIVE, transactionUrl);
    execReceiveOneFlowFile(transaction);
    assertEquals("Client sends nothing as payload to receive flow files.", 0, clientRequest.toByteArray().length);
    verify(apiClient).commitReceivingFlowFiles(transactionUrl, CONFIRM_TRANSACTION, "3680976076");
}