com.linecorp.armeria.common.util.SafeCloseable Java Examples

/**
 * Returns a {@link CompletionStage} that makes sure the current {@link RequestContext} is set and
 * then invokes the input {@code stage}.
 */
default <T> CompletionStage<T> makeContextAware(CompletionStage<T> stage) {
    final CompletableFuture<T> future = JavaVersionSpecific.get().newRequestContextAwareFuture(this);
    stage.handle((result, cause) -> {
        try (SafeCloseable ignored = push()) {
            if (cause != null) {
                future.completeExceptionally(cause);
            } else {
                future.complete(result);
            }
        } catch (Throwable t) {
            future.completeExceptionally(t);
        }
        return null;
    });
    return future;
}
 

@Test
public void testNestedManipulation() throws Exception {
    // Split the secret into two pieces.
    final String secretA = SECRET.substring(0, SECRET.length() >>> 1);
    final String secretB = SECRET.substring(secretA.length());

    final HelloService.Iface client = newClient();
    try (SafeCloseable ignored = Clients.withHttpHeader(AUTHORIZATION, secretA)) {
        // Should fail with the first half of the secret.
        assertAuthorizationFailure(client, secretA);
        try (SafeCloseable ignored2 = Clients.withHttpHeaders(builder -> {
            builder.set(AUTHORIZATION, builder.get(AUTHORIZATION) + secretB);
        })) {
            // Should pass if both manipulators worked.
            assertThat(client.hello("foobar")).isEqualTo("Hello, foobar!");
        }
        // Should fail again with the first half of the secret.
        assertAuthorizationFailure(client, secretA);
    }

    // Ensure that the header manipulator set in the thread-local variable has been cleared.
    assertAuthorizationFailure(client, null);
}
 

@Test
public void httpResponseHeaderContainsFoo() throws TException {
    final Iface client =
            Clients.builder(server.httpUri(BINARY) + "/hello")
                   .decorator((delegate, ctx, req) -> {
                       final HttpResponse res = delegate.execute(ctx, req);
                       return new FilteredHttpResponse(res) {
                           @Override
                           protected HttpObject filter(HttpObject obj) {
                               if (obj instanceof HttpHeaders) {
                                   final HttpHeaders headers = (HttpHeaders) obj;
                                   assertThat(headers.get("foo")).isEqualTo("bar");
                               }
                               return obj;
                           }
                       };
                   })
                   .build(Iface.class);
    try (SafeCloseable ignored = Clients.withHttpHeader(AUTHORIZATION, SECRET)) {
        assertThat(client.hello("trustin")).isEqualTo("Hello, trustin!");
    }
}
 

private void addCallbackAndFlush(Throwable cause, State oldState, ChannelFuture future, boolean isReset) {
    if (oldState != State.DONE) {
        future.addListener(f -> {
            try (SafeCloseable ignored = RequestContextUtil.pop()) {
                if (f.isSuccess() && !isReset) {
                    maybeLogFirstResponseBytesTransferred();
                }
                // Write an access log always with a cause. Respect the first specified cause.
                if (tryComplete()) {
                    logBuilder().endResponse(cause);
                    reqCtx.log().whenComplete().thenAccept(reqCtx.config().accessLogWriter()::log);
                }
            }
        });
    }
    ctx.flush();
}
 

@Override
public void subscribe(Subscriber<? super T>[] s) {
    if (!validate(s)) {
        return;
    }
    final int n = s.length;
    @SuppressWarnings("unchecked")
    final Subscriber<? super T>[] parents = new Subscriber[n];
    for (int i = 0; i < n; i++) {
        final Subscriber<? super T> z = s[i];
        if (z instanceof ConditionalSubscriber) {
            parents[i] = new RequestContextConditionalSubscriber<>(
                    (ConditionalSubscriber<? super T>) z, assemblyContext
            );
        } else {
            parents[i] = new RequestContextSubscriber<>(z, assemblyContext);
        }
    }
    try (SafeCloseable ignored = assemblyContext.push()) {
        source.subscribe(parents);
    }
}
 

