python Pool+Queue 实现消费者-生产者模型

python Pool+Queue 实现消费者-生产者模型

之前做了一个邮件分类的小task，需要解析的邮件有80w封左右，单一进程解析速度太慢，又因为python的GIL特性，所以选择了多进程的方式，提高cpu多核的利用率来加速邮件解析的速度。

生产者与消费者模型

生产者与消费者模型常用于处理并发问题，生产者用于生产数据，消费者用于获得数据，处理数据。

#encoding: utf-8
from multiprocessing import Process, Manager
import time

def create_data(q):
    # data_source用于模拟我们有1w条数据。
    data_source = (i for i in range(10000))
    # 读取数据
    for data in data_source:
        time.sleep(0.001)
        print 'get data %s' % (data,)
        # 将读取的数据放入队列之中，，以方便别人来处理。
        q.put(data)
    # 1w条数据读取完后，将None作为结束信号放入队列中，告诉别人没数据了，可以结束了。
    q.put(None)

def handle_data(q):
    # 读取队列里的数据，并且处理他们。
    while True:
        # 队列的get方法，可以从队列里获取数据，他是阻塞的方法，如果队列中没有数据，他会一直等待。
        # 当然你也可以设置他的超时时间，一旦超时，会抛出Empty异常。非阻塞式的调用可以使用
        # q.get(False) 或者 q.get_nowait(), 当然队列没数据的话，他们也会抛Empty异常。
        data = q.get()
        # 是否处理完成
        if data == None:
            print 'end'
            break
        time.sleep(0.002)
        # 处理数据，这里只是模拟了一下，象征性的每个数据乘了2，在具体的业务中，你可能会有其他处理数据的操作，
        # 比如将数据存入数据库等。
        result = data * 2
        print 'handle data %s， result: %s' % (data, result)


if __name__ == '__main__':
    data_queue = Manager().Queue()
    # 创建一个生产者的进程，用于生产数据。
    producer = Process(target=create_data, args=(data_queue,))
    # 创建一个消费者的进程，用于处理数据。
    consumer = Process(target=handle_data, args=(data_queue, ))
    # 启动进程
    producer.start()
    consumer.start()
    # 使当前进程，也就是main进程，进入等待状态，当p进程都执行完毕后，
    # 当前进程才会继续下去运行。
    producer.join()
    consumer.join()

通过上面的小例子，可以看出采用多进程的方式能够提高程序的运行速度，不过因为木桶效应，整个程序执行完成的时间取决于运行速度最慢的那一个进程。生产者生产一个数据耗时0.001秒，而消费者处理一个数据需要0.002秒，只有消费者执行完毕后，整个程序才能执行完成。可以看出，我们如果能加速消费者进程的执行，便能让整个程序更快的执行完成。优化的方式有很多种，比如优化处理数据的算法，缩短处理每个数据的时间，如果处理单个数据的时间无法缩短，可以再多开一个进程去帮助消费者进程处理数据。

下面的小例子是通过使用进程池+队列+生产者与消费者模型来并发的处理邮件，因为处理数据的速度慢，所以多开了几个进程去处理数据，进程也不宜开多，数量取决于自己电脑的cpu核心数。

#encoding: utf-8
from multiprocessing import Pool, Manager
from time import sleep, time
import csv

# 模拟一共有1W封邮件
EMAIL_SOURCE = (i for i in range(0,10000))

EMAIL_WORKER_DATA = {}
# worker进程数，别设置多了，和自己的cpu核心数有关。
MAX_WORKER_NUM = 3
RESULT_DATA = Manager().Queue()

# 生产者与消费者模型的应用

def _init_worker_data():
    # 为每一个worker进程分配自己的队列
    for i in range(0, MAX_WORKER_NUM):
        EMAIL_WORKER_DATA[i] = Manager().Queue()

# 读取email的数据，并放入每个worker对应的队列里。
def _read_email_data():
    index = 0
    for data in EMAIL_SOURCE:
        sleep(0.0003)
        print 'load data: %s' % (data)
        q = EMAIL_WORKER_DATA[index % MAX_WORKER_NUM]
        q.put(data)
        index += 1
    for q in EMAIL_WORKER_DATA.itervalues():
        q.put(None)

# 处理队列里的数据，将结果录入结果队列里
def _read_worker_data(worker_id):
    while True:
        sleep(0.003)
        q = EMAIL_WORKER_DATA[worker_id]
        data = q.get()
        if data == None:
            print 'finished'
            RESULT_DATA.put(data)
            break
        print '%s gets data: %s' % (worker_id, data)
        RESULT_DATA.put(data)

# 从结果队列里读取数据，并写入csv文件里，方便以后分析数据。
def _read_result_data():
    file_name = 'test.csv'
    f = open(file_name, 'w')
    writer = csv.writer(f)
    writer.writerow('data')
    while True:
        data = RESULT_DATA.get()
        if data == None:
            print 'finished'
            f.close()
            break
        writer.writerow([data])
        print 'handled result data:%s' % (data,)

if __name__ == '__main__':
    start_time = time()
    _init_worker_data()
    p = Pool(MAX_WORKER_NUM + 2)
    # apply_async，开启一个子进程去执行任务，他是非阻塞式的，
    # 执行到了这条语句后，主程序会继续向下运行，不会被阻塞。
    p.apply_async(_read_email_data)
    for i in range(0, MAX_WORKER_NUM):
        p.apply_async(_read_worker_data, args=(i,))
    p.apply_async(_read_result_data)
    # 关闭进程池，不再接受新的进程
    p.close()
    # 主进程阻塞等待子进程的退出，close要先于join使用。
    p.join()
    end_time = time()
    print 'handle time:%ss' % (end_time - start_time)