优化带超时命令执行函数：修复高负载下的竞态条件问题

问题背景 #

我正在优化run_with_timeout函数，该函数接收命令（如ping google.com）与超时时间（如10秒），需实时捕获并展示进程控制台输出。正常场景下运行正常，参考Python核心开发者的并发实现方案后，采用线程+队列架构：print_manager线程专门负责打印输出，queue_line线程负责读取进程输出并写入队列。

现有实现代码：

def queue_line(data, print_queue):
    for line in iter(data.readline, ''):
        print_queue.put(line)
    data.close()

def print_manager(print_queue):
    while True:
        line = print_queue.get()
        print(line, end='')
        # Inform `print_queue` that the job is done
        print_queue.task_done()

def run_with_timeout_v2(command, timeout):
    # Create the process
    process = subprocess.Popen(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True,
        text=True,
        )
    
    process.timed_out = False
    process.output_buffer = ''

    print_queue = Queue()
    print_thread = threading.Thread(target=print_manager, args=(print_queue,))
    print_thread.daemon = True
    print_thread.start()
    del print_thread

    worker_thread = threading.Thread(target=queue_line, args=(process.stdout, print_queue))
    worker_thread.start()

    try:
        print_queue.put('In Try')
        process.wait(timeout=timeout)
    except subprocess.TimeoutExpired:
        print_queue.put('In Except')
        process.timed_out = True
        process.kill()
    finally:
        print_queue.put('In Finally')
        worker_thread.join()
        print_queue.join()

存在的问题 #

通过模糊测试模拟高负载环境时，竞态条件被放大，出现两个核心异常：

1. 输出顺序混乱：手动插入的控制文本（如'In Try'）和进程输出的顺序不可控
2. 超时后仍有进程输出：进程被kill后，queue_line线程可能继续读取系统缓冲区中的残留数据，导致超时提示后还出现进程输出

修复方案 #

1. 给工作线程添加终止信号 #

引入threading.Event作为停止标记，超时触发时立即通知queue_line线程停止读取，避免进程被kill后继续输出缓冲数据。

修改queue_line函数：

def queue_line(data, print_queue, stop_event):
    # 循环读取直到收到停止信号或无数据
    while not stop_event.is_set():
        line = data.readline()
        if not line:
            break
        # 标记消息类型为输出，便于后续区分
        print_queue.put(('output', line))
    data.close()

2. 区分队列消息类型，统一处理逻辑 #

将队列中的消息分为控制消息（如'In Try'这类提示文本）、输出消息（进程控制台输出）和退出消息（通知打印线程终止），让print_manager按类型处理，避免输出混排。

修改print_manager函数：

def print_manager(print_queue):
    while True:
        msg_type, content = print_queue.get()
        # 根据消息类型处理
        if msg_type == 'control':
            print(content)
        elif msg_type == 'output':
            print(content, end='')
        # 标记任务完成
        print_queue.task_done()
        # 收到退出信号时终止线程
        if msg_type == 'exit':
            break

3. 主函数完善资源清理与信号同步 #

在主函数中添加停止事件，超时触发时立即触发停止信号，同时完善线程的终止逻辑，确保所有资源被正确清理。

修改run_with_timeout_v2函数：

def run_with_timeout_v2(command, timeout):
    process = subprocess.Popen(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,  # Python3.7+中text是universal_newlines的别名，只用一个即可
        )
    
    process.timed_out = False
    print_queue = Queue()
    # 创建停止事件，用于通知工作线程终止
    stop_event = threading.Event()

    # 启动打印线程，移除daemon标记，手动控制终止
    print_thread = threading.Thread(target=print_manager, args=(print_queue,))
    print_thread.start()

    # 启动工作线程，传入停止事件
    worker_thread = threading.Thread(target=queue_line, args=(process.stdout, print_queue, stop_event))
    worker_thread.start()

    try:
        # 发送控制消息，标记类型为control
        print_queue.put(('control', 'In Try'))
        process.wait(timeout=timeout)
    except subprocess.TimeoutExpired:
        print_queue.put(('control', 'In Except'))
        process.timed_out = True
        process.kill()
        # 设置停止事件，立即通知工作线程停止读取
        stop_event.set()
    finally:
        print_queue.put(('control', 'In Finally'))
        # 等待工作线程完全终止
        worker_thread.join()
        # 发送退出信号给打印线程，等待其结束
        print_queue.put(('exit', ''))
        print_thread.join()
        # 等待队列中所有任务处理完成
        print_queue.join()

额外优化说明 #

• 移除del print_thread：手动删除线程对象无意义，Python的垃圾回收机制会自动处理
• 合并universal_newlines=True与text=True：两者是同一参数的不同命名（Python3.7+），保留一个即可
• 移除无用的process.output_buffer：原代码中未使用该变量，可删除

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