简体   繁体   中英

wait for asyncio coroutines to finish in specific order

I have a daemon program with a bunch of asyncio.coroutine 's that could be summed up to something like this

import asyncio
import signal

class Daemon:

    def __init__(self, loop=asyncio.get_event_loop()):
        self.loop = loop
        self.running = False
        self.tasks = {
            'coroutine1': asyncio.ensure_future(self.coroutine1()),
            'coroutine2': asyncio.ensure_future(self.coroutine2()),
        }

    def run(self):
        self.running = True
        for task in self.tasks.values():
            task.add_done_callback(self.task_done_callback)
        # gracefuly close everything when SIGINT (could be ^C) is received
        self.loop.add_signal_handler(signal.SIGINT, self.close)
        self.loop.run_forever()

    def close(self):
        self.running = False
        self.loop.run_until_complete(self.tasks['coroutine1'])
        self.loop.run_until_complete(self.tasks['coroutine2'])

    def task_done_callback(self, future):
        for task in self.tasks.values():
            if not task.done():
                return
        self.loop.stop()

    @asyncio.coroutine
    def coroutine1(self):
        while self.running:
            print('coroutine1: do stuff')
            yield from asyncio.sleep(1)

    @asyncio.coroutine
    def coroutine2(self):
        while self.running:
            print('coroutine2: do some other stuff')
            yield from asyncio.sleep(3)

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    daemon = Daemon(loop)
    daemon.run()
    loop.close()

The daemon is gracefully shut down when a SIGINT is received by the program. When this happens the close() method is triggered, which is responsible for notifying all tasks running that they should finish their work and stop. This is done simply by setting running to False . Whenever a task is done, Daemon.task_done_callback is triggered. It checks if all tasks are done, if that's the case, then it stops the loop.

The problem here is that the close() method is not working. That's because I'm calling loop.run_until_complete when the loop is already running (via run_forever ). This is yielding a RuntimeError('This event loop is already running') .

The important thing: coroutine1 needs to finish before coroutine2 because coroutine1 might have problems doing its stuff if coroutine2 is not doing its stuff anymore.

My question here is how do I make sure coroutine1 finishes before coroutine2 ?

this is a way to achieven that. (i remove the signal part and just limited the runtime of coroutine1 ).

import asyncio

class Daemon:

    def __init__(self, loop=asyncio.get_event_loop()):
        self.loop = loop
        self.tasks = {
            'coroutine1': asyncio.ensure_future(self.coroutine1()),
            'coroutine2': asyncio.ensure_future(self.coroutine2())}

    def run(self):
        for task in self.tasks.values():
            task.add_done_callback(self.task_done_callback)
        self.loop.run_forever()

    def task_done_callback(self, future):
        if all(task.done() for task in self.tasks.values()):
            self.loop.stop()

    @asyncio.coroutine
    def coroutine1(self):
        for _ in range(5):
            print('coroutine1: doing stuff')
            yield from asyncio.sleep(0.2)
        print('coroutine1: done!')

    @asyncio.coroutine
    def coroutine2(self):
        while not self.tasks['coroutine1'].done():
            print('coroutine2: doing stuff while coro1 is running')
            yield from asyncio.sleep(0.2)
        print('coroutine2: doing stuff after coro1 has ended')
        yield from asyncio.sleep(1)
        print('coroutine2: done!')

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    daemon = Daemon(loop)
    daemon.run()

the main thing here is to check whether coroutine1 is still running (your self.tasks attribute can be queried for just that).

now to integrate that with signal to stop coroutine1 i'd suggest you register a simple function that sets a flag (eg self.signal_flag ). and then loop inside coroutine1 with something like while not self.signal_flag: ... . these are the changes i would suggest for the complete solution:

class Daemon:

    def __init__(self, loop=asyncio.get_event_loop()):
        ...
        self.signal_flag = False

    def run(self):
        ...
        self.loop.add_signal_handler(signal.SIGINT, self.set_signal_flag)
        self.loop.run_forever()

    def set_signal_flag(self):
        print('caught signal!')
        self.signal_flag = True

    async def coroutine1(self):
        while not self.signal_flag:
            print('coroutine1: doing stuff')
            await asyncio.sleep(0.2)
        print('coroutine1: done!')

also note that there is no need to close the loop. in an IDE where the interpreter keeps running this will make things more complicated than they need to be...

starting from python 3.5 you can (and probably should) use this syntax for your coroutines:

async def coroutine1(self):
    for _ in range(5):
        print('coroutine1: doing stuff')
        await asyncio.sleep(0.2)
    print('coroutine1: done!')

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM