[grid] Simplify httpClientCache manages in handle session of Router

Signed-off-by: Viet Nguyen Duc <[email protected]>

Author: VietND96

java/src/org/openqa/selenium/grid/router/BUILD.bazel

@@ -24,5 +24,6 @@ java_library(
         "//java/src/org/openqa/selenium/remote",
         "//java/src/org/openqa/selenium/status",
         artifact("com.google.guava:guava"),
+        artifact("com.github.ben-manes.caffeine:caffeine"),
     ],
 )

java/src/org/openqa/selenium/grid/router/GridStatusHandler.java

@@ -35,6 +35,8 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeoutException;
+import java.util.logging.Logger;
+import org.openqa.selenium.concurrent.ExecutorServices;
 import org.openqa.selenium.grid.data.DistributorStatus;
 import org.openqa.selenium.grid.distributor.Distributor;
 import org.openqa.selenium.internal.Require;
@@ -47,8 +49,9 @@
 import org.openqa.selenium.remote.tracing.Status;
 import org.openqa.selenium.remote.tracing.Tracer;
 
-class GridStatusHandler implements HttpHandler {
+class GridStatusHandler implements HttpHandler, AutoCloseable {
 
+  private static final Logger LOG = Logger.getLogger(GridStatusHandler.class.getName());
   private static final ExecutorService EXECUTOR_SERVICE =
       Executors.newCachedThreadPool(
           r -> {
@@ -165,4 +168,10 @@ public HttpResponse execute(HttpRequest req) {
       return res;
     }
   }
+
+  @Override
+  public void close() {
+    LOG.info("Shutting down GridStatusHandler executor service");
+    ExecutorServices.shutdownGracefully("Grid status executor", EXECUTOR_SERVICE);
+  }
 }

java/src/org/openqa/selenium/grid/router/HandleSession.java

@@ -26,23 +26,16 @@
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_REQUEST_EVENT;
 import static org.openqa.selenium.remote.tracing.Tags.HTTP_RESPONSE;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.RemovalListener;
 import java.io.Closeable;
 import java.net.URI;
-import java.time.Instant;
-import java.time.temporal.ChronoUnit;
-import java.util.Iterator;
+import java.time.Duration;
 import java.util.concurrent.Callable;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.openqa.selenium.NoSuchSessionException;
-import org.openqa.selenium.concurrent.ExecutorServices;
-import org.openqa.selenium.concurrent.GuardedRunnable;
 import org.openqa.selenium.grid.sessionmap.SessionMap;
 import org.openqa.selenium.grid.web.ReverseProxyHandler;
 import org.openqa.selenium.internal.Require;
@@ -64,88 +57,47 @@ class HandleSession implements HttpHandler, Closeable {
 
   private static final Logger LOG = Logger.getLogger(HandleSession.class.getName());
 
-  private static class CacheEntry {
-    private final HttpClient httpClient;
-    private final AtomicLong inUse;
-    // volatile as the ConcurrentMap will not take care of synchronization
-    private volatile Instant lastUse;
-
-    public CacheEntry(HttpClient httpClient, long initialUsage) {
-      this.httpClient = httpClient;
-      this.inUse = new AtomicLong(initialUsage);
-      this.lastUse = Instant.now();
-    }
-  }
-
-  private static class UsageCountingReverseProxyHandler extends ReverseProxyHandler
+  private static class ReverseProxyHandlerCloseable extends ReverseProxyHandler
       implements Closeable {
-    private final CacheEntry entry;
 
-    public UsageCountingReverseProxyHandler(
-        Tracer tracer, HttpClient httpClient, CacheEntry entry) {
+    public ReverseProxyHandlerCloseable(Tracer tracer, HttpClient httpClient) {
       super(tracer, httpClient);
-
-      this.entry = entry;
     }
 
     @Override
     public void close() {
-      // set the last use here, to ensure we have to calculate the real inactivity of the client
-      entry.lastUse = Instant.now();
-      entry.inUse.decrementAndGet();
+      // No operation needed - cache management is handled by Cache builder
     }
   }
 
   private final Tracer tracer;
   private final HttpClient.Factory httpClientFactory;
   private final SessionMap sessions;
-  private final ConcurrentMap<URI, CacheEntry> httpClients;
-  private final ScheduledExecutorService cleanUpHttpClientsCacheService;
+  private final Cache<URI, HttpClient> httpClientCache;
 
   HandleSession(Tracer tracer, HttpClient.Factory httpClientFactory, SessionMap sessions) {
     this.tracer = Require.nonNull("Tracer", tracer);
     this.httpClientFactory = Require.nonNull("HTTP client factory", httpClientFactory);
     this.sessions = Require.nonNull("Sessions", sessions);
 
-    this.httpClients = new ConcurrentHashMap<>();
-
-    Runnable cleanUpHttpClients =
-        () -> {
-          Instant staleBefore = Instant.now().minus(2, ChronoUnit.MINUTES);
-          Iterator<CacheEntry> iterator = httpClients.values().iterator();
-
-          while (iterator.hasNext()) {
-            CacheEntry entry = iterator.next();
-
-            if (entry.inUse.get() != 0) {
-              // the client is currently in use
-              return;
-            } else if (!entry.lastUse.isBefore(staleBefore)) {
-              // the client was recently used
-              return;
-            } else {
-              // the client has not been used for a while, remove it from the cache
-              iterator.remove();
-
-              try {
-                entry.httpClient.close();
-              } catch (Exception ex) {
-                LOG.log(Level.WARNING, "failed to close a stale httpclient", ex);
-              }
-            }
-          }
-        };
-
-    this.cleanUpHttpClientsCacheService =
-        Executors.newSingleThreadScheduledExecutor(
-            r -> {
-              Thread thread = new Thread(r);
-              thread.setDaemon(true);
-              thread.setName("HandleSession - Clean up http clients cache");
-              return thread;
-            });
-    cleanUpHttpClientsCacheService.scheduleAtFixedRate(
-        GuardedRunnable.guard(cleanUpHttpClients), 1, 1, TimeUnit.MINUTES);
+    // Create Cache with 2 minute expiry after last access
+    // and a removal listener to close HTTP clients
+    this.httpClientCache =
+        Caffeine.newBuilder()
+            .expireAfterAccess(Duration.ofMinutes(2))
+            .removalListener(
+                (RemovalListener<URI, HttpClient>)
+                    (uri, httpClient, cause) -> {
+                      if (httpClient != null) {
+                        try {
+                          LOG.fine("Closing HTTP client for " + uri + ", removal cause: " + cause);
+                          httpClient.close();
+                        } catch (Exception ex) {
+                          LOG.log(Level.WARNING, "Failed to close HTTP client for " + uri, ex);
+                        }
+                      }
+                    })
+            .build();
   }
 
   @Override
@@ -179,7 +131,7 @@ public HttpResponse execute(HttpRequest req) {
       try {
         HttpTracing.inject(tracer, span, req);
         HttpResponse res;
-        try (UsageCountingReverseProxyHandler handler = loadSessionId(tracer, span, id).call()) {
+        try (ReverseProxyHandlerCloseable handler = loadSessionId(tracer, span, id).call()) {
           res = handler.execute(req);
         }
 
@@ -216,46 +168,34 @@ public HttpResponse execute(HttpRequest req) {
     }
   }
 
-  private Callable<UsageCountingReverseProxyHandler> loadSessionId(
+  private Callable<ReverseProxyHandlerCloseable> loadSessionId(
       Tracer tracer, Span span, SessionId id) {
     return span.wrap(
         () -> {
-          CacheEntry cacheEntry =
-              httpClients.compute(
-                  sessions.getUri(id),
-                  (sessionUri, entry) -> {
-                    if (entry != null) {
-                      entry.inUse.incrementAndGet();
-                      return entry;
-                    }
-
-                    ClientConfig config =
-                        ClientConfig.defaultConfig().baseUri(sessionUri).withRetries();
-                    HttpClient httpClient = httpClientFactory.createClient(config);
-
-                    return new CacheEntry(httpClient, 1);
+          URI sessionUri = sessions.getUri(id);
+
+          // Get or create the HTTP client from cache (this also updates the "last access" time)
+          HttpClient httpClient =
+              httpClientCache.get(
+                  sessionUri,
+                  uri -> {
+                    LOG.fine("Creating new HTTP client for " + uri);
+                    ClientConfig config = ClientConfig.defaultConfig().baseUri(uri).withRetries();
+                    return httpClientFactory.createClient(config);
                   });
 
           try {
-            return new UsageCountingReverseProxyHandler(tracer, cacheEntry.httpClient, cacheEntry);
+            return new ReverseProxyHandlerCloseable(tracer, httpClient);
           } catch (Throwable t) {
-            // ensure we do not keep the http client when an unexpected throwable is raised
-            cacheEntry.inUse.decrementAndGet();
             throw t;
           }
         });
   }
 
   @Override
   public void close() {
-    ExecutorServices.shutdownGracefully(
-        "HandleSession - Clean up http clients cache", cleanUpHttpClientsCacheService);
-    httpClients
-        .values()
-        .removeIf(
-            (entry) -> {
-              entry.httpClient.close();
-              return true;
-            });
+    // This will trigger the removal listener for all entries, which will close all HTTP clients
+    httpClientCache.invalidateAll();
+    httpClientCache.cleanUp();
   }
 }