python实现自定义线程池

线程池ThreadPool对象的几个关键方法:

• get_idle_num(): 获得当前空闲线程的数量
• submit(task:callable): 把一个任务(实际上就是一个函数)提交到线程池中执行. 如果没有空闲线程则阻塞.
• wait_idle(): 阻塞, 直到有空闲线程
• stop(): 停止线程池中的所有线程. (注意: 非强制停止, 需要现有任务执行完毕.)

源码如下:

import queue
import threading
import time

class ThreadPool:
    '''线程池'''
    def __init__(self) -> None:
        self.threads_num = 3
        '''线程的数量'''
        self.idle_count = self.threads_num
        '''空闲线程的数量'''
        self._lock = threading.Lock()
        '''线程锁'''
        self._semaphore = threading.Semaphore(self.threads_num)
        '''信号量, 用于控制线程是否空闲'''
        self.task_queue = queue.Queue()
        '''任务队列'''
        self.threads:list[threading.Thread] = []
        '''线程列表'''
        self._start()
    def get_idle_num(self)->int:
        '''获取空闲线程的数量'''
        return self.idle_count
    def join(self):
        '''等待所有线程执行完毕'''
        for thread in self.threads:thread.join()
    def _start(self):
        '''启动线程池'''
        for _ in range(self.threads_num):
            thread = threading.Thread(target=self._worker, daemon=True)  # 创建线程
            thread.start()
            self.threads.append(thread)
    def stop(self):
        '''停止线程池'''
        # 通过向任务队列中添加None，来终止线程
        for _ in range(self.threads_num): self.task_queue.put(None)
    def submit(self, task:callable):
        '''提交任务到线程池. 如果没有空闲线程, 则会阻塞'''
        self._semaphore.acquire()  # 等待有线程空闲
        self.task_queue.put(task)
    def wait_idle(self):
        '''等待有线程空闲'''
        # 通过信号量控制，当空闲线程数量为0时，信号量会阻塞
        with self._semaphore: pass
    def _worker(self):
        '''线程工作函数'''
        while True:
            task = self.task_queue.get()  # 从队列中取出一个任务执行
            if task is None: 
                self._semaphore.release()  # 释放信号量，通知其他线程
                break  # 用于终止线程的方式
            # 空闲的线程数量减1
            with self._lock: self.idle_count -= 1
            task()  # 执行任务
            # 任务执行完毕，空闲的线程数量加1
            with self._lock: self.idle_count += 1
            #self.task_queue.task_done()  # 通知任务队列，任务完成
            self._semaphore.release()  # 释放信号量，通知其他线程
            
if __name__ == '__main__':
    '''测试线程池'''
    pool = ThreadPool()
    pool.wait_idle()
    print('ok')
    pool.submit(lambda: time.sleep(10))
    pool.submit(lambda: time.sleep(15))
    pool.submit(lambda: time.sleep(20))
    #time.sleep(0.001)
    pool.wait_idle()
    print('ok1')
    print(pool.get_idle_num())
    pool.stop()
    pool.join()
    print('ok2')