Java Code Examples for com.linecorp.armeria.client.ClientRequestContext#push()

@Test
void current() {
    assertThatThrownBy(ServiceRequestContext::current).isInstanceOf(IllegalStateException.class)
                                                      .hasMessageContaining("unavailable");

    final ServiceRequestContext sctx = serviceRequestContext();
    try (SafeCloseable unused = sctx.push()) {
        assertThat(ServiceRequestContext.current()).isSameAs(sctx);
        final ClientRequestContext cctx = clientRequestContext();
        try (SafeCloseable unused1 = cctx.push()) {
            assertThat(ServiceRequestContext.current()).isSameAs(sctx);
            assertThat(ClientRequestContext.current()).isSameAs(cctx);
            assertThat((ClientRequestContext) RequestContext.current()).isSameAs(cctx);
        }
        assertCurrentCtx(sctx);
    }
    assertCurrentCtx(null);

    try (SafeCloseable unused = clientRequestContext().push()) {
        assertThatThrownBy(ServiceRequestContext::current)
                .isInstanceOf(IllegalStateException.class)
                .hasMessageContaining("not a server-side context");
    }
}
 

@Test
void currentOrNull() {
    assertThat(ServiceRequestContext.currentOrNull()).isNull();

    final ServiceRequestContext sctx = serviceRequestContext();
    try (SafeCloseable unused = sctx.push()) {
        assertThat(ServiceRequestContext.currentOrNull()).isSameAs(sctx);
        final ClientRequestContext cctx = clientRequestContext();
        try (SafeCloseable unused1 = cctx.push()) {
            assertThat(ServiceRequestContext.currentOrNull()).isSameAs(sctx);
            assertThat(ClientRequestContext.current()).isSameAs(cctx);
            assertThat((ClientRequestContext) RequestContext.current()).isSameAs(cctx);
        }
        assertCurrentCtx(sctx);
    }
    assertCurrentCtx(null);

    try (SafeCloseable unused = clientRequestContext().push()) {
        assertThatThrownBy(ServiceRequestContext::currentOrNull)
                .isInstanceOf(IllegalStateException.class)
                .hasMessageContaining("not a server-side context");
    }
}
 

@Test
void pushWithOldClientCtxWhoseRootIsThisServiceCtx() {
    final ServiceRequestContext sctx = serviceRequestContext();
    try (SafeCloseable ignored = sctx.push()) {
        assertCurrentCtx(sctx);
        // The root of ClientRequestContext is sctx.
        final ClientRequestContext cctx = clientRequestContext();
        try (SafeCloseable ignored1 = cctx.push()) {
            assertCurrentCtx(cctx);
            try (SafeCloseable ignored2 = sctx.push()) {
                assertCurrentCtx(sctx);
            }
            assertCurrentCtx(cctx);
        }
        assertCurrentCtx(sctx);
    }
    assertCurrentCtx(null);
}
 

@Test
void testClientContextWithMinimalLogs() throws Exception {
    final List<ILoggingEvent> events = prepare(a -> {
        // Export all properties.
        for (BuiltInProperty p : BuiltInProperty.values()) {
            a.addBuiltIn(p);
        }
    });

    final ClientRequestContext ctx = newClientContext("/foo", "type=bar");
    try (SafeCloseable ignored = ctx.push()) {
        final ILoggingEvent e = log(events);
        final Map<String, String> mdc = e.getMDCPropertyMap();
        assertThat(mdc).containsEntry("local.host", "client.com")
                       .containsEntry("local.ip", "5.6.7.8")
                       .containsEntry("local.port", "5678")
                       .containsEntry("remote.host", "server.com")
                       .containsEntry("remote.ip", "1.2.3.4")
                       .containsEntry("remote.port", "8080")
                       .containsEntry("req.direction", "OUTBOUND")
                       .containsEntry("req.authority", "server.com:8080")
                       .containsEntry("req.method", "GET")
                       .containsEntry("req.path", "/foo")
                       .containsEntry("req.query", "type=bar")
                       .containsEntry("scheme", "unknown+h2")
                       .containsEntry("tls.session_id", "0101020305080d15")
                       .containsEntry("tls.proto", "TLSv1.2")
                       .containsEntry("tls.cipher", "some-cipher")
                       .containsKey("req.id")
                       .hasSize(16);
    }
}
 

