Tee and avoiding memory leaks

[tsoos99dev]

Hi. What would be a good way to make sure the iterators returned by tee don't get too far out of sync? I'd love to have a version where I can supply the max number of items the most and least advanced iterators can differ by.

[maxfischer2781]

The tee implementation doesn't really have a neat way to check the buffer lengths, especially with #182 on the way.

You're probably best off writing your own group-of-iterators class where each iterator tracks their index and can compare it against its peers. For example, a (draft, untested) utility class would look like this:

from typing import AsyncIterator

class CohesiveAsyncIterators[T]:
    def __init__(self, gap_limit: int) -> None:
        self.gap_limit = gap_limit
        self._peers: dict[object, int] = {}

    async def attach(
        self, peer: AsyncIterator[T], /, index: int = 0
    ) -> AsyncIterator[T]:
        token = object()
        self._peers[token] = index
        try:
            async for item in peer:
                self._peers[token] += 1
                yield item
                # check that other peers are not too far behind before proceeding
                if self._peers[token] > min(self._peers.values()) - self.gap_limit:
                    # could also be a warning or simply waiting for an Event to show we are back in sync
                    raise RuntimeError(f"Iterator {peer} advanced more than {self.gap_limit} from peers")
        finally:
            del self._peers[token]

One would just attach whatever iterators need to stay within the limit. So for the case of tee, that would be something like:

import asyncstdlib as a

def bounded_tee(iterable, n: int, bound: int):
    return map(CohesiveAsyncIterators(gap_limit =bound).attach, a.tee(iterable, n=n))