Shared connection support + test coverage

Port of 1edffdd

Author: ozangunalp

smallrye-reactive-messaging-rabbitmq-og/src/main/java/io/smallrye/reactive/messaging/rabbitmq/og/ConnectionHolder.java

@@ -5,6 +5,8 @@
 
 import java.io.IOException;
 import java.time.Duration;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
@@ -25,11 +27,13 @@ public class ConnectionHolder {
 
     private final ConnectionFactory factory;
     private final String channelName;
+    private final String connectionName;
     private final Vertx vertx;
     private final Context context;
     private final AtomicBoolean connected = new AtomicBoolean(false);
     private final int reconnectAttempts;
     private final int reconnectInterval;
+    private final Set<String> channels = ConcurrentHashMap.newKeySet();
 
     private volatile Connection connection;
     private Consumer<Connection> onConnectionEstablished;
@@ -38,22 +42,22 @@ public ConnectionHolder(
             ConnectionFactory factory,
             String channelName,
             io.vertx.mutiny.core.Vertx mutinyVertx) {
-        this(factory, channelName, mutinyVertx, 0, 10);
+        this(factory, channelName, channelName, mutinyVertx, 0, 10);
     }
 
     public ConnectionHolder(
             ConnectionFactory factory,
             String channelName,
+            String connectionName,
             io.vertx.mutiny.core.Vertx mutinyVertx,
             int reconnectAttempts,
             int reconnectInterval) {
         this.factory = factory;
         this.channelName = channelName;
+        this.connectionName = connectionName;
         this.vertx = mutinyVertx.getDelegate();
         this.reconnectAttempts = reconnectAttempts;
         this.reconnectInterval = reconnectInterval;
-        // We need event loop context for each connection, so the duplicated context will be running on different event
-        // loop threads, but that's fine as long as we use the same context for all operations related to this connection
         this.context = ((VertxInternal) vertx).createEventLoopContext();
     }
 
@@ -74,7 +78,7 @@ public Uni<Connection> connect() {
             try {
                 log.establishingConnection(channelName);
 
-                Connection conn = factory.newConnection();
+                Connection conn = factory.newConnection(connectionName);
 
                 // Set connection before recovery listeners so createChannel() works in recovery callbacks
                 connection = conn;
@@ -182,4 +186,18 @@ public void close() {
             }
         }
     }
+
+    public Set<String> channels() {
+        return channels;
+    }
+
+    public ConnectionHolder retain(String channel) {
+        channels.add(channel);
+        return this;
+    }
+
+    public boolean release(String channel) {
+        channels.remove(channel);
+        return channels.isEmpty();
+    }
 }

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

@@ -137,6 +137,10 @@ public IncomingRabbitMQMetadata getRabbitMQMetadata() {
         return rabbitMQMetadata;
     }
 
