[grid] Standalone should also shutdown when it has been drained

diemol · diemol · commit e3d3da37922b · 2022-04-07T15:20:42.000+02:00
diff --git a/java/src/org/openqa/selenium/grid/commands/BUILD.bazel b/java/src/org/openqa/selenium/grid/commands/BUILD.bazel
@@ -20,6 +20,7 @@ java_library(
         "//java/src/org/openqa/selenium/events",
         "//java/src/org/openqa/selenium/grid:base-command",
         "//java/src/org/openqa/selenium/grid/config",
+        "//java/src/org/openqa/selenium/grid/data",
         "//java/src/org/openqa/selenium/grid/distributor",
         "//java/src/org/openqa/selenium/grid/distributor/config",
         "//java/src/org/openqa/selenium/grid/distributor/local",
diff --git a/java/src/org/openqa/selenium/grid/commands/Standalone.java b/java/src/org/openqa/selenium/grid/commands/Standalone.java
@@ -27,6 +27,7 @@
 import org.openqa.selenium.grid.TemplateGridServerCommand;
 import org.openqa.selenium.grid.config.Config;
 import org.openqa.selenium.grid.config.Role;
+import org.openqa.selenium.grid.data.NodeDrainComplete;
 import org.openqa.selenium.grid.distributor.Distributor;
 import org.openqa.selenium.grid.distributor.config.DistributorOptions;
 import org.openqa.selenium.grid.distributor.local.LocalDistributor;
@@ -204,6 +205,27 @@ protected Handlers createHandlers(Config config) {
     combinedHandler.addHandler(node);
     distributor.add(node);
 
+    bus.addListener(NodeDrainComplete.listener(nodeId -> {
+      if (!node.getId().equals(nodeId)) {
+        return;
+      }
+
+      // Wait a beat before shutting down so the final response from the
+      // node can escape.
+      new Thread(
+        () -> {
+          try {
+            Thread.sleep(1000);
+          } catch (InterruptedException e) {
+            // Swallow, the next thing we're doing is shutting down
+          }
+          LOG.info("Shutting down");
+          System.exit(0);
+        },
+        "Standalone shutdown: " + nodeId)
+        .start();
+    }));
+
     return new Handlers(httpHandler, new ProxyNodeWebsockets(clientFactory, node));
   }
 
