From 8a04527feb4744a9b9802d2095435158ffe0ec07 Mon Sep 17 00:00:00 2001
From: Viet Nguyen Duc <nguyenducviet4496@gmail.com>
Date: Wed, 25 Dec 2024 04:13:42 +0530
Subject: [PATCH] [grid] Distributor listen and handle the NodeRestartedEvent
 (#14938)

Signed-off-by: Viet Nguyen Duc <nguyenducviet4496@gmail.com>
---
 .../distributor/local/LocalDistributor.java   | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java b/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java
index fafa4d806bae7..926dceb864c03 100644
--- 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);