[grid] Node flag register-shutdown-on-failure

Signed-off-by: Viet Nguyen Duc <[email protected]>

Author: VietND96

java/src/org/openqa/selenium/grid/node/config/NodeFlags.java

@@ -198,6 +198,15 @@ public class NodeFlags implements HasRoles {
   @ConfigValue(section = NODE_SECTION, name = "register-period", example = "120")
   public int registerPeriod = DEFAULT_REGISTER_PERIOD;
 
+  @Parameter(
+      names = "--register-shutdown-on-failure",
+      description =
+          "If this flag is enabled, the Node will shut down after the register period is completed."
+              + " This is useful for container environments to restart and register again. If"
+              + " restarted multiple times, the Node container status will be CrashLoopBackOff.")
+  @ConfigValue(section = NODE_SECTION, name = "register-shutdown-on-failure", example = "false")
+  public boolean registerShutdownOnFailure = false;
+
   @Parameter(
       names = "--heartbeat-period",
       description =

java/src/org/openqa/selenium/grid/node/config/NodeOptions.java

@@ -205,6 +205,10 @@ public Duration getRegisterPeriod() {
     return Duration.ofSeconds(seconds);
   }
 
+  public boolean getRegisterShutdownOnFailure() {
+    return config.getBool(NODE_SECTION, "register-shutdown-on-failure").orElse(false);
+  }
+
   public Duration getHeartbeatPeriod() {
     // If the user sets 0 or less, we default to 1s.
     int seconds =

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

@@ -216,6 +216,18 @@ public NettyServer start() {
                 .withMaxDuration(nodeOptions.getRegisterPeriod())
                 .withDelay(nodeOptions.getRegisterCycle())
                 .handleResultIf(result -> true)
+                .onFailure(
+                    event -> {
+                      LOG.severe(
+                          String.format(
+                              "Registration event failed after period of %s seconds. Node will not"
+                                  + " attempt to register again",
+                              nodeOptions.getRegisterPeriod().getSeconds()));
+                      if (nodeOptions.getRegisterShutdownOnFailure()) {
+                        LOG.severe("Shutting down");
+                        System.exit(1);
+                      }
+                    })
                 .build();
 
         LOG.info("Starting registration process for Node " + node.getUri());