Chapter 35: Cache GraphQL responses per user (#11)

* Cache graphql responses per user

* Propagate Correatlion ID with TaskExecutor

* Reuse executor factory task executor bean

* Pass MDC context with AsyncTaskExecutor

Author: philip-jvm

pom.xml

@@ -7,7 +7,7 @@
   <parent>
     <groupId>org.springframework.boot</groupId>
     <artifactId>spring-boot-starter-parent</artifactId>
-    <version>2.3.3.RELEASE</version>
+    <version>2.6.1</version>
   </parent>
   <groupId>org.example</groupId>
   <artifactId>learn-graphql-java</artifactId>
@@ -19,13 +19,21 @@
   <properties>
     <java.version>11</java.version>
     <pmd.version>3.11.0</pmd.version>
-    <graphql.version>11.0.0</graphql.version>
+    <graphql.version>12.0.0</graphql.version>
     <!-- 16.0.0 not on central -->
-    <graphql.extended.scalars.version>15.0.0</graphql.extended.scalars.version>
+    <graphql.extended.scalars.version>17.0</graphql.extended.scalars.version>
     <graphql.extended.validation.version>16.0.0</graphql.extended.validation.version>
   </properties>
 
   <dependencies>
+    <dependency>
+      <groupId>org.springframework.boot</groupId>
+      <artifactId>spring-boot-starter-web</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>com.github.ben-manes.caffeine</groupId>
+      <artifactId>caffeine</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.graphql-java-kickstart</groupId>
       <artifactId>graphql-spring-boot-starter</artifactId>
@@ -66,18 +74,6 @@
       <artifactId>lombok</artifactId>
       <scope>provided</scope>
     </dependency>
-    <dependency>
-      <groupId>com.graphql-java-kickstart</groupId>
-      <artifactId>playground-spring-boot-starter</artifactId>
-      <version>${graphql.version}</version>
-      <scope>runtime</scope>
-    </dependency>
-    <dependency>
-      <groupId>com.graphql-java-kickstart</groupId>
-      <artifactId>voyager-spring-boot-starter</artifactId>
-      <version>${graphql.version}</version>
-      <scope>runtime</scope>
-    </dependency>
     <dependency>
       <groupId>io.projectreactor</groupId>
       <artifactId>reactor-core</artifactId>

src/main/java/com/learn/graphql/cache/RequestKey.java

@@ -0,0 +1,12 @@
+package com.learn.graphql.cache;
+
+import java.util.List;
+import lombok.Value;
+
+@Value
+public class RequestKey {
+
+  String userId;
+  List<String> queries;
+
+}

src/main/java/com/learn/graphql/cache/ResponseCacheManager.java

@@ -0,0 +1,59 @@
+package com.learn.graphql.cache;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.learn.graphql.config.security.GraphQLSecurityConfig;
+import graphql.kickstart.execution.input.GraphQLInvocationInput;
+import graphql.kickstart.servlet.cache.CachedResponse;
+import graphql.kickstart.servlet.cache.GraphQLResponseCacheManager;
+import javax.servlet.http.HttpServletRequest;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.stereotype.Component;
+
+/**
+ * Response Query Caching: Chapter 35
+ */
+@Slf4j
+@Component
+@RequiredArgsConstructor
+public class ResponseCacheManager implements GraphQLResponseCacheManager {
+
+  private final Cache<RequestKey, CachedResponse> responseCache;
+
+  @Override
+  public CachedResponse get(HttpServletRequest request, GraphQLInvocationInput invocationInput) {
+    return responseCache.getIfPresent(getRequestKey(request, invocationInput));
+  }
+
+  @Override
+  public boolean isCacheable(HttpServletRequest request, GraphQLInvocationInput invocationInput) {
+    // Do not cache introspection query
+    return invocationInput.getQueries()
+        .stream()
+        .noneMatch(this::isIntrospectionQuery);
+  }
+
+  @Override
+  public void put(HttpServletRequest request, GraphQLInvocationInput invocationInput,
+      CachedResponse cachedResponse) {
+    responseCache.put(getRequestKey(request, invocationInput), cachedResponse);
+  }
+
+  private RequestKey getRequestKey(HttpServletRequest request,
+      GraphQLInvocationInput invocationInput) {
+    return new RequestKey(getUserId(request), invocationInput.getQueries());
+  }
+
+  private String getUserId(HttpServletRequest request) {
+    var userId = request.getHeader(GraphQLSecurityConfig.USER_ID_PRE_AUTH_HEADER);
+    if (userId == null) {
+      throw new IllegalArgumentException("User Id is null. Cannot read from ResponseCacheManager.");
+    }
+    return userId;
+  }
+
+  private boolean isIntrospectionQuery(String query) {
+    return query.contains("Introspection");
+  }
+
+}

src/main/java/com/learn/graphql/config/AsyncExecutorConfig.java

@@ -0,0 +1,21 @@
+package com.learn.graphql.config;
+
+import com.learn.graphql.util.ExecutorFactory;
+import java.util.concurrent.Executor;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration
+public class AsyncExecutorConfig {
+
+  @Bean
+  public Executor balanceExecutor(ExecutorFactory executorFactory) {
+    return executorFactory.newExecutor();
+  }
+
+  @Bean
+  public Executor bankAccountExecutor(ExecutorFactory executorFactory) {
+    return executorFactory.newExecutor();
+  }
+
+}

src/main/java/com/learn/graphql/config/CacheConfig.java

@@ -0,0 +1,27 @@
+package com.learn.graphql.config;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.learn.graphql.cache.RequestKey;
+import graphql.kickstart.servlet.cache.CachedResponse;
+import java.time.Duration;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+@Slf4j
+@Configuration
+public class CacheConfig {
+
+  @Bean
+  public Cache<RequestKey, CachedResponse> responseCache() {
+    return Caffeine.newBuilder()
+        .expireAfterWrite(Duration.ofMinutes(1L))
+        .maximumSize(100L)
+        .removalListener((key, value, cause) ->
+            log.info("Key {} with value {} was removed from the response cache. Cause {}",
+                key, value, cause))
+        .build();
+  }
+
+}

src/main/java/com/learn/graphql/config/ScalarConfig.java

@@ -1,10 +1,14 @@
 package com.learn.graphql.config.security;
 
+import static com.learn.graphql.config.security.GraphQLSecurityConfig.CORRELATION_ID;
+
 import graphql.kickstart.execution.subscriptions.SubscriptionSession;
 import graphql.kickstart.execution.subscriptions.apollo.ApolloSubscriptionConnectionListener;
 import graphql.kickstart.execution.subscriptions.apollo.OperationMessage;
 import java.util.Map;
+import java.util.UUID;
 import lombok.extern.slf4j.Slf4j;
+import org.slf4j.MDC;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.context.SecurityContextHolder;
 import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
@@ -59,13 +63,15 @@ If a session has a frame available, the session will be passed to another thread
 
     var token = new PreAuthenticatedAuthenticationToken(userId, null, grantedAuthorities);
     session.getUserProperties().put(AUTHENTICATION, token);
+    session.getUserProperties().put(CORRELATION_ID, UUID.randomUUID().toString());
   }
 
   @Override
   public void onStart(SubscriptionSession session, OperationMessage message) {
     log.info("onStart with payload {}", message.getPayload());
     var authentication = (Authentication) session.getUserProperties().get(AUTHENTICATION);
     SecurityContextHolder.getContext().setAuthentication(authentication);
+    MDC.put(CORRELATION_ID, (String) session.getUserProperties().get(CORRELATION_ID));
   }
 
   @Override
