Membatalkan hanya tugas utama dalam program asyncio

Biasanya, jika coroutine dimulai menggunakan fungsi asyncio.run(coroutine), interupsi keyboard (CTRL + C) atau SIGINT akan membatalkan semua tugas yang tertunda pada loop peristiwa. Saya mencari cara agar hanya tugas utama (yang diteruskan ke asyncio.run(coroutine)) yang akan dibatalkan. Idenya adalah bahwa tugas utama kemudian akan mengatur pembatalan semua sub-tugas dalam urutan apa pun yang dianggap sesuai.

Perhatikan sebuah contoh:

import asyncio


async def main():
    foo_task = asyncio.create_task(foo())
    try:
        await asyncio.sleep(10)
        print('main finished')
    finally:
        print('ensuring foo task is finished')
        await foo_task


async def foo():
    await asyncio.sleep(10)
    print('foo finished')


try:
    asyncio.run(main())
except KeyboardInterrupt:
    pass

Saya ingin mengubah kode di atas sehingga jika ada interupsi keyboard atau SIGINT yang dikirim di tengah eksekusi, foo_task akan tetap selesai. Itu harus mencetak yang berikut ini:

ensuring foo task is finished
foo finished

Saya tidak ingin menggunakan pelindung (asyncio.shield(coroutine)) karena saya ingin tugas utama memiliki kontrol penuh atas urutan pembatalan/eksekusi subtugasnya.


python python-3.9 python-asyncio
person Jaanus Varus    schedule 15.03.2021    source sumber

Jawaban (1)


arrow_upward
1
arrow_downward

Saya tidak tahu apakah itu ide yang bagus, tetapi pada sistem operasi mirip Unix Anda dapat mencapai perilaku yang diinginkan dengan penangan sinyal.

import asyncio
from asyncio import tasks
import signal
from typing import Coroutine, Set

to_cancel: Set[Coroutine] = set()  # little workaround to detect the main task

async def main():
    loop = asyncio.get_event_loop()
    loop.add_signal_handler(signal.SIGINT, cancel_main)
    loop.add_signal_handler(signal.SIGTERM, cancel_main)

    foo_task = asyncio.create_task(foo())

    try:
        print("main sleeping")
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        print("main cancelled")
    finally:
        print('ensuring foo task is finished')
        await foo_task
        print('main finished')


async def foo():
    print("foo sleeping")
    await asyncio.sleep(10)
    print("foo finished")


def cancel_main():
    for task in tasks.all_tasks():
        # task.get_coro() for python >= 3.8 else task._coro
        if task.get_coro() in to_cancel and not task.cancelled():
            task.cancel()

if __name__ == "__main__":
    coro = main()
    to_cancel.add(coro)
    asyncio.run(coro)

Hasil

main sleeping
foo sleeping
^C
main cancelled
ensuring foo task is finished
foo finished
main finished
person NobbyNobbs    schedule 15.03.2021