Track next_msg and async for generators the same way

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

Author: wallyqs

nats/src/nats/aio/client.py

@@ -754,12 +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()
-            # 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:
-                for fut in sub._pending_next_msgs_calls.values():
-                    fut.cancel()
-                sub._pending_next_msgs_calls.clear()
+            # For sync subs, stop processing will send sentinels to unblock any waiting consumers
+            sub._stop_processing()
         self._subs.clear()
 
         if self._transport is not None:
@@ -1802,9 +1798,9 @@ async def _process_msg(
                 await sub._jsi.check_for_sequence_mismatch(msg)
 
         # Send sentinel after reaching max messages for non-callback subscriptions.
-        if max_msgs_reached and not sub._cb and sub._active_generators > 0:
-            # Send one sentinel per active generator to unblock them all.
-            for _ in range(sub._active_generators):
+        if max_msgs_reached and not sub._cb and sub._active_consumers is not None and sub._active_consumers > 0:
+            # Send one sentinel per active consumer to unblock them all.
+            for _ in range(sub._active_consumers):
                 try:
                     sub._pending_queue.put_nowait(None)
                 except Exception:

nats/src/nats/aio/subscription.py

@@ -20,7 +20,6 @@
     AsyncIterator,
     Awaitable,
     Callable,
-    Dict,
     Optional,
 )
 
@@ -80,18 +79,16 @@ 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
         self._pending_bytes_limit = pending_bytes_limit
         self._pending_queue: asyncio.Queue[Msg] = asyncio.Queue(maxsize=pending_msgs_limit)
-        # If no callback, then this is a sync subscription which will
-        # require tracking the next_msg calls inflight for cancelling.
+        # Track active consumers (both async generators and next_msg calls) for non-callback subscriptions.
         if cb is None:
-            self._pending_next_msgs_calls: Optional[Dict[str, asyncio.Task]] = {}
+            self._active_consumers = 0  # Counter of active consumers waiting for messages
         else:
-            self._pending_next_msgs_calls = None
+            self._active_consumers = None
         self._pending_size = 0
         self._wait_for_msgs_task = None
 
@@ -138,7 +135,8 @@ async def _message_generator(self) -> AsyncIterator[Msg]:
         Async generator that yields messages directly from the subscription queue.
         """
         yielded_count = 0
-        self._active_generators += 1
+        if self._active_consumers is not None:
+            self._active_consumers += 1
         try:
             while True:
                 # Check if subscription was cancelled/closed.
@@ -171,7 +169,8 @@ async def _message_generator(self) -> AsyncIterator[Msg]:
         except asyncio.CancelledError:
             pass
         finally:
-            self._active_generators -= 1
+            if self._active_consumers is not None:
+                self._active_consumers -= 1
 
     @property
     def pending_msgs(self) -> int:
@@ -225,6 +224,10 @@ async def next_msg(self, timeout: Optional[float] = 1.0) -> Msg:
         if self._cb:
             raise errors.Error("nats: next_msg cannot be used in async subscriptions")
 
+        # Track this next_msg call
+        if self._active_consumers is not None:
+            self._active_consumers += 1
+
         try:
             if timeout == 0 or timeout is None:
                 # Wait forever for a message
@@ -240,13 +243,25 @@ async def next_msg(self, timeout: Optional[float] = 1.0) -> Msg:
             if self._conn.is_closed:
                 raise errors.ConnectionClosedError
             raise
-        else:
-            self._pending_size -= len(msg.data)
-            # For sync subscriptions we will consider a message
-            # to be done once it has been consumed by the client
-            # regardless of whether it has been processed.
+        finally:
+            # Untrack this next_msg call.
+            if self._active_consumers is not None:
+                self._active_consumers -= 1
+
+        # Check for sentinel value which signals to stop
+        if msg is None:
             self._pending_queue.task_done()
-            return msg
+            if self._conn.is_closed:
+                raise errors.ConnectionClosedError
+            raise errors.TimeoutError
+
+        self._pending_size -= len(msg.data)
+
+        # NOTE: For sync subscriptions we will consider a message
+        # to be done once it has been consumed by the client
+        # regardless of whether it has been processed.
+        self._pending_queue.task_done()
+        return msg
 
     def _start(self, error_cb):
         """
@@ -337,11 +352,12 @@ 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()
 
-        # Only put sentinel if there are active async generators
+        # Send sentinels to unblock waiting consumers
         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)
+            if self._pending_queue and self._active_consumers is not None and self._active_consumers > 0:
+                # Send one sentinel for each active consumer (both generators and next_msg calls)
+                for _ in range(self._active_consumers):
+                    self._pending_queue.put_nowait(None)
         except Exception:
             # Queue might be closed or full, that's ok
             pass

nats/tests/test_client.py

@@ -800,6 +800,246 @@ async def test_subscribe_async_generator_with_drain(self):
 
         await nc.close()
 
