concurrent tee iterators are safe and ordered

Author: maxfischer2781

asyncstdlib/itertools.py

@@ -380,11 +380,25 @@ async def __anext__(self) -> T:
             async with self._lock:
                 # Check if another peer produced an item while we were waiting for the lock
                 if not next_node:
-                    next_node[:] = await self._iterator.__anext__(), []
+                    await self._extend_buffer(next_node)
+        if not next_node:
+            raise StopAsyncIteration()
         # for any other TeePeer, the node is already some [value, [...]]
         value, self._buffer = next_node  # type: ignore
         return value
 
+    async def _extend_buffer(self, next_node: "_TeeNode[T]") -> None:
+        """Extend the buffer by fetching a new item from the iterable"""
+        try:
+            next_value = await self._iterator.__anext__()
+        except StopAsyncIteration:
+            return
+        # another peer may have filled the buffer while we waited
+        # seek the last node that needs to be filled
+        while next_node:
+            _, next_node = next_node  # type: ignore
+        next_node[:] = next_value, []
+
     async def aclose(self) -> None:
         self._tee_peers.discard(self._tee_idx)
         if not self._tee_peers and isinstance(self._iterator, ACloseable):

unittests/test_itertools.py

@@ -1,3 +1,4 @@
+from typing import AsyncIterator
 import itertools
 import sys
 import platform
@@ -341,7 +342,7 @@ async def test_tee():
 
 @sync
 async def test_tee_concurrent_locked():
-    """Test that properly uses a lock for synchronisation"""
+    """Test that tee properly uses a lock for synchronisation"""
     items = [1, 2, 3, -5, 12, 78, -1, 111]
 
     async def iter_values():
@@ -393,6 +394,41 @@ async def test_peer(peer_tee):
         await test_peer(this)
 
 
+@pytest.mark.parametrize("size", [2, 3, 5, 9, 12])
+@sync
+async def test_tee_concurrent_ordering(size: int):
+    """Test that tee respects concurrent ordering for all peers"""
+
+    class ConcurrentInvertedIterable:
+        """Helper that concurrently iterates with earlier items taking longer"""
+
+        def __init__(self, count: int) -> None:
+            self.count = count
+            self._counter = itertools.count()
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            value = next(self._counter)
+            if value >= self.count:
+                raise StopAsyncIteration()
+            await Switch(self.count - value)
+            return value
+
+    async def test_peer(peer_tee: AsyncIterator[int]):
+        # consume items from the tee with a delay so that slower items can arrive
+        seen_items: list[int] = []
+        async for item in peer_tee:
+            seen_items.append(item)
+            await Switch()
+        assert seen_items == expected_items
+
+    expected_items = list(range(size)[::-1])
+    peers = a.tee(ConcurrentInvertedIterable(size), n=size)
+    await Schedule(*map(test_peer, peers))
+
+
 @sync
 async def test_pairwise():
     assert await a.list(a.pairwise(range(5))) == [(0, 1), (1, 2), (2, 3), (3, 4)]