PYTHON-2395 Consider connection pool health during server selection

pymongo/pool.py

@@ -1111,15 +1111,20 @@ def __init__(self, address, options, handshake=True):
         if self.enabled_for_cmap:
             self.opts.event_listeners.publish_pool_created(
                 self.address, self.opts.non_default_options)
+        # Similar to active_sockets but includes threads in the wait queue.
+        self.operation_count = 0
 
     def _reset(self, close):
         with self.lock:
             if self.closed:
                 return
             self.generation += 1
-            self.pid = os.getpid()
+            newpid = os.getpid()
+            if self.pid != newpid:
+                self.pid = newpid
+                self.active_sockets = 0
+                self.operation_count = 0
             sockets, self.sockets = self.sockets, collections.deque()
-            self.active_sockets = 0
             if close:
                 self.closed = True
 
@@ -1300,6 +1305,9 @@ def _get_socket(self, all_credentials):
                 'Attempted to check out a connection from closed connection '
                 'pool')
 
+        with self.lock:
+            self.operation_count += 1
+
         # Get a free socket or create one.
         if self.opts.wait_queue_timeout:
             deadline = _time() + self.opts.wait_queue_timeout
@@ -1396,6 +1404,7 @@ def return_socket(self, sock_info):
         self._socket_semaphore.release()
         with self.lock:
             self.active_sockets -= 1
