Repository: maxfischer2781/asyncstdlib

Issue:
`tee()`'s iterators throw RuntimeError when consumed in parallel
Thanks for this library, I've been using it on a project recently and found it really useful.

The `tee()` function seems to have a bug. If the source is an async generator that awaits in-between yields, and the tee's iterators are consumed in parallel, the tee's iterators will throw `RuntimeError: anext(): asynchronous generator is already running`.

What's happening is that one of the tee()'s peer iterators calls `__anext__()` on the source, then while that call is being awaited, a second peer iterator also calls `__anext__()`. If the source is an async generator, it will throw a RuntimeError, as overlapping calls are not allowed.

Here's a test case that triggers the error: https://github.com/h4l/asyncstdlib/commit/680429354fce7ad22c535df74ff503c9f38ec313

And an asyncio program that triggers it:
```python
# tee_bug_asyncio_repo.py
import asyncio
import sys

import asyncstdlib as a


async def slow_iterable():
    for x in range(3):
        await asyncio.sleep(0.1)
        yield x


async def main(*, n):
    print(f"tee() with {n=}")
    async with a.tee(slow_iterable(), n=n) as iterables:
        results = await asyncio.gather(
            *(asyncio.create_task(a.list(iterable)) for iterable in iterables)
        )
        print(results)


if __name__ == "__main__":
    n = 2 if len(sys.argv) < 2 else int(sys.argv[1])
    asyncio.run(main(n=n))
```

```
$ python tee_bug_asyncio_repo.py 1
tee() with n=1
[[0, 1, 2]]

$ python tee_bug_asyncio_repo.py 2
tee() with n=2
Traceback (most recent call last):
  File "/workspaces/asyncstdlib/tee_bug_asyncio_repo.py", line 16, in main
    results = await asyncio.gather(
  File "/workspaces/asyncstdlib/asyncstdlib/builtins.py", line 656, in list
    return [element async for element in aiter(iterable)]
  File "/workspaces/asyncstdlib/asyncstdlib/builtins.py", line 656, in <listcomp>
    return [element async for element in aiter(iterable)]
  File "/workspaces/asyncstdlib/asyncstdlib/itertools.py", line 307, in tee_peer
    item = await iterator.__anext__()
RuntimeError: anext(): asynchronous generator is already running

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/workspaces/asyncstdlib/tee_bug_asyncio_repo.py", line 24, in <module>
    asyncio.run(main(n=n))
  File "/usr/local/lib/python3.10/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/usr/local/lib/python3.10/asyncio/base_events.py", line 646, in run_until_complete
    return future.result()
  File "/workspaces/asyncstdlib/tee_bug_asyncio_repo.py", line 15, in main
    async with a.tee(slow_iterable(), n=n) as iterables:
  File "/workspaces/asyncstdlib/asyncstdlib/itertools.py", line 390, in __aexit__
    await self.aclose()
  File "/workspaces/asyncstdlib/asyncstdlib/itertools.py", line 395, in aclose
    await child.aclose()
RuntimeError: aclose(): asynchronous generator is already running
```

Comments:

Patch:
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 3bfabfc..f55528c 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -18,9 +18,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: '3.x'
     - name: Install dependencies
diff --git a/.github/workflows/verification.yml b/.github/workflows/verification.yml
index 4190131..f7627d3 100644
--- a/.github/workflows/verification.yml
+++ b/.github/workflows/verification.yml
@@ -6,9 +6,9 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: '3.9'
     - name: Install dependencies
diff --git a/asyncstdlib/itertools.py b/asyncstdlib/itertools.py
index 5c64ab7..44f060c 100644
--- a/asyncstdlib/itertools.py
+++ b/asyncstdlib/itertools.py
@@ -17,7 +17,7 @@ from typing import (
 )
 from collections import deque
 
