《asyncio 系列》12. 详解 asyncio 支持的多种队列

楔子

在设计应用程序来处理事件或其他类型的数据时,经常需要一种机制来存储这些事件,并将它们分发给一组 worker。然后这些 worker 可根据这些事件同时执行我们需要执行的任何操作,从而节省时间。asyncio 提供了一个异步队列,可以让我们实现这一点,可将数据块添加到队列中,并让多个 worker 同时运行,从队列中提取数据并在可用时对其进行处理。

这些通常称为 producer-consumer 工作流,某些情况会产生我们需要处理的数据或事件,而处理这些工作内容可能需要很长时间。队列负责帮助我们传输长时间运行的任务,同时保持用户界面持续对外界进行响应。我们可将一个项目放在队列中以供日后处理,并通知用户我们已经在后台开始了这项工作。异步队列还有一个额外优势,就是它提供了一种限制并发的机制,因为每个队列通常允许有限数量的 worker 任务。

异步队列基本知识

队列是一种先进先出的数据结构,这与在杂货店结账时的队列没有太大区别。在结账时,你加入队列,并排在队尾,等待收银员为你前面的所有人结账。一旦收银员为前面的顾客结完账,你就会在队列中移动,而在你之后加入的人会在你身后等待。然后,当你排在队列的第一个位置时,收银员将为你结账。结账后,你将离开队列。

正如我们所描述的,结账队列是一个同步工作流,一名收银员一次为一名顾客结账。如果我们重新设计队列,从而更好地利用并发性,并依旧使用超市收银的例子会怎样?这将意味着多个收银员和一个队列,而不是一个收银员。只要有收银员,他们就可以将下一个顾客引导到收银台,这意味着除了多个收银员同时为客户结账,还有多个收银员同时从队列中引导客户。

这是异步队列的核心内容,我们将多个等待处理的工作项添加到队列中,然后让多个 worker 从队列中提取项目并执行。

import asyncio
from asyncio import Queue
from random import sample, randint
from typing import List

class Product:
    """
    商品
    """
    def __init__(self, name: str, checkout_time: float):
        self.name = name  # 商品名称
        self.checkout_time = checkout_time  # 结算需要的时间

class Customer:
    """
    客户
    """
    def __init__(self, customer_id, products: List[Product]):
        self.customer_id = customer_id  # 客户的 id
        self.products = products  # 客户购买的商品

async def checkout_customer(queue: Queue, cashier_id: int):
    # 检查队列中是否有客户
    while not queue.empty():
        customer: Customer = await queue.get()
        print(f"收银员 {cashier_id} 开始对客户 {customer.customer_id} 的商品进行结算")
        for product in customer.products:
            print(f"收银员 {cashier_id} 正在结算客户 {customer.customer_id} 的商品: {product.name}")
            await asyncio.sleep(product.checkout_time)
        print(f"收银员 {cashier_id} 已完成对客户 {customer.customer_id} 商品的结算")
        queue.task_done()  # 这行代码后续解释

async def main():
    customer_queue = Queue()
    all_products = [Product("苹果", 2), Product("香蕉", .5),
                    Product("草莓", 1), Product("蓝莓", .2)]
    # 创建 4 个客户,并用随机产品进行填充。
    for i in range(1, 5):
        products = sample(all_products, randint(1, 4))
        await customer_queue.put(Customer(i, products))
    # 创建三个收银员,从队列中取出客户,进行服务
    cashiers = [asyncio.create_task(checkout_customer(customer_queue, i)) for i in range(1, 4)]
    await asyncio.gather(customer_queue.join(), *cashiers)

asyncio.run(main())
"""
收银员 1 开始对客户 1 的商品进行结算
收银员 1 正在结算客户 1 的商品: 蓝莓
收银员 2 开始对客户 2 的商品进行结算
收银员 2 正在结算客户 2 的商品: 草莓
收银员 3 开始对客户 3 的商品进行结算
收银员 3 正在结算客户 3 的商品: 草莓
收银员 1 正在结算客户 1 的商品: 香蕉
收银员 1 已完成对客户 1 商品的结算
收银员 1 开始对客户 4 的商品进行结算
收银员 1 正在结算客户 4 的商品: 蓝莓
收银员 1 正在结算客户 4 的商品: 苹果
收银员 2 已完成对客户 2 商品的结算
收银员 3 正在结算客户 3 的商品: 苹果
收银员 1 已完成对客户 4 商品的结算
收银员 3 正在结算客户 3 的商品: 蓝莓
收银员 3 已完成对客户 3 商品的结算
"""