+            self.operation_count -= 1
 
     def _perished(self, sock_info):
         """Return True and close the connection if it is "perished".

pymongo/topology.py

@@ -238,9 +238,15 @@ def select_server(self,
                       server_selection_timeout=None,
                       address=None):
         """Like select_servers, but choose a random server if several match."""
-        return random.choice(self.select_servers(selector,
-                                                 server_selection_timeout,
-                                                 address))
+        servers = self.select_servers(
+            selector, server_selection_timeout, address)
+        if len(servers) == 1:
+            return servers[0]
+        server1, server2 = random.sample(servers, 2)
+        if server1.pool.operation_count <= server2.pool.operation_count:
+            return server1
+        else:
+            return server2
 
     def select_server_by_address(self, address,
                                  server_selection_timeout=None):

test/server_selection/in_window/equilibrium.json

@@ -0,0 +1,46 @@
+{
+  "description": "When in equilibrium selection is evenly distributed",
+  "topology_description": {
+    "type": "Sharded",
+    "servers": [
+      {
+        "address": "a:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "b:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "c:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      }
+    ]
+  },
+  "mocked_topology_state": [
+    {
+      "address": "a:27017",
+      "operation_count": 5
+    },
+    {
+      "address": "b:27017",
+      "operation_count": 5
+    },
+    {
+      "address": "c:27017",
+      "operation_count": 5
+    }
+  ],
+  "iterations": 2000,
+  "outcome": {
+    "tolerance": 0.05,
+    "expected_frequencies": {
+      "a:27017": 0.33,
+      "b:27017": 0.33,
+      "c:27017": 0.33
+    }
+  }
+}

test/server_selection/in_window/many-choices.json

@@ -0,0 +1,106 @@
+{
+  "description": "Selections from many choices occur at correct frequencies",
+  "topology_description": {
+    "type": "Sharded",
+    "servers": [
+      {
+        "address": "a:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "b:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "c:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "d:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "e:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "f:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "g:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "h:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "i:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      }
+    ]
+  },
+  "mocked_topology_state": [
+    {
+      "address": "a:27017",
+      "operation_count": 0
+    },
+    {
+      "address": "b:27017",
+      "operation_count": 5
+    },
+    {
+      "address": "c:27017",
+      "operation_count": 5
+    },
+    {
+      "address": "d:27017",
+      "operation_count": 10
+    },
+    {
+      "address": "e:27017",
+      "operation_count": 10
+    },
+    {
+      "address": "f:27017",
+      "operation_count": 20
+    },
+    {
+      "address": "g:27017",
+      "operation_count": 20
+    },
+    {
+      "address": "h:27017",
+      "operation_count": 50
+    },
+    {
+      "address": "i:27017",
+      "operation_count": 60
+    }
+  ],
+  "iterations": 10000,
+  "outcome": {
+    "tolerance": 0.03,
+    "expected_frequencies": {
+      "a:27017": 0.22,
+      "b:27017": 0.18,
+      "c:27017": 0.18,
+      "d:27017": 0.125,
+      "e:27017": 0.125,
+      "f:27017": 0.074,
+      "g:27017": 0.074,
+      "h:27017": 0.0277,
+      "i:27017": 0
+    }
+  }
+}

test/server_selection/in_window/one-least-two-tied.json

@@ -0,0 +1,46 @@
+{
+  "description": "Least operations gets most selections, two tied share the rest",
+  "topology_description": {
+    "type": "Sharded",
+    "servers": [
+      {
+        "address": "a:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "b:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      },
+      {
+        "address": "c:27017",
+        "avg_rtt_ms": 35,
+        "type": "Mongos"
+      }
+    ]
+  },
+  "mocked_topology_state": [
+    {
+      "address": "a:27017",
+      "operation_count": 16
+    },
+    {
+      "address": "b:27017",
+      "operation_count": 10
+    },
+    {
+      "address": "c:27017",
+      "operation_count": 16
+    }
+  ],
+  "iterations": 2000,
+  "outcome": {
+    "tolerance": 0.05,
+    "expected_frequencies": {
+      "a:27017": 0.165,
+      "b:27017": 0.66,
+      "c:27017": 0.165
+    }
+  }
+}

test/server_selection/in_window/rs-equilibrium.json

@@ -0,0 +1,46 @@
+{
+  "description": "When in equilibrium selection is evenly distributed (replica set)",
+  "topology_description": {
+    "type": "ReplicaSetWithPrimary",
+    "servers": [
+      {
+        "address": "a:27017",
+        "avg_rtt_ms": 35,
+        "type": "RSPrimary"
+      },
+      {
+        "address": "b:27017",
+        "avg_rtt_ms": 35,
+        "type": "RSSecondary"
+      },
+      {
+        "address": "c:27017",
+        "avg_rtt_ms": 35,
+        "type": "RSSecondary"
+      }
+    ]
+  },
+  "mocked_topology_state": [
+    {
+      "address": "a:27017",
+      "operation_count": 6
+    },
+    {
+      "address": "b:27017",
+      "operation_count": 6
+    },
+    {
+      "address": "c:27017",
+      "operation_count": 6
+    }
+  ],
+  "iterations": 2000,
+  "outcome": {
+    "tolerance": 0.05,
+    "expected_frequencies": {
+      "a:27017": 0.33,
+      "b:27017": 0.33,
+      "c:27017": 0.33
+    }
+  }
+}

test/server_selection/in_window/rs-three-choices.json

@@ -0,0 +1,46 @@
+{
+  "description": "Selections from three servers occur at proper distributions (replica set)",
+  "topology_description": {
+    "type": "ReplicaSetWithPrimary",
+    "servers": [
+      {
+        "address": "a:27017",
+        "avg_rtt_ms": 35,
+        "type": "RSPrimary"
+      },
+      {
+        "address": "b:27017",
+        "avg_rtt_ms": 35,
+        "type": "RSSecondary"
+      },
+      {
+        "address": "c:27017",
+        "avg_rtt_ms": 35,
+        "type": "RSSecondary"
+      }
+    ]
+  },
+  "mocked_topology_state": [
+    {
+      "address": "a:27017",
+      "operation_count": 3
+    },
+    {
+      "address": "b:27017",
+      "operation_count": 6
+    },
+    {
+      "address": "c:27017",
+      "operation_count": 20
+    }
+  ],
+  "iterations": 2000,
+  "outcome": {
+    "tolerance": 0.05,
+    "expected_frequencies": {
+      "a:27017": 0.66,
+      "b:27017": 0.33,
+      "c:27017": 0
+    }
+  }
+}