Refactor event loop to accept several clients

Used only in the recovery listeners, this allows sharing the event loop
between connections and saving some threads.

Author: acogoluegnes

src/main/java/com/rabbitmq/client/amqp/impl/AmqpConnection.java

@@ -281,7 +281,7 @@ TopologyListener createTopologyListener(AmqpConnectionBuilder builder) {
     if (builder.recoveryConfiguration().topology()) {
       RecordingTopologyListener rtl =
           new RecordingTopologyListener(
-              "topology-listener-connection-" + this.name(), this.executorService());
+              "topology-listener-connection-" + this.name(), this.environment.recoveryEventLoop());
       this.entityRecovery = new EntityRecovery(this, rtl);
       topologyListener = rtl;
     } else {

src/main/java/com/rabbitmq/client/amqp/impl/AmqpEnvironment.java

@@ -46,6 +46,7 @@ class AmqpEnvironment implements Environment {
   private final MetricsCollector metricsCollector;
   private final ObservationCollector observationCollector;
   private ConnectionUtils.AffinityCache affinityCache = new ConnectionUtils.AffinityCache();
+  private final EventLoop recoveryEventLoop;
 
   AmqpEnvironment(
       ExecutorService executorService,
@@ -71,6 +72,8 @@ class AmqpEnvironment implements Environment {
         metricsCollector == null ? NoOpMetricsCollector.INSTANCE : metricsCollector;
     this.observationCollector =
         observationCollector == null ? Utils.NO_OP_OBSERVATION_COLLECTOR : observationCollector;
+    // TODO start the event loop on first use
+    this.recoveryEventLoop = new EventLoop(this.executorService);
   }
 
   DefaultConnectionSettings<?> connectionSettings() {
@@ -95,6 +98,7 @@ public void close() {
     if (this.closed.compareAndSet(false, true)) {
       this.connectionManager.close();
       this.client.close();
+      this.recoveryEventLoop.close();
       if (this.internalExecutor) {
         this.executorService.shutdownNow();
       }
@@ -130,6 +134,10 @@ ConnectionUtils.AffinityCache affinityCache() {
     return this.affinityCache;
   }
 
+  EventLoop recoveryEventLoop() {
+    return this.recoveryEventLoop;
+  }
+
   AmqpConnection connection(AmqpConnectionBuilder builder) {
     return this.connectionManager.connection(builder);
   }

src/main/java/com/rabbitmq/client/amqp/impl/EventLoop.java

@@ -19,40 +19,46 @@
 
 import com.rabbitmq.client.amqp.AmqpException;
 import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
+import java.util.function.UnaryOperator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-final class EventLoop<S> implements AutoCloseable {
+final class EventLoop implements AutoCloseable {
 
   private static final Duration TIMEOUT = Duration.ofSeconds(60);
   private static final Logger LOGGER = LoggerFactory.getLogger(EventLoop.class);
 
   private final AtomicBoolean closed = new AtomicBoolean(false);
-  private final String label;
   private final Future<?> loop;
   private final AtomicReference<Thread> loopThread = new AtomicReference<>();
-  private final AtomicReference<S> stateReference = new AtomicReference<>();
-  private final BlockingQueue<Consumer<S>> taskQueue = new ArrayBlockingQueue<>(100);
+  private final BlockingQueue<ClientTaskContext<Object>> taskQueue = new ArrayBlockingQueue<>(100);
 
-  EventLoop(Supplier<S> stateSupplier, String label, ExecutorService executorService) {
-    this.label = label;
+  EventLoop(ExecutorService executorService) {
     CountDownLatch loopThreadSetLatch = new CountDownLatch(1);
     this.loop =
         executorService.submit(
             () -> {
-              S state = stateSupplier.get();
               loopThread.set(Thread.currentThread());
-              stateReference.set(state);
               loopThreadSetLatch.countDown();
+              Map<Long, Object> states = new HashMap<>();
               while (!Thread.currentThread().isInterrupted()) {
                 try {
-                  Consumer<S> task = this.taskQueue.take();
-                  task.accept(state);
+                  ClientTaskContext<Object> context = this.taskQueue.take();
+                  Object state = states.get(context.client.id);
+                  state = context.task.apply(state);
+                  if (state == null) {
+                    states.remove(context.client.id);
+                  } else {
+                    states.put(context.client.id, state);
+                  }
                 } catch (InterruptedException e) {
                   return;
                 } catch (Exception e) {
@@ -70,34 +76,49 @@ final class EventLoop<S> implements AutoCloseable {
     }
   }
 
-  void submit(Consumer<S> task) {
+  <S> Client<S> register(Supplier<S> stateSupplier) {
+    Client<S> client = new Client<>(this);
+    this.submit(
+        client,
+        nullState -> {
+          S state = stateSupplier.get();
+          client.stateReference.set(state);
+          return state;
+        });
+    return client;
+  }
+
+  private <ST> void submit(Client<ST> client, UnaryOperator<ST> task) {
     if (this.closed.get()) {
       throw new IllegalStateException("Event loop is closed");
     } else {
       if (Thread.currentThread().equals(this.loopThread.get())) {
-        task.accept(this.stateReference.get());
+        task.apply(client.stateReference.get());
       } else {
         CountDownLatch latch = new CountDownLatch(1);
         try {
-          boolean added =
-              this.taskQueue.offer(
+          ClientTaskContext<ST> context =
+              new ClientTaskContext<>(
+                  client,
                   state -> {
                     try {
-                      task.accept(state);
+                      return task.apply(state);
                     } catch (Exception e) {
-                      LOGGER.info("Error during {} task", this.label, e);
+                      LOGGER.info("Error during task", e);
                     } finally {
                       latch.countDown();
                     }
-                  },
-                  TIMEOUT.toMillis(),
-                  TimeUnit.MILLISECONDS);
+                    return null;
+                  });
+          boolean added =
+              this.taskQueue.offer(
+                  (ClientTaskContext<Object>) context, TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
           if (!added) {
-            throw new AmqpException("Enqueueing of %s task timed out", this.label);
+            throw new AmqpException("Enqueueing of task timed out");
           }
         } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
-          throw new AmqpException(this.label + " task enqueueing has been interrupted", e);
+          throw new AmqpException("Task enqueueing has been interrupted", e);
         }
         try {
           boolean completed = latch.await(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
@@ -106,20 +127,66 @@ void submit(Consumer<S> task) {
           }
         } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
-          throw new AmqpException(this.label + " Topology task processing has been interrupted", e);
+          throw new AmqpException("Topology task processing has been interrupted", e);
         }
       }
     }
   }
 
-  S state() {
-    return this.stateReference.get();
-  }
-
   @Override
   public void close() {
     if (this.closed.compareAndSet(false, true)) {
       this.loop.cancel(true);
     }
   }
+
+  private static final AtomicLong CLIENT_ID_SEQUENCE = new AtomicLong();
+
+  static class Client<S> implements AutoCloseable {
+
+    private final long id;
+    private final AtomicReference<S> stateReference = new AtomicReference<>();
+    private final EventLoop loop;
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
+    private Client(EventLoop loop) {
+      this.id = CLIENT_ID_SEQUENCE.getAndIncrement();
+      this.loop = loop;
+    }
+
+    void submit(Consumer<S> task) {
+      if (this.closed.get()) {
+        throw new IllegalStateException("Event loop is closed");
+      } else {
+        this.loop.submit(
+            this,
+            s -> {
+              task.accept(s);
+              return s;
+            });
+      }
+    }
+
+    @Override
+    public void close() {
+      if (this.closed.compareAndSet(false, true)) {
+        this.loop.submit(this, s -> null);
+      }
+    }
+
+    S state() {
+      return this.stateReference.get();
+    }
+  }
+
+  private static class ClientTaskContext<S> {
+
+    private final Client<S> client;
+    private final UnaryOperator<S> task;
+
+    private ClientTaskContext(Client<S> client, UnaryOperator<S> task) {
+      this.client = client;
+      this.task = task;
+    }
+  }
 }

src/main/java/com/rabbitmq/client/amqp/impl/RecordingTopologyListener.java

@@ -18,7 +18,6 @@
 package com.rabbitmq.client.amqp.impl;
 
 import java.util.*;
-import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
@@ -30,13 +29,13 @@ final class RecordingTopologyListener implements TopologyListener, AutoCloseable
   private static final Logger LOGGER = LoggerFactory.getLogger(RecordingTopologyListener.class);
 
   private final String label;
-  private final EventLoop<State> eventLoop;
+  private final EventLoop.Client<State> eventLoopClient;
 
   private final AtomicBoolean closed = new AtomicBoolean(false);
 
-  RecordingTopologyListener(String label, ExecutorService executorService) {
+  RecordingTopologyListener(String label, EventLoop eventLoop) {
     this.label = label;
-    this.eventLoop = new EventLoop<>(State::new, label, executorService);
+    this.eventLoopClient = eventLoop.register(State::new);
   }
 
   private static class State {
@@ -117,13 +116,13 @@ public void consumerDeleted(long id, String queue) {
   @Override
   public void close() {
     if (this.closed.compareAndSet(false, true)) {
-      this.eventLoop.close();
+      this.eventLoopClient.close();
     }
   }
 
   private void submit(Consumer<State> task) {
     if (!this.closed.get()) {
-      this.eventLoop.submit(task);
+      this.eventLoopClient.submit(task);
     }
   }
 
@@ -342,7 +341,7 @@ interface Visitor {
   // for test assertions
 
   private State state() {
-    return this.eventLoop.state();
+    return this.eventLoopClient.state();
   }
 
   Map<String, ExchangeSpec> exchanges() {

src/test/java/com/rabbitmq/client/amqp/impl/EventLoopTest.java

@@ -28,7 +28,8 @@
 public class EventLoopTest {
 
   static ExecutorService executorService;
-  EventLoop<State> loop;
+  EventLoop loop;
+  EventLoop.Client<State> client;
 
   @BeforeAll
   static void beforeAll() {
@@ -37,7 +38,8 @@ static void beforeAll() {
 
   @BeforeEach
   void beforeEach() {
-    loop = new EventLoop<>(State::new, "test", executorService);
+    loop = new EventLoop(executorService);
+    client = loop.register(State::new);
   }
 
   @AfterEach
@@ -53,22 +55,27 @@ static void afterAll() {
   @Test
   void submitTasks() {
     AtomicInteger buffer = new AtomicInteger();
-    loop.submit(s -> s.a = 42);
-    loop.submit(s -> buffer.set(s.a));
+    client.submit(s -> s.a = 42);
+    client.submit(
+        s -> {
+          buffer.set(s.a);
+        });
     assertThat(buffer).hasValue(42);
 
-    loop.submit(
+    client.submit(
         s -> {
           s.a = 1;
           s.b = 2;
         });
-    loop.submit(s -> buffer.set(s.a));
+    client.submit(s -> buffer.set(s.a));
     assertThat(buffer).hasValue(1);
-    loop.submit(s -> buffer.set(s.b));
+    client.submit(s -> buffer.set(s.b));
     assertThat(buffer).hasValue(2);
 
+    client.close();
+    assertThatThrownBy(() -> client.submit(s -> {})).isInstanceOf(IllegalStateException.class);
     loop.close();
-    assertThatThrownBy(() -> loop.submit(s -> {})).isInstanceOf(IllegalStateException.class);
+    assertThatThrownBy(() -> loop.register(State::new));
   }
 
   static class State {

src/test/java/com/rabbitmq/client/amqp/impl/RecordingTopologyListenerTest.java

@@ -36,7 +36,7 @@ public class RecordingTopologyListenerTest {
   @BeforeEach
   void init() {
     executorService = Executors.newSingleThreadExecutor();
-    recovery = new RecordingTopologyListener("", executorService);
+    recovery = new RecordingTopologyListener("", new EventLoop(executorService));
   }
 
   @AfterEach