[grid] Distributor listen and handle the NodeRestartedEvent (#14938)

Signed-off-by: Viet Nguyen Duc <[email protected]>
SeleniumHQ · Dec 24, 2024 · 8a04527 · 8a04527
1 parent 7854e62
commit 8a04527
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java b/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java
@@ -77,6 +77,7 @@
 import org.openqa.selenium.grid.data.NodeHeartBeatEvent;
 import org.openqa.selenium.grid.data.NodeId;
 import org.openqa.selenium.grid.data.NodeRemovedEvent;
+import org.openqa.selenium.grid.data.NodeRestartedEvent;
 import org.openqa.selenium.grid.data.NodeStatus;
 import org.openqa.selenium.grid.data.NodeStatusEvent;
 import org.openqa.selenium.grid.data.RequestId;
@@ -205,6 +206,7 @@ public LocalDistributor(
 
     bus.addListener(NodeStatusEvent.listener(this::register));
     bus.addListener(NodeStatusEvent.listener(model::refresh));
+    bus.addListener(NodeRestartedEvent.listener(this::handleNodeRestarted));
     bus.addListener(NodeRemovedEvent.listener(nodeStatus -> remove(nodeStatus.getNodeId())));
     bus.addListener(
         NodeHeartBeatEvent.listener(
@@ -327,6 +329,25 @@ private void register(NodeStatus status) {
     }
   }
 
+  private void handleNodeRestarted(NodeStatus status) {
+    Require.nonNull("Node", status);
+    Lock writeLock = lock.writeLock();
+    writeLock.lock();
+    try {
+      if (!nodes.containsKey(status.getNodeId())) {
+        return;
+      }
+      if (!getNodeFromURI(status.getExternalUri()).isDraining()) {
+        LOG.info(
+            String.format(
+                "Node %s has restarted. Setting availability to DOWN.", status.getNodeId()));
+        model.setAvailability(status.getNodeId(), DOWN);
+      }
+    } finally {
+      writeLock.unlock();
+    }
+  }
+
   @Override
   public LocalDistributor add(Node node) {
     Require.nonNull("Node", node);