code: // Licensed to the Software Freedom Conservancy (SFC) under one
  // or more contributor license agreements.  See the NOTICE file
  // distributed with this work for additional information
  // regarding copyright ownership.  The SFC licenses this file
  // to you under the Apache License, Version 2.0 (the
  // "License"); you may not use this file except in compliance
  // with the License.  You may obtain a copy of the License at
  //
  //   http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing,
  // software distributed under the License is distributed on an
  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, either express or implied.  See the License for the
  // specific language governing permissions and limitations
  // under the License.
  
  package org.openqa.selenium.events.zeromq;
  
  import static java.nio.charset.StandardCharsets.UTF_8;
  
  import com.google.common.collect.EvictingQueue;
  import dev.failsafe.Failsafe;
  import dev.failsafe.RetryPolicy;
  import java.net.Inet6Address;
  import java.net.InetAddress;
  import java.net.URI;
  import java.net.URISyntaxException;
  import java.net.UnknownHostException;
  import java.time.temporal.ChronoUnit;
  import java.util.ArrayList;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Queue;
  import java.util.UUID;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.function.Consumer;
  import java.util.logging.Level;
  import java.util.logging.Logger;
  import org.openqa.selenium.events.Event;
  import org.openqa.selenium.events.EventBus;
  import org.openqa.selenium.events.EventListener;
  import org.openqa.selenium.events.EventName;
  import org.openqa.selenium.grid.security.Secret;
  import org.openqa.selenium.internal.Require;
  import org.openqa.selenium.json.Json;
  import org.openqa.selenium.json.JsonException;
  import org.openqa.selenium.json.JsonOutput;
  import org.zeromq.SocketType;
  import org.zeromq.ZContext;
  import org.zeromq.ZMQ;
  
  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 AtomicBoolean pollingStarted = new AtomicBoolean(false);
    private final ExecutorService socketPollingExecutor;
    private final ExecutorService socketPublishingExecutor;
    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;
    private ZMQ.Poller poller;
  
    private ZMQ.Socket pub;
    private ZMQ.Socket sub;
  
    UnboundZmqEventBus(
        ZContext context, String publishConnection, String subscribeConnection, Secret secret) {
      Require.nonNull("Secret", secret);
      StringBuilder builder = new StringBuilder();
      try (JsonOutput out = JSON.newOutput(builder)) {
        out.setPrettyPrint(false).writeClassName(false).write(secret);
      }
      this.encodedSecret = builder.toString();
  
      this.socketPollingExecutor =
          Executors.newSingleThreadExecutor(
              r -> {
                Thread thread = new Thread(r);
                thread.setName("Event Bus Poller");
                thread.setDaemon(true);
                return thread;
              });
  
      this.socketPublishingExecutor =
          Executors.newSingleThreadExecutor(
              r -> {
                Thread thread = new Thread(r);
                thread.setName("Event Bus Publisher");
                thread.setDaemon(true);
                return thread;
              });
  
      this.listenerNotificationExecutor =
          Executors.newFixedThreadPool(
              Math.max(Runtime.getRuntime().availableProcessors() / 2, 2), // At least two threads
              r -> {
                Thread thread = new Thread(r);
                thread.setName("Event Bus Listener Notifier");
                thread.setDaemon(true);
                return thread;
              });
  
      String connectionMessage =
          String.format("Connecting to %s and %s", publishConnection, subscribeConnection);
      LOG.info(connectionMessage);
  
      RetryPolicy<Object> retryPolicy =
          RetryPolicy.builder()
              .withMaxAttempts(5)
              .withDelay(5, 10, ChronoUnit.SECONDS)
              .onFailedAttempt(
                  e -> LOG.log(Level.WARNING, String.format("%s failed", connectionMessage)))
              .onRetry(
                  e ->
                      LOG.log(
                          Level.WARNING,
                          String.format("Failure #%s. Retrying.", e.getAttemptCount())))
              .onRetriesExceeded(e -> LOG.log(Level.WARNING, "Connection aborted."))
              .build();
  
      // Access to the zmq socket is safe here: no threads.
      Failsafe.with(retryPolicy)
          .run(
              () -> {
                sub = context.createSocket(SocketType.SUB);
                sub.setIPv6(isSubAddressIPv6(publishConnection));
                sub.connect(publishConnection);
                sub.subscribe(new byte[0]);
  
                pub = context.createSocket(SocketType.PUB);
                pub.setIPv6(isSubAddressIPv6(subscribeConnection));
                pub.connect(subscribeConnection);
              });
      // Connections are already established
      this.poller = context.createPoller(1);
      this.poller.register(Objects.requireNonNull(sub), ZMQ.Poller.POLLIN);
  
      LOG.info("Sockets created");
  
      socketPollingExecutor.submit(new PollingRunnable(secret));
  
      // 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
      // be up yet.
      while (!pollingStarted.get()) {
        try {
          Thread.sleep(1000);
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          throw new RuntimeException(e);
        }
      }
  
      LOG.info("Event bus ready");
    }
  
    @Override
    public boolean isReady() {
      return !socketPollingExecutor.isShutdown();
    }
  
    private boolean isSubAddressIPv6(String connection) {
      try {
        URI uri = new URI(connection);
  
        if ("inproc".equals(uri.getScheme())) {
          return false;
        }
  
        return InetAddress.getByName(uri.getHost()) instanceof Inet6Address;
      } catch (UnknownHostException | URISyntaxException e) {
        LOG.log(
            Level.WARNING,
            String.format("Could not determine if the address %s is IPv6 or IPv4", connection),
            e);
      }
      return false;
    }
  
    @Override
    public void addListener(EventListener<?> listener) {
      Require.nonNull("Listener", listener);
  
      List<Consumer<Event>> typeListeners =
          listeners.computeIfAbsent(listener.getEventName(), t -> new LinkedList<>());
      typeListeners.add(listener);
    }
  
    @Override
    public void fire(Event event) {
      Require.nonNull("Event to send", event);
  
      socketPublishingExecutor.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() {
      socketPollingExecutor.shutdownNow();
      socketPublishingExecutor.shutdownNow();
      listenerNotificationExecutor.shutdownNow();
      poller.close();
  
      if (sub != null) {
        sub.close();
      }
      if (pub != null) {
        pub.close();
      }
    }
  
    private class PollingRunnable implements Runnable {
      private final Secret secret;
  
      public PollingRunnable(Secret secret) {
        this.secret = secret;
      }
  
      @Override
      public void run() {
        while (!Thread.currentThread().isInterrupted()) {
          try {
            int count = poller.poll(150);
  
            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(), UTF_8));
                // Processing only events we are listening to
                if (!listeners.containsKey(eventName)) {
                  continue;
                }
  
                Secret eventSecret;
                String receivedEventSecret = new String(socket.recv(), UTF_8);
                try {
                  eventSecret = JSON.toType(receivedEventSecret, Secret.class);
                } catch (JsonException ignore) {
                  rejectEvent(
                      eventName,
                      receivedEventSecret,
                      "Could not parse event secret, rejecting event.");
                  continue;
                }
  
                UUID id;
                String eventId = new String(socket.recv(), UTF_8);
                try {
                  id = UUID.fromString(eventId);
                } catch (IllegalArgumentException ignore) {
                  rejectEvent(eventName, eventId, "Could not parse event id, rejecting event.");
                  continue;
                }
  
                String data = new String(socket.recv(), UTF_8);
  
                // Don't bother doing more work if we've seen this message.
                if (recentMessages.contains(id)) {
                  continue;
                }
  
                Object converted = JSON.toType(data, Object.class);
                Event event = new Event(id, eventName, converted);
  
                recentMessages.add(id);
  
                if (!Secret.matches(secret, eventSecret)) {
                  rejectEvent(eventName, data, "Rejecting message without a valid secret");
                  continue;
                }
  
                notifyListeners(eventName, event);
              }
            }
          } catch (Exception e) {
            if (!(e.getCause() instanceof AssertionError)) {
              LOG.log(
                  Level.WARNING,
                  e,
                  () -> "Caught exception while polling for event bus messages: " + e.getMessage());
            }
          }
        }
      }
  
      private void rejectEvent(EventName eventName, String data, String message) {
        Event rejectedEvent =
            new Event(REJECTED_EVENT, new ZeroMqEventBus.RejectedEvent(eventName, data));
        LOG.log(Level.SEVERE, "{0}. {1}", new Object[] {message, rejectedEvent});
  
        notifyListeners(REJECTED_EVENT, rejectedEvent);
      }
  
      private void notifyListeners(EventName eventName, Event event) {
        List<Consumer<Event>> eventListeners = listeners.getOrDefault(eventName, new ArrayList<>());
        eventListeners.forEach(
            listener ->
                listenerNotificationExecutor.submit(
                    () -> {
                      try {
                        listener.accept(event);
                      } catch (Exception e) {
                        LOG.log(
                            Level.WARNING, e, () -> "Caught exception from listener: " + listener);
                      }
                    }));
      }
    }
  }
  