move kill_cursor stuff to connection level. add CursorManager class to handle killing cursors after they've been closed. default manager kills immediately 1 at a time. also add batch manager which is used by the 'enhanced' driver

connection.py

@@ -7,6 +7,7 @@ import traceback
 
 from errors import ConnectionFailure, InvalidName
 from database import Database
+from cursor_manager import CursorManager
 
 class Connection(object):
     """A connection to Mongo.
@@ -30,9 +31,24 @@ class Connection(object):
         self.__host = host
         self.__port = port
         self.__id = 1
+        self.__cursor_manager = CursorManager(self)
 
         self.__connect()
 
+    def set_cursor_manager(self, manager_class):
+        """Set this connections cursor manager.
+
+        Raises TypeError if manager_class is not a subclass of CursorManager. A
+        cursor manager handles closing cursors. Different managers can implement
+        different policies in terms of when to actually kill a cursor that has
+        been closed.
+        """
+        manager = manager_class(self)
+        if not isinstance(manager, CursorManager):
+            raise TypeError("manager_class must be a subclass of CursorManager")
+
+        self.__cursor_manager = manager
+
     def host(self):
         """Get the connection host.
         """
@@ -136,6 +152,37 @@ class Connection(object):
         """
         return self.__getattr__(name)
 
+    def close_cursor(self, cursor_id):
+        """Close a single database cursor.
+
+        Raises TypeError if cursor_id is not an instance of (int, long). What
+        closing the cursor actually means depends on this connection's cursor
+        manager.
+
+        Arguments:
+        - `cursor_id`: cursor id to close
+        """
+        if not isinstance(cursor_id, (types.IntType, types.LongType)):
+            raise TypeError("cursor_id must be an instance of (int, long)")
+
+        self.__cursor_manager.close(cursor_id)
+
+    def kill_cursors(self, cursor_ids):
+        """Kill database cursors with the given ids.
+
+        Raises TypeError if cursor_ids is not an instance of list.
+
+        Arguments:
+        - `cursor_ids`: list of cursor ids to kill
+        """
+        if not isinstance(cursor_ids, types.ListType):
+            raise TypeError("cursor_ids must be a list")
+        message = "\x00\x00\x00\x00"
+        message += struct.pack("<i", len(cursor_ids))
+        for cursor_id in cursor_ids:
+            message += struct.pack("<q", cursor_id)
+        self.send_message(2007, message)
+
     # TODO implement and test this
     def database_names(self):
         """Get a list of all database names.

cursor.py

@@ -32,9 +32,9 @@ class Cursor(object):
             self.__die()
 
     def __die(self):
-        """Kills this cursor.
+        """Closes this cursor.
         """
-        self.__collection.database()._kill_cursor(self.__id)
+        self.__collection.database().connection().close_cursor(self.__id)
         self.__killed = True
 
     def __query_spec(self):

cursor_manager.py

@@ -0,0 +1,71 @@
+"""Different managers to handle when cursors are killed after they are closed.
+
+New cursor managers should be defined as subclasses of CursorManager and can be
+installed on a connection by calling `Connection.set_cursor_manager`."""
+
+import types
+
+class CursorManager(object):
+    """The default cursor manager.
+
+    This manager will kill cursors one at a time as they are closed.
+    """
+    def __init__(self, connection):
+        """Instantiate the manager.
+
+        Arguments:
+        - `connection`: a Mongo Connection
+        """
+        self.__connection = connection
+
+    def close(self, cursor_id):
+        """Close a cursor by killing it immediately.
+
+        Raises TypeError if cursor_id is not an instance of (int, long).
+
+        Arguments:
+        - `cursor_id`: cursor id to close
+        """
+        if not isinstance(cursor_id, (types.IntType, types.LongType)):
+            raise TypeError("cursor_id must be an instance of (int, long)")
+
+        print "killing cursors"
+        self.__connection.kill_cursors([cursor_id])
+
+class BatchCursorManager(CursorManager):
+    """A cursor manager that kills cursors in batches.
+    """
+    def __init__(self, connection):
+        """Instantiate the manager.
+
+        Arguments:
+        - `connection`: a Mongo Connection
+        """
+        self.__dying_cursors = []
+        self.__max_dying_cursors = 20
+        self.__connection = connection
+
+        CursorManager.__init__(self, connection)
+
+    def __del__(self):
+        """Cleanup - be sure to kill any outstanding cursors.
+        """
+        self.__connection.kill_cursors(self.__dying_cursors)
+
+    def close(self, cursor_id):
+        """Close a cursor by killing it in a batch.
+
+        Raises TypeError if cursor_id is not an instance of (int, long).
+
+        Arguments:
+        - `cursor_id`: cursor id to close
+        """
+        if not isinstance(cursor_id, (types.IntType, types.LongType)):
+            raise TypeError("cursor_id must be an instance of (int, long)")
+
+        self.__dying_cursors.append(cursor_id)
+
+        if len(self.__dying_cursors) > self.__max_dying_cursors:
+            print "killing cursors"
+            self.__connection.kill_cursors(self.__dying_cursors)
+            self.__dying_cursors = []

mongo.py

@@ -12,8 +12,7 @@ from connection import Connection
 from son import SON
 from objectid import ObjectId
 from dbref import DBRef
-
-_MAX_DYING_CURSORS = 20
+from cursor_manager import BatchCursorManager
 
 ASCENDING = 1
 DESCENDING = -1
@@ -45,25 +44,13 @@ class Mongo(Database):
         if not isinstance(settings, types.DictType):
             raise TypeError("settings must be an instance of dict")
 
-        self.__dying_cursors = []
         self.__auto_dereference = settings.get("auto_dereference", False)
         self.__auto_reference = settings.get("auto_reference", False)
 
-        Database.__init__(self, Connection(host, port), name)
+        connection = Connection(host, port)
+        connection.set_cursor_manager(BatchCursorManager)
 
-    def _kill_cursors(self):
-        message = "\x00\x00\x00\x00"
-        message += struct.pack("<i", len(self.__dying_cursors))
-        for cursor_id in self.__dying_cursors:
-            message += struct.pack("<q", cursor_id)
-        self.connection().send_message(2007, message)
-        self.__dying_cursors = []
-
-    def _kill_cursor(self, cursor_id):
-        self.__dying_cursors.append(cursor_id)
-
-        if len(self.__dying_cursors) > _MAX_DYING_CURSORS:
-            self._kill_cursors()
+        Database.__init__(self, connection, name)
 
     def __repr__(self):
         return "Mongo(%r, %r, %r)" % (self.name(), self.connection().host(), self.connection().port())

test/test_mongo.py

@@ -10,7 +10,7 @@ from dbref import DBRef
 from son import SON
 from errors import InvalidOperation, ConnectionFailure
 from collection import SYSTEM_INDEX_COLLECTION
-from mongo import Mongo, ASCENDING, DESCENDING, _MAX_DYING_CURSORS
+from mongo import Mongo, ASCENDING, DESCENDING
 
 class TestMongo(unittest.TestCase):
     def setUp(self):
@@ -124,7 +124,7 @@ class TestMongo(unittest.TestCase):
         self.assertEqual(1000, count)
 
         # test that kill cursors doesn't assert or anything
-        for _ in xrange(3 * _MAX_DYING_CURSORS + 2):
+        for _ in xrange(62):
             for _ in db.test.find():
                 break