Initial commit of ReplicaSetConnection PYTHON-196

behackett · behackett · commit 375f75adff93 · 2011-10-02T19:56:00.000-07:00
This is alpha level code and should not be used in
production. Issues that still need attention:

- Auto-authenticate existing sockets in all pools.
- A host of failover issues.
- A lot of error handling.
- Unit tests...
- More documentation and examples.
- Possibly a background thread to maintain the set.
diff --git a/doc/api/pymongo/index.rst b/doc/api/pymongo/index.rst
@@ -23,6 +23,8 @@ Sub-modules:
    errors
    master_slave_connection
    message
+   pool
+   replica_set_connection
    son_manipulator
    cursor_manager
 
diff --git a/doc/api/pymongo/pool.rst b/doc/api/pymongo/pool.rst
@@ -0,0 +1,6 @@
+:mod:`pool` -- Pool module for use with a MongoDB connection.
+==============================================================
+
+.. automodule:: pymongo.pool
+   :synopsis: Pool module for use with a MongoDB connection.
+   :members:
diff --git a/doc/api/pymongo/replica_set_connection.rst b/doc/api/pymongo/replica_set_connection.rst
@@ -0,0 +1,33 @@
+:mod:`replica_set_connection` -- Tools for connecting to a MongoDB replica set
+==============================================================================
+
+.. automodule:: pymongo.replica_set_connection
+   :synopsis: Tools for connecting to a MongoDB replica set
+
+   .. autoclass:: pymongo.replica_set_connection.ReplicaSetConnection([hosts_or_uri[, max_pool_size=10[, document_class=dict[, tz_aware=False[, **kwargs]]]]])
+
+      .. automethod:: disconnect
+      .. automethod:: close
+
+      .. describe:: c[db_name] || c.db_name
+
+         Get the `db_name` :class:`~pymongo.database.Database` on :class:`ReplicaSetConnection` `c`.
+
+         Raises :class:`~pymongo.errors.InvalidName` if an invalid database name is used.
+
+      .. autoattribute:: seeds
+      .. autoattribute:: hosts
+      .. autoattribute:: arbiters
+      .. autoattribute:: read_preference
+      .. autoattribute:: max_pool_size
+      .. autoattribute:: document_class
+      .. autoattribute:: tz_aware
+      .. autoattribute:: slave_okay
+      .. autoattribute:: safe
+      .. automethod:: get_lasterror_options
+      .. automethod:: set_lasterror_options
+      .. automethod:: unset_lasterror_options
+      .. automethod:: database_names
+      .. automethod:: drop_database
+      .. automethod:: copy_database(from_name, to_name[, from_host=None[, username=None[, password=None]]])
+      .. automethod:: close_cursor
diff --git a/pymongo/__init__.py b/pymongo/__init__.py
@@ -14,7 +14,6 @@
 
 """Python driver for MongoDB."""
 
-from pymongo.connection import Connection as PyMongo_Connection
 
 ASCENDING = 1
 """Ascending sort order."""
@@ -38,6 +37,13 @@
 ALL = 2
 """Profile all operations."""
 
+PRIMARY = 1
+"""Send all reads to the primary in a ReplicaSetConnction."""
+SECONDARY = 2
+"""Send reads to all active members of a ReplicaSetConnection."""
+SECONDARY_ONLY = 3
+"""Send all reads to secondaries in a ReplicaSetConnection."""
+
 version_tuple = (2, 0, 1, '+')
 
 def get_version_string():
@@ -48,9 +54,8 @@ def get_version_string():
 version = get_version_string()
 """Current version of PyMongo."""
 
-Connection = PyMongo_Connection
-"""Alias for :class:`pymongo.connection.Connection`."""
-
+from pymongo.connection import Connection
+from pymongo.replica_set_connection import ReplicaSetConnection
 
 def has_c():
     """Is the C extension installed?
diff --git a/pymongo/pool.py b/pymongo/pool.py
@@ -0,0 +1,100 @@
+# Copyright 2009-2011 10gen, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you
+# may not use this file except in compliance with the License.  You
+# may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.  See the License for the specific language governing
+# permissions and limitations under the License.
+
+import os
+import socket
+import threading
+
+
+class Pool(threading.local):
+    """A simple connection pool.
+
+    Uses thread-local socket per thread. By calling return_socket() a
+    thread can return a socket to the pool.
+    """
+
+    # Non thread-locals
+    __slots__ = ["pid", "host", "max_size",
+                 "net_timeout", "conn_timeout", "sockets"]
+
+    # thread-local default
+    sock = None
+
+    def __init__(self, host, max_size, net_timeout, conn_timeout):
+        self.pid = os.getpid()
+        self.host = host
+        self.max_size = max_size
+        self.net_timeout = net_timeout
+        self.conn_timeout = conn_timeout
+        if not hasattr(self, "sockets"):
+            self.sockets = []
+
+    def connect(self):
+        """Connect to Mongo and return a new (connected) socket.
+        """
+        try:
+            # Prefer IPv4. If there is demand for an option
+            # to specify one or the other we can add it later.
+            s = socket.socket(socket.AF_INET)
+            s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+            s.settimeout(self.conn_timeout or 20.0)
+            s.connect(self.host)
+            s.settimeout(self.net_timeout)
+            return s
+        except socket.gaierror:
+            # If that fails try IPv6
+            s = socket.socket(socket.AF_INET6)
+            s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+            s.settimeout(self.conn_timeout or 20.0)
+            s.connect(self.host)
+            s.settimeout(self.net_timeout)
+            return s
+
+    def get_socket(self):
+        """Get a socket from the pool. Returns a new socket if the
+        pool is empty.
+        """
+        # We use the pid here to avoid issues with fork / multiprocessing.
+        # 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], True)
+
+        try:
+            self.sock = (pid, self.sockets.pop())
+            return (self.sock[1], True)
+        except IndexError:
+            self.sock = (pid, self.connect())
+            return (self.sock[1], False)
+
+    def return_socket(self):
+        """Return the socket currently in use to the pool. If the
+        pool is full the socket will be discarded.
+        """
+        if self.sock is not None and self.sock[0] == os.getpid():
+            # 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.max_size:
+                self.sockets.append(self.sock[1])
+            else:
+                self.sock[1].close()
+        self.sock = None
diff --git a/pymongo/replica_set_connection.py b/pymongo/replica_set_connection.py