@@ -76,6 +82,7 @@ public void onStop(SubscriptionSession session, OperationMessage message) {
   @Override
   public void onTerminate(SubscriptionSession session, OperationMessage message) {
     log.info("onTerminate with payload {}", message.getPayload());
+    MDC.clear();
   }
 
 }

src/main/java/com/learn/graphql/config/security/AuthenticationConnectionListener.java

@@ -37,6 +37,7 @@ public class GraphQLSecurityConfig extends WebSecurityConfigurerAdapter {
    */
   public static final String USER_ID_PRE_AUTH_HEADER = "user_id";
   public static final String USER_ROLES_PRE_AUTH_HEADER = "user_roles";
+  public static final String CORRELATION_ID = "correlation_id";
 
   private final PreAuthenticatedAuthenticationProvider preAuthenticatedAuthenticationProvider;
 

src/main/java/com/learn/graphql/config/security/GraphQLSecurityConfig.java

@@ -1,16 +1,20 @@
 package com.learn.graphql.context;
 
+import static com.learn.graphql.config.security.GraphQLSecurityConfig.CORRELATION_ID;
+
 import com.learn.graphql.context.dataloader.DataLoaderRegistryFactory;
 import graphql.kickstart.execution.context.GraphQLContext;
 import graphql.kickstart.servlet.context.DefaultGraphQLServletContext;
 import graphql.kickstart.servlet.context.DefaultGraphQLWebSocketContext;
 import graphql.kickstart.servlet.context.GraphQLServletContextBuilder;
+import java.util.UUID;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import javax.websocket.Session;
 import javax.websocket.server.HandshakeRequest;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.slf4j.MDC;
 import org.springframework.stereotype.Component;
 
 @Slf4j
@@ -24,6 +28,8 @@ public class CustomGraphQLContextBuilder implements GraphQLServletContextBuilder
   public GraphQLContext build(HttpServletRequest httpServletRequest,
       HttpServletResponse httpServletResponse) {
 
+    MDC.put(CORRELATION_ID, UUID.randomUUID().toString());
+
     var userId = httpServletRequest.getHeader("user_id");
 
     var context = DefaultGraphQLServletContext.createServletContext()

src/main/java/com/learn/graphql/context/CustomGraphQLContextBuilder.java

@@ -1,7 +1,6 @@
 package com.learn.graphql.context.dataloader;
 
 import com.learn.graphql.service.BalanceService;
-import com.learn.graphql.util.ExecutorFactory;
 import java.math.BigDecimal;
 import java.util.Map;
 import java.util.Set;
@@ -18,10 +17,10 @@
 @RequiredArgsConstructor
 public class DataLoaderRegistryFactory {
 
-  private final BalanceService balanceService;
-
   public static final String BALANCE_DATA_LOADER = "BALANCE_DATA_LOADER";
-  private static final Executor balanceThreadPool = ExecutorFactory.newExecutor();
+
+  private final BalanceService balanceService;
+  private final Executor balanceExecutor;
 
   public DataLoaderRegistry create(String userId) {
     var registry = new DataLoaderRegistry();
@@ -34,7 +33,7 @@ private DataLoader<UUID, BigDecimal> createBalanceDataLoader(String userId) {
         .newMappedDataLoader((Set<UUID> bankAccountIds, BatchLoaderEnvironment environment) ->
             CompletableFuture.supplyAsync(() ->
                     balanceService.getBalanceFor((Map) environment.getKeyContexts(), userId),
-                balanceThreadPool));
+                balanceExecutor));
   }
 
 }

src/main/java/com/learn/graphql/context/dataloader/DataLoaderRegistryFactory.java

@@ -10,36 +10,29 @@
 import java.time.Instant;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
-import org.slf4j.MDC;
 import org.springframework.stereotype.Component;
 
 @Slf4j
 @Component
 @RequiredArgsConstructor
 public class RequestLoggingInstrumentation extends SimpleInstrumentation {
 
-  public static final String CORRELATION_ID = "correlation_id";
-
   private final Clock clock;
 
   @Override
   public InstrumentationContext<ExecutionResult> beginExecution(
       InstrumentationExecutionParameters parameters) {
     var start = Instant.now(clock);
-    // Add the correlation ID to the NIO thread
-    MDC.put(CORRELATION_ID, parameters.getExecutionInput().getExecutionId().toString());
-
     log.info("Query: {} with variables: {}", parameters.getQuery(), parameters.getVariables());
     return SimpleInstrumentationContext.whenCompleted((executionResult, throwable) -> {
+      // This callback will occur in the resolver thread.
+
       var duration = Duration.between(start, Instant.now(clock));
       if (throwable == null) {
         log.info("Completed successfully in: {}", duration);
       } else {
         log.warn("Failed in: {}", duration, throwable);
       }
-      // If we have async resolvers, this callback can occur in the thread-pool and not the NIO thread.
-      // In that case, the LoggingListener will be used as a fallback to clear the NIO thread.
-      MDC.clear();
     });
   }
 

src/main/java/com/learn/graphql/instrumentation/RequestLoggingInstrumentation.java

@@ -3,37 +3,27 @@
 import graphql.kickstart.servlet.core.GraphQLServletListener;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
-import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
-import org.slf4j.MDC;
 import org.springframework.stereotype.Component;
 
 @Slf4j
 @Component
-@RequiredArgsConstructor
 public class LoggingListener implements GraphQLServletListener {
 
+  /**
+   * Override other default methods to provide GraphQL life-cycle callbacks.
+   */
+  private static final RequestCallback ON_FINALLY_LISTENER = new RequestCallback() {
+    @Override
+    public void onFinally(HttpServletRequest request, HttpServletResponse response) {
+      // The final callback in the GraphQL life-cycle
+      log.info("OnFinally: GraphQL query complete");
+    }
+  };
+
   @Override
   public RequestCallback onRequest(HttpServletRequest request, HttpServletResponse response) {
-    return new RequestCallback() {
-      @Override
-      public void onSuccess(HttpServletRequest request, HttpServletResponse response) {
-        // no-op
-      }
-
-      @Override
-      public void onError(HttpServletRequest request, HttpServletResponse response,
-          Throwable throwable) {
-        log.error("Caught exception in listener.", throwable);
-      }
-
-      @Override
-      public void onFinally(HttpServletRequest request, HttpServletResponse response) {
-        // This callback will be called post graphql lifecycle.
-        // If we are multi-threading we can clear the original NIO thread MDC variables here.
-        MDC.clear();
-      }
-    };
+    return ON_FINALLY_LISTENER;
   }
 
 }