[grid] Implement custom cache weigher and enhance query handling in GraphqlHandler

Author: VietND96

java/src/org/openqa/selenium/grid/graphql/GraphqlHandler.java

@@ -30,6 +30,7 @@
 
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.Weigher;
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.GraphQL;
@@ -43,6 +44,7 @@
 import java.io.InputStream;
 import java.io.UncheckedIOException;
 import java.net.URI;
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -60,17 +62,19 @@
 import org.openqa.selenium.remote.tracing.Status;
 import org.openqa.selenium.remote.tracing.Tracer;
 
-public class GraphqlHandler implements HttpHandler {
+public class GraphqlHandler implements HttpHandler, AutoCloseable {
 
  public static final String GRID_SCHEMA =
  "/org/openqa/selenium/grid/graphql/selenium-grid-schema.graphqls";
  public static final Json JSON = new Json();
+ private static final long MAX_QUERY_SIZE_BYTES = 1024 * 1024; // 1MB limit
  private final Tracer tracer;
  private final Distributor distributor;
  private final NewSessionQueue newSessionQueue;
  private final URI publicUri;
  private final String version;
  private final GraphQL graphQl;
+ private final Cache<String, CompletableFuture<PreparsedDocumentEntry>> cache;
 
  public GraphqlHandler(
  Tracer tracer,
@@ -83,25 +87,45 @@ public GraphqlHandler(
  this.publicUri = Require.nonNull("Uri", publicUri);
  this.version = Require.nonNull("GridVersion", version);
  this.tracer = Require.nonNull("Tracer", tracer);
+
+ // Container-aware cache sizing: 5% of heap memory
  long maxMemory = Runtime.getRuntime().maxMemory();
- long cacheSize = maxMemory / 1024 / 1024;
+ long cacheWeightLimit = (long) (maxMemory * 0.05);
+
+ // Custom weigher to prevent single entries from exceeding 10% of cache weight
+ Weigher<String, CompletableFuture<PreparsedDocumentEntry>> weigher =
+ new QueryCacheWeigher(cacheWeightLimit);
+
+ this.cache =
+ Caffeine.newBuilder()
+ .maximumWeight(cacheWeightLimit)
+ .weigher(weigher)
+ .expireAfterAccess(Duration.ofMinutes(30)) // Time-based eviction
+ .build();
 
  GraphQLSchema schema =
  new SchemaGenerator()
  .makeExecutableSchema(buildTypeDefinitionRegistry(), buildRuntimeWiring());
 
- Cache<String, CompletableFuture<PreparsedDocumentEntry>> cache =
- Caffeine.newBuilder().maximumSize(cacheSize).build();
-
  graphQl =
  GraphQL.newGraphQL(schema)
  .preparsedDocumentProvider(
- (executionInput, computeFunction) ->
- cache.get(
- executionInput.getQuery(),
- key ->
- CompletableFuture.supplyAsync(
- () -> computeFunction.apply(executionInput))))
+ (executionInput, computeFunction) -> {
+ String query = executionInput.getQuery();
+
+ // Query size validation with bypass for oversized queries
+ if (query.getBytes().length > MAX_QUERY_SIZE_BYTES) {
+ // Bypass cache for oversized queries to prevent cache pollution
+ return CompletableFuture.supplyAsync(
+ () -> computeFunction.apply(executionInput));
+ }
+
+ return cache.get(
+ query,
+ key ->
+ CompletableFuture.supplyAsync(
+ () -> computeFunction.apply(executionInput)));
+ })
  .build();
  }
 
@@ -180,6 +204,24 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
  }
  }
 
+ @Override
+ public void close() {
+ // Cancel pending CompletableFutures
+ cache
+ .asMap()
+ .values()
+ .forEach(
+ future -> {
+ if (!future.isDone()) {
+ future.cancel(true);
+ }
+ });
+
+ // Invalidate and clean cache
+ cache.invalidateAll();
+ cache.cleanUp();
+ }
+
  private RuntimeWiring buildRuntimeWiring() {
  GridData gridData = new GridData(distributor, newSessionQueue, publicUri, version);
  return RuntimeWiring.newRuntimeWiring()
@@ -203,4 +245,28 @@ private TypeDefinitionRegistry buildTypeDefinitionRegistry() {
  throw new UncheckedIOException(e);
  }
  }
+
+ // Custom weigher class for query cache
+ private static class QueryCacheWeigher
+ implements Weigher<String, CompletableFuture<PreparsedDocumentEntry>> {
+
+ private final long maxSingleEntryWeight;
+
+ public QueryCacheWeigher(long cacheWeightLimit) {
+ this.maxSingleEntryWeight = (long) (cacheWeightLimit * 0.1); // 10% limit
+ }
+
+ @Override
+ public int weigh(String key, CompletableFuture<PreparsedDocumentEntry> value) {
+ // Estimate memory usage including CompletableFuture overhead
+ long keyWeight = key.length() * 2; // UTF-16 encoding
+ long futureOverhead = 200; // Estimated CompletableFuture overhead
+ long documentOverhead = 500; // Estimated PreparsedDocumentEntry overhead
+
+ long totalWeight = keyWeight + futureOverhead + documentOverhead;
+
+ // Bounds checking to prevent single entries from dominating cache
+ return (int) Math.min(totalWeight, maxSingleEntryWeight);
+ }
+ }
 }