Merge branch 'trunk' into renovate/bazel-7.x

Author: diemol

java/src/org/openqa/selenium/grid/distributor/GridModel.java

@@ -400,7 +400,12 @@ public void release(SessionId id) {
     }
   }
 
-  public void reserve(NodeStatus status, Slot slot) {
+  /**
+   * A helper to reserve a slot of a node. The writeLock must be acquired outside to ensure the view
+   * of the NodeStatus is the current state, otherwise concurrent calls to amend will work with an
+   * outdated view of slots.
+   */
+  private void reserve(NodeStatus status, Slot slot) {
     Instant now = Instant.now();
 
     Slot reserved =
@@ -490,30 +495,29 @@ public void updateHealthCheckCount(NodeId id, Availability availability) {
     }
   }
 
+  /**
+   * A helper to replace the availability and a slot of a node. The writeLock must be acquired
+   * outside to ensure the view of the NodeStatus is the current state, otherwise concurrent calls
+   * to amend will work with an outdated view of slots.
+   */
   private void amend(Availability availability, NodeStatus status, Slot slot) {
     Set<Slot> newSlots = new HashSet<>(status.getSlots());
     newSlots.removeIf(s -> s.getId().equals(slot.getId()));
     newSlots.add(slot);
 
     NodeStatus node = getNode(status.getNodeId());
 
-    Lock writeLock = lock.writeLock();
-    writeLock.lock();
-    try {
-      nodes.remove(node);
-      nodes.add(
-          new NodeStatus(
-              status.getNodeId(),
-              status.getExternalUri(),
-              status.getMaxSessionCount(),
-              newSlots,
-              availability,
-              status.getHeartbeatPeriod(),
-              status.getSessionTimeout(),
-              status.getVersion(),
-              status.getOsInfo()));
-    } finally {
-      writeLock.unlock();
-    }
+    nodes.remove(node);
+    nodes.add(
+        new NodeStatus(
+            status.getNodeId(),
+            status.getExternalUri(),
+            status.getMaxSessionCount(),
+            newSlots,
+            availability,
+            status.getHeartbeatPeriod(),
+            status.getSessionTimeout(),
+            status.getVersion(),
+            status.getOsInfo()));
   }
 }

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

@@ -355,6 +355,7 @@ public LocalDistributor add(Node node) {
     // An exception occurs if Node heartbeat has started but the server is not ready.
     // Unhandled exception blocks the event-bus thread from processing any event henceforth.
     NodeStatus initialNodeStatus;
+    Runnable healthCheck;
     try {
       initialNodeStatus = node.getStatus();
       if (initialNodeStatus.getAvailability() != UP) {
@@ -363,8 +364,17 @@ public LocalDistributor add(Node node) {
         // We do not need to add this Node for now.
         return this;
       }
-      model.add(initialNodeStatus);
-      nodes.put(node.getId(), node);
+      // Extract the health check
+      healthCheck = asRunnableHealthCheck(node);
+      Lock writeLock = lock.writeLock();
+      writeLock.lock();
+      try {
+        nodes.put(node.getId(), node);
+        model.add(initialNodeStatus);
+        allChecks.put(node.getId(), healthCheck);
+      } finally {
+        writeLock.unlock();
+      }
     } catch (Exception e) {
       LOG.log(
           Debug.getDebugLogLevel(),
@@ -373,10 +383,6 @@ public LocalDistributor add(Node node) {
       return this;
     }
 
-    // Extract the health check
-    Runnable healthCheck = asRunnableHealthCheck(node);
-    allChecks.put(node.getId(), healthCheck);
-
     updateNodeStatus(initialNodeStatus, healthCheck);
 
     LOG.info(
@@ -415,7 +421,15 @@ private void updateNodeStatus(NodeStatus status, Runnable healthCheck) {
 
   private Runnable runNodeHealthChecks() {
     return () -> {
-      ImmutableMap<NodeId, Runnable> nodeHealthChecks = ImmutableMap.copyOf(allChecks);
+      ImmutableMap<NodeId, Runnable> nodeHealthChecks;
+      Lock readLock = this.lock.readLock();
+      readLock.lock();
+      try {
+        nodeHealthChecks = ImmutableMap.copyOf(allChecks);
+      } finally {
+        readLock.unlock();
+      }
+
       for (Runnable nodeHealthCheck : nodeHealthChecks.values()) {
         GuardedRunnable.guard(nodeHealthCheck).run();
       }

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

@@ -377,8 +377,10 @@ public boolean isReady() {
   @VisibleForTesting
   @ManagedAttribute(name = "CurrentSessions")
   public int getCurrentSessionCount() {
+    // we need the exact size, see javadoc of Cache.size
+    long n = currentSessions.asMap().values().stream().count();
     // It seems wildly unlikely we'll overflow an int
-    return Math.toIntExact(currentSessions.size());
+    return Math.toIntExact(n);
   }
 
   @VisibleForTesting
@@ -1005,6 +1007,9 @@ public HealthCheck getHealthCheck() {
   public void drain() {
     bus.fire(new NodeDrainStarted(getId()));
     draining = true;
+    // Ensure the pendingSessions counter will not be decremented by timed out sessions not included
+    // in the currentSessionCount and the NodeDrainComplete will be raised to early.
+    currentSessions.cleanUp();
     int currentSessionCount = getCurrentSessionCount();
     if (currentSessionCount == 0) {
       LOG.info("Firing node drain complete message");