重构中间件拦截器抽取会话代码后出现IllegalStateChangeError异常排查

核心问题分析 #

出现IllegalStateChangeError（提示close()无法调用）的本质是数据库会话（AsyncSession）的生命周期被错误管理，具体触发点集中在以下几点：

1. 异步生成器未正确收尾 #

原逻辑中get_async_session是通过async with创建会话并yield的异步生成器：

async def get_async_session():
    async with AsyncSession(engine) as session:
        yield session

当你用await anext(gen)获取会话后，生成器处于挂起状态。如果抽取后的函数没有保留生成器引用，也没让生成器完成后续上下文收尾，生成器被垃圾回收时async with会自动执行__aexit__调用会话close()，但此时会话可能已被其他逻辑关闭或状态不可用，从而抛出异常。

2. 依赖覆盖判断逻辑误判 #

如果判断get_async_session是否被依赖覆盖的逻辑有漏洞，会导致同一请求中同时触发两种会话获取路径：既调用了覆盖后的依赖方法，又通过生成器获取会话，造成会话被重复关闭或状态冲突。

3. 中间件会话使用时机冲突 #

中间件拦截器的执行时机早于FastAPI的依赖注入生命周期，若抽取后的函数在中间件中直接操作会话，会干扰依赖系统对会话的管理（比如依赖系统后续尝试关闭已被中间件影响的会话）。

修复步骤与代码示例 #

步骤1：保证异步生成器完整生命周期 #

不要获取会话后就丢弃生成器，需保留引用直到请求结束，让生成器完成上下文收尾：

async def get_session_from_request_dependencies_or_generator(request):
    # 判断是否存在依赖覆盖（示例逻辑，根据实际场景调整）
    if hasattr(request.state, "overridden_session"):
        return request.state.overridden_session
    else:
        gen = get_async_session()
        session = await anext(gen)
        # 将生成器存入request状态，后续收尾用
        request.state.session_generator = gen
        return session

# 在请求结束钩子中处理生成器收尾
@app.middleware("http")
async def session_middleware(request, call_next):
    try:
        session = await get_session_from_request_dependencies_or_generator(request)
        request.state.db_session = session
        response = await call_next(request)
        return response
    finally:
        # 执行生成器剩余逻辑，让async with完成会话关闭
        if hasattr(request.state, "session_generator"):
            try:
                await anext(request.state.session_generator)
            except StopAsyncIteration:
                pass

步骤2：修正依赖覆盖判断逻辑 #

确保判断逻辑准确，避免重复获取会话。以FastAPI的dependency_overrides为例：

async def get_session_from_request_dependencies_or_generator(app):
    if get_async_session in app.dependency_overrides:
        # 调用覆盖后的依赖方法
        return await app.dependency_overrides[get_async_session]()
    else:
        gen = get_async_session()
        session = await anext(gen)
        # 返回会话和生成器，方便后续收尾
        return session, gen

步骤3：避免中间件直接接管会话生命周期 #

优先让FastAPI的依赖注入系统管理会话，中间件仅通过依赖获取会话：

# 定义标准会话依赖
async def get_db_session():
    async with AsyncSession(engine) as session:
        yield session

# 中间件中通过依赖系统获取会话
@app.middleware("http")
async def session_middleware(request, call_next):
    # 优先使用覆盖后的依赖，否则用默认实现
    dep_func = request.app.dependency_overrides.get(get_db_session, get_db_session)
    # 处理生成器类型的依赖
    if inspect.isasyncgenfunction(dep_func):
        gen = dep_func()
        session = await anext(gen)
        request.state.session_generator = gen
    else:
        session = await dep_func()
    
    request.state.db_session = session
    response = await call_next(request)
    
    # 收尾生成器
    if hasattr(request.state, "session_generator"):
        try:
            await anext(request.state.session_generator)
        except StopAsyncIteration:
            pass
    
    return response

关键注意事项 #

• 异步生成器的async with上下文会在生成器耗尽时自动关闭会话，必须保证生成器能完整执行yield后的逻辑。
• 禁止在多个地方重复调用AsyncSession.close()，否则会触发状态错误。
• 中间件生命周期与依赖注入不同步，直接操作生成器极易导致会话状态混乱。

内容的提问来源于stack exchange，提问作者Myrto Pirli