使用 run_in_executor 和 asyncio 時的超時處理

Question

我正在使用 asyncio 來運行一段這樣的阻塞代碼：

result = await loop.run_in_executor(None, long_running_function)

我的問題是：我可以為long_running_function的執行施加超時嗎？

基本上我不希望long_running_function持續超過 2 秒，並且我無法在其中進行適當的超時處理，因為該函數來自第三方庫。

Answer 1

關於取消長時間運行功能的警告：

盡管使用asyncio.wait_for調用將loop.run_in_executor返回的Future包裝為允許事件循環在某個x秒后停止等待long_running_function ，但它不一定會停止底層的long_running_function 。 這是concurrent.futures的缺點之一，據我所知，沒有簡單的方法可以取消concurrent.futures.Future 。

Answer 2

您可以使用asyncio.wait_for ：

future = loop.run_in_executor(None, long_running_function)
result = await asyncio.wait_for(future, timeout, loop=loop)

Answer 3

雖然沒有使用run_in_executor ，但我有一些關於“用超時處理異步包裝塊函數”的解決方法

import asyncio
import threading
import time
import ctypes


def terminate_thread(t: threading.Thread, exc_type=SystemExit):
    if not t.is_alive(): return
    try:
        tid = next(tid for tid, tobj in threading._active.items() if tobj is t)
    except StopIteration:
        raise ValueError("tid not found")
    if ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exc_type)) != 1:
        raise SystemError("PyThreadState_SetAsyncExc failed")


class AsyncResEvent(asyncio.Event):
    def __init__(self):
        super().__init__()
        self.res = None
        self.is_exc = False
        self._loop = asyncio.get_event_loop()

    def set(self, data=None) -> None:
        self.res = data
        self.is_exc = False
        self._loop.call_soon_threadsafe(super().set)

    def set_exception(self, exc) -> None:
        self.res = exc
        self.is_exc = True
        self._loop.call_soon_threadsafe(super().set)

    async def wait(self, timeout: float | None = None):
        await asyncio.wait_for(super().wait(), timeout)
        if self.is_exc:
            raise self.res
        else:
            return self.res


async def sub_thread_async(func, *args, _timeout: float | None = None, **kwargs):
    res = AsyncResEvent()

    def f():
        try:
            res.set(func(*args, **kwargs))
        except Exception as e:
            res.set_exception(e)
        except SystemExit:
            res.set_exception(TimeoutError)

    (t := threading.Thread(target=f)).start()
    try:
        return await res.wait(_timeout)
    except TimeoutError:
        raise TimeoutError
    finally:
        if not res.is_set():
            terminate_thread(t)


_lock = threading.Lock()


def test(n):
    _tid = threading.get_ident()
    for i in range(n):
        with _lock:
            print(f'print from thread {_tid} ({i})')
        time.sleep(1)
    return n


async def main():
    res_normal = await asyncio.gather(*(sub_thread_async(test, 5) for _ in range(2)))
    print(res_normal)  # [5,5]
    res_normal_2 = await asyncio.gather(*(sub_thread_async(test, 2, _timeout=3) for _ in range(2)))
    print(res_normal_2)  # [2,2]
    res_should_not_get = await asyncio.gather(*(sub_thread_async(test, 5, _timeout=3) for _ in range(2)))
    print(res_should_not_get)  # timeout error


if __name__ == '__main__':
    asyncio.new_event_loop().run_until_complete(main())