Improve perf of async iterators #752

wallyqs · 2025-10-14T08:36:25Z

Rework implementation of async for msg in sub.messages and sub.next_msg to get better perf.

Fixes #703

nats/src/nats/aio/subscription.py

caspervonb

Better than earlier results discussed in #732, which peaked at 20k per second.

uv run ../nats-client/tools/bench.py --client=aio --msgs=1000000

Starting pub/sub benchmark with nats.aio [msgs=1,000,000, size=128 B]

Publisher results:
Test completed: 1,000,000 messages, 128,000,000 bytes, 3.84 seconds
  Throughput: 260,387 msgs/sec, 31.79 MB/sec
  Latency: (min/avg/max/std) = 0.00/0.00/52.37/0.36 ms

Subscriber results:
Test completed: 1,000,000 messages, 128,000,000 bytes, 3.88 seconds
  Throughput: 257,908 msgs/sec, 31.48 MB/sec
  Latency: (min/avg/max/std) = 99.45/2014.35/3976.80/1104.45 ms

nats/src/nats/aio/subscription.py

caspervonb

Tests were relaxed, so potential timing issues?
Technical debt for the whole comparability thing.
Why generator, instead of just calling queue.get() etc in anext?

wallyqs · 2025-10-14T16:55:08Z

Removed the compatibility wrapper with previous iterator since no longer used by the object store. Had to tweak the defaults from test_pull_subscribe_limits since it now processes more messages than before. About anext, I would have to take a closer look whether can use that instead.

Signed-off-by: Waldemar Quevedo <wally@nats.io>

caspervonb

Would be good if next and anext had the same call path, no?

caspervonb · 2025-10-17T10:05:11Z

nats/src/nats/aio/subscription.py

        # require tracking the next_msg calls inflight for cancelling.
        if cb is None:
-            self._pending_next_msgs_calls = {}
+            self._pending_next_msgs_calls: Optional[Dict[str, asyncio.Task]] = {}


So ideally, both next_msg and anext are basically the same call, aside from options like timeout etc.

think it could do sub.next_msg(timeout=0) though it is only the peeking from the queue part that they have in common, it should be fine probably

Signed-off-by: Waldemar Quevedo <wally@nats.io>

Copilot

Pull Request Overview

This PR refactors the NATS subscription message iteration mechanism to use a more efficient async generator pattern. The changes replace the custom _SubscriptionMessageIterator class with a simpler generator-based approach while maintaining compatibility across Python 3.8-3.13.

Key changes:

Replaces _SubscriptionMessageIterator with _message_generator() for cleaner async iteration
Adds Python 3.13+ QueueShutDown support with backwards compatibility for older versions
Implements sentinel-based queue termination for Python < 3.13
Adds comprehensive test coverage for concurrent consumption patterns and edge cases

Reviewed Changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 17 comments.

Show a summary per file

