[grid] Use heartbeat to register the Node (#9321)

* [grid] Use heartbeat to register the Node

Fixes #9297

Co-authored-by: Diego Molina <[email protected]>
Co-authored-by: Diego Molina <[email protected]>

Author: pujagani

java/server/src/org/openqa/selenium/grid/data/NodeHeartBeatEvent.java

@@ -28,13 +28,13 @@ public class NodeHeartBeatEvent extends Event {
 
   private static final EventName NODE_HEARTBEAT = new EventName("node-heartbeat");
 
-  public NodeHeartBeatEvent(NodeId nodeId) {
-    super(NODE_HEARTBEAT, Require.nonNull("Node id", nodeId));
+  public NodeHeartBeatEvent(NodeStatus status) {
+    super(NODE_HEARTBEAT, Require.nonNull("Node status", status));
   }
 
-  public static EventListener<NodeId> listener(Consumer<NodeId> handler) {
+  public static EventListener<NodeStatus> listener(Consumer<NodeStatus> handler) {
     Require.nonNull("Handler", handler);
 
-    return new EventListener<NodeId>(NODE_HEARTBEAT, NodeId.class, handler);
+    return new EventListener<NodeStatus>(NODE_HEARTBEAT, NodeStatus.class, handler);
   }
 }

java/server/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java

@@ -79,6 +79,7 @@
 import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -130,14 +131,20 @@ public LocalDistributor(
     this.clientFactory = Require.nonNull("HTTP client factory", clientFactory);
     this.sessions = Require.nonNull("Session map", sessions);
     this.model = new GridModel(bus);
-    this.nodes = new HashMap<>();
+    this.nodes = new ConcurrentHashMap<>();
     this.sessionRequests = Require.nonNull("New Session Request Queue", sessionRequests);
     this.registrationSecret = Require.nonNull("Registration secret", registrationSecret);
     this.healthcheckInterval = Require.nonNull("Health check interval", healthcheckInterval);
 
     bus.addListener(NodeStatusEvent.listener(this::register));
     bus.addListener(NodeStatusEvent.listener(model::refresh));
-    bus.addListener(NodeHeartBeatEvent.listener(model::touch));
+    bus.addListener(NodeHeartBeatEvent.listener(nodeStatus -> {
+      if (nodes.containsKey(nodeStatus.getId())) {
+        model.touch(nodeStatus.getId());
+      } else {
+        register(nodeStatus);
+      }
+    }));
     bus.addListener(NodeDrainComplete.listener(this::remove));
     bus.addListener(NewSessionRequestEvent.listener(requestIds::offer));
 

java/server/src/org/openqa/selenium/grid/node/httpd/NodeServer.java

@@ -17,14 +17,6 @@
 
 package org.openqa.selenium.grid.node.httpd;
 
-import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
-import static java.net.HttpURLConnection.HTTP_NO_CONTENT;
-import static org.openqa.selenium.grid.config.StandardGridRoles.EVENT_BUS_ROLE;
-import static org.openqa.selenium.grid.config.StandardGridRoles.HTTPD_ROLE;
-import static org.openqa.selenium.grid.config.StandardGridRoles.NODE_ROLE;
-import static org.openqa.selenium.grid.data.Availability.DOWN;
-import static org.openqa.selenium.remote.http.Route.get;
-
 import com.google.auto.service.AutoService;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.MediaType;
@@ -67,6 +59,14 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
 
+import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR;
+import static java.net.HttpURLConnection.HTTP_NO_CONTENT;
+import static org.openqa.selenium.grid.config.StandardGridRoles.EVENT_BUS_ROLE;
+import static org.openqa.selenium.grid.config.StandardGridRoles.HTTPD_ROLE;
+import static org.openqa.selenium.grid.config.StandardGridRoles.NODE_ROLE;
+import static org.openqa.selenium.grid.data.Availability.DOWN;
+import static org.openqa.selenium.remote.http.Route.get;
+
 @AutoService(CliCommand.class)
 public class NodeServer extends TemplateGridServerCommand {
 

java/server/src/org/openqa/selenium/grid/node/local/LocalNode.java

@@ -50,9 +50,9 @@
 import org.openqa.selenium.grid.data.NodeAddedEvent;
 import org.openqa.selenium.grid.data.NodeDrainComplete;
 import org.openqa.selenium.grid.data.NodeDrainStarted;
-import org.openqa.selenium.grid.data.NodeHeartBeatEvent;
 import org.openqa.selenium.grid.data.NodeId;
 import org.openqa.selenium.grid.data.NodeStatus;
+import org.openqa.selenium.grid.data.NodeHeartBeatEvent;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.data.Slot;
@@ -181,7 +181,7 @@ private LocalNode(
         // Lets avoid to create more than one "Regularly" when the Node registers again.
         if (!heartBeatStarted.getAndSet(true)) {
           regularly.submit(
-            () -> bus.fire(new NodeHeartBeatEvent(getId())), heartbeatPeriod, heartbeatPeriod);
+            () -> bus.fire(new NodeHeartBeatEvent(getStatus())), heartbeatPeriod, heartbeatPeriod);
         }
       }
     }));

java/server/test/org/openqa/selenium/grid/distributor/AddingNodesTest.java

@@ -19,6 +19,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.openqa.selenium.Capabilities;

java/server/test/org/openqa/selenium/grid/distributor/DistributorTest.java

@@ -187,9 +187,9 @@ public void shouldStartHeartBeatOnNodeRegistration() {
     AtomicBoolean heartbeatStarted = new AtomicBoolean();
     CountDownLatch latch = new CountDownLatch(1);
 
-    bus.addListener(NodeHeartBeatEvent.listener(nodeId -> {
+    bus.addListener(NodeHeartBeatEvent.listener(nodeStatus -> {
       latch.countDown();
-      if (node.getId().equals(nodeId)) {
+      if (node.getId().equals(nodeStatus.getId())) {
         heartbeatStarted.set(true);
       }
     }));

java/server/test/org/openqa/selenium/grid/distributor/local/LocalDistributorTest.java

@@ -18,6 +18,7 @@
 package org.openqa.selenium.grid.distributor.local;
 
 import com.google.common.collect.ImmutableMap;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.openqa.selenium.Capabilities;