Hi folks! Today, I want to illustrate how confusing it might be to gracefully shut down an app without losing anything important. asyncio When shutting down an application, it's crucial to make sure all running tasks have finished execution before closing down the application. A hard exit can cause data loss and other unforeseen issues. Therefore, it's recommended to use a graceful shutdown mechanism that allows running tasks to complete before shutting down the application. asyncio To accomplish this, the module provides a function that can be used to ensure a task completes its execution even if the application is shutting down. , in this article, I will illustrate that . asyncio shield() However it is not as simple as it may seem 1. Without Any Protection Let's start with a simple program that runs multiple tasks in the background and waits until they're done. Then, I will try to interrupt it in the middle. # example1.py

import asyncio


async def worker(n: int) -> None:
    print(f"[{n}] Started!")
    try:
        # this is a task that shouldn't be canceled in the middle
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        print(f"[{n}] Canceled (this is bad)!")
    else:
        print(f"[{n}] Successfully done!")


async def main() -> None:
    # create 6 unprotected tasks
    tasks = []
    for n in range(6):
        tasks.append(asyncio.create_task(worker(n)))

    # wait for all tasks to finish
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("App was interrupted")
    else:
        print("App was finished gracefully") Let's run and interrupt it afterwards: example1.py > python3 example1.py
[0] Started!
[1] Started!
[2] Started!
[3] Started!
[4] Started!
[5] Started!
^C[0] Canceled (this is bad)!
[1] Canceled (this is bad)!
[2] Canceled (this is bad)!
[3] Canceled (this is bad)!
[4] Canceled (this is bad)!
[5] Canceled (this is bad)!
App was interrupted As you can see, when I interrupted the script (the sign appears when I press ), all tasks were immediately cancelled without waiting for their completion. But this is quite expected since we didn't even bother to protect the tasks. ^C ctrl+C 2. With “shield” Let's update by using to prevent 3 of 6 tasks from being cancelled, as described in the . main() asyncio.shield() documentation # example2.py

import asyncio


async def worker(n: int) -> None:
    ...


async def main() -> None:
    # create 6 tasks, shield only first 3
    tasks = []
    for n in range(6):
        task = asyncio.create_task(worker(n))
        if n < 3:
            # protect task from being canceled 
			# (spoiler: it will be canceled anyway)
            task = asyncio.shield(task)

        tasks.append(task)

    # wait for all tasks to finish
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    ... > python3 example2.py
[0] Started!
[1] Started!
[2] Started!
[3] Started!
[4] Started!
[5] Started!
^C[3] Canceled (this is bad)!
[4] Canceled (this is bad)!
[5] Canceled (this is bad)!
[2] Canceled (this is bad)!
[0] Canceled (this is bad)!
[1] Canceled (this is bad)!
App was interrupted Do you see the difference from the previous example ( )? There isn't any. Nothing has changed. Why is it so? example1.py It is because protects a coroutine if the parent coroutine (inside which the is used) gets cancelled. So it doesn't protect the shielded coroutine from direct cancellation. shield() only shield() Let me explain. If you shut down an app, you'll likely do it the conventional way by following these steps: - retrieve all tasks tasks = asyncio.all_tasks() - cancel all tasks [t.cancel() for t in tasks] - wait for the cancelled tasks to complete loop.run_until_complete(gather(*tasks)) Since works by creating an inner task that is also included in the call, it also receives a cancellation exception, just like everything else. shield() all_tasks() Now, let's take a look at a from (Python 3.11.2) that is called before closing the loop: piece of code asyncio def _cancel_all_tasks(loop):
    to_cancel = tasks.all_tasks(loop)
    if not to_cancel:
        return

    for task in to_cancel:
        task.cancel()

    loop.run_until_complete(tasks.gather(*to_cancel, return_exceptions=True)) As we can see, nothing fancy, just the same three steps. 3. Handle Signals Now it's time to actually protect tasks from unexpected cancellation. The way we do this is by implementing the following steps: Make a set of tasks that we want to protect. Handle basic interruption signals (SIGHUP, SIGTERM, SIGINT) to implement our own shutdown logic. Within the shutdown function, cancel only unprotected tasks. # example3.py

import asyncio
import signal

# tasks that shouldn't be canceled
_DO_NOT_CANCEL_TASKS: set[asyncio.Task] = set()


def protect(task: asyncio.Task) -> None:
    _DO_NOT_CANCEL_TASKS.add(task)


def shutdown(sig: signal.Signals) -> None:
    print(f"Received exit signal {sig.name}")

    all_tasks = asyncio.all_tasks()
    tasks_to_cancel = all_tasks - _DO_NOT_CANCEL_TASKS

    for task in tasks_to_cancel:
        task.cancel()

    print(f"Cancelled {len(tasks_to_cancel)} out of {len(all_tasks)} tasks")


def setup_signal_handler() -> None:
    loop = asyncio.get_running_loop()

    for sig in (signal.SIGHUP, signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, shutdown, sig)

async def worker(n: int) -> None:
    ...


async def main() -> None:
    # setup graceful shutdown
    setup_signal_handler()

    # protect main task from being canceled, 
		# otherwise it will cancel all other tasks
    protect(asyncio.current_task())

    # create 6 tasks, shield only first 3
    tasks = []
    for n in range(6):
        task = asyncio.create_task(worker(n))
        if n < 3:
            protect(task)

        tasks.append(task)

    # wait for all tasks to finish
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    ... > python3 example3.py
[0] Started!
[1] Started!
[2] Started!
[3] Started!
[4] Started!
[5] Started!
^CReceived exit signal SIGINT
Cancelled 3 out of 7 tasks
[5] Canceled (this is bad)!
[3] Canceled (this is bad)!
[4] Canceled (this is bad)!
[0] Successfully done!
[1] Successfully done!
[2] Successfully done!
App was finished gracefully Here we go! Now the group of protected tasks has not been interrupted in the middle. However, ensure that the tasks you are waiting for are tasks. Otherwise, there is a risk of the application being (SIGKILL). not long-running forcefully terminated Conclusion To sum up, when it comes to shutting down an application, it's important to protect against potential unforeseen problems. While provides a function to ensure task completion, this alone is to ensure a graceful shutdown. Instead, custom shutdown logic that protects tasks from unexpected cancellation is necessary. By , you can guarantee that your app will be shut down gracefully. asyncio asyncio shield() not enough handling interruption signals asyncio P.S. All the code mentioned in this article can be obtained from this . repository

This story contains new, firsthand information uncovered by the writer.

The code in this story is for educational purposes. The readers are solely responsible for whatever they build with it.

The writer is smart, but don't just like, take their word for it. #DoYourOwnResearch

Asyncio: How to Say Goodbye Without Losing Your Data

Too Long; Didn't Read

@taras

Credibility

RELATED STORIES