Add parameter to create a session per publisher/consumer

The option is not public for now. Name is temporary.

Author: acogoluegnes

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

@@ -33,6 +33,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 import org.apache.qpid.protonj2.client.ConnectionOptions;
 import org.apache.qpid.protonj2.client.DisconnectionEvent;
 import org.apache.qpid.protonj2.client.Session;
@@ -68,14 +69,18 @@ final class AmqpConnection extends ResourceBase implements Connection {
   private final AtomicBoolean recoveringConnection = new AtomicBoolean(false);
   private final DefaultConnectionSettings<?> connectionSettings =
       DefaultConnectionSettings.instance();
+  private Supplier<SessionHandler> sessionHandlerSupplier;
 
   AmqpConnection(AmqpConnectionBuilder builder) {
     super(builder.listeners());
     this.id = ID_SEQUENCE.getAndIncrement();
     this.environment = builder.environment();
     builder.connectionSettings().copyTo(this.connectionSettings);
     this.connectionSettings.consolidate();
-
+    this.sessionHandlerSupplier =
+        builder.isolateResources()
+            ? () -> new SessionHandler.SingleSessionSessionHandler(this)
+            : () -> new SessionHandler.ConnectionNativeSessionSessionHandler(this);
     BiConsumer<org.apache.qpid.protonj2.client.Connection, DisconnectionEvent> disconnectHandler;
     AmqpConnectionBuilder.AmqpRecoveryConfiguration recoveryConfiguration =
         builder.recoveryConfiguration();
@@ -491,6 +496,10 @@ ObservationCollector observationCollector() {
     return this.environment.observationCollector();
   }
 
+  SessionHandler createSessionHandler() {
+    return this.sessionHandlerSupplier.get();
+  }
+
   Publisher createPublisher(AmqpPublisherBuilder builder) {
     // TODO copy the builder properties to create the publisher
     AmqpPublisher publisher = new AmqpPublisher(builder);

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

@@ -33,6 +33,7 @@ class AmqpConnectionBuilder implements ConnectionBuilder {
   private final List<Resource.StateListener> listeners = new ArrayList<>();
   private String name;
   private TopologyListener topologyListener;
+  private boolean isolateResources = false;
 
   AmqpConnectionBuilder(AmqpEnvironment environment) {
     this.environment = environment;
@@ -115,6 +116,15 @@ public RecoveryConfiguration recovery() {
     return this.recoveryConfiguration;
   }
 
+  AmqpConnectionBuilder isolateResources(boolean isolateResources) {
+    this.isolateResources = isolateResources;
+    return this;
+  }
+
+  boolean isolateResources() {
+    return isolateResources;
+  }
+
   @Override
   public Connection build() {
     // TODO copy the recovery configuration to keep the settings

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

@@ -74,7 +74,7 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
             .subscribe(builder.queue(), builder.messageHandler());
     this.address = "/queue/" + builder.queue();
     this.connection = builder.connection();
-    this.sessionHandler = new SessionHandler.ConnectionNativeSessionSessionHandler(this.connection);
+    this.sessionHandler = this.connection.createSessionHandler();
     this.nativeReceiver = this.createNativeReceiver(this.sessionHandler.session(), this.address);
     this.initStateFromNativeReceiver(this.nativeReceiver);
     this.metricsCollector = this.connection.metricsCollector();

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

@@ -62,7 +62,7 @@ final class AmqpPublisher extends ResourceBase implements Publisher {
     this.destinationSpec = builder.destination();
     this.connection = builder.connection();
     this.publishTimeout = builder.publishTimeout();
-    this.sessionHandler = new SessionHandler.ConnectionNativeSessionSessionHandler(this.connection);
+    this.sessionHandler = this.connection.createSessionHandler();
     this.sender = this.createSender(sessionHandler.session(), this.address, this.publishTimeout);
     this.metricsCollector = this.connection.metricsCollector();
     this.observationCollector = this.connection.observationCollector();

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

@@ -17,7 +17,13 @@
 // [email protected].
 package com.rabbitmq.client.amqp.impl;
 
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+import org.apache.qpid.protonj2.client.Connection;
 import org.apache.qpid.protonj2.client.Session;
+import org.apache.qpid.protonj2.client.exceptions.ClientException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 interface SessionHandler extends AutoCloseable {
 
@@ -48,4 +54,50 @@ public Session sessionNoCheck() {
     @Override
     public void close() {}
   }
+
+  class SingleSessionSessionHandler implements SessionHandler {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(SingleSessionSessionHandler.class);
+
+    private final Supplier<Connection> connection;
+    private final AtomicReference<Session> session = new AtomicReference<>();
+
+    public SingleSessionSessionHandler(AmqpConnection connection) {
+      this.connection = connection::nativeConnection;
+    }
+
+    @Override
+    public Session session() {
+      closeCurrentSession();
+      try {
+        Session session = this.connection.get().openSession();
+        this.session.set(ExceptionUtils.wrapGet(session.openFuture()));
+        return this.session.get();
+      } catch (ClientException e) {
+        this.session.set(null);
+        throw ExceptionUtils.convert(e);
+      }
+    }
+
+    @Override
+    public Session sessionNoCheck() {
+      return this.session();
+    }
+
+    @Override
+    public void close() {
+      closeCurrentSession();
+    }
+
+    private void closeCurrentSession() {
+      Session previousSession = session.getAndSet(null);
+      if (previousSession != null) {
+        try {
+          previousSession.close();
+        } catch (RuntimeException e) {
+          LOGGER.debug("Error while closing session: {}", e.getMessage());
+        }
+      }
+    }
+  }
 }

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

@@ -39,6 +39,8 @@
 import java.util.concurrent.atomic.AtomicReference;
 import org.assertj.core.api.ThrowableAssert.ThrowingCallable;
 import org.junit.jupiter.api.*;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 @DisabledIfRabbitMqCtlNotSet
 public class AmqpConnectionRecoveryTest {
@@ -73,8 +75,10 @@ static void afterAll() {
     environment.close();
   }
 
-  @Test
-  void connectionShouldRecoverAfterClosingIt(TestInfo info) throws Exception {
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void connectionShouldRecoverAfterClosingIt(boolean isolateResources, TestInfo info)
+      throws Exception {
     String q = name(info);
     String connectionName = UUID.randomUUID().toString();
     Map<Resource.State, CountDownLatch> stateLatches = new ConcurrentHashMap<>();
@@ -84,6 +88,7 @@ void connectionShouldRecoverAfterClosingIt(TestInfo info) throws Exception {
         (AmqpConnectionBuilder)
             new AmqpConnectionBuilder(environment)
                 .name(connectionName)
+                .isolateResources(isolateResources)
                 .listeners(
                     context -> {
                       if (stateLatches.containsKey(context.currentState())) {
@@ -170,8 +175,9 @@ void connectionShouldRecoverAfterClosingIt(TestInfo info) throws Exception {
     }
   }
 
-  @Test
-  void connectionShouldRecoverAfterBrokerStopStart(TestInfo info) {
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void connectionShouldRecoverAfterBrokerStopStart(boolean isolateResources, TestInfo info) {
     String q = name(info);
     String connectionName = UUID.randomUUID().toString();
     Map<Resource.State, CountDownLatch> stateLatches = new ConcurrentHashMap<>();
@@ -181,6 +187,7 @@ void connectionShouldRecoverAfterBrokerStopStart(TestInfo info) {
         (AmqpConnectionBuilder)
             new AmqpConnectionBuilder(environment)
                 .name(connectionName)
+                .isolateResources(isolateResources)
                 .listeners(
                     context -> {
                       if (stateLatches.containsKey(context.currentState())) {

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

@@ -33,6 +33,8 @@
 import java.util.stream.IntStream;
 import org.junit.jupiter.api.*;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 @ExtendWith(AmqpTestInfrastructureExtension.class)
 public class RpcTest {
@@ -201,8 +203,9 @@ void rpcUseCorrelationIdRequestProperty() {
     }
   }
 
-  @Test
-  void rpcShouldRecoverAfterConnectionIsClosed()
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void rpcShouldRecoverAfterConnectionIsClosed(boolean isolateResources)
       throws ExecutionException, InterruptedException, TimeoutException {
     String clientConnectionName = UUID.randomUUID().toString();
     CountDownLatch clientConnectionLatch = new CountDownLatch(1);
@@ -213,6 +216,7 @@ void rpcShouldRecoverAfterConnectionIsClosed()
     Connection serverConnection =
         connectionBuilder()
             .name(serverConnectionName)
+            .isolateResources(isolateResources)
             .listeners(recoveredListener(serverConnectionLatch))
             .recovery()
             .backOffDelayPolicy(backOffDelayPolicy)
@@ -222,6 +226,7 @@ void rpcShouldRecoverAfterConnectionIsClosed()
     try (Connection clientConnection =
         connectionBuilder()
             .name(clientConnectionName)
+            .isolateResources(isolateResources)
             .listeners(recoveredListener(clientConnectionLatch))
             .recovery()
             .backOffDelayPolicy(backOffDelayPolicy)

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

@@ -93,24 +93,26 @@ public static Duration waitAtMost(
   public static Duration waitAtMost(
       Duration timeout, CallableBooleanSupplier condition, Supplier<String> message)
       throws Exception {
+    long start = System.nanoTime();
     if (condition.getAsBoolean()) {
       return Duration.ZERO;
     }
-    int waitTime = 100;
-    int waitedTime = 0;
-    int timeoutInMs = (int) timeout.toMillis();
+    Duration waitTime = Duration.ofMillis(100);
+    Duration waitedTime = Duration.ofNanos(System.nanoTime() - start);
     Exception exception = null;
-    while (waitedTime <= timeoutInMs) {
-      Thread.sleep(waitTime);
-      waitedTime += waitTime;
+    while (waitedTime.compareTo(timeout) <= 0) {
+      Thread.sleep(waitTime.toMillis());
+      waitedTime = waitedTime.plus(waitTime);
+      start = System.nanoTime();
       try {
         if (condition.getAsBoolean()) {
-          return Duration.ofMillis(waitedTime);
+          return waitedTime;
         }
         exception = null;
       } catch (Exception e) {
         exception = e;
       }
+      waitedTime = waitedTime.plus(Duration.ofNanos(System.nanoTime() - start));
     }
     String msg;
     if (message == null) {
@@ -123,7 +125,7 @@ public static Duration waitAtMost(
     } else {
       fail(msg, exception);
     }
-    return Duration.ofMillis(waitedTime);
+    return waitedTime;
   }
 
   public static class CountDownLatchReferenceConditions {