+    @async_test
+    async def test_subscribe_concurrent_next_msg(self):
+        """Test multiple concurrent next_msg() calls on the same subscription"""
+        nc = NATS()
+        await nc.connect()
+
+        sub = await nc.subscribe("test.concurrent.next")
+
+        # Publish messages
+        num_msgs = 12
+        for i in range(num_msgs):
+            await nc.publish("test.concurrent.next", f"msg-{i}".encode())
+        await nc.flush()
+
+        # Track results from concurrent next_msg calls
+        consumer_results = {}
+
+        async def consumer_task(consumer_id: str, msg_count: int):
+            """Consumer task that uses next_msg() to get messages"""
+            import random
+
+            received = []
+            try:
+                for _ in range(msg_count):
+                    msg = await sub.next_msg(timeout=2.0)
+                    received.append(msg.data.decode())
+                    # Add random processing delay
+                    await asyncio.sleep(random.uniform(0.01, 0.03))
+            except Exception as e:
+                consumer_results[consumer_id] = f"Error: {e}"
+                return
+            consumer_results[consumer_id] = received
+
+        # Start multiple concurrent consumers using next_msg()
+        tasks = [
+            asyncio.create_task(consumer_task("consumer_A", 3)),
+            asyncio.create_task(consumer_task("consumer_B", 5)),
+            asyncio.create_task(consumer_task("consumer_C", 4)),
+        ]
+
+        # Wait for all consumers to finish
+        await asyncio.gather(*tasks)
+
+        # Verify results
+        consumer_A_msgs = consumer_results.get("consumer_A", [])
+        consumer_B_msgs = consumer_results.get("consumer_B", [])
+        consumer_C_msgs = consumer_results.get("consumer_C", [])
+
+        # All consumers should have finished without errors
+        self.assertIsInstance(consumer_A_msgs, list, f"Consumer A failed: {consumer_A_msgs}")
+        self.assertIsInstance(consumer_B_msgs, list, f"Consumer B failed: {consumer_B_msgs}")
+        self.assertIsInstance(consumer_C_msgs, list, f"Consumer C failed: {consumer_C_msgs}")
+
+        # Each consumer should get exactly what they requested
+        self.assertEqual(len(consumer_A_msgs), 3, f"Consumer A got {len(consumer_A_msgs)} messages, expected 3")
+        self.assertEqual(len(consumer_B_msgs), 5, f"Consumer B got {len(consumer_B_msgs)} messages, expected 5")
+        self.assertEqual(len(consumer_C_msgs), 4, f"Consumer C got {len(consumer_C_msgs)} messages, expected 4")
+
+        # All messages should be unique (no duplicates across consumers)
+        all_received = consumer_A_msgs + consumer_B_msgs + consumer_C_msgs
+        self.assertEqual(
+            len(all_received),
+            len(set(all_received)),
+            f"Found duplicate messages: {[msg for msg in all_received if all_received.count(msg) > 1]}",
+        )
+
+        # All received messages should be from our published set
+        expected_msgs = {f"msg-{i}" for i in range(num_msgs)}
+        received_msgs = set(all_received)
+        self.assertTrue(received_msgs.issubset(expected_msgs))
+
+        # Total should be exactly 12 messages consumed
+        self.assertEqual(len(received_msgs), 12)
+
+        await nc.close()
+
+    @async_test
+    async def test_subscribe_concurrent_next_msg_with_unsubscribe_limit(self):
+        """Test concurrent next_msg() calls with unsubscribe limit"""
+        nc = NATS()
+        await nc.connect()
+
+        sub = await nc.subscribe("test.concurrent.next.limit")
+        await sub.unsubscribe(limit=8)  # Auto-unsubscribe after 8 messages
+
+        # Publish more messages than the limit
+        num_msgs = 15
+        for i in range(num_msgs):
+            await nc.publish("test.concurrent.next.limit", f"msg-{i}".encode())
+        await nc.flush()
+
+        # Track results from concurrent next_msg calls
+        consumer_results = {}
+
+        async def consumer_task(consumer_id: str, max_attempts: int):
+            """Consumer that keeps calling next_msg until timeout or limit reached"""
+            import random
+
+            received = []
+            try:
+                for attempt in range(max_attempts):
+                    try:
+                        msg = await sub.next_msg(timeout=0.5)
+                        received.append(msg.data.decode())
+                        # Add random processing delay
+                        await asyncio.sleep(random.uniform(0.005, 0.02))
+                    except Exception as e:
+                        # Expected when subscription reaches limit
+                        break
+            except Exception as e:
+                consumer_results[consumer_id] = f"Error: {e}"
+                return
+            consumer_results[consumer_id] = received
+
+        # Start multiple concurrent consumers
+        tasks = [
+            asyncio.create_task(consumer_task("consumer_A", 10)),
+            asyncio.create_task(consumer_task("consumer_B", 10)),
+            asyncio.create_task(consumer_task("consumer_C", 10)),
+        ]
+
+        # Wait for all consumers to finish
+        await asyncio.gather(*tasks)
+
+        # Verify results
+        consumer_A_msgs = consumer_results.get("consumer_A", [])
+        consumer_B_msgs = consumer_results.get("consumer_B", [])
+        consumer_C_msgs = consumer_results.get("consumer_C", [])
+
+        # All consumers should have finished without errors
+        self.assertIsInstance(consumer_A_msgs, list, f"Consumer A failed: {consumer_A_msgs}")
+        self.assertIsInstance(consumer_B_msgs, list, f"Consumer B failed: {consumer_B_msgs}")
+        self.assertIsInstance(consumer_C_msgs, list, f"Consumer C failed: {consumer_C_msgs}")
+
+        # Total messages across all consumers should be exactly 8 (the unsubscribe limit)
+        all_received = consumer_A_msgs + consumer_B_msgs + consumer_C_msgs
+        self.assertEqual(len(all_received), 8, f"Expected 8 total messages, got {len(all_received)}: {all_received}")
+
+        # All messages should be unique (no duplicates)
+        self.assertEqual(
+            len(all_received),
+            len(set(all_received)),
+            f"Found duplicate messages: {[msg for msg in all_received if all_received.count(msg) > 1]}",
+        )
+
+        # All received messages should be from our published set
+        expected_msgs = {f"msg-{i}" for i in range(num_msgs)}
+        received_msgs = set(all_received)
+        self.assertTrue(received_msgs.issubset(expected_msgs))
+
+        # Verify subscription reached its limit
+        self.assertEqual(sub._received, 8)
+        self.assertEqual(sub._max_msgs, 8)
+
+        await nc.close()
+
+    @async_test
+    async def test_subscribe_concurrent_next_msg_with_timeout(self):
+        """Test concurrent next_msg() calls with different timeout behaviors"""
+        nc = NATS()
+        await nc.connect()
+
+        sub = await nc.subscribe("test.concurrent.next.timeout")
+
+        # Publish only a few messages (less than what consumers will request)
+        num_msgs = 3
+        for i in range(num_msgs):
+            await nc.publish("test.concurrent.next.timeout", f"msg-{i}".encode())
+        await nc.flush()
+
+        # Track results and timing
+        consumer_results = {}
+
+        async def consumer_task(consumer_id: str, requests: int, timeout: float):
+            """Consumer that requests more messages than available"""
+            import time
+
+            received = []
+            timeouts = 0
+            start_time = time.time()
+
+            try:
+                for _ in range(requests):
+                    try:
+                        msg = await sub.next_msg(timeout=timeout)
+                        received.append(msg.data.decode())
+                    except Exception as e:
+                        if "timeout" in str(e).lower():
+                            timeouts += 1
+                        else:
+                            break
+
+                end_time = time.time()
+                consumer_results[consumer_id] = {
+                    "received": received,
+                    "timeouts": timeouts,
+                    "duration": end_time - start_time,
+                }
+            except Exception as e:
+                consumer_results[consumer_id] = f"Error: {e}"
+
+        # Start consumers with different timeout strategies
+        tasks = [
+            asyncio.create_task(consumer_task("fast_timeout", 5, 0.1)),  # Fast timeout
+            asyncio.create_task(consumer_task("medium_timeout", 5, 0.3)),  # Medium timeout
+            asyncio.create_task(consumer_task("slow_timeout", 5, 0.5)),  # Slow timeout
+        ]
+
+        # Wait for all consumers to finish
+        await asyncio.gather(*tasks)
+
+        # Verify results - collect all data first
+        all_received = []
+        total_timeouts = 0
+        consumers_with_msgs = 0
+
+        for consumer_id, result in consumer_results.items():
+            self.assertIsInstance(result, dict, f"Consumer {consumer_id} failed: {result}")
+
+            received = result["received"]
+            timeouts = result["timeouts"]
+
+            all_received.extend(received)
+            total_timeouts += timeouts
+
+            if len(received) > 0:
+                consumers_with_msgs += 1
+
+        # With only 3 messages and 3 consumers requesting 5 each, some distribution is expected
+        # But the key thing is that all 3 messages should be consumed
+        self.assertEqual(len(set(all_received)), 3, f"Expected 3 unique messages, got {set(all_received)}")
+
+        # There should be timeouts since we're requesting more messages than available
+        self.assertGreater(total_timeouts, 0, "Should have some timeouts when requesting more messages than available")
+
+        # At least one consumer should get messages (but due to race conditions, not necessarily all)
+        self.assertGreater(consumers_with_msgs, 0, "At least one consumer should receive messages")
+
+        await nc.close()
+
     @async_test
     async def test_subscribe_iterate_unsub_comprehension(self):
         nc = NATS()