PYCBC-1774: Add cancellation/close support to the C-extension streaming result

Changes
-------
* rows_queue: add a cancelled flag with cancel()/is_cancelled();
  get() returns a null sentinel once cancelled and the queue drains,
  unblocking a waiting worker (already-queued rows still drain first)
* pycbc_streamed_result: expose cancel()/is_cancelled()
* _finalize: cancel the streamed result on abort only (exc_val is not
  None), so normal completion can still read trailing metadata
* acouchbase range scan: add _abort() to cancel the in-flight core
  scan (existing cancel_scan()) and shut down the executor on the
  error/cancellation paths
* _core.pyi: add cancel()/is_cancelled() stubs (and the missing
  scan-iterator cancel_scan()/is_cancelled())
* tests: helper wiring tests (cancel on abort, skip on normal
  completion) + a live query_t test asserting CancelledError
  propagates and the streamed result is cancelled end-to-end

Change-Id: I0c344ac5daebf4be9c6ec3113b8a82e3bca689f2
Reviewed-on: https://review.couchbase.org/c/couchbase-python-client/+/246063
Reviewed-by: Dimitris Christodoulou <dimitris.christodoulou@couchbase.com>
Tested-by: Build Bot <build@couchbase.com>

Author: thejcfactor

acouchbase/analytics.py

@@ -115,6 +115,12 @@ def _finalize(self, exc_val=None):
         teardown and leak the span/meter or orphan the streaming result and its executor thread.
         """
         self._process_core_span(exc_val=exc_val)
+        if exc_val is not None and self._streaming_result is not None:
+            # Cancellation/error: unblock a worker still waiting on the C++ core so its executor
+            # thread is released promptly instead of waiting for the whole server-side operation
+            # to finish.  Normal completion (exc_val is None) skips this so trailing metadata can
+            # still be read from the streaming result.
+            self._streaming_result.cancel()
         self._shutdown_executor()
 
     async def __anext__(self):

acouchbase/kv_range_scan.py

@@ -65,14 +65,25 @@ def _shutdown_executor(self):
             self._executor_shutdown = True
             self._tp_executor.shutdown(wait=False)
 
+    def _abort(self):
+        """
+        **INTERNAL**
+
+        Cancel the in-flight core scan so a worker still waiting on the C++ core unwinds promptly,
+        then release the executor.  Used on the error/cancellation paths only.
+        """
+        if self._scan_iterator is not None and self._scan_iterator.is_cancelled() is False:
+            self._scan_iterator.cancel_scan()
+        self._shutdown_executor()
+
     async def __anext__(self):
         try:
             return await self._loop.run_in_executor(self._tp_executor, self._get_next_row)
         # We can stop iterator when we receive RangeScanCompletedException
         except asyncio.QueueEmpty:
             exc_cls = PYCBC_ERROR_MAP.get(ExceptionMap.InternalSDKException.value, CouchbaseException)
             excptn = exc_cls('Unexpected QueueEmpty exception caught when doing N1QL query.')
-            self._shutdown_executor()
+            self._abort()
             raise excptn
         except RangeScanCompletedException:
             self._done_streaming = True
@@ -83,16 +94,16 @@ async def __anext__(self):
             self._shutdown_executor()
             raise
         except CouchbaseException as ex:
-            self._shutdown_executor()
+            self._abort()
             raise ex
         except Exception as ex:
             exc_cls = PYCBC_ERROR_MAP.get(ExceptionMap.InternalSDKException.value, CouchbaseException)
             excptn = exc_cls(str(ex))
-            self._shutdown_executor()
+            self._abort()
             raise excptn
         except BaseException:
             # asyncio.CancelledError (and KeyboardInterrupt/SystemExit) derive from BaseException,
-            # not Exception, so they fall through every handler above.  Release the executor so the
-            # cancelled scan does not orphan its worker thread, then re-raise unchanged.
-            self._shutdown_executor()
+            # not Exception, so they fall through every handler above.  Cancel the in-flight scan
+            # and release the executor so the worker thread is not orphaned, then re-raise unchanged.
+            self._abort()
             raise

acouchbase/n1ql.py

@@ -117,6 +117,12 @@ def _finalize(self, exc_val=None):
         teardown and leak the span/meter or orphan the streaming result and its executor thread.
         """
         self._process_core_span(exc_val=exc_val)
+        if exc_val is not None and self._streaming_result is not None:
+            # Cancellation/error: unblock a worker still waiting on the C++ core so its executor
+            # thread is released promptly instead of waiting for the whole server-side operation
+            # to finish.  Normal completion (exc_val is None) skips this so trailing metadata can
+            # still be read from the streaming result.
+            self._streaming_result.cancel()
         self._shutdown_executor()
 
     async def __anext__(self):

acouchbase/search.py