-from ._typing import T, R, T1, T2, T3, T4, T5, AnyIterable, ADD
+from ._typing import T, R, T1, T2, T3, T4, T5, AnyIterable, ADD, AsyncContextManager
 from ._utility import public_module
 from ._core import (
     ScopedIter,
@@ -294,27 +294,45 @@ async def takewhile(
                 break
 
 
+class NoLock:
+    """Dummy lock that provides the proper interface but no protection"""
+
+    async def __aenter__(self) -> None:
+        pass
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        return False
+
+
 async def tee_peer(
     iterator: AsyncIterator[T],
+    # the buffer specific to this peer
     buffer: Deque[T],
+    # the buffers of all peers, including our own
     peers: List[Deque[T]],
+    lock: AsyncContextManager[Any],
 ) -> AsyncGenerator[T, None]:
     """An individual iterator of a :py:func:`~.tee`"""
     try:
         while True:
             if not buffer:
-                try:
-                    item = await iterator.__anext__()
-                except StopAsyncIteration:
-                    break
-                else:
-                    # Append to all buffers, including our own. We'll fetch our
-                    # item from the buffer again, instead of yielding it directly.
-                    # This ensures the proper item ordering if any of our peers
-                    # are fetching items concurrently. They may have buffered their
-                    # item already.
-                    for peer_buffer in peers:
-                        peer_buffer.append(item)
+                async with lock:
+                    # Another peer produced an item while we were waiting for the lock.
+                    # Proceed with the next loop iteration to yield the item.
+                    if buffer:
+                        continue
+                    try:
+                        item = await iterator.__anext__()
+                    except StopAsyncIteration:
+                        break
+                    else:
+                        # Append to all buffers, including our own. We'll fetch our
+                        # item from the buffer again, instead of yielding it directly.
+                        # This ensures the proper item ordering if any of our peers
+                        # are fetching items concurrently. They may have buffered their
+                        # item already.
+                        for peer_buffer in peers:
+                            peer_buffer.append(item)
             yield buffer.popleft()
     finally:
         # this peer is done – remove its buffer
@@ -322,6 +340,7 @@ async def tee_peer(
             if peer_buffer is buffer:
                 peers.pop(idx)
                 break
+        # if we are the last peer, try and close the iterator
         if not peers and hasattr(iterator, "aclose"):
             await iterator.aclose()  # type: ignore
 
@@ -355,14 +374,23 @@ class Tee(Generic[T]):
     If ``iterable`` is an iterator and read elsewhere, ``tee`` will *not*
     provide these items. Also, ``tee`` must internally buffer each item until the
     last iterator has yielded it; if the most and least advanced iterator differ
-    by most data, using a :py:class:`list` is faster (but not lazy).
+    by most data, using a :py:class:`list` is more efficient (but not lazy).
 
     If the underlying iterable is concurrency safe (``anext`` may be awaited
     concurrently) the resulting iterators are concurrency safe as well. Otherwise,
     the iterators are safe if there is only ever one single "most advanced" iterator.
+    To enforce sequential use of ``anext``, provide a ``lock``
+    - e.g. an :py:class:`asyncio.Lock` instance in an :py:mod:`asyncio` application -
+    and access is automatically synchronised.
     """
 
-    def __init__(self, iterable: AnyIterable[T], n: int = 2):
+    def __init__(
+        self,
+        iterable: AnyIterable[T],
+        n: int = 2,
+        *,
+        lock: Optional[AsyncContextManager[Any]] = None,
+    ):
         self._iterator = aiter(iterable)
         self._buffers: List[Deque[T]] = [deque() for _ in range(n)]
         self._children = tuple(
@@ -370,6 +398,7 @@ class Tee(Generic[T]):
                 iterator=self._iterator,
                 buffer=buffer,
                 peers=self._buffers,
+                lock=lock if lock is not None else NoLock(),
             )
             for buffer in self._buffers
         )
@@ -377,7 +406,17 @@ class Tee(Generic[T]):
     def __len__(self) -> int:
         return len(self._children)
 
+    @overload
     def __getitem__(self, item: int) -> AsyncIterator[T]:
+        ...
+
+    @overload
+    def __getitem__(self, item: slice) -> Tuple[AsyncIterator[T], ...]:
+        ...
+
+    def __getitem__(
+        self, item: Union[int, slice]
+    ) -> Union[AsyncIterator[T], Tuple[AsyncIterator[T], ...]]:
         return self._children[item]
 
     def __iter__(self) -> Iterator[AnyIterable[T]]:
diff --git a/docs/source/api/itertools.rst b/docs/source/api/itertools.rst
index 739573c..2610ed0 100644
--- a/docs/source/api/itertools.rst
+++ b/docs/source/api/itertools.rst
@@ -65,9 +65,13 @@ Iterator transforming
 Iterator splitting
 ==================
 
-.. autofunction:: tee(iterable: (async) iter T, n: int = 2)
+.. autofunction:: tee(iterable: (async) iter T, n: int = 2, [*, lock: async with Any])
     :for: :(async iter T, ...)
 
+    .. versionadded:: 3.10.5
+
+        The ``lock`` keyword parameter.
+
 .. autofunction:: pairwise(iterable: (async) iter T)
     :async-for: :(T, T)
 


Test patch:
diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index fcb4308..889ece1 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -10,9 +10,9 @@ jobs:
         python-version: ['3.6', '3.7', '3.8', '3.9', '3.10', 'pypy-3.6', 'pypy-3.7']
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
diff --git a/unittests/test_itertools.py b/unittests/test_itertools.py
index 1a80065..904378d 100644
--- a/unittests/test_itertools.py
+++ b/unittests/test_itertools.py
@@ -1,10 +1,11 @@
 import itertools
+import sys
 
 import pytest
 
 import asyncstdlib as a
 
-from .utility import sync, asyncify, awaitify
+from .utility import sync, asyncify, awaitify, multi_sync, Schedule, Switch, Lock
 
 
 @sync
@@ -210,6 +211,56 @@ async def test_tee():
             assert await a.list(iterator) == iterable
 
 
+@multi_sync
+async def test_tee_concurrent_locked():
+    """Test that properly uses a lock for synchronisation"""
+    items = [1, 2, 3, -5, 12, 78, -1, 111]
+
+    async def iter_values():
+        for item in items:
+            # switch to other tasks a few times to guarantees another runs
+            for _ in range(5):
+                await Switch()
+            yield item
+
+    async def test_peer(peer_tee):
+        assert await a.list(peer_tee) == items
+
+    head_peer, *peers = a.tee(iter_values(), n=len(items) // 2, lock=Lock())
+    await Schedule(*map(test_peer, peers))
+    await Switch()
+    results = [item async for item in head_peer]
+    assert results == items
+
+
+# see https://github.com/python/cpython/issues/74956
+@pytest.mark.skipif(
+    sys.version_info < (3, 8),
+    reason="async generators only protect against concurrent access since 3.8",
+)
+@multi_sync
+async def test_tee_concurrent_unlocked():
+    """Test that does not prevent concurrency without a lock"""
+    items = list(range(12))
+
+    async def iter_values():
+        for item in items:
+            # switch to other tasks a few times to guarantees another runs
+            for _ in range(5):
+                await Switch()
+            yield item
+
+    async def test_peer(peer_tee):
+        assert await a.list(peer_tee) == items
+
+    this, peer = a.tee(iter_values(), n=2)
+    await Schedule(test_peer(peer))
+    await Switch()
+    # underlying generator raises RuntimeError when `__anext__` is interleaved
+    with pytest.raises(RuntimeError):
+        await test_peer(this)
+
+
 @sync
 async def test_pairwise():
     assert await a.list(a.pairwise(range(5))) == [(0, 1), (1, 2), (2, 3), (3, 4)]
diff --git a/unittests/utility.py b/unittests/utility.py
index 2710fe1..b475eda 100644
--- a/unittests/utility.py
+++ b/unittests/utility.py
@@ -85,12 +85,34 @@ class Switch:
         yield self
 
 
+class Lock:
+    def __init__(self):
+        self._owned = False
+        self._waiting = []
+
+    async def __aenter__(self):
+        if self._owned:
+            # wait until it is our turn to take the lock
+            token = object()
+            self._waiting.append(token)
+            while self._owned or self._waiting[0] is not token:
+                await Switch()
+            # take the lock and remove our wait claim
+            self._owned = True
+            self._waiting.pop(0)
+        self._owned = True
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._owned = False
+
+
 def multi_sync(test_case: Callable[..., Coroutine]):
     """
-    Mark an ``async def`` test case to be run synchronously with chicldren
+    Mark an ``async def`` test case to be run synchronously with children
 
     This emulates a primitive "event loop" which only responds
-    to the :py:class:`PingPong`, :py:class:`Schedule` and :py:class:`Switch`.
+    to the :py:class:`PingPong`, :py:class:`Schedule`, :py:class:`Switch`
+    and :py:class:`Lock`.
     """
 
     @wraps(test_case)
@@ -103,7 +125,7 @@ def multi_sync(test_case: Callable[..., Coroutine]):
                 event = coro.send(event)
             except StopIteration as e:
                 result = e.args[0] if e.args else None
-                assert result is None
+                assert result is None, f"got '{result!r}' expected 'None'"
             else:
                 if isinstance(event, PingPong):
                     run_queue.appendleft((coro, event))