CSHARP-2288: Add option for applications to register a custom server selector

Author: MikalaiMazurenka

src/MongoDB.Driver.Core/Core/Clusters/Cluster.cs

@@ -67,6 +67,7 @@ internal abstract class Cluster : ICluster
         private ClusterDescription _description;
         private TaskCompletionSource<bool> _descriptionChangedTaskCompletionSource;
         private readonly object _descriptionLock = new object();
+        private readonly LatencyLimitingServerSelector _latencyLimitingServerSelector;
         private Timer _rapidHeartbeatTimer;
         private readonly object _serverSelectionWaitQueueLock = new object();
         private int _serverSelectionWaitQueueSize;
@@ -91,6 +92,7 @@ protected Cluster(ClusterSettings settings, IClusterableServerFactory serverFact
             _clusterId = new ClusterId();
             _description = ClusterDescription.CreateInitial(_clusterId, _settings.ConnectionMode);
             _descriptionChangedTaskCompletionSource = new TaskCompletionSource<bool>();
+            _latencyLimitingServerSelector = new LatencyLimitingServerSelector(settings.LocalThreshold);
 
             _rapidHeartbeatTimer = new Timer(RapidHeartbeatTimerCallback, null, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
 
@@ -507,25 +509,23 @@ public void WaitingForDescriptionToChange()
             private IServerSelector DecorateSelector(IServerSelector selector)
             {
                 var settings = _cluster.Settings;
-                if (settings.PreServerSelector != null || settings.PostServerSelector != null)
+                var allSelectors = new List<IServerSelector>();
+
+                if (settings.PreServerSelector != null)
                 {
-                    var allSelectors = new List<IServerSelector>();
-                    if (settings.PreServerSelector != null)
-                    {
-                        allSelectors.Add(settings.PreServerSelector);
-                    }
+                    allSelectors.Add(settings.PreServerSelector);
+                }
 
-                    allSelectors.Add(selector);
+                allSelectors.Add(selector);
 
-                    if (settings.PostServerSelector != null)
-                    {
-                        allSelectors.Add(settings.PostServerSelector);
-                    }
-
-                    return new CompositeServerSelector(allSelectors);
+                if (settings.PostServerSelector != null)
+                {
+                    allSelectors.Add(settings.PostServerSelector);
                 }
 
-                return selector;
+                allSelectors.Add(_cluster._latencyLimitingServerSelector);
+
+                return new CompositeServerSelector(allSelectors);
             }
         }
 

src/MongoDB.Driver.Core/Core/Configuration/ClusterSettings.cs

@@ -38,6 +38,7 @@ public class ClusterSettings
         private readonly ClusterConnectionMode _connectionMode;
         private readonly IReadOnlyList<EndPoint> _endPoints;
         private readonly IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> _kmsProviders;
+        private readonly TimeSpan _localThreshold;
         private readonly int _maxServerSelectionWaitQueueSize;
         private readonly string _replicaSetName;
         private readonly IReadOnlyDictionary<string, BsonDocument> _schemaMap;
@@ -53,6 +54,7 @@ public class ClusterSettings
         /// <param name="connectionMode">The connection mode.</param>
         /// <param name="endPoints">The end points.</param>
         /// <param name="kmsProviders">The kms providers.</param>
+        /// <param name="localThreshold">The local threshold.</param>
         /// <param name="maxServerSelectionWaitQueueSize">Maximum size of the server selection wait queue.</param>
         /// <param name="replicaSetName">Name of the replica set.</param>
         /// <param name="serverSelectionTimeout">The server selection timeout.</param>
@@ -64,6 +66,7 @@ public ClusterSettings(
             Optional<ClusterConnectionMode> connectionMode = default(Optional<ClusterConnectionMode>),
             Optional<IEnumerable<EndPoint>> endPoints = default(Optional<IEnumerable<EndPoint>>),
             Optional<IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>>> kmsProviders = default(Optional<IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>>>),
+            Optional<TimeSpan> localThreshold = default,
             Optional<int> maxServerSelectionWaitQueueSize = default(Optional<int>),
             Optional<string> replicaSetName = default(Optional<string>),
             Optional<TimeSpan> serverSelectionTimeout = default(Optional<TimeSpan>),
@@ -75,6 +78,7 @@ public ClusterSettings(
             _connectionMode = connectionMode.WithDefault(ClusterConnectionMode.Automatic);
             _endPoints = Ensure.IsNotNull(endPoints.WithDefault(__defaultEndPoints), "endPoints").ToList();
             _kmsProviders = kmsProviders.WithDefault(null);
+            _localThreshold = Ensure.IsGreaterThanOrEqualToZero(localThreshold.WithDefault(TimeSpan.FromMilliseconds(15)), "localThreshold");
             _maxServerSelectionWaitQueueSize = Ensure.IsGreaterThanOrEqualToZero(maxServerSelectionWaitQueueSize.WithDefault(500), "maxServerSelectionWaitQueueSize");
             _replicaSetName = replicaSetName.WithDefault(null);
             _serverSelectionTimeout = Ensure.IsGreaterThanOrEqualToZero(serverSelectionTimeout.WithDefault(TimeSpan.FromSeconds(30)), "serverSelectionTimeout");
@@ -118,6 +122,17 @@ public IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>> KmsProvi
             get { return _kmsProviders; }
         }
 
+        /// <summary>
+        /// Gets the local threshold.
+        /// </summary>
+        /// <value>
+        /// The local threshold.
+        /// </value>
+        public TimeSpan LocalThreshold
+        {
+            get { return _localThreshold; }
+        }
+
         /// <summary>
         /// Gets the maximum size of the server selection wait queue.
         /// </summary>
@@ -202,6 +217,7 @@ public IServerSelector PostServerSelector
         /// <param name="connectionMode">The connection mode.</param>
         /// <param name="endPoints">The end points.</param>
         /// <param name="kmsProviders">The kms providers.</param>
+        /// <param name="localThreshold">The local threshold.</param>
         /// <param name="maxServerSelectionWaitQueueSize">Maximum size of the server selection wait queue.</param>
         /// <param name="replicaSetName">Name of the replica set.</param>
         /// <param name="serverSelectionTimeout">The server selection timeout.</param>
@@ -214,6 +230,7 @@ public ClusterSettings With(
             Optional<ClusterConnectionMode> connectionMode = default(Optional<ClusterConnectionMode>),
             Optional<IEnumerable<EndPoint>> endPoints = default(Optional<IEnumerable<EndPoint>>),
             Optional<IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>>> kmsProviders = default(Optional<IReadOnlyDictionary<string, IReadOnlyDictionary<string, object>>>),
+            Optional<TimeSpan> localThreshold = default(Optional<TimeSpan>),
             Optional<int> maxServerSelectionWaitQueueSize = default(Optional<int>),
             Optional<string> replicaSetName = default(Optional<string>),
             Optional<TimeSpan> serverSelectionTimeout = default(Optional<TimeSpan>),
@@ -226,6 +243,7 @@ public ClusterSettings With(
                 connectionMode: connectionMode.WithDefault(_connectionMode),
                 endPoints: Optional.Enumerable(endPoints.WithDefault(_endPoints)),
                 kmsProviders: Optional.Create(kmsProviders.WithDefault(_kmsProviders)),
+                localThreshold: localThreshold.WithDefault(_localThreshold),
                 maxServerSelectionWaitQueueSize: maxServerSelectionWaitQueueSize.WithDefault(_maxServerSelectionWaitQueueSize),
                 replicaSetName: replicaSetName.WithDefault(_replicaSetName),
                 serverSelectionTimeout: serverSelectionTimeout.WithDefault(_serverSelectionTimeout),

src/MongoDB.Driver/ClusterRegistry.cs

@@ -85,10 +85,10 @@ private ClusterSettings ConfigureCluster(ClusterSettings settings, ClusterKey cl
                 connectionMode: clusterKey.ConnectionMode.ToCore(),
                 endPoints: Optional.Enumerable(endPoints),
                 kmsProviders: Optional.Create(clusterKey.KmsProviders),
+                localThreshold: clusterKey.LocalThreshold,
                 replicaSetName: clusterKey.ReplicaSetName,
                 maxServerSelectionWaitQueueSize: clusterKey.WaitQueueSize,
                 serverSelectionTimeout: clusterKey.ServerSelectionTimeout,
-                postServerSelector: new LatencyLimitingServerSelector(clusterKey.LocalThreshold),
                 schemaMap: Optional.Create(clusterKey.SchemaMap),
                 scheme: clusterKey.Scheme);
         }

tests/MongoDB.Driver.Core.Tests/Core/Clusters/ClusterTests.cs

@@ -37,14 +37,13 @@ namespace MongoDB.Driver.Core.Clusters
 {
     public class ClusterTests
     {
-        private EventCapturer _capturedEvents;
-        private Mock<IClusterableServerFactory> _mockServerFactory;
+        private readonly EventCapturer _capturedEvents;
+        private readonly Mock<IClusterableServerFactory> _mockServerFactory;
         private ClusterSettings _settings;
 
         public ClusterTests()
         {
-            _settings = new ClusterSettings(serverSelectionTimeout: TimeSpan.FromSeconds(2),
-                postServerSelector: new LatencyLimitingServerSelector(TimeSpan.FromMinutes(2)));
+            _settings = new ClusterSettings(serverSelectionTimeout: TimeSpan.FromSeconds(2));
             _mockServerFactory = new Mock<IClusterableServerFactory>();
             _mockServerFactory.Setup(f => f.CreateServer(It.IsAny<ClusterId>(), It.IsAny<IClusterClock>(), It.IsAny<EndPoint>()))
                 .Returns((ClusterId clusterId, IClusterClock clusterClock, EndPoint endPoint) =>
@@ -481,6 +480,59 @@ public void SelectServer_should_apply_both_pre_and_post_server_selectors(
             _capturedEvents.Any().Should().BeFalse();
         }
 
+        [Theory]
+        [ParameterAttributeData]
+        public void SelectServer_should_call_custom_selector(
+            [Values(true, false)] bool withEligibleServers,
+            [Values(true, false)] bool async)
+        {
+            int numberOfCustomServerSelectorCalls = 0;
+            var customServerSelector = new DelegateServerSelector((c, s) =>
+            {
+                numberOfCustomServerSelectorCalls++;
+                return s.Skip(1);
+            });
+
+            var settings = _settings.With(postServerSelector: customServerSelector);
+            var subject = new StubCluster(settings, _mockServerFactory.Object, _capturedEvents);
+
+            subject.Initialize();
+            subject.SetServerDescriptions(
+                ServerDescriptionHelper.Connected(subject.Description.ClusterId, new DnsEndPoint("localhost", 27019)),
+                ServerDescriptionHelper.Connected(subject.Description.ClusterId, new DnsEndPoint("localhost", 27020)));
+            _capturedEvents.Clear();
+
+            if (withEligibleServers)
+            {
+                var selectedServer = SelectServerAttempt(
+                    subject,
+                    new DelegateServerSelector((c, s) => s), // do not filter servers
+                    async);
+
+                var selectedServerPort = ((DnsEndPoint)selectedServer.EndPoint).Port;
+                selectedServerPort.Should().Be(27020);
+                _capturedEvents.Next().Should().BeOfType<ClusterSelectingServerEvent>();
+                _capturedEvents.Next().Should().BeOfType<ClusterSelectedServerEvent>();
+            }
+            else
+            {
+                var exception = Record.Exception(
+                    () =>
+                        SelectServerAttempt(
+                            subject,
+                            new DelegateServerSelector((c, s) => new ServerDescription[0]), // no eligible servers
+                            async));
+
+                exception.Should().BeOfType<TimeoutException>();
+                _capturedEvents.Next().Should().BeOfType<ClusterSelectingServerEvent>();
+                _capturedEvents.Next().Should().BeOfType<ClusterSelectingServerFailedEvent>();
+            }
+
+            numberOfCustomServerSelectorCalls.Should().Be(1);
+            _capturedEvents.Any().Should().BeFalse();
+        }
+
+        // private methods
         private StubCluster CreateSubject(ClusterConnectionMode connectionMode = ClusterConnectionMode.Automatic, TimeSpan? serverSelectionTimeout = null)
         {
             _settings = _settings.With(connectionMode: connectionMode);
@@ -492,6 +544,22 @@ private StubCluster CreateSubject(ClusterConnectionMode connectionMode = Cluster
             return new StubCluster(_settings, _mockServerFactory.Object, _capturedEvents);
         }
 
+        private IServer SelectServerAttempt(Cluster cluster, IServerSelector operationSelector, bool async)
+        {
+            if (async)
+            {
+                return cluster
+                    .SelectServerAsync(operationSelector, CancellationToken.None)
+                    .GetAwaiter()
+                    .GetResult();
+            }
+            else
+            {
+                return cluster.SelectServer(operationSelector, CancellationToken.None);
+            }
+        }
+
+        // nested types
         private class StubCluster : Cluster
         {
             private Dictionary<EndPoint, IClusterableServer> _servers = new Dictionary<EndPoint, IClusterableServer>();