org.asynchttpclient.AsyncHandler Java Examples

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onBeforeExecute(@Advice.Argument(value = 0) Request request,
                                    @Advice.Argument(value = 1) AsyncHandler<?> asyncHandler,
                                    @Advice.Local("span") Span span) {
    if (tracer == null || tracer.getActive() == null) {
        return;
    }
    ElasticApmAgent.ensureInstrumented(asyncHandler.getClass(), ASYNC_HANDLER_INSTRUMENTATIONS);

    final AbstractSpan<?> parent = tracer.getActive();
    Uri uri = request.getUri();
    span = HttpClientHelper.startHttpClientSpan(parent, request.getMethod(), uri.toUrl(), uri.getScheme(), uri.getHost(), uri.getPort());

    if (span != null) {
        span.activate();
        TextHeaderSetter<Request> headerSetter = null;
        if (headerSetterManager != null) {
            headerSetter = headerSetterManager.getForClassLoaderOfClass(Request.class);
        }
        if (headerSetter != null) {
            span.propagateTraceContext(request, headerSetter);
        }
        handlerSpanMap.put(asyncHandler, span);
    }
}
 

public static Object enter(final Object request, final Object handler) {
  final Request req = (Request)request;
  final Tracer tracer = GlobalTracer.get();
  final Span span = tracer
    .buildSpan(req.getMethod())
    .withTag(Tags.COMPONENT.getKey(), COMPONENT_NAME)
    .withTag(Tags.HTTP_METHOD.getKey(), req.getMethod())
    .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
    .withTag(Tags.HTTP_URL.getKey(), req.getUrl()).start();

  tracer.inject(span.context(), Format.Builtin.HTTP_HEADERS, new TextMap() {
    @Override
    public Iterator<Entry<String,String>> iterator() {
      throw new UnsupportedOperationException("iterator not supported with Tracer.inject()");
    }

    @Override
    public void put(final String key, final String value) {
      req.getHeaders().add(key, value);
    }
  });

  return WrapperProxy.wrap(handler, new TracingAsyncHandler(tracer, (AsyncHandler<?>)handler, span));
}
 

@Override
public <T> ListenableFuture<T> executeRequest(org.asynchttpclient.Request request,
                                              AsyncHandler<T> handler) {

  WebhookOnThrowableHandler<T> wrapped = new WebhookOnThrowableHandler<>(
      handler,
      new CountDownLatch(1));
  ListenableFuture<T> future = super.executeRequest(request, wrapped);
  try {
    future.get();
    future.done();
  } catch (InterruptedException | ExecutionException e) {
    // The future threw an exception, wait for onThrowable to complete.
    wrapped.waitForOnThrowableToFinish();
  }
  return future;
}
 

@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
private static void onAfterExecute(@Advice.Local("span") @Nullable Span span,
                                   @Advice.Argument(value = 1) AsyncHandler<?> asyncHandler,
                                   @Advice.Thrown @Nullable Throwable t) {
    if (span != null) {
        span.deactivate();
        if (t != null) {
            handlerSpanMap.remove(asyncHandler);
            span.captureException(t).end();
        }
    }
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span) {
    span = handlerSpanMap.remove(asyncHandler);
    if (span != null) {
        span.activate();
    }
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span) {
    span = handlerSpanMap.remove(asyncHandler);
    if (span != null) {
        span.activate();
    }
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span, @Advice.Argument(0) HttpResponseStatus status) {
    span = handlerSpanMap.get(asyncHandler);
    if (span != null) {
        span.activate();
    }
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span) {
    span = handlerSpanMap.get(asyncHandler);
    if (span != null) {
        span.activate();
    }
}
 

@Override
public AsyncHandler.State onBodyPartReceived(HttpResponseBodyPart content) throws Exception {
  if (isErrorResponse) {
    return super.onBodyPartReceived(content);
  }

  outputBuffer.writeBytes(content.getBodyByteBuffer());
  return AsyncHandler.State.CONTINUE;
}
 

