0

I have Python program that runs multiple asyncio tasks using asyncio.TaskGroup.

All of the tasks are waiting for some event to happen and I want them being cancelled when any of them completes.

I want that the program can also be closed by sending an interrupt signal (Ctrl+C in terminal), but to prevent it being closed accidentally, that should succeed only when it is sent two times within two seconds.

try-except blocks can be used around asyncio.run to catch KeyboardInterrupt, but that closes all the tasks and I would need to recreate them if interrupt happened only once and I want to continue the tasks.

I found a solution already that seems to work, but I am not 100% sure if it's the correct way.

Afkaaja
  • 47
  • 5

1 Answers1

0

One way is to pass list to each task as a container where references to them can be kept. This way any of the tasks can stop all the other tasks after it completes.

By utilizing asyncio.loop.add_signal_handler and keeping track of when the last SIGINT was caught it is possible determine whether there were two interrupts within two seconds or not.

import asyncio
from asyncio import Task
from random import random
from signal import SIGINT
from time import time


def cancel_all(tasks: list[Task]):
    for task in tasks:
        if task != asyncio.current_task() and not task.done():
            task.cancel()


def ask_exit(tasks: list[Task], last_sigint: list[float]):
    time_now = time()

    if time_now < last_sigint[0] + 2:
        cancel_all(tasks)
        return

    last_sigint[0] = time_now
    print("Are you sure you want to exit?")


async def some_task(tasks: list[Task]):
    await asyncio.sleep(5 + random() * 5)  # event to happen
    cancel_all(tasks)


async def run_tasks() -> None:
    async with asyncio.TaskGroup() as group:
        tasks: list[Task] = []
        tasks.append(group.create_task(some_task(tasks)))
        tasks.append(group.create_task(some_task(tasks)))

        last_sigint: list[float] = [0.0]

        asyncio.get_running_loop().add_signal_handler(
            SIGINT, ask_exit, tasks, last_sigint
        )

    asyncio.get_running_loop().remove_signal_handler(SIGINT)
    print("All tasks done.")


asyncio.run(run_tasks())
Afkaaja
  • 47
  • 5