兩萬字長文:基於 Python 協程的併發編程實踐

???? Python貓” ,一個值得加星標的公衆號

花下貓語:對多數人來說,併發編程都是非常難啃的硬骨頭。你是否也想系統而深入地瞭解這方面的內容呢?今天給大家分享一篇長文,建議收藏細讀!

作者:咬定青松 | 來源:碼上觀世界公衆號

前言

假設有一批小文件,每個文件都可以通過 mysql load 的方式導入數據庫,請問如何操作可以取得較小的時間和資源消耗?

關於這個需求,我們自然會想到各種併發實現方式,比如多進程和多線程。由於衆所周知的多進程切換的高昂代價以及在某些場合下需要考慮多進程之間的協調和通信,如果情非得已,恐怕很少會使用到多進程。然而在本文討論的 python 世界中,多線程可能也不是一個好的選擇。詳見下文論述。

線程模型

我們知道操作系統的任務調度是基於內核調度實體(KSE,Kernel Scheduling Entity),所以線程的實現也是基於內核調度實體,也就是通過跟內核調度實體綁定實現自身的調度。根據線程與內核實體的對應關係上的區別,線程的實現模型大致可以分爲兩大類:內核級線程和用戶級線程。

- 內核級線程模型

線程與內核線程 KSE 是一對一(1 : 1)的映射模型,也就是每一個用戶線程綁定一個實際的內核線程,而線程的調度則完全交付給操作系統內核去做,應用程序對線程的創建、終止以及同步都基於內核提供的系統調用來完成,大部分編程語言的線程庫 (比如 Java 的 java.lang.Thread、C++ 的 std::thread 等等) 都屬於內核級線程模型。這種模型的優勢和劣勢同樣明顯:優勢是實現簡單,直接藉助操作系統內核的線程以及調度器,所以 CPU 可以快速切換調度線程,於是多個線程可以同時運行,因此相較於用戶級線程模型它真正做到了並行處理;但它的劣勢是,由於直接藉助了操作系統內核來創建、銷燬和以及多個線程之間的上下文切換和調度,因此資源成本大幅上漲,且對性能影響很大。

- 用戶級線程模型

線程與內核線程 KSE 是多對一(N : 1)的映射模型,多個用戶線程的一般從屬於單個進程並且多線程的調度是由用戶自己的線程庫來完成,線程的創建、銷燬以及多線程之間的協調等操作都是由用戶自己的線程庫來負責而無須藉助系統調用來實現。許多語言實現的協程庫 基本上都屬於這種方式(比如 python 的 gevent)。由於線程調度是在用戶層面完成的,避免了系統調用和 CPU 在用戶態和內核態之間切換的開銷,因此對系統資源的消耗會小很多,然而該模型有個問題:假設在某個用戶進程上的某個用戶線程因爲一個阻塞調用(比如 I/O 阻塞)而被 CPU 給中斷(搶佔式調度)了,整個進程將被掛起。因此,該模型並不能做到真正意義上的併發。

python 線程

我們廣泛使用的 python 是基於 CPython 實現,然而由於 CPython 的內存管理不是線程安全的,於是引入了一個全局解釋鎖(Global Interpreter Lock)來保障 Python 的線程安全。正是因爲 GIL 的存在,每個線程執行前都需要獲取鎖,從而導致多線程的併發性大大削弱,完全無法發揮多核的優勢。同時 python 的線程切換是基於字節碼指令的條數,因此對於 I/O 密集型計算密集型任務勉強還有用武之地,然而對於計算密集型任務,多線程切換的開銷將使多線程成爲雞肋,執行效率反而不如單線程。以下是一個驗證例子:
順序執行的單線程 (single_thread.py)

#! /usr/bin/python
 
from threading import Thread
import time
 
def my_counter():
    i = 0
    for _ in range(100000000):
        i = i + 1
    return True
 
def main():
    thread_array = {}
    start_time = time.time()
    for tid in range(2):
        t = Thread(target=my_counter)
        t.start()
        t.join()
    end_time = time.time()
    print("Total time: {}".format(end_time - start_time))
 
if __name__ == '__main__':
    main()

同時執行的兩個併發線程 (multi_thread.py)

#! /usr/bin/python
 
from threading import Thread
import time
 
def my_counter():
    i = 0
    for _ in range(100000000):
        i = i + 1
    return True
 
def main():
    thread_array = {}
    start_time = time.time()
    for tid in range(2):
        t = Thread(target=my_counter)
        t.start()
        thread_array[tid] = t
    for i in range(2):
        thread_array[i].join()
    end_time = time.time()
    print("Total time: {}".format(end_time - start_time))
 
if __name__ == '__main__':
    main()

