Simplify the ClientHolder, context and shared client mechanism

Author: ozangunalp

smallrye-reactive-messaging-rabbitmq/src/main/java/io/smallrye/reactive/messaging/rabbitmq/ClientHolder.java

@@ -3,7 +3,10 @@
 import static io.smallrye.reactive.messaging.rabbitmq.i18n.RabbitMQExceptions.ex;
 import static io.smallrye.reactive.messaging.rabbitmq.i18n.RabbitMQLogging.log;
 
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
@@ -13,51 +16,34 @@
 import io.smallrye.mutiny.Uni;
 import io.smallrye.reactive.messaging.providers.helpers.VertxContext;
 import io.vertx.mutiny.core.Context;
-import io.vertx.mutiny.core.Vertx;
 import io.vertx.mutiny.rabbitmq.RabbitMQClient;
 
 public class ClientHolder {
 
     private final RabbitMQClient client;
 
-    private final AtomicBoolean connected = new AtomicBoolean(false);
-    private final AtomicReference<CurrentConnection> connectionHolder = new AtomicReference<>();
-    private final AtomicReference<Context> rootContext;
-    private final AtomicReference<CompletionStage<RabbitMQClient>> connectionStage = new AtomicReference<>();
+    private final AtomicBoolean hasBeenConnected = new AtomicBoolean(false);
+    private final AtomicReference<CompletableFuture<RabbitMQClient>> ongoingConnection = new AtomicReference<>();
     private final Uni<RabbitMQClient> connection;
+    private final Set<String> channels = ConcurrentHashMap.newKeySet();
 
-    private final Vertx vertx;
-
-    public ClientHolder(RabbitMQClient client,
-            String channel,
-            Vertx vertx,
-            Context root) {
+    public ClientHolder(RabbitMQClient client) {
         this.client = client;
-        this.vertx = vertx;
-        this.rootContext = new AtomicReference<>(root);
         this.connection = Uni.createFrom().deferred(() -> client.start()
                 .onSubscription().invoke(() -> {
-                    connected.set(true);
-                    log.connectionEstablished(channel);
+                    hasBeenConnected.set(true);
+                    log.connectionEstablished(String.join(", ", channels));
                 })
                 .onItem().transform(ignored -> {
-                    Context ctx = rootContext.get();
-                    connectionHolder
-                            .set(new CurrentConnection(client, ctx == null ? Vertx.currentContext() : ctx));
-
                     // handle the case we are already disconnected.
-                    if (!client.isConnected() || connectionHolder.get() == null) {
+                    if (!client.isConnected()) {
                         // Throwing the exception would trigger a retry.
-                        connectionHolder.set(null);
                         throw ex.illegalStateConnectionDisconnected();
                     }
                     return client;
                 })
-                .onFailure().invoke(log::unableToConnectToBroker)
-                .onFailure().invoke(t -> {
-                    connectionHolder.set(null);
-                    log.unableToRecoverFromConnectionDisruption(t);
-                }));
+                .onFailure().invoke(log::unableToConnectToBroker))
+                .memoize().until(() -> !client.isConnected());
     }
 
     public static CompletionStage<Void> runOnContext(Context context, IncomingRabbitMQMessage<?> msg,
@@ -76,32 +62,12 @@ public static CompletionStage<Void> runOnContextAndReportFailure(Context context
         });
     }
 
-    public Context getContext() {
-        CurrentConnection connection = connectionHolder.get();
-        if (connection != null) {
-            return connection.context;
-        } else {
-            return null;
-        }
-    }
-
-    public void ensureContext(Context context) {
-        if (context == null) {
-            return;
-        }
-        rootContext.compareAndSet(null, context);
-        CurrentConnection connection = connectionHolder.get();
-        if (connection != null && connection.context == null) {
-            connectionHolder.compareAndSet(connection, new CurrentConnection(connection.client, context));
-        }
-    }
-
     public RabbitMQClient client() {
         return client;
     }
 
     public boolean hasBeenConnected() {
-        return connected.get();
+        return hasBeenConnected.get();
     }
 
     @CheckReturnValue
@@ -113,46 +79,46 @@ public Function<Throwable, Uni<Void>> getNack(final long deliveryTag, final bool
         return t -> client.basicNack(deliveryTag, false, requeue);
     }
 
-    public Vertx getVertx() {
-        return vertx;
-    }
-
     @CheckReturnValue
     public Uni<RabbitMQClient> getOrEstablishConnection() {
-        CompletionStage<RabbitMQClient> existing = connectionStage.get();
+        return Uni.createFrom().deferred(this::establishConnection);
+    }
+
+    private Uni<RabbitMQClient> establishConnection() {
+        CompletableFuture<RabbitMQClient> existing = ongoingConnection.get();
         if (existing != null) {
-            if (!existing.toCompletableFuture().isDone() || client.isConnected()) {
+            if (!existing.isDone() || client.isConnected()) {
                 return Uni.createFrom().completionStage(existing);
             }
-            connectionStage.compareAndSet(existing, null);
+            ongoingConnection.compareAndSet(existing, null);
         }
 
-        for (;;) {
-            CompletionStage<RabbitMQClient> current = connectionStage.get();
-            if (current != null) {
-                return Uni.createFrom().completionStage(current);
-            }
-            CompletionStage<RabbitMQClient> created = connection.subscribeAsCompletionStage();
-            if (connectionStage.compareAndSet(null, created)) {
-                created.whenComplete((result, error) -> {
-                    if (error != null) {
-                        connectionStage.compareAndSet(created, null);
-                    }
-                });
-                return Uni.createFrom().completionStage(created);
-            }
+        CompletableFuture<RabbitMQClient> placeholder = new CompletableFuture<>();
+        CompletableFuture<RabbitMQClient> current = ongoingConnection.compareAndExchange(null, placeholder);
+        if (current != null) {
+            return Uni.createFrom().completionStage(current);
         }
+        connection.subscribe().with(placeholder::complete, placeholder::completeExceptionally);
+        placeholder.whenComplete((result, error) -> {
+            if (error != null) {
+                ongoingConnection.compareAndSet(placeholder, null);
+            }
+        });
+        return Uni.createFrom().completionStage(placeholder);
     }
 
-    private static class CurrentConnection {
+    public Set<String> channels() {
+        return channels;
+    }
 
-        final RabbitMQClient client;
-        final Context context;
+    public ClientHolder retain(String channel) {
+        channels.add(channel);
+        return this;
+    }
 
-        private CurrentConnection(RabbitMQClient client, Context context) {
-            this.client = client;
-            this.context = context;
-        }
+    public boolean release(String channel) {
+        channels.remove(channel);
+        return channels.isEmpty();
     }
 
 }

smallrye-reactive-messaging-rabbitmq/src/main/java/io/smallrye/reactive/messaging/rabbitmq/IncomingRabbitMQMessage.java

@@ -60,12 +60,6 @@ public <V> CompletionStage<Void> handle(IncomingRabbitMQMessage<V> message, Meta
     private final String contentTypeOverride;
     private final T payload;
 
-    public IncomingRabbitMQMessage(RabbitMQMessage delegate, ClientHolder holder,
-            RabbitMQFailureHandler onNack,
-            RabbitMQAckHandler onAck, String contentTypeOverride) {
-        this(delegate.getDelegate(), holder, holder.getContext(), onNack, onAck, contentTypeOverride);
-    }
-
     public IncomingRabbitMQMessage(RabbitMQMessage delegate, ClientHolder holder, Context context,
             RabbitMQFailureHandler onNack,
             RabbitMQAckHandler onAck, String contentTypeOverride) {
@@ -77,7 +71,7 @@ public IncomingRabbitMQMessage(RabbitMQMessage delegate, ClientHolder holder, Co
         this.message = msg;
         this.deliveryTag = msg.envelope().getDeliveryTag();
         this.holder = holder;
-        this.context = context != null ? context : holder.getContext();
+        this.context = context;
         this.contentTypeOverride = contentTypeOverride;
         this.rabbitMQMetadata = new IncomingRabbitMQMetadata(this.message, contentTypeOverride);
         this.onNack = onNack;

smallrye-reactive-messaging-rabbitmq/src/main/java/io/smallrye/reactive/messaging/rabbitmq/RabbitMQConnector.java

@@ -6,14 +6,12 @@
 import static io.smallrye.reactive.messaging.rabbitmq.i18n.RabbitMQExceptions.ex;
 import static io.smallrye.reactive.messaging.rabbitmq.i18n.RabbitMQLogging.log;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.Flow;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import jakarta.annotation.Priority;
 import jakarta.enterprise.context.ApplicationScoped;
@@ -169,10 +167,11 @@ public class RabbitMQConnector implements InboundConnector, OutboundConnector, H
     @Inject
     @Any
     Instance<RabbitMQFailureHandler.Factory> failureHandlerFactories;
-    private List<IncomingRabbitMQChannel> incomings = new CopyOnWriteArrayList<>();
-    private List<OutgoingRabbitMQChannel> outgoings = new CopyOnWriteArrayList<>();
-    private Map<String, ClientRegistration> clientRegistrations = new ConcurrentHashMap<>();
-    private Map<String, SharedClient> sharedClients = new ConcurrentHashMap<>();
+    private final List<IncomingRabbitMQChannel> incomings = new CopyOnWriteArrayList<>();
+    private final List<OutgoingRabbitMQChannel> outgoings = new CopyOnWriteArrayList<>();
+    private final Map<String, ClientHolder> clients = new ConcurrentHashMap<>();
+    // connection-name to fingerprint map to check against same connection-name but different options
+    private final Map<String, String> connectionFingerprints = new ConcurrentHashMap<>();
 
     @Inject
     @Any
@@ -271,11 +270,11 @@ public void terminate(
             outgoing.terminate();
         }
 
-        List<String> registeredChannels = new ArrayList<>(clientRegistrations.keySet());
-        for (String channel : registeredChannels) {
-            releaseClient(channel);
+        for (Map.Entry<String, ClientHolder> entry : clients.entrySet()) {
+            stopClient(entry.getValue().client(), true);
         }
-        sharedClients.clear();
+        clients.clear();
+        connectionFingerprints.clear();
     }
 
     public Vertx vertx() {
@@ -284,16 +283,7 @@ public Vertx vertx() {
 
     public void reportIncomingFailure(String channel, Throwable reason) {
         log.failureReported(channel, reason);
-        ClientRegistration registration = clientRegistrations.remove(channel);
-        if (registration == null) {
-            return;
-        }
-
-        if (registration.shared) {
-            releaseSharedClient(registration.key, false);
-        } else {
-            stopClient(registration.holder.client(), false);
-        }
+        releaseClient(channel, false);
     }
 
     public Instance<RabbitMQFailureHandler.Factory> failureHandlerFactories() {
@@ -316,71 +306,30 @@ public Instance<CredentialsProvider> credentialsProviders() {
         return configMaps;
     }
 
-    public ClientHolder getClientHolder(RabbitMQConnectorCommonConfiguration config, io.vertx.mutiny.core.Context context) {
-        ClientRegistration existing = clientRegistrations.get(config.getChannel());
-        if (existing != null) {
-            return existing.holder;
-        }
-
-        return config.getSharedConnectionName()
-                .map(name -> getOrCreateSharedHolder(config, context, name))
-                .orElseGet(() -> createAndRegisterHolder(config, context, config.getChannel(), false));
-    }
-
-    private ClientHolder createAndRegisterHolder(RabbitMQConnectorCommonConfiguration config,
-            io.vertx.mutiny.core.Context context, String key, boolean shared) {
-        ClientHolder holder = new ClientHolder(RabbitMQClientHelper.createClient(this, config), config.getChannel(), vertx(),
-                context);
-        clientRegistrations.put(config.getChannel(), new ClientRegistration(holder, shared, key));
-        return holder;
-    }
-
-    private ClientHolder getOrCreateSharedHolder(RabbitMQConnectorCommonConfiguration config,
-            io.vertx.mutiny.core.Context context, String name) {
+    public ClientHolder getClientHolder(RabbitMQConnectorCommonConfiguration config) {
+        String channel = config.getChannel();
         RabbitMQOptions options = RabbitMQClientHelper.buildClientOptions(this, config);
+        String connectionName = options.getConnectionName();
         String fingerprint = RabbitMQClientHelper.computeConnectionFingerprint(options);
-        SharedClient shared = sharedClients.compute(name, (key, existing) -> {
-            if (existing != null) {
-                if (!existing.fingerprint.equals(fingerprint)) {
-                    throw ex.illegalStateSharedConnectionConfigMismatch(name);
-                }
-                existing.retain();
-                if (context != null) {
-                    existing.holder.ensureContext(context);
-                }
-                return existing;
-            }
-            return new SharedClient(name, new ClientHolder(
-                    RabbitMQClient.create(vertx(), options),
-                    config.getChannel(),
-                    vertx(),
-                    context), fingerprint);
-        });
-        clientRegistrations.put(config.getChannel(), new ClientRegistration(shared.holder, true, name));
-        return shared.holder;
-    }
-
-    public void releaseClient(String channel) {
-        ClientRegistration registration = clientRegistrations.remove(channel);
-        if (registration == null) {
-            return;
-        }
-
-        if (registration.shared) {
-            releaseSharedClient(registration.key, true);
-        } else {
-            stopClient(registration.holder.client(), true);
+        String existing = connectionFingerprints.putIfAbsent(connectionName, fingerprint);
+        if (existing != null && !existing.equals(fingerprint)) {
+            throw ex.illegalStateSharedConnectionConfigMismatch(connectionName);
         }
+        return clients.compute(fingerprint,
+                (key, current) -> (current == null ? new ClientHolder(RabbitMQClient.create(vertx(), options)) : current)
+                        .retain(channel));
     }
 
-    private void releaseSharedClient(String sharedName, boolean await) {
-        SharedClient shared = sharedClients.get(sharedName);
-        if (shared == null) {
-            return;
-        }
-        if (shared.release()) {
-            sharedClients.remove(sharedName, shared);
-            stopClient(shared.holder.client(), await);
+    public void releaseClient(String channel, boolean await) {
+        for (var e : clients.entrySet()) {
+            ClientHolder shared = e.getValue();
+            if (shared.channels().contains(channel)) {
+                if (clients.computeIfPresent(e.getKey(), (k, c) -> c.release(channel) ? null : c) == null) {
+                    connectionFingerprints.values().remove(e.getKey());
+                    stopClient(shared.client(), await);
+                }
+                return;
+            }
         }
     }
 
@@ -395,36 +344,4 @@ private void stopClient(RabbitMQClient client, boolean await) {
         }
     }
 
-    private static final class ClientRegistration {
-        final ClientHolder holder;
-        final boolean shared;
-        final String key;
-
-        private ClientRegistration(ClientHolder holder, boolean shared, String key) {
-            this.holder = holder;
-            this.shared = shared;
-            this.key = key;
-        }
-    }
-
-    private static final class SharedClient {
-        final String name;
-        final ClientHolder holder;
-        final String fingerprint;
-        final AtomicInteger references = new AtomicInteger(1);
-
-        private SharedClient(String name, ClientHolder holder, String fingerprint) {
-            this.name = name;
-            this.holder = holder;
-            this.fingerprint = fingerprint;
-        }
-
-        private void retain() {
-            references.incrementAndGet();
-        }
-
-        private boolean release() {
-            return references.decrementAndGet() == 0;
-        }
-    }
 }