python小demo-01: 线程池+多进程实现cpu密集型操作

起因：

公司有一个小项目，大概逻辑如下：

　　服务器A会不断向队列中push消息，消息主要内容是视频的地址，服务器B则需要不断从队列中pop消息，然后将该视频进行剪辑最终将剪辑后的视频保存到云服务器。个人主要实现B服务器逻辑。

实现思路：

　　1 线程池+多进程

　　　　要求点一：主进程要以daemon的方式运行。
　　　　要求点二：利用线程池，设置最大同时运行的worker，每一个线程通过调用subprocess中的Popen来运行wget ffprobe ffmpeg等命令处理视频。

　　2 消息队列采用redis的list实现

　　3 主线程从队列中获取到消息后，从线程池中获取空闲从线程(在这里，非主线程统称为从线程，下同)，从线程对该消息做一些逻辑上的处理后，然后生成进程对视频进行剪辑，最后上传视频。

　　　　要求点三：为了让daemon能在收到signint信号时，处理完当前正在进行的worker后关闭，且不能浪费队列中的数据，需要让主进程在有空闲worker时才从队列中获取数据。

大概就是这样：

基本上主要资源耗费在视频下载以及视频处理上，且同时运行的worker(从线程)不会太多(一般cpu有几个就设置几个worker)。

上面一共有三个要求点，其中要求点二并不费事。所以忽略。

实现

要求点一实现：

# -*- coding: utf8 -*-
import os
import sys
import time
import signal
import traceback


# from shadowsocks
def write_pid_file(pid_file, pid):
    import fcntl
    import stat

    try:
        fd = os.open(pid_file, os.O_RDWR | os.O_CREAT,
                     stat.S_IRUSR | stat.S_IWUSR)
    except OSError:
        traceback.print_exc()
        return -1
    flags = fcntl.fcntl(fd, fcntl.F_GETFD)
    assert flags != -1
    flags |= fcntl.FD_CLOEXEC
    r = fcntl.fcntl(fd, fcntl.F_SETFD, flags)
    assert r != -1
    # There is no platform independent way to implement fcntl(fd, F_SETLK, &fl)
    # via fcntl.fcntl. So use lockf instead
    try:
        fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB, 0, 0, os.SEEK_SET)
    except IOError:
        r = os.read(fd, 32)
        if r:
            print('already started at pid %s' % (r))
        else:
            print('already started')
        os.close(fd)
        return -1
    os.ftruncate(fd, 0)
    os.write(fd, (str(pid)))
    return 0


def freopen(f, mode, stream):
    oldf = open(f, mode)
    oldfd = oldf.fileno()
    newfd = stream.fileno()
    os.close(newfd)
    os.dup2(oldfd, newfd)


def daemon_start(settings, main_process_handler):
    def handle_exit(signum, _):
        if signum == signal.SIGTERM:
            sys.exit(0)
        sys.exit(1)

    signal.signal(signal.SIGINT, handle_exit)
    signal.signal(signal.SIGTERM, handle_exit)
    pid = os.fork()
    assert pid != -1

    # Parent
    if pid:
        time.sleep(3)
        sys.exit(0)

    print("child has forked")
    # child signals its parent to exit
    ppid = os.getppid()
    pid = os.getpid()
    if write_pid_file(settings.PID_FILE, pid) != 0:
        os.kill(ppid, signal.SIGINT)
        sys.exit(1)

    # set self to process-group-leader
    os.setsid()

    signal.signal(signal.SIGHUP, signal.SIG_IGN)

    print('started')
    os.kill(ppid, signal.SIGTERM)

    # octal 022
    os.umask(18)
    sys.stdin.close()
    try:
        freopen(settings.DEBUG_LOG_PATH, 'a', sys.stdout)
        freopen(settings.DEBUG_LOG_PATH, 'a', sys.stderr)
    except IOError:
        print(traceback.print_exc())
        sys.exit(1)

    main_process_handler()


