Introduce a way to unwrap RequestContexts

brave/src/main/java/com/linecorp/armeria/client/brave/BraveClient.java

@@ -25,7 +25,6 @@
 import org.slf4j.LoggerFactory;
 
 import com.linecorp.armeria.client.ClientRequestContext;
-import com.linecorp.armeria.client.DefaultClientRequestContext;
 import com.linecorp.armeria.client.HttpClient;
 import com.linecorp.armeria.client.SimpleDecoratingHttpClient;
 import com.linecorp.armeria.common.HttpRequest;
@@ -34,6 +33,7 @@
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext;
 import com.linecorp.armeria.common.logging.ClientConnectionTimings;
+import com.linecorp.armeria.internal.common.RequestContextExtension;
 import com.linecorp.armeria.internal.common.brave.SpanTags;
 
 import brave.Span;
@@ -121,10 +121,10 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
         req = req.withHeaders(newHeaders);
         ctx.updateRequest(req);
 
-        if (currentTraceContext != null && !span.isNoop() && ctx instanceof DefaultClientRequestContext) {
-            final DefaultClientRequestContext defaultCtx = (DefaultClientRequestContext) ctx;
+        final RequestContextExtension ctxExtension = ctx.as(RequestContextExtension.class);
+        if (currentTraceContext != null && !span.isNoop() && ctxExtension != null) {
             // Make the span the current span and run scope decorators when the ctx is pushed.
-            defaultCtx.hook(() -> currentTraceContext.newScope(span.context()));
+            ctxExtension.hook(() -> currentTraceContext.newScope(span.context()));
         }
 
         maybeAddTagsToSpan(ctx, braveReq, span);

brave/src/main/java/com/linecorp/armeria/server/brave/BraveService.java

@@ -23,8 +23,8 @@
 import com.linecorp.armeria.common.HttpRequest;
 import com.linecorp.armeria.common.HttpResponse;
 import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext;
+import com.linecorp.armeria.internal.common.RequestContextExtension;
 import com.linecorp.armeria.internal.common.brave.SpanTags;
-import com.linecorp.armeria.server.DefaultServiceRequestContext;
 import com.linecorp.armeria.server.HttpService;
 import com.linecorp.armeria.server.ServiceRequestContext;
 import com.linecorp.armeria.server.SimpleDecoratingHttpService;
@@ -100,11 +100,10 @@ public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exc
         final HttpServerRequest braveReq = ServiceRequestContextAdapter.asHttpServerRequest(ctx);
         final Span span = handler.handleReceive(braveReq);
 
-        if (currentTraceContext.scopeDecoratorAdded() && !span.isNoop() &&
-            ctx instanceof DefaultServiceRequestContext) {
-            final DefaultServiceRequestContext defaultCtx = (DefaultServiceRequestContext) ctx;
+        final RequestContextExtension ctxExtension = ctx.as(RequestContextExtension.class);
+        if (currentTraceContext.scopeDecoratorAdded() && !span.isNoop() && ctxExtension != null) {
             // Run the scope decorators when the ctx is pushed to the thread local.
-            defaultCtx.hook(() -> currentTraceContext.decorateScope(span.context(),
+            ctxExtension.hook(() -> currentTraceContext.decorateScope(span.context(),
                                                                     SERVICE_REQUEST_DECORATING_SCOPE));
         }
 

core/src/main/java/com/linecorp/armeria/client/ClientRequestContextBuilder.java

@@ -33,6 +33,7 @@
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.logging.ClientConnectionTimings;
 import com.linecorp.armeria.common.util.SystemInfo;
+import com.linecorp.armeria.internal.client.DefaultClientRequestContext;
 import com.linecorp.armeria.internal.common.CancellationScheduler;
 import com.linecorp.armeria.internal.common.CancellationScheduler.CancellationTask;
 

core/src/main/java/com/linecorp/armeria/client/Clients.java

@@ -32,6 +32,7 @@
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.util.SafeCloseable;
 import com.linecorp.armeria.common.util.Unwrappable;
+import com.linecorp.armeria.internal.client.ClientThreadLocalState;
 
 /**
  * Creates a new client that connects to a specified {@link URI}.

core/src/main/java/com/linecorp/armeria/client/DefaultWebClient.java

@@ -16,14 +16,13 @@
 
 package com.linecorp.armeria.client;
 
+import static com.linecorp.armeria.internal.client.ClientUtil.pathWithQuery;
 import static com.linecorp.armeria.internal.common.ArmeriaHttpUtil.concatPaths;
 import static com.linecorp.armeria.internal.common.ArmeriaHttpUtil.isAbsoluteUri;
 import static java.util.Objects.requireNonNull;
 
 import java.net.URI;
 
-import com.google.common.base.Strings;
-
 import com.linecorp.armeria.client.endpoint.EndpointGroup;
 import com.linecorp.armeria.common.ExchangeType;
 import com.linecorp.armeria.common.HttpRequest;
@@ -150,14 +149,4 @@ public RestClient asRestClient() {
     public HttpClient unwrap() {
         return (HttpClient) super.unwrap();
     }
-
-    static String pathWithQuery(URI uri, @Nullable String query) {
-        String path = uri.getRawPath();
-        if (Strings.isNullOrEmpty(path)) {
-            path = query == null ? "/" : "/?" + query;
-        } else if (query != null) {
-            path = path + '?' + query;
-        }
-        return path;
-    }
 }

core/src/main/java/com/linecorp/armeria/client/HttpResponseDecoder.java

@@ -35,6 +35,7 @@
 import com.linecorp.armeria.common.stream.CancelledSubscriptionException;
 import com.linecorp.armeria.common.stream.StreamWriter;
 import com.linecorp.armeria.common.util.Exceptions;
+import com.linecorp.armeria.internal.client.ClientRequestContextExtension;
 import com.linecorp.armeria.internal.common.CancellationScheduler;
 import com.linecorp.armeria.internal.common.CancellationScheduler.CancellationTask;
 import com.linecorp.armeria.internal.common.InboundTrafficController;
@@ -354,9 +355,11 @@ private void cancelAction(@Nullable Throwable cause) {
         private void cancelTimeoutOrLog(@Nullable Throwable cause, boolean cancel) {
 
             CancellationScheduler responseCancellationScheduler = null;
-            if (ctx instanceof DefaultClientRequestContext) {
-                responseCancellationScheduler =
-                        ((DefaultClientRequestContext) ctx).responseCancellationScheduler();
+            if (ctx != null) {
+                final ClientRequestContextExtension ctxExtension = ctx.as(ClientRequestContextExtension.class);
+                if (ctxExtension != null) {
+                    responseCancellationScheduler = ctxExtension.responseCancellationScheduler();
+                }
             }
 
             if (responseCancellationScheduler == null || !responseCancellationScheduler.isFinished()) {
@@ -398,9 +401,13 @@ private void cancelTimeoutOrLog(@Nullable Throwable cause, boolean cancel) {
         }
 
         void initTimeout() {
-            if (ctx instanceof DefaultClientRequestContext) {
+            if (ctx == null) {
+                return;
+            }
+            final ClientRequestContextExtension ctxExtension = ctx.as(ClientRequestContextExtension.class);
+            if (ctxExtension != null) {
                 final CancellationScheduler responseCancellationScheduler =
-                        ((DefaultClientRequestContext) ctx).responseCancellationScheduler();
+                        ctxExtension.responseCancellationScheduler();
                 responseCancellationScheduler.init(
                         ctx.eventLoop(), newCancellationTask(),
                         TimeUnit.MILLISECONDS.toNanos(responseTimeoutMillis), /* server */ false);

core/src/main/java/com/linecorp/armeria/client/UserClient.java

@@ -40,6 +40,7 @@
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.util.AbstractUnwrappable;
 import com.linecorp.armeria.common.util.SystemInfo;
+import com.linecorp.armeria.internal.client.DefaultClientRequestContext;
 
 import io.micrometer.core.instrument.MeterRegistry;
 

core/src/main/java/com/linecorp/armeria/common/RequestContext.java

@@ -47,6 +47,7 @@
 import com.linecorp.armeria.common.logging.RequestLogAccess;
 import com.linecorp.armeria.common.logging.RequestLogBuilder;
 import com.linecorp.armeria.common.util.SafeCloseable;
+import com.linecorp.armeria.common.util.Unwrappable;
 import com.linecorp.armeria.internal.common.JavaVersionSpecific;
 import com.linecorp.armeria.internal.common.RequestContextUtil;
 import com.linecorp.armeria.server.ServiceRequestContext;
@@ -61,7 +62,7 @@
  * A server-side {@link Request} has a {@link ServiceRequestContext} and
  * a client-side {@link Request} has a {@link ClientRequestContext}.
  */
-public interface RequestContext {
+public interface RequestContext extends Unwrappable {
 
     /**
      * Returns the context of the {@link Request} that is being handled in the current thread.

core/src/main/java/com/linecorp/armeria/common/RequestContextWrapper.java

@@ -27,6 +27,7 @@
 import com.linecorp.armeria.common.annotation.Nullable;
 import com.linecorp.armeria.common.logging.RequestLogAccess;
 import com.linecorp.armeria.common.logging.RequestLogBuilder;
+import com.linecorp.armeria.common.util.AbstractUnwrappable;
 import com.linecorp.armeria.server.ServiceRequestContext;
 
 import io.micrometer.core.instrument.MeterRegistry;
@@ -38,180 +39,178 @@
  *
  * @param <T> the self type
  */
-public abstract class RequestContextWrapper<T extends RequestContext> implements RequestContext {
-
-    private final T delegate;
-
+public abstract class RequestContextWrapper<T extends RequestContext>
+        extends AbstractUnwrappable<T> implements RequestContext {
     /**
      * Creates a new instance.
      */
     protected RequestContextWrapper(T delegate) {
-        this.delegate = requireNonNull(delegate, "delegate");
+        super(requireNonNull(delegate, "delegate"));
     }
 
     /**
      * Returns the delegate context.
      */
     protected final T delegate() {
-        return delegate;
+        return unwrap();
     }
 
     @Nullable
     @Override
     public ServiceRequestContext root() {
-        return delegate().root();
+        return unwrap().root();
     }
 
     @Nullable
     @Override
     public <V> V attr(AttributeKey<V> key) {
-        return delegate().attr(key);
+        return unwrap().attr(key);
     }
 
     @Nullable
     @Override
     public <V> V ownAttr(AttributeKey<V> key) {
-        return delegate().ownAttr(key);
+        return unwrap().ownAttr(key);
     }
 
     @Override
     public boolean hasAttr(AttributeKey<?> key) {
-        return delegate().hasAttr(key);
+        return unwrap().hasAttr(key);
     }
 
     @Override
     public boolean hasOwnAttr(AttributeKey<?> key) {
-        return delegate().hasOwnAttr(key);
+        return unwrap().hasOwnAttr(key);
     }
 
     @Override
     public Iterator<Entry<AttributeKey<?>, Object>> attrs() {
-        return delegate().attrs();
+        return unwrap().attrs();
     }
 
     @Override
     public Iterator<Entry<AttributeKey<?>, Object>> ownAttrs() {
-        return delegate().ownAttrs();
+        return unwrap().ownAttrs();
     }
 
     @Override
     public <V> V setAttr(AttributeKey<V> key, @Nullable V value) {
-        return delegate().setAttr(key, value);
+        return unwrap().setAttr(key, value);
     }
 
     @Override
     public HttpRequest request() {
-        return delegate().request();
+        return unwrap().request();
     }
 
     @Nullable
     @Override
     public RpcRequest rpcRequest() {
-        return delegate().rpcRequest();
+        return unwrap().rpcRequest();
     }
 
     @Override
     public void updateRequest(HttpRequest req) {
-        delegate().updateRequest(req);
+        unwrap().updateRequest(req);
     }
 
     @Override
     public void updateRpcRequest(RpcRequest rpcReq) {
-        delegate().updateRpcRequest(rpcReq);
+        unwrap().updateRpcRequest(rpcReq);
     }
 
     @Override
     public SessionProtocol sessionProtocol() {
-        return delegate().sessionProtocol();
+        return unwrap().sessionProtocol();
     }
 
     @Nullable
     @Override
     public <A extends SocketAddress> A remoteAddress() {
-        return delegate().remoteAddress();
+        return unwrap().remoteAddress();
     }
 
     @Nullable
     @Override
     public <A extends SocketAddress> A localAddress() {
-        return delegate().localAddress();
+        return unwrap().localAddress();
     }
 
     @Nullable
     @Override
     public SSLSession sslSession() {
-        return delegate().sslSession();
+        return unwrap().sslSession();
     }
 
     @Override
     public RequestId id() {
-        return delegate().id();
+        return unwrap().id();
     }
 
     @Override
     public HttpMethod method() {
-        return delegate().method();
+        return unwrap().method();
     }
 
     @Override
     public String path() {
-        return delegate().path();
+        return unwrap().path();
     }
 
     @Override
     public String decodedPath() {
-        return delegate().decodedPath();
+        return unwrap().decodedPath();
     }
 
     @Override
     public String query() {
-        return delegate().query();
+        return unwrap().query();
     }
 
     @Override
     public RequestLogAccess log() {
-        return delegate().log();
+        return unwrap().log();
     }
 
     @Override
     public RequestLogBuilder logBuilder() {
-        return delegate().logBuilder();
+        return unwrap().logBuilder();
     }
 
     @Override
     public MeterRegistry meterRegistry() {
-        return delegate().meterRegistry();
+        return unwrap().meterRegistry();
     }
 
     @Override
     public void cancel(Throwable cause) {
-        delegate().cancel(cause);
+        unwrap().cancel(cause);
     }
 
     @Override
     public void cancel() {
-        delegate().cancel();
+        unwrap().cancel();
     }
 
     @Override
     public void timeoutNow() {
-        delegate().timeoutNow();
+        unwrap().timeoutNow();
     }
 
     @Override
     @Nullable
     public Throwable cancellationCause() {
-        return delegate().cancellationCause();
+        return unwrap().cancellationCause();
     }
 
     @Override
     public ContextAwareEventLoop eventLoop() {
-        return delegate().eventLoop();
+        return unwrap().eventLoop();
     }
 
     @Override
     public ByteBufAllocator alloc() {
-        return delegate().alloc();
+        return unwrap().alloc();
     }
 
     @Override
@@ -221,6 +220,6 @@ public ExchangeType exchangeType() {
 
     @Override
     public String toString() {
-        return delegate().toString();
+        return unwrap().toString();
     }
 }