We can simulate shielding the main coroutine from cancellation by using a wrapper coroutine that consumes cancellation requests.
In this tutorial, you will discover how to shield the main coroutine from cancellation.
Let’s get started.
Need to Shield The Main Coroutine From Cancellation
Asyncio tasks can be canceled at any time.
The main coroutine is the entry point into the asyncio program.
You can learn more about the main coroutine in the tutorial:
We may want to protect or shield the main coroutine from cancellation.
This may be for many reasons, such as:
In asyncio, shielding the main coroutine from cancellation can be useful in scenarios where you want to ensure that critical operations or cleanup tasks are completed before the coroutine is terminated. Here are some reasons why you might want to shield the main coroutine from cancellation:
- Cleanup Tasks: We may have critical cleanup tasks that need to be performed before the coroutine exits, such as releasing resources, closing connections, or saving state. By shielding the main coroutine from cancellation, we ensure that these cleanup tasks are executed regardless of whether the coroutine is canceled or not.
- Data Integrity: If the main coroutine is performing operations that could leave data in an inconsistent or corrupted state if abruptly terminated, shielding it from cancellation can help maintain data integrity. This is particularly important in applications where data consistency is crucial, such as databases or file systems.
- Graceful Shutdown: Shielding the main coroutine allows us to implement a graceful shutdown mechanism for our application. Instead of abruptly terminating the coroutine and potentially leaving resources in an inconsistent state, we can use the shielding mechanism to ensure that the application shuts down cleanly, completing any necessary cleanup tasks before exiting.
How can we shield the main coroutine from cancellation in asyncio?
Run loops using all CPUs, download your FREE book to learn how.
How to Shield The Main Coroutine From Cancellation
Asyncio tasks can be shielded.
This can be achieved via the asyncio.shield() function that wraps the target task in a task that absorbs any requests for cancellation.
For example:
|
1 2 3 |
... # create a shield shield = asyncio.shield(task) |
You can learn more about shielding asyncio tasks in the tutorial:
The problem is, the asyncio.shield() function cannot be used to shield the main coroutine.
For example, if we call asyncio.run() to start our asyncio program and attempt to pass it a shielded coroutine, it fails with an exception.
|
1 2 3 |
... # run the asyncio event loop asyncio.run(asyncio.shield(main())) |
This fails because it attempts to convert the main() coroutine into a task before shielding it, and we cannot create a task until the event loop has been started.
One solution is to define a new wrapper coroutine that in turn creates a task for the main() coroutine and shields it. It can then specifically consume any cancellation requests, allowing the true main coroutine to continue running.
For example:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
# main coroutine async def shielded_main(): try: # run the main coroutine as a shielded task await asyncio.shield(asyncio.create_task(main())) except asyncio.CancelledError: # get the current task current = asyncio.current_task() # wait for all other tasks to be done for task in asyncio.all_tasks(): # ensure the task does not await itself if task is not current: # await the other task await task |
This can then be used via asyncio.run().
|
1 2 3 |
... # run the asyncio event loop asyncio.run(shielded_main()) |
This protects the main coroutine from a single cancellation request.
A more robust solution may protect the main coroutine from repeated cancellation requests by using a loop.
For example:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
# main coroutine async def shielded_main(): # schedule the main coroutine as a task task = asyncio.shield(asyncio.create_task(main())) # get the current task current = asyncio.current_task() # loop forever while True: try: # get all tasks all_tasks = asyncio.all_tasks() # remove the current task all_tasks.remove(current) # check for no other tasks if not all_tasks: break # wait for all other tasks to be done for task in all_tasks: # await the other task await task except asyncio.CancelledError: pass |
This is perhaps overly aggressive.
You can learn more about consuming or suppressing the CancelledError exceptions in the tutorial:
Now that we know how to shield the main coroutine from cancellation, let’s look at some worked examples.
Example of Canceling the Main Coroutine
Before we explore how to protect the main coroutine from cancellation, let’s explore an example of canceling the main coroutine.
In this example, we will define a task that gets a list of all running tasks and finds and cancels the main coroutine, in this case named “Task-1”.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
# definition of other task that cancels the main task async def other_task(): # get all tasks tasks = asyncio.all_tasks() # cancel the main task for task in tasks: # check if it is the main task if task.get_name() == 'Task-1': # report a message print(f'Canceling task: {task.get_coro()}') # cancel the task result = task.cancel() # report success of cancellation print(f'Cancellation successful: {result}') # suspend a while await asyncio.sleep(1) # report a final message print('Other task done.') |
The main coroutine will first create and schedule our other_task() in the event loop, then suspend for a moment to allow the new task to run. It will catch the CancelledError and report a message, to confirm that the main coroutine was canceled.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
# main coroutine async def main(): try: # report a message print('Main is running') # schedule the other task task = asyncio.create_task(other_task()) # suspend a while await asyncio.sleep(2) # report a message print('Main is done.') except asyncio.CancelledError: print('Main task cancelled') |
Finally, we can start the asyncio event loop and run our main coroutine.
|
1 2 3 |
... # run the asyncio event loop asyncio.run(main()) |
Tying this together, the complete example is listed below.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# SuperFastPython.com # example of canceling the main coroutine import asyncio # definition of other task that cancels the main task async def other_task(): # get all tasks tasks = asyncio.all_tasks() # cancel the main task for task in tasks: # check if it is the main task if task.get_name() == 'Task-1': # report a message print(f'Canceling task: {task.get_coro()}') # cancel the task result = task.cancel() # report success of cancellation print(f'Cancellation successful: {result}') # suspend a while await asyncio.sleep(1) # report a final message print('Other task done.') # main coroutine async def main(): try: # report a message print('Main is running') # schedule the other task task = asyncio.create_task(other_task()) # suspend a while await asyncio.sleep(2) # report a message print('Main is done.') except asyncio.CancelledError: print('Main task cancelled') # run the asyncio event loop asyncio.run(main()) |
Running the example first starts the asyncio event loop and runs our main coroutine.
The main coroutine runs, reports a message, and then creates and schedules a new task to run our other_task() task before suspending with a sleep.
The other_task() task runs and gets a list of all running tasks. It then searches the list by task name for the main task and cancels it, reporting the details of the main task and that the cancellation request was successful.
The main task resumes and catches the raised CancelledError, confirming that it was canceled.
This highlights how the main task may be canceled.
|
1 2 3 4 |
Main is running Canceling task: <coroutine object main at 0x10e2fc040> Cancellation successful: True Main task cancelled |
Next, let’s look at a naive approach to shielding the main coroutine.
Free Python Asyncio Course
Download your FREE Asyncio PDF cheat sheet and get BONUS access to my free 7-day crash course on the Asyncio API.
Discover how to use the Python asyncio module including how to define, create, and run new coroutines and how to use non-blocking I/O.
Example Cannot use asyncio.shield() For The Main Coroutine
We can explore a naive case of protecting the main coroutine from cancellation.
In this case, we can update the above example to wrap the main() coroutine in a call to asyncio.shield().
|
1 2 3 |
... # run the asyncio event loop asyncio.run(asyncio.shield(main())) |
The hope is that the main coroutine would then be shielded from direct cancellation.
Tying this together, the complete example is listed below.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# SuperFastPython.com # example of shielding the main coroutine from cancellation import asyncio # definition of other task that cancels the main task async def other_task(): # get all tasks tasks = asyncio.all_tasks() # cancel the main task for task in tasks: # check if it is the main task if task.get_name() == 'Task-1': # report a message print(f'Canceling task: {task.get_coro()}') # cancel the task result = task.cancel() # report success of cancellation print(f'Cancellation successful: {result}') # suspend a while await asyncio.sleep(1) # report a final message print('Other task done.') # main coroutine async def main(): try: # report a message print('Main is running') # schedule the other task task = asyncio.create_task(other_task()) # suspend a while await asyncio.sleep(2) # report a message print('Main is done.') except asyncio.CancelledError: print('Main task cancelled') # run the asyncio event loop asyncio.run(asyncio.shield(main())) |
Running the example fails with an exception.
The program fails before the asyncio event loop can even run the main coroutine.
It fails because the call to asyncio.shield() attempts to wrap the main() coroutine in an asyncio.Task, and creating a task requires that an asyncio event loop is running.
This highlights that the naive approach of shielding the main coroutine from cancellation with asyncio.shield() does not work.
|
1 2 3 4 |
ValueError: a coroutine was expected, got <Future pending cb=[shield.<locals>._outer_done_callback() at .../asyncio/tasks.py:908]> Task was destroyed but it is pending! task: <Task pending name='Task-1' coro=<main() running at ...> cb=[shield.<locals>._inner_done_callback() at .../asyncio/tasks.py:891]> sys:1: RuntimeWarning: coroutine 'main' was never awaited |
Next, let’s look at an alternative approach.
Overwhelmed by the python concurrency APIs?
Find relief, download my FREE Python Concurrency Mind Maps
Example of Shielded Main Coroutine
We can explore how to manually shield the main coroutine from cancellation.
This requires developing a main coroutine wrapper that consumes cancellation requests and allows the main coroutine to run.
This approach adds a level of indirection, meaning that the main coroutine is not the actual main coroutine.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
# main coroutine async def shielded_main(): try: # run the main coroutine as a shielded task await asyncio.shield(asyncio.create_task(main())) except asyncio.CancelledError: # get the current task current = asyncio.current_task() # wait for all other tasks to be done for task in asyncio.all_tasks(): # ensure the task does not await itself if task is not current: # await the other task await task |
It then requires that the asyncio.run() call be executed with our shielded_main() coroutine instead of the main() coroutine.
|
1 2 3 |
... # run the asyncio event loop asyncio.run(shielded_main()) |
Canceling the shielded_main() coroutine will consume the CancelledError and allow the true main coroutine main() to continue running.
Tying this together, the complete example is listed below.
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# SuperFastPython.com # example of shielding the main coroutine from cancellation import asyncio # definition of other task that cancels the main task async def other_task(): # cancel the main task for task in asyncio.all_tasks(): # check if it is the main task if task.get_name() == 'Task-1': # report a message print(f'Canceling task: {task.get_coro()}') # cancel the task result = task.cancel() # report success of cancellation print(f'Cancellation successful: {result}') # suspend a while await asyncio.sleep(1) # report a final message print('Other task done.') # main coroutine async def main(): # report a message print('Main is running') # schedule the other task task = asyncio.create_task(other_task()) # suspend a while await asyncio.sleep(2) # report a message print('Main is done.') # main coroutine async def shielded_main(): try: # run the main coroutine as a shielded task await asyncio.shield(asyncio.create_task(main())) except asyncio.CancelledError: # get the current task current = asyncio.current_task() # wait for all other tasks to be done for task in asyncio.all_tasks(): # ensure the task does not await itself if task is not current: # await the other task await task # run the asyncio event loop asyncio.run(shielded_main()) |
Running the example first starts and runs the asyncio event loop with our shielded_main() coroutine.
The shielded_main() coroutine runs and creates, schedules, and awaits the main() coroutine as a shielded task.
The main() task runs and reports a message, then creates and schedules the other_task() coroutine before suspending for two seconds.
The other_task() coroutine runs and attempts to cancel the main coroutine. It finds the shielded_main() and requests it to cancel, reports the details of the coroutine, and reports that the cancellation request was successful.
The shielded_main() resumes and catches a CancelledError. It then awaits all running tasks.
The main() coroutine is not canceled and resumes before reporting a final message and terminating normally.
The shielded_main() coroutine resumes and terminates.
This highlights how we can shield the main coroutine using a wrapper coroutine.
|
1 2 3 4 5 |
Main is running Canceling task: <coroutine object shielded_main at 0x10d9035b0> Cancellation successful: True Other task done. Main is done. |
Takeaways
You now know how to shield the main coroutine from cancellation.
Did I make a mistake? See a typo?
I’m a simple humble human. Correct me, please!
Do you have any additional tips?
I’d love to hear about them!
Do you have any questions?
Ask your questions in the comments below and I will do my best to answer.
Photo by Aditya Chinchure on Unsplash
Do you have any questions?