python asyncio graceful shutdown with SIGINT, SIGTERM

python_asyncio_graceful_shutdown.py

import signal
import asyncio
import logging
   
def register_graceful_shutdown(loop: asyncio.AbstractEventLoop):
    """ Register signal handlers """
    for s in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(s, _signal_exit, s, loop) # pass s(signal), loop to _signal_exit
        
def _signal_exit(s, loop: asyncio.AbstractEventLoop):
    logging.info(f"SIGNAL ({s.name})")
    t = loop.create_task(_close_all_tasks())
    t.add_done_callback(lambda _: loop.stop())
    
async def _close_all_tasks():
    """ Close all tasks """
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
    logging.info(f"Closing {len(tasks)} tasks...")
    for task in tasks:
        task.cancel()
    await asyncio.wait(*tasks)


async def dumb_task():
    """ A dumb task """
    try:
        while True:
            logging.info(f"dumb task")
            await asyncio.sleep(1)
    except asyncio.CancelledError:
        logging.info("dumb task is cleaning up itself")
        await asyncio.sleep(1)
        raise
    
if __name__ == "__main__":
    logging.basicConfig(level=logging.DEBUG, format=f"%(asctime)-15s | %(levelname)-8s | [%(name)s] %(message)s" )
    
    loop = asyncio.new_event_loop()
    register_graceful_shutdown(loop)

    t = loop.create_task(dumb_task())
    
    try:        
        loop.run_forever()
    finally:
        loop.close()
        logging.info("Goodbye!")
   

""" On SIGINT
2022-10-28 11:23:58,233 | DEBUG | [asyncio] Using selector: EpollSelector
2022-10-28 11:23:58,234 | INFO | [root] dumb task
2022-10-28 11:23:59,236 | INFO | [root] dumb task
^C2022-10-28 11:23:59,745 | INFO | [root] SIGNAL (SIGINT)
2022-10-28 11:23:59,745 | INFO | [root] Closing tasks...
2022-10-28 11:23:59,746 | INFO | [root] Closing 1 tasks...
2022-10-28 11:23:59,746 | INFO | [root] dumb task is cleaning up itself
2022-10-28 11:24:00,748 | INFO | [root] Goodbye!
"""

""" On SIGTERM
2022-10-28 11:24:45,044 | DEBUG | [asyncio] Using selector: EpollSelector
2022-10-28 11:24:45,045 | INFO | [root] dumb task
2022-10-28 11:24:46,047 | INFO | [root] dumb task
2022-10-28 11:24:47,049 | INFO | [root] dumb task
2022-10-28 11:24:48,050 | INFO | [root] dumb task
2022-10-28 11:24:48,609 | INFO | [root] SIGNAL (SIGTERM)
2022-10-28 11:24:48,609 | INFO | [root] Closing tasks...
2022-10-28 11:24:48,609 | INFO | [root] Closing 1 tasks...
2022-10-28 11:24:48,609 | INFO | [root] dumb task is cleaning up itself
2022-10-28 11:24:49,611 | INFO | [root] Goodbye!
"""