最近写了一个压力测试脚本,做的事情很简单就是不断发送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,这样才能保证队列不至于一直堆积搞垮内存。