在 mac os,4 核 8G 內存 1.8MHz python3.7 上測試執行,多線程比單線程慢 2 秒!

另外值得一提的是儘管存在 GIL,但 python 多線程仍然不是線程安全的,對於共享狀態的場合仍然需要藉助鎖同步。既然 python 多線程如此之糟,有沒有一種線程切換代價更小和佔用資源更低的技術呢?下面該輪到協程閃亮登場了!

協程

協程(Coroutine)又稱微線程,屬於用戶級線程。上文中提到的 gevent 就是一種協程實現方式,除了 gevent 還有 asyncio。下文詳細介紹。上文中,我們介紹了用戶級線程就是在一個內核調度實體上映射出來的多個用戶線程,用戶線程的創建、調度和銷燬完全由用戶程序控制, 對內核調度透明:內核一旦將 cpu 分配給了線程,該 cpu 的使用權就歸該線程所有,線程可以再次按照比如時間片輪轉等常規調度算法分配給每個微線程,從而實現更大的併發自由度,但所有的微線程只能在該 cpu 上運行,無法做到並行。爲了便於理解,我們這裏把協程看作這些映射出來的“微線程”。用戶程序控制的協程需要解決線程的掛起和喚醒、現場保護等問題,然而區別於線程的是協程不需要處理鎖和同步問題,因爲多個協程是在一個用戶級線程內進行的,但需要處理因爲單個協程阻塞導致整個線程(進程)阻塞的問題。下圖展示線程和協程的對照關係:

生成器 - 協程基礎

理解協程的掛起和喚醒,不得不提到生成器。生成器也是函數,但跟普通的函數稍有區別,請看下面定義的生成器:

def countdown(n):
    while n> 0:
        yield n
        n -= 1

調用 countdown 並不會執行,如果 print 該函數,會發現返回的是 generator 實例對象。

只有通過 next()函數來執行生成器函數。yield 命令產生了一個值,然後掛起函數,直到下一個 next() 函數。當生成器函數遇到 return 或結束,停止迭代數據。除了 next,還可以使用 send 激活生成器,兩者可以交替使用。比如下面生成斐波那契數列的生成器:

def myfibbo(num):
    a,b=0,1
    count=0
    while count<num:
        a,b=a+b,a
    #yield b 是將b 返回給外部調用程序。
    #ret = yield 可以接收外部程序通過send()發送的信息,並賦值給ret
        ret = yield b
        count+=1
        print("step into here,count={},ret={}".format(count,:ret))

第一次當生成器處於 started 狀態時,只能 send(None),否則會報錯,當生成器 while 條件不滿足退出時,會拋出異常 StopIteration, 如果生成器有返回值,會保存在 exception 的 value 屬性中。

生成器首先是個迭代器,因此生成器可以嵌套調用子生成器。

def reader():
    # 模擬從文件讀取數據的生成器,for表達式可以簡寫爲:yield from range(4)
    #for i in range(4):
    #    yield i
    yield from range(4)
    
def reader_wrapper():
    yield from reader()
    
wrap = reader_wrapper()
for i in wrap:
    print(i)

在這裏 yield from 同時起到了一個提供了一個調用者和子生成器之間的透明的雙向通道的作用: 從子生成器獲取數據以及向子生成器傳送數據。通過上述生成器的例子中,我們已經大體感知到協程的影子了,但還是不夠直觀,而且不是正在意義上的協程,只是實現的代碼執行過程中的掛起,喚醒操作。我們再介紹一個真正的協程實現庫 greelet, 知名的網絡併發框架如 eventlet,gevent 都是基於它實現的。

greenlet
from greenlet import greenlet

def test1():
    print(12)
    gr2.switch()
    print(34)

def test2():
    print(56)
    gr1.switch()
    print(78)

gr1 = greenlet(test1)
gr2 = greenlet(test2)
gr1.switch()

上例中創建了兩個 greenlet 協程對象,gr1 和 gr2,分別對應於函數 test1()和 test2()。從中我們可以看出,使用 switch 方法切換協程,確實比 yield, next/send 組合要直觀得多,從輸出結果看,greenlet 協程的運行是交叉執行的,(本質是串行的)所以它不是真正意義上的併發,因此也無法發揮 CPU 多核的優勢。

創建協程對象的方法其實有兩個參數 greenlet(run=None, parent=None)。參數 run 就是其要調用的方法,比如上例中的函數 test1()和 test2();參數 parent 定義了該協程對象的父協程,也就是說,greenlet 協程之間是可以有父子關係的。如果不設或設爲空,則其父協程就是程序默認的”main”主協程。這個”main”協程不需要用戶創建,它所對應的方法就是主程序,而所有用戶創建的協程都是其子孫。大家可以把 greenlet 協程集看作一顆樹,樹的根節點就是”main”,上例中的 gr1 和 gr2 就是其兩個字節點。