代码逻辑不难理解,我们创建了两个数据类:一个用于产品,一个用于客户。Product 类由产品名称和收银员结算该项目所需的时间(以秒为单位)组成,Customer 类由客户 id 和所选择的商品组成(客户可以选择多个产品)。而 checkout_customer 协程所做的事情就是从队列中取出客户,然后对他选择的商品进行结算,比较简单。

然后在主协程函数中,我们创建一个可用的产品列表并生成 4 位客户,每位客户都带有随机的产品,并将其放入到队列中。然后我们创建了 3 个收银员,分别去处理这些客户,对他们的选择的商品进行结算。最后使用 gather 等待收银员 checkout_customer 任务与 customer_queue.join() 协程一起完成。

这里需要说一下队列的 join 方法,首先队列内部有一个 _unfinished_tasks 属性,初始值为 0,当往队列塞入元素时,该属性自增 1。但是从队列中取出元素时,这个值并不会自减 1,而是需要手动调用 task_done 方法,该值才会减 1。而当值不为 0 时,调用 join 会阻塞,如果值为 0 则解除阻塞。

所以我们也要对 queue.join() 进行等待,目的是使队列中的元素全部被消费后,才能解除阻塞。

异步队列源码解析

我们之前介绍了锁、信号量、事件、条件,还有这里的队列,它们都有等待、唤醒的功能,比如:

  • 锁:当锁被获取时,再次获取会阻塞,直到锁被释放,解除阻塞;
  • 事件:当标志位没有设置为 True 时,调用 wait 会阻塞,直接标志位设置为 True,解除阻塞;
  • 队列:当队列为空(已满)时,获取(添加)元素会阻塞,直到队列里面有元素(可用空间),解除阻塞;

那么这背后的等待、唤醒对应的机制是如何实现的呢?我们就来分析一下,顺便介绍一下 asyncio.Queue 的源码,来让这个系列的文章更有深度一些。并且这些同步原语在等待和唤醒机制的实现上是相似的,背后的思想都是一样的,了解了 Queue 是怎么做的,其它的也就都会了。

这里我们直接将 Queue 的源码贴出来进行讲解,为了方便解释,这里将内部方法的顺序做了一些调整。