@@ -113,6 +113,12 @@ def _finalize(self, exc_val=None):
         teardown and leak the span/meter or orphan the streaming result and its executor thread.
         """
         self._process_core_span(exc_val=exc_val)
+        if exc_val is not None and self._streaming_result is not None:
+            # Cancellation/error: unblock a worker still waiting on the C++ core so its executor
+            # thread is released promptly instead of waiting for the whole server-side operation
+            # to finish.  Normal completion (exc_val is None) skips this so trailing metadata can
+            # still be read from the streaming result.
+            self._streaming_result.cancel()
         self._shutdown_executor()
 
     async def __anext__(self):

acouchbase/tests/query_t.py

@@ -187,6 +187,7 @@ class QueryTestSuite:
         'test_mixed_positional_parameters',
         'test_non_blocking',
         'test_preserve_expiry',
+        'test_query_cancellation',
         'test_query_error_context',
         'test_query_metadata',
         'test_query_raw_options',
@@ -216,6 +217,33 @@ async def test_bad_query(self, cb_env):
         with pytest.raises(ParsingFailedException):
             await cb_env.cluster.query("I'm not N1QL!").execute()
 
+    @pytest.mark.asyncio
+    async def test_query_cancellation(self, cb_env):
+        # PYCBC-1774: cancelling an in-flight streaming query must propagate CancelledError
+        # unchanged (not swallow or convert it) and cancel the underlying C++ streamed result so
+        # the executor worker is released promptly instead of waiting for the whole server-side
+        # operation to finish.
+        result = cb_env.cluster.query(f"SELECT * FROM `{cb_env.bucket.name}` LIMIT 100")
+
+        async def consume():
+            rows = []
+            async for row in result:
+                rows.append(row)
+            return rows
+
+        task = asyncio.ensure_future(consume())
+        # let the task start streaming (submit the query) and reach the first, blocking __anext__
+        while result._request._streaming_result is None and not task.done():
+            await asyncio.sleep(0)
+        # cancel while the worker is still waiting on the C++ core for the first row
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        # the cancellation must have reached through to the C++ streamed result
+        assert result._request._streaming_result is not None
+        assert result._request._streaming_result.is_cancelled() is True
+
     @pytest.mark.asyncio
     async def test_mixed_named_parameters(self, cb_env):
         batch_id = cb_env.get_batch_id()

acouchbase/tests/streaming_helpers_t.py

@@ -19,10 +19,21 @@
 
 import pytest
 
+from acouchbase.n1ql import AsyncN1QLRequest
 from couchbase.exceptions import CouchbaseException, InternalSDKException
 from couchbase.logic.streaming import stream_anext
 
 
+class _FakeStreamingResult:
+    """Stand-in for the C-extension ``pycbc_streamed_result`` exposing the new cancel hook."""
+
+    def __init__(self):
+        self.cancel_calls = 0
+
+    def cancel(self):
+        self.cancel_calls += 1
+
+
 class _FakeAsyncStreamingRequest:
     """Minimal stand-in implementing the contract ``stream_anext`` relies on, so the streaming
     iterator teardown/error handling can be exercised without a live cluster.  Records the
@@ -76,6 +87,8 @@ class StreamingAnextTestSuite:
         'test_queue_empty_converted_with_op_name',
         'test_keyboard_interrupt_propagates_unconverted',
         'test_cancellation_finalizes_and_propagates',
+        'test_finalize_cancels_streaming_result_on_error',
+        'test_finalize_skips_cancel_on_normal_completion',
     ]
 
     @pytest.mark.asyncio
@@ -158,6 +171,30 @@ async def test_cancellation_finalizes_and_propagates(self):
         assert isinstance(req.finalize_calls[0], asyncio.CancelledError)
         assert req.executor_shutdown is True
 
+    def test_finalize_cancels_streaming_result_on_error(self):
+        # _finalize is exercised on a real request to verify the cancel() wiring (Phase III).
+        loop = asyncio.new_event_loop()
+        try:
+            req = AsyncN1QLRequest(None, loop, {}, obs_handler=None)
+            req._streaming_result = _FakeStreamingResult()
+            req._finalize(exc_val=CouchbaseException('boom'))
+            # abort path: the C++ streamed result is cancelled so a blocked worker can unwind
+            assert req._streaming_result.cancel_calls == 1
+            assert req._executor_shutdown is True
+        finally:
+            loop.close()
+
+    def test_finalize_skips_cancel_on_normal_completion(self):
+        loop = asyncio.new_event_loop()
+        try:
+            req = AsyncN1QLRequest(None, loop, {}, obs_handler=None)
+            req._streaming_result = _FakeStreamingResult()
+            req._finalize()  # exc_val is None -> normal completion must NOT cancel (metadata follows)
+            assert req._streaming_result.cancel_calls == 0
+            assert req._executor_shutdown is True
+        finally:
+            loop.close()
+
 
 class AsyncStreamingAnextTests(StreamingAnextTestSuite):
     @pytest.fixture(scope='class', autouse=True)

acouchbase/views.py

