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()