asyncio的Semaphore和Queue

asyncio.Semaphore & asyncio.Queue区别

组件	控制什么	保护什么	类比
Semaphore(n)	某个操作的并发数	下游资源（API/DB）	扫码枪数量
Queue(maxsize)	排队等待的任务数	本机内存	仓库容量
Worker 数量	消费者协程数	事件循环负载	工人数量

Semaphore

• 适用场景
  • 限制 API 并发调用
  • 限制数据库连接池
  • 限制文件 IO 并发
  • 多函数共享同一资源（跨模块传递同一个 sem）

sem = asyncio.Semaphore(3)  # 最多允许 3 个协程同时进入

async def fetch(sem, url):
    async with sem:                # 进入 -1，退出 +1，减到 0 阻塞
        return await http_get(url)

Queue

方法	作用	满/空时行为
await queue.put(item)	放入数据	满时阻塞
item = await queue.get()	取出数据	空时阻塞
queue.task_done()	标记一个任务完成	计数器 -1
await queue.join()	等待所有任务完成	计数器归零时返回

说明

put()       → 计数器 +1
get()       → 计数器不变
task_done() → 计数器 -1
join()      → 等待计数器 == 0

案列

import asyncio
import random

async def fetch_api(sem, item_id):
    async with sem:
        print(f"请求 item={item_id}")
        await asyncio.sleep(random.uniform(0.5, 1.5))
        result = f"result_{item_id}"
        print(f"完成 item={item_id}")
        return result

async def worker(queue, sem, results):
    while True:
        item = await queue.get()        # 队列空时阻塞
        if item is None:                # 哨兵，退出
            queue.task_done()
            break
        try:
            result = await fetch_api(sem, item)
            results.append(result)
        except Exception as e:
            print(f"失败 {item}: {e}")
        finally:
            queue.task_done()           # 标记完成，继续循环

async def main():
    item_ids = list(range(1, 21))
    MAX_WORKERS = 5
    MAX_CONCURRENT_API = 3
    QUEUE_SIZE = 6

    sem = asyncio.Semaphore(MAX_CONCURRENT_API)
    queue = asyncio.Queue(maxsize=QUEUE_SIZE)
    results = []

    # 创建固定数量 worker（创建即启动）
    workers = [
        asyncio.create_task(worker(queue, sem, results))
        for _ in range(MAX_WORKERS)
    ]

    # 放入任务
    for item_id in item_ids:
        await queue.put(item_id)       # 队列满时阻塞

    # 放入哨兵
    for _ in workers:
        await queue.put(None)

    # 等待完成
    await queue.join()                 # 等任务全部 task_done
    await asyncio.gather(*workers)     # 等协程全部退出

    print(f"完成 {len(results)} 个任务")

asyncio.run(main())

Worker的生命周期

worker 启动
  │
  └─→ while True:
        item = await queue.get()    ← 队列空则挂起等待
        if item is None:            ← 哨兵，退出
            task_done()
            break
        process(item)               ← 干活
        task_done()                 ← 标记完成，回到循环开头

join()和gather()关系

	queue.join()	asyncio.gather(*workers)
等待什么	计数器归零（所有任务 task_done	协程退出（break / return）
返回时机	最后一个 task_done() 调用后	每个协程彻底结束后
能捕获异常	无法捕获	可以捕获

# 时间线
queue.task_done()    ← join() 被唤醒
break                ← 协程还在运行
return               ← gather() 被唤醒

Worker返回结果的两种方式

共享列表（传参）

results = []

async def worker(queue, sem, results):
    while True:
        ...
        results.append(result)

await asyncio.gather(*workers)
print(results)

gather 收集返回值

async def worker(queue, sem):
    my_results = []
    while True:
        ...
        my_results.append(result)
    return my_results

# all_results = [[...], [...]] 
all_results = await asyncio.gather(*workers)
results = [x for sub in all_results for x in sub]  # 展平

技巧

Queue 管排队，Semaphore 管进门
Worker 是永动机，哨兵来了才下班
task_done 划清单，join 等清单清空
gather 等人走，异常一个不落

Semaphore与Queue的关系

案例

async def worker(queue, sem):
    while True:
        item = await queue.get()
        if item is None:
            break

        # ===== 不需要 Semaphore =====
        data = preprocess(item)      # 10 个 worker 可同时做

        # ===== 需要 Semaphore =====
        async with sem:              # 只有 3 个能同时进入
            result = await call_api(data)

        # ===== 不需要 Semaphore =====
        save_to_db(result)           # 10 个 worker 可同时做
        queue.task_done()

只有Semaphore，没有Queue

# 瞬间创建 10000 个协程，9990 个在内存中排队
tasks = [asyncio.create_task(process(sem, item)) for item in huge_list]
# → 内存爆炸 💥

只有Queue，没有Semaphore

# Queue 控制 worker 数 = 并发数，但粒度太粗
# 无法对 worker 内部某个操作单独限流

两者配合

Queue 保护内存（最多排队 40 个）
Worker 固定数量（5 个协程）
Semaphore 保护下游（最多 3 个并发 API）