在asyncio中实现上下文切换机制：解决单会话多账号冲突问题

需求 #

实现上下文切换机制，确保共享资源（连接网站的单会话对象）同一时间仅被一个账号使用，且账号切换（连接操作）成本较高。

初始实现机制 #

为了在耗时后台操作时释放上下文（切换账号），设计了基于计数器的方案：

counter = 0  # 对应当前账号的ContextVar
def increase_counter():
  global counter
  counter += 1

def decrease_counter():
  global counter
  counter -= 1

async def run_operation():
  while True:
    operation = operation_queue.get()
    increase_counter()

    task = asyncio.create_task(operation())
    task.add_done_callback(lambda fut: decrease_counter())

    await wait_for_free() # 等待counter == 0
    switch_context()  # 切换到其他账号登录

async def operation():
  # 执行前置操作
  decrease_counter()
  sleep(60) # 耗时后台操作
  await wait_for_context()  # 等待上下文返回且counter为1
  # 继续后续操作

该机制在单操作链下运行正常，但多操作并行时失效：

多操作并行的问题代码 #

async def operation():
  increase_counter()
  task1 = asyncio.create_task(sub_op())
  task1.add_done_callback(decrease_counter)
  increase_counter()
  task2 = asyncio.create_task(sub_op())
  task2.add_done_callback(decrease_counter)

  decrease_counter()
  await asyncio.gather(task1, task2)
  await wait_for_context()

async def sub_op():
  decrease_counter()
  sleep(60)
  await wait_for_context()

计数器变化流程 #

run_op (+1 = 1)
task1_creation (+1 = 2)
task2_creation (+1 = 3)
gather_release (-1 = 2)
task1_sleep (-1 = 1)
task2_sleep (-1 = 0)  # 上下文被释放
task1_resume (+1 = 1)
task2_resume (+1 = 2)
task1_done_callback (-1 = 1)
task2_done_callback (-1 = 0)  # 无意义的上下文释放
gather_resume (+1 = 1)

疑问 #

能否避免这种无意义的上下文释放？若无法避免，是否存在替代计数器的机制解决该问题？

问题根源 #

当前的计数器仅统计"活跃引用数"，但无法区分上下文是否还有未完成的待恢复任务。当子任务全部进入睡眠时计数器归0，触发上下文切换；但这些子任务后续还需要恢复使用原上下文，它们完成后的回调又会让计数器归0，导致不必要的切换。

方案1：令牌锁机制（最简洁可靠） #

把上下文看作一个可复用的令牌，用锁来控制整个操作链的上下文占用周期：

• 每个账号对应一把锁，初始为可用状态
• 启动操作链时先获取锁，锁的生命周期覆盖整个操作链所有需要上下文的环节
• 只有当操作链彻底完成（包括子任务恢复后的所有操作），锁才会释放，此时才允许切换上下文

简化示例代码：

from contextvars import ContextVar
import asyncio

# 存储每个账号的锁：账号ID -> asyncio.Lock
account_locks = {}
current_account = ContextVar("current_account")

async def run_operation():
    while True:
        operation, account_id = operation_queue.get()
        # 初始化账号锁（如果不存在）
        if account_id not in account_locks:
            account_locks[account_id] = asyncio.Lock()
        
        # 持有锁期间，上下文完全归当前操作链使用
        async with account_locks[account_id]:
            current_account.set(account_id)
            await operation()
        # 锁自动释放，此时可安全切换到其他账号

async def operation():
    # 执行需要上下文的前置操作
    # 启动子任务时无需额外获取锁（父操作链已持有）
    task1 = asyncio.create_task(sub_op())
    task2 = asyncio.create_task(sub_op())
    await asyncio.gather(task1, task2)
    # 后续需要上下文的操作直接执行（锁仍被持有）

async def sub_op():
    # 耗时后台操作，无需上下文，不影响锁状态
    await asyncio.sleep(60)
    # 恢复后需要上下文，直接使用即可（父操作链仍持有锁）
    # 执行后续操作

这个方案从根源避免了无意义的上下文释放，因为锁的生命周期完全匹配操作链的上下文需求周期。

方案2：任务组跟踪机制 #

为每个上下文维护一个任务集合，只有当所有关联任务彻底完成（包括恢复后的后续操作），才允许释放上下文：

• 用ContextVar存储当前上下文的任务跟踪集合
• 每个需要使用上下文的任务，在启动时注册到集合，彻底完成时注销
• wait_for_free()改为等待集合中无任何活跃任务

这种方式比计数器精准，但实现复杂度稍高，需要跟踪每个任务的完整生命周期。

方案3：借用-归还标记模型 #

为每个操作链维护"借用计数"和"待恢复标记"：

• 任务需要使用上下文时调用borrow_context()，增加借用计数
• 任务进入耗时后台操作时调用return_context()，减少借用计数但标记"后续还会借用"
• 只有当借用计数为0且无待恢复标记时，才允许释放上下文

该方式精准度高，但需要额外维护任务状态，实现成本较高。

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