Improve object store get usage

Signed-off-by: Waldemar Quevedo <[email protected]>

Author: wallyqs

nats/benchmark/sub_perf_messages.py

@@ -80,4 +80,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())

nats/src/nats/aio/subscription.py

@@ -80,6 +80,7 @@ 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
@@ -138,14 +139,19 @@ 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():
+                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.
@@ -171,11 +177,13 @@ async def _message_generator(self) -> AsyncIterator[Msg]:
                 # 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:
+                    if hasattr(self, "_message_iterator") and self._message_iterator:
                         self._message_iterator._cancel()
                     break
         except asyncio.CancelledError:
             pass
+        finally:
+            self._active_generators -= 1
 
     @property
     def pending_msgs(self) -> int:
@@ -334,15 +342,13 @@ 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 hasattr(self, "_message_iterator") and self._message_iterator:
             self._message_iterator._cancel()
-        
-        # Put a sentinel in the queue to wake up any generators waiting on get()
-        # This ensures they see the _closed flag and exit cleanly
+
+        # Only put sentinel if there are active async generators
         try:
-            if self._pending_queue:
-                # Put a None sentinel to wake up the generator
-                # The generator will check _closed and exit
+            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
@@ -358,12 +364,12 @@ 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:
@@ -394,12 +400,11 @@ class _CompatibilityIteratorWrapper:
     Compatibility wrapper that provides the same interface as the old _SubscriptionMessageIterator
     but uses the more efficient generator internally.
     """
+
     def __init__(self, sub: Subscription) -> None:
         self._sub = sub
         self._unsubscribed_future: asyncio.Future[bool] = asyncio.Future()
 
     def _cancel(self) -> None:
         if not self._unsubscribed_future.done():
             self._unsubscribed_future.set_result(True)
-
-

nats/src/nats/js/client.py

@@ -883,6 +883,7 @@ 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

nats/src/nats/js/object_store.py

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

nats/tests/test_client.py

@@ -560,7 +560,7 @@ async def iterator_func(sub):
         # Wait for iterator to complete after drain
         await asyncio.wait_for(fut, 1)
         await iterator_task  # Ensure task cleanup
-        
+
         self.assertEqual(5, len(msgs))
         self.assertEqual("tests.1", msgs[1].subject)
         self.assertEqual("tests.3", msgs[3].subject)
@@ -575,55 +575,55 @@ async def test_subscribe_async_generator(self):
         """Test the optimized async generator implementation for sub.messages"""
         nc = NATS()
         await nc.connect()
-        
+
         # Test basic async generator functionality
         sub = await nc.subscribe("test.generator")
-        
+
         # Publish messages
         num_msgs = 10
         for i in range(num_msgs):
             await nc.publish("test.generator", f"msg-{i}".encode())
         await nc.flush()
-        
+
         # Consume messages using async generator
         received_msgs = []
         async for msg in sub.messages:
             received_msgs.append(msg)
             if len(received_msgs) >= num_msgs:
                 break
-        
+
         # Verify all messages received correctly
         self.assertEqual(len(received_msgs), num_msgs)
         for i, msg in enumerate(received_msgs):
             self.assertEqual(msg.data, f"msg-{i}".encode())
             self.assertEqual(msg.subject, "test.generator")
-        
+
         await nc.close()
 
     @async_test
     async def test_subscribe_async_generator_with_drain(self):
         """Test async generator with drain functionality"""
         nc = NATS()
         await nc.connect()
-        
+
         sub = await nc.subscribe("test.drain")
-        
+
         # Publish messages
         for i in range(5):
             await nc.publish("test.drain", f"drain-msg-{i}".encode())
-        
+
         # Start consuming messages
         received_msgs = []
         async for msg in sub.messages:
             received_msgs.append(msg)
             # Drain after receiving all messages
             if len(received_msgs) == 5:
                 await sub.drain()
-        
+
         # Verify correct number of messages and drain worked
         self.assertEqual(len(received_msgs), 5)
         self.assertEqual(sub.pending_bytes, 0)
-        
+
         await nc.close()
 
     @async_test