嵌套asyncio协程控制流回传及锁释放时机问题

问题背景 #

需要将一系列异步步骤封装为方法，每步生成子结果后触发后续异步步骤直至完成。当前核心问题：run_detection方法中，锁在detector.cooldown()完成后立即释放，而非等到detector.detect()执行完毕，同时需保证控制流正常回传给调用方。

现有定义代码 #

import asyncio
from dataclasses import dataclass

# 假设基础父类已定义
class Detector:
    pass
class AugmentationProcedure:
    pass
class TestCase:
    pass

# 假设类型标识类已定义
class DetectorIdentifier:
    pass
class InputFileIdentifier:
    pass
class AugmentationProcedureIdentifier:
    pass

# 假设输入输出数据类已定义
class TextDetectorInput:
    def __init__(self, language, text):
        self.language = language
        self.text = text
class TextDetectorResult:
    pass
class AugmentationProcedureTextInput:
    pass
class AugmentationProcedureTextOutput:
    def __init__(self, language, text):
        self.language = language
        self.text = text
class TextInputFile:
    def __init__(self, file_name, language):
        self.file_name = file_name
        self.language = language

class TextDetector(Detector):
    lock: asyncio.Lock = asyncio.Lock()

    async def detect(self, input: TextDetectorInput) -> TextDetectorResult: 
        pass

class TextAugmentationProcedure(AugmentationProcedure):
    async def augment(
        self, input: AugmentationProcedureTextInput
    ) -> AugmentationProcedureTextOutput:
            pass
    
TextDetectorCoolDownAndResultType = asyncio.Awaitable[asyncio.Awaitable[TextDetectorResult]]
RunTextDetectionType = asyncio.AsyncIterator[tuple[DetectorIdentifier, TextDetector, TextDetectorCoolDownAndResultType]]
GetFilesTextType = asyncio.AsyncIterator[tuple[InputFileIdentifier, TextInputFile, AugmentationProcedureTextOutput, RunTextDetectionType]]
CheckTextType = asyncio.AsyncIterator[tuple[AugmentationProcedureIdentifier, TextAugmentationProcedure, GetFilesTextType]]

@dataclass
class TextTestCase(TestCase):
    inputs: dict[InputFileIdentifier, TextInputFile]
    augmentation_procedures: dict[AugmentationProcedureIdentifier, TextAugmentationProcedure]
    detectors: dict[DetectorIdentifier, TextDetector]
    
    async def get_files(self, augmentation_procedure: TextAugmentationProcedure) -> GetFilesTextType:
        for input_identifier, text_file in self.inputs.items():
            # 假设read_text_file已实现
            augmentation_input = await read_text_file(
                text_file.file_name, text_file.language
            )

            augmentation_output = await augmentation_procedure.augment(augmentation_input)
        
            detector_input = TextDetectorInput(
                        augmentation_output.language, augmentation_output.text
            )

            yield (input_identifier, text_file, augmentation_output, self.run_detection(detector_input))

    async def run_detection(self, detector_input: TextDetectorInput) -> RunTextDetectionType:
        for detector_identifier, detector in self.detectors.items():
            async def cooldown_and_detect(detector: TextDetector, detector_input: TextDetectorInput):
                # 原代码错误：asyncio.Lock需用async with，且锁未覆盖detect流程
                with detector.lock:
                    # Cooldown
                    await detector.cooldown()
                    
                    return detector.detect(detector_input)
            
            yield (detector_identifier, detector, cooldown_and_detect(detector, detector_input))
    
    async def check(self) -> CheckTextType:
        for augmentation_procedure_identifier, augmentation_procedure in self.augmentation_procedures.items():
            yield (augmentation_procedure_identifier, augmentation_procedure, self.get_files(augmentation_procedure))

现有调用代码 #

test_case = TextTestCase()

async def execute_test():
    results = []
    async for (augmentation_procedure_identifier, augmentation_procedure, augmentation_results) in test_case.check():
        async for (file_identifier, image_file, augmentation_output_awaitable, detectors) in augmentation_results:
            
            results.append([image_file.file_name, str(augmentation_procedure), "Augmenting...", ""])
            # 假设live.update和update_table_with_results已实现
            live.update(update_table_with_results(results))

            async for (detector_identifier, detector, cooldown_awaitable) in detectors:
                try:
                    detection_awaitable = await cooldown_awaitable
                    detection_result = await detection_awaitable
                
                    # TODO: 处理检测结果
                except:
                    pass
                    # TODO: 错误处理

解决方案 #

核心修改点 #

1. 修复异步锁的使用：asyncio.Lock必须用async with而非普通with，否则无法正确获取异步锁。
2. 延长锁的生命周期：在锁的上下文内等待detect完成，而非返回detect的可等待对象，确保锁在检测全流程结束后才释放。
3. 简化调用逻辑：修改后无需两次await，直接获取最终检测结果。

修改后的run_detection方法 #

async def run_detection(self, detector_input: TextDetectorInput) -> RunTextDetectionType:
    for detector_identifier, detector in self.detectors.items():
        async def cooldown_and_detect(detector: TextDetector, detector_input: TextDetectorInput):
            # 使用async with获取异步锁，锁会在代码块结束时自动释放
            async with detector.lock:
                # 执行冷却步骤
                await detector.cooldown()
                # 等待检测完成并返回结果，锁保持至检测结束
                return await detector.detect(detector_input)
        
        yield (detector_identifier, detector, cooldown_and_detect(detector, detector_input))

修改后的调用代码 #

async def execute_test():
    results = []
    async for (augmentation_procedure_identifier, augmentation_procedure, augmentation_results) in test_case.check():
        async for (file_identifier, image_file, augmentation_output, detectors) in augmentation_results:
            
            results.append([image_file.file_name, str(augmentation_procedure), "Augmenting...", ""])
            live.update(update_table_with_results(results))

            async for (detector_identifier, detector, cooldown_and_detect_task) in detectors:
                try:
                    # 直接await得到最终检测结果，无需两次await
                    detection_result = await cooldown_and_detect_task
                
                    # 处理检测结果
                    # TODO: Do stuff here
                except Exception as e:
                    # 捕获并处理异常，避免裸except吞掉系统级错误
                    print(f"Detector {detector_identifier} failed: {str(e)}")
                    # TODO: 错误处理逻辑

额外修正 #

• 修正类型定义中的CheckImageType为CheckTextType，与文本测试用例的逻辑一致
• 修正变量名augmentation_output_awaitable为augmentation_output（原代码中已经是await后的结果）
• 替换裸except为except Exception，避免捕获系统退出等特殊异常

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