如何在FastAPI的StreamingResponse生成器中关闭Socket？

你的代码里，gen_video用with socket.socket()管理连接，但同步Socket的recv是阻塞调用——当客户端断开连接时，FastAPI的StreamingResponse会停止迭代生成器，但生成器此时卡在recv上，没法退出循环，导致with块的自动关闭逻辑无法执行，Socket也就不能及时释放。

这里提供两种可行的解决办法：

方法一：给同步Socket设置超时，捕获生成器停止信号 #

给Socket设置超时，让recv定期跳出阻塞，同时捕获GeneratorExit异常（生成器被停止时抛出），主动触发Socket关闭：

import fastapi
import uvicorn
import socket

def gen_video(port):
    with socket.socket() as s:
        s.settimeout(1.0)  # 设置1秒超时，让recv定期返回
        s.connect(("127.0.0.1", port))
        try:
            while True:
                try:
                    data = s.recv(1024)
                    if not data:  # 数据源主动关闭连接
                        break
                    yield data
                except socket.timeout:
                    continue  # 超时后继续循环，检查是否需要退出
        except GeneratorExit:
            # 客户端断开，生成器被终止，主动关闭Socket（with块最终也会关，这里是提前释放）
            s.close()

app = fastapi.FastAPI()

@app.get("/stream")
async def stream():
    return fastapi.responses.StreamingResponse(gen_video(1234), media_type="video/ogg")

if __name__ == "__main__":
    uvicorn.run(app, port=8000)

方法二：改用异步Socket配合异步生成器（推荐） #

FastAPI是异步框架，用asyncio的异步Socket更适配，能更优雅地处理客户端断开的情况：

import fastapi
import uvicorn
import asyncio

async def gen_video(port):
    # 建立异步连接
    reader, writer = await asyncio.open_connection("127.0.0.1", port)
    try:
        while True:
            # 异步读取数据，不会阻塞事件循环
            data = await reader.read(1024)
            if not data:
                break
            yield data
    finally:
        # 无论是否异常，都关闭连接
        writer.close()
        await writer.wait_closed()

app = fastapi.FastAPI()

@app.get("/stream")
async def stream():
    return fastapi.responses.StreamingResponse(gen_video(1234), media_type="video/ogg")

if __name__ == "__main__":
    uvicorn.run(app, port=8000)

异步方式下，当客户端断开，StreamingResponse会停止迭代异步生成器，finally块会立即执行，关闭Socket连接，无需依赖超时机制。

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