Implemented alternative socket pooling to support greenlets (gevent and eventlet). This is entirely based on mgood's fork, commit 2388960.

Author: jbenet

pymongo/connection.py

@@ -38,7 +38,7 @@
 import select
 import socket
 import struct
-import threading
+import thread
 import time
 import warnings
 
@@ -60,6 +60,21 @@
 _CONNECT_TIMEOUT = 20.0
 
 
+try:
+    from greenlet import greenlet
+except ImportError:
+    def _thread_identifier():
+        """Return the identifier of the current thread-of-execution."""
+        return os.getpid(), thread.get_ident()
+else:
+    def _thread_identifier():
+        """Return the identifier of the current thread-of-execution.
+        Supports greenlets.
+        """
+        return os.getpid(), thread.get_ident(), greenlet.getcurrent()
+
+
+
 def _closed(sock):
     """Return True if we know socket has been closed, False otherwise.
     """
@@ -84,25 +99,20 @@ def _partition_node(node):
     return host, port
 
 
-class _Pool(threading.local):
+
+class _Pool(object):
     """A simple connection pool.
 
-    Uses thread-local socket per thread. By calling return_socket() a
-    thread can return a socket to the pool.
+    Uses thread-local socket per thread (including greenlets).
+    By calling return_socket() a thread can return a socket to the pool.
     """
 
-    # Non thread-locals
-    __slots__ = ["sockets", "pool_size", "pid"]
-
-    # thread-local default
-    sock = None
-
     def __init__(self, pool_size, network_timeout):
         self.pid = os.getpid()
         self.pool_size = pool_size
         self.network_timeout = network_timeout
-        if not hasattr(self, "sockets"):
-            self.sockets = []
+        self.sockets = []
+        self.active_sockets = {}
 
     def connect(self, host, port):
         """Connect to Mongo and return a new (connected) socket.
@@ -126,36 +136,30 @@ def connect(self, host, port):
             return s
 
     def get_socket(self, host, port):
-        # We use the pid here to avoid issues with fork / multiprocessing.
+        # We use the _thread_identifier here to avoid issues with multiple
+        # threads of execution (processes, proper threads, greenlets)
         # See test.test_connection:TestConnection.test_fork for an example of
         # what could go wrong otherwise
-        pid = os.getpid()
-
-        if pid != self.pid:
-            self.sock = None
-            self.sockets = []
-            self.pid = pid
-
-        if self.sock is not None and self.sock[0] == pid:
-            return self.sock[1]
+        sock_id = _thread_identifier()
 
         try:
-            self.sock = (pid, self.sockets.pop())
-        except IndexError:
-            self.sock = (pid, self.connect(host, port))
-
-        return self.sock[1]
+            sock = self.active_sockets[sock_id]
+        except KeyError:
+            try:
+                sock = self.sockets.pop()
+            except IndexError:
+                sock = self.connect(host, port)
+            self.active_sockets[sock_id] = sock
+        return sock
 
     def return_socket(self):
-        if self.sock is not None and self.sock[0] == os.getpid():
+        sock = self.active_sockets.pop(_thread_identifier(), None)
+        if sock is not None:
             # There's a race condition here, but we deliberately
             # ignore it.  It means that if the pool_size is 10 we
             # might actually keep slightly more than that.
             if len(self.sockets) < self.pool_size:
-                self.sockets.append(self.sock[1])
-            else:
-                self.sock[1].close()
-        self.sock = None
+                self.sockets.append(sock)
 
 
 class Connection(common.BaseObject):  # TODO support auth for pooling

test/test_pooling.py

@@ -159,13 +159,16 @@ def test_no_disconnect(self):
     def test_simple_disconnect(self):
         self.c.test.stuff.find()
         self.assertEqual(0, len(self.c._Connection__pool.sockets))
-        self.assertNotEqual(None, self.c._Connection__pool.sock)
+        if hasattr(self.c._Connection__pool, 'sock'):
+          self.assertNotEqual(None, self.c._Connection__pool.sock)
         self.c.end_request()
         self.assertEqual(1, len(self.c._Connection__pool.sockets))
-        self.assertEqual(None, self.c._Connection__pool.sock)
+        if hasattr(self.c._Connection__pool, 'sock'):
+          self.assertEqual(None, self.c._Connection__pool.sock)
         self.c.disconnect()
         self.assertEqual(0, len(self.c._Connection__pool.sockets))
-        self.assertEqual(None, self.c._Connection__pool.sock)
+        if hasattr(self.c._Connection__pool, 'sock'):
+          self.assertEqual(None, self.c._Connection__pool.sock)
 
     def test_disconnect(self):
         run_cases(self, [SaveAndFind, Disconnect, Unique])
@@ -271,6 +274,43 @@ def loop(pipe):
         self.assert_(b_sock != c_sock)
         self.assertEqual(a_sock, a._Connection__pool.get_socket(a.host, a.port))
 
+    def test_pool_with_greenlets(self):
+        try:
+            from greenlet import greenlet
+        except ImportError:
+            raise SkipTest()
+
+        c = get_connection()
+        c.test.test.find_one()
+        c.end_request()
+        self.assertEqual(1, len(c._Connection__pool.sockets))
+        a_sock = c._Connection__pool.sockets[0]
+
+        def loop(name, pipe):
+            c.test.test.find_one()
+            self.assertEqual(0, len(c._Connection__pool.sockets))
+            greenlet.getcurrent().parent.switch()
+            c.end_request()
+            pipe.append(c._Connection__pool.sockets[-1])
+
+        ga1 = []
+        ga2 = []
+
+        g1 = greenlet(loop)
+        g2 = greenlet(loop)
+
+        g1.switch('g1', ga1)
+        g2.switch('g2', ga2)
+        g1.switch()
+        g2.switch()
+
+        b_sock = ga1[0]
+        c_sock = ga2[0]
+        self.assert_(a_sock is b_sock)
+        self.assert_(a_sock is not c_sock)
+        self.assert_(b_sock is not c_sock)
+
+
     def test_max_pool_size(self):
         c = get_connection(max_pool_size=4)