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

Example 16 with CommunicationsSession

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

the class SocketClientProtocol method shutdown.

@Override
public void shutdown(final Peer peer) throws IOException {
    readyForFileTransfer = false;
    final CommunicationsSession commsSession = peer.getCommunicationsSession();
    final DataOutputStream dos = new DataOutputStream(commsSession.getOutput().getOutputStream());
    logger.debug("{} Shutting down with {}", this, peer);
    // Indicate that we would like to have some data
    RequestType.SHUTDOWN.writeRequestType(dos);
    dos.flush();
}

Example 17 with CommunicationsSession

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

the class SocketClientProtocol method negotiateCodec.

@Override
public FlowFileCodec negotiateCodec(final Peer peer) throws IOException, ProtocolException {
    if (!handshakeComplete) {
        throw new IllegalStateException("Handshake has not been performed");
    }
    logger.debug("{} Negotiating Codec with {}", this, peer);
    final CommunicationsSession commsSession = peer.getCommunicationsSession();
    final DataInputStream dis = new DataInputStream(commsSession.getInput().getInputStream());
    final DataOutputStream dos = new DataOutputStream(commsSession.getOutput().getOutputStream());
    RequestType.NEGOTIATE_FLOWFILE_CODEC.writeRequestType(dos);
    FlowFileCodec codec = new StandardFlowFileCodec();
    try {
        codec = (FlowFileCodec) RemoteResourceInitiator.initiateResourceNegotiation(codec, dis, dos);
    } catch (HandshakeException e) {
        throw new ProtocolException(e.toString());
    }
    logger.debug("{} negotiated FlowFileCodec {} with {}", new Object[] { this, codec, commsSession });
    return codec;
}

Example 18 with CommunicationsSession

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

the class SiteToSiteRestApiClient method openConnectionForSend.

