[grid] Improve race conditions in Grid session distribution (#16939)

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

Author: VietND96

java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java

@@ -441,27 +441,34 @@ private CreateSessionResponse startSession(
   }
 
   private SlotId reserveSlot(RequestId requestId, Capabilities caps) {
-    Lock writeLock = lock.writeLock();
-    writeLock.lock();
+    // Use read lock for slot selection to allow concurrent reads
+    // This reduces contention compared to using write lock for the entire operation
+    Set<SlotId> slotIds;
+    Lock readLock = lock.readLock();
+    readLock.lock();
     try {
-      Set<SlotId> slotIds = slotSelector.selectSlot(caps, getAvailableNodes(), slotMatcher);
-      if (slotIds.isEmpty()) {
-        LOG.log(
-            getDebugLogLevel(),
-            String.format("No slots found for request %s and capabilities %s", requestId, caps));
-        return null;
-      }
-
-      for (SlotId slotId : slotIds) {
-        if (reserve(slotId)) {
-          return slotId;
-        }
-      }
+      slotIds = slotSelector.selectSlot(caps, getAvailableNodes(), slotMatcher);
+    } finally {
+      readLock.unlock();
+    }
 
+    if (slotIds.isEmpty()) {
+      LOG.log(
+          getDebugLogLevel(),
+          String.format("No slots found for request %s and capabilities %s", requestId, caps));
       return null;
-    } finally {
-      writeLock.unlock();
     }
+
+    // Try to reserve each candidate slot
+    // The reserve() method uses write lock briefly and atomic operations ensure thread safety
+    // Multiple threads may select the same slot but only one will successfully reserve it
+    for (SlotId slotId : slotIds) {
+      if (reserve(slotId)) {
+        return slotId;
+      }
+    }
+
+    return null;
   }
 
   private boolean isNotSupported(Capabilities caps) {

java/src/org/openqa/selenium/grid/distributor/local/LocalGridModel.java

@@ -337,13 +337,16 @@ public boolean reserve(SlotId slotId) {
       }
 
       Optional<Slot> maybeSlot =
-          node.getSlots().stream().filter(slot -> slotId.equals(slot.getId())).findFirst();
+          node.getSlots().stream()
+              .filter(slot -> slotId.equals(slot.getId()))
+              .filter(slot -> slot.getSession() == null)
+              .findFirst();
 
       if (maybeSlot.isEmpty()) {
-        LOG.warning(
+        LOG.fine(
             String.format(
-                "Asked to reserve slot on node %s, but no slot with id %s found",
-                node.getNodeId(), slotId));
+                "Asked to reserve slot %s on node %s, but slot not found or already reserved",
+                slotId, node.getNodeId()));
         return false;
       }
 

java/src/org/openqa/selenium/grid/node/local/SessionSlot.java

@@ -62,7 +62,8 @@ public class SessionSlot
   private final boolean supportingCdp;
   private final boolean supportingBiDi;
   private final AtomicLong connectionCounter;
-  private ActiveSession currentSession;
+  // volatile ensures memory visibility across threads when session is set after reservation
+  private volatile ActiveSession currentSession;
 
   public SessionSlot(EventBus bus, Capabilities stereotype, SessionFactory factory) {
     this.bus = Require.nonNull("Event bus", bus);
@@ -131,16 +132,17 @@ public void stop(SessionClosedReason reason) {
     }
 
     SessionId id = currentSession.getId();
+    LOG.info(String.format("Stopping session %s (reason: %s)", id, reason));
     try {
       currentSession.stop();
+      LOG.info(String.format("Session stopped successfully: %s", id));
     } catch (Exception e) {
-      LOG.log(Level.WARNING, "Unable to cleanly close session", e);
+      LOG.log(Level.WARNING, String.format("Unable to cleanly close session %s", id), e);
     }
     currentSession = null;
     connectionCounter.set(0);
     release();
     bus.fire(new SessionClosedEvent(id, reason));
-    LOG.info(String.format("Stopping session %s (reason: %s)", id, reason));
   }
 
   @Override

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

@@ -33,7 +33,6 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
-import java.util.Iterator;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -118,28 +117,29 @@ public void close() {
     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);
-              }
-            }
-          }
+
+          // Use removeIf for safe and efficient removal from ConcurrentHashMap
+          httpClients
+              .entrySet()
+              .removeIf(
+                  entry -> {
+                    CacheEntry cacheEntry = entry.getValue();
+                    if (cacheEntry.inUse.get() != 0) {
+                      // the client is currently in use
+                      return false;
+                    }
+                    if (!cacheEntry.lastUse.isBefore(staleBefore)) {
+                      // the client was recently used
+                      return false;
+                    }
+                    // the client has not been used for a while, close and remove it
+                    try {
+                      cacheEntry.httpClient.close();
+                    } catch (Exception ex) {
+                      LOG.log(Level.WARNING, "failed to close a stale httpclient", ex);
+                    }
+                    return true;
+                  });
         };
 
     this.cleanUpHttpClientsCacheService =