1. 异步上下文管理器基础概念
在Python异步编程中，上下文管理器是一个非常重要的概念。传统的同步上下文管理器通过__enter__和__exit__方法实现，而异步上下文管理器则使用__aenter__和__aexit__方法。
python复制class AsyncContextManager:
async def __aenter__(self):
# 初始化资源
return resource
async def __aexit__(self, exc_type, exc_val, exc_tb):
# 清理资源
pass
异步上下文管理器的典型使用方式是async with语句：
python复制async with AsyncContextManager() as resource:
# 使用资源
await do_something(resource)
2. 原生asyncio.timeout的实现分析
Python 3.11引入了asyncio.timeout()，这是一个内置的异步上下文管理器，用于为代码块设置超时限制。它的基本用法如下：
python复制async def main():
try:
async with asyncio.timeout(1.0): # 1秒超时
await long_running_operation()
except TimeoutError:
print("操作超时")
asyncio.timeout()的工作原理：
- 创建一个Timeout对象，记录当前时间和超时时间
- 在进入上下文时启动计时
- 如果代码块执行时间超过限制，会取消当前任务
- 将
asyncio.CancelledError转换为TimeoutError
3. 实现支持超时的自定义上下文管理器
让我们实现一个基础版的超时上下文管理器：
python复制class TimeoutContext:
def __init__(self, timeout):
self.timeout = timeout
self._task = None
self._cancelled = False
async def __aenter__(self):
self._task = asyncio.current_task()
if self._task is None:
raise RuntimeError("没有正在运行的任务")
self._deadline = asyncio.get_running_loop().time() + self.timeout
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
if exc_type is asyncio.CancelledError and self._cancelled:
raise TimeoutError("操作超时") from None
return False
这个基础版本还缺少实际的超时检查功能，我们将在下一节完善它。
4. 添加超时检查机制
要实现真正的超时功能，我们需要在后台启动一个检查任务：
python复制class TimeoutContext:
def __init__(self, timeout):
if timeout is not None and timeout < 0:
raise ValueError("超时时间不能为负")
self.timeout = timeout
self._task = None
self._timeout_handle = None
async def __aenter__(self):
self._task = asyncio.current_task()
if self._task is None:
raise RuntimeError("没有正在运行的任务")
if self.timeout is not None:
loop = asyncio.get_running_loop()
when = loop.time() + self.timeout
self._timeout_handle = loop.call_at(when, self._cancel_task)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
if self._timeout_handle:
self._timeout_handle.cancel()
if exc_type is asyncio.CancelledError and hasattr(self, '_timed_out'):
raise TimeoutError(f"操作超时（{self.timeout}秒）") from None
return False
def _cancel_task(self):
if not self._task.done():
self._timed_out = True
self._task.cancel("操作超时")
这个实现已经可以处理基本的超时情况，但还缺少一些重要功能，比如动态调整超时时间。
5. 实现取消功能
为了让我们的上下文管理器更加强大，我们添加取消功能：
python复制class TimeoutContext:
def __init__(self, timeout):
self.timeout = timeout
self._task = None
self._timeout_handle = None
self._cancel_requested = False
async def __aenter__(self):
self._task = asyncio.current_task()
if self._task is None:
raise RuntimeError("没有正在运行的任务")
self._schedule_timeout()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
self._cancel_timeout()
if exc_type is asyncio.CancelledError:
if hasattr(self, '_timed_out'):
raise TimeoutError(f"操作超时（{self.timeout}秒）") from None
if self._cancel_requested:
raise CancelledError("操作被取消") from None
return False
def _schedule_timeout(self):
if self.timeout is not None:
loop = asyncio.get_running_loop()
when = loop.time() + self.timeout
self._timeout_handle = loop.call_at(when, self._timeout_task)
def _cancel_timeout(self):
if self._timeout_handle:
self._timeout_handle.cancel()
self._timeout_handle = None
def _timeout_task(self):
if not self._task.done():
self._timed_out = True
self._task.cancel()
def cancel(self):
"""主动取消操作"""
if not self._task.done():
self._cancel_requested = True
self._task.cancel()
现在我们可以这样使用：
python复制async def operation():
async with TimeoutContext(2.0) as timeout:
try:
await long_running_task()
except TimeoutError:
print("任务超时")
# 也可以主动取消
if some_condition:
timeout.cancel()
6. 动态调整超时时间
为了更灵活地控制超时，我们可以添加动态调整功能：
python复制class TimeoutContext:
# ... 之前的代码 ...
def reschedule(self, new_timeout):
"""重新设置超时时间"""
if new_timeout is not None and new_timeout < 0:
raise ValueError("超时时间不能为负")
self._cancel_timeout()
self.timeout = new_timeout
self._schedule_timeout()
def remaining(self):
"""返回剩余时间"""
if self.timeout is None or not hasattr(self, '_deadline'):
return None
return max(0, self._deadline - asyncio.get_running_loop().time())
使用示例：
python复制async def adaptive_operation():
async with TimeoutContext(10.0) as timeout:
# 第一阶段操作
await phase_one()
# 根据第一阶段结果调整超时
if needs_more_time:
timeout.reschedule(20.0)
# 第二阶段操作
await phase_two()
print(f"剩余时间: {timeout.remaining():.1f}秒")
7. 完整实现与测试
下面是完整的实现代码，包含所有功能和错误处理：
python复制import asyncio
from typing import Optional, Union
class TimeoutError(Exception):
"""操作超时异常"""
pass
class TimeoutContext:
def __init__(self, timeout: Optional[Union[float, int]]):
if timeout is not None and timeout < 0:
raise ValueError("超时时间不能为负")
self._timeout = timeout
self._task = None
self._timeout_handle = None
self._cancel_requested = False
self._timed_out = False
self._deadline = None
async def __aenter__(self):
self._task = asyncio.current_task()
if self._task is None:
raise RuntimeError("没有正在运行的任务")
if self._timeout is not None:
loop = asyncio.get_running_loop()
self._deadline = loop.time() + self._timeout
self._timeout_handle = loop.call_at(self._deadline, self._timeout_task)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
self._cancel_timeout()
if exc_type is asyncio.CancelledError:
if self._timed_out:
raise TimeoutError(f"操作超时（{self._timeout}秒）") from None
if self._cancel_requested:
raise asyncio.CancelledError("操作被取消") from None
return False
def _schedule_timeout(self):
if self._timeout is not None and not self._timeout_handle:
loop = asyncio.get_running_loop()
self._deadline = loop.time() + self._timeout
self._timeout_handle = loop.call_at(self._deadline, self._timeout_task)
def _cancel_timeout(self):
if self._timeout_handle:
self._timeout_handle.cancel()
self._timeout_handle = None
def _timeout_task(self):
if not self._task.done():
self._timed_out = True
self._task.cancel()
def cancel(self):
"""主动取消操作"""
if not self._task.done():
self._cancel_requested = True
self._task.cancel()
def reschedule(self, new_timeout: Optional[Union[float, int]]):
"""重新设置超时时间"""
if new_timeout is not None and new_timeout < 0:
raise ValueError("超时时间不能为负")
self._cancel_timeout()
self._timeout = new_timeout
self._timed_out = False
self._schedule_timeout()
def remaining(self) -> Optional[float]:
"""返回剩余时间，如果没有设置超时则返回None"""
if self._timeout is None or self._deadline is None:
return None
remaining = self._deadline - asyncio.get_running_loop().time()
return max(0, remaining)
@property
def timeout(self) -> Optional[float]:
"""获取当前超时设置"""
return self._timeout
@property
def expired(self) -> bool:
"""检查是否已经超时"""
return self._timed_out
测试用例：
python复制async def test_timeout():
# 测试正常完成
async with TimeoutContext(1.0) as timeout:
await asyncio.sleep(0.5)
assert not timeout.expired
# 测试超时
try:
async with TimeoutContext(0.5) as timeout:
await asyncio.sleep(1.0)
except TimeoutError:
print("超时测试通过")
else:
assert False, "应该触发超时"
# 测试动态调整
async with TimeoutContext(1.0) as timeout:
await asyncio.sleep(0.5)
timeout.reschedule(1.0) # 总共1.5秒
await asyncio.sleep(0.9) # 总共1.4秒
assert not timeout.expired
# 测试取消
try:
async with TimeoutContext(10.0) as timeout:
await asyncio.sleep(0.1)
timeout.cancel()
await asyncio.sleep(1.0)
except asyncio.CancelledError:
print("取消测试通过")
else:
assert False, "应该触发取消"
asyncio.run(test_timeout())
8. 实际应用场景与最佳实践
8.1 网络请求超时控制
python复制async def fetch_with_timeout(url, timeout=5.0):
async with TimeoutContext(timeout):
try:
async with aiohttp.ClientSession() as session:
async with session.get(url) as response:
return await response.text()
except TimeoutError:
print(f"获取 {url} 超时")
return None
8.2 数据库操作超时
python复制async def query_database(query, params, timeout=3.0):
async with TimeoutContext(timeout):
try:
conn = await asyncpg.connect(DATABASE_URL)
try:
return await conn.execute(query, *params)
finally:
await conn.close()
except TimeoutError:
print("数据库查询超时")
raise
8.3 组合多个操作
python复制async def complex_operation():
# 第一阶段：快速操作
async with TimeoutContext(1.0) as timeout:
await fast_operation()
# 第二阶段：可能需要更多时间
if needs_more_time():
timeout.reschedule(5.0)
await slow_operation()
# 检查剩余时间
if timeout.remaining() < 2.0:
print("警告：剩余时间不足")
最佳实践
-
合理设置超时时间：根据操作类型设置合理的超时时间，网络请求通常2-5秒，数据库操作1-3秒，CPU密集型操作可能需要更长。
-
资源清理：确保在超时或取消时正确释放资源，使用
try/finally块。 -
错误处理：区分超时错误和其他类型的错误，提供有意义的错误信息。
-
日志记录：记录超时事件，帮助调试和性能分析。
-
避免嵌套：尽量避免嵌套多个超时上下文管理器，这可能导致复杂的行为。
-
测试边界条件：特别测试刚好在超时边界完成的操作。