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)
