Adjust async streaming response stop() method to have a timeout argument, and use it when exiting context manager.

Things sometimes hang in CI and this might help with that.

Also adjusted when _ack_queue.task_done() is called (not that anybody seems to be joining on the the queue?).

Author: johnbywater

kurrentdbclient/common.py

@@ -167,7 +167,7 @@ def __init__(self, grpc_streamers: AsyncGrpcStreamers) -> None:
         self._is_stopped_lock = asyncio.Lock()
 
     @abstractmethod
-    async def stop(self) -> None:
+    async def stop(self, *, timeout: float | None = None) -> None:
         """
         Stops the iterator(s) of streaming call.
         """
@@ -188,7 +188,7 @@ class AsyncGrpcStreamers(BaseGrpcStreamers[AsyncGrpcStreamer]):
     async def close(self) -> None:
         for async_grpc_streamer in self:
             # print("closing streamer")
-            await async_grpc_streamer.stop()
+            await async_grpc_streamer.stop(timeout=1)
             # print("closed streamer")
 
 
@@ -577,7 +577,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._is_context_manager_active = False
 
     @abstractmethod
-    async def stop(self) -> None:
+    async def stop(self, *, timeout: float | None = None) -> None:
         pass  # pragma: no cover
 
     def __aiter__(self) -> Self:
@@ -589,7 +589,7 @@ async def __aenter__(self) -> Self:
 
     async def __aexit__(self, *args: object, **kwargs: Any) -> None:
         self._is_context_manager_active = False
-        await self.stop()
+        await self.stop(timeout=5)
 
     def _set_iter_error_for_testing(self) -> None:
         # This, because I can't find a good way to inspire an error during iterating

kurrentdbclient/instrumentation/opentelemetry/spanners.py

@@ -956,8 +956,8 @@ async def __anext__(self) -> RecordedEvent:
                 span.end()
                 self._current_span = None
 
-    async def stop(self) -> None:
-        await self.response.stop()
+    async def stop(self, *, timeout: float | None = None) -> None:
+        await self.response.stop(timeout=timeout)
 
     async def __aenter__(self) -> Self:
         await self.response.__aenter__()

kurrentdbclient/persistent.py

