使用生成器把Kafka寫入速度提高1000倍G

通過本文你會知道Python裡面什麼時候用yield最合適。本文不會給你講生成器是什麼，所以你需要先了解Python的yield，再來看本文。

疑惑

多年以前，當我剛剛開始學習Python協程的時候，我看到絕大多數的文章都舉了一個生產者-消費者的例子，用來表示在生產者內部可以隨時呼叫消費者，達到和多執行緒相同的效果。這裡憑記憶簡單還原一下當年我看到的程式碼：

import time


def consumer():
    product = None
    while True:
        if product is not None:
            print('consumer: {}'.format(product))
        product = yield None


def producer():
    c = consumer()
    next(c)
    for i in range(10):
        c.send(i)

start = time.time()
producer()
end = time.time()
print(f'直到把所有資料塞入Kafka，一共耗時：{end - start}秒')
複製程式碼

執行效果如下圖所示。

這些文章的說法，就像統一好了口徑一樣，說這樣寫可以減少執行緒切換開銷，從而大大提高程式的執行效率。但是當年我始終想不明白，這種寫法與直接呼叫函式有什麼區別，如下圖所示。

直到後來我需要操作Kafka的時候，我明白了使用yield的好處。

探索

為了便於理解，我會把實際場景做一些簡化，以方便說明事件的產生髮展和解決過程。事件的起因是我需要把一些資訊寫入到Kafka中，我的程式碼一開始是這樣的：

import time
from pykafka import KafkaClient

client = KafkaClient(hosts="127.0.0.1:9092")
topic = client.topics[b'test']


def consumer(product):
    with topic.get_producer(delivery_reports=True) as producer:
        producer.produce(str(product).encode())


def feed():
    for i in range(10):
        consumer(i)


start = time.time()
feed()
end = time.time()
print(f'直到把所有資料塞入Kafka，一共耗時：{end - start}秒')
複製程式碼

這段程式碼的執行效果如下圖所示。

寫入10條資料需要100秒，這樣的龜速顯然是有問題的。問題就出在這一句程式碼：

with topic.get_producer(delivery_reports=True) as producer
複製程式碼

獲得Kafka生產者物件是一個非常耗費時間的過程，每獲取一次都需要10秒鐘才能完成。所以寫入10個資料就獲取十次生產者物件。這消耗的100秒主要就是在獲取生產者物件，而真正寫入資料的時間短到可以忽略不計。

由於生產者物件是可以複用的，於是我對程式碼作了一些修改：

import time
from pykafka import KafkaClient

client = KafkaClient(hosts="127.0.0.1:9092")
topic = client.topics[b'test']
products = []


def consumer(product_list):
    with topic.get_producer(delivery_reports=True) as producer:
        for product in product_list:
            producer.produce(str(product).encode())


def feed():
    for i in range(10):
        products.append(i)
    consumer(products)


start = time.time()
feed()
end = time.time()
print(f'直到把所有資料塞入Kafka，一共耗時：{end - start}秒')
複製程式碼

首先把所有資料存放在一個列表中，最後再一次性給consumer函式。在一個Kafka生產者物件中展開列表，再把資料一條一條塞入Kafka。這樣由於只需要獲取一次生產者物件，所以需要耗費的時間大大縮短，如下圖所示。

這種寫法在資料量小的時候是沒有問題的，但資料量一旦大起來，如果全部先放在一個列表裡面的話，伺服器記憶體就爆了。

於是我又修改了程式碼。每100條資料儲存一次，並清空暫存的列表：

import time
from pykafka import KafkaClient

client = KafkaClient(hosts="127.0.0.1:9092")
topic = client.topics[b'test']


def consumer(product_list):
    with topic.get_producer(delivery_reports=True) as producer:
        for product in product_list:
            producer.produce(str(product).encode())


def feed():
    products = []
    for i in range(1003):
        products.append(i)
        if len(products) >= 100:
            consumer(products)
            products = []

    if products:
        consumer(products)


start = time.time()
feed()
end = time.time()
print(f'直到把所有資料塞入Kafka，一共耗時：{end - start}秒')
複製程式碼

由於最後一輪迴圈可能無法湊夠100條資料，所以feed函式裡面，迴圈結束以後還需要判斷products列表是否為空，如果不為空，還要再消費一次。這樣的寫法，在上面這段程式碼中，一共1003條資料，每100條資料獲取一次生產者物件，那麼需要獲取11次生產者物件，耗時至少為110秒。

