Awaiting a asyncio Future after Cancelling it - python

Looking at the asyncio docs, I came across this example
async def main():
# Create a "cancel_me" Task
task = asyncio.create_task(cancel_me())
# Wait for 1 second
await asyncio.sleep(1)
task.cancel()
try:
await task
except asyncio.CancelledError:
print("main(): cancel_me is cancelled now")
asyncio.run(main())
After task.cancel(), what is the purpose of doing await task? Is this to wait for the future to be finished if it was ever shielded from cancellation?
In other words, why not:
async def main():
# Create a "cancel_me" Task
task = asyncio.create_task(cancel_me())
# Wait for 1 second
await asyncio.sleep(1)
task.cancel()
asyncio.run(main())

From the documentation of cancel() (under asyncio.Task):
This arranges for a CancelledError exception to be thrown into the wrapped coroutine on the next cycle of the event loop.
The coroutine then has a chance to clean up or even deny the request by suppressing the exception with a try … … except CancelledError … finally block.
When asyncio.CancelledError is thrown into cancel_me(), execution resumes in the except asyncio.CancelledError block. For the snippet provided with the documentation, it does not in fact make any difference whether cancel_me() is awaited or not after cancelling, because the exception handling block executes synchronously.
On the other hand, if the exception handling block did perform asynchronous operations, the difference would become visible:
async def cancel_me():
print('cancel_me(): before sleep')
try:
# Wait for 1 hour
await asyncio.sleep(3600)
except asyncio.CancelledError:
await asyncio.sleep(1)
print('cancel_me(): cancel sleep, this never gets printed')
raise
finally:
print('cancel_me(): after sleep')
async def main():
# Create a "cancel_me" Task
task = asyncio.create_task(cancel_me())
# Wait for 1 second
await asyncio.sleep(1)
task.cancel()
print("main(): cancel_me is cancelled now")
asyncio.run(main())
# Expected output:
#
# cancel_me(): before sleep
# main(): cancel_me is cancelled now
# cancel_me(): after sleep
The last, surprising print takes place because of the following:
main() returns after its last print
asyncio.run() tries to cancel all pending tasks
cancel_me(), albeit already cancelled, is still pending, awaiting on the exception block sleep
the finally clause in cancel_me() is executed and the even loop terminates
Also worth noting: given that asyncio.run() throws a CancelledError into all the tasks that are still pending, if cancel_me() had not been cancelled already, the except asyncio.CancelledError block would execute in its entirety.

Related

Correctly adding a signal handler to Asyncio code

