[grid] Migrate cache builder to use Caffeine

VietND96 · VietND96 · commit 54ead6a78b34 · 2025-04-02T06:59:38.000+07:00
Signed-off-by: Viet Nguyen Duc &lt;nguyenducviet4496@gmail.com&gt;
diff --git a/MODULE.bazel b/MODULE.bazel
@@ -177,6 +177,7 @@ maven.install(
         "com.google.code.findbugs:jsr305:3.0.2",
         "com.google.code.gson:gson:2.12.1",
         "com.google.guava:guava:33.4.5-jre",
+        "com.github.ben-manes.caffeine:caffeine:3.2.0",
         "com.google.auto:auto-common:1.2.2",
         "com.google.auto.service:auto-service:1.1.1",
         "com.google.auto.service:auto-service-annotations:1.1.1",
diff --git a/java/maven_install.json b/java/maven_install.json
@@ -1,7 +1,7 @@
 {
   "__AUTOGENERATED_FILE_DO_NOT_MODIFY_THIS_FILE_MANUALLY": "THERE_IS_NO_DATA_ONLY_ZUUL",
-  "__INPUT_ARTIFACTS_HASH": 466809341,
-  "__RESOLVED_ARTIFACTS_HASH": -983027519,
+  "__INPUT_ARTIFACTS_HASH": -95711414,
+  "__RESOLVED_ARTIFACTS_HASH": -1591732796,
   "artifacts": {
     "com.beust:jcommander": {
       "shasums": {
@@ -59,6 +59,13 @@
       },
       "version": "2.18.2"
     },
+    "com.github.ben-manes.caffeine:caffeine": {
+      "shasums": {
+        "jar": "ec411dfdf0c03f25218648ce89861630b71680e5858a9a7278ebac8e55cab3d7",
+        "sources": "67e14ef5c04c193a7fcafa788b55b89a079fd584b202469721ce6d2d6c753090"
+      },
+      "version": "3.2.0"
+    },
     "com.github.javaparser:javaparser-core": {
       "shasums": {
         "jar": "a24c4fa7799ffe0c7a9af11d4eecd757098ed4498f86067bf28b46b2bfea1833",
@@ -810,6 +817,10 @@
       "com.fasterxml.jackson.core:jackson-databind",
       "org.yaml:snakeyaml"
     ],
+    "com.github.ben-manes.caffeine:caffeine": [
+      "com.google.errorprone:error_prone_annotations",
+      "org.jspecify:jspecify"
+    ],
     "com.github.spotbugs:spotbugs": [
       "com.github.spotbugs:spotbugs-annotations",
       "com.github.stephenc.jcip:jcip-annotations",
@@ -1199,6 +1210,10 @@
       "com.fasterxml.jackson.dataformat.yaml.snakeyaml.error",
       "com.fasterxml.jackson.dataformat.yaml.util"
     ],
+    "com.github.ben-manes.caffeine:caffeine": [
+      "com.github.benmanes.caffeine.cache",
+      "com.github.benmanes.caffeine.cache.stats"
+    ],
     "com.github.javaparser:javaparser-core": [
       "com.github.javaparser",
       "com.github.javaparser.ast",
@@ -2910,6 +2925,8 @@
       "com.fasterxml.jackson.core:jackson-databind:jar:sources",
       "com.fasterxml.jackson.dataformat:jackson-dataformat-yaml",
       "com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:jar:sources",
+      "com.github.ben-manes.caffeine:caffeine",
+      "com.github.ben-manes.caffeine:caffeine:jar:sources",
       "com.github.javaparser:javaparser-core",
       "com.github.javaparser:javaparser-core:jar:sources",
       "com.github.spotbugs:spotbugs",
diff --git a/java/src/org/openqa/selenium/grid/graphql/BUILD.bazel b/java/src/org/openqa/selenium/grid/graphql/BUILD.bazel
@@ -21,5 +21,6 @@ java_library(
         "//java/src/org/openqa/selenium/remote/http",
         artifact("com.google.guava:guava"),
         artifact("com.graphql-java:graphql-java"),
+        artifact("com.github.ben-manes.caffeine:caffeine"),
     ],
 )
diff --git a/java/src/org/openqa/selenium/grid/graphql/GraphqlHandler.java b/java/src/org/openqa/selenium/grid/graphql/GraphqlHandler.java
@@ -28,8 +28,8 @@
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_RESPONSE;
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_RESPONSE_EVENT;
 
-import com.google.common.cache.Cache;
-import com.google.common.cache.CacheBuilder;
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.GraphQL;
@@ -46,7 +46,6 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
 import org.openqa.selenium.grid.distributor.Distributor;
 import org.openqa.selenium.grid.sessionqueue.NewSessionQueue;
 import org.openqa.selenium.internal.Require;
@@ -84,33 +83,25 @@ public GraphqlHandler(
     this.publicUri = Require.nonNull("Uri", publicUri);
     this.version = Require.nonNull("GridVersion", version);
     this.tracer = Require.nonNull("Tracer", tracer);
+    long maxMemory = Runtime.getRuntime().maxMemory();
+    long cacheSize = maxMemory / 1024 / 1024;
 
     GraphQLSchema schema =
         new SchemaGenerator()
             .makeExecutableSchema(buildTypeDefinitionRegistry(), buildRuntimeWiring());
 
     Cache<String, CompletableFuture<PreparsedDocumentEntry>> cache =
-        CacheBuilder.newBuilder().maximumSize(1024).build();
+        Caffeine.newBuilder().maximumSize(cacheSize).build();
 
     graphQl =
         GraphQL.newGraphQL(schema)
             .preparsedDocumentProvider(
-                (executionInput, computeFunction) -> {
-                  try {
-                    return cache.get(
+                (executionInput, computeFunction) ->
+                    cache.get(
                         executionInput.getQuery(),
-                        () ->
+                        key ->
                             CompletableFuture.supplyAsync(
-                                () -> computeFunction.apply(executionInput)));
-                  } catch (ExecutionException e) {
-                    if (e.getCause() instanceof RuntimeException) {
-                      throw (RuntimeException) e.getCause();
-                    } else if (e.getCause() != null) {
-                      throw new RuntimeException(e.getCause());
-                    }
-                    throw new RuntimeException(e);
-                  }
-                })
+                                () -> computeFunction.apply(executionInput))))
             .build();
   }
 
diff --git a/java/src/org/openqa/selenium/grid/node/local/BUILD.bazel b/java/src/org/openqa/selenium/grid/node/local/BUILD.bazel
@@ -28,5 +28,6 @@ java_library(
         "//java/src/org/openqa/selenium/json",
         "//java/src/org/openqa/selenium/remote",
         artifact("com.google.guava:guava"),
+        artifact("com.github.ben-manes.caffeine:caffeine"),
     ],
 )
diff --git a/java/src/org/openqa/selenium/grid/node/local/LocalNode.java b/java/src/org/openqa/selenium/grid/node/local/LocalNode.java
@@ -30,13 +30,11 @@
 import static org.openqa.selenium.remote.http.Contents.string;
 import static org.openqa.selenium.remote.http.HttpMethod.DELETE;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.RemovalCause;
+import com.github.benmanes.caffeine.cache.Ticker;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Ticker;
-import com.google.common.cache.Cache;
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.RemovalCause;
-import com.google.common.cache.RemovalListener;
-import com.google.common.cache.RemovalNotification;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.io.Closeable;
@@ -56,7 +54,6 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -200,35 +197,35 @@ protected LocalNode(
     // Do not clear this cache automatically using a timer.
     // It will be explicitly cleaned up, as and when "currentSessions" is auto cleaned.
     this.uploadsTempFileSystem =
-        CacheBuilder.newBuilder()
+        Caffeine.newBuilder()
             .removalListener(
-                (RemovalListener<SessionId, TemporaryFilesystem>)
-                    notification ->
-                        Optional.ofNullable(notification.getValue())
-                            .ifPresent(
-                                tempFS -> {
-                                  tempFS.deleteTemporaryFiles();
-                                  tempFS.deleteBaseDir();
-                                }))
+                (SessionId key, TemporaryFilesystem tempFS, RemovalCause cause) -> {
+                  Optional.ofNullable(tempFS)
+                      .ifPresent(
+                          fs -> {
+                            fs.deleteTemporaryFiles();
+                            fs.deleteBaseDir();
+                          });
+                })
             .build();
 
     // Do not clear this cache automatically using a timer.
     // It will be explicitly cleaned up, as and when "currentSessions" is auto cleaned.
     this.downloadsTempFileSystem =
-        CacheBuilder.newBuilder()
+        Caffeine.newBuilder()
             .removalListener(
-                (RemovalListener<SessionId, TemporaryFilesystem>)
-                    notification ->
-                        Optional.ofNullable(notification.getValue())
-                            .ifPresent(
-                                fs -> {
-                                  fs.deleteTemporaryFiles();
-                                  fs.deleteBaseDir();
-                                }))
+                (SessionId key, TemporaryFilesystem tempFS, RemovalCause cause) -> {
+                  Optional.ofNullable(tempFS)
+                      .ifPresent(
+                          fs -> {
+                            fs.deleteTemporaryFiles();
+                            fs.deleteBaseDir();
+                          });
+                })
             .build();
 
     this.currentSessions =
