Python ThreadPoolExecutor 任务队列阻塞提交

分享一个多线程中阻塞提交任务队列的小例子，防止队列中任务过多导致内存占用过大，同时保证充分利用线程资源。

问题描述

  由于ThreadPoolExecutor默认采用的是无界队列，如果需要处理的任务量特别大，在生产速度大于消费速度时，可能会耗光系统资源，希望找到一种方式避免这种情况。

代码

  先不解释，直接上代码

# !/usr/bin/env python
# -*- coding: utf-8 -*-

from concurrent.futures import ThreadPoolExecutor, as_completed
import time
from random import randint


def test(num):
    for _ in range(2):
        time.sleep(randint(1, 5))   # 设置随机等待时间，
    return num


if __name__ == '__main__':
    max_pool = 5	# 线程池最大线程数
    executor = ThreadPoolExecutor(max_pool)
    all_task = []
    for i in range(50):		# 假设任务规模为50, 当然实际可能有上百上千万
        if len(all_task) < max_pool:
            all_task.append(executor.submit(test, i))
            print('创建：', i)
        else:
            for future in as_completed(all_task):
                all_task.remove(future)
                all_task.append(executor.submit(test, i))
                print('完成：', future.result(), '   创建：', i)
                break
    for future in as_completed(all_task):
        print('完成：', future.result())

解释

  使用 concurrent.futures.ThreadPoolExecutor 处理多线程任务的过程是首先将任务提交到一个任务队列中，若线程池中有线程执行完毕或者存在空闲线程时， 则从任务队列中拿取一个新的任务。若生产速度大于消费速度，会导致任务队列积压越来越多，占用过多内存。在消费速度无法提高时，可以通过降低生产速度即提交到任务队列的速度来解决该问题。因此本方法的核心就是监听线程池中有无任务执行完成，当线程池中有任务完成时再提交新的任务到队列中。

for future in as_completed(all_task):	# 返回已经执行结束的future实例
	all_task.remove(future)		# 从任务队列中移除
    all_task.append(executor.submit(test, i))	# 提交新的任务
    print('完成：', future.result(), '   创建：', i)
    break	# 这里只要监听到一个任务完成就跳出循环，否则上面提交的任务就重复了

  假设我们通过submit方法(返回一个Future实例)提交的任务都放在all_task列表中，使用as_completed来实现类似监听线程状态的功能。

concurrent.futures.as_completed(fs, timeout=None)
  Returns an iterator over the Future instances (possibly created by different Executor instances) given by fs that yields futures as they complete (finished or cancelled futures). Any futures given by fs that are duplicated will be returned once. Any futures that completed before as_completed() is called will be yielded first. The returned iterator raises a TimeoutError if __next__() is called and the result isn’t available after timeout seconds from the original call to as_completed(). timeout can be an int or float. If timeout is not specified or None, there is no limit to the wait time.

  简单来说，as_completed传入一个Futures类列表，会返回一个futures类的迭代器，通过迭代该迭代器可以获得已完成的futures。而且该方法一个比较好的特性是：as_completed会按线程完成的顺序返回，若没有线程完成，则会阻塞主进程，直到有线程结束执行。利用这个特性，我们就可以实现类似监听的功能，只要保证all_task中的线程和线程池中正在执行的线程保持同步就行，因此，当as_completed返回一个Future实例也就是线程执行结束时，我们从all_task中将该线程移除并提交一个新的线程，然后停止监听(break)进入下一个循环重新监听。
  通过这种方式，线程池中执行的线程始终保持在最大设定数，充分利用了线程，有剩余任务时不会有空闲线程存在，在任务队列等待的线程相当于没有，就防止了主线程一直往任务队列塞导致内存消耗过大。
  如果有问题恳请指正，有更好的实现方式欢迎评论讨论~

参考

• 对于该问题另外两种解决方式参考这里：
  https://blog.51cto.com/walkerqt/2317632
• 更多关于ThreadPoolExecutor、submit、as_completed以及Future实例的内容参考这里：
  https://docs.python.org/zh-cn/3/library/concurrent.futures.html