public void openConnectionForSend(final String transactionUrl, final Peer peer) throws IOException {
    final CommunicationsSession commSession = peer.getCommunicationsSession();
    final String flowFilesPath = transactionUrl + "/flow-files";
    final HttpPost post = createPost(flowFilesPath);
    // Set uri so that it'll be used as transit uri.
    ((HttpCommunicationsSession) peer.getCommunicationsSession()).setDataTransferUrl(post.getURI().toString());
    post.setHeader("Content-Type", "application/octet-stream");
    post.setHeader("Accept", "text/plain");
    post.setHeader(HttpHeaders.PROTOCOL_VERSION, String.valueOf(transportProtocolVersionNegotiator.getVersion()));
    setHandshakeProperties(post);
    final CountDownLatch initConnectionLatch = new CountDownLatch(1);
    final URI requestUri = post.getURI();
    final PipedOutputStream outputStream = new PipedOutputStream();
    final PipedInputStream inputStream = new PipedInputStream(outputStream, DATA_PACKET_CHANNEL_READ_BUFFER_SIZE);
    final ReadableByteChannel dataPacketChannel = Channels.newChannel(inputStream);
    final HttpAsyncRequestProducer asyncRequestProducer = new HttpAsyncRequestProducer() {

        private final ByteBuffer buffer = ByteBuffer.allocate(DATA_PACKET_CHANNEL_READ_BUFFER_SIZE);

        private int totalRead = 0;

        private int totalProduced = 0;

        private boolean requestHasBeenReset = false;

        @Override
        public HttpHost getTarget() {
            return URIUtils.extractHost(requestUri);
        }

        @Override
        public HttpRequest generateRequest() throws IOException, HttpException {
            // Pass the output stream so that Site-to-Site client thread can send
            // data packet through this connection.
            logger.debug("sending data to {} has started...", flowFilesPath);
            ((HttpOutput) commSession.getOutput()).setOutputStream(outputStream);
            initConnectionLatch.countDown();
            final BasicHttpEntity entity = new BasicHttpEntity();
            entity.setChunked(true);
            entity.setContentType("application/octet-stream");
            post.setEntity(entity);
            return post;
        }

        private final AtomicBoolean bufferHasRemainingData = new AtomicBoolean(false);

        /**
         * If the proxy server requires authentication, the same POST request has to be sent again.
         * The first request will result 407, then the next one will be sent with auth headers and actual data.
         * This method produces a content only when it's need to be sent, to avoid producing the flow-file contents twice.
         * Whether we need to wait auth is determined heuristically by the previous POST request which creates transaction.
         * See {@link SiteToSiteRestApiClient#initiateTransactionForSend(HttpPost)} for further detail.
         */
        @Override
        public void produceContent(final ContentEncoder encoder, final IOControl ioControl) throws IOException {
            if (shouldCheckProxyAuth() && proxyAuthRequiresResend.get() && !requestHasBeenReset) {
                logger.debug("Need authentication with proxy server. Postpone producing content.");
                encoder.complete();
                return;
            }
            if (bufferHasRemainingData.get()) {
                // If there's remaining buffer last time, send it first.
                writeBuffer(encoder);
                if (bufferHasRemainingData.get()) {
                    return;
                }
            }
            int read;
            // or corresponding outputStream is closed.
            if ((read = dataPacketChannel.read(buffer)) > -1) {
                logger.trace("Read {} bytes from dataPacketChannel. {}", read, flowFilesPath);
                totalRead += read;
                buffer.flip();
                writeBuffer(encoder);
            } else {
                final long totalWritten = commSession.getOutput().getBytesWritten();
                logger.debug("sending data to {} has reached to its end. produced {} bytes by reading {} bytes from channel. {} bytes written in this transaction.", flowFilesPath, totalProduced, totalRead, totalWritten);
                if (totalRead != totalWritten || totalProduced != totalWritten) {
                    final String msg = "Sending data to %s has reached to its end, but produced : read : wrote byte sizes (%d : %d : %d) were not equal. Something went wrong.";
                    throw new RuntimeException(String.format(msg, flowFilesPath, totalProduced, totalRead, totalWritten));
                }
                transferDataLatch.countDown();
                encoder.complete();
                dataPacketChannel.close();
            }
        }

        private void writeBuffer(ContentEncoder encoder) throws IOException {
            while (buffer.hasRemaining()) {
                final int written = encoder.write(buffer);
                logger.trace("written {} bytes to encoder.", written);
                if (written == 0) {
                    logger.trace("Buffer still has remaining. {}", buffer);
                    bufferHasRemainingData.set(true);
                    return;
                }
                totalProduced += written;
            }
            bufferHasRemainingData.set(false);
            buffer.clear();
        }

        @Override
        public void requestCompleted(final HttpContext context) {
            logger.debug("Sending data to {} completed.", flowFilesPath);
            debugProxyAuthState(context);
        }

        @Override
        public void failed(final Exception ex) {
            final String msg = String.format("Failed to send data to %s due to %s", flowFilesPath, ex.toString());
            logger.error(msg, ex);
            eventReporter.reportEvent(Severity.WARNING, EVENT_CATEGORY, msg);
        }

        @Override
        public boolean isRepeatable() {
            // In order to pass authentication, request has to be repeatable.
            return true;
        }

        @Override
        public void resetRequest() throws IOException {
            logger.debug("Sending data request to {} has been reset...", flowFilesPath);
            requestHasBeenReset = true;
        }

        @Override
        public void close() throws IOException {
            logger.debug("Closing sending data request to {}", flowFilesPath);
            closeSilently(outputStream);
            closeSilently(dataPacketChannel);
            stopExtendingTtl();
        }
    };
    postResult = getHttpAsyncClient().execute(asyncRequestProducer, new BasicAsyncResponseConsumer(), null);
    try {
        // Need to wait the post request actually started so that we can write to its output stream.
        if (!initConnectionLatch.await(connectTimeoutMillis, TimeUnit.MILLISECONDS)) {
            throw new IOException("Awaiting initConnectionLatch has been timeout.");
        }
        // Started.
        transferDataLatch = new CountDownLatch(1);
        startExtendingTtl(transactionUrl, dataPacketChannel, null);
    } catch (final InterruptedException e) {
        throw new IOException("Awaiting initConnectionLatch has been interrupted.", e);
    }
}

Example 19 with CommunicationsSession

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

the class AbstractTransaction method confirm.

