Fixes request context lost in Infinispan cache get/getasync

Author: karesti

extensions/infinispan-cache/runtime/src/main/java/io/quarkus/cache/infinispan/runtime/InfinispanCacheImpl.java

@@ -11,10 +11,11 @@
 import java.util.function.Function;
 import java.util.function.Predicate;
 
+import org.eclipse.microprofile.context.ManagedExecutor;
+import org.eclipse.microprofile.context.ThreadContext;
 import org.infinispan.client.hotrod.RemoteCache;
 import org.infinispan.client.hotrod.impl.protocol.Codec27;
 import org.infinispan.commons.util.NullValue;
-import org.infinispan.commons.util.concurrent.CompletionStages;
 import org.reactivestreams.FlowAdapters;
 
 import io.quarkus.arc.Arc;
@@ -36,6 +37,8 @@
  */
 public class InfinispanCacheImpl extends AbstractCache implements Cache {
 
+    private final ThreadContext threadContext;
+    private final ManagedExecutor managedExecutor;
     private final RemoteCache remoteCache;
     private final InfinispanCacheInfo cacheInfo;
     private final Map<Object, CompletableFuture> computationResults = new ConcurrentHashMap<>();
@@ -47,6 +50,8 @@ public InfinispanCacheImpl(InfinispanCacheInfo cacheInfo, RemoteCache remoteCach
         this.remoteCache = remoteCache;
         this.lifespan = cacheInfo.lifespan.map(l -> l.toMillis()).orElse(-1L);
         this.maxIdle = cacheInfo.maxIdle.map(m -> m.toMillis()).orElse(-1L);
+        this.threadContext = Arc.container().select(ThreadContext.class).get();
+        this.managedExecutor = Arc.container().select(ManagedExecutor.class).get();
     }
 
     public InfinispanCacheImpl(InfinispanCacheInfo cacheInfo,
@@ -81,16 +86,11 @@ private <T> T decodeNull(Object value) {
 
     @Override
     public <K, V> Uni<V> get(K key, Function<K, V> valueLoader) {
-        return Uni.createFrom()
-                .completionStage(() -> CompletionStages.handleAndCompose(remoteCache.getAsync(key), (v1, ex1) -> {
-                    if (ex1 != null) {
-                        return CompletableFuture.failedFuture(ex1);
-                    }
-
+        CompletableFuture<V> infinispanGet = threadContext.withContextCapture(remoteCache.getAsync(key))
+                .thenApplyAsync(v1 -> {
                     if (v1 != null) {
                         return CompletableFuture.completedFuture(decodeNull(v1));
                     }
-
                     CompletableFuture<V> resultAsync = new CompletableFuture<>();
                     CompletableFuture<V> computedValue = computationResults.putIfAbsent(key, resultAsync);
                     if (computedValue != null) {
@@ -111,18 +111,15 @@ public <K, V> Uni<V> get(K key, Function<K, V> valueLoader) {
                                 computationResults.remove(key);
                             });
                     return resultAsync;
-                }));
+                }, managedExecutor).thenCompose(f -> f);
+
+        return Uni.createFrom().completionStage(infinispanGet);
     }
 
     @Override
     public <K, V> Uni<V> getAsync(K key, Function<K, Uni<V>> valueLoader) {
         Context context = Vertx.currentContext();
-
-        return Uni.createFrom().completionStage(CompletionStages.handleAndCompose(remoteCache.getAsync(key), (v1, ex1) -> {
-            if (ex1 != null) {
-                return CompletableFuture.failedFuture(ex1);
-            }
-
+        CompletableFuture<V> infinispanGet = threadContext.withContextCapture(remoteCache.getAsync(key)).thenApplyAsync(v1 -> {
             if (v1 != null) {
                 return CompletableFuture.completedFuture(decodeNull(v1));
             }
@@ -152,7 +149,9 @@ public <K, V> Uni<V> getAsync(K key, Function<K, Uni<V>> valueLoader) {
                         }
                     });
             return resultAsync;
-        })).emitOn(new Executor() {
+        }, managedExecutor).thenCompose(f -> f);
+
+        return Uni.createFrom().completionStage(infinispanGet).emitOn(new Executor() {
             // We need make sure we go back to the original context when the cache value is computed.
             // Otherwise, we would always emit on the context having computed the value, which could
             // break the duplicated context isolation.

integration-tests/infinispan-cache/src/main/java/io/quarkus/it/cache/infinispan/ClientRequestFilter.java

@@ -0,0 +1,23 @@
+package io.quarkus.it.cache.infinispan;
+
+import jakarta.inject.Inject;
+import jakarta.ws.rs.client.ClientRequestContext;
+import jakarta.ws.rs.ext.Provider;
+
+@Provider
+public class ClientRequestFilter implements jakarta.ws.rs.client.ClientRequestFilter {
+
+    ClientRequestService requestService;
+
+    @Inject
+    public ClientRequestFilter(ClientRequestService requestService) {
+        this.requestService = requestService;
+    }
+
+    @Override
+    public void filter(ClientRequestContext requestContext) {
+        if (requestService != null && requestService.data() != null) {
+            requestContext.getHeaders().add("extra", requestService.data());
+        }
+    }
+}

integration-tests/infinispan-cache/src/main/java/io/quarkus/it/cache/infinispan/ClientRequestService.java

@@ -0,0 +1,17 @@
+package io.quarkus.it.cache.infinispan;
+
+import jakarta.enterprise.context.RequestScoped;
+
+@RequestScoped
+public class ClientRequestService {
+    String data;
+
+    public String data() {
+        return data;
+    }
+
+    public ClientRequestService setData(String data) {
+        this.data = data;
+        return this;
+    }
+}

integration-tests/infinispan-cache/src/main/java/io/quarkus/it/cache/infinispan/ExpensiveResource.java

@@ -2,39 +2,71 @@
 
 import java.util.concurrent.atomic.AtomicInteger;
 
+import jakarta.inject.Inject;
+import jakarta.ws.rs.DELETE;
 import jakarta.ws.rs.GET;
-import jakarta.ws.rs.POST;
 import jakarta.ws.rs.Path;
 import jakarta.ws.rs.PathParam;
 import jakarta.ws.rs.QueryParam;
+import jakarta.ws.rs.core.Response;
 
 import org.infinispan.protostream.GeneratedSchema;
 import org.infinispan.protostream.annotations.Proto;
 import org.infinispan.protostream.annotations.ProtoSchema;
 
+import io.quarkus.cache.CacheInvalidate;
 import io.quarkus.cache.CacheInvalidateAll;
 import io.quarkus.cache.CacheKey;
 import io.quarkus.cache.CacheResult;
+import io.smallrye.mutiny.Uni;
 
 @Path("/expensive-resource")
 public class ExpensiveResource {
 
     private final AtomicInteger invocations = new AtomicInteger(0);
 
+    @Inject
+    ClientRequestService requestService;
+
     @GET
     @Path("/{keyElement1}/{keyElement2}/{keyElement3}")
     @CacheResult(cacheName = "expensiveResourceCache")
     public ExpensiveResponse getExpensiveResponse(@PathParam("keyElement1") @CacheKey String keyElement1,
-            @PathParam("keyElement2") @CacheKey String keyElement2, @PathParam("keyElement3") @CacheKey String keyElement3,
+            @PathParam("keyElement2") @CacheKey String keyElement2,
+            @PathParam("keyElement3") @CacheKey String keyElement3,
             @QueryParam("foo") String foo) {
         invocations.incrementAndGet();
+        requestService.setData("getExpensiveResponse " + foo);
         return new ExpensiveResponse(keyElement1 + " " + keyElement2 + " " + keyElement3 + " too!");
     }
 
-    @POST
+    @GET
+    @Path("/async/{keyElement1}/{keyElement2}/{keyElement3}")
+    @CacheResult(cacheName = "expensiveResourceCache")
+    public Uni<ExpensiveResponse> getExpensiveResponseAsync(@PathParam("keyElement1") @CacheKey String keyElement1,
+            @PathParam("keyElement2") @CacheKey String keyElement2,
+            @PathParam("keyElement3") @CacheKey String keyElement3,
+            @QueryParam("foo") String foo) {
+        invocations.incrementAndGet();
+        requestService.setData("getExpensiveResponseAsync " + foo);
+        return Uni.createFrom()
+                .item(new ExpensiveResponse(keyElement1 + " " + keyElement2 + " " + keyElement3 + " async too!"));
+    }
+
+    @DELETE
+    @Path("/{keyElement1}/{keyElement2}/{keyElement3}")
+    @CacheInvalidate(cacheName = "expensiveResourceCache")
+    public Response resetExpensiveResponse(@PathParam("keyElement1") @CacheKey String keyElement1,
+            @PathParam("keyElement2") @CacheKey String keyElement2, @PathParam("keyElement3") @CacheKey String keyElement3,
+            @QueryParam("foo") String foo) {
+        requestService.setData("invalidate");
+        return Response.ok().build();
+    }
+
+    @DELETE
     @CacheInvalidateAll(cacheName = "expensiveResourceCache")
     public void invalidateAll() {
-
+        requestService.setData("invalidateAll");
     }
 
     @GET

integration-tests/infinispan-cache/src/main/java/io/quarkus/it/cache/infinispan/SunriseRestServerResource.java

@@ -1,6 +1,9 @@
 package io.quarkus.it.cache.infinispan;
 
+import java.util.concurrent.atomic.AtomicInteger;
+
 import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
 import jakarta.ws.rs.DELETE;
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.Path;
@@ -12,30 +15,37 @@
 @Path("sunrise")
 public class SunriseRestServerResource {
 
-    private int sunriseTimeInvocations;
+    private final AtomicInteger invocations = new AtomicInteger(0);
+
+    @Inject
+    ClientRequestService requestService;
 
     @GET
     @Path("time/{city}")
     public String getSunriseTime(@RestPath String city, @RestQuery String date) {
-        sunriseTimeInvocations++;
+        invocations.incrementAndGet();
+        requestService.setData(city);
         return "2020-12-20T10:15:30";
     }
 
     @GET
     @Path("invocations")
     public Integer getSunriseTimeInvocations() {
-        return sunriseTimeInvocations;
+        requestService.setData("invocations");
+        return invocations.get();
     }
 
     @DELETE
     @Path("invalidate/{city}")
     public void invalidate(@RestPath String city, @RestQuery String notPartOfTheCacheKey, @RestQuery String date) {
         // Do nothing. We only need to test the caching annotation on the client side.
+        requestService.setData("invalidate " + city);
     }
 
     @DELETE
     @Path("invalidate")
     public void invalidateAll() {
         // Do nothing. We only need to test the caching annotation on the client side.
+        requestService.setData("invalidate all");
     }
 }

integration-tests/infinispan-cache/src/test/java/io/quarkus/it/cache/infinispan/CacheTest.java

@@ -12,22 +12,46 @@ public class CacheTest {
 
     @Test
     public void testCache() {
-        runExpensiveRequest();
-        runExpensiveRequest();
-        runExpensiveRequest();
+        when().get("/expensive-resource/invocations").then().statusCode(200).body(is("0"));
+        runGetExpensiveRequest();
+        runGetExpensiveRequest();
+        runGetExpensiveRequest();
         when().get("/expensive-resource/invocations").then().statusCode(200).body(is("1"));
 
+        runDeleteExpensiveRequest();
+
+        runGetExpensiveRequestAsync();
+        runGetExpensiveRequestAsync();
+        runGetExpensiveRequestAsync();
+
+        when().get("/expensive-resource/invocations").then().statusCode(200).body(is("2"));
+
         when()
-                .post("/expensive-resource")
+                .delete("/expensive-resource")
                 .then()
                 .statusCode(204);
     }
 
-    private void runExpensiveRequest() {
+    private void runGetExpensiveRequest() {
         when()
                 .get("/expensive-resource/I/love/Quarkus?foo=bar")
                 .then()
                 .statusCode(200)
                 .body("result", is("I love Quarkus too!"));
     }
+
+    private void runGetExpensiveRequestAsync() {
+        when()
+                .get("/expensive-resource/async/I/love/Quarkus?foo=bar")
+                .then()
+                .statusCode(200)
+                .body("result", is("I love Quarkus async too!"));
+    }
+
+    private void runDeleteExpensiveRequest() {
+        when()
+                .delete("/expensive-resource/I/love/Quarkus?foo=bar")
+                .then()
+                .statusCode(200);
+    }
 }