Reduce memory footprint of await many dispatch

[sevdog]

Currently the await_many_dispatch method keeps references of completed task and results until a new task is "done", this slows down the gargage collection of result (which is also referenced by task).

This PR points to reduce the delay in garbage collection of result by removing the support variable results and by cleaning the task variable after use.

[carltongibson]

Hey @sevdog. Thanks for the suggestion. I'm wondering, did you measure this at all? What made you think to suggest the change?

[sevdog]

Hi @carltongibson, I found this while working at django/channels_redis#426 when I was testing some AI tools to see if it coould find something useful to help with #2094.

Here is a benchmark script to test memory utilization with various implementations of `await_many_dispatch`

import argparse
import asyncio
import gc
import tracemalloc

# ---------------------------------------------------------------------------
# variants of channels/utils.py
# ---------------------------------------------------------------------------
async def _cleanup(tasks):
    for task in tasks:
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass


async def original(consumer_callables, dispatch):
    tasks = [asyncio.ensure_future(cc()) for cc in consumer_callables]
    try:
        while True:
            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for i, task in enumerate(tasks):
                if task.done():
                    result = task.result()
                    await dispatch(result)
                    tasks[i] = asyncio.ensure_future(consumer_callables[i]())
    finally:
        await _cleanup(tasks)


async def inline(consumer_callables, dispatch):
    tasks = [asyncio.ensure_future(cc()) for cc in consumer_callables]
    try:
        while True:
            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for i, task in enumerate(tasks):
                if task.done():
                    await dispatch(task.result())
                    tasks[i] = asyncio.ensure_future(consumer_callables[i]())
    finally:
        await _cleanup(tasks)


async def indexed(consumer_callables, dispatch):
    tasks = [asyncio.ensure_future(cc()) for cc in consumer_callables]
    try:
        while True:
            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for i in range(len(tasks)):
                if tasks[i].done():
                    await dispatch(tasks[i].result())
                    tasks[i] = asyncio.ensure_future(consumer_callables[i]())
    finally:
        await _cleanup(tasks)


async def enumerate_and_del(consumer_callables, dispatch):
    tasks = [asyncio.ensure_future(cc()) for cc in consumer_callables]
    try:
        while True:
            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for i, task in enumerate(tasks):
                if task.done():
                    await dispatch(task.result())
                    del task
                    tasks[i] = asyncio.ensure_future(consumer_callables[i]())
    finally:
        await _cleanup(tasks)


VARIANTS = {
    "original": original,
    "inline": inline,
    "indexed": indexed,
    "enumerate_and_del": enumerate_and_del,
}

# ---------------------------------------------------------------------------
# Simulate a "connection": big message, forever idle
# ---------------------------------------------------------------------------
def make_source(payload_bytes):
    sent = False

    async def source():
        nonlocal sent
        if not sent:
            sent = True
            return {"type": "websocket.receive", "bytes": bytes(payload_bytes)}
        await asyncio.Future()  # idle: never resolved

    return source


async def noop_dispatch(message):
    # fast handler which keeps nothing
    return


async def main(variant, n_conns, payload_kib, ramp_batches, batch_delay):
    loop_fn = VARIANTS[variant]
    payload = payload_kib * 1024
    conns = []
    per_batch = max(1, n_conns // ramp_batches)
    for _ in range(ramp_batches):
        for _ in range(per_batch):
            conns.append(
                asyncio.ensure_future(
                    loop_fn([make_source(payload)], noop_dispatch)
                )
            )
        await asyncio.sleep(batch_delay)  # let process msg and go idle
        gc.collect()

    # plateau: here we measure allocated memory
    gc.collect()
    current, peak = tracemalloc.get_traced_memory()
    expected = n_conns * payload / 1e6
    print(
        f"[{variant}] conns={n_conns} payload={payload_kib}KiB -> "
        f"kept={current / 1e6:7.1f} MB | ppeak={peak / 1e6:7.1f} MB "
        f"| expected~{expected:.1f} MB if 1 msg/conn kept"
    )

    await asyncio.sleep(2.0)  # keep plateau for memray
    for c in conns:
        c.cancel()
    await asyncio.gather(*conns, return_exceptions=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("variant", choices=list(VARIANTS))
    ap.add_argument("--conns", type=int, default=2000)
    ap.add_argument("--payload-kib", type=int, default=64)
    ap.add_argument("--ramp-batches", type=int, default=10)
    ap.add_argument("--batch-delay", type=float, default=0.4)
    args = ap.parse_args()

    tracemalloc.start()
    asyncio.run(
        main(args.variant, args.conns, args.payload_kib, args.ramp_batches, args.batch_delay)
    )

The first two implementations shows the same memory footprint:

While the last two shows this memory footprint:

The reason is because when task is a local variable it holds a reference to result (because it is a Future and thus has a reference to it when it is available), which cannot be cleared until the for loop re-evaluates and the variable is assigned again.

Obviously the memory depends on message size, with small messages and frequent updates this may not be noted and could be considered just a corner-case. This comes into hand with low-frequencey and huge messages.

[carltongibson]

Thank you @sevdog — that's exactly the context I was hoping for. 🎁

Let me take a look.