@Test
void pushWithOldClientCtxWhoseRootIsSameServiceCtx_derivedCtx() {
    final ServiceRequestContext sctx = serviceRequestContext();
    try (SafeCloseable ignored = sctx.push()) {
        assertCurrentCtx(sctx);
        final ClientRequestContext cctx1 = clientRequestContext();
        final ClientRequestContext derived = cctx1.newDerivedContext(cctx1.id(), cctx1.request(),
                                                                     cctx1.rpcRequest());
        try (SafeCloseable ignored1 = derived.push()) {
            assertCurrentCtx(derived);
            final ClientRequestContext cctx2 = clientRequestContext();
            assertThat(derived.root()).isSameAs(cctx2.root());

            try (SafeCloseable ignored2 = cctx2.push()) {
                assertCurrentCtx(cctx2);
            }
            assertCurrentCtx(derived);
        }
        assertCurrentCtx(sctx);
    }
    assertCurrentCtx(null);
}
 

@Test
public void testSimpleManipulationAsync() throws Exception {
    final HelloService.AsyncIface client = Clients.newClient(
            server.httpUri(BINARY) + "/hello", HelloService.AsyncIface.class);

    final BlockingQueue<Object> result = new ArrayBlockingQueue<>(1);
    final Callback callback = new Callback(result);

    try (SafeCloseable ignored = Clients.withHttpHeader(AUTHORIZATION, SECRET)) {
        client.hello("armeria", callback);
    }

    assertThat(result.poll(10, TimeUnit.SECONDS)).isEqualTo("Hello, armeria!");

    // Ensure that the header manipulator set in the thread-local variable has been cleared.
    client.hello("bar", callback);
    assertThat(result.poll(10, TimeUnit.SECONDS))
            .isInstanceOf(TException.class)
            .matches(o -> ((Throwable) o).getMessage().contains("not authorized"),
                     "must fail with authorization failure");
}
 

@Test
void testMdcPropertyPreservation() throws Exception {
    final List<ILoggingEvent> events = prepare(a -> a.addBuiltIn(BuiltInProperty.REQ_DIRECTION));

    MDC.put("some-prop", "some-value");
    final ServiceRequestContext ctx = newServiceContext("/foo", null);
    try (SafeCloseable ignored = ctx.push()) {
        final ILoggingEvent e = log(events);
        final Map<String, String> mdc = e.getMDCPropertyMap();
        assertThat(mdc).containsEntry("req.direction", "INBOUND")
                       .containsEntry("some-prop", "some-value")
                       .hasSize(2);
    } finally {
        MDC.remove("some-prop");
    }
}
 

@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 contextAwareDoesNotThrowException() {
    final WebClient client =
            WebClient.builder(server.httpUri())
                     .responseTimeoutMillis(100)
                     .decorator(RetryingClient.builder(RetryRule.failsafe())
                                              .maxTotalAttempts(2)
                                              .newDecorator())
                     .build();

    final ServiceRequestContext dummyCtx = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    try (SafeCloseable ignored = dummyCtx.push()) {
        final CompletableFuture<AggregatedHttpResponse> future = client.get("/").aggregate();
        assertThatThrownBy(() -> dummyCtx.makeContextAware(future).join()).hasCauseInstanceOf(
                ResponseTimeoutException.class);
    }
}
 

@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");
    }
}
 

@Override
public void info(String format, Object arg) {
    if (isInfoEnabled()) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.info(decorate(format), arg);
        }
    }
}
 

@SuppressWarnings("unchecked")
@Override
public T get() throws Throwable {
    try (SafeCloseable ignored = assemblyContext.push()) {
        return ((Supplier<T>) source).get();
    }
}
 

@Override
public void info(String msg) {
    if (isInfoEnabled()) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.info(decorate(msg));
        }
    }
}
 

@Override
@SuppressWarnings("unchecked")
public void operationComplete(Future future) throws Exception {
    try (SafeCloseable ignored = ctx.push()) {
        listener.operationComplete(future);
    }
}
 

@Override
public void debug(Marker marker, String format, Object arg1, Object arg2) {
    if (isDebugEnabled(marker)) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.debug(marker, decorate(format), arg1, arg2);
        }
    }
}
 

@Test
public void testSimpleManipulation() throws Exception {
    final HelloService.Iface client = newClient();
    try (SafeCloseable ignored = Clients.withHttpHeader(AUTHORIZATION, SECRET)) {
        assertThat(client.hello("trustin")).isEqualTo("Hello, trustin!");
    }

    // Ensure that the header manipulator set in the thread-local variable has been cleared.
    assertAuthorizationFailure(client, null);
}
 