在子協程執行完畢後,會自動返回父協程。比如上例中 test1() 函數退出,代碼會返回到主程序。

eventlet

eventlet 在 Greenlet 的基礎上實現了自己的 GreenThread,實際上就是 greenlet 類的擴展封裝,而與 Greenlet 的不同是,Eventlet 實現了自己調度器稱爲 Hub,Hub 類似於 Tornado 的 IOLoop,是單實例的。在 Hub 中有一個 event loop,根據不同的事件來切換到對應的 GreenThread。同時 eventlet 還實現了一系列的補丁來使 Python 標準庫中的 socket 等 module 來支持 GreenThread 的切換。

eventlet 的 Hub 可以被定製來實現自己調度過程。

eventlet 使用舉例:

import eventlet
from eventlet.green.urllib import request

urls = [
    "http://www.baidu.com",
    "http://www.tmall.com",
    "http://www.tencent.com",
]

def fetch(url):
    print("opening", url)
    body =request.urlopen(url).read()
    print("done with", url)
    return url, body

pool = eventlet.GreenPool(2)
for url, body in pool.imap(fetch, urls):
    print("got body from", url, "of length", len(body))

示例代碼中引入 GreenPool 協程池來控制併發度。

gevent

gevent 是基於 libev(Linux 上 epoll,FreeBSD 上 kqueue)和 greenlet 實現的 Python 網絡庫。libev 是一個事件循環器:向 libev 註冊感興趣的事件,比如 socket 可讀事件,libev 會對所註冊的事件的源進行管理,並在事件發生時觸發相應的程序。也就是說 libev 提供了指定文件描述符事件發生時調用回調函數的機制。而 libev 依賴的 epoll 是 Linux 內核爲處理大批量文件描述符而作了改進的 poll,是 Linux 下多路複用 IO 接口 select/poll 的增強版本,它能顯著提高程序在大量併發連接中只有少量活躍的情況下的系統 CPU 利用率。爲了將 python 標準庫改造成支持 gevent 的非阻塞庫,gevent 使用了 monkey_patch(俗稱“猴子補丁”)的辦法對大部分標準庫包括 socket、ssl、threading 和 select 等模塊做了改寫。所謂“猴子補丁”就是不改變源代碼而對功能進行追加和變更,所以“猴子補丁”並不是 Python 中專有的,一方面它充分利用了動態語言的靈活性,可以對現有的語言 Api 進行追加,替換,修改 Bug,甚至性能優化等,另一方面也給系統維護帶來了一些風險。

gevent 使用舉例:

from gevent import socket

urls = ['www.baidu.com', 'www.example.com', 'www.python.org']
jobs = [gevent.spawn(socket.gethostbyname, url) for url in urls]
gevent.joinall(jobs, timeout=2)
result = [job.value for job in jobs]
print(result)

gevent.spawn()方法 spawn 一些 jobs,然後通過 gevent.joinall 將 jobs 加入到微線程執行隊列中等待其完成,設置超時爲 2 秒。執行後的結果通過檢查 gevent.Greenlet.value 值來收集。gevent.socket.gethostbyname() 函數與標準的 socket.gethotbyname() 有相同的接口,但它不會阻塞主線程。

from gevent import monkey;monkey.patch_all()
import gevent
import requests
import time
import pymysql
from gevent.pool import Pool
    
def query(sql):
    db = pymysql.connect(host ='rm-bp1ek8zy4654v7216zo.mysql.rds.aliyuncs.com', user = 'public_admin',passwd= 'xxxxxxyyyyy', db= 'performance_test')
    cursor = db.cursor()
    data = cursor.execute(sql)
    cursor.close()
    db.close()

if __name__ == '__main__':
    p = Pool(5)
    sql_list=['select  * from seller_payments_report_v2 limit 10' for i in range(50)]
    p.map(query,sql_list)
    p.join()

示例代碼中,引入了協程池來控制併發,通過 mysql 終端 show processlist;可以看到 gevent 實現了對數據庫的併發查詢。值得注意的是這裏簡單的查詢沒有發生阻塞,但複雜的操作比如 load file 就不一定了。感興趣的讀者可以自行驗證。

異步編程框架

