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

[DeltaEpsilon7787]

Observed behavior

The way async iterators are implemented, they have an unacceptable drop in throughput compared to callback-based subscription. Checking the code of how they're implemented, they appear to be doing a lot of unnecessary actions for each response in __anext__.

Using code in Steps to Reproduce, on my PC I get 1.29 seconds for callbacks and 4.6 seconds for async iter, which translates to 310k msg/sec for callbacks and 87k msg/sec for built-in async its. This is a 3.5x drop in throughput.

Expected behavior

This should either be remedied or explicitly mentioned somewhere as this is relatively unexpected behavior. One would expect the async iterator to, well, be a thin wrapper around the pending queue.

In the code provided in Steps to Reproduce, I implement a very basic custom async iterator that uses callbacks to fill in an internal queue that the async it then gets results from. Given that callback impl and async it impl are mutually exclusive in current API, this does not break any API preconditions to my knowledge.
This implementation only takes 1.4 seconds which is only ~5% slower than just callbacks, although it is missing various boilerplate for stopping the iteration.

Server and client version

nats-server: v2.10.5
nats-py: v2.7.0

Host environment

No response

Steps to reproduce

from time import perf_counter

import asyncio as aio
import multiprocessing as mp

import nats

TEST_AMOUNT = 400000


class CbMessageConsumer:
    def __init__(self, target):
        self.counter = 0

        self._target = target
        self._completed = aio.Event()

    async def cb(self, msg):
        self.counter += 1
        if self.counter >= self._target:
            self._completed.set()

    async def completion_observer(self):
        await self._completed.wait()


class IterMessageConsumer:
    def __init__(self, target):
        self.counter = 0

        self._target = target
        self._completed = aio.Event()

    async def do_work(self, source):
        async for msg in source:
            self.counter += 1
            if self.counter >= self._target:
                self._completed.set()
                break

    async def completion_observer(self):
        await self._completed.wait()


# Custom `async for` using callbacks
class FasterAsyncIt:
    def __init__(self):
        self._queue = aio.Queue()

    async def _queue_cb(self, msg):
        self._queue.put_nowait(msg)

    def __aiter__(self):
        return self

    async def __anext__(self):
        # Ommitted checking for closed channel and such
        return await self._queue.get()


def cb_receiver():
    async def _():
        nc = await nats.connect("127.0.0.1:4222")

        cb_message_consumer = CbMessageConsumer(TEST_AMOUNT)
        cb_observer_task = aio.create_task(cb_message_consumer.completion_observer())
        cb_observer_task.add_done_callback(
            lambda _: print(f"CB: {perf_counter() - start_time}")
        )

        start_time = perf_counter()
        await nc.subscribe("test.cb", cb=cb_message_consumer.cb)

        await cb_observer_task
        await nc.close()

    aio.run(_())


def iter_receiver():
    async def _():
        nc = await nats.connect("127.0.0.1:4222")

        it_message_consumer = IterMessageConsumer(TEST_AMOUNT)
        it_observer_task = aio.create_task(it_message_consumer.completion_observer())
        it_observer_task.add_done_callback(
            lambda _: print(f"IT: {perf_counter() - start_time}")
        )

        it_based_sub = await nc.subscribe("test.iter")
        start_time = perf_counter()
        aio.create_task(it_message_consumer.do_work(it_based_sub.messages))

        await it_observer_task
        await nc.close()

    aio.run(_())

def custom_iter_receiver():
    async def _():
        nc = await nats.connect("127.0.0.1:4222")

        custom_it = FasterAsyncIt()

        custom_it_message_consumer = IterMessageConsumer(TEST_AMOUNT)
        custom_it_observer_task = aio.create_task(
            custom_it_message_consumer.completion_observer()
        )
        custom_it_observer_task.add_done_callback(
            lambda _: print(f"CUSTOM IT: {perf_counter() - start_time}")
        )

        await nc.subscribe("test.custom_iter", cb=custom_it._queue_cb)
        start_time = perf_counter()
        aio.create_task(custom_it_message_consumer.do_work(custom_it))
        await custom_it_observer_task
        await nc.close()

    aio.run(_())


def sender_side():
    async def _():
        nc = await nats.connect("127.0.0.1:4222")

        source = iter(range(TEST_AMOUNT))

        try:
            while True:
                for _ in range(30000):
                    next(source)
                    await nc.publish("test.cb", b"")
                    await nc.publish("test.iter", b"")
                    await nc.publish("test.custom_iter", b"")
                # Throttle to avoid slow consumers
                await aio.sleep(0.01)

        except StopIteration:
            pass

        await nc.close()

    aio.run(_())


if __name__ == "__main__":
    sender = mp.Process(target=sender_side)
    receiver_1 = mp.Process(target=cb_receiver)
    receiver_2 = mp.Process(target=iter_receiver)
    receiver_3 = mp.Process(target=custom_iter_receiver)

    receiver_1.start()
    receiver_2.start()
    receiver_3.start()
    sender.start()

    sender.join()
    receiver_1.join()
    receiver_2.join()
    receiver_3.join()

[chrisprobst]

Checking the implementation it is clear that asyncio.get_running_loop().create_task(self._queue.get()) is the problem here. Every iteration creates a new tasks. That's relatively heavy for python to do (similar to starting a goroutine per read).

Might be interesting if the impl. can be improved. The problem seems to be that waiting on the next read should be interruptable by the self._unsubscribed_future.

It's easy to verify this by changing the impl. to:

    async def __anext__(self) -> Msg:        
        msg = await self._queue.get()
        sub = self._sub    
        self._queue.task_done()
        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

If you try this, the performance is the same as with callback style but you won't be interrupted if unsubscribed. Maybe some internal developer can give guidance on potential improvements?

[chrisprobst]

After some digging, I think a very simple "fast path" already solves this "high bandwidth" scenario:

    def _process_msg(self, msg: Msg) -> Msg:
        self._queue.task_done()
        self._sub._pending_size -= len(msg.data)

        # Unblock the iterator in case it has already received enough messages.
        sub = self._sub
        if sub._max_msgs > 0 and sub._received >= sub._max_msgs:
            self._cancel()
        return msg

    async def __anext__(self) -> Msg:
        if not self._queue.empty(): # This is an enormous speed up in case messages are already queued.
            return self._process_msg(self._queue.get_nowait())

        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
        )

        if get_task in finished:
            return self._process_msg(get_task.result())            
        elif self._unsubscribed_future.done():
            get_task.cancel()

        raise StopAsyncIteration

Seems simple to solve.

[chrisprobst]

This also applies to sub.next_msg. I could do a PR for this, let me know if this is a wanted change or not.

@wallyqs

EDIT: The object_store.get internally also uses the iterator. Therefore, improving the iterator performance would automatically help with high bandwidth object downloads as well.

[wallyqs]

Addressed in this PR: #752