[grid] Add event bus heartbeat to prevent steal connection

VietND96 · VietND96 · commit 7a60828b850b · 2025-10-16T02:11:01.000+07:00
Signed-off-by: Viet Nguyen Duc &lt;nguyenducviet4496@gmail.com&gt;
diff --git a/java/src/org/openqa/selenium/events/zeromq/BoundZmqEventBus.java b/java/src/org/openqa/selenium/events/zeromq/BoundZmqEventBus.java
@@ -20,6 +20,7 @@
 import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
+import java.time.Duration;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.logging.Level;
@@ -43,7 +44,11 @@ class BoundZmqEventBus implements EventBus {
   private final ExecutorService executor;
 
   BoundZmqEventBus(
-      ZContext context, String publishConnection, String subscribeConnection, Secret secret) {
+      ZContext context,
+      String publishConnection,
+      String subscribeConnection,
+      Secret secret,
+      Duration heartbeatPeriod) {
     String address = new NetworkUtils().getHostAddress();
     Addresses xpubAddr = deriveAddresses(address, publishConnection);
     Addresses xsubAddr = deriveAddresses(address, subscribeConnection);
@@ -53,11 +58,13 @@ class BoundZmqEventBus implements EventBus {
     xpub = context.createSocket(SocketType.XPUB);
     xpub.setIPv6(xpubAddr.isIPv6);
     xpub.setImmediate(true);
+    configureHeartbeat(xpub, heartbeatPeriod, "XPUB");
     xpub.bind(xpubAddr.bindTo);
 
     xsub = context.createSocket(SocketType.XSUB);
     xsub.setIPv6(xsubAddr.isIPv6);
     xsub.setImmediate(true);
+    configureHeartbeat(xsub, heartbeatPeriod, "XSUB");
     xsub.bind(xsubAddr.bindTo);
 
     executor =
@@ -69,7 +76,24 @@ class BoundZmqEventBus implements EventBus {
             });
     executor.submit(() -> ZMQ.proxy(xsub, xpub, null));
 
-    delegate = new UnboundZmqEventBus(context, xpubAddr.advertise, xsubAddr.advertise, secret);
+    delegate =
+        new UnboundZmqEventBus(
+            context, xpubAddr.advertise, xsubAddr.advertise, secret, heartbeatPeriod);
+  }
+
+  private void configureHeartbeat(ZMQ.Socket socket, Duration heartbeatPeriod, String socketType) {
+    if (heartbeatPeriod != null && !heartbeatPeriod.isZero() && !heartbeatPeriod.isNegative()) {
+      int heartbeatIvl = (int) heartbeatPeriod.toMillis();
+      socket.setHeartbeatIvl(heartbeatIvl);
+      // Set heartbeat timeout to 3x the interval
+      socket.setHeartbeatTimeout(heartbeatIvl * 3);
+      // Set heartbeat TTL to 6x the interval
+      socket.setHeartbeatTtl(heartbeatIvl * 6);
+      LOG.info(
+          String.format(
+              "Event bus %s socket heartbeat configured: interval=%dms, timeout=%dms, ttl=%dms",
+              socketType, heartbeatIvl, heartbeatIvl * 3, heartbeatIvl * 6));
+    }
   }
 
   @Override
diff --git a/java/src/org/openqa/selenium/events/zeromq/UnboundZmqEventBus.java b/java/src/org/openqa/selenium/events/zeromq/UnboundZmqEventBus.java
@@ -27,6 +27,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.UnknownHostException;
+import java.time.Duration;
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.LinkedList;
@@ -76,7 +77,11 @@ class UnboundZmqEventBus implements EventBus {
   private ZMQ.Socket sub;
 
   UnboundZmqEventBus(
-      ZContext context, String publishConnection, String subscribeConnection, Secret secret) {
+      ZContext context,
+      String publishConnection,
+      String subscribeConnection,
+      Secret secret,
+      Duration heartbeatPeriod) {
     Require.nonNull("Secret", secret);
     StringBuilder builder = new StringBuilder();
     try (JsonOutput out = JSON.newOutput(builder)) {
@@ -136,11 +141,13 @@ class UnboundZmqEventBus implements EventBus {
             () -> {
               sub = context.createSocket(SocketType.SUB);
               sub.setIPv6(isSubAddressIPv6(publishConnection));
+              configureHeartbeat(sub, heartbeatPeriod, "SUB");
               sub.connect(publishConnection);
               sub.subscribe(new byte[0]);
 
               pub = context.createSocket(SocketType.PUB);
               pub.setIPv6(isSubAddressIPv6(subscribeConnection));
+              configureHeartbeat(pub, heartbeatPeriod, "PUB");
               pub.connect(subscribeConnection);
             });
     // Connections are already established
@@ -172,6 +179,21 @@ public boolean isReady() {
     return !socketPollingExecutor.isShutdown();
   }
 
+  private void configureHeartbeat(ZMQ.Socket socket, Duration heartbeatPeriod, String socketType) {
+    if (heartbeatPeriod != null && !heartbeatPeriod.isZero() && !heartbeatPeriod.isNegative()) {
+      int heartbeatIvl = (int) heartbeatPeriod.toMillis();
+      socket.setHeartbeatIvl(heartbeatIvl);
+      // Set heartbeat timeout to 3x the interval
+      socket.setHeartbeatTimeout(heartbeatIvl * 3);
+      // Set heartbeat TTL to 6x the interval
+      socket.setHeartbeatTtl(heartbeatIvl * 6);
+      LOG.info(
+          String.format(
+              "Event bus %s socket heartbeat configured: interval=%dms, timeout=%dms, ttl=%dms",
+              socketType, heartbeatIvl, heartbeatIvl * 3, heartbeatIvl * 6));
+    }
+  }
+
   private boolean isSubAddressIPv6(String connection) {
     try {
       URI uri = new URI(connection);
diff --git a/java/src/org/openqa/selenium/events/zeromq/ZeroMqEventBus.java b/java/src/org/openqa/selenium/events/zeromq/ZeroMqEventBus.java
@@ -21,6 +21,7 @@
 
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.function.Consumer;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.events.EventListener;
@@ -45,10 +46,20 @@ private ZeroMqEventBus() {
 
   public static EventBus create(
       ZContext context, String publish, String subscribe, boolean bind, Secret secret) {
+    return create(context, publish, subscribe, bind, secret, Duration.ofSeconds(60));
+  }
+
+  public static EventBus create(
+      ZContext context,
+      String publish,
+      String subscribe,
+      boolean bind,
+      Secret secret,
+      Duration heartbeatPeriod) {
     if (bind) {
-      return new BoundZmqEventBus(context, publish, subscribe, secret);
+      return new BoundZmqEventBus(context, publish, subscribe, secret, heartbeatPeriod);
     }
-    return new UnboundZmqEventBus(context, publish, subscribe, secret);
+    return new UnboundZmqEventBus(context, publish, subscribe, secret, heartbeatPeriod);
   }
 
   public static EventBus create(Config config) {
@@ -85,10 +96,19 @@ public static EventBus create(Config config) {
                 });
 
     boolean bind = config.getBool(EVENTS_SECTION, "bind").orElse(false);
+    int heartbeatPeriodSeconds =
+        config.getInt(EVENTS_SECTION, "eventbus-heartbeat-period").orElse(60);
+    Duration heartbeatPeriod = Duration.ofSeconds(heartbeatPeriodSeconds);
 
     SecretOptions secretOptions = new SecretOptions(config);
 
-    return create(new ZContext(), publish, subscribe, bind, secretOptions.getRegistrationSecret());
+    return create(
+        new ZContext(),
+        publish,
+        subscribe,
+        bind,
+        secretOptions.getRegistrationSecret(),
+        heartbeatPeriod);
   }
 
   private static String mungeUri(URI base, String scheme, int port) {
diff --git a/java/src/org/openqa/selenium/grid/server/EventBusFlags.java b/java/src/org/openqa/selenium/grid/server/EventBusFlags.java
@@ -65,6 +65,12 @@ public class EventBusFlags implements HasRoles {
       example = "\"org.openqa.selenium.events.zeromq.ZeroMqEventBus\"")
   private String implementation;
 
+  @Parameter(
+      names = {"--eventbus-heartbeat-period"},
+      description = "How often, in seconds, will the EventBus socket send heartbeats")
+  @ConfigValue(section = EVENTS_SECTION, name = "eventbus-heartbeat-period", example = "30")
+  private int eventbusHeartbeatPeriod;
+
   @Override
   public Set<Role> getRoles() {
     return ImmutableSet.of(EVENT_BUS_ROLE);
diff --git a/java/src/org/openqa/selenium/grid/server/EventBusOptions.java b/java/src/org/openqa/selenium/grid/server/EventBusOptions.java
@@ -17,6 +17,7 @@
 
 package org.openqa.selenium.grid.server;
 
+import java.time.Duration;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
 import org.openqa.selenium.internal.Require;
@@ -25,6 +26,7 @@ public class EventBusOptions {
 
   static final String EVENTS_SECTION = "events";
   private static final String DEFAULT_CLASS = "org.openqa.selenium.events.zeromq.ZeroMqEventBus";
+  private static final int DEFAULT_HEARTBEAT_PERIOD = 60;
   private final Config config;
   private volatile EventBus bus;
 
@@ -47,6 +49,12 @@ public EventBus getEventBus() {
     return localBus;
   }
 
+  public Duration getHeartbeatPeriod() {
+    int period =
+        config.getInt(EVENTS_SECTION, "eventbus-heartbeat-period").orElse(DEFAULT_HEARTBEAT_PERIOD);
+    return Duration.ofSeconds(period);
+  }
+
   private EventBus createBus() {
     return config.getClass(EVENTS_SECTION, "implementation", EventBus.class, DEFAULT_CLASS);
   }
