[grid] Querying Node status only once at registration

diemol · diemol · commit fb69a4f434c0 · 2022-02-22T08:12:43.000+01:00
We were querying the Node twice during its registration
process, this saves a network call which could be relevant
when many Nodes are being registrated.
diff --git a/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java b/java/src/org/openqa/selenium/grid/distributor/local/LocalDistributor.java
@@ -20,6 +20,7 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static org.openqa.selenium.grid.data.Availability.DOWN;
 import static org.openqa.selenium.grid.data.Availability.DRAINING;
+import static org.openqa.selenium.grid.data.Availability.UP;
 import static org.openqa.selenium.internal.Debug.getDebugLogLevel;
 import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES;
 import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES_EVENT;
@@ -44,6 +45,7 @@
 import org.openqa.selenium.concurrent.GuardedRunnable;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
+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.DistributorStatus;
@@ -91,6 +93,7 @@
 
 import java.io.Closeable;
 import java.io.UncheckedIOException;
+import java.net.URI;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -297,8 +300,10 @@ public LocalDistributor add(Node node) {
 
     // An exception occurs if Node heartbeat has started but the server is not ready.
     // Unhandled exception blocks the event-bus thread from processing any event henceforth.
+    NodeStatus initialNodeStatus;
     try {
-      model.add(node.getStatus());
+      initialNodeStatus = node.getStatus();
+      model.add(initialNodeStatus);
       nodes.put(node.getId(), node);
     } catch (Exception e) {
       LOG.log(
@@ -309,21 +314,10 @@ public LocalDistributor add(Node node) {
     }
 
     // Extract the health check
-    Runnable runnableHealthCheck = asRunnableHealthCheck(node);
-    allChecks.put(node.getId(), runnableHealthCheck);
-
-    // Running the health check right after the Node registers itself. We retry the
-    // execution because the Node might on a complex network topology. For example,
-    // Kubernetes pods with IPs that take a while before they are reachable.
-    RetryPolicy<Object> initialHealthCheckPolicy = new RetryPolicy<>()
-      .withMaxAttempts(-1)
-      .withMaxDuration(Duration.ofSeconds(90))
-      .withDelay(Duration.ofSeconds(15))
-      .abortIf(result -> true);
-
-    LOG.log(getDebugLogLevel(), "Running initial health check for Node " + node.getUri());
-    Executors.newSingleThreadExecutor().submit(
-      () -> Failsafe.with(initialHealthCheckPolicy).run(runnableHealthCheck::run));
+    Runnable healthCheck = asRunnableHealthCheck(node);
+    allChecks.put(node.getId(), healthCheck);
+
+    updateNodeStatus(initialNodeStatus, healthCheck);
 
     LOG.info(String.format(
       "Added node %s at %s. Health check every %ss",
@@ -336,6 +330,27 @@ public LocalDistributor add(Node node) {
     return this;
   }
 
+  private void updateNodeStatus(NodeStatus status, Runnable healthCheck) {
+    // Setting the Node as available if the initial call to status was successful.
+    // Otherwise, retry to have it available as soon as possible.
+    if (status.getAvailability() == UP) {
+      updateNodeAvailability(status.getExternalUri(), status.getNodeId(), status.getAvailability());
+    } else {
+      // Running the health check right after the Node registers itself. We retry the
+      // execution because the Node might on a complex network topology. For example,
+      // Kubernetes pods with IPs that take a while before they are reachable.
+      RetryPolicy<Object> initialHealthCheckPolicy = new RetryPolicy<>()
+        .withMaxAttempts(-1)
+        .withMaxDuration(Duration.ofSeconds(90))
+        .withDelay(Duration.ofSeconds(15))
+        .abortIf(result -> true);
+
+      LOG.log(getDebugLogLevel(), "Running health check for Node " + status.getExternalUri());
+      Executors.newSingleThreadExecutor().submit(
+        () -> Failsafe.with(initialHealthCheckPolicy).run(healthCheck::run));
+    }
+  }
+
   private Runnable runNodeHealthChecks() {
     return () -> {
       ImmutableMap<NodeId, Runnable> nodeHealthChecks = ImmutableMap.copyOf(allChecks);
@@ -363,26 +378,27 @@ private Runnable asRunnableHealthCheck(Node node) {
         failedCheckException = e;
       }
 
-      Lock writeLock = lock.writeLock();
-      writeLock.lock();
-      try {
-        LOG.log(
-          getDebugLogLevel(),
-          String.format(
-            "Health check result for %s was %s",
-            node.getUri(),
-            result.getAvailability()));
-        model.setAvailability(id, result.getAvailability());
-        model.updateHealthCheckCount(id, result.getAvailability());
-      } finally {
-        writeLock.unlock();
-      }
+      updateNodeAvailability(node.getUri(), id, result.getAvailability());
       if (checkFailed) {
         throw new HealthCheckFailedException("Node " + id, failedCheckException);
       }
     };
   }
 
+  private void updateNodeAvailability(URI nodeUri, NodeId id, Availability availability) {
+    Lock writeLock = lock.writeLock();
+    writeLock.lock();
+    try {
+      LOG.log(
+        getDebugLogLevel(),
+        String.format("Health check result for %s was %s", nodeUri, availability));
+      model.setAvailability(id, availability);
+      model.updateHealthCheckCount(id, availability);
+    } finally {
+      writeLock.unlock();
+    }
+  }
+
   @Override
   public boolean drain(NodeId nodeId) {
     Node node = nodes.get(nodeId);
