CSHARP-724: Make MongoServer Primary property robust.

Author: rstam

MongoDB.Driver/Communication/Proxies/DiscoveringMongoServerProxy.cs

@@ -145,6 +145,12 @@ public MongoServerState State
             }
         }
 
+        // internal properties
+        internal IMongoServerProxy WrappedProxy
+        {
+            get { return _serverProxy; }
+        }
+
         // public methods
         /// <summary>
         /// Chooses the server instance.

MongoDB.Driver/Communication/Proxies/MultipleInstanceMongoServerProxy.cs

@@ -398,6 +398,13 @@ protected void MakeInstancesMatchAddresses(IEnumerable<MongoServerAddress> addre
         protected virtual void ProcessConnectedInstanceStateChange(MongoServerInstance instance)
         { }
 
+        /// <summary>
+        /// Processes the disconnected instance state change.
+        /// </summary>
+        /// <param name="instance">The instance.</param>
+        protected virtual void ProcessDisconnectedInstanceStateChange(MongoServerInstance instance)
+        { }
+
         // private methods
         private void AddInstance(MongoServerInstance instance)
         {
@@ -472,6 +479,7 @@ private void ProcessInstanceStateChange(MongoServerInstance instance)
                     else
                     {
                         _connectedInstances.Remove(instance);
+                        ProcessDisconnectedInstanceStateChange(instance);
                     }
                 }
 

MongoDB.Driver/Communication/Proxies/ReplicaSetMongoServerProxy.cs

@@ -28,6 +28,7 @@ internal sealed class ReplicaSetMongoServerProxy : MultipleInstanceMongoServerPr
         // private fields
         private readonly Random _random = new Random();
         private readonly object _randomLock = new object();
+        private MongoServerInstance _primary;
         private string _replicaSetName;
 
         // constructors
@@ -53,6 +54,17 @@ public ReplicaSetMongoServerProxy(MongoServerSettings serverSettings, IEnumerabl
         { }
 
         // public properties
+        /// <summary>
+        /// Gets the primary.
+        /// </summary>
+        /// <value>
+        /// The primary.
+        /// </value>
+        public MongoServerInstance Primary
+        {
+            get { return _primary; }
+        }
+
         /// <summary>
         /// Gets the type of the proxy.
         /// </summary>
@@ -200,6 +212,15 @@ protected override void ProcessConnectedInstanceStateChange(MongoServerInstance
             }
         }
 
+        /// <summary>
+        /// Processes the disconnected instance state change.
+        /// </summary>
+        /// <param name="instance">The instance.</param>
+        protected override void ProcessDisconnectedInstanceStateChange(MongoServerInstance instance)
+        {
+            Interlocked.CompareExchange(ref _primary, null, instance);
+        }
+
         // private methods
         /// <summary>
         /// Gets a randomly selected matching instance.
@@ -254,6 +275,7 @@ private MongoServerInstance GetMatchingInstance(List<ConnectedInstanceCollection
 
         private void ProcessConnectedPrimaryStateChange(MongoServerInstance instance)
         {
+            Interlocked.Exchange(ref _primary, instance);
             Interlocked.CompareExchange(ref _replicaSetName, instance.ReplicaSetInformation.Name, null);
 
             var members = instance.ReplicaSetInformation.Members;

MongoDB.Driver/MongoServer.cs

@@ -275,7 +275,36 @@ public virtual MongoServerInstance Primary
         {
             get
             {
-                return _serverProxy.Instances.SingleOrDefault(x => x.IsPrimary);
+                var serverProxy = _serverProxy;
+
+                var discoveringServerProxy = serverProxy as DiscoveringMongoServerProxy;
+                if (discoveringServerProxy != null)
+                {
+                    serverProxy = discoveringServerProxy.WrappedProxy;
+                    if (serverProxy == null)
+                    {
+                        return null;
+                    }
+                }
+
+                var directProxy = serverProxy as DirectMongoServerProxy;
+                if (directProxy != null)
+                {
+                    var instance = directProxy.Instances[0];
+                    if (instance.IsPrimary)
+                    {
+                        return instance;
+                    }
+                    return null;
+                }
+
+                var replicaSetProxy = serverProxy as ReplicaSetMongoServerProxy;
+                if (replicaSetProxy != null)
+                {
+                    return replicaSetProxy.Primary;
+                }
+
+                return null;
             }
         }