Merge pull request #3 from FichteForks/pr/sync-spawn-map_n

Make spawn_n and map_n synchronous

Author: gistart

README.md

@@ -72,7 +72,7 @@ async def spawn_n_usage(todo=[range(1,51), range(51,101), range(101,200)]):
                 # Returns quickly for all tasks, does not wait for pool space.
                 # Workers are not spawned, they wait for pool space in their
                 # own background tasks.
-                fut = await pool.spawn_n(worker(i))
+                fut = pool.spawn_n(worker(i))
                 futures.append(fut)
         # At this point not a single worker should start.
 
@@ -147,7 +147,7 @@ async def callbacks_usage():
 
     async with AioPool(size=2) as pool:
         for i in todo:
-            fut = await pool.spawn_n(wrk(i), cb, (pool, i))
+            fut = pool.spawn_n(wrk(i), cb, (pool, i))
             futures.append(fut)
 
     results = []
@@ -173,7 +173,7 @@ async def callbacks_usage():
 
 async def exec_usage(todo=range(1,11)):
     async with AioPool(size=4) as pool:
-        futures = await pool.map_n(worker, todo)
+        futures = pool.map_n(worker, todo)
 
         # While other workers are waiting or active, you can "synchronously"
         # execute one task. It does not interrupt  others, just waits for pool
@@ -195,9 +195,9 @@ async def cancel_usage():
 
     pool = AioPool(size=2)
 
-    f_quick = await pool.spawn_n(aio.sleep(0.1))
-    f12 = await pool.spawn(wrk()), await pool.spawn_n(wrk())
-    f35 = await pool.map_n(wrk, range(3))
+    f_quick = pool.spawn_n(aio.sleep(0.1))
+    f12 = await pool.spawn(wrk()), pool.spawn_n(wrk())
+    f35 = pool.map_n(wrk, range(3))
 
     # At this point, if you cancel futures, returned by pool methods,
     # you just won't be able to retrieve spawned task results, task
@@ -232,9 +232,9 @@ async def details(todo=range(1,11)):
     # This code:
     f1 = []
     for i in todo:
-        f1.append(await pool.spawn_n(worker(i)))
+        f1.append(pool.spawn_n(worker(i)))
     # is equivalent to one call of `map_n`:
-    f2 = await pool.map_n(worker, todo)
+    f2 = pool.map_n(worker, todo)
 
     # Afterwards you can await for any given future:
     try:

asyncio_pool/base_pool.py

@@ -106,7 +106,7 @@ async def _wrap(self, coro, future, cb=None, ctx=None):
         res, exc, tb = None, None, None
         try:
             res = await coro
-        except Exception as _exc:
+        except BaseException as _exc:
             exc = _exc
             tb = traceback.format_exc()
         finally:
@@ -140,7 +140,7 @@ async def _spawn(self, future, coro, cb=None, ctx=None):
         acq_error = False
         try:
             await self.semaphore.acquire()
-        except Exception as e:
+        except BaseException as e:
             acq_error = True
             if not future.done():
                 future.set_exception(e)
@@ -182,7 +182,7 @@ async def spawn(self, coro, cb=None, ctx=None):
         self._waiting[future] = self.loop.create_future()  # as a placeholder
         return await self._spawn(future, coro, cb=cb, ctx=ctx)
 
-    async def spawn_n(self, coro, cb=None, ctx=None):
+    def spawn_n(self, coro, cb=None, ctx=None):
         '''Creates waiting task for given `coro` regardless of pool space. If
         pool is not full, this task will be executed very soon. Main difference
         is that `spawn_n` does not block and returns future very quickly.
@@ -203,15 +203,15 @@ async def exec(self, coro, cb=None, ctx=None):
         '''
         return await (await self.spawn(coro, cb, ctx))
 
-    async def map_n(self, fn, iterable, cb=None, ctx=None):
+    def map_n(self, fn, iterable, cb=None, ctx=None):
         '''Creates coroutine with `fn` function for each item in `iterable`,
         spawns each of them with `spawn_n`, returning futures.
 
         Read more about callbacks in `spawn` docstring.
         '''
         futures = []
         for it in iterable:
-            fut = await self.spawn_n(fn(it), cb, ctx)
+            fut = self.spawn_n(fn(it), cb, ctx)
             futures.append(fut)
         return futures
 

