[grid] Decrementing pending sessions for draining to stop method

diemol · diemol · commit be30b6241d89 · 2022-04-07T12:21:32.000+02:00
Using a listener for this was an overhead as this could
be done when the session is being stopped in the Node.

Additionally, we were not checking in the listener if the
session actually belonged to the Node.
diff --git a/java/src/org/openqa/selenium/grid/node/local/LocalNode.java b/java/src/org/openqa/selenium/grid/node/local/LocalNode.java
@@ -42,7 +42,6 @@
 import org.openqa.selenium.grid.data.NodeId;
 import org.openqa.selenium.grid.data.NodeStatus;
 import org.openqa.selenium.grid.data.Session;
-import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.data.Slot;
 import org.openqa.selenium.grid.data.SlotId;
 import org.openqa.selenium.grid.jmx.JMXHelper;
@@ -217,18 +216,6 @@ private LocalNode(
       heartbeatPeriod.getSeconds(),
       TimeUnit.SECONDS);
 
-    bus.addListener(SessionClosedEvent.listener(id -> {
-      // Shouldn't we check that the session actually belongs to this Node?
-      // Listen to session terminated events, so we know when to fire the NodeDrainComplete event
-      if (this.isDraining()) {
-        int done = pendingSessions.decrementAndGet();
-        if (done <= 0) {
-          LOG.info("Firing node drain complete message");
-          bus.fire(new NodeDrainComplete(this.getId()));
-        }
-      }
-    }));
-
     Runtime.getRuntime().addShutdownHook(new Thread(this::stopAllSessions));
     new JMXHelper().register(this);
   }
@@ -487,6 +474,15 @@ 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() {