顯然，要解決這個問題，最直接的辦法就是減少獲取Kafka生產者物件的次數並最大限度複用生產者物件。如果讀者舉一反三的能力比較強，那麼根據開關檔案的兩種寫法：

# 寫法一
with open('test.txt', 'w', encoding='utf-8') as f:
    f.write('xxx')
    
# 寫法二
f = open('test.txt', 'w', encoding='utf-8')
f.write('xxx')
f.close()
複製程式碼

可以推測出獲取Kafka生產者物件的另一種寫法：

# 寫法二
producer = topic.get_producer(delivery_reports=True)
producer.produce(b'xxxx')
producer.close()
複製程式碼

這樣一來，只要獲取一次生產者物件並把它作為全域性變數就可以一直使用了。

然而，pykafka的官方文件中使用的是第一種寫法，通過上下文管理器with來獲得生產者物件。暫且不論第二種方式是否會報錯，只從寫法上來說，第二種方式必需要手動關閉物件。開發者經常會出現開了忘記關的情況，從而導致很多問題。而且如果中間出現了異常，使用上下文管理器的第一種方式會自動關閉生產者物件，但第二種方式仍然需要開發者手動關閉。

函式VS生成器

但是如果使用第一種方式，怎麼能在一個上下文裡面接收生產者傳進來的資料呢？這個時候才是yield派上用場的時候。

首先需要明白，使用yield以後，函式就變成了一個生成器。生成器與普通函式的不同之處可以通過下面兩段程式碼來進行說明：

def funciton(i):
    print('進入')
    print(i)
    print('結束')

for i in range(5):
    funciton(i)

複製程式碼

執行效果如下圖所示。

函式在被呼叫的時候，函式會從裡面的第一行程式碼一直執行到某個return或者函式的最後一行才會退出。

而生成器可以從中間開始執行，從中間跳出。例如下面的程式碼：

def generator():
    print('進入')
    i = None
    while True:
        if i is not None:
            print(i)
        print('跳出')
        i = yield None

g = generator()
next(g)
for i in range(5):
    g.send(i)
複製程式碼

執行效果如下圖所示。

從圖中可以看到，進入只列印了一次。程式碼執行到i = yield None後就跳到外面，外面的資料可以通過g.send(i)的形式傳進生成器，生成器內部拿到外面傳進來的資料以後繼續執行下一輪while迴圈，列印出被傳進來的內容，然後到i = yield None的時候又跳出。如此反覆。

所以回到最開始的Kafka問題。如果把with topic.get_producer(delivery_reports=True) as producer寫在上面這一段程式碼的print('進入')這個位置上，那豈不是隻需要獲取一次Kafka生產者物件，然後就可以一直使用了？

根據這個邏輯，設計如下程式碼：

import time
from pykafka import KafkaClient

client = KafkaClient(hosts="127.0.0.1:9092")
topic = client.topics[b'test']


def consumer():
    with topic.get_producer(delivery_reports=True) as producer:
        print('init finished..')
        next_data = ''
        while True:
            if next_data:
                producer.produce(str(next_data).encode())
            next_data = yield True


def feed():
    c = consumer()
    next(c)
    for i in range(1000):
        c.send(i)

start = time.time()
feed()
end = time.time()
print(f'直到把所有資料塞入Kafka，一共耗時：{end - start}秒')
複製程式碼

這一次直接插入1000條資料，總共只需要10秒鐘，相比於每插入一次都獲取一次Kafka生產者物件的方法，效率提高了1000倍。執行效果如下圖所示。

後記

讀者如果仔細對比第一段程式碼和最後一段程式碼，就會發現他們本質上是一回事。但是第一段程式碼，也就是網上很多人講yield的時候舉的生產者-消費者的例子之所以會讓人覺得毫無用處，就在於他們的消費者幾乎就是秒執行，這樣看不出和函式呼叫的差別。而我最後這一段程式碼，它的消費者分成兩個部分，第一部分是獲取Kafka生產者物件，這個過程非常耗時；第二部分是把資料通過Kafka生產者物件插入Kafka，這一部分執行速度極快。在這種情況下，使用生成器把這個消費者程式碼分開，讓耗時長的部分只執行一次，讓耗時短的反覆執行，這樣就能體現出生成器的優勢。

我的公眾號：未聞Code