+    public java.util.Optional<String> getCorrelationId() {
+        return java.util.Optional.ofNullable(rabbitMQMetadata.getCorrelationId());
+    }
+
     @Override
     public Supplier<CompletionStage<Void>> getAck() {
         return () -> ackHandler.handle(this);

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

@@ -3,9 +3,11 @@
 import static io.smallrye.reactive.messaging.annotations.ConnectorAttribute.Direction.INCOMING;
 import static io.smallrye.reactive.messaging.annotations.ConnectorAttribute.Direction.INCOMING_AND_OUTGOING;
 import static io.smallrye.reactive.messaging.annotations.ConnectorAttribute.Direction.OUTGOING;
+import static io.smallrye.reactive.messaging.rabbitmq.og.i18n.RabbitMQExceptions.ex;
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.Flow;
 
@@ -33,6 +35,7 @@
 import io.smallrye.reactive.messaging.health.HealthReport;
 import io.smallrye.reactive.messaging.health.HealthReporter;
 import io.smallrye.reactive.messaging.providers.connectors.ExecutionHolder;
+import io.smallrye.reactive.messaging.rabbitmq.og.internals.RabbitMQClientHelper;
 import io.vertx.core.Vertx;
 
 @ApplicationScoped
@@ -57,6 +60,7 @@
 @ConnectorAttribute(name = "reconnect-interval", direction = INCOMING_AND_OUTGOING, description = "The interval (in seconds) between two reconnection attempts", type = "int", alias = "rabbitmq-reconnect-interval", defaultValue = "10")
 @ConnectorAttribute(name = "network-recovery-interval", direction = INCOMING_AND_OUTGOING, description = "How long (ms) will automatic recovery wait before attempting to reconnect", type = "int", defaultValue = "5000")
 @ConnectorAttribute(name = "user", direction = INCOMING_AND_OUTGOING, description = "The user name to use when connecting to the broker", type = "string", defaultValue = "guest")
+@ConnectorAttribute(name = "shared-connection-name", direction = INCOMING_AND_OUTGOING, description = "Optional identifier allowing multiple channels to share the same RabbitMQ connection when set to the same value", type = "string")
 @ConnectorAttribute(name = "include-properties", direction = INCOMING_AND_OUTGOING, description = "Whether to include properties when a broker message is passed on the event bus", type = "boolean", defaultValue = "false")
 @ConnectorAttribute(name = "requested-channel-max", direction = INCOMING_AND_OUTGOING, description = "The initially requested maximum channel number", type = "int", defaultValue = "2047")
 @ConnectorAttribute(name = "requested-heartbeat", direction = INCOMING_AND_OUTGOING, description = "The initially requested heartbeat interval (seconds), zero for none", type = "int", defaultValue = "60")
@@ -161,6 +165,8 @@ public class RabbitMQConnector implements InboundConnector, OutboundConnector, H
 
     private final List<io.smallrye.reactive.messaging.rabbitmq.og.internals.IncomingRabbitMQChannel> incomings = new CopyOnWriteArrayList<>();
     private final List<io.smallrye.reactive.messaging.rabbitmq.og.internals.OutgoingRabbitMQChannel> outgoings = new CopyOnWriteArrayList<>();
+    private final Map<String, ConnectionHolder> connections = new ConcurrentHashMap<>();
+    private final Map<String, String> connectionFingerprints = new ConcurrentHashMap<>();
 
     RabbitMQConnector() {
         // used for proxies
@@ -170,13 +176,7 @@ public class RabbitMQConnector implements InboundConnector, OutboundConnector, H
     public Flow.Publisher<? extends Message<?>> getPublisher(Config config) {
         RabbitMQConnectorIncomingConfiguration ic = new RabbitMQConnectorIncomingConfiguration(config);
 
-        // Create ConnectionFactory
-        ConnectionFactory factory = io.smallrye.reactive.messaging.rabbitmq.og.internals.RabbitMQClientHelper
-                .createConnectionFactory(ic, connectionFactories, credentialsProviders, configCustomizers);
-
-        // Create ConnectionHolder
-        ConnectionHolder holder = new ConnectionHolder(factory, ic.getChannel(),
-                executionHolder.vertx(), ic.getReconnectAttempts(), ic.getReconnectInterval());
+        ConnectionHolder holder = getOrCreateConnectionHolder(ic);
 
         // Create IncomingRabbitMQChannel
         io.smallrye.reactive.messaging.rabbitmq.og.internals.IncomingRabbitMQChannel channel = new io.smallrye.reactive.messaging.rabbitmq.og.internals.IncomingRabbitMQChannel(
@@ -192,13 +192,7 @@ public Flow.Publisher<? extends Message<?>> getPublisher(Config config) {
     public Flow.Subscriber<? extends Message<?>> getSubscriber(Config config) {
         RabbitMQConnectorOutgoingConfiguration oc = new RabbitMQConnectorOutgoingConfiguration(config);
 
-        // Create ConnectionFactory
-        ConnectionFactory factory = io.smallrye.reactive.messaging.rabbitmq.og.internals.RabbitMQClientHelper
-                .createConnectionFactory(oc, connectionFactories, credentialsProviders, configCustomizers);
-
-        // Create ConnectionHolder
-        ConnectionHolder holder = new ConnectionHolder(factory, oc.getChannel(),
-                executionHolder.vertx(), oc.getReconnectAttempts(), oc.getReconnectInterval());
+        ConnectionHolder holder = getOrCreateConnectionHolder(oc);
 
         // Create OutgoingRabbitMQChannel
         io.smallrye.reactive.messaging.rabbitmq.og.internals.OutgoingRabbitMQChannel channel = new io.smallrye.reactive.messaging.rabbitmq.og.internals.OutgoingRabbitMQChannel(
@@ -297,6 +291,45 @@ public void terminate(
             }
         }
         outgoings.clear();
+
+        for (ConnectionHolder holder : connections.values()) {
+            holder.close();
+        }
+        connections.clear();
+        connectionFingerprints.clear();
+    }
+
+    public ConnectionHolder getOrCreateConnectionHolder(RabbitMQConnectorCommonConfiguration config) {
+        String channel = config.getChannel();
+        ConnectionFactory factory = RabbitMQClientHelper
+                .createConnectionFactory(config, connectionFactories, credentialsProviders, configCustomizers);
+        String connectionName = RabbitMQClientHelper.resolveConnectionName(config);
+        String fingerprint = RabbitMQClientHelper.computeConnectionFingerprint(factory);
+        String existing = connectionFingerprints.putIfAbsent(connectionName, fingerprint);
+        if (existing != null && !existing.equals(fingerprint)) {
+            throw ex.illegalStateSharedConnectionConfigMismatch(connectionName);
+        }
+        return connections.compute(fingerprint,
+                (key, current) -> {
+                    if (current == null) {
+                        current = new ConnectionHolder(factory, channel, connectionName,
+                                executionHolder.vertx(), config.getReconnectAttempts(), config.getReconnectInterval());
+                    }
+                    return current.retain(channel);
+                });
+    }
+
+    public void releaseClient(String channel) {
+        for (var e : connections.entrySet()) {
+            ConnectionHolder holder = e.getValue();
+            if (holder.channels().contains(channel)) {
+                if (connections.computeIfPresent(e.getKey(), (k, c) -> c.release(channel) ? null : c) == null) {
+                    connectionFingerprints.values().remove(e.getKey());
+                    holder.close();
+                }
+                return;
+            }
+        }
     }
 
     public Vertx vertx() {

smallrye-reactive-messaging-rabbitmq-og/src/main/java/io/smallrye/reactive/messaging/rabbitmq/og/i18n/RabbitMQExceptions.java

@@ -50,4 +50,7 @@ public interface RabbitMQExceptions {
 
     @Message(id = 19012, value = "RabbitMQ channel is closed")
     IllegalStateException illegalStateChannelClosed();
+
+    @Message(id = 19013, value = "Shared connection '%s' has mismatched configuration; ensure all channels using the same shared-connection-name have identical connection settings")
+    IllegalStateException illegalStateSharedConnectionConfigMismatch(String sharedConnectionName);
 }

smallrye-reactive-messaging-rabbitmq-og/src/main/java/io/smallrye/reactive/messaging/rabbitmq/og/internals/IncomingRabbitMQChannel.java

@@ -37,6 +37,7 @@
 import io.smallrye.reactive.messaging.rabbitmq.og.tracing.RabbitMQOpenTelemetryInstrumenter;
 import io.smallrye.reactive.messaging.rabbitmq.og.tracing.RabbitMQTrace;
 import io.vertx.core.Context;
+import io.vertx.core.internal.VertxInternal;
 
 /**
  * Incoming RabbitMQ channel that consumes messages from a queue.
@@ -49,6 +50,7 @@ public class IncomingRabbitMQChannel {
     private final Instance<java.util.Map<String, ?>> configMaps;
     private final RabbitMQOpenTelemetryInstrumenter instrumenter;
 
+    private final Context incomingContext;
     private final AtomicBoolean subscribed = new AtomicBoolean(false);
     private final AtomicInteger outstandingMessages = new AtomicInteger(0);
     private final AtomicReference<Channel> channelRef = new AtomicReference<>();
@@ -65,10 +67,10 @@ public IncomingRabbitMQChannel(
             Instance<java.util.Map<String, ?>> configMaps,
             Instance<OpenTelemetry> openTelemetryInstance) {
 
-        System.out.println("Creating IncomingRabbitMQChannel for channel: " + configuration.getChannel());
         this.connectionHolder = connectionHolder;
         this.configuration = configuration;
         this.configMaps = configMaps;
+        this.incomingContext = ((VertxInternal) connectionHolder.getVertx()).createEventLoopContext();
 
         // Initialize tracing if enabled
         if (configuration.getTracingEnabled()) {
@@ -94,7 +96,7 @@ public Multi<? extends Message<?>> getStream() {
     private Multi<? extends Message<?>> createStream() {
         // Determine if broadcast mode is needed
         boolean broadcast = configuration.getBroadcast();
-        Context context = connectionHolder.getContext();
+        Context context = incomingContext;
 
         Multi<Message<?>> messageStream;
 
@@ -190,7 +192,7 @@ private void setupConsumerOnConnection(
         Channel channel = connectionHolder.createChannel();
         channelRef.set(channel);
 
-        Context context = connectionHolder.getContext();
+        Context context = incomingContext;
 
         // Set up QoS for backpressure (prefetch count)
         // Note: In auto-ack mode, messages are immediately acknowledged upon delivery,

smallrye-reactive-messaging-rabbitmq-og/src/main/java/io/smallrye/reactive/messaging/rabbitmq/og/internals/OutgoingRabbitMQChannel.java

@@ -29,6 +29,7 @@
 import io.smallrye.reactive.messaging.rabbitmq.og.tracing.RabbitMQOpenTelemetryInstrumenter;
 import io.smallrye.reactive.messaging.rabbitmq.og.tracing.RabbitMQTrace;
 import io.vertx.core.Context;
+import io.vertx.core.internal.VertxInternal;
 
 /**
  * Outgoing RabbitMQ channel that publishes messages to an exchange.
@@ -45,8 +46,8 @@ public class OutgoingRabbitMQChannel implements Subscriber<Message<?>> {
     private final AtomicBoolean completed = new AtomicBoolean(false);
     private final AtomicLong inflightMessages = new AtomicLong(0);
 
+    private final Context outgoingContext;
     private Channel channel;
-    private Context context;
     private Subscription subscription;
 
     private final boolean publisherConfirms;
@@ -68,6 +69,7 @@ public OutgoingRabbitMQChannel(
         this.connectionHolder = connectionHolder;
         this.configuration = configuration;
         this.configMaps = configMaps;
+        this.outgoingContext = ((VertxInternal) connectionHolder.getVertx()).createEventLoopContext();
 
         // Initialize tracing if enabled
         if (configuration.getTracingEnabled()) {
@@ -92,7 +94,6 @@ private void initialize() {
                                 try {
                                     // Create channel for publishing
                                     channel = connectionHolder.createChannel();
-                                    context = connectionHolder.getContext();
 
                                     // Set up topology
                                     setupTopology();
@@ -128,11 +129,11 @@ private void setupTopology() throws IOException {
 
     private void setupConfirmListeners() throws IOException {
         ConfirmCallback ackCallback = (sequenceNumber, multiple) -> {
-            context.runOnContext(v -> handleConfirm(sequenceNumber, multiple, true));
+            outgoingContext.runOnContext(v -> handleConfirm(sequenceNumber, multiple, true));
         };
 
         ConfirmCallback nackCallback = (sequenceNumber, multiple) -> {
-            context.runOnContext(v -> handleConfirm(sequenceNumber, multiple, false));
+            outgoingContext.runOnContext(v -> handleConfirm(sequenceNumber, multiple, false));
         };
 
         channel.addConfirmListener(ackCallback, nackCallback);
@@ -221,7 +222,7 @@ private Uni<Void> publishWithRetry(Message<?> message, int remainingAttempts) {
                 throw new RuntimeException("Failed to publish message", e);
             }
         })
-                .runSubscriptionOn(command -> context.runOnContext(x -> command.run()))
+                .runSubscriptionOn(command -> outgoingContext.runOnContext(x -> command.run()))
                 .onFailure().retry()
                 .withBackOff(Duration.ofSeconds(retryInterval))
                 .atMost(remainingAttempts);

smallrye-reactive-messaging-rabbitmq-og/src/main/java/io/smallrye/reactive/messaging/rabbitmq/og/internals/RabbitMQClientHelper.java

@@ -8,9 +8,11 @@
 
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.security.KeyManagementException;
 import java.security.KeyStore;
 import java.security.KeyStoreException;
+import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.security.cert.CertificateException;
 import java.time.Duration;
@@ -87,9 +89,7 @@ static ConnectionFactory getConnectionFactory(
             RabbitMQConnectorCommonConfiguration config,
             Instance<CredentialsProvider> credentialsProviders) {
 
-        String connectionName = String.format("%s (%s)",
-                config.getChannel(),
-                config instanceof RabbitMQConnectorIncomingConfiguration ? "Incoming" : "Outgoing");
+        String connectionName = resolveConnectionName(config);
 
         Address[] addresses = config.getAddresses()
                 .map(Address::parseAddresses)
@@ -195,6 +195,42 @@ private static SSLContext createSSLContext(RabbitMQConnectorCommonConfiguration
         }
     }
 
+    public static String resolveConnectionName(RabbitMQConnectorCommonConfiguration config) {
+        return config.getSharedConnectionName()
+                .orElseGet(() -> String.format("%s (%s)",
+                        config.getChannel(),
+                        config instanceof RabbitMQConnectorIncomingConfiguration ? "Incoming" : "Outgoing"));
+    }
+
+    public static String computeConnectionFingerprint(ConnectionFactory factory) {
+        String raw = factory.getHost()
+                + ":" + factory.getPort()
+                + ":" + factory.getVirtualHost()
+                + ":" + factory.getUsername()
+                + ":" + factory.isAutomaticRecoveryEnabled()
+                + ":" + factory.getConnectionTimeout()
+                + ":" + factory.getHandshakeTimeout()
+                + ":" + factory.getRequestedChannelMax()
+                + ":" + factory.getRequestedHeartbeat()
+                + ":" + factory.getNetworkRecoveryInterval()
+                + ":" + factory.isSSL();
+        return sha256(raw);
+    }
+
+    private static String sha256(String value) {
+        try {
+            MessageDigest digest = MessageDigest.getInstance("SHA-256");
+            byte[] hash = digest.digest(value.getBytes(StandardCharsets.UTF_8));
+            StringBuilder hex = new StringBuilder(hash.length * 2);
+            for (byte b : hash) {
+                hex.append(String.format("%02x", b));
+            }
+            return hex.toString();
+        } catch (NoSuchAlgorithmException e) {
+            throw new IllegalStateException("Unable to compute SHA-256 hash", e);
+        }
+    }
+
     public static String serverQueueName(String name) {
         if (name.equals("(server.auto)")) {
             return "";