Examples with BasicAsyncResponseConsumer org.apache.http.nio.protocol.BasicAsyncResponseConsumer used on opensource projects

Example 1 with BasicAsyncResponseConsumer

use of org.apache.http.nio.protocol.BasicAsyncResponseConsumer in project nifi by apache.

the class SiteToSiteRestApiClient method initiateTransactionForSend.

/**
 * <p>
 * Initiate a transaction for sending data.
 * </p>
 *
 * <p>
 * If a proxy server requires auth, the proxy server returns 407 response with available auth schema such as basic or digest.
 * Then client has to resend the same request with its credential added.
 * This mechanism is problematic for sending data from NiFi.
 * </p>
 *
 * <p>
 * In order to resend a POST request with auth param,
 * NiFi has to either read flow-file contents to send again, or keep the POST body somewhere.
 * If we store that in memory, it would causes OOM, or storing it on disk slows down performance.
 * Rolling back processing session would be overkill.
 * Reading flow-file contents only when it's ready to send in a streaming way is ideal.
 * </p>
 *
 * <p>
 * Additionally, the way proxy authentication is done is vary among Proxy server software.
 * Some requires 407 and resend cycle for every requests, while others keep a connection between a client and
 * the proxy server, then consecutive requests skip auth steps.
 * The problem is, that how should we behave is only told after sending a request to the proxy.
 * </p>
 *
 * In order to handle above concerns correctly and efficiently, this method do the followings:
 *
 * <ol>
 * <li>Send a GET request to controller resource, to initiate an HttpAsyncClient. The instance will be used for further requests.
 *      This is not required by the Site-to-Site protocol, but it can setup proxy auth state safely.</li>
 * <li>Send a POST request to initiate a transaction. While doing so, it captures how a proxy server works.
 * If 407 and resend cycle occurs here, it implies that we need to do the same thing again when we actually send the data.
 * Because if the proxy keeps using the same connection and doesn't require an auth step, it doesn't do so here.</li>
 * <li>Then this method stores whether the final POST request should wait for the auth step.
 * So that {@link #openConnectionForSend} can determine when to produce contents.</li>
 * </ol>
 *
 * <p>
 * The above special sequence is only executed when a proxy instance is set, and its username is set.
 * </p>
 *
 * @param post a POST request to establish transaction
 * @return POST request response
 * @throws IOException thrown if the post request failed
 */
private HttpResponse initiateTransactionForSend(final HttpPost post) throws IOException {
    if (shouldCheckProxyAuth()) {
        final CloseableHttpAsyncClient asyncClient = getHttpAsyncClient();
        final HttpGet get = createGetControllerRequest();
        final Future<HttpResponse> getResult = asyncClient.execute(get, null);
        try {
            final HttpResponse getResponse = getResult.get(readTimeoutMillis, TimeUnit.MILLISECONDS);
            logger.debug("Proxy auth check has done. getResponse={}", getResponse.getStatusLine());
        } catch (final ExecutionException e) {
            logger.debug("Something has happened at get controller requesting thread for proxy auth check. {}", e.getMessage());
            throw toIOException(e);
        } catch (TimeoutException | InterruptedException e) {
            throw new IOException(e);
        }
    }
    final HttpAsyncRequestProducer asyncRequestProducer = new HttpAsyncRequestProducer() {

        private boolean requestHasBeenReset = false;

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

        @Override
        public HttpRequest generateRequest() throws IOException, HttpException {
            final BasicHttpEntity entity = new BasicHttpEntity();
            post.setEntity(entity);
            return post;
        }

        @Override
        public void produceContent(ContentEncoder encoder, IOControl ioctrl) throws IOException {
            encoder.complete();
            if (shouldCheckProxyAuth() && requestHasBeenReset) {
                logger.debug("Produced content again, assuming the proxy server requires authentication.");
                proxyAuthRequiresResend.set(true);
            }
        }

        @Override
        public void requestCompleted(HttpContext context) {
            debugProxyAuthState(context);
        }

        @Override
        public void failed(Exception ex) {
            final String msg = String.format("Failed to create transaction for %s", post.getURI());
            logger.error(msg, ex);
            eventReporter.reportEvent(Severity.WARNING, EVENT_CATEGORY, msg);
        }

        @Override
        public boolean isRepeatable() {
            return true;
        }

        @Override
        public void resetRequest() throws IOException {
            requestHasBeenReset = true;
        }

        @Override
        public void close() throws IOException {
        }
    };
    final Future<HttpResponse> responseFuture = getHttpAsyncClient().execute(asyncRequestProducer, new BasicAsyncResponseConsumer(), null);
    final HttpResponse response;
    try {
        response = responseFuture.get(readTimeoutMillis, TimeUnit.MILLISECONDS);
    } catch (final ExecutionException e) {
        logger.debug("Something has happened at initiate transaction requesting thread. {}", e.getMessage());
        throw toIOException(e);
    } catch (TimeoutException | InterruptedException e) {
        throw new IOException(e);
    }
    return response;
}

Example 2 with BasicAsyncResponseConsumer

use of org.apache.http.nio.protocol.BasicAsyncResponseConsumer 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);
    }
}