PYTHON-1408 Cursor iteration should complete when another thread closes the cursor.

Closing a cursor should not raise an error when killCursors fails.

Author: ShaneHarvey

pymongo/command_cursor.py

@@ -71,7 +71,9 @@ def __del__(self):
  def __die(self, synchronous=False):
  """Closes this cursor.
  """
- if self.__id and not self.__killed:
+ already_killed = self.__killed
+ self.__killed = True
+ if self.__id and not already_killed:
  address = _CursorAddress(
  self.__address, self.__collection.full_name)
  if synchronous:
@@ -81,7 +83,6 @@ def __die(self, synchronous=False):
  # The cursor will be closed later in a different session.
  self.__collection.database.client.close_cursor(
  self.__id, address)
- self.__killed = True
  self.__end_session(synchronous)
 
  def __end_session(self, synchronous):

pymongo/cursor.py

@@ -293,7 +293,9 @@ def _clone_base(self, session):
  def __die(self, synchronous=False):
  """Closes this cursor.
  """
- if self.__id and not self.__killed:
+ already_killed = self.__killed
+ self.__killed = True
+ if self.__id and not already_killed:
  if self.__exhaust and self.__exhaust_mgr:
  # If this is an exhaust cursor and we haven't completely
  # exhausted the result set we *must* close the socket
@@ -311,7 +313,6 @@ def __die(self, synchronous=False):
  self.__id, address)
  if self.__exhaust and self.__exhaust_mgr:
  self.__exhaust_mgr.close()
- self.__killed = True
  if self.__session and not self.__explicit_session:
  self.__session._end_session(lock=synchronous)
  self.__session = None

pymongo/mongo_client.py

@@ -57,10 +57,10 @@
  ConfigurationError,
  ConnectionFailure,
  InvalidOperation,
- InvalidURI,
  NetworkTimeout,
  NotMasterError,
  OperationFailure,
+ PyMongoError,
  ServerSelectionTimeoutError)
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_selectors import (writable_preferred_server_selector,
@@ -1150,8 +1150,12 @@ def _close_cursor_now(self, cursor_id, address=None, session=None):
  if self.__cursor_manager is not None:
  self.__cursor_manager.close(cursor_id, address)
  else:
- self._kill_cursors(
- [cursor_id], address, self._get_topology(), session)
+ try:
+ self._kill_cursors(
+ [cursor_id], address, self._get_topology(), session)
+ except PyMongoError:
+ # Make another attempt to kill the cursor later.
+ self.__kill_cursors_queue.append((address, [cursor_id]))
 
  def kill_cursors(self, cursor_ids, address=None):
  """DEPRECATED - Send a kill cursors message soon with the given ids.

test/test_change_stream.py

@@ -169,6 +169,21 @@ def test_aggregate_cursor_blocks(self):
  self.assertEqual(changes[0]['operationType'], 'insert')
  self.assertEqual(changes[0]['fullDocument'], inserted_doc)
 
+ def test_concurrent_close(self):
+ """Ensure a ChangeStream can be closed from another thread."""
+ # Use a short await time to speed up the test.
+ with self.coll.watch(max_await_time_ms=250) as change_stream:
+ def iterate_cursor():
+ for change in change_stream:
+ pass
+ t = threading.Thread(target=iterate_cursor)
+ t.start()
+ self.coll.insert_one({})
+ time.sleep(1)
+ change_stream.close()
+ t.join(3)
+ self.assertFalse(t.is_alive())
+
  def test_update_resume_token(self):
  """ChangeStream must continuously track the last seen resumeToken."""
  with self.coll.watch() as change_stream:

test/test_cursor.py

@@ -19,6 +19,8 @@
 import random
 import re
 import sys
+import time
+import threading
 
 sys.path[0:0] = [""]
 
@@ -1130,6 +1132,27 @@ def test_tailable(self):
  cursor.rewind()
  self.assertEqual([4, 5, 6], [doc["x"] for doc in cursor[0:3]])
 
+ def test_concurrent_close(self):
+ """Ensure a tailable can be closed from another thread."""
+ db = self.db
+ db.drop_collection("test")
+ db.create_collection("test", capped=True, size=1000, max=3)
+ self.addCleanup(db.drop_collection, "test")
+ cursor = db.test.find(cursor_type=CursorType.TAILABLE)
+
+ def iterate_cursor():
+ while cursor.alive:
+ for doc in cursor:
+ pass
+ t = threading.Thread(target=iterate_cursor)
+ t.start()
+ time.sleep(1)
+ cursor.close()
+ self.assertFalse(cursor.alive)
+ t.join(3)
+ self.assertFalse(t.is_alive())
+
+
  def test_distinct(self):
  self.db.drop_collection("test")