I'm trying to modify the graceful shutdown example from RogueLynn to cancel running processes that were spawned by the tasks.
Below is a minimal example to demonstrate the issue I'm facing. With this example, I get a warning message that the callback function isn't awaited and when I do try to terminate the script, the asyncio.gather call doesn't seem to complete. Any idea how to resolve this such that the shutdown callback executes completely?
import asyncio
import functools
import signal
async def run_process(time):
try:
print(f"Starting to sleep for {time} seconds")
await asyncio.sleep(time)
print(f"Completed sleep of {time} seconds")
except asyncio.CancelledError:
print("Received cancellation terminating process")
raise
async def main():
tasks = [run_process(10), run_process(5), run_process(2)]
for future in asyncio.as_completed(tasks):
try:
await future
except Exception as e:
print(f"Caught exception: {e}")
async def shutdown(signal, loop):
# Cancel running tasks on keyboard interrupt
print(f"Running shutdown")
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
[task.cancel() for task in tasks]
await asyncio.gather(*tasks, return_exceptions=True)
print("Finished waiting for cancelled tasks")
loop.stop()
try:
loop = asyncio.get_event_loop()
signals = (signal.SIGINT,)
for sig in signals:
loop.add_signal_handler(sig, functools.partial(asyncio.create_task, shutdown(sig, loop)))
loop.run_until_complete(main())
finally:
loop.close()
Output when run to completion:
Starting to sleep for 2 seconds
Starting to sleep for 10 seconds
Starting to sleep for 5 seconds
Completed sleep of 2 seconds
Completed sleep of 5 seconds
Completed sleep of 10 seconds
/home/git/envs/lib/python3.8/asyncio/unix_events.py:140: RuntimeWarning: coroutine 'shutdown' was never awaited
del self._signal_handlers[sig]
And output when script is interrupted:
Starting to sleep for 2 seconds
Starting to sleep for 10 seconds
Starting to sleep for 5 seconds
Completed sleep of 2 seconds
^CRunning shutdown
Received cancellation terminating process
Received cancellation terminating process
Task was destroyed but it is pending!
task: <Task pending name='Task-5' coro=<shutdown() running at ./test.py:54> wait_for=<_GatheringFuture finished result=[CancelledError(), CancelledError(), CancelledError()]>>
Traceback (most recent call last):
File "./test.py", line 65, in <module>
loop.run_until_complete(main())
File "/home/git/envs/lib/python3.8/asyncio/base_events.py", line 616, in run_until_complete
return future.result()
asyncio.exceptions.CancelledError
The CancelledError you see results from your as_completed for loop.
If you just want to fix it, you could add an exception handling for it, e.g.
[...]
try:
await future
except Exception as e:
print(f"Caught exception: {e}")
except asyncio.CancelledError:
print("task was cancelled")
[...]
Note that, you will find a warning telling you that Task was destroyed but it is pending! for your shutdown task, which you could just ignore. I guess it is because you stop the loop within a task.
Still, I would like to point at the difference between co-routines, tasks, and futures, see https://docs.python.org/3/library/asyncio-task.html.
What you call tasks and future are co-routines.
For the type of problem that you are trying to solve, I would advise to have a look at Asynchronous Context Managers. Graceful shutdown sounds to me like you want to close some database connections or dump some process variables... Here, you could have a look at https://docs.python.org/3/library/contextlib.html#contextlib.asynccontextmanager
However, if things become more complex, you may want to write such a signal handler which adds its own task to the loop. In this case, I would advise to create the relevant tasks explicitly with asyncio.create_task(coro, name="my-task-name") so you can select exactly the tasks you want to cancel first by name, e.g.
tasks = [
task for task in asyncio.all_tasks()
if task.get_name().startswith("my-task")
]
Otherwise, you may accidentally cancel a cleanup-task.

Why does asyncio.sleep alllow for the task to not be immediately cancelled?

import asyncio
async def cancel_me():
print('cancel_me(): before sleep')
try:
# Wait for 1 hour
await asyncio.sleep(3600)
except asyncio.CancelledError:
print('cancel_me(): cancel sleep')
raise
finally:
print('cancel_me(): after sleep')
async def main():
# Create a "cancel_me" Task
task = asyncio.create_task(cancel_me())
# if comment out this line we go directly to -> main(): cancel_me is cancelled now
await asyncio.sleep(1)
# This arranges for a CancelledError exception to be thrown into the wrapped
# coroutine on the next cycle of the event loop.
task.cancel()
try:
await task
except asyncio.CancelledError:
print("main(): cancel_me is cancelled now")
asyncio.run(main())
With await asyncio.sleep(1), we have:
# Output
# ------------------------
# cancel_me(): before sleep
# cancel_me(): cancel sleep
# cancel_me(): after sleep
# main(): cancel_me is cancelled now
Without await asyncio.sleep(1), we have:
# Output
# ------------------------
# main(): cancel_me is cancelled now
Why is there this difference in output just from commenting out a asyncio.sleep command?
Edit (as a reaction to a comment below):
So, await asyncio.sleep allows other tasks to start? This seems counter-intuitive. Does it mean that any await after a create_task allows any created/scheduled tasks to be started?
Also, consider the following picture:
I got this picture by commenting out everything that goes below await asyncio.sleep(1), in the main(), and putting await asyncio.sleep(5), instead. In this picture, we see that cancel_me(): cancel sleep gets printed, even when we have no task cancel... Not only that, but we do the print in the exception, and don't raise any exception.
With the await asyncio.sleep(1) commented out, you cancel the task before it even starts running, so it just doesn't run.

Send Ctrl + C to asyncio Task

