Track RPC client and server instances in connection

Author: acogoluegnes

src/main/java/com/rabbitmq/model/amqp/AmqpConnection.java

@@ -67,6 +67,8 @@ class AmqpConnection extends ResourceBase implements Connection {
   private volatile Session nativeSession;
   private final List<AmqpPublisher> publishers = new CopyOnWriteArrayList<>();
   private final List<AmqpConsumer> consumers = new CopyOnWriteArrayList<>();
+  private final List<RpcClient> rpcClients = new CopyOnWriteArrayList<>();
+  private final List<RpcServer> rpcServers = new CopyOnWriteArrayList<>();
   private final TopologyListener topologyListener;
   private volatile EntityRecovery entityRecovery;
   private final Future<?> recoveryLoop;
@@ -177,13 +179,18 @@ public void close() {
         }
       }
       this.closeManagement();
+      for (RpcClient rpcClient : this.rpcClients) {
+        rpcClient.close();
+      }
+      for (RpcServer rpcServer : this.rpcServers) {
+        rpcServer.close();
+      }
       for (AmqpPublisher publisher : this.publishers) {
         publisher.close();
       }
       for (AmqpConsumer consumer : this.consumers) {
         consumer.close();
       }
-
       try {
         this.nativeConnection.close();
       } catch (Exception e) {
@@ -517,6 +524,26 @@ void removeConsumer(AmqpConsumer consumer) {
     this.topologyListener.consumerDeleted(consumer.id(), consumer.address());
   }
 
+  RpcClient createRpcClient(RpcSupport.AmqpRpcClientBuilder builder) {
+    RpcClient rpcClient = new AmqpRpcClient(builder);
+    this.rpcClients.add(rpcClient);
+    return rpcClient;
+  }
+
+  void removeRpcClient(RpcClient rpcClient) {
+    this.rpcClients.remove(rpcClient);
+  }
+
+  RpcServer createRpcServer(RpcSupport.AmqpRpcServerBuilder builder) {
+    RpcServer rpcServer = new AmqpRpcServer(builder);
+    this.rpcServers.add(rpcServer);
+    return rpcServer;
+  }
+
+  void removeRpcServer(RpcServer rpcServer) {
+    this.rpcServers.remove(rpcServer);
+  }
+
   private void changeStateOfPublishers(State newState, Throwable failure) {
     this.changeStateOfResources(this.publishers, newState, failure);
   }

src/main/java/com/rabbitmq/model/amqp/AmqpRpcClient.java

@@ -22,34 +22,42 @@
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 class AmqpRpcClient implements RpcClient {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(AmqpRpcClient.class);
+
   private static final Publisher.Callback NO_OP_CALLBACK = ctx -> {};
 
+  private final AmqpConnection connection;
   private final Publisher publisher;
   private final Consumer consumer;
   private final Map<Object, CompletableFuture<Message>> outstandingRequests =
       new ConcurrentHashMap<>();
   private final Supplier<Object> correlationIdSupplier;
   private final BiFunction<Message, Object, Message> requestPostProcessor;
   private final Function<Message, Object> correlationIdExtractor;
+  private final AtomicBoolean closed = new AtomicBoolean(false);
 
   AmqpRpcClient(RpcSupport.AmqpRpcClientBuilder builder) {
-    AmqpConnection connection = builder.connection();
+    this.connection = builder.connection();
 
-    AmqpPublisherBuilder publisherBuilder = (AmqpPublisherBuilder) connection.publisherBuilder();
+    AmqpPublisherBuilder publisherBuilder =
+        (AmqpPublisherBuilder) this.connection.publisherBuilder();
     ((DefaultAddressBuilder<?>) builder.requestAddress()).copyTo(publisherBuilder.addressBuilder());
     this.publisher = publisherBuilder.build();
 
     String replyTo = builder.replyToQueue();
     if (replyTo == null) {
       Management.QueueInfo queueInfo =
-          connection.management().queue().exclusive(true).autoDelete(true).declare();
+          this.connection.management().queue().exclusive(true).autoDelete(true).declare();
       replyTo = queueInfo.name();
     }
     if (builder.correlationIdExtractor() == null) {
@@ -58,7 +66,7 @@ class AmqpRpcClient implements RpcClient {
       this.correlationIdExtractor = builder.correlationIdExtractor();
     }
     this.consumer =
-        connection
+        this.connection
             .consumerBuilder()
             .queue(replyTo)
             .messageHandler(
@@ -114,7 +122,18 @@ public CompletableFuture<Message> publish(Message message) {
 
   @Override
   public void close() {
-    this.publisher.close();
-    this.consumer.close();
+    if (this.closed.compareAndSet(false, true)) {
+      this.connection.removeRpcClient(this);
+      try {
+        this.publisher.close();
+      } catch (Exception e) {
+        LOGGER.warn("Error while closing RPC client publisher: {}", e.getMessage());
+      }
+      try {
+        this.consumer.close();
+      } catch (Exception e) {
+        LOGGER.warn("Error while closing RPC client consumer: {}", e.getMessage());
+      }
+    }
   }
 }

src/main/java/com/rabbitmq/model/amqp/AmqpRpcServer.java

@@ -21,23 +21,30 @@
 import com.rabbitmq.model.Message;
 import com.rabbitmq.model.Publisher;
 import com.rabbitmq.model.RpcServer;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiFunction;
 import java.util.function.Function;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 class AmqpRpcServer implements RpcServer {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(AmqpRpcServer.class);
+
   private static final Publisher.Callback NO_OP_CALLBACK = ctx -> {};
 
+  private final AmqpConnection connection;
   private final Publisher publisher;
   private final Consumer consumer;
   private final Function<Message, Object> correlationIdExtractor;
   private final BiFunction<Message, Object, Message> replyPostProcessor;
+  private final AtomicBoolean closed = new AtomicBoolean(false);
 
   AmqpRpcServer(RpcSupport.AmqpRpcServerBuilder builder) {
-    AmqpConnection connection = builder.connection();
+    this.connection = builder.connection();
     Handler handler = builder.handler();
 
-    this.publisher = connection.publisherBuilder().build();
+    this.publisher = this.connection.publisherBuilder().build();
 
     Context context =
         new Context() {
@@ -62,7 +69,7 @@ public Message message(byte[] body) {
       this.replyPostProcessor = builder.replyPostProcessor();
     }
     this.consumer =
-        connection
+        this.connection
             .consumerBuilder()
             .queue(builder.requestQueue())
             .messageHandler(
@@ -81,7 +88,18 @@ public Message message(byte[] body) {
 
   @Override
   public void close() {
-    this.consumer.close();
-    this.publisher.close();
+    if (this.closed.compareAndSet(false, true)) {
+      this.connection.removeRpcServer(this);
+      try {
+        this.consumer.close();
+      } catch (Exception e) {
+        LOGGER.warn("Error while closing RPC server consumer: {}", e.getMessage());
+      }
+      try {
+        this.publisher.close();
+      } catch (Exception e) {
+        LOGGER.warn("Error while closing RPC server publisher: {}", e.getMessage());
+      }
+    }
   }
 }

src/main/java/com/rabbitmq/model/amqp/RpcSupport.java

@@ -78,7 +78,7 @@ Function<Message, Object> correlationIdExtractor() {
 
     @Override
     public RpcClient build() {
-      return new AmqpRpcClient(this);
+      return this.connection.createRpcClient(this);
     }
 
     AmqpConnection connection() {
@@ -168,7 +168,7 @@ public RpcServerBuilder replyPostProcessor(
 
     @Override
     public RpcServer build() {
-      return new AmqpRpcServer(this);
+      return this.connection.createRpcServer(this);
     }
 
     AmqpConnection connection() {