class Queue:

    def __init__(self, maxsize=0, *, loop=None):
        # asyncio 的这些同步原语都支持 loop 参数,用于指定事件循环
        # 如果没指定,那么采用当前所在的事件循环,但从 3.10 开始这个参数就被移除了
        if loop is None:
            self._loop = events.get_event_loop()
        else:
            self._loop = loop
            warnings.warn("The loop argument is deprecated since Python 3.8, "
                          "and scheduled for removal in Python 3.10.",
                          DeprecationWarning, stacklevel=2)
        # 队列允许设置一个最大值,如果 maxsize 小于等于 0,那么表示队列的容量是无限的
        self._maxsize = maxsize

        # 两个双端队列, 用于存储 Future 对象,一会说
        self._getters = collections.deque()
        self._putters = collections.deque()
        # 未完成的任务数,每当往队列添加一个任务,该值就会自增 1
        self._unfinished_tasks = 0
        # 事件,当调用 await queue.join() 时,内部会执行 await self._finished.wait()
        self._finished = locks.Event(loop=loop)
        # 设置标志位,初始状态为 True
        self._finished.set()
        # 以 maxsize 为参数,调用 _init 方法
        self._init(maxsize)

    def _init(self, maxsize):
        # 初始化一个双端队列,该队列是真正用来存储元素的
        # 所以 asyncio.Queue 在存储元素上,是直接使用的 dequeue
        # 而且这个 maxsize 参数并没有用上
        self._queue = collections.deque()

    def _get(self):
        # 获取元素的具体实现,从双端队列的开头弹出一个元素
        return self._queue.popleft()

    def _put(self, item):
        # 添加元素的具体实现,往双端队列的尾部添加一个元素        
        self._queue.append(item)

    def qsize(self):
        # 获取内部的元素个数
        return len(self._queue)

    @property
    def maxsize(self):
        # 获取队列允许存储的最大元素个数,小于等于 0 则表示没有限制
        return self._maxsize
    
    def empty(self):
        # 队列是否为空
        return not self._queue
    
    def full(self):
        # 队列是否已满
        # 如果 self._maxsize <= 0,表示队列容量无限,那么永远不可能满,直接返回 False
        if self._maxsize <= 0:
            return False
        # 否则就看当前的元素个数是否超过了 self._maxsize
        else:
            return self.qsize() >= self._maxsize    
       
    async def put(self, item):
        # 将一个元素添加到队列中,如果队列已满(要求 self._maxsize 必须大于 0)
        while self.full():
            # 那么创建一个 future,并将其添加到 self._putters 中
            putter = self._loop.create_future()
            self._putters.append(putter)
            try:
                # 此处会发生阻塞
                await putter
            except:
                putter.cancel()  # 出现异常了,将 future 取消掉
                try:
                    # 并从双端队列中移除
                    self._putters.remove(putter)
                except ValueError:
                    # 出现异常,则说明该 future 之前就已经被移除了
                    pass     
                if not self.full() and not putter.cancelled():
                    self._wakeup_next(self._putters)
                raise
        # 到这里说明队列有剩余空间,那么通过 self.put_nowait 添加进去
        return self.put_nowait(item)
    
    def put_nowait(self, item):
        # 往队列里面添加一个元素(不阻塞),如果队列已满,那么立即返回 QueueFull 异常
        if self.full():
            raise QueueFull
        # 调用 self._put,往双端队列里面添加元素
        self._put(item)
        # 未完成任务 +1
        self._unfinished_tasks += 1
        # 清空标志位,因为队列有元素了,那么 await queue.join() 要处于阻塞
        self._finished.clear()
        # 队列满时调用 put 会创建一个 future 添加到 self._putters 中
        # 队列空时调用 get 会创建一个 future 添加到 self._getters 中
        # 现在队列里面已经有元素,那么是不是应该让调用 get 处于阻塞的任务解除阻塞呢
        # 该方法会从 self._getters 的头部弹出一个 future,让它对应的任务解除阻塞
        self._wakeup_next(self._getters)    
    
    def _wakeup_next(self, waiters):
        # 如果 waiters 里面有处于等待的 future
        while waiters:
            # 那么从头部弹出一个
            waiter = waiters.popleft()
            # 调用 set_result(None),那么该 future 对应的任务会解除阻塞
            if not waiter.done():
                waiter.set_result(None)
                break

    async def get(self):
        # 从队列里面获取元素,如果队列为空
        while self.empty():
            # 创建一个 future 并添加到 self._getter 中
            getter = self._loop.create_future()
            self._getters.append(getter)
            try:
                # 此处会陷入阻塞,什么时候解除阻塞呢?显然上面已经给出答案了
                # 当有人往队列里面添加元素,并从 self._getters 里面它的 future 并调用 set_result 的时候
                await getter
            except:
                getter.cancel() 
                try:
                    self._getters.remove(getter)
                except ValueError:
                    pass
                if not self.empty() and not getter.cancelled():
                    self._wakeup_next(self._getters)
                raise
        # 到这里说明队列有元素,那么通过 self.get_nowait 获取
        return self.get_nowait()

    def get_nowait(self):
        # 从队列里面获取一个元素(不阻塞),如果队列为空,那么立即返回 QueueEmpty 异常
        if self.empty():
            raise QueueEmpty
        # 调用 self._get,从双端队列里面获取元素
        item = self._get()
        # 队列的元素被取走了,证明有位置了,那么调用 put 阻塞的任务是不是也应该解除阻塞呢?
        self._wakeup_next(self._putters)
        # 返回元素
        # 所以要注意:get() 和 put() 是协程方法,而 get_nowait() 和 put_nowait() 是普通方法
        # 我们使用协程方法即可,当然,如果你能确定队列不为空或者未满,那么调用 *_nowait 普通方法也可以
        return item

    def task_done(self):
        # 调用 task_done,内部的 _unfinished_tasks 为自减 1
        if self._unfinished_tasks <= 0:
            raise ValueError('task_done() called too many times')
        self._unfinished_tasks -= 1
        # 如果已经为 0 了,那么要设置标志位,对应 await queue.join() 解除阻塞
        # 所以,如果需要使用队列的 join 机制,那么别忘记在取走元素时调用 task_done 方法
        if self._unfinished_tasks == 0:
            self._finished.set()

    async def join(self):
        # 如果 _unfinished_tasks 大于 0,那么 await self._finished.wait()
        # 否则解除阻塞,直接返回
        if self._unfinished_tasks > 0:
            await self._finished.wait()

    def __repr__(self):
        return f'<{type(self).__name__} at {id(self):#x} {self._format()}>'

    def __str__(self):
        return f'<{type(self).__name__} {self._format()}>'

    def _format(self):
        result = f'maxsize={self._maxsize!r}'
        if getattr(self, '_queue', None):
            result += f' _queue={list(self._queue)!r}'
        if self._getters:
            result += f' _getters[{len(self._getters)}]'
        if self._putters:
            result += f' _putters[{len(self._putters)}]'
        if self._unfinished_tasks:
            result += f' tasks={self._unfinished_tasks}'
        return result            

