optimization

jrhee17 · Dec 4, 2024 · 2b50dd6 · 2b50dd6
1 parent 3beeb33
commit 2b50dd6
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 33 deletions.
diff --git a/core/src/main/java/com/linecorp/armeria/client/DefaultRequestExecution.java b/core/src/main/java/com/linecorp/armeria/client/DefaultRequestExecution.java
diff --git a/core/src/main/java/com/linecorp/armeria/client/RequestExecution.java b/core/src/main/java/com/linecorp/armeria/client/RequestExecution.java
@@ -32,7 +32,9 @@ public interface RequestExecution {
      * TBU.
      */
     static RequestExecution of(ClientRequestContext ctx) {
-        return new DefaultRequestExecution(ctx);
+        final ClientRequestContextExtension ctxExt = ctx.as(ClientRequestContextExtension.class);
+        checkArgument(ctxExt != null, "ctx (%s) should be created from 'ClientRequestContextBuilder'", ctx);
+        return ctxExt;
     }
 
     /**

diff --git a/core/src/main/java/com/linecorp/armeria/internal/client/ClientRequestContextExtension.java b/core/src/main/java/com/linecorp/armeria/internal/client/ClientRequestContextExtension.java
@@ -18,10 +18,14 @@
 
 import java.util.concurrent.CompletableFuture;
 
+import com.linecorp.armeria.client.Client;
 import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.RequestExecution;
 import com.linecorp.armeria.client.endpoint.EndpointGroup;
 import com.linecorp.armeria.common.HttpHeaderNames;
 import com.linecorp.armeria.common.HttpHeaders;
+import com.linecorp.armeria.common.Request;
+import com.linecorp.armeria.common.Response;
 import com.linecorp.armeria.common.logging.RequestLog;
 import com.linecorp.armeria.internal.common.CancellationScheduler;
 import com.linecorp.armeria.internal.common.RequestContextExtension;
@@ -30,7 +34,8 @@
  * This class exposes extension methods for {@link ClientRequestContext}
  * which are used internally by Armeria but aren't intended for public usage.
  */
-public interface ClientRequestContextExtension extends ClientRequestContext, RequestContextExtension {
+public interface ClientRequestContextExtension extends ClientRequestContext, RequestContextExtension,
+                                                       RequestExecution {
 
     /**
      * Returns the {@link CancellationScheduler} used to schedule a response timeout.
@@ -75,4 +80,14 @@ public interface ClientRequestContextExtension extends ClientRequestContext, Req
     HttpHeaders internalRequestHeaders();
 
     boolean initializationTriggered();
+
+    @Override
+    default ClientRequestContext ctx() {
+        return this;
+    }
+
+    @Override
+    default <I extends Request, O extends Response> O execute(Client<I, O> delegate, I req) throws Exception {
+        return RequestExecution.super.execute(delegate, req);
+    }
 }