Java 类org.apache.http.nio.protocol.HttpAsyncRequestProducer 实例源码

@SuppressWarnings({ "unchecked", "rawtypes" })
public void testTooLargeResponse() throws Exception {
    ContentTooLongException tooLong = new ContentTooLongException("too long!");
    CloseableHttpAsyncClient httpClient = mock(CloseableHttpAsyncClient.class);
    when(httpClient.<HttpResponse>execute(any(HttpAsyncRequestProducer.class), any(HttpAsyncResponseConsumer.class),
            any(HttpClientContext.class), any(FutureCallback.class))).then(new Answer<Future<HttpResponse>>() {
        @Override
        public Future<HttpResponse> answer(InvocationOnMock invocationOnMock) throws Throwable {
            HeapBufferedAsyncResponseConsumer consumer = (HeapBufferedAsyncResponseConsumer) invocationOnMock.getArguments()[1];
            FutureCallback callback = (FutureCallback) invocationOnMock.getArguments()[3];
            assertEquals(new ByteSizeValue(100, ByteSizeUnit.MB).bytesAsInt(), consumer.getBufferLimit());
            callback.failed(tooLong);
            return null;
        }
    });
    RemoteScrollableHitSource source = sourceWithMockedClient(true, httpClient);

    AtomicBoolean called = new AtomicBoolean();
    Consumer<Response> checkResponse = r -> called.set(true);
    Throwable e = expectThrows(RuntimeException.class,
            () -> source.doStartNextScroll(FAKE_SCROLL_ID, timeValueMillis(0), checkResponse));
    // Unwrap the some artifacts from the test
    while (e.getMessage().equals("failed")) {
        e = e.getCause();
    }
    // This next exception is what the user sees
    assertEquals("Remote responded with a chunk that was too large. Use a smaller batch size.", e.getMessage());
    // And that exception is reported as being caused by the underlying exception returned by the client
    assertSame(tooLong, e.getCause());
    assertFalse(called.get());
}

/**
 * Verifies the content of the {@link HttpRequest} that's internally created and passed through to the http client
 */
@SuppressWarnings("unchecked")
public void testInternalHttpRequest() throws Exception {
    ArgumentCaptor<HttpAsyncRequestProducer> requestArgumentCaptor = ArgumentCaptor.forClass(HttpAsyncRequestProducer.class);
    int times = 0;
    for (String httpMethod : getHttpMethods()) {
        HttpUriRequest expectedRequest = performRandomRequest(httpMethod);
        verify(httpClient, times(++times)).<HttpResponse>execute(requestArgumentCaptor.capture(),
                any(HttpAsyncResponseConsumer.class), any(HttpClientContext.class), any(FutureCallback.class));
        HttpUriRequest actualRequest = (HttpUriRequest)requestArgumentCaptor.getValue().generateRequest();
        assertEquals(expectedRequest.getURI(), actualRequest.getURI());
        assertEquals(expectedRequest.getClass(), actualRequest.getClass());
        assertArrayEquals(expectedRequest.getAllHeaders(), actualRequest.getAllHeaders());
        if (expectedRequest instanceof HttpEntityEnclosingRequest) {
            HttpEntity expectedEntity = ((HttpEntityEnclosingRequest) expectedRequest).getEntity();
            if (expectedEntity != null) {
                HttpEntity actualEntity = ((HttpEntityEnclosingRequest) actualRequest).getEntity();
                assertEquals(EntityUtils.toString(expectedEntity), EntityUtils.toString(actualEntity));
            }
        }
    }
}

@Before
@SuppressWarnings("unchecked")
public void createRestClient() throws IOException {
    CloseableHttpAsyncClient httpClient = mock(CloseableHttpAsyncClient.class);
    when(httpClient.<HttpResponse>execute(any(HttpAsyncRequestProducer.class), any(HttpAsyncResponseConsumer.class),
           any(HttpClientContext.class), any(FutureCallback.class))).thenAnswer(new Answer<Future<HttpResponse>>() {
        @Override
        public Future<HttpResponse> answer(InvocationOnMock invocationOnMock) throws Throwable {
            HttpAsyncRequestProducer requestProducer = (HttpAsyncRequestProducer) invocationOnMock.getArguments()[0];
            HttpUriRequest request = (HttpUriRequest)requestProducer.generateRequest();
            HttpHost httpHost = requestProducer.getTarget();
            HttpClientContext context = (HttpClientContext) invocationOnMock.getArguments()[2];
            assertThat(context.getAuthCache().get(httpHost), instanceOf(BasicScheme.class));
            FutureCallback<HttpResponse> futureCallback = (FutureCallback<HttpResponse>) invocationOnMock.getArguments()[3];
            //return the desired status code or exception depending on the path
            if (request.getURI().getPath().equals("/soe")) {
                futureCallback.failed(new SocketTimeoutException(httpHost.toString()));
            } else if (request.getURI().getPath().equals("/coe")) {
                futureCallback.failed(new ConnectTimeoutException(httpHost.toString()));
            } else if (request.getURI().getPath().equals("/ioe")) {
                futureCallback.failed(new IOException(httpHost.toString()));
            } else {
                int statusCode = Integer.parseInt(request.getURI().getPath().substring(1));
                StatusLine statusLine = new BasicStatusLine(new ProtocolVersion("http", 1, 1), statusCode, "");
                futureCallback.completed(new BasicHttpResponse(statusLine));
            }
            return null;
        }
    });
    int numHosts = RandomNumbers.randomIntBetween(getRandom(), 2, 5);
    httpHosts = new HttpHost[numHosts];
    for (int i = 0; i < numHosts; i++) {
        httpHosts[i] = new HttpHost("localhost", 9200 + i);
    }
    failureListener = new HostsTrackingFailureListener();
    restClient = new RestClient(httpClient, 10000, new Header[0], httpHosts, null, failureListener);
}