以上就是 asyncio.Queue 的源码部分,可以说非常简单,其它的同步原语也是类似的实现机制。

异步队列的实际使用场景

现在我们了解了异步队列如何工作的基础知识,并阅读了源码,但由于我们通常不会在日常工作中构建超市模型,所以让我们看一些现实中的场景,看看如何将其应用到实际应用程序中。

Web 应用程序中的队列

当有一个可以在后台运行的潜在耗时操作的时候,队列在 Web 应用程序中将很有帮助。如果在 Web 请求的主协程中运行此操作,将阻止对用户的响应(直到操作完成),这可能会给最终用户留下一个缓慢、无响应的页面,降低用户的使用体验。

设想我们就职于一家电子商务公司,并使用缓慢的订单管理系统进行操作。处理订单可能需要几秒钟的时间,但我们不想让用户在下单时进行等待。此外,订单管理系统不能很好地处理负载,所以我们想限制同时向它发出的请求数量。这种情况下,队列可以解决这两个问题。正如之前看到的,在添加更多块或抛出异常之前,队列可以拥有允许的最大元素数量。这为并发性提供了天然限制。

队列还解决了用户等待响应时间过长的问题,将元素放到队列中是立即发生的,这意味着可通知用户他们的订单已经被接收了,从而提供快捷的用户体验。当然在现实世界中,这可能导致后台任务在没有通知用户的情况下失败,因此需要某种形式的数据持久性和逻辑来应对这种情况。

为验证这一点,我们 FastAPI 创建一个简单的 Web 应用程序,它使用一个队列来运行后台任务,这里通过 asyncio.sleep 来模拟与慢速订单管理系统的交互。在现实世界的微服务体系结构中,你可能通过 aiohttp(或类似的库)调用一个 REST API 进行通信,但为了简单,这里就使用 sleep。

Web 应用所做的事情如下:FastAPI 启动之后通过 hook 创建一个队列和一组 worker 任务,这些任务负责与慢速订单服务交互。然后创建一个 HTTP POST 订单端点,它将在队列上放置一个订单,一旦将订单放入队列中,将返回一个 HTTP 200 和一条消息,表明已经完成下单。此外还将在 FastAPI 的关闭 hook 中添加一些安全的关闭逻辑,因为当应用程序关闭时可能仍有一些订单正在被处理。在关闭 hook 中,将等到所有忙碌的 worker 完成它们的工作。

import asyncio
from asyncio import Queue
from functools import partial
from random import randint
from fastapi import FastAPI, APIRouter
from fastapi.requests import Request
from fastapi.responses import Response
import uvicorn

router = APIRouter()

async def process_order_worker(worker_id: int, order_queue: Queue):
    # 从队列中获取订单,然后处理它
    while True:
        print(f"Worker {worker_id}, 等待订单")
        order = await order_queue.get()
        print(f"Worker {worker_id}, 开始处理订单")
        await asyncio.sleep(order)
        print(f"Worker {worker_id}, 订单处理完毕")
        order_queue.task_done()

@router.post("/order")
async def place_order(request: Request):
    order_queue: Queue = request.app.state.order_queue
    # 将订单放入队列,并立即响应用户,这里用一个 sleep 时长来模拟订单处理
    await order_queue.put(randint(1, 4))
    return Response("订单已被接收, 等待后台处理", media_type="text/plain")

async def create_order_queue(app: FastAPI):
    # 程序启动时执行的钩子函数
    print("创建队列和任务")
    # 创建一个最多容纳 10 个元素的队列,并创建 4 个 worker
    order_queue: Queue = asyncio.Queue(10)
    app.state.order_queue = order_queue
    # 程序启动之后会创建 4 个 worker,然后不断地从队列中取出订单,进行处理
    app.state.order_tasks = [asyncio.create_task(process_order_worker(i, order_queue))
                             for i in range(1, 5)]

async def destroy_order_queue(app: FastAPI):
    # 程序结束时执行的钩子函数
    order_tasks = app.state.order_tasks
    # 等待所有繁忙的任务完成
    order_queue: Queue = app.state.order_queue
    print("等待尚未完成的任务执行完毕, 但只有 10 秒的机会")
    try:
        # 10 秒内未完成,那么程序直接结束,不再等待
        await asyncio.wait_for(order_queue.join(), timeout=10)
    except asyncio.TimeoutError:
        print("程序结束, 但还有任务尚未完成, 这里直接取消")
    finally:
        [task.cancel() for task in order_tasks]


app = FastAPI()
app.include_router(router)
app.router.on_startup.append(partial(create_order_queue, app))
app.router.on_shutdown.append(partial(destroy_order_queue, app))

