PYTHON-829 Call ismaster on each new connection.

ajdavis · ajdavis · commit 9b208caed96b · 2015-01-26T16:54:14.000-05:00
Call ismaster on each new connection and store the results on the SocketInfo
instance.

The upcoming Authentication Spec says: "If credentials exist, upon opening a
socket, drivers MUST send an isMaster command immediately. This allows a driver
to determine whether the server is an Arbiter. Calling ismaster additionally
allows the driver to know if the default authentication method for each socket
is MONGODB-CR or SCRAM-SHA-1, avoiding races when the driver repopulates the
pool after a disconnect."

In theory we could choose not to call ismaster if there are no credentials, but
it's simpler always to call ismaster, and paves the way for future breaking
changes to the wire protocol besides the current breaking change to
authentication.
diff --git a/pymongo/monitor.py b/pymongo/monitor.py
@@ -137,7 +137,7 @@ def _check_once(self):
 
         Returns a ServerDescription, or raises an exception.
         """
-        with self._pool.get_socket({}, 0, 0) as sock_info:
+        with self._pool.get_socket({}) as sock_info:
             response, round_trip_time = self._check_with_socket(sock_info)
             self._avg_round_trip_time.add_sample(round_trip_time)
             sd = ServerDescription(
diff --git a/pymongo/network.py b/pymongo/network.py
@@ -20,6 +20,8 @@
 from pymongo import helpers, message
 from pymongo.errors import AutoReconnect
 
+_UNPACK_INT = struct.Struct("<i").unpack
+
 
 def command(sock, dbname, spec):
     """Execute a command over the socket, or raise socket.error.
@@ -43,15 +45,15 @@ def command(sock, dbname, spec):
 def receive_message(sock, operation, request_id):
     """Receive a raw BSON message or raise socket.error."""
     header = _receive_data_on_socket(sock, 16)
-    length = struct.unpack("<i", header[:4])[0]
+    length = _UNPACK_INT(header[:4])[0]
 
     # No request_id for exhaust cursor "getMore".
     if request_id is not None:
-        response_id = struct.unpack("<i", header[8:12])[0]
+        response_id = _UNPACK_INT(header[8:12])[0]
         assert request_id == response_id, "ids don't match %r %r" % (
             request_id, response_id)
 
-    assert operation == struct.unpack("<i", header[12:])[0]
+    assert operation == _UNPACK_INT(header[12:])[0]
     return _receive_data_on_socket(sock, length - 16)
 
 
diff --git a/pymongo/pool.py b/pymongo/pool.py
@@ -20,6 +20,7 @@
 from bson.py3compat import u, itervalues
 from pymongo import auth, thread_util
 from pymongo.errors import ConnectionFailure
