Support cancelling in flight REST requests

Fixes: #41971
quarkusio · Feb 20, 2025 · df4da57 · df4da57
1 parent 6bc9e3c
commit df4da57
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/...e/src/main/java/org/jboss/resteasy/reactive/client/handlers/ClientSendRequestHandler.java b/...e/src/main/java/org/jboss/resteasy/reactive/client/handlers/ClientSendRequestHandler.java
@@ -98,6 +98,8 @@ public void handle(RestClientRequestContext requestContext) {
         future.subscribe().with(new Consumer<>() {
             @Override
             public void accept(HttpClientRequest httpClientRequest) {
+                requestContext.setHttpClientRequest(httpClientRequest);
+
                 // adapt headers to HTTP/2 depending on the underlying HTTP connection
                 ClientSendRequestHandler.this.adaptRequest(httpClientRequest);
 

diff --git a/...tive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/impl/UniInvoker.java b/...tive/client/runtime/src/main/java/org/jboss/resteasy/reactive/client/impl/UniInvoker.java
@@ -1,6 +1,8 @@
 package org.jboss.resteasy.reactive.client.impl;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -10,10 +12,11 @@
 import jakarta.ws.rs.core.Response;
 
 import io.smallrye.mutiny.Uni;
+import io.vertx.core.http.HttpClientRequest;
 
 public class UniInvoker extends AbstractRxInvoker<Uni<?>> {
 
-    private InvocationBuilderImpl invocationBuilder;
+    private final InvocationBuilderImpl invocationBuilder;
 
     public UniInvoker(InvocationBuilderImpl invocationBuilder) {
         this.invocationBuilder = invocationBuilder;
@@ -22,10 +25,16 @@ public UniInvoker(InvocationBuilderImpl invocationBuilder) {
     @Override
     public <R> Uni<R> method(String name, Entity<?> entity, GenericType<R> responseType) {
         AsyncInvokerImpl invoker = (AsyncInvokerImpl) invocationBuilder.rx();
+        AtomicReference<RestClientRequestContext> restClientRequestContextRef = new AtomicReference<>();
         return Uni.createFrom().completionStage(new Supplier<CompletionStage<R>>() {
             @Override
             public CompletionStage<R> get() {
-                return invoker.method(name, entity, responseType);
+                RestClientRequestContext restClientRequestContext = invoker.performRequestInternal(name, entity,
+                        responseType == null ? new GenericType<>(String.class) : responseType,
+                        true);
+                restClientRequestContextRef.set(restClientRequestContext);
+                CompletableFuture response = restClientRequestContext.getResult();
+                return invoker.mapResponse(response, responseType == null ? String.class : responseType.getRawType());
             }
         }).onFailure().transform(new Function<>() {
             @Override
@@ -35,6 +44,18 @@ public Throwable apply(Throwable t) {
                 }
                 return t;
             }
+        }).onCancellation().invoke(new Runnable() {
+            @Override
+            public void run() {
+                // be very defensive here as things could have been nulled out when the application is being torn down
+                RestClientRequestContext restClientRequestContext = restClientRequestContextRef.get();
+                if (restClientRequestContext != null) {
+                    HttpClientRequest httpClientRequest = restClientRequestContext.getHttpClientRequest();
+                    if (httpClientRequest != null) {
+                        httpClientRequest.reset();
+                    }
+                }
+            }
         });
     }