if __name__ == '__main__':
    uvicorn.run(app, port=9999)

在代码中首先创建了一个 process_order_worker 协程,这会从队列中获得一个项目,在本例中是一个整数,然后休眠一段时间,以模拟使用缓慢的订单管理系统。这个协程将永远循环,不断从队列中取出项目并处理它们。

然后分别创建用于设置和删除队列的协程 create_order_queue 和 destroy_order_queue。创建队列很简单,因为我们创建了一个最多 10 个元素的 asyncio 队列,并创建了 4 个 worker 任务,将它们存储在 FastAPI 实例中。

销毁队列有点复杂,首先使用 order_queue.join 等待队列完成对其所有元素的处理。由于应用程序正在关闭,它将不再提供任何 HTTP 请求,因此没有其他订单可以进入队列。这意味着任何已经在队列中的内容都将由一个 worker 处理,并且 worker 当前正在处理的所有任务也将完成。然后还将 join 包装在 wait_for 中,超时时间为 10 秒。这是一个好主意,因为我们不希望一个失控的任务长时间阻止应用程序关闭。

最后创建一个 POST 端点 /order,该端点创建一个随机延迟,并将其添加到队列中。将订单添加到队列后,会使用 HTTP 200 状态代码和一条短消息来响应用户。注意,我们在添加订单的时候直接调用队列的 put 方法,这意味着如果队列已满,则请求将被阻塞,直到消息位于队列中(这可能需要一些时间)。如果你希望立即返回,那么可以使用 put_nowait,然后以 HTTP 500 错误或其他错误代码进行响应,表示当前队列已满、无法再添加订单,要求调用者稍后再试。因此这里就需要对一些耗时的请求进行权衡,是一直等待直到返回 200,还是在无法返回 200 的时候立即返回 500,则取决于服务需求。

总之在 Web 应用程序中使用异步队列时,要牢记的一件事是队列的故障模式。如果一个 API 实例由于某种原因(如内存不足)崩溃了,或者我们需要重新启动服务器以重新部署应用程序,该怎么办?这种情况下,将丢失队列中所有未处理的订单,因为它们只存储在内存中。有时,在队列中丢失信息并不是什么大问题,但对于客户订单,后果可能比较严重。

另外 Python asyncio 的异步队列没有提供任务持久性或队列持久性的概念,如果希望队列中的任务对这些类型的故障具有鲁棒性,需要在某处引入一种方法(如使用数据库或缓存)来保存任务。然而,更好的选择是在支持任务持久性的 asyncio 之外使用单独的队列,而 Celery 和 RabbitMQ 是可以将任务队列持久保存到磁盘的两种方式。

当然,使用单独的队列架构会增加复杂性,在具有持久任务的持久队列中,它还带来了需要持久存储到磁盘的性能挑战。要为应用程序确定最佳架构,你需要仔细考虑仅在内存中使用 asyncio 队列与使用单独的架构组件之间的关系。

启动服务之后,会打印如下输出,然后我们可以通过 post 请求添加任务,worker 会立即处理。比如我们调用 6 次 /order:

此时订单已经开始处理了,如果处理速度非常快,那么在代码中直接使用 put 是没有问题的,因为队列始终有空余位置。但如果速度没那么快,可以考虑增大 worker 数量或队列的容量。

网络爬虫队列

让我们构建一个爬虫,首先创建一个无边界队列(如果你担心内存溢出,也可以进行边界设定)来保存 URL 以供下载,然后 worker 将从队列中提取 URL 并使用 aiohttp 下载它们。下载后将使用流行的 HTML 解器 PyQuery 来提取链接,从而放回队列中。

import asyncio
from asyncio import Queue
from aiohttp import ClientSession
from pyquery import PyQuery

class WorkItem:

    def __init__(self, url: str, item_depth: int):
        """
        我们不想扫描整个互联网,所以只扫描 root 页面之外的一组页面
        我们将其称为最大深度, 即 max_depth
        如果最大深度设置为 3,这意味着将只获取 root 页面以下 3 层的链接
        而 item_depth 表示当前所在的深度
        :param url:
        :param item_depth:
        """
        self.url = url
        self.item_depth = item_depth

async def worker(worker_id: int, queue: Queue, session: ClientSession, max_depth):
    print(f"Worker {worker_id} 开始工作")
    while True:
        # 从队列中获取一个 URL 进行处理,然后 开始下载它
        work_item: WorkItem = await queue.get()
        print(f"Worker {worker_id} 开始处理 {work_item.url}")
        await process_page(work_item, queue, session, max_depth)
        print(f"Worker {worker_id} 已处理 {work_item.url} 完毕")
        queue.task_done()