@Override
public void sendBatchActivities(List<Activity> activities) throws PXException, IOException {
    HttpAsyncRequestProducer producer = null;
    try {
        String requestBody = JsonUtils.writer.writeValueAsString(activities);
        logger.info("Sending Activity: {}", requestBody);
        HttpPost post = new HttpPost(this.pxConfiguration.getServerURL() + Constants.API_ACTIVITIES);
        post.setEntity(new StringEntity(requestBody, UTF_8));
        post.setConfig(PXCommonUtils.getRequestConfig(pxConfiguration.getConnectionTimeout(),pxConfiguration.getApiTimeout()));
        post.setHeader(HttpHeaders.CONTENT_TYPE, "application/json");
        post.setHeader(HttpHeaders.AUTHORIZATION, "Bearer " + pxConfiguration.getAuthToken());
        producer = HttpAsyncMethods.create(post);
        asyncHttpClient.execute(producer, new BasicAsyncResponseConsumer(), new PxClientAsyncHandler());
    } catch (Exception e) {
        throw new PXException(e);
    } finally {
        if (producer != null) {
            producer.close();
        }
    }
}

@Override
public void sendEnforcerTelemetry(EnforcerTelemetry enforcerTelemetry) throws PXException, IOException{
    HttpAsyncRequestProducer producer = null;
    try {
        String requestBody = JsonUtils.writer.writeValueAsString(enforcerTelemetry);
        logger.info("Sending enforcer telemetry: {}", requestBody);
        HttpPost post = new HttpPost(this.pxConfiguration.getServerURL() + Constants.API_ENFORCER_TELEMETRY);
        post.setEntity(new StringEntity(requestBody, UTF_8));
        PXCommonUtils.getDefaultHeaders(pxConfiguration.getAuthToken());
        post.setHeader(HttpHeaders.CONTENT_TYPE, "application/json");
        post.setHeader(HttpHeaders.AUTHORIZATION, "Bearer " + pxConfiguration.getAuthToken());
        post.setConfig(PXCommonUtils.getRequestConfig(pxConfiguration.getConnectionTimeout(),pxConfiguration.getApiTimeout()));
        producer = HttpAsyncMethods.create(post);
        asyncHttpClient.execute(producer, new BasicAsyncResponseConsumer(), new PxClientAsyncHandler());
    } catch (Exception e) {
        e.printStackTrace();
    } finally {
        if (producer != null) {
            producer.close();
        }
    }
}

private Future<Response> executeRequest(Request request, HttpRequestBase method, HttpContext context, HTTPCallback<HttpResponse> callback) {
    if (request.isDownload()) {
        HttpAsyncRequestProducer producer = HttpAsyncMethods.create(method);
        HttpAsyncResponseConsumer<HttpResponse> consumer = new BasicAsyncResponseConsumer();
        return executeRequest(producer, consumer, context, callback);
    } else
        return executeRequest(method, context, callback);
}

/**
 * Send HttpRequest to Client
 * @param uri - FPC Client Uri
 */
public void connectToClient(String uri){
    this.clientUri = uri;
       try{
           client.start();
           HttpAsyncRequestProducer get = HttpAsyncMethods.createGet(this.clientUri);
           client.execute(get, new MyResponseConsumer(this.clientUri), null);
       } catch (Exception e) {
        ErrorLog.logError(e.getStackTrace());
       }
}

@Override public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
    Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {

    HttpAsyncRequestProducer producer = (HttpAsyncRequestProducer)allArguments[0];
    String uri = producer.generateRequest().getRequestLine().getUri();
    String requestMethod = producer.generateRequest().getRequestLine().getMethod();
    AbstractSpan span = ContextManager.createLocalSpan("httpasyncclient/" + method.getName());
    Tags.HTTP.METHOD.set(span, requestMethod);
    span.setComponent(ComponentsDefine.HTTP_ASYNC_CLIENT).setLayer(SpanLayer.HTTP);
    Tags.URL.set(span, uri);

}

private HttpRequest getHttpRequest(final Object target) {
    try {
        if (!(target instanceof RequestProducerGetter)) {
            return null;
        }

        final HttpAsyncRequestProducer requestProducer = ((RequestProducerGetter)target)._$PINPOINT$_getRequestProducer();
        return requestProducer.generateRequest();
    } catch (Exception e) {
        return null;
    }
}

private NameIntValuePair<String> getHost(final Object target) {
    if (!(target instanceof RequestProducerGetter)) {
        return null;
    }

    final HttpAsyncRequestProducer producer = ((RequestProducerGetter)target)._$PINPOINT$_getRequestProducer();
    final HttpHost httpHost = producer.getTarget();
    if(httpHost != null) {
        return new NameIntValuePair<String>(httpHost.getHostName(), httpHost.getPort());
    } else {
        return null;
    }
}

private Future<Response> executeRequest(HttpAsyncRequestProducer producer, HttpAsyncResponseConsumer<HttpResponse> consumer, HttpContext httpContext, HTTPCallback<HttpResponse> callback) {
    client.execute(producer, consumer, httpContext, new HTTPCCallback(callback));

    return callback.getFuture();
}

@Override
public Future<HttpResponse> execute(final HttpAsyncClient httpClient) throws FileNotFoundException {
    final HttpAsyncResponseConsumer<HttpResponse> consumer = new BasicAsyncResponseConsumer();
    final HttpAsyncRequestProducer producer = this.getProducer();
    return httpClient.execute(producer, consumer, null);
}

HttpAsyncRequestProducer _$PINPOINT$_getRequestProducer();