PYTHON-1954 Stop holding the topology lock while creating new connections

Author: prashantmital

pymongo/pool.py

@@ -1033,14 +1033,19 @@ def reset(self):
     def close(self):
         self._reset(close=True)
 
-    def remove_stale_sockets(self):
-        """Removes stale sockets then adds new ones if pool is too small."""
+    def remove_stale_sockets(self, reference_pool_id):
+        """Removes stale sockets then adds new ones if pool is too small and
+        has not been reset. The `reference_pool_id` argument specifies the
+        `pool_id` at the point in time this operation was requested on the
+        pool.
+        """
         if self.opts.max_idle_time_seconds is not None:
             with self.lock:
                 while (self.sockets and
                        self.sockets[-1].idle_time_seconds() > self.opts.max_idle_time_seconds):
                     sock_info = self.sockets.pop()
                     sock_info.close_socket(ConnectionClosedReason.IDLE)
+
         while True:
             with self.lock:
                 if (len(self.sockets) + self.active_sockets >=
@@ -1054,6 +1059,11 @@ def remove_stale_sockets(self):
             try:
                 sock_info = self.connect()
                 with self.lock:
+                    # Close connection and return if the pool was reset during
+                    # socket creation or while acquiring the pool lock.
+                    if self.pool_id != reference_pool_id:
+                        sock_info.close_socket()
+                        break
                     self.sockets.appendleft(sock_info)
             finally:
                 self._socket_semaphore.release()

pymongo/topology.py

@@ -429,9 +429,13 @@ def mark_server_unknown_and_request_check(self, address):
 
     def update_pool(self):
         # Remove any stale sockets and add new sockets if pool is too small.
+        servers = []
         with self._lock:
             for server in self._servers.values():
-                server._pool.remove_stale_sockets()
+                servers.append((server, server._pool.pool_id))
+
+        for server, pool_id in servers:
+            server._pool.remove_stale_sockets(pool_id)
 
     def close(self):
         """Clear pools and terminate monitors. Topology reopens on demand."""

test/test_command_monitoring_spec.py

@@ -26,8 +26,7 @@
 from pymongo.errors import OperationFailure
 from pymongo.write_concern import WriteConcern
 from test import unittest, client_context
-from test.utils import single_client, wait_until, EventListener
-from test.utils_selection_tests import parse_read_preference
+from test.utils import single_client, wait_until, EventListener, parse_read_preference
 
 # Location of JSON test specifications.
 _TEST_PATH = os.path.join(

test/test_discovery_and_monitoring.py

@@ -30,43 +30,14 @@
 from pymongo.settings import TopologySettings
 from pymongo.uri_parser import parse_uri
 from test import unittest
+from test.utils import MockPool
 
 
 # Location of JSON test specifications.
 _TEST_PATH = os.path.join(
     os.path.dirname(os.path.realpath(__file__)), 'discovery_and_monitoring')
 
 
-class MockSocketInfo(object):
-    def close(self):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-
-class MockPool(object):
-    def __init__(self, *args, **kwargs):
-        self.pool_id = 0
-        self._lock = threading.Lock()
-
-    def _reset(self):
-        with self._lock:
-            self.pool_id += 1
-
-    def reset(self):
-        self._reset()
-
-    def close(self):
-        self._reset()
-
-    def update_is_writable(self, is_writable):
-        pass
-
-
 class MockMonitor(object):
     def __init__(self, server_description, topology, pool, topology_settings):
         self._server_description = server_description
@@ -81,7 +52,7 @@ def request_check(self):
     def close(self):
         pass
 
-    def remove_stale_sockets(self):
+    def remove_stale_sockets(self, reference_pool_id):
         pass
 
 

test/test_heartbeat_monitoring.py

@@ -22,49 +22,9 @@
 from pymongo.errors import ConnectionFailure
 from pymongo.ismaster import IsMaster
 from pymongo.monitor import Monitor
-from pymongo.pool import PoolOptions
 from test import unittest, client_knobs
-from test.utils import HeartbeatEventListener, single_client, wait_until
-
-
-class MockSocketInfo(object):
-    def close(self):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-
-class MockPool(object):
-    def __init__(self, *args, **kwargs):
-        self.pool_id = 0
-        self._lock = threading.Lock()
-        self.opts = PoolOptions()
-
-    def get_socket(self, all_credentials):
-        return MockSocketInfo()
-
-    def return_socket(self, *args, **kwargs):
-        pass
-
-    def _reset(self):
-        with self._lock:
-            self.pool_id += 1
-
-    def reset(self):
-        self._reset()
-
-    def close(self):
-        self._reset()
-
-    def update_is_writable(self, is_writable):
-        pass
-
-    def remove_stale_sockets(self):
-        pass
+from test.utils import (HeartbeatEventListener, MockPool, single_client,
+                        wait_until)
 
 
 class TestHeartbeatMonitoring(unittest.TestCase):

test/test_topology.py

@@ -37,47 +37,7 @@
                                       writable_server_selector)
 from pymongo.settings import TopologySettings
 from test import client_knobs, unittest
-from test.utils import wait_until
-
-
-class MockSocketInfo(object):
-    def close(self):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-
-class MockPool(object):
-    def __init__(self, *args, **kwargs):
-        self.pool_id = 0
-        self._lock = threading.Lock()
-        self.opts = PoolOptions()
-
-    def get_socket(self, all_credentials):
-        return MockSocketInfo()
-
-    def return_socket(self, *args, **kwargs):
-        pass
-
-    def _reset(self):
-        with self._lock:
-            self.pool_id += 1
-
-    def reset(self):
-        self._reset()
-
-    def close(self):
-        self._reset()
-
-    def update_is_writable(self, is_writable):
-        pass
-
-    def remove_stale_sockets(self):
-        pass
+from test.utils import MockPool, wait_until
 
 
 class MockMonitor(object):

test/utils.py

@@ -32,9 +32,10 @@
 from bson.objectid import ObjectId
 
 from pymongo import (MongoClient,
-                     monitoring)
+                     monitoring, read_preferences)
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.monitoring import _SENSITIVE_COMMANDS, ConnectionPoolListener
+from pymongo.pool import PoolOptions
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_selectors import (any_server_selector,
@@ -44,8 +45,6 @@
 from test import (client_context,
                   db_user,
                   db_pwd)
-from test.utils_selection_tests import parse_read_preference
-
 
 IMPOSSIBLE_WRITE_CONCERN = WriteConcern(w=1000)
 
@@ -185,6 +184,46 @@ def failed(self, event):
         self.results.append(event)
 
 
+class MockSocketInfo(object):
+    def close(self):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+class MockPool(object):
+    def __init__(self, *args, **kwargs):
+        self.pool_id = 0
+        self._lock = threading.Lock()
+        self.opts = PoolOptions()
+
+    def get_socket(self, all_credentials):
+        return MockSocketInfo()
+
+    def return_socket(self, *args, **kwargs):
+        pass
+
+    def _reset(self):
+        with self._lock:
+            self.pool_id += 1
+
+    def reset(self):
+        self._reset()
+
+    def close(self):
+        self._reset()
+
+    def update_is_writable(self, is_writable):
+        pass
+
+    def remove_stale_sockets(self, reference_pool_id):
+        pass
+
+
 class ScenarioDict(dict):
     """Dict that returns {} for any unknown key, recursively."""
     def __init__(self, data):
@@ -811,3 +850,14 @@ def run(self):
         except BaseException as exc:
             self.exc = exc
             raise
+
+
+def parse_read_preference(pref):
+    # Make first letter lowercase to match read_pref's modes.
+    mode_string = pref.get('mode', 'primary')
+    mode_string = mode_string[:1].lower() + mode_string[1:]
+    mode = read_preferences.read_pref_mode_from_name(mode_string)
+    max_staleness = pref.get('maxStalenessSeconds', -1)
+    tag_sets = pref.get('tag_sets')
+    return read_preferences.make_read_preference(
+        mode, tag_sets=tag_sets, max_staleness=max_staleness)

test/utils_selection_tests.py

@@ -21,7 +21,6 @@
 sys.path[0:0] = [""]
 
 from bson import json_util
-from pymongo import read_preferences
 from pymongo.common import clean_node, HEARTBEAT_FREQUENCY
 from pymongo.errors import AutoReconnect, ConfigurationError
 from pymongo.ismaster import IsMaster
@@ -30,34 +29,7 @@
 from pymongo.server_selectors import writable_server_selector
 from pymongo.topology import Topology
 from test import unittest
-
-
-class MockSocketInfo(object):
-    def close(self):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-
-class MockPool(object):
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def reset(self):
-        pass
-
-    def close(self):
-        pass
-
-    def update_is_writable(self, is_writable):
-        pass
-
-    def remove_stale_sockets(self):
-        pass
+from test.utils import MockPool, parse_read_preference
 
 
 class MockMonitor(object):
@@ -288,14 +260,3 @@ class TestAllScenarios(unittest.TestCase):
             setattr(TestAllScenarios, new_test.__name__, new_test)
 
     return TestAllScenarios
-
-
-def parse_read_preference(pref):
-    # Make first letter lowercase to match read_pref's modes.
-    mode_string = pref.get('mode', 'primary')
-    mode_string = mode_string[:1].lower() + mode_string[1:]
-    mode = read_preferences.read_pref_mode_from_name(mode_string)
-    max_staleness = pref.get('maxStalenessSeconds', -1)
-    tag_sets = pref.get('tag_sets')
-    return read_preferences.make_read_preference(
-        mode, tag_sets=tag_sets, max_staleness=max_staleness)

test/utils_spec_runner.py

@@ -46,8 +46,7 @@
                         camel_to_upper_camel,
                         CompareType,
                         OvertCommandListener,
-                        rs_client)
-from test.utils_selection_tests import parse_read_preference
+                        rs_client, parse_read_preference)
 
 
 class SpecRunner(IntegrationTest):