+from pymongo.ismaster import IsMaster
 from pymongo.monotonic import time as _time
 from pymongo.network import (command,
                              receive_message,
@@ -107,16 +108,21 @@ def socket_keepalive(self):
 
 
 class SocketInfo(object):
-    """Store a socket with some metadata."""
-    def __init__(self, sock, pool, host):
+    """Store a socket with some metadata.
+
+    :Parameters:
+      - `sock`: a raw socket object
+      - `pool`: a Pool instance
+      - `ismaster`: an IsMaster instance, response to ismaster call on `sock`
+      - `host`: a string, the server's hostname (without port)
+    """
+    def __init__(self, sock, pool, ismaster, host):
         self.sock = sock
         self.host = host
         self.authset = set()
         self.closed = False
         self.last_checkout = _time()
-
-        self._min_wire_version = None
-        self._max_wire_version = None
+        self.ismaster = ismaster
 
         # The pool's pool_id changes with each reset() so we can close sockets
         # created before the last reset.
@@ -196,19 +202,13 @@ def close(self):
         except:
             pass
         
-    def set_wire_version_range(self, min_wire_version, max_wire_version):
-        self._min_wire_version = min_wire_version
-        self._max_wire_version = max_wire_version
-        
     @property
     def min_wire_version(self):
-        assert self._min_wire_version is not None
-        return self._min_wire_version
+        return self.ismaster.min_wire_version
         
     @property
     def max_wire_version(self):
-        assert self._max_wire_version is not None
-        return self._max_wire_version
+        return self.ismaster.max_wire_version
 
     def __eq__(self, other):
         return self.sock == other.sock
@@ -353,26 +353,24 @@ def connect(self):
            return_socket() when you're done with it.
         """
         sock = _configured_socket(self.address, self.opts)
-        return SocketInfo(sock, self, host=self.address[0])
+        try:
+            ismaster = IsMaster(command(sock, 'admin', {'ismaster': 1}))
+        except:
+            sock.close()
+            raise
+
+        return SocketInfo(sock, self, ismaster, host=self.address[0])
 
     @contextlib.contextmanager
-    def get_socket(
-            self,
-            all_credentials,
-            min_wire_version,
-            max_wire_version,
-            checkout=False):
+    def get_socket(self, all_credentials, checkout=False):
         """Get a socket from the pool. Use with a "with" statement.
 
         Returns a :class:`SocketInfo` object wrapping a connected
         :class:`socket.socket`.
 
         This method should always be used in a with-statement::
 
-            with pool.get_socket(credentials,
-                                 min_wire_version,
-                                 max_wire_version,
-                                 checkout) as socket_info:
+            with pool.get_socket(credentials, checkout) as socket_info:
                 socket_info.send_message(msg)
                 data = socket_info.receive_message(op_code, request_id)
 
@@ -382,14 +380,11 @@ def get_socket(
 
         :Parameters:
           - `all_credentials`: dict, maps auth source to MongoCredential.
-          - `min_wire_version`: int, minimum protocol the server supports.
-          - `max_wire_version`: int, maximum protocol the server supports.
           - `checkout` (optional): keep socket checked out.
         """
         # First get a socket, then attempt authentication. Simplifies
         # semaphore management in the face of network errors during auth.
         sock_info = self._get_socket_no_auth()
-        sock_info.set_wire_version_range(min_wire_version, max_wire_version)
         try:
             sock_info.check_auth(all_credentials)
             yield sock_info
diff --git a/pymongo/server.py b/pymongo/server.py
@@ -103,11 +103,7 @@ def send_message_with_response(
 
     @contextlib.contextmanager
     def get_socket(self, all_credentials, checkout=False):
-        sd = self.description
-        with self.pool.get_socket(all_credentials,
-                                  sd.min_wire_version,
-                                  sd.max_wire_version,
-                                  checkout) as sock_info:
+        with self.pool.get_socket(all_credentials, checkout) as sock_info:
             yield sock_info
 
     @property
diff --git a/test/pymongo_mocks.py b/test/pymongo_mocks.py
@@ -42,12 +42,7 @@ def __init__(self, client, pair, *args, **kwargs):
                       PoolOptions(connect_timeout=20))
 
     @contextlib.contextmanager
-    def get_socket(
-            self,
-            all_credentials,
-            min_wire_version,
-            max_wire_version,
-            checkout=False):
+    def get_socket(self, all_credentials, checkout=False):
         client = self.client
         host_and_port = '%s:%s' % (self.mock_host, self.mock_port)
         if host_and_port in client.mock_down_hosts:
@@ -58,10 +53,7 @@ def get_socket(
             + client.mock_members
             + client.mock_mongoses), "bad host: %s" % host_and_port
 
-        with Pool.get_socket(self,
-                             all_credentials,
-                             min_wire_version,
-                             max_wire_version) as sock_info:
+        with Pool.get_socket(self, all_credentials) as sock_info:
             sock_info.mock_host = self.mock_host
             sock_info.mock_port = self.mock_port
             yield sock_info
diff --git a/test/test_pooling.py b/test/test_pooling.py
@@ -124,7 +124,7 @@ def run_mongo_thread(self):
         self.state = 'get_socket'
 
         # Pass 'checkout' so we can hold the socket.
-        with self.pool.get_socket({}, 0, 0, checkout=True) as sock:
+        with self.pool.get_socket({}, checkout=True) as sock:
             self.sock = sock
 
         self.state = 'sock'
@@ -189,10 +189,10 @@ def test_pool_reuses_open_socket(self):
         # Test Pool's _check_closed() method doesn't close a healthy socket.
         cx_pool = self.create_pool(max_pool_size=10)
         cx_pool._check_interval_seconds = 0  # Always check.
-        with cx_pool.get_socket({}, 0, 0) as sock_info:
+        with cx_pool.get_socket({}) as sock_info:
             pass
 
-        with cx_pool.get_socket({}, 0, 0) as new_sock_info:
+        with cx_pool.get_socket({}) as new_sock_info:
             self.assertEqual(sock_info, new_sock_info)
 
         self.assertEqual(1, len(cx_pool.sockets))
@@ -201,11 +201,11 @@ def test_get_socket_and_exception(self):
         # get_socket() returns socket after a non-network error.
         cx_pool = self.create_pool(max_pool_size=1, wait_queue_timeout=1)
         with self.assertRaises(ZeroDivisionError):
-            with cx_pool.get_socket({}, 0, 0) as sock_info:
+            with cx_pool.get_socket({}) as sock_info:
                 1 / 0
 
         # Socket was returned, not closed.
-        with cx_pool.get_socket({}, 0, 0) as new_sock_info:
+        with cx_pool.get_socket({}) as new_sock_info:
             self.assertEqual(sock_info, new_sock_info)
 
         self.assertEqual(1, len(cx_pool.sockets))
@@ -214,7 +214,7 @@ def test_pool_removes_closed_socket(self):
         # Test that Pool removes explicitly closed socket.
         cx_pool = self.create_pool()
 
-        with cx_pool.get_socket({}, 0, 0) as sock_info:
+        with cx_pool.get_socket({}) as sock_info:
             # Use SocketInfo's API to close the socket.
             sock_info.close()
 
@@ -226,25 +226,25 @@ def test_pool_removes_dead_socket(self):
         cx_pool = self.create_pool(max_pool_size=1, wait_queue_timeout=1)
         cx_pool._check_interval_seconds = 0  # Always check.
 
-        with cx_pool.get_socket({}, 0, 0) as sock_info:
+        with cx_pool.get_socket({}) as sock_info:
             # Simulate a closed socket without telling the SocketInfo it's
             # closed.
             sock_info.sock.close()
             self.assertTrue(socket_closed(sock_info.sock))
 
-        with cx_pool.get_socket({}, 0, 0) as new_sock_info:
+        with cx_pool.get_socket({}) as new_sock_info:
             self.assertEqual(0, len(cx_pool.sockets))
             self.assertNotEqual(sock_info, new_sock_info)
 
         self.assertEqual(1, len(cx_pool.sockets))
 
         # Semaphore was released.
-        with cx_pool.get_socket({}, 0, 0):
+        with cx_pool.get_socket({}):
             pass
 
     def test_return_socket_after_reset(self):
         pool = self.create_pool()
-        with pool.get_socket({}, 0, 0) as sock:
+        with pool.get_socket({}) as sock:
             pool.reset()
 
         self.assertTrue(sock.closed)
@@ -258,20 +258,20 @@ def test_pool_check(self):
                                    wait_queue_timeout=1)
         cx_pool._check_interval_seconds = 0  # Always check.
 
-        with cx_pool.get_socket({}, 0, 0) as sock_info:
+        with cx_pool.get_socket({}) as sock_info:
             # Simulate a closed socket without telling the SocketInfo it's
             # closed.
             sock_info.sock.close()
 
         # Swap pool's address with a bad one.
         address, cx_pool.address = cx_pool.address, ('foo.com', 1234)
         with self.assertRaises(socket.error):
-            with cx_pool.get_socket({}, 0, 0):
+            with cx_pool.get_socket({}):
                 pass
 
         # Back to normal, semaphore was correctly released.
         cx_pool.address = address
-        with cx_pool.get_socket({}, 0, 0, checkout=True):
+        with cx_pool.get_socket({}, checkout=True):
             pass
 
     def test_pool_with_fork(self):
@@ -326,18 +326,18 @@ def loop(pipe):
         self.assertTrue(b_sock != c_sock)
 
         # a_sock, created by parent process, is still in the pool
-        with get_pool(a).get_socket({}, 0, 0) as d_sock:
+        with get_pool(a).get_socket({}) as d_sock:
             self.assertEqual(a_sock, d_sock)
 
     def test_wait_queue_timeout(self):
         wait_queue_timeout = 2  # Seconds
         pool = self.create_pool(
             max_pool_size=1, wait_queue_timeout=wait_queue_timeout)
         
-        with pool.get_socket({}, 0, 0) as sock_info:
+        with pool.get_socket({}) as sock_info:
             start = time.time()
             with self.assertRaises(ConnectionFailure):
-                with pool.get_socket({}, 0, 0):
+                with pool.get_socket({}):
                     pass
 
         duration = time.time() - start
@@ -353,7 +353,7 @@ def test_no_wait_queue_timeout(self):
         pool = self.create_pool(max_pool_size=1)
         
         # Reach max_size.
-        with pool.get_socket({}, 0, 0) as s1:
+        with pool.get_socket({}) as s1:
             t = SocketGetter(self.c, pool)
             t.start()
             while t.state != 'get_socket':
@@ -375,8 +375,8 @@ def test_wait_queue_multiple(self):
             max_pool_size=2, wait_queue_multiple=wait_queue_multiple)
 
         # Reach max_size sockets.
-        with pool.get_socket({}, 0, 0):
-            with pool.get_socket({}, 0, 0):
+        with pool.get_socket({}):
+            with pool.get_socket({}):
 
                 # Reach max_size * wait_queue_multiple waiters.
                 threads = []
@@ -390,7 +390,7 @@ def test_wait_queue_multiple(self):
                     self.assertEqual(t.state, 'get_socket')
 
                 with self.assertRaises(ExceededMaxWaiters):
-                    with pool.get_socket({}, 0, 0):
+                    with pool.get_socket({}):
                         pass
 
     def test_no_wait_queue_multiple(self):
@@ -399,7 +399,7 @@ def test_no_wait_queue_multiple(self):
         socks = []
         for _ in range(2):
             # Pass 'checkout' so we can hold the socket.
-            with pool.get_socket({}, 0, 0, checkout=True) as sock:
+            with pool.get_socket({}, checkout=True) as sock:
                 socks.append(sock)
 
         threads = []
@@ -498,7 +498,7 @@ def test_max_pool_size_with_connection_failure(self):
         # socket from pool" instead of the socket.error.
         for i in range(2):
             with self.assertRaises(socket.error):
-                with test_pool.get_socket({}, 0, 0, checkout=True):
+                with test_pool.get_socket({}, checkout=True):
                     pass
 
 
diff --git a/test/test_topology.py b/test/test_topology.py
@@ -56,7 +56,7 @@ def __init__(self, *args, **kwargs):
         self.pool_id = 0
         self._lock = threading.Lock()
 
-    def get_socket(self, all_credentials, min_wire_version, max_wire_version):
+    def get_socket(self, all_credentials):
         return MockSocketInfo()
 
     def return_socket(self, _):
