Tee locking (#87)

* optional lock-protection for tee

* test loop supports locks

* testing tee with and without locks

* skip concurrency test for old Python versions

* documented lock parameter

* added annotations for slicing a tee

* update actions versions

Author: maxfischer2781

.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

.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

.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

asyncstdlib/itertools.py

@@ -17,7 +17,7 @@
 )
 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,34 +294,53 @@ 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
         for idx, peer_buffer in enumerate(peers):  # pragma: no branch
             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,29 +374,49 @@ async def derivative(sensor_data):
     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(
             tee_peer(
                 iterator=self._iterator,
                 buffer=buffer,
                 peers=self._buffers,
+                lock=lock if lock is not None else NoLock(),
             )
             for buffer in self._buffers
         )
 
     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]]:

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)
 

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)]

unittests/utility.py

@@ -85,12 +85,34 @@ def __await__(self):
         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 run_sync(*args, **kwargs):
                 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))