async def details(todo=range(1, 11)): pool = AioPool(size=5) # This code: f1 = [] for i in todo: f1.append(pool.spawn_n(worker(i))) # is equivalent to one call of `map_n`: 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 BaseException: # exception happened in worker (including CancelledError) will be re-raised pass # Or use `asyncio.wait` to handle results in batches (see `iterwait` also): important_res = 0 more_important = [f1[1], f2[1], f2[2]] while more_important: done, more_important = await aio.wait(more_important, timeout=0.5) # handle result, note it will re-raise exceptions important_res += sum(f.result() for f in done) assert important_res == 2 + 2 + 3 # But you need to join, to allow all spawned workers to finish # (of course you can `asyncio.wait` all of the futures if you want to) await pool.join() assert all(f.done() for f in itertools.chain(f1, f2)) # this is guaranteed assert 2 * sum(todo) == sum(f.result() for f in itertools.chain(f1, f2))
async def test_cancel(): async def wrk(*arg, **kw): await aio.sleep(0.5) return 1 async def wrk_safe(*arg, **kw): try: await aio.sleep(0.5) except aio.CancelledError: await aio.sleep(0.1) # simulate cleanup return 1 pool = AioPool(size=5) 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(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) # not interrupted and finished successfully assert f_quick.done() and f_quick.result() is None cancelled, results = await pool.cancel() # all 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 (basically no-op)
async def cancel_usage(): async def wrk(*arg, **kw): await aio.sleep(0.5) return 1 pool = AioPool(size=2) 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 # themselves will continue working. Don't do this: # f_quick.cancel() # use `pool.cancel` instead: # 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) # cancel all others await aio.sleep(0.1) # not interrupted and finished successfully 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 await pool.join() # joins successfully
class Worker(BaseClient): def __init__(self, enabled_tasks=[]): BaseClient.__init__(self, TYPE_WORKER, self._message_callback, self._on_connected) self.defrsps = {} self.lockers = {} self._tasks = {} self._broadcast_tasks = [] self.enabled_tasks = enabled_tasks self._pool = None self.grab_agents = {} self.grab_queue = None self.executor = None def set_enable_tasks(self, enabled_tasks): self.enabled_tasks = enabled_tasks def set_executor(self, executor): '''executer for sync process''' self.executor = executor def is_enabled(self, func): if len(self.enabled_tasks) == 0: return True return func in self.enabled_tasks async def _on_connected(self): for func in self._tasks.keys(): if func in self._broadcast_tasks: await self._broadcast(func) else: await self._add_func(func) await asyncio.sleep(1) async def _add_func(self, func): if not self.is_enabled(func): return logger.info(f'Add {func}') agent = await self.gen_agent() await agent.send(cmd.CanDo(self._add_prefix_subfix(func))) self.remove_agent(agent) async def add_func(self, func, task, defrsp=DoneResponse(), locker=None): if self.connected: await self._add_func(func) self._tasks[func] = task self.defrsps[func] = defrsp self.lockers[func] = locker async def _broadcast(self, func): if not self.is_enabled(func): return logger.info(f'Broadcast {func}') agent = await self.gen_agent() await agent.send(cmd.Broadcast(self._add_prefix_subfix(func))) self.remove_agent(agent) async def broadcast(self, func, task, defrsp=DoneResponse(), locker=None): if self.connected: await self._broadcast(func) self._tasks[func] = task self.defrsps[func] = defrsp self.lockers[func] = locker self._broadcast_tasks.append(func) async def remove_func(self, func): logger.info(f'Remove {func}') agent = await self.gen_agent() await agent.send(cmd.CantDo(self._add_prefix_subfix(func))) self.remove_agent(agent) self._tasks.pop(func, None) self.defrsps.pop(func, None) self.lockers.pop(func, None) if func in self._broadcast_tasks: self._broadcast_tasks.remove(func) async def next_grab(self): if self._pool.is_full: return for _ in range(self.grab_queue.qsize()): agent = await self.grab_queue.get() await self.grab_queue.put(agent) if agent.is_timeout(): await agent.safe_send() break async def work(self, size): self._pool = AioPool(size=size) agents = [await self.gen_agent() for _ in range(size)] self.grab_queue = asyncio.Queue() for agent in agents: item = GrabAgent(agent) await self.grab_queue.put(item) self.grab_agents[agent.msgid] = item while True: await self.next_grab() await asyncio.sleep(1) async def _message_callback(self, payload, msgid): await self.next_grab() self._pool.spawn_n(self.run_task(payload, msgid)) async def run_task(self, payload, msgid): try: job = Job(payload[1:], self) await self.process_job(job) except asyncio.exceptions.CancelledError: pass finally: agent = self.grab_agents[msgid] await agent.safe_send() async def process_job(self, job): task = self._tasks.get(job.func_name) if not task: await self.remove_func(job.func_name) await job.fail() else: async def process(): await self._process_job(job, task) locker = self.lockers.get(job.func_name) if locker: locker_name, count = locker(job) if locker_name: await job.with_lock(locker_name, count, process) return await process() async def _process_job(self, job, task): try: if asyncio.iscoroutinefunction(task): ret = await task(job) else: if self.executor: loop = asyncio.get_running_loop() task = loop.run_in_executor(self.executor, task, job) await asyncio.wait([task]) ret = task.result() else: ret = task(job) except Exception as e: logger.exception(e) ret = self.defrsps.get(job.func_name, FailResponse()) if not job.finished: if isinstance(ret, str): await job.done(bytes(ret, 'utf-8')) elif isinstance(ret, bytes): await job.done(ret) elif isinstance(ret, DoneResponse): await job.done(ret.buf) elif isinstance(ret, FailResponse): await job.fail() elif isinstance(ret, SchedLaterResponse): await job.sched_later(ret.delay, ret.count) else: await job.done() # decorator def func(self, func_name, broadcast=False, defrsp=DoneResponse(), locker=None): def _func(task): self._tasks[func_name] = task self.defrsps[func_name] = defrsp self.lockers[func_name] = locker if broadcast: self._broadcast_tasks.append(func_name) return task return _func def blueprint(self, app): app.set_worker(self) self._tasks.update(app.tasks) self.defrsps.update(app.defrsps) self.lockers.update(app.lockers) for btsk in app.broadcast_tasks: self._broadcast_tasks.append(btsk)