PYTHON-978 - Use killCursors command

Author: aherlihy

pymongo/mongo_client.py

@@ -937,35 +937,48 @@ def _process_kill_cursors_queue(self):
  server = topology.select_server(
  writable_server_selector)
 
- if publish:
- start = datetime.datetime.now()
- data = message.kill_cursors(cursor_ids)
- if publish:
- duration = datetime.datetime.now() - start
- try:
- dbname, collname = address.namespace.split(".", 1)
- except AttributeError:
- dbname = collname = 'OP_KILL_CURSORS'
- command = SON([('killCursors', collname),
- ('cursors', cursor_ids)])
- listeners.publish_command_start(
- command, dbname, data[0], address)
- start = datetime.datetime.now()
  try:
- server.send_message(data, self.__all_credentials)
- except Exception as exc:
- if publish:
- dur = (datetime.datetime.now() - start) + duration
- listeners.publish_command_failure(
- dur, message._convert_exception(exc),
- 'killCursors', data[0], address)
- raise
- if publish:
- duration = (datetime.datetime.now() - start) + duration
- # OP_KILL_CURSORS returns no reply, fake one.
- reply = {'cursorsUnknown': cursor_ids, 'ok': 1}
- listeners.publish_command_success(
- duration, reply, 'killCursors', data[0], address)
+ namespace = address.namespace
+ db, coll = namespace.split('.', 1)
+ except AttributeError:
+ namespace = None
+ db = coll = "OP_KILL_CURSORS"
+
+ spec = SON([('killCursors', coll),
+ ('cursors', cursor_ids)])
+ with server.get_socket(self.__all_credentials) as sock_info:
+ if (sock_info.max_wire_version >= 4 and
+ namespace is not None):
+ sock_info.command(db, spec, slave_ok=True)
+ else:
+ if publish:
+ start = datetime.datetime.now()
+ request_id, msg = message.kill_cursors(cursor_ids)
+ if publish:
+ duration = datetime.datetime.now() - start
+ listeners.publish_command_start(
+ spec, db, request_id, address)
+ start = datetime.datetime.now()
+
+ try:
+ sock_info.send_message(msg, 0)
+ except Exception as exc:
+ if publish:
+ dur = ((datetime.datetime.now() - start)
+ + duration)
+ listeners.publish_command_failure(
+ dur, message._convert_exception(exc),
+ 'killCursors', request_id, address)
+ raise
+
+ if publish:
+ duration = ((datetime.datetime.now() - start)
+ + duration)
+ # OP_KILL_CURSORS returns no reply, fake one.
+ reply = {'cursorsUnknown': cursor_ids, 'ok': 1}
+ listeners.publish_command_success(
+ duration, reply, 'killCursors', request_id,
+ address)
 
  except ConnectionFailure as exc:
  warnings.warn("couldn't close cursor on %s: %s"

test/test_client.py

@@ -42,6 +42,7 @@
  CursorNotFound,
  NetworkTimeout,
  InvalidURI)
+from pymongo.message import _CursorAddress
 from pymongo.mongo_client import MongoClient
 from pymongo.pool import SocketInfo
 from pymongo.read_preferences import ReadPreference
@@ -707,7 +708,7 @@ def test_operation_failure(self):
  new_sock_info = next(iter(pool.sockets))
  self.assertEqual(old_sock_info, new_sock_info)
 
- def test_kill_cursors(self):
+ def test_kill_cursors_with_cursoraddress(self):
  if (client_context.is_mongos
  and not client_context.version.at_least(2, 4, 7)):
  # Old mongos sends incorrectly formatted error response when
@@ -720,7 +721,40 @@ def test_kill_cursors(self):
  self.collection.insert_many([{'_id': i} for i in range(200)])
  cursor = self.collection.find().batch_size(1)
  next(cursor)
- self.client.kill_cursors([cursor.cursor_id])
+ self.client.kill_cursors(
+ [cursor.cursor_id],
+ _CursorAddress(self.client.address, self.collection.full_name))
+
+ # Prevent killcursors from reaching the server while a getmore is in
+ # progress -- the server logs "Assertion: 16089:Cannot kill active
+ # cursor."
+ time.sleep(2)
+
+ def raises_cursor_not_found():
+ try:
+ next(cursor)
+ return False
+ except CursorNotFound:
+ return True
+
+ wait_until(raises_cursor_not_found, 'close cursor')
+
+ def test_kill_cursors_with_tuple(self):
+ if (client_context.is_mongos
+ and not client_context.version.at_least(2, 4, 7)):
+ # Old mongos sends incorrectly formatted error response when
+ # cursor isn't found, see SERVER-9738.
+ raise SkipTest("Can't test kill_cursors against old mongos")
+
+ self.collection = self.client.pymongo_test.test
+ self.collection.drop()
+
+ self.collection.insert_many([{'_id': i} for i in range(200)])
+ cursor = self.collection.find().batch_size(1)
+ next(cursor)
+ self.client.kill_cursors(
+ [cursor.cursor_id],
+ self.client.address)
 
  # Prevent killcursors from reaching the server while a getmore is in
  # progress -- the server logs "Assertion: 16089:Cannot kill active

test/test_cursor_manager.py

@@ -20,6 +20,7 @@
 
 from pymongo.cursor_manager import CursorManager
 from pymongo.errors import CursorNotFound
+from pymongo.message import _CursorAddress
 from test import (client_context,
  client_knobs,
  unittest,
@@ -74,7 +75,9 @@ def close(self, cursor_id, address):
  # is sent after the killCursors message.
  cursor = client.pymongo_test.test.find().batch_size(1)
  next(cursor)
- client.close_cursor(cursor.cursor_id)
+ client.close_cursor(
+ cursor.cursor_id,
+ _CursorAddress(self.client.address, self.collection.full_name))
 
  def raises_cursor_not_found():
  try:

test/test_monitoring.py

@@ -170,7 +170,7 @@ def test_find_and_get_more(self):
  self.assertEqual('find', succeeded.command_name)
  self.assertTrue(isinstance(succeeded.request_id, int))
  self.assertEqual(cursor.address, succeeded.connection_id)
- csr = succeeded.reply[u'cursor']
+ csr = succeeded.reply["cursor"]
  self.assertEqual(csr["id"], cursor_id)
  self.assertEqual(csr["ns"], "pymongo_test.test")
  self.assertEqual(csr["firstBatch"], [{} for _ in range(4)])
@@ -527,7 +527,8 @@ def test_kill_cursors(self):
  self.assertEqual(cursor.address, succeeded.connection_id)
  # There could be more than one cursor_id here depending on
  # when the thread last ran.
- self.assertIn(cursor_id, succeeded.reply['cursorsUnknown'])
+ self.assertTrue(cursor_id in succeeded.reply['cursorsUnknown']
+ or cursor_id in succeeded.reply['cursorsKilled'])
 
  def test_non_bulk_writes(self):
  coll = self.client.pymongo_test.test