So I have two tasks running until complete in my event loop. I want to handle KeyboardInterrupt when running those Tasks and send the same signal to the Tasks in the event loop when receiving it.
loop = asyncio.get_event_loop()
try:
loop.run_until_complete(asyncio.gather(
task_1(),
task_2()
))
except KeyboardInterrupt:
pass
# Here I want to send sigterm to the tasks in event loop
Is there any way to do this?
Use asyncio.run() instead. asyncio.run() cancels all tasks/futures when receiving exceptions.
asyncio.gather() would then cancel the tanks when itself is cancelled. You may capture asyncio.CancelledError around asyncio.sleep() to see. Discard it would keep the task running.
async def task_1():
try:
await asyncio.sleep(10)
except asyncio.CancelledError as ex:
print('task1', type(ex))
raise
async def task_2():
try:
await asyncio.sleep(10)
except asyncio.CancelledError as ex:
print('task2', type(ex))
raise
async def main():
await asyncio.gather(task_1(), task_2())
if __name__ == '__main__':
asyncio.run(main())
Otherwise, you need to keep references to tasks, cancel them, and re-run the event loop to propagate CancelledError, like what asynio.run() do under the hood.

How can I run an asyncio loop as long as there are pending cancellation-shielded tasks left but no longer?

I'm trying to add some code to my existing asyncio loop to provide for a clean shutdown on Ctrl-C. Below is an abstraction of the sort of thing it's doing.
import asyncio, signal
async def task1():
print("Starting simulated task1")
await asyncio.sleep(5)
print("Finished simulated task1")
async def task2():
print("Starting simulated task2")
await asyncio.sleep(5)
print("Finished simulated task2")
async def tasks():
await task1()
await task2()
async def task_loop():
try:
while True:
await asyncio.shield(tasks())
await asyncio.sleep(60)
except asyncio.CancelledError:
print("Shutting down task loop")
raise
async def aiomain():
loop = asyncio.get_running_loop()
task = asyncio.Task(task_loop())
loop.add_signal_handler(signal.SIGINT, task.cancel)
await task
def main():
try:
asyncio.run(aiomain())
except asyncio.CancelledError:
pass
#def main():
# try:
# loop = asyncio.get_event_loop()
# loop.create_task(aiomain())
# loop.run_forever()
# except asyncio.CancelledError:
# pass
if __name__ == '__main__':
main()
In this example, imagine that the sequence of task1 and task2 needs to be finished once it's started, or some artifacts will be left in an inconsistent state. (Hence the asyncio.shield wrapper around calling tasks.)
With the code as above, if I interrupt the script soon after it starts and it's just printed Starting simulated task1 then the loop stops and task2 never gets started. If I try switching to the version of main that's commented out then that one never exits, even though the loop is properly cancelled and nothing further happens at least for several minutes. It does have a bit of progress in that it at least finishes any in-progress sequence of task1 and task2.
Some possible solutions from brainstorming, though I still get the feeling there must be something simpler that I'm missing:
Create a wrapper around asyncio.shield which increments a variable synchronized by an asyncio.Condition object, runs the shielded function, then decrements the variable. Then, in aiomain in a CancelledError handler, wait for the variable to reach zero before reraising the exception. (In an implementation, I would probably go for combining all the parts of this into one class with __aexit__ implementing the wait for zero on CancelledError logic.)
Skip using asyncio's cancellation mechanism entirely, and instead use an asyncio.Event or similar to allow for interruption points or interruptible sleeps. Though this does seem like it would be more invasive requiring me to specify what points are considered interruptible, as opposed to declaring what sequences need to be shielded from cancellation.
This is a very good question. I have learned some things while working out an answer, so I hope you are still monitoring this thread.
The first thing to investigate is, how does the shield() method work? On this point, the docs are confusing to say the least. I couldn't figure it out until I read the standard library test code in test_tasks.py. Here is my understanding:
Consider this code fragment:
async def coro_a():
await asyncio.sheild(task_b())
...
task_a = asyncio.create_task(coro_a())
task_a.cancel()
When the task_a.cancel() statement is executed, task_a is indeed cancelled. The await statement throws a CancelledError immediately, without waiting for task_b to finish. But task_b continues to run. The outer task (a) stops but the inner task (b) doesn't.
Here is a modified version of your program that illustrates this. The major change is to insert a wait in your CancelledError exception handler, to keep your program alive a few seconds longer. I'm running on Windows and that's why I changed your signal handler a little bit also, but that's a minor point. I also added time stamps to the print statements.
import asyncio
import signal
import time
async def task1():
print("Starting simulated task1", time.time())
await asyncio.sleep(5)
print("Finished simulated task1", time.time())
async def task2():
print("Starting simulated task2", time.time())
await asyncio.sleep(5)
print("Finished simulated task2", time.time())
async def tasks():
await task1()
await task2()
async def task_loop():
try:
while True:
await asyncio.shield(tasks())
await asyncio.sleep(60)
except asyncio.CancelledError:
print("Shutting down task loop", time.time())
raise
async def aiomain():
task = asyncio.create_task(task_loop())
KillNicely(task)
try:
await task
except asyncio.CancelledError:
print("Caught CancelledError", time.time())
await asyncio.sleep(5.0)
raise
class KillNicely:
def __init__(self, cancel_me):
self.cancel_me = cancel_me
self.old_sigint = signal.signal(signal.SIGINT,
self.trap_control_c)
def trap_control_c(self, signum, stack):
if signum != signal.SIGINT:
self.old_sigint(signum, stack)
else:
print("Got Control-C", time.time())
print(self.cancel_me.cancel())
def main():
try:
asyncio.run(aiomain())
except asyncio.CancelledError:
print("Program exit, cancelled", time.time())
# Output when ctrlC is struck during task1
#
# Starting simulated task1 1590871747.8977509
# Got Control-C 1590871750.8385916
# True
# Shutting down task loop 1590871750.8425908
# Caught CancelledError 1590871750.8435903
# Finished simulated task1 1590871752.908434
# Starting simulated task2 1590871752.908434
# Program exit, cancelled 1590871755.8488846
if __name__ == '__main__':
main()
You can see that your program didn't work because it exited as soon as task_loop was cancelled, before task1 and task2 had a chance to finish. They were still there all along (or rather they would have been there, if the program continued to run).
This illustrates how shield() and cancel() interact, but it doesn't actually solve your stated problem. For that, I think, you need to have an awaitable object that you can use to keep the program alive until the vital tasks are finished. This object needs to be created at the top level and passed down the stack to the place where the vital tasks are executing. Here is a program that is similar to yours, but preforms the way you want.
I did three runs: (1) control-C during task1, (2) control-C during task2, (3) control-C after both tasks were finished. In the first two cases the program continued until task2 was finished. In the third case it ended immediately.
import asyncio
import signal
import time
async def task1():
print("Starting simulated task1", time.time())
await asyncio.sleep(5)
print("Finished simulated task1", time.time())
async def task2():
print("Starting simulated task2", time.time())
await asyncio.sleep(5)
print("Finished simulated task2", time.time())
async def tasks(kwrap):
fut = asyncio.get_running_loop().create_future()
kwrap.awaitable = fut
await task1()
await task2()
fut.set_result(1)
async def task_loop(kwrap):
try:
while True:
await asyncio.shield(tasks(kwrap))
await asyncio.sleep(60)
except asyncio.CancelledError:
print("Shutting down task loop", time.time())
raise
async def aiomain():
kwrap = KillWrapper()
task = asyncio.create_task(task_loop(kwrap))
KillNicely(task)
try:
await task
except asyncio.CancelledError:
print("Caught CancelledError", time.time())
await kwrap.awaitable
raise
class KillNicely:
def __init__(self, cancel_me):
self.cancel_me = cancel_me
self.old_sigint = signal.signal(signal.SIGINT,
self.trap_control_c)
def trap_control_c(self, signum, stack):
if signum != signal.SIGINT:
self.old_sigint(signum, stack)
else:
print("Got Control-C", time.time())
print(self.cancel_me.cancel())
class KillWrapper:
def __init__(self):
self.awaitable = asyncio.get_running_loop().create_future()
self.awaitable.set_result(0)
def main():
try:
asyncio.run(aiomain())
except asyncio.CancelledError:
print("Program exit, cancelled", time.time())
# Run 1 Control-C during task1
# Starting simulated task1 1590872408.6737766
# Got Control-C 1590872410.7344952
# True
# Shutting down task loop 1590872410.7354996
# Caught CancelledError 1590872410.7354996
# Finished simulated task1 1590872413.6747622
# Starting simulated task2 1590872413.6747622
# Finished simulated task2 1590872418.6750958
# Program exit, cancelled 1590872418.6750958
#
# Run 1 Control-C during task2
# Starting simulated task1 1590872492.927735
# Finished simulated task1 1590872497.9280624
# Starting simulated task2 1590872497.9280624
# Got Control-C 1590872499.5973852
# True
# Shutting down task loop 1590872499.5983844
# Caught CancelledError 1590872499.5983844
# Finished simulated task2 1590872502.9274273
# Program exit, cancelled 1590872502.9287038
#
# Run 1 Control-C after task2 -> immediate exit
# Starting simulated task1 1590873694.2925708
# Finished simulated task1 1590873699.2928336
# Starting simulated task2 1590873699.2928336
# Finished simulated task2 1590873704.2938952
# Got Control-C 1590873706.0790765
# True
# Shutting down task loop 1590873706.0804725
# Caught CancelledError 1590873706.0804725
# Program exit, cancelled 1590873706.0814824
Here is what I ended up using:
import asyncio, signal
async def _shield_and_wait_body(coro, finish_event):
try:
await coro
finally:
finish_event.set()
async def shield_and_wait(coro):
finish_event = asyncio.Event()
task = asyncio.shield(_shield_and_wait_body(coro, finish_event))
try:
await task
except asyncio.CancelledError:
await finish_event.wait()
raise
def shield_and_wait_decorator(coro_fn):
return lambda *args, **kwargs: shield_and_wait(coro_fn(*args, **kwargs))
async def task1():
print("Starting simulated task1")
await asyncio.sleep(5)
print("Finished simulated task1")
async def task2():
print("Starting simulated task2")
await asyncio.sleep(5)
print("Finished simulated task2")
#shield_and_wait_decorator
async def tasks():
await task1()
await task2()
async def task_loop():
try:
while True:
# Alternative to applying #shield_and_wait_decorator to tasks()
#await shield_and_wait(tasks())
await tasks()
await asyncio.sleep(60)
except asyncio.CancelledError:
print("Shutting down task loop")
raise
def sigint_handler(task):
print("Cancelling task loop")
task.cancel()
async def aiomain():
loop = asyncio.get_running_loop()
task = asyncio.Task(task_loop())
loop.add_signal_handler(signal.SIGINT, sigint_handler, task)
await task
def main():
try:
asyncio.run(aiomain())
except asyncio.CancelledError:
pass
if __name__ == '__main__':
main()
Similar to the answer by Paul Cornelius, this inserts a wait for the subtask to finish before allowing the CancelledError to propagate up the call chain. However, it does not require touching the code other than at the point you would be calling asyncio.shield.
(In my actual use case, I had three loops running simultaneously, using an asyncio.Lock to make sure one task or sequence of tasks finished before another would start. I also had an asyncio.Condition on that lock communicating from one coroutine to another. When I tried the approach of waiting in aiomain or main for all shielded tasks to be done, I ran into an issue where a cancelled parent released the lock, then a shielded task tried to signal the condition variable using that lock, giving an error. It also didn't make sense to move acquiring and releasing the lock into the shielded task - that would result in task B still running in the sequence: shielded task A starts, coroutine for task B expires its timer and blocks waiting for the lock, Control+C. By putting the wait at the point of the shield_and_wait call, on the other hand, it neatly avoided prematurely releasing the lock.)
One caveat: it seems that shield_and_wait_decorator doesn't work properly on class methods.

asyncio: prevent task from being cancelled twice

Sometimes, my coroutine cleanup code includes some blocking parts (in the asyncio sense, i.e. they may yield).
I try to design them carefully, so they don't block indefinitely. So "by contract", coroutine must never be interrupted once it's inside its cleanup fragment.
Unfortunately, I can't find a way to prevent this, and bad things occur when it happens (whether it's caused by actual double cancel call; or when it's almost finished by itself, doing cleanup, and happens to be cancelled from elsewhere).
Theoretically, I can delegate cleanup to some other function, protect it with a shield, and surround it with try-except loop, but it's just ugly.
Is there a Pythonic way to do so?
#!/usr/bin/env python3
import asyncio
#asyncio.coroutine
def foo():
"""
This is the function in question,
with blocking cleanup fragment.
"""
try:
yield from asyncio.sleep(1)
except asyncio.CancelledError:
print("Interrupted during work")
raise
finally:
print("I need just a couple more seconds to cleanup!")
try:
# upload results to the database, whatever
yield from asyncio.sleep(1)
except asyncio.CancelledError:
print("Interrupted during cleanup :(")
else:
print("All cleaned up!")
#asyncio.coroutine
def interrupt_during_work():
# this is a good example, all cleanup
# finishes successfully
t = asyncio.async(foo())
try:
yield from asyncio.wait_for(t, 0.5)
except asyncio.TimeoutError:
pass
else:
assert False, "should've been timed out"
t.cancel()
# wait for finish
try:
yield from t
except asyncio.CancelledError:
pass
#asyncio.coroutine
def interrupt_during_cleanup():
# here, cleanup is interrupted
t = asyncio.async(foo())
try:
yield from asyncio.wait_for(t, 1.5)
except asyncio.TimeoutError:
pass
else:
assert False, "should've been timed out"
t.cancel()
# wait for finish
try:
yield from t
except asyncio.CancelledError:
pass
#asyncio.coroutine
def double_cancel():
# cleanup is interrupted here as well
t = asyncio.async(foo())
try:
yield from asyncio.wait_for(t, 0.5)
except asyncio.TimeoutError:
pass
else:
assert False, "should've been timed out"
t.cancel()
try:
yield from asyncio.wait_for(t, 0.5)
except asyncio.TimeoutError:
pass
else:
assert False, "should've been timed out"
# although double cancel is easy to avoid in
# this particular example, it might not be so obvious
# in more complex code
t.cancel()
# wait for finish
try:
yield from t
except asyncio.CancelledError:
pass
#asyncio.coroutine
def comain():
print("1. Interrupt during work")
yield from interrupt_during_work()
print("2. Interrupt during cleanup")
yield from interrupt_during_cleanup()
print("3. Double cancel")
yield from double_cancel()
def main():
loop = asyncio.get_event_loop()
task = loop.create_task(comain())
loop.run_until_complete(task)
if __name__ == "__main__":
main()
I ended up writing a simple function that provides a stronger shield, so to speak.
Unlike asyncio.shield, which protects the callee, but raises CancelledError in its caller, this function suppresses CancelledError altogether.
The drawback is that this function doesn't allow you to handle CancelledError later. You won't see whether it has ever happened. Something slightly more complex would be required to do so.
#asyncio.coroutine
def super_shield(arg, *, loop=None):
arg = asyncio.async(arg)
while True:
try:
return (yield from asyncio.shield(arg, loop=loop))
except asyncio.CancelledError:
continue
I found WGH's solution when encountering a similar problem. I'd like to await a thread, but regular asyncio cancellation (with or without shield) will just cancel the awaiter and leave the thread floating around, uncontrolled. Here is a modification of super_shield that optionally allows reacting on cancel requests and also handles cancellation from within the awaitable:
await protected(aw, lambda: print("Cancel request"))
This guarantees that the awaitable has finished or raised CancelledError from within. If your task could be cancelled by other means (e.g. setting a flag observed by a thread), you can use the optional cancel callback to enable cancellation.
Implementation:
async def protect(aw, cancel_cb: typing.Callable = None):
"""
A variant of `asyncio.shield` that protects awaitable as well
as the awaiter from being cancelled.
Cancellation events from the awaiter are turned into callbacks
for handling cancellation requests manually.
:param aw: Awaitable.
:param cancel_cb: Optional cancellation callback.
:return: Result of awaitable.
"""
task = asyncio.ensure_future(aw)
while True:
try:
return await asyncio.shield(task)
except asyncio.CancelledError:
if task.done():
raise
if cancel_cb is not None:
cancel_cb()

Categories

Resources