Perf: make do_cancel quicker. (#8846)

### What problem does this PR solve?

### Type of change

- [x] Performance Improvement
This commit is contained in:
Kevin Hu
2025-07-15 14:35:00 +08:00
committed by GitHub
parent 5fa6f2f151
commit 24c41d2a61
3 changed files with 36 additions and 29 deletions

View File

@ -606,6 +606,7 @@ TimeoutException = Union[Type[BaseException], BaseException]
OnTimeoutCallback = Union[Callable[..., Any], Coroutine[Any, Any, Any]]
def timeout(
seconds: float |int = None,
attempts: int = 2,
*,
exception: Optional[TimeoutException] = None,
on_timeout: Optional[OnTimeoutCallback] = None
@ -625,41 +626,46 @@ def timeout(
thread.daemon = True
thread.start()
try:
result = result_queue.get(timeout=seconds)
if isinstance(result, Exception):
raise result
return result
except queue.Empty:
raise TimeoutError(f"Function '{func.__name__}' timed out after {seconds} seconds")
for a in range(attempts):
try:
result = result_queue.get(timeout=seconds)
if isinstance(result, Exception):
raise result
return result
except queue.Empty:
pass
raise TimeoutError(f"Function '{func.__name__}' timed out after {seconds} seconds and {attempts} attempts.")
@wraps(func)
async def async_wrapper(*args, **kwargs) -> Any:
if seconds is None:
return await func(*args, **kwargs)
try:
with trio.fail_after(seconds):
return await func(*args, **kwargs)
except trio.TooSlowError:
if on_timeout is not None:
if callable(on_timeout):
result = on_timeout()
if isinstance(result, Coroutine):
return await result
return result
return on_timeout
for a in range(attempts):
try:
with trio.fail_after(seconds):
return await func(*args, **kwargs)
except trio.TooSlowError:
if a < attempts -1:
continue
if on_timeout is not None:
if callable(on_timeout):
result = on_timeout()
if isinstance(result, Coroutine):
return await result
return result
return on_timeout
if exception is None:
raise TimeoutError(f"Operation timed out after {seconds} seconds")
if exception is None:
raise TimeoutError(f"Operation timed out after {seconds} seconds and {attempts} attempts.")
if isinstance(exception, BaseException):
raise exception
if isinstance(exception, BaseException):
raise exception
if isinstance(exception, type) and issubclass(exception, BaseException):
raise exception(f"Operation timed out after {seconds} seconds")
if isinstance(exception, type) and issubclass(exception, BaseException):
raise exception(f"Operation timed out after {seconds} seconds and {attempts} attempts.")
raise RuntimeError("Invalid exception type provided")
raise RuntimeError("Invalid exception type provided")
if asyncio.iscoroutinefunction(func):
return async_wrapper