上面介紹的 eventlet,gevent 都是從一種同步 IO 模型的角度來實現的,這裏介紹一種異步的實現方式。所謂異步 IO,是跟同步 IO 相對的,異步 IO 是計算機操作系統對輸入輸出的一種處理方式:發起 IO 請求的線程不等 IO 操作完成,就繼續執行隨後的代碼,IO 結果用其他方式 (回調) 通知發起 IO 請求的程序。這樣通過異步 IO,應用程序在發起 IO 請求完成之前,不必等待 IO 完成,就可以繼續去幹其他事情,等待操作系統完成 IO 再通知應用程序去處理。現代操作系統已經將這些 IO 狀態包裝成基本的事件,如可讀事件,可寫事件,並且提供應用程序可以接收這些事件的系統模塊。比如 select 模塊。在 python 中 select 模塊就是 selectors,selectors 是對底層 select/poll/epoll/kqueue 的封裝。DefaultSelector 類會根據 OS 環境自動選擇最佳的模塊,最新的 Linux 系統中基本都是基於 epoll 實現。在詳細介紹 asyncio 之前,先通過一個網絡爬蟲的例子,從最基本的 select 模塊講起。

基於 selector 的異步實現

首先我們把要抓取的 url 地址簡化爲一個 host 主機地址列表:

host_to_access = {'www.baidu.com', 'www.taobao.com', 'www.tencent.com', 'www.toutiao.com', 'www.meituan.com', 'www.tmall.com'}

然後實現一個 Fetcher 類,用於跟一個 url 地址綁定,每個 url 對應一個 Fetcher,用於對 url 的連接和讀取響應:

class Fetcher:
    def __init__(self,host):
        self.response = b'' # Empty array of bytes.
        self.host = host
        self.sock = None

    def fetch(self):
        self.sock = socket.socket()
        self.sock.setblocking(False)
        try:
            self.sock.connect((self.host, 80))
        except BlockingIOError:
            pass
        # Register next callback.
        selector.register(self.sock.fileno(),
                      EVENT_WRITE,
                      self.connected)

    def connected(self, key, mask):
        selector.unregister(self.sock.fileno())
        request='GET / HTTP/1.0\r\nHost: {}\r\n\r\n'.format(self.host)
        print('{} connected'.format(self.host))
        self.sock.send(request.encode('ascii'))
        # Register the next callback.
        selector.register(key.fd,
                      EVENT_READ,
                      self.read_response)

    def read_response(self, key, mask):
        global stopped
        chunk = self.sock.recv(4096)  # 4k chunk size.
        if chunk:
            self.response += chunk
        else:
            selector.unregister(key.fd)  # Done reading.
            host_to_access.remove(self.host)
            if not host_to_access:
                stopped = True
            print("key:{},mask:{},read response:{}".format(key,mask,self.response))

Fetcher 使用非阻塞 socket,這樣主程序就不需要等待 IO 立即返回,這一點在大量 IO 請求的場景下至關重要。爲了讓 IO 請求可讀可寫的時候應用程序能夠去處理,我們註冊了兩個回調函數:connected 和 read_response。前者在建立連接之後,也就是噹噹前 IO 可寫的時候調用;後者在發送請求到服務器之後,服務器有響應內容,客戶端當前 IO 請求可讀的時候調用。接下來的問題是,客戶端同時發起若干個請求,如何知道哪些請求可讀可寫呢?換句話說客戶端如何獲取到操作系統的 IO 事件通知呢?這裏就需要用到 selectors 模塊中的 select 機制了:select 返回當前可讀可寫的事件列表,如果當前沒有事件發生,當前操作將會被阻塞。明顯,這裏需要一種事件循環方式去輪訓檢測可讀可寫的事件,然後調用組冊的回調函數,直到所有請求都處理完畢,示例中的 callback 就是應用程序事先註冊的回調函數入口:

def loop():
    while not stopped:
        events = selector.select()
        for event_key, event_mask in events:
            callback = event_key.data
            callback(event_key, event_mask)

最後在主程序中依次對每個 url 通過 Fetch 的 fetch 方法觸發請求,之後的處理就交給事件循環和回調函數了。

if __name__ == '__main__':
    for host in host_to_access:
        fetcher = Fetcher(host)
        fetcher.fetch()
    loop()

從上面的例子中我們看到異步 IO 方式跟前文中的同步方式有一個共同點,那就是都在一個線程中併發實現多任務處理的。不同點就是再看不到“猴子補丁”的身影了,從代碼安全性上似乎好了不少,但是事實並非如此:在上文中,我們通過 Fetcher 類保存了當前請求的 socket、host 和響應內容。而且這個保存當前應用程序的狀態是不得不爲的,因爲不像同步調用程序那樣,接下來要處理的步驟是確定的,異步調用的方式會在 I/O 操作完成之前返回並清除棧幀,然後在未來某個時刻繼續未完之事。隨着應用程序需要保存的狀態逐步增多,維護應用程序的代價也越大。此其一,其二是回調函數缺少上下文,應用程序的維護者很難從問題現象中迅速查詢被調用函數從哪裏發起,然後又流轉到哪裏。特別是當回調函數嵌套調用回調函數的時候,這種“堆棧撕裂”的問題將變得更加棘手。那有沒有一種更好的方式既保留着回調的優勢又能避免它的問題呢?還記得前文中介紹的生成器嗎?接下來我們使用生成器來重寫上面爬蟲的例子。

