Java Code Examples for com.linecorp.armeria.client.Clients#newContextCaptor()

@Override
protected void get(WebClient client, String path, BiConsumer<Integer, Throwable> callback) {
    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        final HttpResponse res = client.get(path);
        final ClientRequestContext ctx = ctxCaptor.get();
        res.aggregate().handle((response, cause) -> {
            try (SafeCloseable ignored = ctx.push()) {
                if (cause == null) {
                    callback.accept(response.status().code(), null);
                } else {
                    callback.accept(null, cause);
                }
            }
            return null;
        });
    }
}
 

@Test
void pushedContextOnAsyncMethodCallback() throws Exception {
    final AtomicReference<ClientRequestContext> ctxHolder = new AtomicReference<>();
    final AsyncIface client = Clients.newClient(server.httpUri(BINARY) + "/hello", AsyncIface.class);

    final ClientRequestContext ctx;
    final CountDownLatch latch = new CountDownLatch(1);
    try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
        client.hello("foo", new AsyncMethodCallback<String>() {
            @Override
            public void onComplete(String response) {
                assertThat(response).isEqualTo("Hello, foo!");
                ctxHolder.set(RequestContext.currentOrNull());
                latch.countDown();
            }

            @Override
            public void onError(Exception exception) {}
        });
        ctx = captor.get();
    }

    latch.await();
    assertThat(ctx).isSameAs(ctxHolder.get());
}
 

private static void checkContextOnAsyncMethodCallbackOnError(AsyncIface client) throws Exception {
    final AtomicReference<ClientRequestContext> ctxHolder = new AtomicReference<>();
    final ClientRequestContext ctx;
    final CountDownLatch latch = new CountDownLatch(1);
    try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
        client.hello("foo", new AsyncMethodCallback<String>() {
            @Override
            public void onComplete(String response) {}

            @Override
            public void onError(Exception exception) {
                ctxHolder.set(RequestContext.currentOrNull());
                latch.countDown();
            }
        });
        ctx = captor.get();
    }

    latch.await();
    assertThat(ctx).isSameAs(ctxHolder.get());
}
 

@ParameterizedTest
@ArgumentsSource(ParametersProvider.class)
void contextCaptorSync(
        ClientOptions clientOptions, SerializationFormat format, SessionProtocol protocol)
        throws Exception {
    final HelloService.Iface client = Clients.builder(uri(Handlers.HELLO, format, protocol))
                                             .options(clientOptions)
                                             .build(Handlers.HELLO.iface());
    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        client.hello("kukuman");
        final ClientRequestContext ctx = ctxCaptor.get();
        final RpcRequest rpcReq = ctx.rpcRequest();
        assertThat(rpcReq).isNotNull();
        assertThat(rpcReq.method()).isEqualTo("hello");
        assertThat(rpcReq.params()).containsExactly("kukuman");
    }
}
 

@ParameterizedTest
@ArgumentsSource(ParametersProvider.class)
void contextCaptorAsync(
        ClientOptions clientOptions, SerializationFormat format, SessionProtocol protocol)
        throws Exception {
    final HelloService.AsyncIface client =
            Clients.builder(uri(Handlers.HELLO, format, protocol))
                   .options(clientOptions)
                   .build(Handlers.HELLO.asyncIface());

    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        client.hello("kukuman", new ThriftFuture<>());
        final ClientRequestContext ctx = ctxCaptor.get();
        final RpcRequest rpcReq = ctx.rpcRequest();
        assertThat(rpcReq).isNotNull();
        assertThat(rpcReq.method()).isEqualTo("hello");
        assertThat(rpcReq.params()).containsExactly("kukuman");
    }
}
 

@Test
void respectsHttpClientUri_endpointGroup() throws Exception {
    final EndpointGroup group = EndpointGroup.of(Endpoint.of("127.0.0.1", server.httpPort()),
                                                 Endpoint.of("127.0.0.1", server.httpPort()));

    final Service service = ArmeriaRetrofit.builder("http", group)
                                           .addConverterFactory(converterFactory)
                                           .build()
                                           .create(Service.class);

    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        final Response<Pojo> response = service.postForm("Cony", 26).get();

        final RequestLog log = ctxCaptor.get().log().whenComplete().join();
        assertThat(log.sessionProtocol()).isSameAs(SessionProtocol.H2C);
        assertThat(log.requestHeaders().authority()).isEqualTo("127.0.0.1:" + server.httpPort());

        // TODO(ide) Use the actual `host:port`. See https://github.com/line/armeria/issues/379
        final HttpUrl url = response.raw().request().url();
        assertThat(url.scheme()).isEqualTo("http");
        assertThat(url.host()).startsWith("armeria-group-");
        assertThat(url.pathSegments()).containsExactly("postForm");
    }
}
 