private static <I extends Request, O extends Response, U extends Client<I, O>>
O pushAndExecute(U delegate, ClientRequestContext ctx) throws Exception {
    @SuppressWarnings("unchecked")
    final I req = (I) firstNonNull(ctx.request(), ctx.rpcRequest());
    try (SafeCloseable ignored = ctx.push()) {
        return delegate.execute(ctx, req);
    }
}
 

@Test
void mapCurrent() {
    assertThat(ServiceRequestContext.mapCurrent(ctx -> "foo", () -> "defaultValue"))
            .isEqualTo("defaultValue");
    assertThat(ServiceRequestContext.mapCurrent(Function.identity(), null)).isNull();

    final ServiceRequestContext sctx = serviceRequestContext();
    try (SafeCloseable unused = sctx.push()) {
        assertThat(ServiceRequestContext.mapCurrent(c -> c == sctx ? "foo" : "bar",
                                                    () -> "defaultValue"))
                .isEqualTo("foo");
        assertThat(ServiceRequestContext.mapCurrent(Function.identity(), null)).isSameAs(sctx);
        final ClientRequestContext cctx = clientRequestContext();
        try (SafeCloseable unused1 = cctx.push()) {
            assertThat(ServiceRequestContext.mapCurrent(c -> c == sctx ? "foo" : "bar",
                                                        () -> "defaultValue"))
                    .isEqualTo("foo");
            assertThat(ClientRequestContext.mapCurrent(c -> c == cctx ? "baz" : "qux",
                                                       () -> "defaultValue"))
                    .isEqualTo("baz");
            assertThat(ServiceRequestContext.mapCurrent(Function.identity(), null)).isSameAs(sctx);
            assertThat(ClientRequestContext.mapCurrent(Function.identity(), null)).isSameAs(cctx);
            assertThat(RequestContext.mapCurrent(Function.identity(), null)).isSameAs(cctx);
        }
        assertCurrentCtx(sctx);
    }
    assertCurrentCtx(null);

    try (SafeCloseable unused = clientRequestContext().push()) {
        assertThatThrownBy(() -> ServiceRequestContext.mapCurrent(c -> "foo", () -> "bar"))
                .isInstanceOf(IllegalStateException.class)
                .hasMessageContaining("not a server-side context");
    }
}
 

@Test
void pushWithOldIrrelevantClientCtx() {
    final ClientRequestContext cctx = clientRequestContext();
    try (SafeCloseable ignored = cctx.push()) {
        assertCurrentCtx(cctx);
        final ServiceRequestContext sctx = serviceRequestContext();
        assertThatThrownBy(sctx::push).isInstanceOf(IllegalStateException.class);
    }
    assertCurrentCtx(null);
}
 

@Test
void get() {
    final ServiceRequestContext ctx = newContext();
    RequestScopedMdc.put(ctx, "foo", "1");
    assertThat(RequestScopedMdc.get(ctx, "foo")).isEqualTo("1");
    assertThat(MDC.get("foo")).isNull();

    try (SafeCloseable ignored = ctx.push()) {
        assertThat(MDC.get("foo")).isEqualTo("1");
        // Request-scoped property should have priority over thread-local one.
        MDC.put("foo", "2");
        assertThat(MDC.get("foo")).isEqualTo("1");

        // A client context should expose the properties from the root context.
        final ClientRequestContext cctx = ClientRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
        assertThat(cctx.root()).isSameAs(ctx);
        assertThat(RequestScopedMdc.get(cctx, "foo")).isEqualTo("1");

        // A client context can override the property from the root context,
        // but it shouldn't affect the root context's own property.
        RequestScopedMdc.put(cctx, "foo", "3");
        assertThat(RequestScopedMdc.get(ctx, "foo")).isEqualTo("1");
        assertThat(RequestScopedMdc.get(cctx, "foo")).isEqualTo("3");

        try (SafeCloseable ignored2 = cctx.push()) {
            // If both ctx and cctx do not have 'foo' set, thread-local property should be retrieved.
            RequestScopedMdc.remove(ctx, "foo");
            RequestScopedMdc.remove(cctx, "foo");
            assertThat(MDC.get("foo")).isEqualTo("2");
        }
    }
}
 

