Add a test to ensure concurrent access to the event bus is relatively safe

Author: shs96c

java/server/src/org/openqa/selenium/events/zeromq/BoundZmqEventBus.java

@@ -68,8 +68,6 @@ class BoundZmqEventBus implements EventBus {
     executor.submit(() -> ZMQ.proxy(xsub, xpub, null));
 
     delegate = new UnboundZmqEventBus(context, xpubAddr.advertise, xsubAddr.advertise, secret);
-
-    LOG.info("Event bus ready");
   }
 
   @Override

java/server/src/org/openqa/selenium/events/zeromq/UnboundZmqEventBus.java

@@ -48,6 +48,9 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.logging.Level;
@@ -60,7 +63,9 @@ class UnboundZmqEventBus implements EventBus {
   static final EventName REJECTED_EVENT = new EventName("selenium-rejected-event");
   private static final Logger LOG = Logger.getLogger(EventBus.class.getName());
   private static final Json JSON = new Json();
-  private final ExecutorService executor;
+  private final ScheduledExecutorService socketPollingExecutor;
+  private final ExecutorService listenerNotificationExecutor;
+
   private final Map<EventName, List<Consumer<Event>>> listeners = new ConcurrentHashMap<>();
   private final Queue<UUID> recentMessages = EvictingQueue.create(128);
   private final String encodedSecret;
@@ -76,12 +81,16 @@ class UnboundZmqEventBus implements EventBus {
     }
     this.encodedSecret = builder.toString();
 
-    executor = Executors.newSingleThreadExecutor(r -> {
+    ThreadFactory threadFactory = r -> {
       Thread thread = new Thread(r);
       thread.setName("Event Bus");
       thread.setDaemon(true);
       return thread;
-    });
+    };
+    this.socketPollingExecutor = Executors.newSingleThreadScheduledExecutor(threadFactory);
+    this.listenerNotificationExecutor = Executors.newFixedThreadPool(
+      Math.max(Runtime.getRuntime().availableProcessors() / 2, 2), // At least two threads
+      threadFactory);
 
     String connectionMessage = String.format("Connecting to %s and %s", publishConnection, subscribeConnection);
     LOG.info(connectionMessage);
@@ -93,6 +102,7 @@ class UnboundZmqEventBus implements EventBus {
       .onRetry(e -> LOG.log(Level.WARNING, String.format("Failure #%s. Retrying.", e.getAttemptCount())))
       .onRetriesExceeded(e -> LOG.log(Level.WARNING, "Connection aborted."));
 
+    // Access to the zmq socket is safe here: no threads.
     Failsafe.with(retryPolicy).run(
       () -> {
         sub = context.createSocket(SocketType.SUB);
@@ -113,54 +123,11 @@ class UnboundZmqEventBus implements EventBus {
 
     AtomicBoolean pollingStarted = new AtomicBoolean(false);
 
-    executor.submit(() -> {
-      LOG.info("Bus started");
-      while (!Thread.currentThread().isInterrupted()) {
-        try {
-          poller.poll(150);
-          pollingStarted.lazySet(true);
-
-          if (poller.pollin(0)) {
-            ZMQ.Socket socket = poller.getSocket(0);
-
-            EventName eventName = new EventName(new String(socket.recv(ZMQ.DONTWAIT), UTF_8));
-            Secret eventSecret = JSON.toType(new String(socket.recv(ZMQ.DONTWAIT), UTF_8), Secret.class);
-            UUID id = UUID.fromString(new String(socket.recv(ZMQ.DONTWAIT), UTF_8));
-            String data = new String(socket.recv(ZMQ.DONTWAIT), UTF_8);
-
-            Object converted = JSON.toType(data, Object.class);
-            Event event = new Event(id, eventName, converted);
-
-            if (recentMessages.contains(id)) {
-              continue;
-            }
-            recentMessages.add(id);
-
-            if (!Secret.matches(secret, eventSecret)) {
-              LOG.severe(String.format("Received message without a valid secret. Rejecting. %s -> %s", event, data));
-              Event rejectedEvent = new Event(REJECTED_EVENT, new ZeroMqEventBus.RejectedEvent(eventName, data));
-
-              listeners.getOrDefault(REJECTED_EVENT, new ArrayList<>())
-                .forEach(listener -> listener.accept(rejectedEvent));
-              return;
-            }
-
-            List<Consumer<Event>> typeListeners = listeners.get(eventName);
-            if (typeListeners == null) {
-              continue;
-            }
-
-            typeListeners.parallelStream().forEach(listener -> listener.accept(event));
-          }
-        } catch (Throwable e) {
-          if (e.getCause() != null && e.getCause() instanceof AssertionError) {
-            // Do nothing.
-          } else {
-            throw e;
-          }
-        }
-      }
-    });
+    socketPollingExecutor.scheduleWithFixedDelay(
+      () -> pollForIncomingEvents(poller, secret, pollingStarted),
+      0,
+      100,
+      TimeUnit.MILLISECONDS);
 
     // 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
@@ -173,11 +140,13 @@ class UnboundZmqEventBus implements EventBus {
         throw new RuntimeException(e);
       }
     }
+
+    LOG.info("Event bus ready");
   }
 
   @Override
   public boolean isReady() {
-    return !executor.isShutdown();
+    return !socketPollingExecutor.isShutdown();
   }
 
   private boolean isSubAddressIPv6(String connection) {
@@ -207,20 +176,83 @@ public void addListener(EventListener<?> listener) {
   public void fire(Event event) {
     Require.nonNull("Event to send", event);
 
-    pub.sendMore(event.getType().getName().getBytes(UTF_8));
-    pub.sendMore(encodedSecret.getBytes(UTF_8));
-    pub.sendMore(event.getId().toString().getBytes(UTF_8));
-    pub.send(event.getRawData().getBytes(UTF_8));
+    socketPollingExecutor.execute(() -> {
+      pub.sendMore(event.getType().getName().getBytes(UTF_8));
+      pub.sendMore(encodedSecret.getBytes(UTF_8));
+      pub.sendMore(event.getId().toString().getBytes(UTF_8));
+      pub.send(event.getRawData().getBytes(UTF_8));
+    });
   }
 
   @Override
   public void close() {
-    executor.shutdown();
+    socketPollingExecutor.shutdownNow();
+    listenerNotificationExecutor.shutdownNow();
+
     if (sub != null) {
       sub.close();
     }
     if (pub != null) {
       pub.close();
     }
   }
+
+  private void pollForIncomingEvents(ZMQ.Poller poller, Secret secret, AtomicBoolean pollingStarted) {
+    try {
+      int count = poller.poll(0);
+
+      pollingStarted.lazySet(true);
+
+      for (int i = 0; i < count; i++) {
+        if (poller.pollin(i)) {
+          ZMQ.Socket socket = poller.getSocket(i);
+
+          EventName eventName = new EventName(new String(socket.recv(ZMQ.DONTWAIT), UTF_8));
+          Secret eventSecret =
+              JSON.toType(new String(socket.recv(ZMQ.DONTWAIT), UTF_8), Secret.class);
+          UUID id = UUID.fromString(new String(socket.recv(ZMQ.DONTWAIT), UTF_8));
+          String data = new String(socket.recv(ZMQ.DONTWAIT), UTF_8);
+
+          Object converted = JSON.toType(data, Object.class);
+          Event event = new Event(id, eventName, converted);
+
+          if (recentMessages.contains(id)) {
+            return;
+          }
+          recentMessages.add(id);
+
+          if (!Secret.matches(secret, eventSecret)) {
+            LOG.severe(
+                String.format(
+                    "Received message without a valid secret. Rejecting. %s -> %s", event, data));
+            Event rejectedEvent =
+                new Event(REJECTED_EVENT, new ZeroMqEventBus.RejectedEvent(eventName, data));
+
+            notifyListeners(REJECTED_EVENT, rejectedEvent);
+
+            return;
+          }
+
+          notifyListeners(eventName, event);
+        }
+      }
+    } catch (Throwable e) {
+      // if the exception escapes, then we never get scheduled again
+      LOG.log(Level.WARNING, e, () -> "Caught and swallowed exception: " + e.getMessage());
+    }
+  }
+
+  private void notifyListeners(EventName eventName, Event event) {
+    List<Consumer<Event>> eventListeners = listeners.getOrDefault(eventName, new ArrayList<>());
+    eventListeners.stream().forEach(listener -> {
+      listenerNotificationExecutor.submit(() -> {
+        try {
+          listener.accept(event);
+
+        } catch (Throwable t) {
+          LOG.log(Level.WARNING, t, () -> "Caught exception from listener: " + listener);
+        }
+      });
+    });
+  }
 }

java/server/test/org/openqa/selenium/events/EventBusTest.java

@@ -31,6 +31,8 @@
 
 import java.util.Collection;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
@@ -103,4 +105,24 @@ public void shouldNotReceiveEventsNotMeantForTheTopic() {
 
     assertThat(count.get()).isEqualTo(0);
   }
+
+  @Test
+  public void shouldBeAbleToFireEventsInParallel() throws InterruptedException {
+    int maxCount = 100;
+    EventName name = new EventName("cheese");
+
+    CountDownLatch count = new CountDownLatch(maxCount);
+    bus.addListener(new EventListener<>(name, Object.class, obj -> count.countDown()));
+
+    ExecutorService executor = Executors.newCachedThreadPool();
+    try {
+      for (int i = 0; i < maxCount; i++) {
+        executor.submit(() -> bus.fire(new Event(name, "")));
+      }
+
+      assertThat(count.await(20, SECONDS)).describedAs(count.toString()).isTrue();
+    } finally {
+      executor.shutdownNow();
+    }
+  }
 }