基於 generator 和 selector 的異步實現

我們在生成器部分的介紹中瞭解到生成器是可以保存當前狀態的,這裏通過示例詳細展示這個功能點, 首先我們重構 Fetcher 類:

class Fetcher:
    def __init__(self,host):
        self.response = b'' # Empty array of bytes.
        self.host = host

    def fetch(self):
        sock = socket.socket()
        sock.setblocking(False)
        try:
            sock.connect((self.host, 80))
        except BlockingIOError:
            pass
        f = Future()

        def on_connected():
            f.set_result(None)
            print('{} connected'.format(self.host))

        selector.register(sock.fileno(), EVENT_WRITE, on_connected)
        yield f
        selector.unregister(sock.fileno())

        request='GET / HTTP/1.0\r\nHost: {}\r\n\r\n'.format(self.host)
        sock.send(request.encode('ascii'))

        global stopped
        while True:
            f = Future()

            def on_readable():
                try:
                    data=sock.recv(4096)
                    f.set_result(data)
                except socket.error as  e:
                    print(e)
            selector.register(sock.fileno(), EVENT_READ, on_readable)
            chunk = yield f
            selector.unregister(sock.fileno())
            print("host:{},read response:{}".format(self.host,chunk))
            if chunk:
                self.response += chunk
            else:
                print(self.host+" removed.")
                host_to_access.remove(self.host)
                if not host_to_access:
                    stopped = True

Fetcher 類與上文的實現有幾點區別:
1. 不需要保存當前 socket 狀態;
2. 回調函數 on_connected 和 on_readable 移到 Fetcher 類裏,且不需要參數!
3. 通過引入 yield 不僅保留了回調的異步實現,而且保持了同步實現的簡明邏輯。
4. 讀取響應內容的完整邏輯被封裝進 fetch 方法裏。
5. 引入 Future 類,使得程序邏輯按照時間線向未來延伸,通過跟 yield 配合,每次 IO 請求阻塞在 Future 上,同時 fetch 由普通方法變成了一個生成器。

接下來需要解決的問題是,如何喚醒在 Future 處阻塞的請求呢?

我們定義一個 Task 任務類,來驅動整個請求,Task 和 Future 類的實現代碼如下:

class Task:
    def __init__(self, coro):
        self.coro = coro
        f = Future()
        f.set_result(None)
        self.step(f)

    def step(self, future):
        try:
            # send會進入到coro執行, 即fetch, 直到下次yield
            # next_future 爲yield返回的對象
            next_future = self.coro.send(future.result)
        except StopIteration:
            return
        next_future.add_done_callback(self.step)
        
class Future:
    def __init__(self):
        self.result = None
        self._callbacks = []

    def add_done_callback(self, fn):
        self._callbacks.append(fn)

    def set_result(self, result):
        self.result = result
        for fn in self._callbacks:
            fn(self)

注意 Task 的 step 方法:每次 send 喚起生成器並得到新的生成器,新的生成器綁定的 step 方法將在有事件到來時候被再次執行。循環不斷,直到響應內容讀取完畢:將第一次通過初始化 Task 調用 step 執行,通過 send(None)方法激活 fetch 生成器,發起 url 連接請求,當連接請求建立後,通過 on_connected 回調方法執行在 Future 中傳入的 step 方法發起讀請求,再次掛起,當讀信號到來時候,再次觸發 step 喚醒生成器,並將讀取的響應內容傳遞給 chunk。。。這樣一來,生成器和 Task 通過 Futrure 串聯起來了。最後該整個流程的關鍵驅動器事件循環上場了:

def loop():
    while not stopped:
        events = selector.select()
        for event_key, event_mask in events:
            callback = event_key.data
            callback()

if __name__ == '__main__':
    import time
    start = time.time()
    for host in host_to_access:
        fetcher = Fetcher(host)
        Task(fetcher.fetch())
    loop()

從上述示例中,我們看到基於生成器的異步實現,回調函數已經不需要關心是誰觸發了事件,而且每個生成器代碼中也不需要維護 socket 狀態了,整個代碼風格非常接近同步代碼。是時候爲實現這段連接並獲取響應的代碼代碼段正名了 - 協程:即協作式的例程。實際上 python2.5 中也確實有基於生成器的協程實現提案:Coroutines via Enhanced Generators。協程擁有自己的幀棧,每次迭代之間,會暫停執行,繼續下次迭代的時候還不會丟失先前的狀態。然而美中不足的是基於 yield 實現的協程還是不夠優雅,我們再次重構來看看。

