Use asyncio.QueueShutdown when available

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

Author: wallyqs

nats/src/nats/aio/client.py

@@ -22,6 +22,7 @@
 import re
 import ssl
 import string
+import sys
 import time
 from collections import UserString
 from dataclasses import dataclass
@@ -59,6 +60,9 @@
 )
 from .transport import TcpTransport, Transport, WebSocketTransport
 
+# Python 3.13+ has QueueShutDown exception for cleaner queue termination
+_HAS_QUEUE_SHUTDOWN = sys.version_info >= (3, 13)
+
 try:
     from importlib.metadata import version
 
@@ -1797,15 +1801,19 @@ async def _process_msg(
             if sub._jsi:
                 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_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:
-                    # Queue might be full or closed, that's ok
-                    break
+        # Unblock waiting consumers after reaching max messages for non-callback subscriptions.
+        if max_msgs_reached and not sub._cb:
+            try:
+                if _HAS_QUEUE_SHUTDOWN:
+                    # Python 3.13+: Use queue shutdown for graceful termination.
+                    sub._pending_queue.shutdown()
+                elif sub._active_consumers is not None and sub._active_consumers > 0:
+                    # Python < 3.13: Send sentinels for each active consumer
+                    for _ in range(sub._active_consumers):
+                        sub._pending_queue.put_nowait(None)
+            except Exception:
+                # Queue might be closed or full.
+                pass
 
     def _build_message(
         self,

nats/src/nats/aio/subscription.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import asyncio
+import sys
 from typing import (
     TYPE_CHECKING,
     AsyncIterator,
@@ -31,6 +32,16 @@
 if TYPE_CHECKING:
     from nats.js import JetStreamContext
 
+# Python 3.13+ has QueueShutDown exception for cleaner queue termination.
+_HAS_QUEUE_SHUTDOWN = sys.version_info >= (3, 13)
+if _HAS_QUEUE_SHUTDOWN:
+    from asyncio import QueueShutDown
+else:
+    # For older Python versions, we'll use a custom exception
+    class QueueShutDown(Exception):
+        pass
+
+
 DEFAULT_SUB_PENDING_MSGS_LIMIT = 512 * 1024
 DEFAULT_SUB_PENDING_BYTES_LIMIT = 128 * 1024 * 1024
 
@@ -84,8 +95,10 @@ def __init__(
         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)
-        # Track active consumers (both async generators and next_msg calls) for non-callback subscriptions.
-        if cb is None:
+
+        # For Python < 3.13, we need to track active consumers for sentinel-based termination
+        # For Python 3.13+, we use QueueShutDown which doesn't require tracking.
+        if not _HAS_QUEUE_SHUTDOWN and cb is None:
             self._active_consumers = 0  # Counter of active consumers waiting for messages
         else:
             self._active_consumers = None
@@ -135,8 +148,10 @@ async def _message_generator(self) -> AsyncIterator[Msg]:
         Async generator that yields messages directly from the subscription queue.
         """
         yielded_count = 0
+
         if self._active_consumers is not None:
             self._active_consumers += 1
+
         try:
             while True:
                 # Check if subscription was cancelled/closed.
@@ -151,6 +166,8 @@ async def _message_generator(self) -> AsyncIterator[Msg]:
                     msg = await self._pending_queue.get()
                 except asyncio.CancelledError:
                     break
+                except QueueShutDown:
+                    break
 
                 # Check for sentinel value which signals generator to stop.
                 if msg is None:
@@ -224,7 +241,6 @@ 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
 
@@ -243,8 +259,11 @@ async def next_msg(self, timeout: Optional[float] = 1.0) -> Msg:
             if self._conn.is_closed:
                 raise errors.ConnectionClosedError
             raise
+        except QueueShutDown:
+            if self._conn.is_closed:
+                raise errors.ConnectionClosedError
+            raise errors.TimeoutError
         finally:
-            # Untrack this next_msg call.
             if self._active_consumers is not None:
                 self._active_consumers -= 1
 
@@ -352,14 +371,17 @@ 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()
 
-        # Send sentinels to unblock waiting consumers
+        # Unblock waiting consumers
         try:
-            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)
+            if self._pending_queue:
+                if _HAS_QUEUE_SHUTDOWN:
+                    # Python 3.13+: Use queue shutdown for graceful termination.
+                    self._pending_queue.shutdown()
+                elif self._active_consumers is not None and self._active_consumers > 0:
+                    # Python < 3.13: Send sentinels for each active consumer.
+                    for _ in range(self._active_consumers):
+                        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:
@@ -401,3 +423,5 @@ async def _wait_for_msgs(self, error_cb) -> None:
                     self._stop_processing()
             except asyncio.CancelledError:
                 break
+            except QueueShutDown:
+                break

nats/src/nats/js/client.py

@@ -883,15 +883,14 @@ def __init__(
             self._cb = sub._cb
             self._future = sub._future
             self._closed = sub._closed
-            self._active_generators = sub._active_generators
+            self._active_consumers = sub._active_consumers
 
             # 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._pending_next_msgs_calls = sub._pending_next_msgs_calls
 
         async def consumer_info(self) -> api.ConsumerInfo:
             """