@Test
void testClientContextWithFullLogs() throws Exception {
    final List<ILoggingEvent> events = prepare(a -> {
        // Export all properties.
        for (BuiltInProperty p : BuiltInProperty.values()) {
            a.addBuiltIn(p);
        }
        // .. and an attribute.
        a.addAttribute("attrs.my_attr_name", MY_ATTR, new CustomObjectNameStringifier());
        a.addAttribute("attrs.my_attr_value", MY_ATTR, new CustomObjectValueStringifier());
        // .. and some HTTP headers.
        a.addRequestHeader(HttpHeaderNames.USER_AGENT);
        a.addResponseHeader(HttpHeaderNames.DATE);
    });

    final ClientRequestContext ctx = newClientContext("/bar", null);
    try (SafeCloseable ignored = ctx.push()) {
        final RequestLogBuilder log = ctx.logBuilder();
        log.serializationFormat(ThriftSerializationFormats.BINARY);
        log.requestLength(64);
        log.requestHeaders(RequestHeaders.of(HttpMethod.GET, "/bar",
                                             HttpHeaderNames.USER_AGENT, "some-client"));
        log.requestContent(RPC_REQ, THRIFT_CALL);
        log.endRequest();
        log.responseLength(128);
        log.responseHeaders(ResponseHeaders.of(HttpStatus.OK,
                                               HttpHeaderNames.DATE, "some-date"));
        log.responseContent(RPC_RES, THRIFT_REPLY);
        log.endResponse();

        final ILoggingEvent e = log(events);
        final Map<String, String> mdc = e.getMDCPropertyMap();
        assertThat(mdc).containsEntry("local.host", "client.com")
                       .containsEntry("local.ip", "5.6.7.8")
                       .containsEntry("local.port", "5678")
                       .containsEntry("remote.host", "server.com")
                       .containsEntry("remote.ip", "1.2.3.4")
                       .containsEntry("remote.port", "8080")
                       .containsEntry("req.direction", "OUTBOUND")
                       .containsEntry("req.authority", "server.com:8080")
                       .containsEntry("req.method", "GET")
                       .containsEntry("req.path", "/bar")
                       .containsEntry("scheme", "tbinary+h2")
                       .containsEntry("req.name", "hello")
                       .containsEntry("req.content_length", "64")
                       .containsEntry("req.name", "hello")
                       .containsEntry("req.serviceName", Object.class.getName())
                       .containsEntry("req.content", "[world]")
                       .containsEntry("res.status_code", "200")
                       .containsEntry("res.content_length", "128")
                       .containsEntry("res.content", "Hello, world!")
                       .containsEntry("req.headers.user-agent", "some-client")
                       .containsEntry("res.headers.date", "some-date")
                       .containsEntry("tls.session_id", "0101020305080d15")
                       .containsEntry("tls.proto", "TLSv1.2")
                       .containsEntry("tls.cipher", "some-cipher")
                       .containsEntry("attrs.my_attr_name", "some-name")
                       .containsEntry("attrs.my_attr_value", "some-value")
                       .containsKey("req.id")
                       .containsKey("elapsed_nanos")
                       .hasSize(27);
    }
}
 

