[grid] Fixing Node drain when session expires

diemol · diemol · commit 582073d3a71e · 2022-05-23T13:39:37.000+02:00
Helps with #9845
diff --git a/java/src/org/openqa/selenium/grid/node/local/LocalNode.java b/java/src/org/openqa/selenium/grid/node/local/LocalNode.java
@@ -161,32 +161,41 @@ private LocalNode(
                            String.format("%s is %s", uri, status.getAvailability()));
                        } : healthCheck;
 
+    this.tempFileSystems = CacheBuilder.newBuilder()
+      .expireAfterAccess(sessionTimeout)
+      .ticker(ticker)
+      .removalListener((RemovalListener<SessionId, TemporaryFilesystem>) notification -> {
+        TemporaryFilesystem tempFS = notification.getValue();
+        tempFS.deleteTemporaryFiles();
+        tempFS.deleteBaseDir();
+      })
+      .build();
+
     this.currentSessions = CacheBuilder.newBuilder()
       .expireAfterAccess(sessionTimeout)
       .ticker(ticker)
       .removalListener((RemovalListener<SessionId, SessionSlot>) notification -> {
         if (notification.getKey() != null && notification.getValue() != null) {
           // Attempt to stop the session
           SessionSlot slot = notification.getValue();
-          if (!slot.isAvailable()) {
-            slot.stop();
+          SessionId sessionId = notification.getKey();
+          slot.stop();
+          // Invalidate temp file system
+          this.tempFileSystems.invalidate(sessionId);
+          // Decrement pending sessions if Node is draining
+          if (this.isDraining()) {
+            int done = pendingSessions.decrementAndGet();
+            if (done <= 0) {
+              LOG.info("Node draining complete!");
+              bus.fire(new NodeDrainComplete(this.getId()));
+            }
           }
         } else {
           LOG.log(Debug.getDebugLogLevel(), "Received stop session notification with null values");
         }
       })
       .build();
 
-    this.tempFileSystems = CacheBuilder.newBuilder()
-      .expireAfterAccess(sessionTimeout)
-      .ticker(ticker)
-      .removalListener((RemovalListener<SessionId, TemporaryFilesystem>) notification -> {
-        TemporaryFilesystem tempFS = notification.getValue();
-        tempFS.deleteTemporaryFiles();
-        tempFS.deleteBaseDir();
-      })
-      .build();
-
     ScheduledExecutorService sessionCleanupNodeService =
       Executors.newSingleThreadScheduledExecutor(
         r -> {
@@ -486,16 +495,6 @@ public void stop(SessionId id) throws NoSuchSessionException {
     }
 
     currentSessions.invalidate(id);
-    tempFileSystems.invalidate(id);
-
-    // Decrement pending sessions if Node is draining
-    if (this.isDraining()) {
-      int done = pendingSessions.decrementAndGet();
-      if (done <= 0) {
-        LOG.info("Node draining complete!");
-        bus.fire(new NodeDrainComplete(this.getId()));
-      }
-    }
   }
 
   private void stopAllSessions() {