private synchronized void check() {
    if (closeable.isClosing()) {
        return;
    }

    final RequestHeaders headers;
    final RequestHeadersBuilder builder =
            RequestHeaders.builder(useGet ? HttpMethod.GET : HttpMethod.HEAD, path)
                          .add(HttpHeaderNames.AUTHORITY, authority);
    if (maxLongPollingSeconds > 0) {
        headers = builder.add(HttpHeaderNames.IF_NONE_MATCH, wasHealthy ? "\"healthy\"" : "\"unhealthy\"")
                         .add(HttpHeaderNames.PREFER, "wait=" + maxLongPollingSeconds)
                         .build();
    } else {
        headers = builder.build();
    }

    try (ClientRequestContextCaptor reqCtxCaptor = Clients.newContextCaptor()) {
        lastResponse = webClient.execute(headers);
        final ClientRequestContext reqCtx = reqCtxCaptor.get();
        lastResponse.subscribeWithPooledObjects(new HealthCheckResponseSubscriber(reqCtx, lastResponse),
                                                reqCtx.eventLoop());
    }
}
 

@Test
void decodedContentPreview() {
    final WebClient client = WebClient.builder(server.httpUri())
                                      .decorator(DecodingClient.newDecorator())
                                      .decorator(ContentPreviewingClient.newDecorator(100))
                                      .build();
    final RequestHeaders headers = RequestHeaders.of(HttpMethod.POST, "/",
                                                     HttpHeaderNames.CONTENT_TYPE, "text/plain");

    final ClientRequestContext context;
    try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
        final AggregatedHttpResponse res = client.execute(headers, "Armeria").aggregate().join();
        assertThat(res.contentUtf8()).isEqualTo("Hello Armeria!");
        assertThat(res.headers().get(HttpHeaderNames.CONTENT_ENCODING)).isEqualTo("gzip");
        context = captor.get();
    }

    final RequestLog requestLog = context.log().whenComplete().join();
    assertThat(requestLog.requestContentPreview()).isEqualTo("Armeria");
    assertThat(requestLog.responseContentPreview()).isEqualTo("Hello Armeria!");
}
 

/**
 * Unlike {@link #decodedContentPreview()}, the content preview of this test is encoded data because
 * the previewing decorator is inserted before {@link DecodingClient}.
 */
@Test
void contentPreviewIsDecodedInPreviewer() {
    final WebClient client = WebClient.builder(server.httpUri())
                                      .decorator(decodingContentPreviewDecorator())
                                      .decorator(DecodingClient.newDecorator())
                                      .build();
    final RequestHeaders headers = RequestHeaders.of(HttpMethod.POST, "/",
                                                     HttpHeaderNames.CONTENT_TYPE,
                                                     MediaType.PLAIN_TEXT_UTF_8);

    final ClientRequestContext context;
    try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
        final AggregatedHttpResponse res = client.execute(headers, "Armeria").aggregate().join();
        assertThat(res.contentUtf8()).isEqualTo("Hello Armeria!");
        assertThat(res.headers().get(HttpHeaderNames.CONTENT_ENCODING)).isEqualTo("gzip");
        context = captor.get();
    }

    final RequestLog requestLog = context.log().whenComplete().join();
    assertThat(requestLog.requestContentPreview()).isEqualTo("Armeria");
    assertThat(requestLog.responseContentPreview()).isEqualTo("Hello Armeria!");
}
 

@Test
void contextCaptorBlocking() {
    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        blockingStub.emptyCall(EMPTY);
        final ClientRequestContext ctx = ctxCaptor.get();
        assertThat(ctx.path()).isEqualTo("/armeria.grpc.testing.TestService/EmptyCall");
    }
}
 

@Test
void contextCaptorAsync() {
    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        asyncStub.emptyCall(EMPTY, StreamRecorder.create());
        final ClientRequestContext ctx = ctxCaptor.get();
        assertThat(ctx.path()).isEqualTo("/armeria.grpc.testing.TestService/EmptyCall");
    }
}
 

