Build a Simple Message Queue (Async Task Queue) Yourself: Python Implementation

Another reinventing-the-wheel series. There are many mature, well-developed async task queue frameworks available for Python, such as Celery or RQ, but none of them come with a built-in message queue service—they all require a message queue like Redis or RabbitMQ. I don't need all that extra stuff for a small project, so I decided to implement one myself.

Approach

1. Add tasks that need to run asynchronously to a queue
2. Automatically take tasks from the queue and create new threads to execute them
3. Save the task execution results and output
4. When a task completes, call a callback function to handle the returned data
5. Use output redirection to handle task output

Below I'll analyze the code for a few key parts.

Priority Queue

Python's standard library already has a priority queue implementation, but when objects of equal priority are passed in, it automatically compares those objects. Our Task class doesn't implement the relevant operator overloads, so they can't be compared.

To solve this, we could either implement operator overloading in Task, or implement a priority queue ourselves. I chose the latter.

class PriorityQueue:
    """
    自己实现的优先级队列，使用了Python的堆
    """
    def __init__(self):
        self._index = 0
        self.queue = []

    def push(self, priority, val):
        heapq.heappush(self.queue, (priority, self._index, val))
        self._index += 1

    def pop(self):
        return heapq.heappop(self.queue)[-1]

    @property
    def empty(self):
        return len(self.queue) == 0

Output Redirection

This part is actually pretty tricky. After all, it's an async task queue, so thread contention is inevitable, and there's no way to control each task's output independently. Still, I went ahead and did it—it's usable for small projects (honestly I was just too lazy to port some existing code).

The code is as follows. It likewise uses a queue to implement an output buffer; you can adjust the size yourself. By default it supports outputting to the console, a file, or returning a list.

If you want custom redirected output, you need to implement an output handling function yourself and assign it directly to the custom attribute of the Redirection class.

class Redirection:
    def __init__(self, buffer_size=512):
        self.buffer = Queue(maxsize=512)
        self._console = sys.stdout
        # 自定义的输出端
        self.custom = None

    def write(self, output_stream):
        # 加入缓冲区队列
        self.buffer.put(output_stream)

    def to_console(self):
        sys.stdout = self._console
        # 出列
        while not self.buffer.empty():
            print(self.buffer.get())

    def to_file(self, file_path):
        with open(file_path, 'w+') as f:
            sys.stdout = f
            while not self.buffer.empty():
                print(self.buffer.get())
            f.close()

    def to_custom(self):
        while not self.buffer.empty():
            self.custom(self.buffer.get())

    def to_list(self):
        data = []
        while not self.buffer.empty():
            data.append(self.buffer.get())
        return data

    def flush(self):
        self.buffer.empty()

    def reset(self):
        sys.stdout = self._console

Task Class

Not much to say. Define the task class, Task. The code is as follows:

class Task:
    def __init__(self, func, callback=None, priority=Priority.MIDDLE, args=(), kwargs={}):
        """
        Args:
            func: 需要执行的函数
            callback:  执行完的回调函数
            priority: 优先级
            *args:
            **kwargs:
        """
        self._id = uuid.uuid4().hex
        self.function = func
        self.callback = callback
        self.priority = priority
        self.args = args
        self.kwargs = kwargs
        # 任务运行过程的输出，stdout的输出
        self._outputs: Redirection = None

    @property
    def id(self):
        return self._id

    @property
    def outputs(self) -> Redirection:
        return self._outputs

    @outputs.setter
    def outputs(self, value: Redirection):
        self._outputs = value

    def run(self):
        try:
            if self.callback:
                # 回调函数原型 callback(task_obj, result)
                result = self.callback(self, self.function(*self.args, **self.kwargs))
            else:
                result = self.function(*self.args, **self.kwargs)
            return result
        except Exception as e:
            if self.callback:
                result = self.callback(self, e)
            else:
                result = e
            return result

It's very simple to use, as follows:

def fun1(num1, num2):
    print(f'num1={num1}')
    print(f'num2={num2}')
    return num1 + num2

