From 20c084f902baa404641c5f3d4907e22dd4b616a2 Mon Sep 17 00:00:00 2001
From: Diego Molina <diemol@gmail.com>
Date: Fri, 28 Jul 2023 17:41:47 +0200
Subject: [PATCH] [grid] Covering use case where a Node is UP after being
 marked DOWN

Fixes #12116
---
 .../openqa/selenium/grid/distributor/GridModel.java | 13 ++++++++++---
 .../grid/distributor/local/LocalDistributor.java    |  2 +-
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/java/src/org/openqa/selenium/grid/distributor/GridModel.java b/java/src/org/openqa/selenium/grid/distributor/GridModel.java
index ab535dc7e9873..c2b3d05898702 100644
--- a/java/src/org/openqa/selenium/grid/distributor/GridModel.java
+++ b/java/src/org/openqa/selenium/grid/distributor/GridModel.java
@@ -178,15 +178,22 @@ public void refresh(NodeStatus status) {
     }
   }
 
-  public void touch(NodeId id) {
-    Require.nonNull("Node ID", id);
+  public void touch(NodeStatus nodeStatus) {
+    Require.nonNull("Node ID", nodeStatus);
 
     Lock writeLock = lock.writeLock();
     writeLock.lock();
     try {
-      NodeStatus node = getNode(id);
+      NodeStatus node = getNode(nodeStatus.getNodeId());
       if (node != null) {
         nodePurgeTimes.put(node.getNodeId(), Instant.now());
+        // Covers the case where the Node might be DOWN in the Grid model (e.g. Node lost
+        // connectivity for a while). The Node reports itself back as UP.
+        if (node.getAvailability() != nodeStatus.getAvailability()
+            && nodeStatus.getAvailability() == UP) {
+          nodes.remove(node);
+          nodes.add(nodeStatus);
+        }
       }
     } finally {
       writeLock.unlock();
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 7ce25d741467a..5d9fff18cb070 100644
--- a/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java
+++ b/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java
@@ -205,7 +205,7 @@ public LocalDistributor(
         NodeHeartBeatEvent.listener(
             nodeStatus -> {
               if (nodes.containsKey(nodeStatus.getNodeId())) {
-                model.touch(nodeStatus.getNodeId());
+                model.touch(nodeStatus);
               } else {
                 register(nodeStatus);
               }