[grid] Expose register status via Node status response (#15448)

Author: VietND96

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

@@ -704,7 +704,8 @@ private SlotId reserveSlot(RequestId requestId, Capabilities caps) {
   }
 
   private boolean isNotSupported(Capabilities caps) {
-    return getAvailableNodes().stream().noneMatch(node -> node.hasCapability(caps, slotMatcher));
+    return getAvailableNodes().stream()
+        .noneMatch(node -> node.hasCapability(caps, slotMatcher) && node.getAvailability() == UP);
   }
 
   private boolean reserve(SlotId id) {
@@ -794,7 +795,7 @@ public void run() {
         // up starving a session request.
         Map<Capabilities, Long> stereotypes =
             getAvailableNodes().stream()
-                .filter(NodeStatus::hasCapacity)
+                .filter(node -> node.hasCapacity() && node.getAvailability() == UP)
                 .flatMap(node -> node.getSlots().stream().map(Slot::getStereotype))
                 .collect(
                     Collectors.groupingBy(ImmutableCapabilities::copyOf, Collectors.counting()));

java/src/org/openqa/selenium/grid/distributor/selector/DefaultSlotSelector.java

@@ -18,6 +18,7 @@
 package org.openqa.selenium.grid.distributor.selector;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static org.openqa.selenium.grid.data.Availability.UP;
 
 import com.google.common.annotations.VisibleForTesting;
 import java.util.Comparator;
@@ -48,7 +49,7 @@ public Set<SlotId> selectSlot(
     // Nodes).
     // After that, Nodes are ordered by their load, last session creation, and their id.
     return nodes.stream()
-        .filter(node -> node.hasCapacity(capabilities, slotMatcher))
+        .filter(node -> node.hasCapacity(capabilities, slotMatcher) && node.getAvailability() == UP)
         .sorted(
             Comparator.comparingLong(this::getNumberOfSupportedBrowsers)
                 // Now sort by node which has the lowest load (natural ordering)

java/src/org/openqa/selenium/grid/node/Node.java

@@ -32,6 +32,7 @@
 import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
@@ -130,7 +131,8 @@ public abstract class Node implements HasReadyState, Routable {
   private final URI uri;
   private final Duration sessionTimeout;
   private final Route routes;
-  protected boolean draining;
+  protected final AtomicBoolean draining = new AtomicBoolean(false);
+  protected final AtomicBoolean registered = new AtomicBoolean(false);
 
   protected Node(
       Tracer tracer, NodeId id, URI uri, Secret registrationSecret, Duration sessionTimeout) {
@@ -271,7 +273,15 @@ public Duration getSessionTimeout() {
   }
 
   public boolean isDraining() {
-    return draining;
+    return draining.get();
+  }
+
+  public boolean isRegistered() {
+    return registered.get();
+  }
+
+  public void register() {
+    registered.set(true);
   }
 
   public abstract void drain();

java/src/org/openqa/selenium/grid/node/StatusHandler.java

@@ -47,6 +47,8 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
                 status.hasCapacity(),
                 "message",
                 status.hasCapacity() ? "Ready" : "No free slots available",
+                "registered",
+                node.isRegistered(),
                 "node",
                 status));
 

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

@@ -37,7 +37,6 @@
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.openqa.selenium.BuildInfo;
@@ -74,7 +73,6 @@
 public class NodeServer extends TemplateGridServerCommand {
 
   private static final Logger LOG = Logger.getLogger(NodeServer.class.getName());
-  private final AtomicBoolean nodeRegistered = new AtomicBoolean(false);
   private Node node;
   private EventBus bus;
   private final Thread shutdownHook =
@@ -130,7 +128,7 @@ protected Handlers createHandlers(Config config) {
 
     HttpHandler readinessCheck =
         req -> {
-          if (node.getStatus().hasCapacity()) {
+          if (node.isReady() && node.getStatus().hasCapacity()) {
             return new HttpResponse()
                 .setStatus(HTTP_OK)
                 .setHeader("Content-Type", MediaType.PLAIN_TEXT_UTF_8.toString())
@@ -147,7 +145,7 @@ protected Handlers createHandlers(Config config) {
         NodeAddedEvent.listener(
             nodeId -> {
               if (node.getId().equals(nodeId)) {
-                nodeRegistered.set(true);
+                node.register();
                 LOG.info("Node has been added");
               }
             }));
@@ -237,7 +235,7 @@ public NettyServer start() {
               Failsafe.with(registrationPolicy)
                   .run(
                       () -> {
-                        if (nodeRegistered.get()) {
+                        if (node.isRegistered()) {
                           throw new InterruptedException("Stopping registration thread.");
                         }
                         HealthCheck.Result check = node.getHealthCheck().check();

java/src/org/openqa/selenium/grid/node/k8s/OneShotNode.java

@@ -415,7 +415,7 @@ public NodeStatus getStatus() {
   @Override
   public void drain() {
     events.fire(new NodeDrainStarted(getId()));
-    draining = true;
+    draining.set(true);
   }
 
   @Override

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

@@ -1015,7 +1015,7 @@ public void drain() {
       AttributeMap attributeMap = tracer.createAttributeMap();
       attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
       bus.fire(new NodeDrainStarted(getId()));
-      draining = true;
+      draining.set(true);
       // Ensure the pendingSessions counter will not be decremented by timed out sessions not
       // included
       // in the currentSessionCount and the NodeDrainComplete will be raised to early.

java/src/org/openqa/selenium/grid/node/remote/RemoteNode.java

@@ -295,7 +295,7 @@ public void drain() {
     HttpResponse res = client.with(addSecret).execute(req);
 
     if (res.getStatus() == HTTP_OK) {
-      draining = true;
+      draining.set(true);
     }
   }
 

java/test/org/openqa/selenium/grid/router/StressTest.java

@@ -137,7 +137,7 @@ void multipleSimultaneousSessions() throws Exception {
               executor);
     }
 
-    CompletableFuture.allOf(futures).get(4, MINUTES);
+    CompletableFuture.allOf(futures).get(6, MINUTES);
   }
 
   @Test
@@ -190,6 +190,6 @@ void multipleSimultaneousSessionsTimedOut() throws Exception {
               executor);
     }
 
-    CompletableFuture.allOf(futures).get(5, MINUTES);
+    CompletableFuture.allOf(futures).get(6, MINUTES);
   }
 }