private static RequestLog testRemoteInvocation(Tracing tracing, @Nullable String remoteServiceName)
        throws Exception {

    HttpTracing httpTracing = HttpTracing.newBuilder(tracing)
                                         .clientRequestParser(ArmeriaHttpClientParser.get())
                                         .clientResponseParser(ArmeriaHttpClientParser.get())
                                         .build();
    if (remoteServiceName != null) {
        httpTracing = httpTracing.clientOf(remoteServiceName);
    }

    // prepare parameters
    final HttpRequest req = HttpRequest.of(RequestHeaders.of(HttpMethod.POST, "/hello/armeria",
                                                             HttpHeaderNames.SCHEME, "http",
                                                             HttpHeaderNames.AUTHORITY, "foo.com"));
    final RpcRequest rpcReq = RpcRequest.of(HelloService.Iface.class, "hello", "Armeria");
    final HttpResponse res = HttpResponse.of(HttpStatus.OK);
    final RpcResponse rpcRes = RpcResponse.of("Hello, Armeria!");
    final ClientRequestContext ctx = ClientRequestContext.builder(req).build();
    final HttpRequest actualReq = ctx.request();
    assertThat(actualReq).isNotNull();

    ctx.logBuilder().requestFirstBytesTransferred();
    ctx.logBuilder().requestContent(rpcReq, actualReq);
    ctx.logBuilder().endRequest();

    try (SafeCloseable ignored = ctx.push()) {
        final HttpClient delegate = mock(HttpClient.class);
        when(delegate.execute(any(), any())).thenReturn(res);

        final BraveClient stub = BraveClient.newDecorator(httpTracing).apply(delegate);
        // do invoke
        final HttpResponse actualRes = stub.execute(ctx, actualReq);

        assertThat(actualRes).isEqualTo(res);

        verify(delegate, times(1)).execute(same(ctx), argThat(arg -> {
            final RequestHeaders headers = arg.headers();
            return headers.contains(HttpHeaderNames.of("x-b3-traceid")) &&
                   headers.contains(HttpHeaderNames.of("x-b3-spanid")) &&
                   headers.contains(HttpHeaderNames.of("x-b3-sampled"));
        }));
    }

    ctx.logBuilder().responseHeaders(ResponseHeaders.of(HttpStatus.OK));
    ctx.logBuilder().responseFirstBytesTransferred();
    ctx.logBuilder().responseContent(rpcRes, res);
    ctx.logBuilder().endResponse();
    return ctx.log().ensureComplete();
}
 

@Test
void getAllNested() {
    MDC.put("foo", "1");
    MDC.put("bar", "2");

    final ServiceRequestContext ctx = newContext();
    try (SafeCloseable ignored = ctx.push()) {
        final ClientRequestContext cctx = ClientRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));

        // When the root context map exists but it's empty:
        try (SafeCloseable ignored2 = cctx.push()) {
            assertThat(RequestScopedMdc.getAll(cctx)).isEmpty();
            assertThat(MDC.getCopyOfContextMap()).containsOnly(
                    Maps.immutableEntry("foo", "1"),
                    Maps.immutableEntry("bar", "2"));
        }

        // When the root context map is not empty:
        RequestScopedMdc.put(ctx, "bar", "3");
        RequestScopedMdc.put(ctx, "baz", "4");
        try (SafeCloseable ignored2 = cctx.push()) {
            // root context's properties should be retrieved.
            assertThat(RequestScopedMdc.getAll(cctx)).containsOnly(
                    Maps.immutableEntry("bar", "3"),
                    Maps.immutableEntry("baz", "4"));
            assertThat(MDC.getCopyOfContextMap()).containsOnly(
                    Maps.immutableEntry("foo", "1"),
                    Maps.immutableEntry("bar", "3"),
                    Maps.immutableEntry("baz", "4"));

            // root context's properties should be overwritten by own properties.
            RequestScopedMdc.put(cctx, "baz", "5");
            RequestScopedMdc.put(cctx, "qux", "6");

            assertThat(RequestScopedMdc.getAll(cctx)).containsOnly(
                    Maps.immutableEntry("bar", "3"),
                    Maps.immutableEntry("baz", "5"),
                    Maps.immutableEntry("qux", "6"));
            assertThat(MDC.getCopyOfContextMap()).containsOnly(
                    Maps.immutableEntry("foo", "1"),
                    Maps.immutableEntry("bar", "3"),
                    Maps.immutableEntry("baz", "5"),
                    Maps.immutableEntry("qux", "6"));
        }
    }
}