ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

深入理解Python @asynccontextmanager:异步上下文管理器原理与实战

深入理解Python @asynccontextmanager:异步上下文管理器原理与实战 我写异步代码也有几年了asynccontextmanager算是我用得最频繁的标准库工具之一。很多刚接触协程的人会写一堆重复的__aenter__/__aexit__样板代码或者干脆在async with里裸写try/finally其实标准库早就给了更优雅的解法。这篇文章就围绕这个装饰器把它的原理、用法、实战场景和容易踩的坑一次讲透。1. 为什么异步代码里需要“上下文管理器”这层抽象1.1 先搞清楚 asyncio 里“资源管理”为什么麻烦在同步编程里with open() as f这种写法大家都很熟进入时获取资源退出时自动释放。到了异步场景资源的获取和释放变成了耗时操作——建立数据库连接要await、关闭 HTTP 会话要await、释放信号量也要await。这时普通的with语句就不够了必须用async with配合实现了__aenter__和__aexit__的对象。问题在于手写一个异步上下文管理器类是这样的class AsyncSessionManager: async def __aenter__(self): self.session await create_session() return self.session async def __aexit__(self, exc_type, exc, tb): await self.session.close()如果每换一种资源就要写一个这样的类代码量会迅速膨胀而且大部分逻辑都是重复的。更麻烦的是__aexit__里还要处理exc_type、exc、traceback三个参数很多人根本用不到却也必须写齐全。1.2 生成器是上下文管理器的“语法糖底子”Python 里同步的contextmanager之所以好用是因为它利用了生成器的特性生成器函数在yield处暂停正好对应“资源使用中”的这个阶段。yield之前是进入逻辑yield之后是退出逻辑。from contextlib import contextmanager contextmanager def managed_resource(): print(enter) yield print(exit)asynccontextmanager就是同一思路的异步版本。它接收的是一个异步生成器函数也就是async def里包含yield的函数。调用时返回一个异步上下文管理器async with进入时执行yield之前的代码退出时执行yield之后的代码。2. asynccontextmanager 的核心用法与设计意图2.1 最小可用示例一个异步资源管理器先看一个最典型的例子管理一个 aiohttp 客户端会话import asyncio from contextlib import asynccontextmanager from aiohttp import ClientSession asynccontextmanager async def http_session_manager(): async with ClientSession() as session: yield session async def main(): async with http_session_manager() as session: async with session.get(https://example.com) as resp: print(resp.status) asyncio.run(main())关键点在于yield语句yield之前的部分负责创建资源yield把资源“吐”给调用方yield之后的部分负责清理资源。这个装饰器内部的实现帮你处理了异常传递的问题——如果async with代码块里抛了异常异常会在yield处重新抛出这样你有机会在退出逻辑里做相应处理。2.2 为什么说它比手动写类更“省心”手写__aenter__和__aexit__最大的痛点是样板代码太多而且容易出错。比如很多人会忘记在__aexit__里return布尔值来控制“是否吞掉异常”或者忘记把__aexit__本身定义成async def。用生成器版本后这些底层细节都交给标准库实现了。还有一点容易忽略asynccontextmanager装饰后的函数返回的是一个_AsyncGeneratorContextManager对象。这个对象同时实现了__aenter__和__aexit__。但注意它本身没有实现__await__所以你不能在await后面直接用。这是很多新手会犯的错后面我单独说。2.3 与 contextmanager 的差异对照同样管理一批资源同步版本和异步版本的写法和适用场景有明显区别对比项contextmanagerasynccontextmanager适用场景同步代码中的资源管理异步代码中涉及 await 的资源管理装饰的函数类型普通生成器函数异步生成器函数async def yield使用方式with 语句async with 语句内部资源创建/释放可以直接调用阻塞方法必须 await否则会卡住事件循环异常处理位置yield 处抛出yield 处抛出同样支持异常传递需要格外注意的是如果在一个async with块内部调用了阻塞的time.sleep(1)而不是await asyncio.sleep(1)整个事件循环都会被卡住。装饰器本身救不了这个问题它只是帮你组织资源生命周期不会替你优化异步调用方式。3. 实战场景事务、连接池与并发控制的优雅封装3.1 数据库事务的自动提交与回滚我在实际项目中用得最多的地方就是封装数据库事务。以前写数据库操作经常要手动维护“成功了 commit失败了 rollback”的逻辑一旦漏写数据就脏了。用asynccontextmanager可以把这套逻辑收敛在一个地方import asyncpg from contextlib import asynccontextmanager asynccontextmanager async def transaction(conn): try: await conn.execute(BEGIN) yield except Exception: await conn.execute(ROLLBACK) raise else: await conn.execute(COMMIT) async def main(): conn await asyncpg.connect(userpostgres, databasetest) try: async with transaction(conn): await conn.execute(INSERT INTO users(name) VALUES($1), 张三) # 如果这里抛异常事务自动回滚不会提交脏数据 finally: await conn.close()这段代码把事务的三个阶段拆得很清楚yield前开启事务yield中间执行业务代码yield后根据异常情况决定提交还是回滚。关键点在except里必须raise把原始异常抛出去不能静默吞掉否则调用方完全不知道事务已经失败了。我在之前一个订单服务里把库存扣减和订单创建放在同一个事务里用这个装饰器封装后业务逻辑里只需要写async with transaction(conn):可读性提升了一大截。而且测试的时候也特别好 mock因为事务边界非常明确。3.2 管理异步信号量控制并发上限异步代码里最常见的需求之一就是“限制同一时间的并发请求数”。很多人直接用asyncio.Semaphore但每次都要手动acquire和release一旦中间忘了释放并发控制会慢慢失效。用装饰器包一层信号量的生命周期自动管理from contextlib import asynccontextmanager import asyncio asynccontextmanager async def rate_limiter(limit): sem asyncio.Semaphore(limit) async with sem: yield async def worker(i, limiter): async with limiter: print(fworker {i} start) await asyncio.sleep(1) print(fworker {i} end) async def main(): limiter rate_limiter(3) await asyncio.gather(*(worker(i, limiter) for i in range(10))) asyncio.run(main())这个写法有个隐藏优势信号量被封装在上下文管理器的创建阶段每次async with limiter都会拿到同一个信号量实例。如果我不想共享信号量、想每批次独立限流只需在装饰器函数里yield asyncio.Semaphore(limit)即可。这种灵活度是手写类结构难以复制的。3.3 异步 HTTP 调用中的超时与会话复用用 aiohttp 时“每请求新建会话”是不推荐的做法因为会话内部维护连接池复用效率更高。但会话的关闭又不能忘记。这里同样适合用asynccontextmanager封装from contextlib import asynccontextmanager from aiohttp import ClientSession, ClientTimeout asynccontextmanager async def http_client(base_url, timeout10): timeout_obj ClientTimeout(totaltimeout) async with ClientSession(base_urlbase_url, timeouttimeout_obj) as session: yield session async def main(): async with http_client(https://api.example.com, timeout5) as session: async with session.get(/users) as resp: data await resp.json() print(data)另外一个小技巧是在async with外层叠asyncio.wait_for来控制整段代码的总耗时try: async with asyncio.timeout(3): async with http_client(https://api.example.com) as session: ... except TimeoutError: print(request timed out)这里用到了 Python 3.11 引入的asyncio.timeout上下文管理器能很方便地给整个资源使用阶段设置总期限。4. 容易踩的坑与排查技巧实录4.1 忘了写 async 关键字同步生成器被装饰后静默失败这是一个非常隐蔽的坑。asynccontextmanager装饰的是“异步生成器函数”也就是async def加yield。如果你粗心写成了普通函数加yieldPython 不一定立刻报错而是在async with的地方抛出奇怪的异常。# 错误的写法def 而不是 async def asynccontextmanager def bad_manager(): yield 123这段代码在定义阶段不会报错但一旦async with bad_manager() as x:解释器会抛出RuntimeError: generator didnt stop after athrow()排查起来很费劲。我的经验是写完装饰器先做一个最小调用测试不要直接嵌入业务代码。4.2 Yield 之前的代码抛异常时资源创建被中断asynccontextmanager的语义是yield之前的代码相当于__aenter__。如果这段代码抛异常async with会直接失败此时执行不到yield之后的清理逻辑。这本身合理但很多人误以为“只要用了这个装饰器退出逻辑一定会执行”。比如这样asynccontextmanager async def risky_manager(): resource await create_resource() # 这里抛了异常 yield resource await resource.close() # 这行不会执行如果create_resource内部已经申请了部分资源又在中途失败了它自己必须做好回滚。装饰器无法帮你自动清理“创建了一半的资源”。这一点和类的__aenter__是一样的逻辑。4.3 在装饰器内部再次使用 async with 导致上下文重叠用一个装饰器封装另一个装饰器的资源本身没问题但要注意别把“资源的创建”和“资源的使用”混淆。看下面这个例子asynccontextmanager async def outer_manager(): async with inner_manager() as res: yield res这个写法是可以正常工作的——外层进入时先创建内层资源yield给调用方外层退出时内层资源也一起释放。但如果你在yield之后再使用res就会触发RuntimeError: async generator ignored GeneratorExit之类的错误因为此时内层资源已经进入退出流程了。提示yield之后不要尝试再次使用你从内层拿到的资源对象。把yield理解成“资源使用区间的终点”终点之后一切都已经开始销毁。4.4 忘记异常会在 yield 处重新抛出当async with代码块内抛出异常时异常会传递到yield那一行并从那里继续执行yield之后的清理代码。所以如果你在退出逻辑里写了asynccontextmanager async def manager(): yield resource await cleanup() # 注意如果 async with 内抛异常这里执行时异常已经“激活”此时cleanup()里的异常会覆盖原始异常吗答案是会的。如果cleanup()自身抛了新异常原始异常会丢失。这会影响排查问题的效率。安全的做法是退出逻辑里尽量捕获所有可能抛出的异常asynccontextmanager async def manager(): yield resource try: await cleanup() except Exception: # 记录日志不要覆盖原始的业务异常 logger.exception(cleanup failed)4.5 并发场景下的非线程安全操作asyncio是单线程的但多个协程可以交错执行。如果多个协程同时使用同一个asynccontextmanager装饰器返回的上下文管理器并且内部共享了可变状态就需要小心了。这个装饰器并不会为每个async with创建独立的“上下文数据”如果你在装饰器函数外部定义了共享变量并发修改时依然存在竞态问题。最常见的例子是统计并发次数# 不推荐shared_counter 的修改不是原子的 counter 0 asynccontextmanager async def counting_manager(): global counter counter 1 try: yield finally: counter - 1在await期间事件循环可能切换协程counter 1可以被打断。虽然 GIL 让这个简单的加法不一定会出问题但更复杂的共享状态就难说了。我的建议是装饰器函数内部尽量只管理资源不维护业务状态。真要统计用itertools.count()配合asyncio.Lock或直接使用contextvars。4.6 Python 3.10 及以下版本的兼容性细节asynccontextmanager从 Python 3.7 加入标准库3.10 之前已经稳定使用。不过有些周边 API 需要注意。3.11 之前没有asyncio.timeout()我用的是# Python 3.10 及以下的超时写法 try: await asyncio.wait_for(coro, timeout3) except asyncio.TimeoutError: ...3.11 之后用asyncio.timeout(3)可以嵌套和管理多个await这对async with组合特别友好。如果你的项目还在用 3.8 或 3.9我不建议升级只为这个装饰器因为它本身在这些版本上稳定运行但asyncio.timeout这类配套工具确实要到 3.11 才香。5. 进阶玩法组合多个上下文管理器5.1 用 AsyncExitStack 管理不定数量的资源如果一个业务逻辑需要同时管理多个异步上下文资源而且数量是动态的嵌套async with会非常痛苦# 这种写法太笨了 async with session_manager() as session: async with transaction(conn) as tx: async with rate_limiter(5) as limiter: ...Python 提供了contextlib.AsyncExitStack来解决动态资源栈的问题。它支持enter_async_context()来动态追加资源并按后进先出的顺序释放from contextlib import AsyncExitStack async def main(): async with AsyncExitStack() as stack: session await stack.enter_async_context(http_client(https://api.example.com)) conn await stack.enter_async_context(transaction(await asyncpg.connect(...))) # 业务逻辑这在写测试夹具和中间件时特别有用比如给 FastAPI 的yield依赖注入多个异步资源AsyncExitStack是标准做法。5.2 自定义“可等待”的上下文管理器有时候我们希望async with之后还能手动await某个资源。这里有个思路是把资源本身设计成可等待对象然后对上下文管理器内部做一层包装asynccontextmanager async def background_task(coro_factory): task asyncio.create_task(coro_factory()) try: yield task finally: if not task.done(): task.cancel() async def main(): async with background_task(lambda: some_worker()) as t: # 可以做其他事情需要时等待任务结果 result await t这种模式在“启动后台任务 资源释放”的场景下非常顺手。不过要小心finally里task.cancel()如果任务已经正常完成cancel()是无害的直接忽略即可。5.3 与 FastAPI 的 yield 依赖结合FastAPI 的依赖注入支持yield它本质上也是上下文管理器。很多人不知道你可以在 FastAPI 依赖里用自己的asynccontextmanager资源然后直接抛给依赖系统管理from fastapi import FastAPI, Depends app FastAPI() asynccontextmanager async def get_db(): conn await asyncpg.connect(...) try: yield conn finally: await conn.close() app.get(/) async def read_root(dbDepends(get_db)): # db 就是 get_db 里 yield 出来的连接对象 ...这个写法的好处是get_db已经帮你完成了资源的创建和释放。FastAPI 会在请求结束时自动关闭连接你不需要在路由函数里写try/finally。6. 常见问题速查表问题现象可能原因解决方案RuntimeError: generator didnt stop after athrow()装饰器装饰的是普通函数而不是async def把函数改成async def确认里面有yieldasync with内部资源使用完但没释放yield之后的退出代码未写或写错了检查装饰器函数体确保yield之后有await resource.close()异常被吞掉排查不到原始错误退出逻辑里抛了新的异常覆盖了原始异常退出逻辑用try/except记录日志不覆盖原始异常并发场景下共享状态数据异常多个协程共享同一个可变对象缺少锁保护把状态放到上下文管理器内部或用asyncio.Lock保护想动态管理多个资源时嵌套过深使用了多个嵌套async with改用AsyncExitStack统一管理在装饰器函数中直接使用await外层变量异步生成器的惰性执行导致状态未初始化确保yield之前的代码能访问到正确作用域的变量必要时用参数传递7. 最后分享一个我自己常用的调试技巧调试asynccontextmanager时我会在装饰器函数里临时加一个打印确认进入和退出是否成对出现asynccontextmanager async def debug_manager(name): print(f[{name}] enter) try: yield finally: print(f[{name}] exit)如果只看到enter而没看到exit说明资源泄漏了优先检查是否有异常导致yield之后的代码没执行。如果看到两次enter但只有一次exit说明你可能不小心创建了同一个上下文管理器变量并在两个async with中复用而它内部又共享了同一个资源。这种情况我会把资源创建放到装饰器函数体内每次调用都返回新资源避免交叉污染。这个装饰器看起来简单但深入用下去会发现它的设计非常精妙。它不是银弹但绝大多数异步资源管理场景下它都是最优解。希望这篇文章能帮你把它用顺手。
返回列表