asyncio_pool/mx_asyncgen.py

@@ -37,7 +37,7 @@ async def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True,
         '''Spawns coroutines created with `fn` for each item in `iterable`, then
         waits for results with `iterwait`. See docs for `map_n` and `iterwait`.
         '''
-        futures = await self.map_n(fn, iterable, cb, ctx)
+        futures = self.map_n(fn, iterable, cb, ctx)
         generator = iterwait(futures, flat=flat, timeout=timeout,
                 get_result=get_result, yield_when=yield_when, loop=self.loop)
         async for batch in generator:

asyncio_pool/mx_asynciter.py

@@ -60,7 +60,7 @@ def __aiter__(_self):
 
             async def __anext__(_self):
                 if not hasattr(_self, 'waiter'):
-                    _self.waiter = mk_waiter(await mk_map())
+                    _self.waiter = mk_waiter(mk_map())
                 return await _self.waiter.__anext__()
 
         return _itermap()

asyncio_pool/results.py

@@ -14,7 +14,7 @@ def result_noraise(future, flat=True):
     try:
         res = future.result()
         return res if flat else (res, None)
-    except Exception as exc:
+    except BaseException as exc:
         return exc if flat else (None, exc)
 
 

examples/_usage.py

@@ -24,7 +24,7 @@ async def spawn_n_usage(todo=[range(1,51), range(51,101), range(101,200)]):
                 # Returns quickly for all tasks, does not wait for pool space.
                 # Workers are not spawned, they wait for pool space in their
                 # own background tasks.
-                fut = await pool.spawn_n(worker(i))
+                fut = pool.spawn_n(worker(i))
                 futures.append(fut)
         # At this point not a single worker should start.
 
@@ -99,7 +99,7 @@ async def cb(res, err, ctx):  # callback
 
     async with AioPool(size=2) as pool:
         for i in todo:
-            fut = await pool.spawn_n(wrk(i), cb, (pool, i))
+            fut = pool.spawn_n(wrk(i), cb, (pool, i))
             futures.append(fut)
 
     results = []
@@ -110,7 +110,7 @@ async def cb(res, err, ctx):  # callback
         #   results.append(getres.flat(fut))
         try:
             results.append(fut.result())
-        except Exception as e:
+        except BaseException as e:
             results.append(e)
 
     # First error happens for n == 0 in wrk, exception of it is passed to
@@ -125,7 +125,7 @@ async def cb(res, err, ctx):  # callback
 
 async def exec_usage(todo=range(1,11)):
     async with AioPool(size=4) as pool:
-        futures = await pool.map_n(worker, todo)
+        futures = pool.map_n(worker, todo)
 
         # While other workers are waiting or active, you can "synchronously"
         # execute one task. It does not interrupt  others, just waits for pool
@@ -147,9 +147,9 @@ async def wrk(*arg, **kw):
 
     pool = AioPool(size=2)
 
-    f_quick = await pool.spawn_n(aio.sleep(0.1))
-    f12 = await pool.spawn(wrk()), await pool.spawn_n(wrk())
-    f35 = await pool.map_n(wrk, range(3))
+    f_quick = pool.spawn_n(aio.sleep(0.1))
+    f12 = await pool.spawn(wrk()), pool.spawn_n(wrk())
+    f35 = pool.map_n(wrk, range(3))
 
     # At this point, if you cancel futures, returned by pool methods,
     # you just won't be able to retrieve spawned task results, task
@@ -184,15 +184,15 @@ async def details(todo=range(1,11)):
     # This code:
     f1 = []
     for i in todo:
-        f1.append(await pool.spawn_n(worker(i)))
+        f1.append(pool.spawn_n(worker(i)))
     # is equivalent to one call of `map_n`:
-    f2 = await pool.map_n(worker, todo)
+    f2 = pool.map_n(worker, todo)
 
     # Afterwards you can await for any given future:
     try:
         assert 3 == await f1[2]  # result of spawn_n(worker(3))
-    except Exception as e:
-        # exception happened in worker (or CancelledError) will be re-raised
+    except BaseException:
+        # exception happened in worker (including CancelledError) will be re-raised
         pass
 
     # Or use `asyncio.wait` to handle results in batches (see `iterwait` also):

tests/loadtest.py

@@ -25,7 +25,7 @@ async def loadtest_spawn_n(tasks, pool_size, duration):
     futures = []
     async with AioPool(size=pool_size) as pool:
         for i in range(tasks):
-            fut = await pool.spawn_n(aio.sleep(duration))
+            fut = pool.spawn_n(aio.sleep(duration))
             futures.append(fut)
 
     return [getres.flat(f) for f in futures]

tests/test_base.py

@@ -26,7 +26,7 @@ async def wrk(n):
 
     pool_size = 5
     async with AioPool(size=pool_size) as pool:
-        futures = await pool.map_n(wrk, todo)
+        futures = pool.map_n(wrk, todo)
 
         await aio.sleep(0.01)
 
@@ -70,7 +70,7 @@ async def outer(n, pool):
     loop = aio.get_event_loop()
     pool = AioPool(size=100)
 
-    tasks = await pool.map_n(inner, todo)
+    pool.map_n(inner, todo)
     joined = [loop.create_task(outer(j, pool)) for j in to_release]
     await pool.join()
 
@@ -86,18 +86,27 @@ async def wrk(*arg, **kw):
         await aio.sleep(0.5)
         return 1
 
-    pool = AioPool(size=2)
+    async def wrk_safe(*arg, **kw):
+        try:
+            await aio.sleep(0.5)
+        except aio.CancelledError:
+            await aio.sleep(0.1)  # simulate cleanup
+            pass
+        return 1
+
+    pool = AioPool(size=5)
 
-    f_quick = await pool.spawn_n(aio.sleep(0.15))
-    f12 = await pool.spawn(wrk()), await pool.spawn_n(wrk())
-    f35 = await pool.map_n(wrk, range(3))
+    f_quick = pool.spawn_n(aio.sleep(0.15))
+    f_safe = await pool.spawn(wrk_safe())
+    f3 = await pool.spawn(wrk())
+    pool.spawn_n(wrk())
+    f567 = pool.map_n(wrk, range(3))
 
     # cancel some
     await aio.sleep(0.1)
-    cancelled, results = await pool.cancel(f12[0], f35[2])  # running and waiting
-    assert 2 == cancelled  # none of them had time to finish
-    assert 2 == len(results) and \
-        all(isinstance(res, aio.CancelledError) for res in results)
+    cancelled, results = await pool.cancel(f3, f567[2])  # running and waiting
+    assert cancelled == len(results) == 2  # none of them had time to finish
+    assert all(isinstance(res, aio.CancelledError) for res in results)
 
     # cancel all others
     await aio.sleep(0.1)
@@ -106,11 +115,12 @@ async def wrk(*arg, **kw):
     assert f_quick.done() and f_quick.result() is None
 
     cancelled, results = await pool.cancel()  # all
-    assert 3 == cancelled
-    assert len(results) == 3 and \
-        all(isinstance(res, aio.CancelledError) for res in results)
+    assert cancelled == len(results) == 4
+    assert f_safe.done() and f_safe.result() == 1  # could recover
+    # the others could not
+    assert sum(isinstance(res, aio.CancelledError) for res in results) == 3
 
-    assert await pool.join()  # joins successfully
+    assert await pool.join()  # joins successfully (basically no-op)
 
 
 @pytest.mark.asyncio

tests/test_callbacks.py

@@ -31,7 +31,7 @@ async def test_spawn_n():
     async with AioPool(size=2) as pool:
         for i in todo:
             ctx = (pool, i)
-            fut = await pool.spawn_n(wrk(i), cb, ctx)
+            fut = pool.spawn_n(wrk(i), cb, ctx)
             futures.append(fut)
 
     results = [getres.flat(f) for f in futures]
@@ -52,7 +52,7 @@ async def test_map():
 async def test_map_n():
     todo = range(2,11)
     async with AioPool(size=3) as pool:
-        futures = await pool.map_n(wrk, todo, cb)
+        futures = pool.map_n(wrk, todo, cb)
 
     results = [getres.flat(f) for f in futures]
     assert 2 * sum(todo) == sum(results)

tests/test_spawn.py

@@ -27,7 +27,7 @@ async def wrk(n):
 
     async with AioPool(size=2) as pool:
         for i in range(1,6):
-            await pool.spawn_n(wrk(i))  # does not wait for pool, just spawns waiting coros
+            pool.spawn_n(wrk(i))  # does not wait for pool, just spawns waiting coros
         assert len(started) == 0  # so atm no worker should be able to start