@Override
public AsyncHandler.State onBodyPartReceived(HttpResponseBodyPart content) throws Exception {
  if (requestFailed) {
    return super.onBodyPartReceived(content);
  }
  outputBuffer.writeBytes(content.getBodyByteBuffer());
  return AsyncHandler.State.CONTINUE;
}
 

public TracingAsyncHandler(final Tracer tracer, final AsyncHandler<?> handler, final Span span) {
  this.tracer = tracer;
  this.handler = handler;
  this.span = span;
}
 

@DataProvider(value = {
    "true   |   true",
    "true   |   false",
    "false  |   true",
    "false  |   false"
}, splitBy = "\\|")
@Test
public void executeAsyncHttpRequest_sets_up_and_executes_call_as_expected(
    boolean performSubspan, boolean currentTracingInfoNull
) {
    // given
    Whitebox.setInternalState(helperSpy, "performSubSpanAroundDownstreamCalls", performSubspan);

    CircuitBreaker<Response> circuitBreakerMock = mock(CircuitBreaker.class);
    doReturn(Optional.of(circuitBreakerMock)).when(helperSpy).getCircuitBreaker(any(RequestBuilderWrapper.class));
    ManualModeTask<Response> cbManualTaskMock = mock(ManualModeTask.class);
    doReturn(cbManualTaskMock).when(circuitBreakerMock).newManualModeTask();

    String url = "http://localhost/some/path";
    String method = "GET";
    BoundRequestBuilder reqMock = mock(BoundRequestBuilder.class);
    RequestBuilderWrapper rbw = new RequestBuilderWrapper(url, method, reqMock, Optional.empty(), false);
    AsyncResponseHandler responseHandlerMock = mock(AsyncResponseHandler.class);

    Span initialSpan = (currentTracingInfoNull) ? null : Tracer.getInstance().startRequestWithRootSpan("foo");
    Deque<Span> initialSpanStack = (currentTracingInfoNull)
                                   ? null
                                   : Tracer.getInstance().getCurrentSpanStackCopy();
    Map<String, String> initialMdc = (currentTracingInfoNull) ? null : MDC.getCopyOfContextMap();
    resetTracingAndMdc();

    // when
    CompletableFuture resultFuture = helperSpy.executeAsyncHttpRequest(
        rbw, responseHandlerMock, initialSpanStack, initialMdc
    );

    // then
    // Verify that the circuit breaker came from the getCircuitBreaker helper method and that its
    //      throwExceptionIfCircuitBreakerIsOpen() method was called.
    verify(helperSpy).getCircuitBreaker(rbw);
    verify(cbManualTaskMock).throwExceptionIfCircuitBreakerIsOpen();

    // Verify that the inner request's execute method was called with a
    //      AsyncCompletionHandlerWithTracingAndMdcSupport for the handler.
    ArgumentCaptor<AsyncHandler> executedHandlerCaptor = ArgumentCaptor.forClass(AsyncHandler.class);
    verify(reqMock).execute(executedHandlerCaptor.capture());
    AsyncHandler executedHandler = executedHandlerCaptor.getValue();
    assertThat(executedHandler).isInstanceOf(AsyncCompletionHandlerWithTracingAndMdcSupport.class);

    // Verify that the AsyncCompletionHandlerWithTracingAndMdcSupport was created with the expected args
    AsyncCompletionHandlerWithTracingAndMdcSupport achwtams =
        (AsyncCompletionHandlerWithTracingAndMdcSupport) executedHandler;
    assertThat(achwtams.completableFutureResponse).isSameAs(resultFuture);
    assertThat(achwtams.responseHandlerFunction).isSameAs(responseHandlerMock);
    assertThat(achwtams.performSubSpanAroundDownstreamCalls).isEqualTo(performSubspan);
    assertThat(achwtams.circuitBreakerManualTask).isEqualTo(Optional.of(cbManualTaskMock));
    if (performSubspan) {
        int initialSpanStackSize = (initialSpanStack == null) ? 0 : initialSpanStack.size();
        assertThat(achwtams.distributedTraceStackToUse).hasSize(initialSpanStackSize + 1);
        Span subspan = (Span) achwtams.distributedTraceStackToUse.peek();
        assertThat(subspan.getSpanName())
            .isEqualTo(initialSpanNameFromStrategy.get());
        if (initialSpan != null) {
            assertThat(subspan.getTraceId()).isEqualTo(initialSpan.getTraceId());
            assertThat(subspan.getParentSpanId()).isEqualTo(initialSpan.getSpanId());
        }
        assertThat(achwtams.mdcContextToUse.get(SpanFieldForLoggerMdc.TRACE_ID.mdcKey)).isEqualTo(subspan.getTraceId());
    }
    else {
        assertThat(achwtams.distributedTraceStackToUse).isSameAs(initialSpanStack);
        assertThat(achwtams.mdcContextToUse).isSameAs(initialMdc);
    }

    // Verify that the trace headers were added (or not depending on state).
    Span spanForDownstreamCall = achwtams.getSpanForCall();
    if (initialSpan == null && !performSubspan) {
        assertThat(spanForDownstreamCall).isNull();
        verifyNoMoreInteractions(reqMock);
    }
    else {
        assertThat(spanForDownstreamCall).isNotNull();
        verify(reqMock).setHeader(TraceHeaders.TRACE_SAMPLED,
                                  convertSampleableBooleanToExpectedB3Value(spanForDownstreamCall.isSampleable()));
        verify(reqMock).setHeader(TraceHeaders.TRACE_ID, spanForDownstreamCall.getTraceId());
        verify(reqMock).setHeader(TraceHeaders.SPAN_ID, spanForDownstreamCall.getSpanId());
        if (spanForDownstreamCall.getParentSpanId() == null) {
            verify(reqMock, never()).setHeader(eq(TraceHeaders.PARENT_SPAN_ID), anyString());
        }
        else {
            verify(reqMock).setHeader(TraceHeaders.PARENT_SPAN_ID, spanForDownstreamCall.getParentSpanId());
        }
        verify(reqMock, never()).setHeader(eq(TraceHeaders.SPAN_NAME), anyString());
    }

    // Verify that any subspan had request tagging performed.
    if (performSubspan) {
        strategyRequestTaggingArgs.get().verifyArgs(spanForDownstreamCall, rbw, wingtipsTagAndNamingAdapterMock);
    }
}
 

@Test
public void givenHttpClient_executeAsyncGetRequestWithAsyncHandler() {

    // execute an unbound GET request
    Request unboundGetRequest = Dsl.get("http://www.baeldung.com").build();

    HTTP_CLIENT.executeRequest(unboundGetRequest, new AsyncHandler<Integer>() {

        int responseStatusCode = -1;

        @Override
        public State onStatusReceived(HttpResponseStatus responseStatus) {
            responseStatusCode = responseStatus.getStatusCode();
            return State.CONTINUE;
        }

        @Override
        public State onHeadersReceived(HttpHeaders headers) {
            return State.CONTINUE;
        }

        @Override
        public State onBodyPartReceived(HttpResponseBodyPart bodyPart) {
            return State.CONTINUE;
        }

        @Override
        public void onThrowable(Throwable t) {

        }

        @Override
        public Integer onCompleted() {
            assertEquals(200, responseStatusCode);
            return responseStatusCode;
        }
    });

    try {
        Thread.sleep(5000);
    } catch (InterruptedException e) {
        e.printStackTrace();
    }
}
 

WebhookOnThrowableHandler(AsyncHandler<T> handler, CountDownLatch latch) {
  this.handler = handler;
  this.latch = latch;
}