async def process_page(work_item: WorkItem, queue: Queue, session: ClientSession, max_depth: int):
    # 下载 URL 页面,并解析页面中的所有链接,将它们放回队列中。
    try:
        response = await asyncio.wait_for(session.get(work_item.url), timeout=3)
        if work_item.item_depth == max_depth:
            print(f"{work_item.url} 已达到最大深度")
        else:
            html = await response.text(encoding="utf-8")
            p = PyQuery(html)
            links = p.find("a")
            for link in links.items():
                queue.put_nowait(WorkItem(link.attr("href"), work_item.item_depth + 1))
    except Exception as e:
        print(f"处理 {work_item.url} 时遇见错误: {e}")

async def main():
    start_url = "https://www.example.com"
    url_queue = asyncio.Queue()
    url_queue.put_nowait(WorkItem(start_url, 0))
    async with ClientSession() as session:
        workers = [asyncio.create_task(worker(i, url_queue, session, 3)) for i in range(1, 101)]
        await url_queue.join()
        [w.cancel() for w in workers]

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

在代码中,我们首先定义了一个 WorkItem 类,这是一个简单的数据类用于保存 URL 和该 URL 的深度。然后定义 worker,它从队列中获取一个 WorkItem 并调用 process_page。如果可以的话,process_page 协程函数会下载 URL 的内容(可能会发生超时或异常,我们只是记录并忽略这些错误)。然后使用 PyQuery 获取所有链接,并将它们添加回队列,供其他 worker 处理。

在主协程中创建队列并使用第一个 WorkItem 引导它,在这个例子中,对 example.com 进行硬编码,因为它是 root 页面,所以深度为 0。然后创建一个 aiohttp 会话,并创建 100 个 worker,这意味着可同时下载 100 个 URL,我们将其最大深度设置为 3。然后等待队列被清空,并且所有 worker 都使用 Queue.join 完成相关工作。完成队列处理后,取消所有 worker 任务。当你运行此代码时,应该会看到 100 个 worker 任务启动,并开始从它下载的每个 URL 中查找链接。

现在已经通过构建一个虚拟的超市结账队伍,以及构建一个订单管理 API 和一个网络爬虫来了解异步队列的基础知识,到目前为止,worker 对队列中的每个元素都给予了同等的优先级,就是把排在最前面的元素取出并进行处理。如果希望一些任务尽快发生,即使它们排在队列的后面,该怎么办呢?让我们看一下优先级队列,看看如何实现这一点。

优先级队列

之前的队列示例以 FIFO(先进先出)的顺序处理队列中的元素,谁先排队,谁先被处理。这在许多情况下都很有效,无论是在软件工程中还是在生活中。

然而在某些应用程序中,将所有任务视为平等关系则未必是可取的。假设我们正在构建一个数据处理管道,其中每个任务都是一个长时间运行的查询,可能需要运行几分钟。假设两个任务大致同时运行,第一个任务是低优先级的数据查询,而第二个任务是关键的数据更新,应尽快处理。使用简单的队列,将首先处理第一个任务,让第二个更重要的任务等待直到第一个任务完成。如果第一个任务需要几个小时,或者如果所有 worker 都很忙,则第二个任务可能会等待很长时间。

为了解决解决这个问题,可使用优先队列,让 worker 首先处理最重要的任务。在后台,优先队列由堆(使用 heapq 模块)支持,而不是像简单队列那样使用 Python 的双端队列。为了创建一个 asyncio 优先队列,我们需要创建一个 asyncio.PriorityOueue 实例。

我们不会在这里过多讨论数据结构的细节,但堆是一棵二叉树,其属性是每个父节点的值都小于其所有子节点的值(小根堆,如果父节点大于子节点则是大根堆)。这与通常用于排序和搜索的二叉搜索树不同,二叉搜索树的唯一属性是节点的左侧子节点小于其父节点,节点的右侧子节点大于其父节点。堆的特性是最顶部的节点总是树中的最小元素,如果总是让最小节点成为最高优先级的节点,那么高优先级的节点将永远是队列中的第一个。

放入队列中的工作项不太可能是纯整数,因此我们需要某种方法来构造具有合理优先级规则的工作项。一种方法是使用元组,其中第一个元素是表示优先级的整数,第二个元素是任何类型的任务数据。默认队列实现通过查找元组的第一个值来决定优先级,其中最低的数字具有最高的优先级。让我们看一个将元组作为工作项的示例,以了解优先队列如何工作。

import asyncio
from asyncio import Queue, PriorityQueue
from typing import Tuple