基於 yield from 和 selector 的異步實現

通過前面介紹,我們已經知道 yield from 也是 Python 的語法,它可以讓嵌套生成器不必通過循環迭代 yield,而是直接 yield from;此外它還打通了生成器和子生成器。直接看代碼:下面將請求和讀取響應的函數封裝如下:

def connect(sock, address):
    f = Future()
    sock.setblocking(False)
    try:
        sock.connect(address)
    except BlockingIOError:
        pass

    def on_connected():
        f.set_result(None)
        print('{} connected'.format(address))

    selector.register(sock.fileno(), EVENT_WRITE, on_connected)
    yield from f
    selector.unregister(sock.fileno())


def read(sock):
    f = Future()

    def on_readable():
        f.set_result(sock.recv(4096))

    selector.register(sock.fileno(), EVENT_READ, on_readable)
    chunk = yield from f
    selector.unregister(sock.fileno())
    return chunk


def read_all(sock):
    response = []
    chunk = yield from read(sock)
    while chunk:
        response.append(chunk)
        chunk = yield from read(sock)
    return b''.join(response)

這樣 Fetcher 類的實現變得更加簡潔:

class Fetcher:
    def __init__(self,host):
        self.response = b'' # Empty array of bytes.
        self.host = host

    def fetch(self):
        global stopped
        sock = socket.socket()
        yield from connect(sock,(self.host, 80))

        request='GET / HTTP/1.0\r\nHost: {}\r\n\r\n'.format(self.host)
        sock.send(request.encode('ascii'))
        self.response = yield from read_all(sock)
        print("{} response:{}".format(self.host,self.response))
        host_to_access.remove(self.host)
        if not host_to_access:
            stopped = True

另外值得一提的是,yield from 必須是可迭代對象,而 yield 可以是普通對象,需要需要實現 Future 的 __iter__ 方法:

    def __iter__(self):
        yield self
        return self.result

從中我們看到用 yield from 改進基於生成器的協程,代碼抽象程度更高。很多知名異步編程框架也是基於 yield from。然而本文是不是到此爲止了呢?且慢,本文接下來介紹真正的主角 asyncio!

asyncio

上面我們實現的爬蟲可以看着簡化的 asyncio。實際上 asyncio 是 Python 3.4 試驗性引入的異步 I/O 框架,提供了基於協程做異步 I/O 編寫單線程併發代碼的基礎設施。其核心組件有事件循環(Event Loop)、協程 (Coroutine)、任務(Task)、未來對象(Future) 以及其他一些擴充和輔助性質的模塊。Python3.5 中新增的 async/await 語法對協程有了明確而顯式的支持,稱之爲原生協程。實際上 async/await 和 yield from 這兩種風格的協程底層複用共同的實現,而且相互兼容。
使用 asyncio 須經過一下幾個步驟:定義協程函數 ->(封裝成 task->)獲取事件循環 -> 將 task 放到事件循環中執行。定義好的協程並不能直接使用,需要將其包裝成爲了一個任務(task 對象),然後放到事件循環中才能被執行。所謂 task 對象是 Future 類的一個子類,保存了協程運行後的狀態,用於未來獲取協程的結果。在上面的步驟中,之所以在封裝 task 這一個步驟上加上括號,是因爲我們也可以選擇直接將協程放到事件循環中,事件循環會自動幫我們完成這一操作。

任務創建

任務創建有多種方式,第一種方式通過 asyncio 提供的 ensure_future() 函數創建 task,如:

import asyncio
import requests

async def scan(url):
    r = requests.get(url).status_code
    print("{}:{}".format(url,r))
    return r

task = asyncio.ensure_future(scan('http://www.baidu.com'))
loop = asyncio.get_event_loop()
loop.run_until_complete(task)
print(task.result())

第二種,直接通過事件循環的 create_task 方法創建

import asyncio
import requests

async def scan(url):
    r = requests.get(url).status_code
    print("{}:{}".format(url,r))
    return r

loop = asyncio.get_event_loop()
task = loop.create_task(scan('http://www.baidu.com')) # 封裝爲task
loop.run_until_complete(task)
print(task.result())

第三種:直接將協程放到事件循環中執行。這種方法並不是說不用將協程封裝爲 task,而是事件循環內部會自動幫我們完成這一步驟。

import asyncio
import requests

async def scan(url):
    r = requests.get(url).status_code
    print("{}:{}".format(url,r))
    return r

loop = asyncio.get_event_loop()
loop.run_until_complete(scan('http://www.baidu.com'))

無論是上述哪一種方法,最終都需要通過 run_until_complete 方法去執行我們定義好的協程。run_until_complete 是一個阻塞(blocking)調用,直到協程運行結束,它才返回

多協程運行