File	Description
nats/src/nats/aio/subscription.py	Refactored message iteration to use async generators with version-specific queue shutdown handling
nats/src/nats/aio/client.py	Updated client cleanup to use new sentinel-based termination approach
nats/src/nats/js/client.py	Updated JetStream subscription to reference new `_active_consumers` field
nats/src/nats/js/object_store.py	Changed from `_message_iterator` to `messages` property
nats/tests/test_client.py	Added extensive test coverage for async generators, concurrent consumption, and edge cases
nats/tests/test_js.py	Adjusted assertions to handle timing variability in tests
nats/benchmark/*.py	Added new benchmark scripts for performance testing
.github/workflows/test.yml	Added additional NATS server versions to test matrix

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-10-30T21:16:48Z

nats/src/nats/aio/subscription.py

+                if self._max_msgs > 0 and yielded_count >= self._max_msgs:
+                    break


This check is redundant as the same condition is checked again at lines 184-185 after yielding the message. Consider removing this pre-check since the subscription should still process any messages already in the queue before stopping, and the post-yield check is sufficient to enforce the limit.

Suggested change

if self._max_msgs > 0 and yielded_count >= self._max_msgs:

break

Copilot · 2025-10-30T21:16:48Z

nats/tests/test_client.py

+            print("WRN", e)
+            pass


This debug print statement should be removed or replaced with proper logging. Using print() in tests can clutter test output and the exception is being silently swallowed after printing.

Suggested change

print("WRN", e)

pass

raise

Copilot · 2025-10-30T21:16:48Z

nats/src/nats/aio/subscription.py

+        # 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:


The condition not _HAS_QUEUE_SHUTDOWN and cb is None means that for Python 3.13+ with callback subscriptions, _active_consumers is set to None. However, this is inconsistent with the check at line 886 in js/client.py which copies this field. Consider documenting why callback-based subscriptions don't need active consumer tracking, or simplify the logic to only check cb is None.

Suggested change

# 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:

# Only synchronous (non-callback) subscriptions need active consumer tracking,

# regardless of Python version. Callback-based subscriptions do not require it.

if cb is None:

Copilot · 2025-10-30T21:16:49Z

nats/tests/test_client.py

+
+        async def consumer_task(consumer_id: str, max_messages: int = None):
+            """Consumer task that processes messages"""
+            import random


The random module should be imported at the top of the file rather than inside a function. Move this import to the file-level imports section for better code organization and to avoid repeated imports in concurrent tasks.

Copilot · 2025-10-30T21:16:49Z

nats/tests/test_client.py

+
+        async def consumer_task(consumer_id: str, max_messages: int = None):
+            """Consumer task that processes messages"""
+            import random


The random module should be imported at the top of the file rather than inside a function. Move this import to the file-level imports section for better code organization and to avoid repeated imports in concurrent tasks.

Copilot · 2025-10-30T21:16:51Z

nats/src/nats/aio/subscription.py

+                sentinels_to_send = max(1, self._active_consumers)
+                for _ in range(sentinels_to_send):
+                    self._pending_queue.put_nowait(None)
+        except Exception:


'except' clause does nothing but pass and there is no explanatory comment.

Copilot · 2025-10-30T21:16:51Z

nats/tests/test_client.py


+        # Wait for iterator to complete after drain
        await asyncio.wait_for(fut, 1)
+        await iterator_task  # Ensure task cleanup


This statement has no effect.

Copilot · 2025-10-30T21:16:52Z

nats/tests/test_client.py

+        self.assertEqual(msg2.data, b"timeout_zero_msg")
+
+        # Clean up
+        await publish_task


This statement has no effect.

Copilot · 2025-10-30T21:16:52Z

nats/tests/test_client.py

+
+        # Clean up
+        await publish_task
+        await publish_task2


This statement has no effect.

Copilot · 2025-10-30T21:16:52Z

nats/tests/test_client.py

+            await asyncio.gather(*tasks)
+        except Exception as e:
+            print("WRN", e)
+            pass


Unnecessary 'pass' statement.

Suggested change

pass

wallyqs force-pushed the sub-next-msg-perf branch from b1a68b3 to 8626364 Compare October 14, 2025 08:42

This was referenced Oct 14, 2025

Add nats-client package #732

Open

Using the async iterator is significantly slower compared to using callbacks in a subscription #541

Open

wallyqs force-pushed the sub-next-msg-perf branch 2 times, most recently from 7e973e9 to 2e1aba4 Compare October 14, 2025 11:04

Jarema requested a review from caspervonb October 14, 2025 12:36

caspervonb reviewed Oct 14, 2025

View reviewed changes

nats/src/nats/aio/subscription.py Outdated Show resolved Hide resolved

caspervonb reviewed Oct 14, 2025

View reviewed changes

nats/src/nats/aio/subscription.py Outdated Show resolved Hide resolved

caspervonb reviewed Oct 14, 2025

View reviewed changes

wallyqs force-pushed the sub-next-msg-perf branch from 00a7074 to 0d14e55 Compare October 14, 2025 16:04

wallyqs force-pushed the sub-next-msg-perf branch 2 times, most recently from cecb078 to ce7dd49 Compare October 17, 2025 09:45

wallyqs added 6 commits October 17, 2025 03:00

Improve AsyncIterator usage

8cf5c87

Signed-off-by: Waldemar Quevedo <wally@nats.io>

Improve object store get usage

9d2a26c

Signed-off-by: Waldemar Quevedo <wally@nats.io>

Add obj get perf bench

e5297b1

Signed-off-by: Waldemar Quevedo <wally@nats.io>

Improve sub.next_msg, allow timeout=0 to block forever for next msg

5cae16f

Signed-off-by: Waldemar Quevedo <wally@nats.io>

Remove _message_iterator since not used by obj store anymore

2dc2031

Signed-off-by: Waldemar Quevedo <wally@nats.io>

Handle cancellation of concurrent async generators

54642f2

Signed-off-by: Waldemar Quevedo <wally@nats.io>

wallyqs force-pushed the sub-next-msg-perf branch from ce7dd49 to 54642f2 Compare October 17, 2025 10:00

caspervonb reviewed Oct 17, 2025

View reviewed changes

Track next_msg and async for generators the same way

adc5c10

Signed-off-by: Waldemar Quevedo <wally@nats.io>

wallyqs force-pushed the sub-next-msg-perf branch from bdb4b87 to 71751ed Compare October 20, 2025 19:38

Use asyncio.QueueShutdown when available

87735f1

Signed-off-by: Waldemar Quevedo <wally@nats.io>

wallyqs force-pushed the sub-next-msg-perf branch from 71751ed to 87735f1 Compare October 20, 2025 23:07

Add more nats-server versions

e245add

Signed-off-by: Waldemar Quevedo <wally@nats.io>

wallyqs force-pushed the sub-next-msg-perf branch from 03aedef to e245add Compare October 21, 2025 04:50

Test fixes

d943ec9

Signed-off-by: Waldemar Quevedo <wally@nats.io>

caspervonb requested a review from Copilot October 30, 2025 21:13

Copilot AI reviewed Oct 30, 2025

View reviewed changes

		if self._max_msgs > 0 and yielded_count >= self._max_msgs:
		break

Improve perf of async iterators #752

Are you sure you want to change the base?

Improve perf of async iterators #752

Uh oh!

Conversation

wallyqs commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

caspervonb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

caspervonb left a comment

Choose a reason for hiding this comment

Uh oh!

wallyqs commented Oct 14, 2025

Uh oh!

caspervonb left a comment

Choose a reason for hiding this comment

Uh oh!

caspervonb Oct 17, 2025

Choose a reason for hiding this comment

Uh oh!

wallyqs Oct 17, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wallyqs commented Oct 14, 2025 •

edited

Loading

caspervonb left a comment •

edited

Loading