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')