Task(
    func=fun1,
    callback=lambda task, result: print(f'task result: {result}'),
)

# 也可以写成这样
Task(
    func=lambda num1, num2: num1 + num2,
    callback=lambda task, result: print(f'task result: {result}'),
    args=[2, 3]
)

Task Queue

Finally, implementing the task queue. It's also simple—just schedule according to task priority.

class TaskQueue:
    """
    基于线程的异步任务队列
    todo 下次做一个基于进程的队列，充分利用多核CPU性能
    """

    def __init__(self, output_redirect=False):
        self.queue = PriorityQueue()
        self.output_redirect = output_redirect
        self._redirect_objs = {}
        self._results = {}

    def put(self, task: Task):
        """
        将task加入任务列表
        Args:
            task:
        Returns:返回task id
        """
        self.queue.push(task.priority, task)
        return task.id

    def get(self):
        return self.queue.pop()

    def run(self):
        while not self.queue.empty:
            task = self.get()
            # 开启新线程
            t = threading.Thread(target=self._task_wrapper, name=task.id, args=[task])
            self._log(f'Start thread {task.id}')
            t.start()

    def get_output(self, task_id: str) -> Redirection:
        return self._redirect_objs.get(task_id, None)

    def get_result(self, task_id: str):
        return self._results.get(task_id, None)

    @staticmethod
    def _log(msg: str):
        """日志输出接口，可以替换为日志组件"""
        print(f'[TaskQueue] {msg}')

    def _task_wrapper(self, task: Task):
        if self.output_redirect:
            if task.id in self._redirect_objs:
                redirect_obj = self._redirect_objs[task.id]
            else:
                redirect_obj = Redirection(2048)
                self._redirect_objs[task.id] = redirect_obj
            # 重定向输出
            sys.stdout = redirect_obj
            task.outputs = redirect_obj
            result = task.run()
            # 恢复默认输出
            redirect_obj.reset()
            self._log(f'Task finished. {task.id}')
        else:
            result = task.run()

        # 保存结果
        self._results[task.id] = result

That's all the project code. It's simple, and there are plenty of shortcomings, but it's barely good enough for small projects~

I'll publish it to pip in a few days. Anyone who needs it can just pip install it. Please follow the blog for updates.

Demo

Here's demo code for using the task queue:

if __name__ == '__main__':
    task_queue = TaskQueue(output_redirect=True)


    def fun1():
        time.sleep(2)
        return 1


    def fun2():
        time.sleep(3)
        return 2


    task_queue.put(Task(
        func=fun1,
        callback=lambda task, result: print(f'task1 result: {result}'),
    ))
    task_queue.put(Task(
        func=fun2,
        callback=lambda task, result: print(f'task2 result: {result}'),
    ))


    def custom_output(msg):
        print(f'[custom] {msg}')


    def fun3(num1, num2):
        print(f'num1={num1}')
        print(f'num2={num2}')
        return num1 + num2


    def callback(task_obj, result):
        print(f'task3 result={result}')
        output = task_obj.outputs
        output.custom = custom_output
        output.reset()
        output.to_custom()


    task_queue.put(Task(func=fun3, callback=callback, args=[2, 3]))

    print('task queue run')
    task_queue.run()
    print('do other things...')

    for i in range(0, 100):
        print(i * i)

References

• Operator overloading https://blog.csdn.net/goodlixueyong/article/details/52589979
• https://www.open-open.com/code/view/1433410658322
• Celery async task queue https://www.cnblogs.com/StitchSun/p/8552488.html
• https://zhuanlan.zhihu.com/p/37637660

Welcome to Discuss

For questions, please leave a message on the WeChat official account backend. I’ll reply to every message~

• WeChat official account: 画星星高手
• Coding livestream: https://live.bilibili.com/11883038
• Zhihu: https://www.zhihu.com/people/dealiaxy
• Jianshu: https://www.jianshu.com/u/965b95853b9f