我們可以將多個協程函數加入事件循環,這時候需要藉助 asyncio.gather 函數或者 asyncio.wait 函數。兩個函數功能極其相似,不同的是,gather 接受的參數是多個協程,而 wait 接受的是一個協程列表。async.wait 會返回兩個值:done 和 pending,done 爲已完成的 task,pending 爲超時未完成的 task。而 async.gather 只返回已完成 task。比如

import asyncio
import requests

async def scan(url):
    r = requests.get(url).status_code
    print("{}:{}".format(url,r))
    return r
    
tasks=[asyncio.ensure_future(scan(urls[i])) for i in range(3)]
loop = asyncio.get_event_loop()
loop.run_until_complete(asyncio.wait(tasks))
for task in tasks:
    print('Task Result:', task.result())
回調函數
import asyncio
import requests
import functools

urls=["http://www.baidu.com","http://www.taobao.com","http://www.toutiao.com"]
total=len(urls)
done=0

async def scan(urls):
    if len(urls)>0:
        url=urls.pop()
        r = requests.get(url).status_code
        return ("{} :{}".format(url,r))

def call_back(loop,future):
    global done
    done=done+1
    print('回調函數,協程返回值爲:{},Done:{}'.format(future.result(),done))
    if done<total:
        task = asyncio.ensure_future(scan(urls))
        task.add_done_callback(functools.partial(call_back,loop))
    else:
        loop.stop()

async def main(loop):
    task = asyncio.ensure_future(scan(urls))
    task.add_done_callback(functools.partial(call_back,loop))

loop = asyncio.get_event_loop()
task = asyncio.ensure_future(main(loop))
loop.run_forever()

示例功能爲通過回調的方式實現對 url 列表的接龍訪問:創建 Task 的同時通過 task.add_done_callback 爲 task 任務增加完成回調函數 call_back,在回調函數中,判斷 url 是否請求完畢,如果沒有請求完畢,繼續新建任務,直到所有 url 請求完畢。如果增加首次創建任務的數量,則可以實現類似協程池的功能。另外我們也看到協程函數可以嵌套調用。

循環退出

從上面示例中我們有兩種辦法退出當前loop,run_until_complete和run_forever。前者會等待task完成後自己再退出,後者會一直運行,直到調用stop。

關於協程函數

在上面的示例中,我們只關心協程的創建和運行,但沒有關注協程函數的實現注意事項:如果協程函數調用了阻塞操作,那麼其他協程和主線程將被阻塞。這意味着協程函數邏輯要麼使用用非阻塞功能,要麼同步調用的功能時間很短,否則無法發揮協程的併發優勢。比如上面的request請求url就是同步調用,無法真正實現併發。幸運的是aio庫中有對應的異步實現:aiohttp。這樣一來,結合asyncio和aiohttp來實現我們的爬蟲將變得非常容易:

import asyncio
import aiohttp

host_to_access = {'www.baidu.com', 'www.taobao.com', 'www.tencent.com', 'www.toutiao.com', 'www.meituan.com', 'www.tmall.com'}

loop = asyncio.get_event_loop()

async def fetch(url):
    async with aiohttp.ClientSession(loop=loop) as session:
        async with session.get(url) as response:
            response = await response.read()
            print('{} response:{}'.format(url,response))
            return response


if __name__ == '__main__':
    tasks = [fetch('http://'+host + '/') for host in host_to_access]
    loop.run_until_complete(asyncio.gather(*tasks))

這裏我們對比生成器版的協程,看到使用asyncio庫帶來的巨大優勢:

  • 沒有了yield 或 yield from,而是async/await

  • 沒有了自造的loop(),取而代之的是asyncio.get_event_loop()

  • 無需自己在socket上做異步操作,不用顯式地註冊和註銷事件,aiohttp庫已經代勞

  • 沒有了顯式的 Future 和 Task,asyncio已封裝

  • 更少量的代碼,更優雅的設計

更重要的是,asyncio帶來了更明顯的性能提升和代碼功能的完善。行文至此,相信我們已經對協程有了更深入的理解,下面我們回到文章開篇提的問題,我們看看用協程如何來解決這個問題?假設這一批小文件存放在aws s3上,這裏需要一個工具類,能夠遍歷和下載s3上的文件到本地,然後再通過異步方式load進mysql,異步方式寫mysql,我們使用aiomysql。

s3工具類
#boto3是適用於 Python 的 AWS 開發工具包
import boto3
import os
import shutil
from tempfile import mkdtemp
from contextlib import contextmanager
from tempfile import NamedTemporaryFile