@@ -121,6 +121,12 @@ def _finalize(self, exc_val=None):
         teardown and leak the span/meter or orphan the streaming result and its executor thread.
         """
         self._process_core_span(exc_val=exc_val)
+        if exc_val is not None and self._streaming_result is not None:
+            # Cancellation/error: unblock a worker still waiting on the C++ core so its executor
+            # thread is released promptly instead of waiting for the whole server-side operation
+            # to finish.  Normal completion (exc_val is None) skips this so trailing metadata can
+            # still be read from the streaming result.
+            self._streaming_result.cancel()
         self._shutdown_executor()
 
     async def __anext__(self):

couchbase/logic/pycbc_core/_core.pyi

@@ -82,13 +82,21 @@ class pycbc_streamed_result(Generic[T]):
 
     def __next__(self) -> Union[T, pycbc_exception, None]: ...
 
+    def cancel(self) -> None: ...
+
+    def is_cancelled(self) -> bool: ...
+
 
 class pycbc_scan_iterator(Generic[T]):
 
     def __iter__(self) -> pycbc_scan_iterator[T]: ...
 
     def __next__(self) -> Union[T, pycbc_exception]: ...
 
+    def cancel_scan(self) -> None: ...
+
+    def is_cancelled(self) -> bool: ...
+
 
 class pycbc_exception:
 

src/result.cxx

@@ -214,6 +214,36 @@ static PyMemberDef pycbc_streamed_result_members[] = {
   { nullptr } // Sentinel
 };
 
+static PyObject*
+pycbc_streamed_result__cancel__(pycbc_streamed_result* self, PyObject* args)
+{
+  if (self->rows) {
+    self->rows->cancel();
+  }
+  Py_RETURN_NONE;
+}
+
+static PyObject*
+pycbc_streamed_result__is_cancelled__(pycbc_streamed_result* self, PyObject* args)
+{
+  if (self->rows && self->rows->is_cancelled()) {
+    Py_RETURN_TRUE;
+  }
+  Py_RETURN_FALSE;
+}
+
+static PyMethodDef pycbc_streamed_result_methods[] = {
+  { "cancel",
+    (PyCFunction)pycbc_streamed_result__cancel__,
+    METH_NOARGS,
+    PyDoc_STR("Cancel the streaming operation so a blocked iterator can unwind") },
+  { "is_cancelled",
+    (PyCFunction)pycbc_streamed_result__is_cancelled__,
+    METH_NOARGS,
+    PyDoc_STR("Check if the streaming operation has been cancelled") },
+  { nullptr } // Sentinel
+};
+
 static PyTypeObject pycbc_streamed_result_type = {
   PyVarObject_HEAD_INIT(nullptr, 0) "pycbc_core.pycbc_streamed_result", // tp_name
   sizeof(pycbc_streamed_result),                                        // tp_basicsize
@@ -241,7 +271,7 @@ static PyTypeObject pycbc_streamed_result_type = {
   0,                                                                    // tp_weaklistoffset
   pycbc_streamed_result__iter__,                                        // tp_iter
   pycbc_streamed_result__iternext__,                                    // tp_iternext
-  nullptr,                                                              // tp_methods
+  pycbc_streamed_result_methods,                                        // tp_methods
   pycbc_streamed_result_members,                                        // tp_members
   nullptr,                                                              // tp_getset
   nullptr,                                                              // tp_base

src/result.hxx

@@ -51,6 +51,7 @@ public:
     : rows_()
     , mut_()
     , cv_()
+    , cancelled_(false)
   {
   }
 
@@ -72,11 +73,33 @@ public:
     cv_.notify_one();
   }
 
+  // Unblock a waiting get(): once cancelled, get() stops waiting for new rows and returns a
+  // null sentinel as soon as the queue drains.  Used to release a streaming iterator's worker
+  // thread when the operation is cancelled/errored, rather than waiting for the entire
+  // server-side operation to complete.  Already-queued rows are still drained first.
+  void cancel()
+  {
+    std::lock_guard<std::mutex> lock(mut_);
+    cancelled_ = true;
+    cv_.notify_all();
+  }
+
+  bool is_cancelled()
+  {
+    std::lock_guard<std::mutex> lock(mut_);
+    return cancelled_;
+  }
+
   T get(std::chrono::milliseconds timeout_ms)
   {
     std::unique_lock<std::mutex> lock(mut_);
 
     while (rows_.empty()) {
+      if (cancelled_) {
+        // Cancelled with no rows remaining: stop waiting and signal end-of-iteration
+        // (a null PyObject* surfaces as StopIteration to the Python iterator).
+        return nullptr;
+      }
       if (cv_.wait_for(lock, timeout_ms) == std::cv_status::timeout) {
         // This timeout (e.g. timeout_ms) is the same timeout we pass to the C++ core.
         // If we timeout on the Python side this means:
@@ -111,6 +134,7 @@ private:
   std::queue<T> rows_;
   std::mutex mut_;
   std::condition_variable cv_;
+  bool cancelled_;
 };
 
 struct pycbc_result {