如何在找到目标哈希匹配后终止Multiprocessing Pool及优化建议

starmap是阻塞方法，必须等所有任务执行完毕才会返回，无法中途终止进程池。要实现找到匹配后立刻停止所有进程，需换用imap_unordered（按任务完成顺序返回结果），并在主进程中监听返回结果，一旦发现匹配就调用terminate()终止整个进程池。

修改后的代码示例 #

import multiprocessing
from Crypto.Hash import keccak
import click

def test(selector, bounds):
    # 预先拆分selector，避免循环内重复执行split
    selector_prefix = selector.split('(')[0]
    selector_suffix = selector.split('(')[1]
    target_hash = 'ae026468'
    
    for i in range(bounds[0], bounds[1]):
        current_str = f"{selector_prefix}_{i}({selector_suffix}"
        k = keccak.new(digest_bits=256)
        k.update(bytes(current_str, 'utf-8'))
        if k.hexdigest()[:8] == target_hash:
            return current_str  # 返回匹配结果，由主进程处理输出
    return None

def callTest():
    # iters为预先拆分好的任务列表，格式示例：[(selector, (0, 1000)), (selector, (1000, 2000)), ...]
    iters = [...]
    a_pool = multiprocessing.Pool()
    
    try:
        # 用imap_unordered迭代任务结果，拿到结果就立即处理
        for result in a_pool.imap_unordered(test, iters):
            if result is not None:
                click.echo(f'Found collision with {result}')
                a_pool.terminate()  # 直接终止所有子进程，停止剩余任务
                break
    finally:
        a_pool.close()
        a_pool.join()

关键说明 #

• 替换starmap为imap_unordered：它会在任务完成后立刻返回结果，无需等待所有任务结束，让主进程能第一时间触发终止操作。
• 子进程仅返回结果：避免子进程直接处理IO（如打印），减少跨进程输出混乱，同时让主进程明确判断是否找到目标。
• 用terminate()终止进程：close()仅禁止提交新任务，会等待现有任务完成；terminate()会直接终止所有子进程，实现立即停止的效果。

• 任务均匀拆分：确保每个子进程分配的任务区间大小尽可能一致，避免出现部分进程早早完成、部分进程长时间运行的失衡情况，提升整体效率。
• 预计算固定逻辑：像示例中预先拆分selector的前缀和后缀，避免在循环内重复执行相同的字符串拆分操作，减少不必要的计算开销。
• 子进程专注计算：子进程只负责哈希计算逻辑，把打印、文件写入等IO操作交给主进程处理，减少跨进程的资源竞争和IO冲突。
• 合理设置进程数：默认multiprocessing.Pool()会使用CPU核心数，CPU密集型任务保持默认即可；若涉及IO等待，可适当增加进程数（不超过核心数的2倍，避免上下文切换开销过大）。
• 优雅退出替代方案（可选）：如果不想粗暴终止进程，可使用multiprocessing.Value创建共享布尔标志，子进程循环时定期检查该标志，一旦主进程将标志设为True，子进程就主动退出循环。这种方式更优雅，但会增加少量同步开销。

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