class S3util:
    def list_keys(bucket,key_prefix):
        s3_client = boto3.client('s3', region_name='us-west-2')
        list_response = s3_client.list_objects(Bucket=bucket, Prefix=key_prefix, )
        keys=[content['Key'] for content in list_response['Contents']]
        return keys
    def download(bucket='etl.data',key):
        with TemporaryDirectory(prefix='s32mysql_') as tmp_dir,\
        NamedTemporaryFile(mode="wb",
                       dir=tmp_dir,
                       suffix=file_ext) as f:
        fname=f.name
        s3.download_file(bucket,key,fname)
        return fname

@contextmanager
def TemporaryDirectory(suffix='', prefix=None, dir=None):
    name = mkdtemp(suffix=suffix, prefix=prefix, dir=dir)
    try:
        yield name
    finally:
        try:
            shutil.rmtree(name)
        except OSError as e:
            # ENOENT - no such file or directory
            if e.errno != errno.ENOENT:
                raise e
        
aiomysql

aiomysql提供了對mysql操作的異步訪問接口,基本使用方法跟同步很相似:

import asyncio
import aiomysql

async def load_files():
    conn = await aiomysql.connect(host="rm-bp1ek8zy4654v7216zo.mysql.rds.aliyuncs.com", port=3306,
                                  user='admin', password='xxxyyy',
                                  db='performance_test', charset='utf8',local_infile=True)
    cursor = await conn.cursor()
    await cursor.execute("LOAD DATA LOCAL INFILE '/tmp/0190000' ignore INTO TABLE bi_supplier_spu_month_detail  FIELDS TERMINATED BY X'01'")
    await cursor.execute("LOAD DATA LOCAL INFILE '/tmp/0190001' ignore INTO TABLE bi_supplier_spu_month_detail  FIELDS TERMINATED BY X'01'")
    await cursor.execute('select * from bi_supplier_spu_month_detail limit 10')
    result = await cursor.fetchall()
    for record in result:
        print("record: ", record)
    await cursor.close()
    conn.close()

loop = asyncio.get_event_loop()
loop.run_until_complete(load_files())

大併發場景下,持續的創建數據庫連接可能會導致連接超時、失敗等問題,創建連接池是個好主義,我們爲上面代碼增加連接池:

async def get_pool():
    '''
    初始化,獲取數據庫連接池
    '''
    try:
        print("aiomysql.create_pool")
        pool = await aiomysql.create_pool(host='rm-bp1ek8zy4654v7216zo.mysql.rds.aliyuncs.com', port=3306,
                                    user='public_admin', password='w3SJ605T1AYtjpCs',
                                    db='performance_test', charset='utf8',local_infile=True)
        return pool
    except asyncio.CancelledError:
        raise asyncio.CancelledError
    except Exception as ex:
        print("mysql數據庫連接失敗:{}".format(ex.args[0]))
        return False

async def getCurosr(pool):
    '''
    獲取db連接和cursor對象,用於db的讀寫操作
    '''
    conn = await pool.acquire()
    cur = await conn.cursor()
    return conn, cur

結合前文中介紹的知識鋪墊,我們很容易實現上述需求:

import asyncio
import functools
import logging

loop =asyncio.get_event_loop()
s3=S3util()
list_keys=s3.list_keys()

total=len(list_keys)
complete=0
fail_list=[]
async def store(keys,table_name):
    key=None
    try:
        key=keys.pop()
        file=s3.download(key)
        pool=await get_pool()
        conn,cursor=await getCurosr(pool)
        await cursor.execute(F"LOAD DATA LOCAL INFILE {file} ignore INTO TABLE {table_name}  FIELDS TERMINATED BY X'01'")
    except Exception as e:
        pass
    finally:
        if key:
            global complete
            complete+=1

def call_back(loop,future,table_name):
    logging.info('return result:{},progress:{}/{}'.format(future.result(),complete,total))
    if len(list_obj)>0:
        task = asyncio.ensure_future(store(list_keys,table_name))
        task.add_done_callback(functools.partial(call_back,loop,table_name))
    if complete>=total:
        loop.stop()
      
if __name__ == '__main__':
    concurrency=10
    table_name='bi_supplier_spu_month_detail'
    for i in range(concurrency):
        task = asyncio.ensure_future(store(list_keys,table_name))
        task.add_done_callback(functools.partial(call_back,loop,table_name))

我們通過在美國西部的aws服務器上拉取數據,同時load進杭州阿里雲的基礎版配置的數據庫上測試,110M的(144個)的文件在6分鐘內完成。相比其他工具,比如datax的並行插入速度有接近一倍的提升。

優質文章,推薦閱讀:

漲見識了,在終端執行 Python 代碼的 6 種方式!

趣漫畫:Java 對 Python 的滲透能成功嗎?

11 個最佳的 Python 編譯器和解釋器

深度剖析爲什麼 Python 中整型不會溢出?

感謝創作者的好文

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章