def daemon_stop(pid_file):
    import errno
    try:
        with open(pid_file) as f:
            pid = buf = f.read()
            if not buf:
                print('not running')
    except IOError as e:
        print(traceback.print_exc())
        if e.errno == errno.ENOENT:
            print("not running")
            # always exit 0 if we are sure daemon is not running
            return
        sys.exit(1)
    pid = int(pid)
    if pid > 0:
        try:
            os.kill(pid, signal.SIGTERM)
        except OSError as e:
            if e.errno == errno.ESRCH:
                print('not running')
                # always exit 0 if we are sure daemon is not running
                return
            print(traceback.print_exc())
            sys.exit(1)
    else:
        print('pid is not positive: %d', pid)

    # sleep for maximum 300s
    for i in range(0, 100):
        try:
            # query for the pid
            os.kill(pid, 0)
        except OSError as e:
            # not found the process
            if e.errno == errno.ESRCH:
                break
        time.sleep(3)
        print("waiting for all threads finished....")
    else:
        print('timed out when stopping pid %d', pid)
        sys.exit(1)
    print('stopped')
    os.unlink(pid_file)


def main():
    args = sys.argv[1:]
    assert len(args) == 2
    if args[0] not in ["stop", "start"]:
        print("only supported: [stop | start]")
        return
    if args[1] not in ["dev", "local", "prod"]:
        print("only supported: [dev | local | prod]")

    from globals import set_settings, initialize_redis
    set_settings(args[1])
    initialize_redis()
    from globals import settings
    import entry

    if args[0] == "stop":
        print("stopping...")
        daemon_stop(settings.PID_FILE)
    elif args[0] == "start":
        print("starting...")
        daemon_start(settings, entry.run)


main()

daemon.py

要求点三实现：

线程池，采用python的futures模块。该模块提供了线程池的机制。稍微说一下他的线程池实现原理吧，ThreadPoolExecutor该类实现了线程池：

　　1 每个实例本身有个_work_queue属性，这是一个Queue对象，里面存储了任务。

　　2 每当我们调用该对象的submit方法时，都会向其_work_queue中放入一个任务，同时生成从线程，直到从线程数达到max_worker所设定的值。

　　3 该线程池实例中所有的从线程会不断的从_work_queue中获取任务，并执行。同时从线程的daemon属性被设置为True

# -*- coding: utf8 -*-
import json
import traceback
import signal
import sys
import time
from threading import Lock
from concurrent.futures import ThreadPoolExecutor
from .globals import settings, video_info_queue




def handler(data):
    # 业务逻辑


running_futures_count = 0


def run():
    global running_futures_count
    count_lock = Lock()

    pool = ThreadPoolExecutor(max_workers=settings.MAX_WORKER)
    try:
        def reduce_count(_):
            global running_futures_count
            with count_lock:
                running_futures_count -= 1

        def handle_exit(_, __):
            print("get SIGINT signal")
            pool.shutdown(False)
            while True:
                if running_futures_count == 0:
                    sys.exit(0)
                time.sleep(1)
                print("now running futures count is %s, please wait" % running_futures_count)

        def handle_data(data):
            global running_futures_count
            with count_lock:
                running_futures_count += 1
            future = pool.submit(handler, data)
            future.add_done_callback(reduce_count)

        signal.signal(signal.SIGINT, handle_exit)
        signal.signal(signal.SIGTERM, handle_exit)

        while not pool._shutdown:
            print(len(pool._work_queue.queue), pool._shutdown)
            while not pool._shutdown and (len(pool._work_queue.queue) < pool._max_workers):
                data = video_info_queue.bpop(20)
                if data:
                    handle_data(data)
                else:
                    data = abnormal_video_info_queue.bpop(1)
                    print("video_info_queue is empty, get data: %s from abnormal_video_info_queue" % data)
                    if data:
                        print("abnormal_video_info_queue")
                        handle_data(data)
            time.sleep(5)
            print("now all the workers is busy, so wait and do not submit")
    finally:
        pool.shutdown(False)