@Override
public final void confirm() throws IOException {
    try {
        try {
            if (state == TransactionState.TRANSACTION_STARTED && !dataAvailable && direction == TransferDirection.RECEIVE) {
                // client requested to receive data but no data available. no need to confirm.
                state = TransactionState.TRANSACTION_CONFIRMED;
                return;
            }
            if (state != TransactionState.DATA_EXCHANGED) {
                throw new IllegalStateException("Cannot confirm Transaction because state is " + state + "; Transaction can only be confirmed when state is " + TransactionState.DATA_EXCHANGED);
            }
            final CommunicationsSession commsSession = peer.getCommunicationsSession();
            if (direction == TransferDirection.RECEIVE) {
                if (dataAvailable) {
                    throw new IllegalStateException("Cannot complete transaction because the sender has already sent more data than client has consumed.");
                }
                // we received a FINISH_TRANSACTION indicator. Send back a CONFIRM_TRANSACTION message
                // to peer so that we can verify that the connection is still open. This is a two-phase commit,
                // which helps to prevent the chances of data duplication. Without doing this, we may commit the
                // session and then when we send the response back to the peer, the peer may have timed out and may not
                // be listening. As a result, it will re-send the data. By doing this two-phase commit, we narrow the
                // Critical Section involved in this transaction so that rather than the Critical Section being the
                // time window involved in the entire transaction, it is reduced to a simple round-trip conversation.
                logger.trace("{} Sending CONFIRM_TRANSACTION Response Code to {}", this, peer);
                final String calculatedCRC = String.valueOf(crc.getValue());
                writeTransactionResponse(ResponseCode.CONFIRM_TRANSACTION, calculatedCRC);
                final Response confirmTransactionResponse;
                try {
                    confirmTransactionResponse = readTransactionResponse();
                } catch (final IOException ioe) {
                    logger.error("Failed to receive response code from {} when expecting confirmation of transaction", peer);
                    if (eventReporter != null) {
                        eventReporter.reportEvent(Severity.ERROR, "Site-to-Site", "Failed to receive response code from " + peer + " when expecting confirmation of transaction");
                    }
                    throw ioe;
                }
                logger.trace("{} Received {} from {}", this, confirmTransactionResponse, peer);
                switch(confirmTransactionResponse.getCode()) {
                    case CONFIRM_TRANSACTION:
                        break;
                    case BAD_CHECKSUM:
                        throw new IOException(this + " Received a BadChecksum response from peer " + peer);
                    default:
                        throw new ProtocolException(this + " Received unexpected Response from peer " + peer + " : " + confirmTransactionResponse + "; expected 'Confirm Transaction' Response Code");
                }
                state = TransactionState.TRANSACTION_CONFIRMED;
            } else {
                logger.debug("{} Sent FINISH_TRANSACTION indicator to {}", this, peer);
                writeTransactionResponse(ResponseCode.FINISH_TRANSACTION);
                final String calculatedCRC = String.valueOf(crc.getValue());
                // we've sent a FINISH_TRANSACTION. Now we'll wait for the peer to send a 'Confirm Transaction' response
                final Response transactionConfirmationResponse = readTransactionResponse();
                if (transactionConfirmationResponse.getCode() == ResponseCode.CONFIRM_TRANSACTION) {
                    // Confirm checksum and echo back the confirmation.
                    logger.trace("{} Received {} from {}", this, transactionConfirmationResponse, peer);
                    final String receivedCRC = transactionConfirmationResponse.getMessage();
                    // CRC was not used before version 4
                    if (protocolVersion > 3) {
                        if (!receivedCRC.equals(calculatedCRC)) {
                            writeTransactionResponse(ResponseCode.BAD_CHECKSUM);
                            throw new IOException(this + " Sent data to peer " + peer + " but calculated CRC32 Checksum as " + calculatedCRC + " while peer calculated CRC32 Checksum as " + receivedCRC + "; canceling transaction and rolling back session");
                        }
                    }
                    writeTransactionResponse(ResponseCode.CONFIRM_TRANSACTION, "");
                } else {
                    throw new ProtocolException("Expected to receive 'Confirm Transaction' response from peer " + peer + " but received " + transactionConfirmationResponse);
                }
                state = TransactionState.TRANSACTION_CONFIRMED;
            }
        } catch (final IOException ioe) {
            throw new IOException("Failed to confirm transaction with " + peer + " due to " + ioe, ioe);
        }
    } catch (final Exception e) {
        error();
        throw e;
    }
}