@Override
public void error(Marker marker, String format, Object... arguments) {
    if (isErrorEnabled(marker)) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.error(marker, decorate(format), arguments);
        }
    }
}
 

@Override
public void warn(String format, Object... arguments) {
    if (isWarnEnabled()) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.warn(decorate(format), arguments);
        }
    }
}
 

@SuppressWarnings("unchecked")
@Override
protected void subscribeActual(Subscriber<? super T> s) {
    try (SafeCloseable ignored = assemblyContext.push()) {
        if (s instanceof ConditionalSubscriber) {
            source.subscribe(new RequestContextConditionalSubscriber<>((ConditionalSubscriber<? super T>) s,
                                                                       assemblyContext));
        } else {
            source.subscribe(new RequestContextSubscriber<>(s, assemblyContext));
        }
    }
}
 

private void makeContextAwareLoggingException0(Runnable action) {
    final SafeCloseable handle;
    try {
        handle = ctx.push();
    } catch (Throwable th) {
        logger.warn("An error occurred while pushing a context", th);
        throw th;
    }

    try {
        action.run();
    } finally {
        handle.close();
    }
}
 

@Override
public void run() {
    isRun = true;

    final ScheduledFuture<?> timeoutFuture = get();
    if (timeoutFuture != null) {
        if (timeoutFuture.isDone() || !timeoutFuture.cancel(false)) {
            // Timeout task ran already or is determined to run.
            numActiveRequests.decrementAndGet();
            return;
        }
    }

    try (SafeCloseable ignored = ctx.replace()) {
        try {
            final O actualRes = unwrap().execute(ctx, req);
            actualRes.whenComplete().handleAsync((unused, cause) -> {
                numActiveRequests.decrementAndGet();
                drain();
                return null;
            }, ctx.eventLoop());
            resFuture.complete(actualRes);
        } catch (Throwable t) {
            numActiveRequests.decrementAndGet();
            resFuture.completeExceptionally(t);
        }
    }
}
 

@Override
protected HttpResponse execute(PooledHttpClient delegate, ClientRequestContext ctx,
                             PooledHttpRequest req) throws Exception {
    // Whether the decorator is applied to an unpooled or pooled delegate, it doesn't matter, we have
    // easy access to the unsafe API.
    return HttpResponse.from(
            delegate.execute(ctx, req).aggregateWithPooledObjects(ctx.eventLoop(), ctx.alloc())
                    .thenApply(agg -> {
                        try (SafeCloseable unused = agg) {
                            return HttpResponse.of(agg.contentUtf8() + " and goodbye!");
                        }
                    }));
}
 

@Override
public void warn(Marker marker, String format, Object... arguments) {
    if (isWarnEnabled(marker)) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.warn(marker, decorate(format), arguments);
        }
    }
}
 

@SuppressWarnings("unchecked")
@Override
public T call() throws Exception {
    try (SafeCloseable ignored = assemblyContext.push()) {
        return ((Callable<T>) source).call();
    }
}
 

@Override
public void error(String format, Object... arguments) {
    if (isErrorEnabled()) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.error(decorate(format), arguments);
        }
    }
}
 

@Override
public void warn(Marker marker, String msg) {
    if (isWarnEnabled(marker)) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.warn(marker, decorate(msg));
        }
    }
}
 

/**
 * Returns a {@link BiConsumer} that makes sure the current {@link RequestContext} is set and then invokes
 * the input {@code action}.
 */
default <T, U> BiConsumer<T, U> makeContextAware(BiConsumer<T, U> action) {
    return (t, u) -> {
        try (SafeCloseable ignored = push()) {
            action.accept(t, u);
        }
    };
}
 

@Override
public void trace(String format, Object... arguments) {
    if (isTraceEnabled()) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.trace(decorate(format), arguments);
        }
    }
}
 

@Override
public void warn(Marker marker, String format, Object arg1, Object arg2) {
    if (isWarnEnabled(marker)) {
        try (SafeCloseable ignored = ctx.push()) {
            logger.warn(marker, decorate(format), arg1, arg2);
        }
    }
}