使用concurrent.future.ThreadPoolExecutor时注意事项

最近写了一个压力测试脚本，做的事情很简单就是不断发送REST API给服务器，发现跑一段时间以后机器就卡死了，检查发现原来内存被占用光了，我的二逼代码简化如下:

from concurrent import futures
import time


def work():
    time.sleep(2)


if __name__ == '__main__':
    pool = futures.ThreadPoolExecutor(128)
    while 1:
        pool.submit(work, )

粗看没啥问题，但是这个脚本如果一直跑下去就会导致内存一直增长，跟同事探讨了一下觉得原因可能是threa消耗的速度远远比不上新加thread（主要是未执行的thread）的数量， 导致保存这些thread对象的队列的size不断增大。

然后翻了一下源代码, 发现submit方法调用时会把thread的具体function放入一个叫_work_queue的属性里面，而这个_work_queue就是我们在线程里面经常用的queue.Queue类的一个实例:

class ThreadPoolExecutor(_base.Executor):
    # ...
    def __init__(self, max_workers=None, thread_name_prefix=''):
        # ...
        self._work_queue = queue.Queue()

    def submit(self, fn, *args, **kwargs):
        with self._shutdown_lock:
            if self._shutdown:
                raise RuntimeError('cannot schedule new futures after shutdown')

            f = _base.Future()
            w = _WorkItem(f, fn, args, kwargs)

            self._work_queue.put(w)
            self._adjust_thread_count()
            return f

ThreadPoolExecutor的实现上虽然限制了最大执行的thread数量，但是并没有一个参数能直接限制加入队列的task的数量，为了验证猜想，于是参考了queue.Queue的源代码，发现里面有个qsize()方法可以参考当前队列的size，修改代码如下：

from concurrent import futures
import time


def work():
    time.sleep(2)


if __name__ == '__main__':
    pool = futures.ThreadPoolExecutor(128)
    while 1:
        pool.submit(work, )
        # check pool queue size here
        print(pool._work_queue.qsize())

运行结果如下：

...
98815
98816
98817
98818
98819
...

再次查阅源代码，发现可以设置maxsize属性来限制队列的size，于是再修改代码如下：

from concurrent import futures
import time


def work():
    time.sleep(2)


if __name__ == '__main__':
    pool = futures.ThreadPoolExecutor(128)
    # limit work queue size to avoid memory size increasing issue.
    pool._work_queue.maxsize = 1000
    while 1:
        pool.submit(work, )
        print(pool._work_queue.qsize())

运行结果如下：

...
998
999
1000
997
998
999
1000
...

运行时可以明显看到在queue的size达到1000时程序会有等待的现象，其实是因为queue的put操作默认是blocking的，如果不设置超时会block知道有free的slot，官方文档的描述：

Queue.put(item[, block[, timeout]])

Put item into the queue. If optional args block is true and timeout is None (the default), block if necessary until a free slot is available.

总结来说，在使用futures.ThreadPoolExecutor时应该注意，如果使用while循环来提交任务，要注意限制一下ThreadPoolExecutor对象的_work_queue.maxsize，这样才能保证队列不至于一直堆积搞垮内存。