PYTHON-1189 Fix race condition in poll.poll

Each Pool guards its poll.poll with a mutex to prevent concurrent calls.
poll and select are retried if they are interrupted with EINTR. Also fixes
PYTHON-1179.

Author: ShaneHarvey

pymongo/network.py

@@ -18,18 +18,22 @@
 import errno
 import select
 import struct
+import threading
 
 _HAS_POLL = True
-_poller = None
 _EVENT_MASK = 0
 try:
     from select import poll
-    _poller = poll()
     _EVENT_MASK = (select.POLLIN | select.POLLPRI | select.POLLERR |
                    select.POLLHUP | select.POLLNVAL)
 except ImportError:
     _HAS_POLL = False
 
+try:
+    from select import error as _SELECT_ERROR
+except ImportError:
+    _SELECT_ERROR = OSError
+
 from pymongo import helpers, message
 from pymongo.common import MAX_MESSAGE_SIZE
 from pymongo.errors import (AutoReconnect,
@@ -159,12 +163,7 @@ def _receive_data_on_socket(sock, length):
         try:
             chunk = sock.recv(length)
         except (IOError, OSError) as exc:
-            err = None
-            if hasattr(exc, 'errno'):
-                err = exc.errno
-            elif exc.args:
-                err = exc.args[0]
-            if err == errno.EINTR:
+            if _errno_from_exception(exc) == errno.EINTR:
                 continue
             raise
         if chunk == b"":
@@ -176,17 +175,56 @@ def _receive_data_on_socket(sock, length):
     return msg
 
 
-def socket_closed(sock):
-    """Return True if we know socket has been closed, False otherwise.
-    """
-    try:
+def _errno_from_exception(exc):
+    if hasattr(exc, 'errno'):
+        return exc.errno
+    elif exc.args:
+        return exc.args[0]
+    else:
+        return None
+
+
+class SocketChecker(object):
+
+    def __init__(self):
         if _HAS_POLL:
-            _poller.register(sock, _EVENT_MASK)
-            rd = _poller.poll(0)
-            _poller.unregister(sock)
+            self._lock = threading.Lock()
+            self._poller = poll()
         else:
-            rd, _, _ = select.select([sock], [], [], 0)
-    # Any exception here is equally bad (select.error, ValueError, etc.).
-    except:
-        return True
-    return len(rd) > 0
+            self._lock = None
+            self._poller = None
+
+    def socket_closed(self, sock):
+        """Return True if we know socket has been closed, False otherwise.
+        """
+        while True:
+            try:
+                if self._poller:
+                    with self._lock:
+                        self._poller.register(sock, _EVENT_MASK)
+                        try:
+                            rd = self._poller.poll(0)
+                        finally:
+                            self._poller.unregister(sock)
+                else:
+                    rd, _, _ = select.select([sock], [], [], 0)
+            except (RuntimeError, KeyError):
+                # RuntimeError is raised during a concurrent poll. KeyError
+                # is raised by unregister if the socket is not in the poller.
+                # These errors should not be possible since we protect the
+                # poller with a mutex.
+                raise
+            except ValueError:
+                # ValueError is raised by register/unregister/select if the
+                # socket file descriptor is negative or outside the range for
+                # select (> 1023).
+                return True
+            except (_SELECT_ERROR, IOError) as exc:
+                if _errno_from_exception(exc) in (errno.EINTR, errno.EAGAIN):
+                    continue
+                return True
+            except:
+                # Any other exceptions should be attributed to a closed
+                # or invalid socket.
+                return True
+            return len(rd) > 0

pymongo/pool.py

@@ -45,7 +45,7 @@ class SSLError(socket.error):
 from pymongo.monotonic import time as _time
 from pymongo.network import (command,
                              receive_message,
-                             socket_closed)
+                             SocketChecker)
 from pymongo.read_concern import DEFAULT_READ_CONCERN
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_type import SERVER_TYPE
@@ -703,6 +703,7 @@ def __init__(self, address, options, handshake=True):
 
         self._socket_semaphore = thread_util.create_semaphore(
             self.opts.max_pool_size, max_waiters)
+        self.socket_checker = SocketChecker()
 
     def reset(self):
         with self.lock:
@@ -879,7 +880,7 @@ def _check(self, sock_info):
                 and (
                     0 == self._check_interval_seconds
                     or age > self._check_interval_seconds)):
-            if socket_closed(sock_info.sock):
+            if self.socket_checker.socket_closed(sock_info.sock):
                 sock_info.close()
                 error = True
 

test/test_pooling.py

@@ -29,7 +29,7 @@
 
 sys.path[0:0] = [""]
 
-from pymongo.network import socket_closed
+from pymongo.network import SocketChecker
 from pymongo.pool import Pool, PoolOptions
 from test import client_context, unittest
 from test.utils import (get_pool,
@@ -236,7 +236,8 @@ def test_pool_removes_dead_socket(self):
             # Simulate a closed socket without telling the SocketInfo it's
             # closed.
             sock_info.sock.close()
-            self.assertTrue(socket_closed(sock_info.sock))
+            self.assertTrue(
+                cx_pool.socket_checker.socket_closed(sock_info.sock))
 
         with cx_pool.get_socket({}) as new_sock_info:
             self.assertEqual(0, len(cx_pool.sockets))
@@ -251,9 +252,28 @@ def test_pool_removes_dead_socket(self):
     def test_socket_closed(self):
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         s.connect((client_context.host, client_context.port))
-        self.assertFalse(socket_closed(s))
+        socket_checker = SocketChecker()
+        self.assertFalse(socket_checker.socket_closed(s))
         s.close()
-        self.assertTrue(socket_closed(s))
+        self.assertTrue(socket_checker.socket_closed(s))
+
+    def test_socket_closed_thread_safe(self):
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect((client_context.host, client_context.port))
+        socket_checker = SocketChecker()
+
+        def check_socket():
+            for _ in range(1000):
+                self.assertFalse(socket_checker.socket_closed(s))
+
+        threads = []
+        for i in range(3):
+            thread = threading.Thread(target=check_socket)
+            thread.start()
+            threads.append(thread)
+
+        for thread in threads:
+            thread.join()
 
     def test_return_socket_after_reset(self):
         pool = self.create_pool()