async def worker(queue: Queue):
    while not queue.empty():
        worker_item: Tuple[int, str] = await queue.get()
        print(f"处理数据项: {worker_item}")
        queue.task_done()

async def main():
    # PriorityQueue 继承自 Queue
    priority_queue = PriorityQueue()
    work_items = [(3, "低优先级"), (1, "高优先级"), (2, "中优先级")]
    worker_task = asyncio.create_task(worker(priority_queue))
    for work_item in work_items:
        priority_queue.put_nowait(work_item)

    await asyncio.gather(priority_queue.join(), worker_task)

asyncio.run(main())
"""
处理数据项: (1, '高优先级')
处理数据项: (2, '中优先级')
处理数据项: (3, '低优先级')
"""

在代码中,我们创建了三个工作项,分别具有高优先级、中等优先级和低优先级。然后乱序将它们添加到优先级队列中,首先插入最低优先级的项目。在普通队列中,这意味着将首先处理优先级最低的项目,但对于优先队列则不是。输出结果表明是按优先级顺序处理工作项,而不是按它们插入队列的顺序进行处理。元组适用于简单的情况,但如果工作项中有大量数据,元组可能变得混乱和难以控制。有没有办法创建一个可按我们想要的顺序处理堆的类?事实上,可以实现这一点,最简单的方法是使用数据类(如果不能选择数据类,也可实现 __lt__、__le__、__gt__、__ge__)。

import asyncio
from asyncio import Queue, PriorityQueue
from dataclasses import dataclass, field

@dataclass(order=True)
class WorkItem:
    priority: int
    data: str = field(compare=False)

async def worker(queue: Queue):
    while not queue.empty():
        worker_item: WorkItem = await queue.get()
        print(f"处理数据项: {worker_item}")
        queue.task_done()

async def main():
    priority_queue = PriorityQueue()
    work_items = [WorkItem(3, "低优先级"), WorkItem(1, "高优先级"), WorkItem(2, "中优先级")]
    worker_task = asyncio.create_task(worker(priority_queue))
    for work_item in work_items:
        priority_queue.put_nowait(work_item)

    await asyncio.gather(priority_queue.join(), worker_task)

asyncio.run(main())
"""
处理数据项: WorkItem(priority=1, data='高优先级')
处理数据项: WorkItem(priority=2, data='中优先级')
处理数据项: WorkItem(priority=3, data='低优先级')
"""

在代码中我们创建了一个 dataclass,并将 order 设置为True。然后添加一个优先级整数和一个字符串数据字段,并将字符串字段从比较中排除。这意味着将这些工作项添加到队列中时,只会根据优先级字段进行排序。运行上面的代码,将按正确顺序进行处理。

现在已经了解了优先级队列的基础知识,让我们回到前面的订单管理 API 示例。假设有一些超级用户,他们在电子商务网站上消费额较高。我们希望确保他们的订单总是首先得到处理,从而确保提供最佳体验,那么此时便可为这些用户使用优先级队列。

另外优先级队列的一个有趣的极端情况是,当同时将两个具有相同优先级的工作项添加到队列中时,会发生什么呢。按插入顺序被 worker 处理吗?让我们用一个简单例子来验证一下。

import asyncio
from asyncio import Queue, PriorityQueue
from dataclasses import dataclass, field

@dataclass(order=True)
class WorkItem:
    priority: int
    data: str = field(compare=False)

async def worker(queue: Queue):
    while not queue.empty():
        worker_item: WorkItem = await queue.get()
        print(f"处理数据项: {worker_item}")
        queue.task_done()

async def main():
    priority_queue = PriorityQueue()
    work_items = [WorkItem(3, "低优先级(1)"), WorkItem(3, "低优先级(2)"),
                  WorkItem(2, "中优先级(1)"), WorkItem(3, "低优先级(3)"),
                  WorkItem(1, "高优先级"), WorkItem(2, "中优先级(2)")]
    worker_task = asyncio.create_task(worker(priority_queue))
    for work_item in work_items:
        priority_queue.put_nowait(work_item)

    await asyncio.gather(priority_queue.join(), worker_task)

asyncio.run(main())
"""
处理数据项: WorkItem(priority=1, data='高优先级')
处理数据项: WorkItem(priority=2, data='中优先级(2)')
处理数据项: WorkItem(priority=2, data='中优先级(1)')
处理数据项: WorkItem(priority=3, data='低优先级(1)')
处理数据项: WorkItem(priority=3, data='低优先级(3)')
处理数据项: WorkItem(priority=3, data='低优先级(2)')
"""

