[grid] Send node heartbeat only on successful node registration (#9213)

* [grid] Send node heartbeat only on successful node registration

* [grid] Update node registration test to wait for heartbeat event

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

Author: 3 people

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

@@ -197,7 +197,7 @@ public NettyServer start() {
         Executors.newSingleThreadExecutor().submit(() -> {
           Failsafe.with(registrationPolicy).run(
             () -> {
-              LOG.fine("Sending registration event");
+              LOG.info("Sending registration event");
               HealthCheck.Result check = node.getHealthCheck().check();
               if (DOWN.equals(check.getAvailability())) {
                 LOG.severe("Node is not alive: " + check.getMessage());

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

@@ -36,6 +36,7 @@
 import org.openqa.selenium.grid.data.Availability;
 import org.openqa.selenium.grid.data.CreateSessionRequest;
 import org.openqa.selenium.grid.data.CreateSessionResponse;
+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;
@@ -171,7 +172,13 @@ private LocalNode(
     this.regularly = new Regularly("Local Node: " + externalUri);
     regularly.submit(currentSessions::cleanUp, Duration.ofSeconds(30), Duration.ofSeconds(30));
     regularly.submit(tempFileSystems::cleanUp, Duration.ofSeconds(30), Duration.ofSeconds(30));
-    regularly.submit(() -> bus.fire(new NodeHeartBeatEvent(getId())), heartbeatPeriod, heartbeatPeriod);
+
+    bus.addListener(NodeAddedEvent.listener(nodeId -> {
+      if (getId().equals(nodeId)) {
+        regularly.submit(() ->
+          bus.fire(new NodeHeartBeatEvent(getId())), heartbeatPeriod, heartbeatPeriod);
+      }
+    }));
 
     bus.addListener(SessionClosedEvent.listener(id -> {
       try {

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.assertj.core.api.AbstractAssert;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
@@ -35,6 +36,7 @@
 import org.openqa.selenium.grid.data.CreateSessionRequest;
 import org.openqa.selenium.grid.data.CreateSessionResponse;
 import org.openqa.selenium.grid.data.DistributorStatus;
+import org.openqa.selenium.grid.data.NodeHeartBeatEvent;
 import org.openqa.selenium.grid.data.NodeRemovedEvent;
 import org.openqa.selenium.grid.data.NodeStatus;
 import org.openqa.selenium.grid.data.Session;
@@ -81,6 +83,7 @@
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -148,6 +151,56 @@ public void creatingANewSessionWithoutANodeEndsInFailure() {
     }
   }
 
+  @Test
+  public void shouldStartHeartBeatOnNodeRegistration() {
+    EventBus bus = new GuavaEventBus();
+    LocalSessionMap sessions = new LocalSessionMap(tracer, bus);
+    LocalNewSessionQueue localNewSessionQueue = new LocalNewSessionQueue(
+      tracer,
+      bus,
+      Duration.ofSeconds(2),
+      Duration.ofSeconds(2));
+    LocalNewSessionQueuer queuer = new LocalNewSessionQueuer(
+      tracer,
+      bus,
+      localNewSessionQueue,
+      registrationSecret);
+    LocalNode node = LocalNode.builder(tracer, bus, routableUri, routableUri, registrationSecret)
+      .add(
+        caps,
+        new TestSessionFactory((id, c) -> new Session(id, nodeUri, stereotype, c, Instant.now())))
+      .build();
+
+    Distributor distributor = new LocalDistributor(
+      tracer,
+      bus,
+      new PassthroughHttpClient.Factory(node),
+      sessions,
+      queuer,
+      registrationSecret);
+    distributor.add(node);
+
+    AtomicBoolean heartbeatStarted = new AtomicBoolean();
+    CountDownLatch latch = new CountDownLatch(1);
+
+    bus.addListener(NodeHeartBeatEvent.listener(nodeId -> {
+      latch.countDown();
+      if (node.getId().equals(nodeId)) {
+        heartbeatStarted.set(true);
+      }
+    }));
+    waitToHaveCapacity(distributor);
+    boolean eventFired = false;
+    try {
+      eventFired = latch.await(30, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      Assert.fail("Thread Interrupted");
+    }
+
+    assertThat(eventFired).isTrue();
+    assertThat(heartbeatStarted).isTrue();
+  }
+
   @Test
   public void shouldBeAbleToAddANodeAndCreateASession() {
     LocalSessionMap sessions = new LocalSessionMap(tracer, bus);