Made the cache lookup be either async or not

bbakerman · bbakerman · commit c42baa71fec5 · 2025-09-28T18:59:14.000+10:00
diff --git a/src/main/java/graphql/execution/preparsed/caching/CachingDocumentProvider.java b/src/main/java/graphql/execution/preparsed/caching/CachingDocumentProvider.java
@@ -10,6 +10,8 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
 
+import static graphql.execution.Async.toCompletableFuture;
+
 /**
  * The CachingDocumentProvider allows previously parsed and validated operations to be cached and
  * hence re-used.  This can lead to significant time savings, especially for large operations.
@@ -56,11 +58,11 @@ public DocumentCache getDocumentCache() {
     public CompletableFuture<PreparsedDocumentEntry> getDocumentAsync(ExecutionInput executionInput, Function<ExecutionInput, PreparsedDocumentEntry> parseAndValidateFunction) {
         if (documentCache.isNoop()) {
             // saves creating keys and doing a lookup that will just call this function anyway
-            return CompletableFuture.completedFuture(parseAndValidateFunction.apply(executionInput));
+            return toCompletableFuture(parseAndValidateFunction.apply(executionInput));
         }
         DocumentCache.DocumentCacheKey cacheKey = new DocumentCache.DocumentCacheKey(executionInput.getQuery(), executionInput.getOperationName(), executionInput.getLocale());
-        PreparsedDocumentEntry cacheEntry = documentCache.get(cacheKey, key -> parseAndValidateFunction.apply(executionInput));
-        return CompletableFuture.completedFuture(cacheEntry);
+        Object cacheEntry = documentCache.get(cacheKey, key -> parseAndValidateFunction.apply(executionInput));
+        return toCompletableFuture(cacheEntry);
     }
 
 }
diff --git a/src/main/java/graphql/execution/preparsed/caching/DocumentCache.java b/src/main/java/graphql/execution/preparsed/caching/DocumentCache.java
@@ -1,5 +1,6 @@
 package graphql.execution.preparsed.caching;
 
+import graphql.DuckTyped;
 import graphql.PublicApi;
 import graphql.execution.preparsed.PreparsedDocumentEntry;
 import org.jspecify.annotations.NullMarked;
@@ -17,14 +18,18 @@
 @NullMarked
 public interface DocumentCache {
     /**
-     * Called to get a document that has previously been parsed ad validated.
+     * Called to get a document that has previously been parsed ad validated.  The return value of this method
+     * can be either a {@link PreparsedDocumentEntry} or a {@link java.util.concurrent.CompletableFuture} promise
+     * to a {@link PreparsedDocumentEntry}.  This allows caches that are in memory to return direct values OR
+     * if the cache is distributed, it can return a promise to a value.
      *
      * @param key             the cache key
      * @param mappingFunction if the value is missing in cache this function can be called to create a value
      *
-     * @return a non null document entry
+     * @return a non-null {@link PreparsedDocumentEntry} or a promise to one via a {@link java.util.concurrent.CompletableFuture}
      */
-    PreparsedDocumentEntry get(DocumentCacheKey key, Function<DocumentCacheKey, PreparsedDocumentEntry> mappingFunction);
+    @DuckTyped(shape = "PreparsedDocumentEntry | CompletableFuture<PreparsedDocumentEntry")
+    Object get(DocumentCacheKey key, Function<DocumentCacheKey, PreparsedDocumentEntry> mappingFunction);
 
     /**
      * @return true if the cache in fact does no caching otherwise false.  This helps the implementation optimise how the cache is used or not.
@@ -35,6 +40,7 @@ public interface DocumentCache {
      * Called to clear the cache.  If your implementation doesn't support this, then just no op the method
      */
     void invalidateAll();
+
     /**
      * This represents the key to the document cache
      */
diff --git a/src/test/groovy/graphql/execution/preparsed/caching/CachingDocumentProviderTest.groovy b/src/test/groovy/graphql/execution/preparsed/caching/CachingDocumentProviderTest.groovy
@@ -1,6 +1,5 @@
 package graphql.execution.preparsed.caching
 
-
 import com.github.benmanes.caffeine.cache.Caffeine
 import com.github.benmanes.caffeine.cache.Ticker
 import graphql.ExecutionInput
@@ -12,9 +11,11 @@ import graphql.parser.Parser
 import spock.lang.Specification
 
 import java.time.Duration
+import java.util.concurrent.CompletableFuture
 import java.util.function.Function
 
 import static graphql.ExecutionInput.newExecutionInput
+import static graphql.execution.preparsed.caching.DocumentCache.DocumentCacheKey
 
 class CachingDocumentProviderTest extends Specification {
     private String heroQuery1
@@ -160,11 +161,14 @@ class CachingDocumentProviderTest extends Specification {
 
         def cache = new DocumentCache() {
             // not really useful in production since its unbounded
-            def map = new HashMap<DocumentCache.DocumentCacheKey,PreparsedDocumentEntry>()
+            def map = new HashMap<DocumentCacheKey, PreparsedDocumentEntry>()
 
             @Override
-            PreparsedDocumentEntry get(DocumentCache.DocumentCacheKey key, Function<DocumentCache.DocumentCacheKey, PreparsedDocumentEntry> mappingFunction) {
-                return map.computeIfAbsent(key,mappingFunction)
+            Object get(DocumentCacheKey key, Function<DocumentCacheKey, PreparsedDocumentEntry> mappingFunction) {
+                // we can have async values
+                return CompletableFuture.supplyAsync {
+                    return map.computeIfAbsent(key, mappingFunction)
+                }
             }
 
             @Override
@@ -289,7 +293,7 @@ class CachingDocumentProviderTest extends Specification {
         def caffeineCache = Caffeine.newBuilder()
                 .ticker(ticker)
                 .expireAfterAccess(Duration.ofMinutes(2))
-                .<DocumentCache.DocumentCacheKey, PreparsedDocumentEntry> build()
+                .<DocumentCacheKey, PreparsedDocumentEntry> build()
         def documentCache = new CaffeineDocumentCache(caffeineCache)
 
         // note this is a custom caffeine instance pass in