@@ -110,6 +110,7 @@ def __init__(
         self._max_ack_delay = max_ack_delay
         self._stopping_grace = stopping_grace
         self._has_requested_options = False
+        self._num_constructed_acks = 0
         self.subscription_id = b""
 
     def _construct_initial_read_req(self) -> persistent_pb2.ReadReq:
@@ -130,10 +131,10 @@ def _construct_initial_read_req(self) -> persistent_pb2.ReadReq:
             options.all.CopyFrom(shared_pb2.Empty())
         return persistent_pb2.ReadReq(options=options)
 
-    @staticmethod
     def _construct_ack_or_nack_read_req(
-        subscription_id: bytes, event_ids: list[UUID], action: str
+        self, subscription_id: bytes, event_ids: list[UUID], action: str
     ) -> persistent_pb2.ReadReq:
+        self._num_constructed_acks = len(event_ids)
         ids = [shared_pb2.UUID(string=str(event_id)) for event_id in event_ids]
         if action == "ack":
             read_req = persistent_pb2.ReadReq(
@@ -210,17 +211,19 @@ async def __anext__(self) -> persistent_pb2.ReadReq:
             # First return read request with options, then return read request
             # with batch of n/acks whenever the batch is full, or when the n/ack
             # actions changes, or periodically, or when stopping.
-
             if not self._has_requested_options:
                 # Return initial read request with options.
                 self._has_requested_options = True
                 return self._construct_initial_read_req()
 
+            # Account on queue for previously returned n/acks.
+            for _ in range(self._num_constructed_acks):
+                self._ack_queue.task_done()
+            self._num_constructed_acks = 0
+
             # Return read request with a batch of n/acks...
 
             # Initialise batch, maybe from held n/ack.
-            for _ in self._batch_ids:
-                self._ack_queue.task_done()
             self._batch_ids = []
             batch_action: str | None = None
             if self._ack_held is not None:
@@ -234,7 +237,6 @@ async def __anext__(self) -> persistent_pb2.ReadReq:
                 if self._is_stopping:
                     # Allow time for server to process last n/acks.
                     await asyncio.sleep(self._stopping_grace)
-                    self._ack_queue.task_done()
                     raise StopAsyncIteration from None
 
                 try:
@@ -316,12 +318,11 @@ async def nack(
         assert action in ["unknown", "park", "retry", "skip", "stop"]
         await self._ack_queue.put((event_id, action))
 
-    async def stop(self, *, wait_until_stopped: bool = True) -> None:
+    async def stop(self, *, timeout: float | None = None) -> None:
         if not self._is_poisoned:
             self._is_poisoned = True
             await self._ack_queue.put((None, "poison"))
-            if wait_until_stopped:
-                await self._is_stopped.wait()
+            await asyncio.wait_for(self._is_stopped.wait(), timeout)
 
 
 class SubscriptionReadReqs(BaseSubscriptionReadReqs):
@@ -364,6 +365,11 @@ def __next__(self) -> persistent_pb2.ReadReq:
                 self._has_requested_options = True
                 return self._construct_initial_read_req()
 
+            # Account on queue for previously returned n/acks.
+            for _ in range(self._num_constructed_acks):
+                self._ack_queue.task_done()
+            self._num_constructed_acks = 0
+
             # Send a batch of n/acks...
 
             # Initialise batch, maybe from held n/ack.
@@ -387,7 +393,6 @@ def __next__(self) -> persistent_pb2.ReadReq:
                     # Wait for next n/ack, timeout with "max ack delay".
                     get_timeout = max(0.0, self._calc_time_until_next_ack_batch())
                     event_id, action = self._ack_queue.get(timeout=get_timeout)
-                    self._ack_queue.task_done()
 
                     # If queue was poisoned, send non-empty batch now.
                     if action == "poison":
@@ -572,7 +577,7 @@ async def init(self) -> None:
                 msg = f"Expected subscription confirmation, got: {first_read_resp}"
                 raise SubscriptionConfirmationError(msg)
         except BaseException:
-            await self.stop(wait_until_stopped=False)
+            await self.stop(timeout=1)
             raise
 
     @property
@@ -600,26 +605,27 @@ async def __anext__(self) -> RecordedEvent:
                 else:  # pragma: no cover
                     pass
         except BaseException:
-            await self.stop(wait_until_stopped=False)
+            await self.stop(timeout=1)
             raise
 
-    async def stop(self, *, wait_until_stopped: bool = True) -> None:
+    async def stop(self, *, timeout: float | None = None) -> None:
         if self._is_context_manager_active:
             self._is_stopping = True
         elif not await self._set_is_stopped():
-            await self._read_reqs.stop(wait_until_stopped=wait_until_stopped)
-            self._read_resp_stream_worker_task.cancel()
-            self._stream_stream_call.cancel()
-            await asyncio.sleep(0.05)
-            self._grpc_streamers.remove(self)
             try:
-                await self._read_resp_stream_worker_task
-            except asyncio.CancelledError:
-                pass
-            except grpc.RpcError as e:
-                raise handle_rpc_error(e) from e
-            if self._read_reqs.errored:
-                raise self._read_reqs.errored
+                await self._read_reqs.stop(timeout=timeout)
+            finally:
+                self._read_resp_stream_worker_task.cancel()
+                self._stream_stream_call.cancel()
+                self._grpc_streamers.remove(self)
+                try:
+                    await self._read_resp_stream_worker_task
+                except asyncio.CancelledError:
+                    pass
+                except grpc.RpcError as e:
+                    raise handle_rpc_error(e) from e
+                if self._read_reqs.errored:
+                    raise self._read_reqs.errored
 
     async def ack(self, item: UUID | RecordedEvent) -> None:
         await self._read_reqs.ack(event_id=self._get_event_id(item))

kurrentdbclient/streams.py

@@ -225,7 +225,7 @@ async def __anext__(self) -> RecordedEvent:
             await self.stop()
             raise
 
-    async def stop(self, *, wait_until_stopped: bool = True) -> None:
+    async def stop(self, *, timeout: float | None = None) -> None:
         if self._is_context_manager_active:
             self._is_stopping = True
         elif not await self._set_is_stopped():

tests/test_client_async.py

@@ -2752,12 +2752,8 @@ async def test_persistent_subscription_read_reqs(self) -> None:
         req5 = await reqs.__anext__()
         self.assertEqual(len(req5.ack.ids), 2)
 
-        # Cover the case of stopping without waiting (wait_until_stopped=False).
-        reqs = AsyncSubscriptionReadReqs("group1", max_ack_batch_size=3)
-        await reqs.stop(wait_until_stopped=False)
-
-        # Cover the case of calling this method twice.
-        await reqs.stop(wait_until_stopped=False)
+        # Cover the case of calling stop() twice.
+        await reqs.stop()
 
         # Iterate until stopped.
         async for _ in reqs:
@@ -2772,10 +2768,20 @@ async def iterate_until_stopped() -> None:
 
         async def sleep_then_stop() -> None:
             await asyncio.sleep(1)
-            await reqs.stop(wait_until_stopped=False)
+            await reqs.stop()
 
         await asyncio.gather(iterate_until_stopped(), sleep_then_stop())
 
+        # Iterate and stop with timeout when request iterator is not being called.
+        reqs = AsyncSubscriptionReadReqs("group1", max_ack_batch_size=3)
+        await reqs.ack(uuid4())
+
+        async def sleep_then_stop_with_timeout() -> None:
+            await reqs.stop(timeout=1)
+
+        with self.assertRaises(TimeoutError):
+            await sleep_then_stop_with_timeout()
+
         # Can't call ack() after stopped.
         with self.assertRaises(ProgrammingError):
             await reqs.ack(uuid4())