Pythonのasyncioでgraceful shutdown

プログラム実行中にハンドリングできない例外が発生した場合に、後処理を行ってからプログラムを終了させたい場合がある (graceful shutdown)。通常はtry~exceptで実装するが、asyncioを使っていると各コルーチンでハンドリングできない例外が送出された場合にtry~exceptで例外を処理することができない。

"python asyncio graceful shutdown"などで検索すると色々な方法が出てくるが、すぐに使い回せるコードが見つからなかったためまとめておく。ここではasyncioの挙動については割愛する（詳細は末尾の参考文献を参照）。

ハンドリングするべき例外はsignalとExceptionがあり、それぞれのハンドラーを作成してからイベントループに登録しループを回せばよい。追加のエラーハンドリングはshutdown()に実装する。

検証環境はCentOS7, Python3.10。指定するsignalはOSによって異なる可能性がある。

import asyncio
import signal


async def task() -> None:
    cnt = 5
    while True:
        print(f"task is running... will raise Exception in {cnt} seconds")
        await asyncio.sleep(1)
        cnt -= 1
        if cnt == 0:
            raise Exception("something wrong happened")


async def shutdown(loop, signal=None) -> None:
    if signal:
        print(f"received exit signal {signal.name}...")
    tasks = [t for t in asyncio.all_tasks() if t is not
             asyncio.current_task()]

    [task.cancel() for task in tasks]

    print(f"cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    print("cancelled tasks")

    # add custom shutdown logic here

    loop.stop()
    print("gracefully shutdown the service.")


def handle_exception(loop, context) -> None:
    # context["message"] will always be there; but context["exception"] may not
    msg = context.get("exception", context["message"])
    print(f"caught exception: {msg}")
    asyncio.create_task(shutdown(loop=loop))


def main() -> None:
    loop = asyncio.get_event_loop()
    signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)  # may want to catch other signals too
    for s in signals:
        loop.add_signal_handler(
            s, lambda s=s: asyncio.create_task(shutdown(loop=loop, signal=s)))

    loop.set_exception_handler(handle_exception)

    try:
        loop.create_task(task())
        loop.run_forever()
    finally:
        loop.close()


if __name__ == "__main__":
    main()

元ネタはこちらの記事。不要な部分を削って必要最小限のコードにした。

https://www.roguelynn.com/words/asyncio-graceful-shutdowns/
https://www.roguelynn.com/words/asyncio-exception-handling/