[grid] stop polling events on close

aguspe · Nov 12, 2024 · 9af4e15 · 9af4e15
1 parent 9ff0f85
commit 9af4e15
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/java/src/org/openqa/selenium/concurrent/BUILD.bazel b/java/src/org/openqa/selenium/concurrent/BUILD.bazel
@@ -4,6 +4,7 @@ java_library(
     name = "concurrent",
     srcs = glob(["*.java"]),
     visibility = [
+        "//java/src/org/openqa/selenium/events:__subpackages__",
         "//java/src/org/openqa/selenium/grid:__subpackages__",
         "//java/src/org/openqa/selenium/remote:__subpackages__",
     ],

diff --git a/java/src/org/openqa/selenium/events/zeromq/BUILD.bazel b/java/src/org/openqa/selenium/events/zeromq/BUILD.bazel
@@ -12,6 +12,7 @@ java_library(
     ],
     deps = [
         "//java/src/org/openqa/selenium:core",
+        "//java/src/org/openqa/selenium/concurrent",
         "//java/src/org/openqa/selenium/events",
         "//java/src/org/openqa/selenium/grid/config",
         "//java/src/org/openqa/selenium/grid/security",

diff --git a/java/src/org/openqa/selenium/events/zeromq/UnboundZmqEventBus.java b/java/src/org/openqa/selenium/events/zeromq/UnboundZmqEventBus.java
@@ -42,6 +42,7 @@
 import java.util.function.Consumer;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import org.openqa.selenium.concurrent.ExecutorServices;
 import org.openqa.selenium.events.Event;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.events.EventListener;
@@ -61,6 +62,7 @@ class UnboundZmqEventBus implements EventBus {
   private static final Logger LOG = Logger.getLogger(EventBus.class.getName());
   private static final Json JSON = new Json();
   private final AtomicBoolean pollingStarted = new AtomicBoolean(false);
+  private final PollingRunnable socketPolling;
   private final ExecutorService socketPollingExecutor;
   private final ExecutorService socketPublishingExecutor;
   private final ExecutorService listenerNotificationExecutor;
@@ -147,7 +149,8 @@ class UnboundZmqEventBus implements EventBus {
 
     LOG.info("Sockets created");
 
-    socketPollingExecutor.submit(new PollingRunnable(secret));
+    socketPolling = new PollingRunnable(secret);
+    socketPollingExecutor.submit(socketPolling);
 
     // Give ourselves up to a second to connect, using The World's Worst heuristic. If we don't
     // manage to connect, it's not the end of the world, as the socket we're connecting to may not
@@ -211,9 +214,11 @@ public void fire(Event event) {
 
   @Override
   public void close() {
-    socketPollingExecutor.shutdownNow();
-    socketPublishingExecutor.shutdownNow();
-    listenerNotificationExecutor.shutdownNow();
+    socketPolling.shutdown = true;
+    ExecutorServices.shutdownGracefully("Event Bus Poller", socketPollingExecutor);
+    ExecutorServices.shutdownGracefully("Event Bus Publisher", socketPublishingExecutor);
+    ExecutorServices.shutdownGracefully(
+        "Event Bus Listener Notifier", listenerNotificationExecutor);
     poller.close();
 
     if (sub != null) {
@@ -226,14 +231,15 @@ public void close() {
 
   private class PollingRunnable implements Runnable {
     private final Secret secret;
+    private volatile boolean shutdown;
 
     public PollingRunnable(Secret secret) {
       this.secret = secret;
     }
 
     @Override
     public void run() {
-      while (!Thread.currentThread().isInterrupted()) {
+      while (!Thread.currentThread().isInterrupted() && !shutdown) {
         try {
           int count = poller.poll(150);