as_completed和wait源碼分析
前言
在ThreadPoolExecutor引導的多線程開發中，有as_completed()和wait()兩個輔助函數。下面結合源碼分析它們各自作用。因後面多次提到事件鎖，也許，你需要對它事先了解Python同步機制（semaphore,event,queue）。
（以下基於Python3.7）
as_complete
def greet():
print("hello world")
if __name__ == "__main__":
executor = ThreadPoolExecutor()
task = executor.submit(greet)
print(type(task)) # 輸出：<class 'concurrent.futures._base.Future'>
results = as_completed([task])
print(type(results)) # 輸出：<class 'generator'>
for result in results:
print(type(result)) # 輸出：<class 'concurrent.futures._base.Future'>
最初知道有as_completed()這個函數我是很不解的，因爲如上面代碼所示，submit()提交任務後拿到的返回值是Future對象，經過as_completed包裝後成了生成器，但是打開生成器一看，結果還是Future對象！這破玩意就這點用處？顯示不是。官方文檔對其做的說明，揭露了它的本質：
Returns an iterator over the Future instances (possibly created by different Executor instances) given by fs that yields futures as they complete (finished or were cancelled). 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 concurrent.futures.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.
在這裏邊有兩句話比較重要：
- 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.
先看第一句：當future重複時，只返回一次。示例如下：
def greet(word):
return word
if __name__ == "__main__"
executor = ThreadPoolExecutor()
tasks = [executor.submit(greet, word) for word in ["hello", "world"]]
tasksDouble = tasks * 2 # futures x 2
for item in as_completed(tasksDouble):
print(item.result())
# 輸出：
hello
world
可以看出，我們對tasks做了乘2操作，但是經手as_completed()之後並沒有重複打印hello或者word。說明在as_completed中有去重操作。Python內部僅僅做了一個很簡單的處理——集合真是強大的去重助理。
# as_completed源碼
def as_completed(fs, timeout=None):
...
fs = set(fs) # 去重操作
...
再看第二句：as_completed會先把該函數調用之前完成的furture依次yeild出去。也就是說，返回結果不會順序了。似乎莫名奇妙，但我們來看看用as_completed和不用的區別。
# 不使用as_completed
def print_num(order):
"""
i 表示線程啓動次序
通過隨機獲取num, 使得線程與線程之間的結束時間可能不同
"""
num = random.randrange(10)
time.sleep(num)
ordict = collections.OrderedDict()
ordict["oroder"] = order
ordict["value"] = num
return ordict # 最後打印調用次序以及線程運行的近似時間
if __name__ == "__main__":
executor = ThreadPoolExecutor()
alltasks = [executor.submit(print_num, i) for i in range(10)]
for task in alltasks:
print(task.result())
# 輸出：
OrderedDict([('oroder', 0), ('value', 5)])
OrderedDict([('oroder', 1), ('value', 3)])
OrderedDict([('oroder', 2), ('value', 1)])
OrderedDict([('oroder', 3), ('value', 1)])
OrderedDict([('oroder', 4), ('value', 1)])
OrderedDict([('oroder', 5), ('value', 7)])
OrderedDict([('oroder', 6), ('value', 0)])
OrderedDict([('oroder', 7), ('value', 0)])
OrderedDict([('oroder', 8), ('value', 6)])
OrderedDict([('oroder', 9), ('value', 1)])
可以看到輸出結果按照了線程啓動的順序，那麼使用as_completed又會怎樣呢？
def print_num(order):
num = random.randrange(10)
time.sleep(num)
ordict = collections.OrderedDict()
ordict["oroder"] = order
ordict["value"] = num
return ordict
if __name__ == "__main__":
executor = ThreadPoolExecutor()
alltasks = [executor.submit(print_num, i) for i in range(10)]
for task in as_completed(alltasks):
print(task.result())
# 輸出：
OrderedDict([('oroder', 9), ('value', 1)])
OrderedDict([('oroder', 1), ('value', 3)])
OrderedDict([('oroder', 2), ('value', 6)])
OrderedDict([('oroder', 4), ('value', 6)])
OrderedDict([('oroder', 0), ('value', 7)])
OrderedDict([('oroder', 3), ('value', 7)])
OrderedDict([('oroder', 5), ('value', 7)])
OrderedDict([('oroder', 6), ('value', 9)])
OrderedDict([('oroder', 7), ('value', 9)])
OrderedDict([('oroder', 8), ('value', 9)])
顯然，先執行完的先輸出。
實現邏輯從as_completed函數的try語句開始，
def as_completed(fs, timeout=None):
...
try:
# 先yield出所有已經結束的future，
# 包括了兩個狀態：CANCELLED_AND_NOTIFIED 和 FINISHED
yield from _yield_finished_futures(finished, waiter,
ref_collect=(fs,))
# pending裏面存放的是還沒有結束的future
while pending:
... # 處理超時邏輯
waiter.event.wait(wait_timeout) # 使用event鎖阻塞程序
# 直到finished_futures中有了內容，程序向下執行
with waiter.lock:
finished = waiter.finished_futures
waiter.finished_futures = []
waiter.event.clear() # 重置event鎖的狀態，爲下次阻塞程序做準備
# reverse to keep finishing order
finished.reverse()
# yield出去已經結束的future
yield from _yield_finished_futures(finished, waiter,
ref_collect=(fs, pending))
finally:
...
單看as_completed中的代碼，並不能直接獲悉上邊註釋裏的內容，因爲涉及到太多函數之間的調用。避免篇幅冗長，撿重要說。
從waiter = _create_and_install_waiters(fs, _AS_COMPLETED)這一句拿到的waiter是**_AsCompletedWaiter**的對象，這個對象的add_result()存在釋放鎖event鎖的行爲，使得程序可以通過event.wait向下執行。
# 從as_completed中的_create_and_install_waiters跳進去
class _AsCompletedWaiter(_Waiter):
"""Used by as_completed()."""
...
def add_result(self, future):
with self.lock:
super(_AsCompletedWaiter, self).add_result(future)
self.event.set() # 釋放event鎖
...
_yield_finished_futures()方法會返回已經完成的future，同時還會把這個future從pending中移除，達到結束while循環的目的。
def _yield_finished_futures(fs, waiter, ref_collect):
# fs表示已經完成的future
# ref_collect爲元組類型，
# 第一個元素爲全部future，
# 第二個元素pendding，也就是as_completed中待解決的future
while fs:
f = fs[-1]
for futures_set in ref_collect:
futures_set.remove(f) # 移除已經完成的future
with f._condition:
f._waiters.remove(waiter)
del f
# Careful not to keep a reference to the popped value
yield fs.pop()
wait
def wait(fs, timeout=None, return_when=ALL_COMPLETED):
...
wait()函數如其名，只爲等待。它返回的對象是可迭代類型（你可以把它當作元組看待，它是collections.namedtuple的實例），裏邊有兩個集合類型的元素，第一個集合裝着已經完成的future，第二個集合裝着未完成的future。
根據return_when接收到的參數不同，決定wait()什麼時候返回。源碼中的docstring對其做了這樣說明：
- FIRST_COMPLETED - Return when any future finishes or is cancelled.
- FIRST_EXCEPTION - Return when any future finishes by raising an exception. If no future raises an exception then it is equivalent to ALL_COMPLETED.
- ALL_COMPLETED - Return when all futures finish or are cancelled.
也就是說，return_when有三種狀態。當值爲FIRST_COMPLETED時，一旦有future完成或者被取消，返回；當值爲FIRST_EXCEPTION時，一旦有future拋出異常就返回，倘若沒有異常，它就等效於ALL_COMPLETED；當值爲ALL_COMPLETED時，直到所有future結束或者被取消纔會返回。
源碼分析如下：
# wait源碼
def wait(fs, timeout=None, return_when=ALL_COMPLETED):
with _AcquireFutures(fs):
# 先篩選出完成或被取消的future，放進done中
done = set(f for f in fs
if f._state in [CANCELLED_AND_NOTIFIED, FINISHED])
not_done = set(fs) - done
# 當return_when=FIRST_COMPLETED，同時done不爲空，返回
if (return_when == FIRST_COMPLETED) and done:
return DoneAndNotDoneFutures(done, not_done)
# 當return_when=FIRST_EXCEPTION，同時done不爲空，也有異常存在，返回
elif (return_when == FIRST_EXCEPTION) and done:
if any(f for f in done
if not f.cancelled() and f.exception() is not None):
return DoneAndNotDoneFutures(done, not_done)
# len(done)等於len(fs)時，說明所有future都結束了，返回
if len(done) == len(fs):
return DoneAndNotDoneFutures(done, not_done)
# -------------------------------------------------------
# 上述代碼對futures做了一遍過濾，我想這是爲了儘快返回結果的策略
# 如果運行到這兒，都還沒退出wait，此時藉助_create_and_install_waiters幫忙
waiter = _create_and_install_waiters(fs, return_when)
waiter.event.wait(timeout) # 程序會在此阻塞，直到事件鎖釋放
for f in fs:
with f._condition:
f._waiters.remove(waiter)
done.update(waiter.finished_futures)
return DoneAndNotDoneFutures(done, set(fs) - done)
_create_and_install_waiters()函數的作用類似於分發任務，講as_completed()時也簡單提到過。其處理邏輯是：當as_completed調用它時，任務交給**_AsCompletedWaiter()；當wait調用它，且return_when=FIRST_COMPLETED是，任務交給_FirstCompletedWaiter()；當wait調用它，且return_when=FIRST_EXCEPTION或者ALL_COMPLETED時，任務交給_AllCompletedWaiter()**。
# _FirstCompletedWaiter源碼
class _FirstCompletedWaiter(_Waiter):
"""Used by wait(return_when=FIRST_COMPLETED)."""
def add_result(self, future):
super().add_result(future)
self.event.set()
def add_exception(self, future):
super().add_exception(future)
self.event.set()
def add_cancelled(self, future):
super().add_cancelled(future)
self.event.set()
event.set()會釋放事件鎖，所以add_result()和add_cancelled()中都調用了。需要注意的是，add_exception也調用了set，所以在FIRST_COMPLETED狀態下，出現第一個異常也能導致wait立馬返回。
# _AllCompletedWaiter源碼
class _AllCompletedWaiter(_Waiter):
"""Used by wait(return_when=FIRST_EXCEPTION and ALL_COMPLETED)."""
def __init__(self, num_pending_calls, stop_on_exception):
self.num_pending_calls = num_pending_calls # 計數器
self.stop_on_exception = stop_on_exception
self.lock = threading.Lock()
super().__init__()
def _decrement_pending_calls(self):
with self.lock:
self.num_pending_calls -= 1
if not self.num_pending_calls:
self.event.set()
def add_exception(self, future):
super().add_exception(future)
if self.stop_on_exception: # 如果是FIRST_EXCEPTION，出現異常就釋放鎖
self.event.set()
else:
self._decrement_pending_calls() # 否則對計數器self.num_pending_calls減1，
# 直到self.num_pending_calls = 0，釋放鎖
...
FIRST_EXCEPTION和ALL_COMPLETED都使用了_AllCompletedWaiter的代碼，區別在於前者self.stop_on_exception=True，後者=False。