Restore nats directory

Author: caspervonb

nats/src/nats/aio/client.py

@@ -754,6 +754,8 @@ async def _close(self, status: int, do_cbs: bool = True) -> None:
             # Async subs use join when draining already so just cancel here.
             if sub._wait_for_msgs_task and not sub._wait_for_msgs_task.done():
                 sub._wait_for_msgs_task.cancel()
+            if sub._message_iterator:
+                sub._message_iterator._cancel()
             # Sync subs may have some inflight next_msg calls that could be blocking
             # so cancel them here to unblock them.
             if sub._pending_next_msgs_calls:

nats/src/nats/aio/subscription.py

@@ -20,9 +20,10 @@
     AsyncIterator,
     Awaitable,
     Callable,
-    Dict,
+    List,
     Optional,
 )
+from uuid import uuid4
 
 from nats import errors
 
@@ -80,7 +81,6 @@ def __init__(
         self._cb = cb
         self._future = future
         self._closed = False
-        self._active_generators = 0  # Track active async generators
 
         # Per subscription message processor.
         self._pending_msgs_limit = pending_msgs_limit
@@ -89,12 +89,11 @@ def __init__(
         # If no callback, then this is a sync subscription which will
         # require tracking the next_msg calls inflight for cancelling.
         if cb is None:
-            self._pending_next_msgs_calls: Optional[Dict[str, asyncio.Task]] = {}
+            self._pending_next_msgs_calls = {}
         else:
             self._pending_next_msgs_calls = None
         self._pending_size = 0
         self._wait_for_msgs_task = None
-        # For compatibility with tests that expect _message_iterator
         self._message_iterator = None
 
         # For JetStream enabled subscriptions.
@@ -130,61 +129,10 @@ def messages(self) -> AsyncIterator[Msg]:
             async for msg in sub.messages:
                 print('Received', msg)
         """
-        if self._cb:
+        if not self._message_iterator:
             raise errors.Error("cannot iterate over messages with a non iteration subscription type")
 
-        return self._message_generator()
-
-    async def _message_generator(self) -> AsyncIterator[Msg]:
-        """
-        Async generator that yields messages directly from the subscription queue.
-        """
-        yielded_count = 0
-        self._active_generators += 1
-        try:
-            while True:
-                # Check if subscription was cancelled/closed.
-                if self._closed:
-                    break
-
-                # Check if wrapper was cancelled (for compatibility with tests).
-                if (
-                    hasattr(self, "_message_iterator")
-                    and self._message_iterator
-                    and self._message_iterator._unsubscribed_future.done()
-                ):
-                    break
-
-                # Check max message limit based on how many we've yielded so far.
-                if self._max_msgs > 0 and yielded_count >= self._max_msgs:
-                    break
-
-                try:
-                    msg = await self._pending_queue.get()
-                except asyncio.CancelledError:
-                    break
-
-                # Check for sentinel value which signals generator to stop.
-                if msg is None:
-                    self._pending_queue.task_done()
-                    break
-
-                self._pending_queue.task_done()
-                self._pending_size -= len(msg.data)
-
-                yield msg
-                yielded_count += 1
-
-                # Check if we should auto-unsubscribe after yielding this message.
-                if self._max_msgs > 0 and yielded_count >= self._max_msgs:
-                    # Cancel the wrapper too for consistency.
-                    if hasattr(self, "_message_iterator") and self._message_iterator:
-                        self._message_iterator._cancel()
-                    break
-        except asyncio.CancelledError:
-            pass
-        finally:
-            self._active_generators -= 1
+        return self._message_iterator
 
     @property
     def pending_msgs(self) -> int:
@@ -212,7 +160,6 @@ def delivered(self) -> int:
     async def next_msg(self, timeout: Optional[float] = 1.0) -> Msg:
         """
         :params timeout: Time in seconds to wait for next message before timing out.
-                        Use 0 or None to wait forever (no timeout).
         :raises nats.errors.TimeoutError:
 
         next_msg can be used to retrieve the next message from a stream of messages using
@@ -221,23 +168,22 @@ async def next_msg(self, timeout: Optional[float] = 1.0) -> Msg:
             sub = await nc.subscribe('hello')
             msg = await sub.next_msg(timeout=1)
 
-            # Wait forever for a message
-            msg = await sub.next_msg(timeout=0)
-
         """
+
+        async def timed_get() -> Msg:
+            return await asyncio.wait_for(self._pending_queue.get(), timeout)
+
         if self._conn.is_closed:
             raise errors.ConnectionClosedError
 
         if self._cb:
             raise errors.Error("nats: next_msg cannot be used in async subscriptions")
 
+        task_name = str(uuid4())
         try:
-            if timeout == 0 or timeout is None:
-                # Wait forever for a message
-                msg = await self._pending_queue.get()
-            else:
-                # Wait with timeout
-                msg = await asyncio.wait_for(self._pending_queue.get(), timeout)
+            future = asyncio.create_task(timed_get())
+            self._pending_next_msgs_calls[task_name] = future
+            msg = await future
         except asyncio.TimeoutError:
             if self._conn.is_closed:
                 raise errors.ConnectionClosedError
@@ -253,6 +199,8 @@ async def next_msg(self, timeout: Optional[float] = 1.0) -> Msg:
             # regardless of whether it has been processed.
             self._pending_queue.task_done()
             return msg
+        finally:
+            self._pending_next_msgs_calls.pop(task_name, None)
 
     def _start(self, error_cb):
         """
@@ -270,9 +218,7 @@ def _start(self, error_cb):
             # Used to handle the single response from a request.
             pass
         else:
-            # For async iteration, we now use a generator directly via the messages property
-            # But we create a compatibility wrapper for tests
-            self._message_iterator = _CompatibilityIteratorWrapper(self)
+            self._message_iterator = _SubscriptionMessageIterator(self)
 
     async def drain(self):
         """
@@ -343,18 +289,9 @@ def _stop_processing(self) -> None:
         """
         if self._wait_for_msgs_task and not self._wait_for_msgs_task.done():
             self._wait_for_msgs_task.cancel()
-        if hasattr(self, "_message_iterator") and self._message_iterator:
+        if self._message_iterator:
             self._message_iterator._cancel()
 
-        # Only put sentinel if there are active async generators
-        try:
-            if self._pending_queue and self._active_generators > 0:
-                # Put a None sentinel to wake up any async generators
-                self._pending_queue.put_nowait(None)
-        except Exception:
-            # Queue might be closed or full, that's ok
-            pass
-
     async def _wait_for_msgs(self, error_cb) -> None:
         """
         A coroutine to read and process messages if a callback is provided.
@@ -365,12 +302,6 @@ async def _wait_for_msgs(self, error_cb) -> None:
         while True:
             try:
                 msg = await self._pending_queue.get()
-
-                # Check for sentinel value (None) which signals task to stop
-                if msg is None:
-                    self._pending_queue.task_done()
-                    break
-
                 self._pending_size -= len(msg.data)
 
                 try:
@@ -396,16 +327,35 @@ async def _wait_for_msgs(self, error_cb) -> None:
                 break
 
 
-class _CompatibilityIteratorWrapper:
-    """
-    Compatibility wrapper that provides the same interface as the old _SubscriptionMessageIterator
-    but uses the more efficient generator internally.
-    """
-
+class _SubscriptionMessageIterator:
     def __init__(self, sub: Subscription) -> None:
-        self._sub = sub
+        self._sub: Subscription = sub
+        self._queue: asyncio.Queue[Msg] = sub._pending_queue
         self._unsubscribed_future: asyncio.Future[bool] = asyncio.Future()
 
     def _cancel(self) -> None:
         if not self._unsubscribed_future.done():
             self._unsubscribed_future.set_result(True)
+
+    def __aiter__(self) -> _SubscriptionMessageIterator:
+        return self
+
+    async def __anext__(self) -> Msg:
+        get_task = asyncio.get_running_loop().create_task(self._queue.get())
+        tasks: List[asyncio.Future] = [get_task, self._unsubscribed_future]
+        finished, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        sub = self._sub
+
+        if get_task in finished:
+            self._queue.task_done()
+            msg = get_task.result()
+            self._sub._pending_size -= len(msg.data)
+
+            # Unblock the iterator in case it has already received enough messages.
+            if sub._max_msgs > 0 and sub._received >= sub._max_msgs:
+                self._cancel()
+            return msg
+        elif self._unsubscribed_future.done():
+            get_task.cancel()
+
+        raise StopAsyncIteration

nats/src/nats/js/client.py

@@ -883,14 +883,14 @@ def __init__(
             self._cb = sub._cb
             self._future = sub._future
             self._closed = sub._closed
-            self._active_generators = sub._active_generators
 
             # Per subscription message processor.
             self._pending_msgs_limit = sub._pending_msgs_limit
             self._pending_bytes_limit = sub._pending_bytes_limit
             self._pending_queue = sub._pending_queue
             self._pending_size = sub._pending_size
             self._wait_for_msgs_task = sub._wait_for_msgs_task
+            self._message_iterator = sub._message_iterator
             self._pending_next_msgs_calls = sub._pending_next_msgs_calls
 
         async def consumer_info(self) -> api.ConsumerInfo:

nats/src/nats/js/object_store.py

@@ -213,7 +213,7 @@ async def get(
             else:
                 executor_fn = writeinto.write
 
-        async for msg in sub.messages:
+        async for msg in sub._message_iterator:
             tokens = msg._get_metadata_fields(msg.reply)
 
             if executor: