├── .github └── workflows │ └── tests.yml ├── .gitignore ├── LICENSE.txt ├── README.md ├── __init__.py ├── asyncio_pool ├── __init__.py ├── base_pool.py ├── mx_asyncgen.py ├── mx_asynciter.py └── results.py ├── build_dist.sh ├── docs ├── _readme_template.md └── build_readme.py ├── examples ├── __init__.py ├── _usage.py └── ls_remote.py ├── reqs-dev.txt ├── reqs-test.txt ├── run_tests.sh ├── setup.py ├── setup_dev.sh └── tests ├── __init__.py ├── loadtest.py ├── test_base.py ├── test_callbacks.py ├── test_map.py └── test_spawn.py /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: test-all 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | branches: [ master ] 8 | workflow_dispatch: 9 | 10 | jobs: 11 | test: 12 | name: run all tests 13 | runs-on: ubuntu-20.04 14 | strategy: 15 | matrix: 16 | python-version: [3.6, 3.7, 3.8, 3.9, '3.10', 'pypy3'] 17 | steps: 18 | - name: checkout 19 | uses: actions/checkout@v2 20 | - name: python${{ matrix.python-version }} 21 | uses: actions/setup-python@v2 22 | with: 23 | python-version: ${{ matrix.python-version }} 24 | - name: pip-requirements 25 | run: pip install --upgrade -r reqs-test.txt 26 | - name: run-test 27 | run: pytest -sv --continue-on-collection-errors ./tests 28 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.swo 2 | *.swp 3 | *.pyc 4 | *.pyo 5 | *.log 6 | *.lock 7 | .env* 8 | .pypyenv* 9 | .pytest_* 10 | .mypy_cache 11 | dist 12 | build 13 | *.egg-info 14 | __pycache__ 15 | local_settings.py 16 | 17 | !.gitkeep 18 | !/.gitignore 19 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2018 Oleg Mihaylov (gistart) 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of 4 | this software and associated documentation files (the "Software"), to deal in 5 | the Software without restriction, including without limitation the rights to 6 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 7 | of the Software, and to permit persons to whom the Software is furnished to do 8 | so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in all 11 | copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | SOFTWARE. 20 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # asyncio-pool 2 | 3 | Pool of asyncio coroutines with familiar interface. Supports python 3.5+ (including PyPy 6+, which is also 3.5 atm) 4 | 5 | AioPool makes sure _no more_ and _no less_ (if possible) than `size` spawned coroutines are active at the same time. _spawned_ means created and scheduled with one of the pool interface methods, _active_ means coroutine function started executing it's code, as opposed to _waiting_ -- which waits for pool space without entering coroutine function. 6 | 7 | ## Interface 8 | 9 | Read [code doctrings](../master/asyncio_pool/base_pool.py) for details. 10 | 11 | #### AioPool(size=4, *, loop=None) 12 | 13 | Creates pool of `size` concurrent tasks. Supports async context manager interface. 14 | 15 | #### spawn(coro, cb=None, ctx=None) 16 | 17 | Waits for pool space, then creates task for `coro` coroutine, returning future for it's result. Can spawn coroutine, created by `cb` with result of `coro` as first argument. `ctx` context is passed to callback as third positinal argument. 18 | 19 | #### exec(coro, cb=None, ctx=None) 20 | 21 | Waits for pool space, then creates task for `coro`, then waits for it to finish, then returns result of `coro` if no callback is provided, otherwise creates task for callback, waits for it and returns result of callback. 22 | 23 | #### spawn_n(coro, cb=None, ctx=None) 24 | 25 | Creates waiting task for `coro`, returns future without waiting for pool space. Task is executed "in pool" when pool space is available. 26 | 27 | #### join() 28 | 29 | Waits for all spawned (active and waiting) tasks to finish. Joining pool from coroutine, spawned by the same pool leads to *deadlock*. 30 | 31 | #### cancel(*futures) 32 | 33 | Cancels spawned tasks (active and waiting), finding them by provided `futures`. If no futures provided -- cancels all spawned tasks. 34 | 35 | #### map(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat) 36 | 37 | Spawns coroutines created by `fn` function for each item in `iterable` with `spawn`, waits for all of them to finish (including callbacks), returns results maintaining order of `iterable`. 38 | 39 | #### map_n(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat) 40 | 41 | Spawns coroutines created by `fn` function for each item in `iterable` with `spawn_n`, returns futures for task results maintaining order of `iterable`. 42 | 43 | #### itermap(fn, iterable, cb=None, ctx=None, *, flat=True, get_result=getres.flat, timeout=None, yield_when=asyncio.ALL_COMPLETED) 44 | 45 | Spawns tasks with `map_n(fn, iterable, cb, ctx)`, then waits for results with `asyncio.wait` function, yielding ready results one by one if `flat` == True, otherwise yielding list of ready results. 46 | 47 | 48 | 49 | ## Usage 50 | 51 | `spawn` and `map` methods is probably what you should use in 99% of cases. Their overhead is minimal (~3% execution time), and even in worst cases memory usage is insignificant. 52 | 53 | `spawn_n`, `map_n` and `itermap` methods give you more control and flexibily, but they come with a price of higher overhead. They spawn all tasks that you want, and most of the tasks wait their turn "in background". If you spawn too much (10**6+ tasks) -- you'll use most of the memory you have in system, also you'll lose a lot of time on "concurrency management" of all the tasks spawned. 54 | 55 | Play with `python tests/loadtest.py -h` to understand what you want to use. 56 | 57 | Usage examples (more in [tests/](../master/tests/) and [examples/](../master/examples/)): 58 | 59 | ```python 60 | 61 | 62 | async def worker(n): # dummy worker 63 | await aio.sleep(1 / n) 64 | return n 65 | 66 | 67 | async def spawn_n_usage(todo=[range(1,51), range(51,101), range(101,200)]): 68 | futures = [] 69 | async with AioPool(size=20) as pool: 70 | for tasks in todo: 71 | for i in tasks: # too many tasks 72 | # Returns quickly for all tasks, does not wait for pool space. 73 | # Workers are not spawned, they wait for pool space in their 74 | # own background tasks. 75 | fut = pool.spawn_n(worker(i)) 76 | futures.append(fut) 77 | # At this point not a single worker should start. 78 | 79 | # Context manager calls `join` at exit, so this will finish when all 80 | # workers return, crash or cancelled. 81 | 82 | assert sum(itertools.chain.from_iterable(todo)) == \ 83 | sum(f.result() for f in futures) 84 | 85 | 86 | async def spawn_usage(todo=range(1,4)): 87 | futures = [] 88 | async with AioPool(size=2) as pool: 89 | for i in todo: # 1, 2, 3 90 | # Returns quickly for 1 and 2, then waits for empty space for 3, 91 | # spawns 3 and returns. Can save some resources I guess. 92 | fut = await pool.spawn(worker(i)) 93 | futures.append(fut) 94 | # At this point some of the workers already started. 95 | 96 | # Context manager calls `join` at exit, so this will finish when all 97 | # workers return, crash or cancelled. 98 | 99 | assert sum(todo) == sum(fut.result() for fut in futures) # all done 100 | 101 | 102 | async def map_usage(todo=range(100)): 103 | pool = AioPool(size=10) 104 | # Waits and collects results from all spawned workers, 105 | # returns them in same order as `todo`, if worker crashes or cancelled: 106 | # returns exception object as a result. 107 | # Basically, it wraps `spawn_usage` code into one call. 108 | results = await pool.map(worker, todo) 109 | 110 | # await pool.join() # is not needed here, bcs no other tasks were spawned 111 | 112 | assert isinstance(results[0], ZeroDivisionError) \ 113 | and sum(results[1:]) == sum(todo) 114 | 115 | 116 | async def itermap_usage(todo=range(1,11)): 117 | result = 0 118 | async with AioPool(size=10) as pool: 119 | # Combines spawn_n and iterwait, which is a wrapper for asyncio.wait, 120 | # which yields results of finished workers according to `timeout` and 121 | # `yield_when` params passed to asyncio.wait (see it's docs for details) 122 | async for res in pool.itermap(worker, todo, timeout=0.5): 123 | result += res 124 | # technically, you can skip join call 125 | 126 | assert result == sum(todo) 127 | 128 | 129 | async def callbacks_usage(): 130 | 131 | async def wrk(n): # custom dummy worker 132 | await aio.sleep(1 / n) 133 | return n 134 | 135 | async def cb(res, err, ctx): # callback 136 | if err: # error handling 137 | exc, tb = err 138 | assert tb # the only purpose of this is logging 139 | return exc 140 | 141 | pool, n = ctx # context can be anything you like 142 | await aio.sleep(1 / (n-1)) 143 | return res + n 144 | 145 | todo = range(5) 146 | futures = [] 147 | 148 | async with AioPool(size=2) as pool: 149 | for i in todo: 150 | fut = pool.spawn_n(wrk(i), cb, (pool, i)) 151 | futures.append(fut) 152 | 153 | results = [] 154 | for fut in futures: 155 | # there are helpers for result extraction. `flat` one will do 156 | # exactly what's written below 157 | # from asyncio_pool import getres 158 | # results.append(getres.flat(fut)) 159 | try: 160 | results.append(fut.result()) 161 | except Exception as e: 162 | results.append(e) 163 | 164 | # First error happens for n == 0 in wrk, exception of it is passed to 165 | # callback, callback returns it to us. Second one happens in callback itself 166 | # and is passed to us by pool. 167 | assert all(isinstance(e, ZeroDivisionError) for e in results[:2]) 168 | 169 | # All n's in `todo` are passed through `wrk` and `cb` (cb adds wrk result 170 | # and # number, passed inside context), except for n == 0 and n == 1. 171 | assert sum(results[2:]) == 2 * (sum(todo) - 0 - 1) 172 | 173 | 174 | async def exec_usage(todo=range(1,11)): 175 | async with AioPool(size=4) as pool: 176 | futures = pool.map_n(worker, todo) 177 | 178 | # While other workers are waiting or active, you can "synchronously" 179 | # execute one task. It does not interrupt others, just waits for pool 180 | # space, then waits for task to finish and then returns it's result. 181 | important_res = await pool.exec(worker(2)) 182 | assert 2 == important_res 183 | 184 | # You can continue working as usual: 185 | moar = await pool.spawn(worker(10)) 186 | 187 | assert sum(todo) == sum(f.result() for f in futures) 188 | 189 | 190 | async def cancel_usage(): 191 | 192 | async def wrk(*arg, **kw): 193 | await aio.sleep(0.5) 194 | return 1 195 | 196 | pool = AioPool(size=2) 197 | 198 | f_quick = pool.spawn_n(aio.sleep(0.1)) 199 | f12 = await pool.spawn(wrk()), pool.spawn_n(wrk()) 200 | f35 = pool.map_n(wrk, range(3)) 201 | 202 | # At this point, if you cancel futures, returned by pool methods, 203 | # you just won't be able to retrieve spawned task results, task 204 | # themselves will continue working. Don't do this: 205 | # f_quick.cancel() 206 | # use `pool.cancel` instead: 207 | 208 | # cancel some 209 | await aio.sleep(0.1) 210 | cancelled, results = await pool.cancel(f12[0], f35[2]) # running and waiting 211 | assert 2 == cancelled # none of them had time to finish 212 | assert 2 == len(results) and \ 213 | all(isinstance(res, aio.CancelledError) for res in results) 214 | 215 | # cancel all others 216 | await aio.sleep(0.1) 217 | 218 | # not interrupted and finished successfully 219 | assert f_quick.done() and f_quick.result() is None 220 | 221 | cancelled, results = await pool.cancel() # all 222 | assert 3 == cancelled 223 | assert len(results) == 3 and \ 224 | all(isinstance(res, aio.CancelledError) for res in results) 225 | 226 | assert await pool.join() # joins successfully 227 | 228 | 229 | async def details(todo=range(1,11)): 230 | pool = AioPool(size=5) 231 | 232 | # This code: 233 | f1 = [] 234 | for i in todo: 235 | f1.append(pool.spawn_n(worker(i))) 236 | # is equivalent to one call of `map_n`: 237 | f2 = pool.map_n(worker, todo) 238 | 239 | # Afterwards you can await for any given future: 240 | try: 241 | assert 3 == await f1[2] # result of spawn_n(worker(3)) 242 | except Exception as e: 243 | # exception happened in worker (or CancelledError) will be re-raised 244 | pass 245 | 246 | # Or use `asyncio.wait` to handle results in batches (see `iterwait` also): 247 | important_res = 0 248 | more_important = [f1[1], f2[1], f2[2]] 249 | while more_important: 250 | done, more_important = await aio.wait(more_important, timeout=0.5) 251 | # handle result, note it will re-raise exceptions 252 | important_res += sum(f.result() for f in done) 253 | 254 | assert important_res == 2 + 2 + 3 255 | 256 | # But you need to join, to allow all spawned workers to finish 257 | # (of course you can `asyncio.wait` all of the futures if you want to) 258 | await pool.join() 259 | 260 | assert all(f.done() for f in itertools.chain(f1,f2)) # this is guaranteed 261 | assert 2 * sum(todo) == sum(f.result() for f in itertools.chain(f1,f2)) 262 | 263 | 264 | ``` 265 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/gistart/asyncio-pool/cd79fe782f97c6a268f61307a80397744d6c3f4b/__init__.py -------------------------------------------------------------------------------- /asyncio_pool/__init__.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | from .results import getres 4 | from .base_pool import BaseAioPool 5 | 6 | 7 | if sys.version_info < (3, 6): # this means 3.5 # TODO test 3.4? 8 | 9 | from .mx_asynciter import MxAsyncIterPool, iterwait 10 | 11 | class AioPool(MxAsyncIterPool, BaseAioPool): pass 12 | 13 | else: 14 | from .mx_asyncgen import MxAsyncGenPool, iterwait 15 | 16 | class AioPool(MxAsyncGenPool, BaseAioPool): pass 17 | -------------------------------------------------------------------------------- /asyncio_pool/base_pool.py: -------------------------------------------------------------------------------- 1 | '''Pool of asyncio coroutines with familiar interface, python3.5+ friendly''' 2 | 3 | import traceback 4 | import asyncio as aio 5 | from .results import getres 6 | 7 | 8 | class BaseAioPool(object): 9 | 10 | def __init__(self, size=1024, *, loop=None): 11 | '''Pool of asyncio coroutines with familiar interface. 12 | 13 | Pool makes sure _no more_ and _no less_ (if possible) than `size` 14 | spawned coroutines are active at the same time. _spawned_ means created 15 | and scheduled with one of the pool interface methods, _active_ means 16 | coroutine function started executing it's code, as opposed to 17 | _waiting_ -- which waits for pool space without entering coroutine 18 | function. 19 | 20 | Support asynchronous context management protocol (`aenter`, `aexit`). 21 | 22 | The main idea behind spwaning methods is -- they return newly created 23 | futures, not "native" ones, returned by `pool.create_task` or used for 24 | `await`. Read more about this in readme and docstrings below. 25 | ''' 26 | 27 | self.loop = loop or _get_loop() 28 | 29 | self.size = size 30 | self._executed = 0 31 | self._joined = set() 32 | self._waiting = {} # future -> task 33 | self._active = {} # future -> task 34 | self.semaphore = aio.Semaphore(value=self.size) 35 | 36 | async def __aenter__(self): 37 | return self 38 | 39 | async def __aexit__(self, ext_type, exc, tb): 40 | await self.join() 41 | 42 | def __len__(self): 43 | return len(self._waiting) + self.n_active 44 | 45 | @property 46 | def n_active(self): 47 | '''Counts active coroutines''' 48 | return self.size - self.semaphore._value 49 | 50 | @property 51 | def is_empty(self): 52 | '''Returns `True` if no coroutines are active or waiting.''' 53 | return 0 == len(self._waiting) == self.n_active 54 | 55 | @property 56 | def is_full(self): 57 | '''Returns `True` if `size` coroutines are already active.''' 58 | return self.size <= len(self) 59 | 60 | async def join(self): 61 | '''Waits (blocks) for all spawned coroutines to finish, both active and 62 | waiting. *Do not `join` inside spawned coroutine*.''' 63 | 64 | if self.is_empty: 65 | return True 66 | 67 | fut = self.loop.create_future() 68 | self._joined.add(fut) 69 | try: 70 | return await fut 71 | finally: 72 | self._joined.remove(fut) 73 | 74 | def _release_joined(self): 75 | if not self.is_empty: 76 | raise RuntimeError() # TODO better message 77 | 78 | for fut in self._joined: 79 | if not fut.done(): 80 | fut.set_result(True) 81 | 82 | def _build_callback(self, cb, res, err=None, ctx=None): 83 | # not sure if this is a safe code( in case any error: 84 | # return cb(res, err, ctx), None 85 | 86 | bad_cb = RuntimeError('cb should accept at least one argument') 87 | to_pass = (res, err, ctx) 88 | 89 | nargs = cb.__code__.co_argcount 90 | if nargs == 0: 91 | return None, bad_cb 92 | 93 | # trusting user here, better ideas? 94 | if cb.__code__.co_varnames[0] in ('self', 'cls'): 95 | nargs -= 1 # class/instance method, skip first arg 96 | 97 | if nargs == 0: 98 | return None, bad_cb 99 | 100 | try: 101 | return cb(*to_pass[:nargs]), None 102 | except Exception as e: 103 | return None, e 104 | 105 | async def _wrap(self, coro, future, cb=None, ctx=None): 106 | res, exc, tb = None, None, None 107 | try: 108 | res = await coro 109 | except BaseException as _exc: 110 | exc = _exc 111 | tb = traceback.format_exc() 112 | finally: 113 | self._executed += 1 114 | 115 | while cb: 116 | err = None if exc is None else (exc, tb) 117 | 118 | _cb, _cb_err = self._build_callback(cb, res, err, ctx) 119 | if _cb_err is not None: 120 | exc = _cb_err # pass to future 121 | break 122 | 123 | wrapped = self._wrap(_cb, future) 124 | self.loop.create_task(wrapped) 125 | return 126 | 127 | self.semaphore.release() 128 | 129 | if not future.done(): 130 | if exc: 131 | future.set_exception(exc) 132 | else: 133 | future.set_result(res) 134 | 135 | del self._active[future] 136 | if self.is_empty: 137 | self._release_joined() 138 | 139 | async def _spawn(self, future, coro, cb=None, ctx=None): 140 | acq_error = False 141 | try: 142 | await self.semaphore.acquire() 143 | except BaseException as e: 144 | acq_error = True 145 | coro.close() 146 | if not future.done(): 147 | future.set_exception(e) 148 | finally: 149 | del self._waiting[future] 150 | 151 | if future.done(): 152 | if not acq_error and future.cancelled(): # outside action 153 | self.semaphore.release() 154 | else: # all good, can spawn now 155 | wrapped = self._wrap(coro, future, cb=cb, ctx=ctx) 156 | task = self.loop.create_task(wrapped) 157 | self._active[future] = task 158 | return future 159 | 160 | async def spawn(self, coro, cb=None, ctx=None): 161 | '''Waits for pool space and creates task for given `coro` coroutine, 162 | returns a future for it's result. 163 | 164 | If callback `cb` coroutine function (not coroutine itself!) is passed, 165 | `coro` result won't be assigned to created future, instead, `cb` will 166 | be executed with it as a first positional argument. Callback function 167 | should accept 1,2 or 3 positional arguments. Full callback sigature is 168 | `cb(res, err, ctx)`. It makes no sense to create a callback without 169 | `coro` result, so first positional argument is mandatory. 170 | 171 | Second positional argument of callback will be error, which `is None` 172 | if coroutine did not crash and wasn't cancelled. In case any exception 173 | was raised during `coro` execution, error will be a tuple containing 174 | (`exc` exception object, `tb` traceback string). if you wish to ignore 175 | errors, you can pass callback without seconds and third positional 176 | arguments. 177 | 178 | If context `ctx` is passed to `spawn`, it will be re-sent to callback 179 | as third argument. If you don't plan to use any context, you can create 180 | callback with positional arguments only for result and error. 181 | ''' 182 | future = self.loop.create_future() 183 | self._waiting[future] = self.loop.create_future() # as a placeholder 184 | return await self._spawn(future, coro, cb=cb, ctx=ctx) 185 | 186 | def spawn_n(self, coro, cb=None, ctx=None): 187 | '''Creates waiting task for given `coro` regardless of pool space. If 188 | pool is not full, this task will be executed very soon. Main difference 189 | is that `spawn_n` does not block and returns future very quickly. 190 | 191 | Read more about callbacks in `spawn` docstring. 192 | ''' 193 | future = self.loop.create_future() 194 | task = self.loop.create_task(self._spawn(future, coro, cb=cb, ctx=ctx)) 195 | self._waiting[future] = task 196 | return future 197 | 198 | async def exec(self, coro, cb=None, ctx=None): 199 | '''Waits for pool space, then waits for `coro` (and it's callback if 200 | passed) to finish, returning result of `coro` or callback (if passed), 201 | or raising error if smth crashed in process or was cancelled. 202 | 203 | Read more about callbacks in `spawn` docstring. 204 | ''' 205 | return await (await self.spawn(coro, cb, ctx)) 206 | 207 | def map_n(self, fn, iterable, cb=None, ctx=None): 208 | '''Creates coroutine with `fn` function for each item in `iterable`, 209 | spawns each of them with `spawn_n`, returning futures. 210 | 211 | Read more about callbacks in `spawn` docstring. 212 | ''' 213 | futures = [] 214 | for it in iterable: 215 | fut = self.spawn_n(fn(it), cb, ctx) 216 | futures.append(fut) 217 | return futures 218 | 219 | async def map(self, fn, iterable, cb=None, ctx=None, *, 220 | get_result=getres.flat): 221 | '''Spawns coroutines, created with `fn` function for each item in 222 | `iterable`, waits for all of them to finish, crash or be cancelled, 223 | returning resuls. 224 | 225 | `get_result` is function, that accepts future as only positional 226 | argument, whose goal is to extract result from future. You can pass 227 | your own, or use inluded `getres` object, that has 3 extractors: 228 | `getres.dont` will return future untouched, `getres.flat` will return 229 | exception object if coroutine crashed or was cancelled, otherwise will 230 | return result of a coroutine (or of the callback), `getres.pair` will 231 | return tuple of (`result', 'exception` object) with None in place of 232 | missing item. 233 | 234 | Read more about callbacks in `spawn` docstring. 235 | ''' 236 | futures = [] 237 | for it in iterable: 238 | fut = await self.spawn(fn(it), cb, ctx) 239 | futures.append(fut) 240 | 241 | await aio.wait(futures) 242 | return [get_result(fut) for fut in futures] 243 | 244 | async def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True, 245 | get_result=getres.flat, timeout=None, yield_when=aio.ALL_COMPLETED): 246 | '''Spawns coroutines created with `fn` for each item in `iterable`, then 247 | waits for results with `iterwait` (implementation specific). See docs 248 | for `map_n` and `iterwait` (in mixins for py3.5 and py3.6+). 249 | ''' 250 | raise NotImplementedError('Use one of mixins') 251 | 252 | async def cancel(self, *futures, get_result=getres.flat): 253 | '''Cancels spawned or waiting tasks, found by their `futures`. If no 254 | `futures` are passed -- cancels all spwaned and waiting tasks. 255 | 256 | Cancelling futures, returned by pool methods, usually won't help you 257 | to cancel executing tasks, so you have to use this method. 258 | 259 | Returns tuple of (`cancelled` count of cancelled tasks, `results` 260 | collected from futures of cancelled tasks). 261 | ''' 262 | tasks, _futures = [], [] 263 | 264 | if not futures: # meaning cancel all 265 | tasks.extend(self._waiting.values()) 266 | tasks.extend(self._active.values()) 267 | _futures.extend(self._waiting.keys()) 268 | _futures.extend(self._active.keys()) 269 | else: 270 | for fut in futures: 271 | task = self._active.get(fut, self._waiting.get(fut)) 272 | if task: 273 | tasks.append(task) 274 | _futures.append(fut) 275 | 276 | cancelled = 0 277 | if tasks: 278 | cancelled = sum(1 for task in tasks if task.cancel()) 279 | await aio.wait(tasks) # let them actually cancel 280 | # need to collect them anyway, to supress warnings 281 | return cancelled, [get_result(fut) for fut in _futures] 282 | 283 | 284 | def _get_loop(): 285 | """ 286 | Backward compatibility w/ py<3.8 287 | """ 288 | 289 | if hasattr(aio, 'get_running_loop'): 290 | return aio.get_running_loop() 291 | return aio.get_event_loop() -------------------------------------------------------------------------------- /asyncio_pool/mx_asyncgen.py: -------------------------------------------------------------------------------- 1 | '''Mixin for BaseAioPool with async generator features, python3.6+''' 2 | 3 | import asyncio as aio 4 | from .results import getres 5 | 6 | 7 | async def iterwait(futures, *, flat=True, get_result=getres.flat, 8 | timeout=None, yield_when=aio.ALL_COMPLETED): 9 | '''Wraps `asyncio.wait` into asynchronous generator, accessible with 10 | `async for` syntax. May be useful in conjunction with `spawn_n`. 11 | 12 | `timeout` and `yield_when` parameters are passed to `asyncio.wait`, see 13 | documentation for this great instrument. 14 | 15 | Returns results for provided futures, as soon as results are ready. If 16 | `flat` is True -- generates one result at a time (per `async for`). If 17 | `flat` is False -- generates a list of ready results. 18 | ''' 19 | _futures = futures[:] 20 | while _futures: 21 | done, _futures = await aio.wait(_futures, timeout=timeout, 22 | return_when=yield_when) 23 | if flat: 24 | for fut in done: 25 | yield get_result(fut) 26 | else: 27 | yield [get_result(fut) for fut in done] 28 | 29 | 30 | class MxAsyncGenPool(object): 31 | # Asynchronous generator wrapper for asyncio.wait. 32 | 33 | async def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True, 34 | get_result=getres.flat, timeout=None, 35 | yield_when=aio.ALL_COMPLETED): 36 | '''Spawns coroutines created with `fn` for each item in `iterable`, then 37 | waits for results with `iterwait`. See docs for `map_n` and `iterwait`. 38 | ''' 39 | futures = self.map_n(fn, iterable, cb, ctx) 40 | generator = iterwait(futures, flat=flat, timeout=timeout, 41 | get_result=get_result, yield_when=yield_when) 42 | async for batch in generator: 43 | yield batch # TODO is it possible to return a generator? 44 | -------------------------------------------------------------------------------- /asyncio_pool/mx_asynciter.py: -------------------------------------------------------------------------------- 1 | '''Mixin for BaseAioPool with async generator _simulation_ for python 3.5''' 2 | 3 | import asyncio as aio 4 | from collections import deque 5 | from functools import partial 6 | from .results import getres 7 | 8 | 9 | class iterwait: 10 | 11 | def __init__(self, futures, *, flat=True, get_result=getres.flat, 12 | timeout=None, yield_when=aio.ALL_COMPLETED, loop=None): 13 | 14 | self.results = deque() 15 | self.flat = flat 16 | self._futures = futures 17 | self._getres = get_result 18 | self._wait = partial(aio.wait, timeout=timeout, loop=loop, 19 | return_when=yield_when) 20 | 21 | def __aiter__(self): 22 | return self 23 | 24 | async def __anext__(self): 25 | if not (self._futures or self.results): 26 | raise StopAsyncIteration() 27 | while not self.results: 28 | await self._wait_next() 29 | return self.results.popleft() 30 | 31 | async def _wait_next(self): 32 | while True: 33 | done, self._futures = await self._wait(self._futures) 34 | if done: 35 | batch = [self._getres(fut) for fut in done] 36 | if self.flat: 37 | self.results.extend(batch) 38 | else: 39 | self.results.append(batch) 40 | break 41 | 42 | 43 | class MxAsyncIterPool(object): 44 | 45 | def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True, 46 | get_result=getres.flat, timeout=None, 47 | yield_when=aio.ALL_COMPLETED): 48 | '''Spawns coroutines created with `fn` for each item in `iterable`, then 49 | waits for results with `iterwait`. See docs for `map_n` and `iterwait`. 50 | ''' 51 | mk_map = partial(self.map_n, fn, iterable, cb=cb, ctx=ctx) 52 | mk_waiter = partial(iterwait, flat=flat, loop=self.loop, 53 | get_result=get_result, timeout=timeout, 54 | yield_when=yield_when) 55 | 56 | class _itermap: 57 | def __aiter__(_self): 58 | return _self 59 | 60 | async def __anext__(_self): 61 | if not hasattr(_self, 'waiter'): 62 | _self.waiter = mk_waiter(mk_map()) 63 | return await _self.waiter.__anext__() 64 | 65 | return _itermap() 66 | -------------------------------------------------------------------------------- /asyncio_pool/results.py: -------------------------------------------------------------------------------- 1 | from functools import partial 2 | 3 | 4 | def result_noraise(future, flat=True): 5 | '''Extracts result from future, never raising an exception. 6 | 7 | If `flat` is True -- returns result or exception instance (including 8 | CancelledError), if `flat` is False -- returns tuple of (`result`, 9 | `exception` object). 10 | 11 | If traceback is needed -- just re-raise returned exception.''' 12 | try: 13 | res = future.result() 14 | return res if flat else (res, None) 15 | except BaseException as exc: 16 | return exc if flat else (None, exc) 17 | 18 | 19 | class getres: 20 | dont = lambda fut: fut 21 | flat = partial(result_noraise, flat=True) 22 | pair = partial(result_noraise, flat=False) 23 | -------------------------------------------------------------------------------- /build_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/sh 2 | 3 | # build 4 | rm -r *.egg_info build 5 | .env/bin/python setup.py sdist bdist_wheel 6 | 7 | # upload 8 | #.env/bin/twine upload -u gistart --skip-existing dist/* 9 | -------------------------------------------------------------------------------- /docs/_readme_template.md: -------------------------------------------------------------------------------- 1 | # asyncio-pool 2 | 3 | Pool of asyncio coroutines with familiar interface. Supports python 3.5+ (including PyPy 6+, which is also 3.5 atm) 4 | 5 | AioPool makes sure _no more_ and _no less_ (if possible) than `size` spawned coroutines are active at the same time. _spawned_ means created and scheduled with one of the pool interface methods, _active_ means coroutine function started executing it's code, as opposed to _waiting_ -- which waits for pool space without entering coroutine function. 6 | 7 | ## Interface 8 | 9 | Read [code doctrings](../master/asyncio_pool/base_pool.py) for details. 10 | 11 | #### AioPool(size=4, *, loop=None) 12 | 13 | Creates pool of `size` concurrent tasks. Supports async context manager interface. 14 | 15 | #### spawn(coro, cb=None, ctx=None) 16 | 17 | Waits for pool space, then creates task for `coro` coroutine, returning future for it's result. Can spawn coroutine, created by `cb` with result of `coro` as first argument. `ctx` context is passed to callback as third positinal argument. 18 | 19 | #### exec(coro, cb=None, ctx=None) 20 | 21 | Waits for pool space, then creates task for `coro`, then waits for it to finish, then returns result of `coro` if no callback is provided, otherwise creates task for callback, waits for it and returns result of callback. 22 | 23 | #### spawn_n(coro, cb=None, ctx=None) 24 | 25 | Creates waiting task for `coro`, returns future without waiting for pool space. Task is executed "in pool" when pool space is available. 26 | 27 | #### join() 28 | 29 | Waits for all spawned (active and waiting) tasks to finish. Joining pool from coroutine, spawned by the same pool leads to *deadlock*. 30 | 31 | #### cancel(*futures) 32 | 33 | Cancels spawned tasks (active and waiting), finding them by provided `futures`. If no futures provided -- cancels all spawned tasks. 34 | 35 | #### map(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat) 36 | 37 | Spawns coroutines created by `fn` function for each item in `iterable` with `spawn`, waits for all of them to finish (including callbacks), returns results maintaining order of `iterable`. 38 | 39 | #### map_n(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat) 40 | 41 | Spawns coroutines created by `fn` function for each item in `iterable` with `spawn_n`, returns futures for task results maintaining order of `iterable`. 42 | 43 | #### itermap(fn, iterable, cb=None, ctx=None, *, flat=True, get_result=getres.flat, timeout=None, yield_when=asyncio.ALL_COMPLETED) 44 | 45 | Spawns tasks with `map_n(fn, iterable, cb, ctx)`, then waits for results with `asyncio.wait` function, yielding ready results one by one if `flat` == True, otherwise yielding list of ready results. 46 | 47 | 48 | 49 | ## Usage 50 | 51 | `spawn` and `map` methods is probably what you should use in 99% of cases. Their overhead is minimal (~3% execution time), and even in worst cases memory usage is insignificant. 52 | 53 | `spawn_n`, `map_n` and `itermap` methods give you more control and flexibily, but they come with a price of higher overhead. They spawn all tasks that you want, and most of the tasks wait their turn "in background". If you spawn too much (10**6+ tasks) -- you'll use most of the memory you have in system, also you'll lose a lot of time on "concurrency management" of all the tasks spawned. 54 | 55 | Play with `python tests/loadtest.py -h` to understand what you want to use. 56 | 57 | Usage examples (more in [tests/](../master/tests/) and [examples/](../master/examples/)): 58 | 59 | ```python 60 | {tmpl_usage_examples} 61 | ``` 62 | -------------------------------------------------------------------------------- /docs/build_readme.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | curr_dir = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.insert(0, os.path.split(curr_dir)[0]) 5 | 6 | 7 | if __name__ == "__main__": 8 | 9 | with open('docs/_readme_template.md') as f: 10 | template = f.read() 11 | with open('examples/_usage.py') as u: 12 | usage = u.read() 13 | 14 | readme_usage, do_append = [], False 15 | for line in usage.splitlines(): 16 | if line.startswith('#') and line.endswith('<>to_here'): 20 | do_append = False 21 | 22 | if do_append: 23 | readme_usage.append(line) 24 | 25 | readme = template.format( 26 | tmpl_usage_examples='\n'.join(readme_usage) 27 | ) 28 | 29 | with open('README.md', 'w') as rd: 30 | rd.write(readme) 31 | -------------------------------------------------------------------------------- /examples/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/gistart/asyncio-pool/cd79fe782f97c6a268f61307a80397744d6c3f4b/examples/__init__.py -------------------------------------------------------------------------------- /examples/_usage.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | curr_dir = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.insert(0, os.path.split(curr_dir)[0]) 5 | 6 | import itertools 7 | import asyncio as aio 8 | from asyncio_pool import AioPool 9 | #<>to_here 215 | if __name__ == "__main__": 216 | aio.get_event_loop().run_until_complete(aio.gather( 217 | spawn_n_usage(), 218 | spawn_usage(), 219 | map_usage(), 220 | itermap_usage(), 221 | callbacks_usage(), 222 | exec_usage(), 223 | cancel_usage(), 224 | details() 225 | )) 226 | -------------------------------------------------------------------------------- /examples/ls_remote.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | curr_dir = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.insert(0, os.path.split(curr_dir)[0]) 5 | 6 | import logging 7 | import asyncio as aio 8 | import json as myformat # just for a demo 9 | 10 | 11 | def parse_res(res): 12 | doc = myformat.loads(res) 13 | return doc.get('files') or [], doc.get('folders') or [] 14 | 15 | 16 | async def ls_cb(res, err, ctx): 17 | apipool, dbpool, api, db, log = ctx 18 | if err: 19 | log.error() 20 | return 21 | 22 | files, folders = parse_res(res) 23 | for folder in folders: 24 | await apipool.spawn(api.ls(folder['path']), cb=ls_cb, ctx=ctx) 25 | await dbpool.spawn(db.save(folder)) 26 | for file in files: 27 | await dbpool.spawn(db.save(file)) 28 | 29 | 30 | async def example_ls(): 31 | loop = aio.get_event_loop() 32 | api = await create_client(loop) 33 | db = await create_connection(dsn, loop) 34 | log = logging.getLogger('example_ls') 35 | 36 | with AioPool(size=10) as dbpool, \ 37 | AioPool(size=10) as apipool: 38 | ctx = (apipool, dbpool, api, db, log) 39 | await apipool.spawn(api.ls('/'), cb=ls_cb, ctx=ctx) 40 | 41 | 42 | if __name__ == "__main__": 43 | aio.get_event_loop().run_until_complete(example_ls()) 44 | -------------------------------------------------------------------------------- /reqs-dev.txt: -------------------------------------------------------------------------------- 1 | -r reqs-test.txt 2 | 3 | ipython 4 | setuptools 5 | wheel 6 | twine 7 | -------------------------------------------------------------------------------- /reqs-test.txt: -------------------------------------------------------------------------------- 1 | async-timeout 2 | pytest 3 | pytest-asyncio -------------------------------------------------------------------------------- /run_tests.sh: -------------------------------------------------------------------------------- 1 | #! /bin/sh 2 | 3 | pypy3=pypy3 4 | py35=python3.5 5 | py36=python3.6 6 | py37=python3.7 7 | py38=python3.8 8 | py39=python3.9 9 | py310=python3.10 10 | 11 | default_env=$py37 12 | 13 | 14 | todo=${@:-"./tests"} 15 | 16 | for py in $pypy3 $py35 $py36 $py37 $py38 $py39 $py310 17 | do 18 | echo "" 19 | if [ -x "$(command -v $py)" ]; then 20 | pyname="$(basename $py)" 21 | envname=".env_$pyname" 22 | 23 | if ! [ -d $envname ]; then 24 | echo "$pyname: virtual env does not exist" 25 | else 26 | echo "$pyname: running for $todo" 27 | $envname/bin/pytest --continue-on-collection-errors $todo 28 | fi 29 | else 30 | echo "$py: not found" 31 | fi 32 | done 33 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | 4 | github_url = 'https://github.com/gistart/asyncio-pool' 5 | 6 | readme_lines = [] 7 | with open('README.md') as fd: 8 | readme_lines = filter(None, fd.read().splitlines()) 9 | readme_lines = list(readme_lines)[:3] 10 | readme_lines.append('Read more at [github page](%s).' % github_url) 11 | readme = '\n\n'.join(readme_lines) 12 | 13 | 14 | setuptools.setup( 15 | name='asyncio_pool', 16 | version='0.6.0', 17 | author='gistart', 18 | author_email='gistart@yandex.ru', 19 | description='Pool of asyncio coroutines with familiar interface', 20 | long_description=readme, 21 | long_description_content_type='text/markdown', 22 | url=github_url, 23 | license='MIT', 24 | packages=['asyncio_pool'], 25 | # install_requires=['asyncio'], # where " openstack/deb-python-trollius asyncio" comes from??? 26 | python_requires='>=3.5', 27 | classifiers=( 28 | "Programming Language :: Python :: 3", 29 | "Programming Language :: Python :: 3.5", 30 | "Programming Language :: Python :: 3.6", 31 | "Programming Language :: Python :: 3.7", 32 | "License :: OSI Approved :: MIT License", 33 | "Operating System :: OS Independent", 34 | "Framework :: AsyncIO", 35 | ) 36 | ) 37 | -------------------------------------------------------------------------------- /setup_dev.sh: -------------------------------------------------------------------------------- 1 | #! /bin/sh 2 | 3 | pypy3=pypy3 4 | py35=python3.5 5 | py36=python3.6 6 | py37=python3.7 7 | py38=python3.8 8 | py39=python3.9 9 | py310=python3.10 10 | 11 | default_env=$py37 12 | 13 | 14 | for py in $pypy3 $py35 $py36 $py37 $py38 $py39 $py310 15 | do 16 | if [ -x "$(command -v $py)" ]; then 17 | pyname="$(basename $py)" 18 | envname=".env_$pyname" 19 | 20 | if ! [ -d $envname ]; then 21 | echo "$pyname: creating new env at $envname" 22 | $py -m venv --system-site-packages $envname 23 | 24 | echo "$pyname: upgrading pip" 25 | $envname/bin/pip install --upgrade pip 26 | fi 27 | 28 | if ! [ -d ".env" ] && [ "$py" = "$default_env" ]; then 29 | ln -s $envname .env 30 | fi 31 | 32 | $envname/bin/pip install --upgrade -r reqs-dev.txt 33 | 34 | else 35 | echo "$py: not found" 36 | fi 37 | done 38 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/gistart/asyncio-pool/cd79fe782f97c6a268f61307a80397744d6c3f4b/tests/__init__.py -------------------------------------------------------------------------------- /tests/loadtest.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | curr_dir = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.insert(0, os.path.split(curr_dir)[0]) 5 | 6 | import time 7 | import argparse 8 | import asyncio as aio 9 | from asyncio_pool import AioPool, getres 10 | 11 | 12 | async def loadtest_spawn(tasks, pool_size, duration): 13 | futures = [] 14 | async with AioPool(size=pool_size) as pool: 15 | for i in range(tasks): 16 | fut = await pool.spawn(aio.sleep(duration)) 17 | futures.append(fut) 18 | 19 | return [getres.flat(fut) for fut in futures] 20 | 21 | 22 | async def loadtest_spawn_n(tasks, pool_size, duration): 23 | futures = [] 24 | async with AioPool(size=pool_size) as pool: 25 | for i in range(tasks): 26 | fut = pool.spawn_n(aio.sleep(duration)) 27 | futures.append(fut) 28 | 29 | return [getres.flat(f) for f in futures] 30 | 31 | 32 | async def loadtest_map(tasks, pool_size, duration): 33 | async def wrk(i): 34 | await aio.sleep(duration) 35 | 36 | async with AioPool(size=pool_size) as pool: 37 | return await pool.map(wrk, range(tasks)) 38 | 39 | 40 | async def loadtest_itermap(tasks, pool_size, duration): 41 | async def wrk(i): 42 | await aio.sleep(duration) 43 | 44 | results = [] 45 | async with AioPool(size=pool_size) as pool: 46 | async for res in pool.itermap(wrk, range(tasks)): 47 | results.append(res) 48 | 49 | return results 50 | 51 | 52 | def print_stats(args, exec_time): 53 | ideal = args.task_duration * (args.tasks / args.pool_size) 54 | overhead = exec_time - ideal 55 | per_task = overhead / args.tasks 56 | overhead_perc = ((exec_time / ideal) - 1) * 100 57 | 58 | print(f'{ideal:15.5f}s -- ideal result') 59 | print(f'{exec_time:15.5f}s -- total executing time') 60 | print(f'{overhead:15.5f}s -- total overhead') 61 | print(f'{per_task:15.5f}s -- overhead per task') 62 | print(f'{overhead_perc:13.3f}% -- overhead total percent') 63 | 64 | 65 | if __name__ == "__main__": 66 | methods = { 67 | 'spawn': loadtest_spawn, 68 | 'spawn_n': loadtest_spawn_n, 69 | 'map': loadtest_map, 70 | 'itermap': loadtest_itermap, 71 | } 72 | 73 | p = argparse.ArgumentParser() 74 | p.add_argument('method', choices=methods.keys()) 75 | p.add_argument('--tasks', '-t', type=int, default=10**5) 76 | p.add_argument('--task-duration', '-d', type=float, default=0.2) 77 | p.add_argument('--pool-size', '-p', type=int, default=10**3) 78 | args = p.parse_args() 79 | 80 | print('>>> Running %d tasks in pool of size=%s, each task takes %.3f sec.' % 81 | (args.tasks, args.pool_size, args.task_duration)) 82 | print('>>> This will run more than %.5f seconds' % 83 | (args.task_duration * (args.tasks / args.pool_size))) 84 | 85 | ts_start = time.perf_counter() 86 | m = methods.get(args.method)(args.tasks, args.pool_size, args.task_duration) 87 | aio.get_event_loop().run_until_complete(m) 88 | exec_time = time.perf_counter() - ts_start 89 | print_stats(args, exec_time) 90 | -------------------------------------------------------------------------------- /tests/test_base.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import asyncio as aio 3 | from asyncio_pool import AioPool 4 | from async_timeout import timeout 5 | 6 | 7 | @pytest.mark.asyncio 8 | async def test_concurrency(): 9 | todo = range(1,21) 10 | coros_active = {n:False for n in todo} 11 | 12 | async def wrk(n): 13 | nonlocal coros_active 14 | coros_active[n] = True 15 | await aio.sleep(1 / n) 16 | coros_active[n] = False 17 | return n 18 | 19 | pool_size = 5 20 | async with AioPool(size=pool_size) as pool: 21 | futures = pool.map_n(wrk, todo) 22 | 23 | await aio.sleep(0.01) 24 | 25 | while not pool.is_empty: 26 | n_active = sum(filter(None, coros_active.values())) 27 | assert n_active <= pool_size 28 | await aio.sleep(0.01) 29 | 30 | assert sum(todo) == sum([f.result() for f in futures]) 31 | 32 | 33 | @pytest.mark.asyncio 34 | async def test_timeout_cancel(): 35 | async def wrk(sem): 36 | async with sem: 37 | await aio.sleep(1) 38 | 39 | sem = aio.Semaphore(value=2) 40 | 41 | async with timeout(0.2): 42 | with pytest.raises(aio.CancelledError): 43 | await aio.gather(*[wrk(sem) for _ in range(3)]) 44 | 45 | 46 | @pytest.mark.asyncio 47 | async def test_outer_join(): 48 | 49 | todo, to_release = range(1,15), range(10) 50 | done, released = [], [] 51 | 52 | async def inner(n): 53 | nonlocal done 54 | await aio.sleep(1 / n) 55 | done.append(n) 56 | 57 | async def outer(n, pool): 58 | nonlocal released 59 | await pool.join() 60 | released.append(n) 61 | 62 | loop = aio.get_event_loop() 63 | pool = AioPool(size=100) 64 | 65 | pool.map_n(inner, todo) 66 | joined = [loop.create_task(outer(j, pool)) for j in to_release] 67 | await pool.join() 68 | 69 | assert len(released) <= len(to_release) 70 | await aio.wait(joined) 71 | assert len(todo) == len(done) and len(released) == len(to_release) 72 | 73 | 74 | @pytest.mark.asyncio 75 | async def test_cancel(): 76 | 77 | async def wrk(*arg, **kw): 78 | await aio.sleep(0.5) 79 | return 1 80 | 81 | async def wrk_safe(*arg, **kw): 82 | try: 83 | await aio.sleep(0.5) 84 | except aio.CancelledError: 85 | await aio.sleep(0.1) # simulate cleanup 86 | return 1 87 | 88 | pool = AioPool(size=5) 89 | 90 | f_quick = pool.spawn_n(aio.sleep(0.15)) 91 | f_safe = await pool.spawn(wrk_safe()) 92 | f3 = await pool.spawn(wrk()) 93 | pool.spawn_n(wrk()) 94 | f567 = pool.map_n(wrk, range(3)) 95 | 96 | # cancel some 97 | await aio.sleep(0.1) 98 | cancelled, results = await pool.cancel(f3, f567[2]) # running and waiting 99 | assert cancelled == len(results) == 2 # none of them had time to finish 100 | assert all(isinstance(res, aio.CancelledError) for res in results) 101 | 102 | # cancel all others 103 | await aio.sleep(0.1) 104 | 105 | # not interrupted and finished successfully 106 | assert f_quick.done() and f_quick.result() is None 107 | 108 | cancelled, results = await pool.cancel() # all 109 | assert cancelled == len(results) == 4 110 | assert f_safe.done() and f_safe.result() == 1 # could recover 111 | # the others could not 112 | assert sum(isinstance(res, aio.CancelledError) for res in results) == 3 113 | 114 | assert await pool.join() # joins successfully (basically no-op) 115 | 116 | 117 | @pytest.mark.asyncio 118 | async def test_internal_join(): 119 | 120 | async def wrk(pool): 121 | return await pool.join() # deadlock 122 | 123 | pool = AioPool(size=10) 124 | fut = await pool.spawn(wrk(pool)) 125 | 126 | await aio.sleep(0.5) 127 | assert not fut.done() # dealocked, will never return 128 | 129 | await pool.cancel(fut) 130 | await pool.join() 131 | -------------------------------------------------------------------------------- /tests/test_callbacks.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import asyncio as aio 3 | from asyncio_pool import AioPool, getres 4 | 5 | 6 | async def cb(res, err, ctx): 7 | if err: 8 | exc, tb = err 9 | assert tb 10 | return exc 11 | 12 | await aio.sleep(1 / (res - 1)) 13 | return res * 2 14 | 15 | 16 | async def wrk(n): 17 | await aio.sleep(1 / n) 18 | return n 19 | 20 | 21 | @pytest.mark.asyncio 22 | async def test_spawn_n(): 23 | todo = range(5) 24 | futures = [] 25 | async with AioPool(size=2) as pool: 26 | for i in todo: 27 | ctx = (pool, i) 28 | fut = pool.spawn_n(wrk(i), cb, ctx) 29 | futures.append(fut) 30 | 31 | results = [getres.flat(f) for f in futures] 32 | assert all(isinstance(e, ZeroDivisionError) for e in results[:2]) 33 | assert sum(results[2:]) == 2 * (sum(todo) - 0 - 1) 34 | 35 | 36 | @pytest.mark.asyncio 37 | async def test_map(): 38 | todo = range(2,11) 39 | async with AioPool(size=3) as pool: 40 | results = await pool.map(wrk, todo, cb) 41 | 42 | assert 2 * sum(todo) == sum(results) 43 | 44 | 45 | @pytest.mark.asyncio 46 | async def test_map_n(): 47 | todo = range(2,11) 48 | async with AioPool(size=3) as pool: 49 | futures = pool.map_n(wrk, todo, cb) 50 | 51 | results = [getres.flat(f) for f in futures] 52 | assert 2 * sum(todo) == sum(results) 53 | 54 | 55 | @pytest.mark.parametrize('timeout', [None, 0.1, 2]) 56 | @pytest.mark.asyncio 57 | async def test_itermap(timeout): 58 | todo = range(2,11) 59 | 60 | res = 0 61 | async with AioPool(size=3) as pool: 62 | async for i in pool.itermap(wrk, todo, cb, timeout=timeout): 63 | res += i 64 | 65 | assert 2 * sum(todo) == res 66 | 67 | 68 | @pytest.mark.asyncio 69 | async def test_callback_types(): 70 | 71 | def cbsync(res, err, ctx): 72 | return res, err, ctx 73 | 74 | async def cb0(): 75 | await aio.sleep(0) 76 | return 'x' 77 | 78 | async def cb1(res): 79 | assert res 80 | await aio.sleep(0) 81 | return res 82 | 83 | async def cb2(res, err): 84 | await aio.sleep(0) 85 | return res, err 86 | 87 | async def cb3(res, err, ctx): 88 | await aio.sleep(0) 89 | return res, err, ctx 90 | 91 | class CbCls: 92 | def cbsync(self, res): 93 | return res 94 | 95 | async def cb0(self): 96 | await aio.sleep(0) 97 | return 'y' 98 | 99 | async def cb1(self, res): 100 | assert res 101 | await aio.sleep(0) 102 | return res 103 | 104 | async def cb2(self, res, err): 105 | await aio.sleep(0) 106 | return res, err 107 | 108 | async def cb3(self, res, err, ctx): 109 | assert res 110 | await aio.sleep(0) 111 | return res, err, ctx 112 | 113 | @staticmethod 114 | async def cb_static(res, err, ctx): 115 | assert res 116 | await aio.sleep(0) 117 | return res, err, ctx 118 | 119 | @classmethod 120 | async def cb_cls(cls, res): 121 | assert res 122 | await aio.sleep(0) 123 | return res 124 | 125 | async def wrk(n): 126 | if isinstance(n, int): 127 | assert n - 500 128 | await aio.sleep(0) 129 | return n 130 | 131 | _r = getres.flat 132 | _ctx = (123, 456, [1,2], {3,4}) 133 | inst = CbCls() 134 | 135 | async with AioPool() as pool: 136 | # non-async callbacks 137 | with pytest.raises(TypeError): 138 | await pool.exec(wrk(1), cbsync, _ctx) 139 | with pytest.raises(TypeError): 140 | await pool.exec(wrk(1), inst.cbsync, _ctx) 141 | 142 | # zero args 143 | with pytest.raises(RuntimeError): 144 | await pool.exec(wrk(1), cb0, _ctx) 145 | with pytest.raises(RuntimeError): 146 | await pool.exec(wrk(1), inst.cb0, _ctx) 147 | 148 | # one arg 149 | assert 123 == await pool.exec(wrk(123), cb1, _ctx)\ 150 | == await pool.exec(wrk(123), inst.cb1, _ctx)\ 151 | == await pool.exec(wrk(123), CbCls.cb_cls, _ctx) 152 | 153 | with pytest.raises(AssertionError): 154 | await pool.exec(wrk(False), cb1, _ctx) 155 | 156 | futs = (await pool.spawn(wrk(500), cb1, _ctx), 157 | await pool.spawn(wrk(False), cb1, _ctx)) 158 | await aio.wait(futs) 159 | assert all(isinstance(_r(f), AssertionError) for f in futs) 160 | 161 | # two args 162 | assert (123, None) == await pool.exec(wrk(123), cb2, _ctx)\ 163 | == await pool.exec(wrk(123), inst.cb2, _ctx) 164 | 165 | assert (False, None) == await pool.exec(wrk(False), inst.cb2, _ctx) 166 | 167 | res, (exc, tb) = await pool.exec(wrk(500), cb2, _ctx) 168 | assert res is None and isinstance(exc, AssertionError) and \ 169 | tb.startswith('Traceback (most recent call last)') 170 | 171 | # three args 172 | assert (123, None, _ctx) == await pool.exec(wrk(123), cb3, _ctx) \ 173 | == await pool.exec(wrk(123), inst.cb3, _ctx) \ 174 | == await pool.exec(wrk(123), CbCls.cb_static, _ctx) 175 | 176 | with pytest.raises(AssertionError): 177 | await pool.exec(wrk(False), inst.cb3, _ctx) 178 | 179 | res, (exc, tb), ctx = await pool.exec(wrk(500), cb3, _ctx) 180 | assert res is None and isinstance(exc, AssertionError) and \ 181 | tb.startswith('Traceback (most recent call last)') 182 | 183 | 184 | def inspect(obj): 185 | names = dir(obj) 186 | pad = len(max(names, key=lambda n: len(n))) 187 | for name in names: 188 | print(name.ljust(pad),':', getattr(obj, name)) 189 | -------------------------------------------------------------------------------- /tests/test_map.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import asyncio as aio 3 | from asyncio_pool import AioPool, getres 4 | 5 | 6 | async def wrk(n): 7 | await aio.sleep(1 / n) 8 | return n*10 9 | 10 | 11 | @pytest.mark.asyncio 12 | async def test_map_simple(): 13 | task = range(1,11) 14 | pool = AioPool(size=7) 15 | res = await pool.map(wrk, task) 16 | assert res == [i*10 for i in task] 17 | 18 | 19 | @pytest.mark.asyncio 20 | async def test_map_crash(): 21 | task = range(5) 22 | pool = AioPool(size=10) 23 | 24 | # exc as result 25 | res = await pool.map(wrk, task, get_result=getres.flat) 26 | assert isinstance(res[0], Exception) 27 | assert res[1:] == [i*10 for i in task[1:]] 28 | 29 | # tuple as result 30 | res = await pool.map(wrk, task, get_result=getres.pair) 31 | assert res[0][0] is None and isinstance(res[0][1], ZeroDivisionError) 32 | assert [r[0] for r in res[1:]] == [i*10 for i in task[1:]] and \ 33 | not any(r[1] for r in res[1:]) 34 | 35 | 36 | @pytest.mark.asyncio 37 | async def test_itermap(): 38 | 39 | async def wrk(n): 40 | await aio.sleep(n) 41 | return n 42 | 43 | async with AioPool(size=3) as pool: 44 | i = 0 45 | async for res in pool.itermap(wrk, [0.5] * 4, flat=False, timeout=0.6): 46 | if i == 0: 47 | assert 15 == int(sum(res) * 10) 48 | elif i == 1: 49 | assert 5 == int(sum(res) * 10) 50 | else: 51 | assert False # should not get here 52 | i += 1 # does not support enumerate btw ( 53 | 54 | 55 | @pytest.mark.asyncio 56 | async def test_itermap_cancel(): 57 | 58 | async def wrk(n): 59 | await aio.sleep(n / 100) 60 | return n 61 | 62 | todo = range(1, 101) 63 | 64 | async with AioPool(5) as pool: 65 | async for res in pool.itermap(wrk, todo, yield_when=aio.FIRST_COMPLETED): 66 | if res == 13: 67 | cancelled, _ = await pool.cancel() 68 | break 69 | assert cancelled == 100 - 13 -------------------------------------------------------------------------------- /tests/test_spawn.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import asyncio as aio 3 | from asyncio_pool import AioPool 4 | 5 | 6 | @pytest.mark.asyncio 7 | async def test_spawns_behaviour(): 8 | started = [] 9 | 10 | async def wrk(n): 11 | nonlocal started 12 | started.append(n) 13 | await aio.sleep(0.1) 14 | 15 | async with AioPool(size=2) as pool: 16 | for i in range(1,6): 17 | await pool.spawn(wrk(i)) # waits for pool to be available 18 | assert len(started) != 0 # so atm some of workers should start 19 | 20 | started.clear() 21 | 22 | async with AioPool(size=2) as pool: 23 | for i in range(1,6): 24 | pool.spawn_n(wrk(i)) # does not wait for pool, just spawns waiting coros 25 | assert len(started) == 0 # so atm no worker should be able to start 26 | 27 | 28 | @pytest.mark.asyncio 29 | async def test_spawn_crash(): 30 | async def wrk(n): 31 | return 1 / n 32 | 33 | futures = [] 34 | async with AioPool(size=1) as pool: 35 | for i in (2, 1, 0): 36 | futures.append(await pool.spawn(wrk(i))) 37 | 38 | with pytest.raises(ZeroDivisionError): 39 | futures[-1].result() 40 | 41 | 42 | @pytest.mark.asyncio 43 | async def test_spawn_and_exec(): 44 | 45 | order = [] 46 | marker = 9999 47 | 48 | async def wrk(n): 49 | nonlocal order 50 | order.append(n) 51 | if n == marker: 52 | await aio.sleep(0.5) 53 | else: 54 | await aio.sleep(1 / n) 55 | order.append(n) 56 | return n 57 | 58 | task = range(1, 11) 59 | futures = [] 60 | async with AioPool(size=7) as pool: 61 | for i in task: 62 | futures.append(await pool.spawn(wrk(i))) 63 | assert marker == await pool.exec(wrk(marker)) 64 | 65 | assert [f.result() for f in futures] == list(task) 66 | assert pool._executed == len(task) + 1 67 | assert order != list(sorted(order)) 68 | 69 | ix = order.index(marker) 70 | iy = order.index(marker, ix+1) 71 | assert iy - ix > 1 72 | --------------------------------------------------------------------------------