如何在 asyncio 引发 TimeoutError 之前由 asyncio 任务本身处理超时异常

Question

Due to one use case, One of my long-running functions executes multiple instructions. But I have to give a maximum time for its execution. If the function is not able to finish its execution within the allocated time, it should clean up the progress and return.

Let's have a look at a sample code below:

import asyncio

async def eternity():
    # Sleep for one hour
    try:
        await asyncio.sleep(3600)
        print('yay!, everything is done..')
    except Exception as e:
        print("I have to clean up lot of thing in case of Exception or not able to finish by the allocated time")


async def main():
    try:
        ref = await asyncio.wait_for(eternity(), timeout=5)
    except asyncio.exceptions.TimeoutError:
        print('timeout!')

asyncio.run(main())

The function eternity is the long-running function. The catch is that, in case of some exception or reaching the maximum allocated time, the function needs to clean up the mess it has made.

PS eternity is an independent function and only it can understand what to clean.

I am looking for a way to raise an exception inside my task just before the timeout, OR send some interrupt or terminate signal to the task and handle it.
Basically, I want to execute some peice of code in my task before asyncio raises the TimeoutError and take control.
Also, I am using Python 3.9.
Hope I was able to explain the problem.

Answer 1

What you need is async context manager:

import asyncio

class MyClass(object):

    async def eternity(self):
        # Sleep for one hour
        await asyncio.sleep(3600)
        print('yay!, everything is done..')

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        print("I have to clean up lot of thing in case of Exception or not able to finish by the allocated time")


async def main():
    try:
        async with MyClass() as my_class:
            ref = await asyncio.wait_for(my_class.eternity(), timeout=5)
    except asyncio.exceptions.TimeoutError:
        print('timeout!')


asyncio.run(main())

And this is the output:

I have to clean up lot of thing in case of Exception or not able to finish by the allocated time
timeout!

For more details take a look at here .