@Test
void sessionProtocolH1C() throws Exception {
    final Service service = ArmeriaRetrofit.builder("h1c://127.0.0.1:" + server.httpPort())
                                           .addConverterFactory(converterFactory)
                                           .build()
                                           .create(Service.class);

    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        final Pojo pojo = service.pojo().get();
        assertThat(pojo).isEqualTo(new Pojo("Cony", 26));

        final RequestLog log = ctxCaptor.get().log().whenComplete().join();
        assertThat(log.sessionProtocol()).isSameAs(SessionProtocol.H1C);
    }
}
 

@Test
void shouldRetryEvenIfEndpointGroupIsEmpty() {
    final int numAttempts = 3;
    final WebClient client =
            WebClient.builder(SessionProtocol.HTTP, EndpointGroup.empty())
                     .decorator(RetryingClient.builder(RetryRule.builder()
                                                                .onUnprocessed()
                                                                .thenBackoff(Backoff.withoutDelay()))
                                              .maxTotalAttempts(numAttempts)
                                              .newDecorator())
                     .build();

    final ClientRequestContext ctx;
    try (ClientRequestContextCaptor ctxCaptor = Clients.newContextCaptor()) {
        client.get("/").aggregate();
        ctx = ctxCaptor.get();
    }

    // Make sure all attempts have failed with `EmptyEndpointGroupException`.
    final RequestLog log = ctx.log().whenComplete().join();
    assertEmptyEndpointGroupException(log);

    assertThat(log.children()).hasSize(numAttempts);
    log.children().stream()
       .map(RequestLogAccess::ensureComplete)
       .forEach(RetryingClientWithEmptyEndpointGroupTest::assertEmptyEndpointGroupException);
}
 

@Override
public void serverStarted(Server server) throws Exception {
    final InstanceInfo newInfo = fillAndCreateNewInfo(instanceInfo, server);

    try (ClientRequestContextCaptor contextCaptor = Clients.newContextCaptor()) {
        final HttpResponse response = client.register(newInfo);
        final ClientRequestContext ctx = contextCaptor.get();
        response.aggregate().handle((res, cause) -> {
            if (closed) {
                return null;
            }
            if (cause != null) {
                logger.warn("Failed to register {} to Eureka: {}",
                            newInfo.getHostName(), client.uri(), cause);
                return null;
            }
            final ResponseHeaders headers = res.headers();
            if (headers.status() != HttpStatus.NO_CONTENT) {
                logger.warn("Failed to register {} to Eureka: {}. (status: {}, content: {})",
                            newInfo.getHostName(), client.uri(), headers.status(), res.contentUtf8());
            } else {
                logger.info("Registered {} to Eureka: {}", newInfo.getHostName(), client.uri());
                scheduleHeartBeat(ctx.eventLoop(), newInfo);
            }
            return null;
        });
    }
}
 

private void fetchRegistry() {
    final PooledHttpResponse response;
    final ClientRequestContext ctx;
    try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) {
        response = webClient.execute(requestHeaders);
        ctx = captor.get();
    }

    final EventLoop eventLoop = ctx.eventLoop();
    response.aggregateWithPooledObjects(eventLoop, ctx.alloc()).handle((aggregatedRes, cause) -> {
        try (SafeCloseable unused = aggregatedRes) {
            if (closed) {
                return null;
            }
            if (cause != null) {
                logger.warn("Unexpected exception while fetching the registry from: {}." +
                            " (requestHeaders: {})", webClient.uri(), requestHeaders, cause);
            } else {
                final HttpStatus status = aggregatedRes.status();
                if (!status.isSuccess()) {
                    logger.warn("Unexpected response from: {}. (status: {}, content: {}, " +
                                "requestHeaders: {})", webClient.uri(), status,
                                aggregatedRes.contentUtf8(), requestHeaders);
                } else {
                    final HttpData content = aggregatedRes.content();
                    try {
                        final List<Endpoint> endpoints = responseConverter.apply(content.array());
                        setEndpoints(endpoints);
                    } catch (Exception e) {
                        logger.warn("Unexpected exception while parsing a response from: {}. " +
                                    "(content: {}, responseConverter: {}, requestHeaders: {})",
                                    webClient.uri(), content.toStringUtf8(),
                                    responseConverter, requestHeaders, e);
                    }
                }
            }
        }
        scheduledFuture = eventLoop.schedule(this::fetchRegistry,
                                             registryFetchIntervalSeconds, TimeUnit.SECONDS);
        return null;
    });
}