事实证明,相同优先级的数据项被插入时,顺序并不是添加的顺序,发生这种情况是因为底层堆排序算法不是稳定的排序算法。如果你希望当优先级相同时,按照添加顺序进行处理,那么可以给数据项再加入一个索引字段,这样当优先级相同时会按照索引进行排序。

import asyncio
from asyncio import Queue, PriorityQueue
from dataclasses import dataclass, field

@dataclass(order=True)
class WorkItem:
    priority: int
    index: int
    data: str = field(compare=False)

async def worker(queue: Queue):
    while not queue.empty():
        worker_item: WorkItem = await queue.get()
        print(f"处理数据项: {worker_item}")
        queue.task_done()

async def main():
    priority_queue = PriorityQueue()
    work_items = [WorkItem(3, 0, "低优先级(1)"), WorkItem(3, 1, "低优先级(2)"),
                  WorkItem(2, 0, "中优先级(1)"), WorkItem(3, 2, "低优先级(3)"),
                  WorkItem(1, 0, "高优先级"), WorkItem(2, 1, "中优先级(2)")]
    worker_task = asyncio.create_task(worker(priority_queue))
    for work_item in work_items:
        priority_queue.put_nowait(work_item)

    await asyncio.gather(priority_queue.join(), worker_task)

asyncio.run(main())
"""
处理数据项: WorkItem(priority=1, index=0, data='高优先级')
处理数据项: WorkItem(priority=2, index=0, data='中优先级(1)')
处理数据项: WorkItem(priority=2, index=1, data='中优先级(2)')
处理数据项: WorkItem(priority=3, index=0, data='低优先级(1)')
处理数据项: WorkItem(priority=3, index=1, data='低优先级(2)')
处理数据项: WorkItem(priority=3, index=2, data='低优先级(3)')
"""

Python 底层在排序的时候也是这么做的,现在我们已经了解了如何以 FIFO 队列顺序和优先队列顺序处理工作项。如果想首先处理最近添加的工作项怎么办?接下来,让我们看看如何使用 LIFO 队列来执行此操作。

LIFO 队列

LIFO 队列在计算机科学领域更常被称为堆栈(说白了就是栈),可将它们想象成一堆扑克筹码:当你下注时,你从筹码顶部取出筹码(或 pop它们),当你赢得新的筹码时,你将新的筹码放回筹码堆顶部(或 push 它们)。如果希望 worker 首先处理最近添加的项目可使用这种队列。

import asyncio
from asyncio import Queue, LifoQueue
from dataclasses import dataclass, field

@dataclass(order=True)
class WorkItem:
    priority: int
    index: int
    data: str = field(compare=False)

async def worker(queue: Queue):
    while not queue.empty():
        worker_item: WorkItem = await queue.get()
        print(f"处理数据项: {worker_item}")
        queue.task_done()

async def main():
    lifo_queue = LifoQueue()
    work_items = [WorkItem(3, 0, "低优先级(1)"), WorkItem(3, 1, "低优先级(2)"),
                  WorkItem(3, 2, "低优先级(3)"), WorkItem(1, 0, "高优先级"),
                  WorkItem(2, 0, "中优先级")]
    worker_task = asyncio.create_task(worker(lifo_queue))
    for work_item in work_items:
        lifo_queue.put_nowait(work_item)

    await asyncio.gather(lifo_queue.join(), worker_task)

asyncio.run(main())
"""
处理数据项: WorkItem(priority=2, index=0, data='中优先级')
处理数据项: WorkItem(priority=1, index=0, data='高优先级')
处理数据项: WorkItem(priority=3, index=2, data='低优先级(3)')
处理数据项: WorkItem(priority=3, index=1, data='低优先级(2)')
处理数据项: WorkItem(priority=3, index=0, data='低优先级(1)')
"""

数据类的前两个字段已经没用了,说白了就是按照先入后出的顺序来处理的,把它理解为栈即可。

代码非常简单。

小结

在本篇文章中,我们学习了以下内容:

  • asyncio 队列是任务队列,在工作流中很有用,有生成数据的协程和负责处理该数据的协程。
  • 队列将数据生成与数据处理进行分离,因为我们可以让 producer 将项目放入队列中,然后多个 worker 同时处理。
  • 可使用优先级队列为某些任务调整优先级,这对于那些更重要的工作以及需要优先处理的工作来说非常重要。
  • asyncio 队列不是分布式的,也不是持久的。如果你需要那些功能的支持,则需要寻找一个单独的架构组件,如 Celery 或 RabbitMQ。
posted @ 2023-05-13 23:54  古明地盆  阅读(1216)  评论(0编辑  收藏  举报