-        CacheBuilder.newBuilder()
+        Caffeine.newBuilder()
             .expireAfterAccess(sessionTimeout)
             .ticker(ticker)
             .removalListener(this::stopTimedOutSession)
@@ -313,19 +310,19 @@ public void close() {
     shutdown.run();
   }
 
-  private void stopTimedOutSession(RemovalNotification<SessionId, SessionSlot> notification) {
+  private void stopTimedOutSession(SessionId key, SessionSlot value, RemovalCause cause) {
     try (Span span = tracer.getCurrentContext().createSpan("node.stop_session")) {
       AttributeMap attributeMap = tracer.createAttributeMap();
       attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
-      if (notification.getKey() != null && notification.getValue() != null) {
-        SessionSlot slot = notification.getValue();
-        SessionId id = notification.getKey();
+      if (key != null && value != null) {
+        SessionSlot slot = value;
+        SessionId id = key;
         attributeMap.put("node.id", getId().toString());
         attributeMap.put("session.slotId", slot.getId().toString());
         attributeMap.put("session.id", id.toString());
         attributeMap.put("session.timeout_in_seconds", getSessionTimeout().toSeconds());
-        attributeMap.put("session.remove.cause", notification.getCause().name());
-        if (notification.wasEvicted() && notification.getCause() == RemovalCause.EXPIRED) {
+        attributeMap.put("session.remove.cause", cause.name());
+        if (cause == RemovalCause.EXPIRED) {
           // Session is timing out, stopping it by sending a DELETE
           LOG.log(Level.INFO, () -> String.format("Session id %s timed out, stopping...", id));
           span.setStatus(Status.CANCELLED);
@@ -334,7 +331,7 @@ private void stopTimedOutSession(RemovalNotification<SessionId, SessionSlot> not
           LOG.log(Level.INFO, () -> String.format("Session id %s is stopping on demand...", id));
           span.addEvent(String.format("Stopping the session %s on demand", id), attributeMap);
         }
-        if (notification.wasEvicted()) {
+        if (cause == RemovalCause.EXPIRED) {
           try {
             slot.execute(new HttpRequest(DELETE, "/session/" + id));
           } catch (Exception e) {
@@ -686,15 +683,11 @@ public Session getSession(SessionId id) throws NoSuchSessionException {
 
   @Override
   public TemporaryFilesystem getUploadsFilesystem(SessionId id) throws IOException {
-    try {
-      return uploadsTempFileSystem.get(
-          id,
-          () ->
-              TemporaryFilesystem.getTmpFsBasedOn(
-                  TemporaryFilesystem.getDefaultTmpFS().createTempDir("session", id.toString())));
-    } catch (ExecutionException e) {
-      throw new IOException(e);
-    }
+    return uploadsTempFileSystem.get(
+        id,
+        key ->
+            TemporaryFilesystem.getTmpFsBasedOn(
+                TemporaryFilesystem.getDefaultTmpFS().createTempDir("session", id.toString())));
   }
 
   @Override
@@ -862,7 +855,7 @@ public void stop(SessionId id) throws NoSuchSessionException {
   }
 
   private void stopAllSessions() {
-    if (currentSessions.size() > 0) {
+    if (currentSessions.estimatedSize() > 0) {
       LOG.info("Trying to stop all running sessions before shutting down...");
       currentSessions.invalidateAll();
     }
diff --git a/rust/Cargo.Bazel.lock b/rust/Cargo.Bazel.lock
@@ -1,5 +1,5 @@
 {
-  "checksum": "f9b7e2302bd7b43ff0899cac357fbbfe0226873f02b610c0eb4246784e14e3b3",
+  "checksum": "a896d7a078f1c885fa00c3b9aa5b8727c4f5bdfc58cdc26836fb74ebbdd3c87c",
   "crates": {
     "addr2line 0.21.0": {
       "name": "addr2line",
@@ -2031,16 +2031,79 @@
               "id": "jobserver 0.1.31",
               "target": "jobserver"
             },
-            {
-              "id": "libc 0.2.168",
-              "target": "libc"
-            },
             {
               "id": "shlex 1.3.0",
               "target": "shlex"
             }
           ],
-          "selects": {}
+          "selects": {
+            "aarch64-apple-darwin": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "aarch64-unknown-linux-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "aarch64-unknown-nixos-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "arm-unknown-linux-gnueabi": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "i686-unknown-linux-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "powerpc-unknown-linux-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "s390x-unknown-linux-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "x86_64-apple-darwin": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "x86_64-unknown-freebsd": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "x86_64-unknown-linux-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ],
+            "x86_64-unknown-nixos-gnu": [
+              {
+                "id": "libc 0.2.168",
+                "target": "libc"
+              }
+            ]
+          }
         },
         "edition": "2018",
         "version": "1.1.30"

Original file line number	Diff line number	Diff line change
`@@ -21,5 +21,6 @@ java_library(`
`21`	`21`	`"//java/src/org/openqa/selenium/remote/http",`
`22`	`22`	`artifact("com.google.guava:guava"),`
`23`	`23`	`artifact("com.graphql-java:graphql-java"),`
	`24`	`+ artifact("com.github.ben-manes.caffeine:caffeine"),`
`24`	`25`	`],`
`25`	`26`	`)`
Original file line number	Diff line number	Diff line change
`@@ -28,5 +28,6 @@ java_library(`
`28`	`28`	`"//java/src/org/openqa/selenium/json",`
`29`	`29`	`"//java/src/org/openqa/selenium/remote",`
`30`	`30`	`artifact("com.google.guava:guava"),`
	`31`	`+ artifact("com.github.ben-manes.caffeine:caffeine"),`
`31`	`32`	`],`
`32`	`33`	`)`