* Add IUriSelector

lukebakken · lukebakken · commit ebce43dcb33e · 2024-11-26T10:01:57.000-08:00
diff --git a/RabbitMQ.AMQP.Client/ConnectionSettings.cs b/RabbitMQ.AMQP.Client/ConnectionSettings.cs
@@ -12,6 +12,19 @@
 
 namespace RabbitMQ.AMQP.Client
 {
+    public interface IUriSelector
+    {
+        Uri Select(ICollection<Uri> uris);
+    }
+
+    public class RandomUriSelector : IUriSelector
+    {
+        public Uri Select(ICollection<Uri> uris)
+        {
+            return uris.Skip(Utils.RandomNext(0, uris.Count)).First();
+        }
+    }
+
     public class ConnectionSettingsBuilder
     {
         private string _host = "localhost";
@@ -27,6 +40,7 @@ public class ConnectionSettingsBuilder
         private TlsSettings? _tlsSettings = null;
         private Uri? _uri;
         private List<Uri>? _uris;
+        private IUriSelector? _uriSelector;
 
         public static ConnectionSettingsBuilder Create()
         {
@@ -125,10 +139,16 @@ public ConnectionSettingsBuilder Uris(IEnumerable<Uri> uris)
             return this;
         }
 
+        public ConnectionSettingsBuilder UriSelector(IUriSelector uriSelector)
+        {
+            _uriSelector = uriSelector;
+            return this;
+        }
+
         public ConnectionSettings Build()
         {
-            ValidateUris();
             // TODO this should do something similar to consolidate in the Java code
+            ValidateUris();
             if (_uri is not null)
             {
                 return new ConnectionSettings(_uri,
@@ -140,6 +160,7 @@ public ConnectionSettings Build()
             else if (_uris is not null)
             {
                 return new ConnectionSettings(_uris,
+                    _uriSelector,
                     _containerId, _saslMechanism,
                     _recoveryConfiguration,
                     _maxFrameSize,
@@ -171,7 +192,9 @@ private void ValidateUris()
     public class ConnectionSettings : IEquatable<ConnectionSettings>
     {
         private readonly Address _address = new("amqp://localhost:5672");
-        private readonly List<Address> _addresses = new();
+        private readonly List<Uri>? _uris;
+        private readonly Dictionary<Uri, Address>? _uriToAddress;
+        private readonly IUriSelector _uriSelector = new RandomUriSelector();
         private readonly string _virtualHost = Consts.DefaultVirtualHost;
         private readonly string _containerId = string.Empty;
         private readonly uint _maxFrameSize = Consts.DefaultMaxFrameSize;
@@ -197,7 +220,6 @@ public ConnectionSettings(Uri uri,
                 password: password,
                 path: "/",
                 scheme: uri.Scheme);
-            _addresses.Add(_address);
 
             if (_address.UseSsl && _tlsSettings == null)
             {
@@ -206,16 +228,27 @@ public ConnectionSettings(Uri uri,
         }
 
         public ConnectionSettings(IEnumerable<Uri> uris,
+            IUriSelector? uriSelector = null,
             string? containerId = null,
             SaslMechanism? saslMechanism = null,
             IRecoveryConfiguration? recoveryConfiguration = null,
             uint? maxFrameSize = null,
             TlsSettings? tlsSettings = null)
             : this(containerId, saslMechanism, recoveryConfiguration, maxFrameSize, tlsSettings)
         {
+            _uris = uris.ToList();
+
+            if (uriSelector is not null)
+            {
+                _uriSelector = uriSelector;
+            }
+
+            _uriToAddress = new(_uris.Count);
+
             string? tmpVirtualHost = null;
 
-            foreach (Uri uri in uris)
+            bool first = true;
+            foreach (Uri uri in _uris)
             {
                 (string? user, string? password) = ProcessUserInfo(uri);
 
@@ -238,10 +271,15 @@ public ConnectionSettings(IEnumerable<Uri> uris,
                     password: password,
                     path: "/",
                     scheme: uri.Scheme);
-                _addresses.Add(address);
-            }
 
-            _address = _addresses[0];
+                _uriToAddress[uri] = address;
+
+                if (first)
+                {
+                    _address = address;
+                    first = false;
+                }
+            }
 
             if (tmpVirtualHost is not null)
             {
@@ -265,7 +303,6 @@ public ConnectionSettings(string scheme,
             _address = new Address(host: host, port: port,
                 user: user, password: password,
                 path: "/", scheme: scheme);
-            _addresses.Add(_address);
 
             if (virtualHost is not null)
             {
@@ -331,23 +368,34 @@ internal Address Address
         {
             get
             {
-                if (_addresses.Count == 1)
+                if (_uris is not null &&
+                    _uriSelector is not null &&
+                    _uriToAddress is not null)
                 {
-                    if (false == Object.ReferenceEquals(_address, _addresses[0]))
-                    {
-                        InternalBugException.CreateAndThrow("_address must be same object as _addresses[0]");
-                    }
-                    return _address;
+                    Uri uri = _uriSelector.Select(_uris);
+                    return _uriToAddress[uri];
                 }
                 else
                 {
-                    // Note: the max value for RandomNext is not inclusive
-                    return _addresses[Utils.RandomNext(0, _addresses.Count)];
+                    return _address;
                 }
             }
         }
 
-        internal List<Address> Addresses => _addresses;
+        internal List<Address>? Addresses
+        {
+            get
+            {
+                if (_uriToAddress is not null)
+                {
+                    return _uriToAddress.Values.ToList();
+                }
+                else
+                {
+                    return null;
+                }
+            }
+        }
 
         public override string ToString()
         {
diff --git a/RabbitMQ.AMQP.Client/PublicAPI.Unshipped.txt b/RabbitMQ.AMQP.Client/PublicAPI.Unshipped.txt
@@ -59,7 +59,7 @@ RabbitMQ.AMQP.Client.ConnectionException.ConnectionException(string! message) ->
 RabbitMQ.AMQP.Client.ConnectionException.ConnectionException(string! message, System.Exception! innerException) -> void
 RabbitMQ.AMQP.Client.ConnectionSettings
 RabbitMQ.AMQP.Client.ConnectionSettings.ConnectionSettings(string! scheme, string! host, int port, string? user = null, string? password = null, string? virtualHost = null, string! containerId = "", RabbitMQ.AMQP.Client.SaslMechanism? saslMechanism = null, RabbitMQ.AMQP.Client.IRecoveryConfiguration? recoveryConfiguration = null, uint? maxFrameSize = null, RabbitMQ.AMQP.Client.TlsSettings? tlsSettings = null) -> void
-RabbitMQ.AMQP.Client.ConnectionSettings.ConnectionSettings(System.Collections.Generic.IEnumerable<System.Uri!>! uris, string? containerId = null, RabbitMQ.AMQP.Client.SaslMechanism? saslMechanism = null, RabbitMQ.AMQP.Client.IRecoveryConfiguration? recoveryConfiguration = null, uint? maxFrameSize = null, RabbitMQ.AMQP.Client.TlsSettings? tlsSettings = null) -> void
+RabbitMQ.AMQP.Client.ConnectionSettings.ConnectionSettings(System.Collections.Generic.IEnumerable<System.Uri!>! uris, RabbitMQ.AMQP.Client.IUriSelector? uriSelector = null, string? containerId = null, RabbitMQ.AMQP.Client.SaslMechanism? saslMechanism = null, RabbitMQ.AMQP.Client.IRecoveryConfiguration? recoveryConfiguration = null, uint? maxFrameSize = null, RabbitMQ.AMQP.Client.TlsSettings? tlsSettings = null) -> void
 RabbitMQ.AMQP.Client.ConnectionSettings.ConnectionSettings(System.Uri! uri, string? containerId = null, RabbitMQ.AMQP.Client.SaslMechanism? saslMechanism = null, RabbitMQ.AMQP.Client.IRecoveryConfiguration? recoveryConfiguration = null, uint? maxFrameSize = null, RabbitMQ.AMQP.Client.TlsSettings? tlsSettings = null) -> void
 RabbitMQ.AMQP.Client.ConnectionSettings.ContainerId.get -> string!
 RabbitMQ.AMQP.Client.ConnectionSettings.Host.get -> string!
@@ -88,6 +88,7 @@ RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.Scheme(string! scheme) -> RabbitM
 RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.TlsSettings(RabbitMQ.AMQP.Client.TlsSettings! tlsSettings) -> RabbitMQ.AMQP.Client.ConnectionSettingsBuilder!
 RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.Uri(System.Uri! uri) -> RabbitMQ.AMQP.Client.ConnectionSettingsBuilder!
 RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.Uris(System.Collections.Generic.IEnumerable<System.Uri!>! uris) -> RabbitMQ.AMQP.Client.ConnectionSettingsBuilder!
+RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.UriSelector(RabbitMQ.AMQP.Client.IUriSelector! uriSelector) -> RabbitMQ.AMQP.Client.ConnectionSettingsBuilder!
 RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.User(string! user) -> RabbitMQ.AMQP.Client.ConnectionSettingsBuilder!
 RabbitMQ.AMQP.Client.ConnectionSettingsBuilder.VirtualHost(string! virtualHost) -> RabbitMQ.AMQP.Client.ConnectionSettingsBuilder!
 RabbitMQ.AMQP.Client.Consts
@@ -671,6 +672,8 @@ RabbitMQ.AMQP.Client.IStreamSpecification.InitialClusterSize(int initialClusterS
 RabbitMQ.AMQP.Client.IStreamSpecification.MaxAge(System.TimeSpan maxAge) -> RabbitMQ.AMQP.Client.IStreamSpecification!
 RabbitMQ.AMQP.Client.IStreamSpecification.MaxSegmentSizeBytes(RabbitMQ.AMQP.Client.ByteCapacity! maxSegmentSize) -> RabbitMQ.AMQP.Client.IStreamSpecification!
 RabbitMQ.AMQP.Client.IStreamSpecification.Queue() -> RabbitMQ.AMQP.Client.IQueueSpecification!
+RabbitMQ.AMQP.Client.IUriSelector
+RabbitMQ.AMQP.Client.IUriSelector.Select(System.Collections.Generic.ICollection<System.Uri!>! uris) -> System.Uri!
 RabbitMQ.AMQP.Client.LifeCycleCallBack
 RabbitMQ.AMQP.Client.MessageHandler
 RabbitMQ.AMQP.Client.MetricsReporter
@@ -714,6 +717,9 @@ RabbitMQ.AMQP.Client.QueueType.STREAM = 2 -> RabbitMQ.AMQP.Client.QueueType
 RabbitMQ.AMQP.Client.QuorumQueueDeadLetterStrategy
 RabbitMQ.AMQP.Client.QuorumQueueDeadLetterStrategy.AtLeastOnce = 1 -> RabbitMQ.AMQP.Client.QuorumQueueDeadLetterStrategy
 RabbitMQ.AMQP.Client.QuorumQueueDeadLetterStrategy.AtMostOnce = 0 -> RabbitMQ.AMQP.Client.QuorumQueueDeadLetterStrategy
+RabbitMQ.AMQP.Client.RandomUriSelector
+RabbitMQ.AMQP.Client.RandomUriSelector.RandomUriSelector() -> void
+RabbitMQ.AMQP.Client.RandomUriSelector.Select(System.Collections.Generic.ICollection<System.Uri!>! uris) -> System.Uri!
 RabbitMQ.AMQP.Client.RecoveryConfiguration
 RabbitMQ.AMQP.Client.RecoveryConfiguration.Activated(bool activated) -> RabbitMQ.AMQP.Client.IRecoveryConfiguration!
 RabbitMQ.AMQP.Client.RecoveryConfiguration.BackOffDelayPolicy(RabbitMQ.AMQP.Client.IBackOffDelayPolicy! backOffDelayPolicy) -> RabbitMQ.AMQP.Client.IRecoveryConfiguration!
diff --git a/Tests/ConnectionTests/ConnectionSettingsTests.cs b/Tests/ConnectionTests/ConnectionSettingsTests.cs
@@ -131,6 +131,7 @@ public void ConnectionSettingsViaUris()
         Assert.Equal(vhost, connectionSettings.VirtualHost);
         Assert.Equal(scheme, connectionSettings.Scheme);
 
+        Assert.NotNull(connectionSettings.Addresses);
         Amqp.Address a0 = connectionSettings.Addresses[0];
         Assert.Equal(host, a0.Host);
         Assert.Equal(5671, a0.Port);
@@ -196,6 +197,7 @@ public void ConnectionSettingsViaBuilderWithUris()
         Assert.Equal(vhost, connectionSettings.VirtualHost);
         Assert.Equal(scheme, connectionSettings.Scheme);
 
+        Assert.NotNull(connectionSettings.Addresses);
         Amqp.Address a0 = connectionSettings.Addresses[0];
         Assert.Equal(host, a0.Host);
         Assert.Equal(5671, a0.Port);
