async concurrent mapping: decouple concurrency and buffersize

Author: ebonnal

streamable/afunctions.py

@@ -235,6 +235,7 @@ def amap(
     return ConcurrentAMapAsyncIterator(
         aiterator,
         transformation,
+        concurrency=concurrency,
         buffersize=concurrency,
         ordered=ordered,
     )

streamable/aiterators.py

@@ -717,15 +717,25 @@ def __init__(
         self,
         iterator: AsyncIterator[T],
         transformation: Callable[[T], Coroutine[Any, Any, U]],
+        concurrency: int,
         buffersize: int,
         ordered: bool,
     ) -> None:
         super().__init__(iterator, buffersize, ordered)
         self.transformation = transformation
+        self.concurrency = concurrency
+        self.semaphore: asyncio.Semaphore
+
+    def _context_manager(self) -> ContextManager:
+        self.semaphore = asyncio.Semaphore(
+            self.concurrency, loop=asyncio.get_running_loop()
+        )
+        return super()._context_manager()
 
     async def _safe_transformation(self, elem: T) -> Union[U, ExceptionContainer]:
         try:
-            return await self.transformation(elem)
+            async with self.semaphore:
+                return await self.transformation(elem)
         except Exception as e:
             return ExceptionContainer(e)
 
@@ -749,13 +759,15 @@ def __init__(
         self,
         iterator: AsyncIterator[T],
         transformation: Callable[[T], Coroutine[Any, Any, U]],
+        concurrency: int,
         buffersize: int,
         ordered: bool,
     ) -> None:
         super().__init__(
             _ConcurrentAMapAsyncIterable(
                 iterator,
                 transformation,
+                concurrency,
                 buffersize,
                 ordered,
             ).__aiter__()

streamable/functions.py

@@ -231,6 +231,7 @@ def amap(
         event_loop,
         iterator,
         transformation,
+        concurrency=concurrency,
         buffersize=concurrency,
         ordered=ordered,
     )

streamable/iterators.py

@@ -667,16 +667,19 @@ def __init__(
         event_loop: asyncio.AbstractEventLoop,
         iterator: Iterator[T],
         transformation: Callable[[T], Coroutine[Any, Any, U]],
+        concurrency: int,
         buffersize: int,
         ordered: bool,
     ) -> None:
         super().__init__(iterator, buffersize, ordered)
         self.transformation = transformation
         self.event_loop = event_loop
+        self.semaphore = asyncio.Semaphore(concurrency, loop=event_loop)
 
     async def _safe_transformation(self, elem: T) -> Union[U, ExceptionContainer]:
         try:
-            return await self.transformation(elem)
+            async with self.semaphore:
+                return await self.transformation(elem)
         except Exception as e:
             return ExceptionContainer(e)
 
@@ -701,6 +704,7 @@ def __init__(
         event_loop: asyncio.AbstractEventLoop,
         iterator: Iterator[T],
         transformation: Callable[[T], Coroutine[Any, Any, U]],
+        concurrency: int,
         buffersize: int,
         ordered: bool,
     ) -> None:
@@ -709,6 +713,7 @@ def __init__(
                 event_loop,
                 iterator,
                 transformation,
+                concurrency,
                 buffersize,
                 ordered,
             ).__iter__()

tests/test_iterators.py

@@ -22,6 +22,7 @@ def test_ConcurrentAMapAsyncIterable(self) -> None:
                 _ConcurrentAMapAsyncIterable(
                     sync_to_async_iter(iter(src)),
                     async_identity,
+                    concurrency=2,
                     buffersize=2,
                     ordered=True,
                 )

tests/utils.py

@@ -156,7 +156,7 @@ class TestError(Exception):
     pass
 
 
-DELTA_RATE = 0.4
+DELTA_RATE = 0.1
 # size of the test collections
 N = 256