异步Rust——构建我们自己的异步队列

418 阅读31分钟

虽然我们已经探讨了基本的异步语法,并通过高层次的异步概念解决了一个问题,但你可能仍然不完全理解任务和未来(futures)到底是什么,以及它们如何在异步运行时中流动。描述 futures 和 tasks 可能很困难,而且它们也不容易理解。本章将通过让你构建自己的异步队列并最小化依赖关系,来巩固你迄今为止对 futures 和 tasks 的理解,以及它们如何在异步运行时中运行。

这个异步运行时将是可定制的,允许你选择队列的数量和处理这些队列的消费线程的数量。实现方式不必统一。例如,我们可以有一个低优先级队列,配有两个消费线程,并有一个高优先级队列,配有五个消费线程。然后,我们可以选择一个 future 将在哪个队列中处理。我们还将能够实现任务偷取(task stealing),即消费线程可以从其他队列中偷取任务,如果它们自己的队列为空。最后,我们将构建自己的宏,以便高层次地使用我们的异步运行时。

在本章结束时,你将能够实现自定义异步队列,并完全理解 futures 和 tasks 如何在异步运行时中流动。你还将具备定制异步运行时的技能,以解决那些标准的开箱即用的运行时环境可能无法处理的特定问题。即使你以后不想再实现自己的异步队列,你也将对异步运行时有更深入的理解,从而更有效地操作高层次的异步 crate 来解决问题。你还将理解即使在实现高层次异步代码时,异步代码的权衡。

我们通过定义任务如何生成来开始构建异步队列的探索,因为任务的生成是进入运行时的入口点。这个异步运行时将是可定制的,允许你选择有多少队列以及多少消费线程来处理这些队列。

构建我们自己的异步队列

在本节中,我们将通过构建自定义的异步队列的过程,来展示如何将 futures 转换为 tasks 并执行它们。如果我们将实现过程分解成几个步骤,我们将直观地了解 futures 如何转化为 tasks 并被执行。

在这个例子中,我们构建一个简单的异步队列,并将处理三个任务。我们将在定义每个任务时详细描述它。

在编写代码之前,我们需要以下依赖项:

[dependencies]
async-task = "4.4.0"
futures-lite = "1.12.0"
flume = "0.10.14"

我们使用的这些依赖项是:

  • async-task
    这个 crate 对于在异步运行时中生成和管理任务是必不可少的。它提供了将 futures 转换为 tasks 所需的核心功能。
  • futures-lite
    这是一个轻量级的 futures 实现。
  • flume
    这是一个多生产者、多消费者的通道,我们将用它来实现我们的异步队列,从而安全地在运行时中传递任务。我们本可以使用 async-channel,但选择了 flume,因为我们希望能够克隆接收器,因为我们将任务分配给多个消费者。此外,flume 提供了无界通道,能够容纳无限数量的消息,并实现了无锁算法。这使得 flume 在高度并发的程序中尤其有用,因为队列可能需要并行处理大量消息,而不像标准库中的通道依赖于阻塞的互斥锁进行同步。

接下来,我们需要在 main.rs 文件中导入以下内容:

use std::{future::Future, panic::catch_unwind, thread};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use std::sync::LazyLock;

use async_task::{Runnable, Task};
use futures_lite::future;

我们将在本章中使用这些导入,并会在使用时为你解释它们的上下文。

每个任务都需要能够被传递到队列中。我们应该从构建任务生成函数开始。这是我们传递 future 到函数中的地方。函数将把 future 转换为任务,并将任务放入队列中以便执行。此时,函数可能看起来比较复杂,所以让我们从这个签名开始:

fn spawn_task<F, T>(future: F) -> Task<T>
where
    F: Future<Output = T> + Send + 'static,
    T: Send + 'static,
{
    . . .
}

这是一个泛型函数,接受实现了 FutureSend 特征的任何类型。这样做的原因是,我们不想将函数限制为只能发送一种类型的 future。Future 特征表明我们的 future 将导致一个错误或值 T。我们的 future 需要 Send 特征,因为我们将把 future 发送到一个不同的线程,该线程是基于队列的。Send 特征强制执行一些约束,确保我们的 future 可以安全地在线程之间共享。

'static 表示我们的 future 不包含任何生命周期比静态生命周期更短的引用。因此,future 可以在程序运行期间使用。确保这个生命周期非常重要,因为我们不能强迫程序员等待一个任务完成。如果开发者永远不等待一个任务,任务可能会在整个程序生命周期内运行。由于我们无法保证任务何时完成,我们必须确保任务的生命周期是静态的。当浏览异步代码时,你可能会看到 async move 被使用。它是将异步闭包中使用的变量的所有权转移到任务上,从而确保生命周期是静态的。

现在我们已经定义了 spawn_task 函数的签名,接下来我们进入函数中的第一块代码,它定义了任务队列:

static QUEUE: LazyLock<flume::Sender<Runnable>> = LazyLock::new(|| {
    . . .
});

通过 static,我们确保队列在程序的整个生命周期内都存在。这是合理的,因为我们希望在程序运行的过程中不断地向队列发送任务。LazyLock 结构体在第一次访问时会被初始化,一旦初始化完成,就不会再进行初始化。因为我们每次调用任务生成函数时都会发送一个 future 到异步运行时。如果我们每次调用 spawn_task 都初始化队列,我们将清空队列中先前的任务。现在,我们有了一个通道的发送端,用于发送 Runnable

Runnable 是一个可运行任务的句柄。每个生成的任务都有一个唯一的 Runnable 句柄,只有在任务被安排运行时,句柄才会存在。该句柄具有 run 函数,用于轮询任务的 future 一次。然后,Runnable 被丢弃。只有当 waker 唤醒任务时,任务才会重新被安排运行。回想一下第二章,如果我们没有将 waker 传递给 future,它将不会再次被轮询。这是因为 future 不能被唤醒以便再次轮询。我们可以构建一个异步运行时来轮询 futures,无论是否存在 waker,我们将在第十章中探讨这一点。

现在我们已经定义了队列的签名,我们可以看看传递给 LazyLock 的闭包。我们需要创建我们的通道以及接收传递给该通道的 futures 的机制:

let (tx, rx) = flume::unbounded::<Runnable>();

thread::spawn(move || {
    while let Ok(runnable) = rx.recv() {
        println!("runnable accepted");
        let _ = catch_unwind(|| runnable.run());
    }
});
tx

在创建通道后,我们启动一个线程,该线程等待传入的任务流量。由于我们正在构建异步队列来处理传入的异步任务,因此等待传入的流量是阻塞的。我们不能在这个线程中依赖异步。当接收到 Runnable 后,我们使用 catch_unwind 函数运行它。我们使用 catch_unwind 是因为我们无法保证传递给异步运行时的代码质量。理想情况下,所有 Rust 开发者都会正确处理可能的错误,但万一他们没有这么做,catch_unwind 会在代码运行时捕捉抛出的任何错误,并根据结果返回 OkErr。这样可以防止编写不当的 future 破坏我们的异步运行时。然后我们返回传输通道,这样我们就可以将 Runnable 发送到线程中进行处理。

现在我们有一个正在运行的线程,它等待任务被发送到该线程进行处理,我们通过以下代码实现这一点:

let schedule = |runnable| QUEUE.send(runnable).unwrap();
let (runnable, task) = async_task::spawn(future, schedule);

在这里,我们创建了一个闭包,接受一个 Runnable 并将其发送到我们的队列。然后,我们通过使用 async_taskspawn 函数来创建 RunnableTask。这个函数会调用一个不安全的函数,将 future 分配到堆上。由 spawn 函数返回的任务和 Runnable 都指向同一个 future。

注意

在本章中,我们不会构建自己的执行器或代码来创建 runnable 或调度任务。我们将在第十章中构建一个完全基于标准库、没有外部依赖的异步服务器时完成这部分内容。

现在,既然 runnabletask 都指向同一个 future,我们必须将 runnable 调度执行,并返回任务:

runnable.schedule();
println!("Here is the queue count: {:?}", QUEUE.len());
return task;

当我们调度 runnable 时,实际上是将任务放入队列等待处理。如果我们没有调度 runnable,任务将不会执行,而且当我们尝试阻塞主线程等待任务执行时,程序会崩溃(因为队列中没有 runnable,但我们仍然返回了任务)。记住,taskrunnable 都指向同一个 future

现在我们已经将 runnable 调度到队列中执行并返回了任务,我们的基本异步运行时就完成了。接下来,我们只需要构建一些基本的 futures。我们将有两种类型的任务。

第一个任务类型是我们的 CounterFuture 我们在第二章中最初探讨过这个任务。这个 future 会递增计数器,并在每次轮询后打印结果,通过 std::thread::sleep 来模拟延迟。代码如下:

struct CounterFuture {
    count: u32,
}

impl Future for CounterFuture {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.count += 1;
        println!("polling with result: {}", self.count);
        std::thread::sleep(Duration::from_secs(1));
        if self.count < 3 {
            cx.waker().wake_by_ref();
            Poll::Pending
        } else {
            Poll::Ready(self.count)
        }
    }
}

第二个任务(任务 3)是使用 async/await 语法创建的异步函数。 该函数会在打印消息前休眠 1 秒。代码如下:

async fn async_fn() {
    std::thread::sleep(Duration::from_secs(1));
    println!("async fn");
}

在这个例子中,我们并没有像在 CounterFuture 中那样手动编写轮询机制。相反,我们使用了 Rust 提供的内置异步功能,该功能会自动处理任务的轮询和调度。需要注意的是,我们在 async_fn 中的 sleep 是阻塞的,因为我们希望看到任务在队列中的处理方式。

在继续之前,我们可以稍微绕道,理解一下非阻塞的异步 sleep 函数是如何工作的。在本章中,我们使用的 sleep 函数会阻塞执行器。这样做是为了教学目的,让我们能够轻松地映射任务如何在运行时中被处理。然而,如果我们希望构建一个高效的异步 sleep 函数,我们需要让执行器来轮询我们的 sleep future,并在时间没有到达时返回 Pending。首先,我们需要 Instant 来计算经过的时间,并需要两个字段来跟踪休眠时间:

use std::time::Instant;

struct AsyncSleep {
    start_time: Instant,
    duration: Duration,
}

impl AsyncSleep {
    fn new(duration: Duration) -> Self {
        Self {
            start_time: Instant::now(),
            duration,
        }
    }
}

然后,我们可以在每次轮询时检查当前时间与 start_time 之间的差异,如果时间不够,就返回 Pending

impl Future for AsyncSleep {
    type Output = bool;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let elapsed_time = self.start_time.elapsed();
        if elapsed_time >= self.duration {
            Poll::Ready(true)
        } else {
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

这样就不会阻塞执行器,避免空闲的休眠时间占用资源。由于休眠只是一个过程的一部分,我们可以在异步块中调用 await 来等待我们的异步 sleep future,如下所示:

let async_sleep = AsyncSleep::new(Duration::from_secs(5));
let async_sleep_handle = spawn_task(async {
    async_sleep.await;
    . . .
});

注意

像编程中的大多数事物一样,永远存在权衡。如果有很多任务排在休眠任务之前,那么异步休眠任务在完成之前的等待时间可能会超过所需的持续时间,因为它可能必须等待其他任务完成才能在每次轮询之间完成。如果你有一个操作要求两个步骤之间间隔 x 秒,使用阻塞休眠可能是更好的选择,但如果你有很多这样的任务,队列会迅速堵塞。

回到我们之前的阻塞示例,我们现在可以使用以下主函数在我们的运行时中运行一些 futures:

fn main() {
    let one = CounterFuture { count: 0 };
    let two = CounterFuture { count: 0 };
    let t_one = spawn_task(one);
    let t_two = spawn_task(two);
    let t_three = spawn_task(async {
        async_fn().await;
        async_fn().await;
        async_fn().await;
        async_fn().await;
    });
    std::thread::sleep(Duration::from_secs(5));
    println!("before the block");
    future::block_on(t_one);
    future::block_on(t_two);
    future::block_on(t_three);
}

这个主函数有一些重复,但这是必要的,以便我们能够感受到我们刚刚构建的异步运行时如何处理 futures。注意,任务 3 包含多次调用 async_fn。这帮助我们看到运行时如何在单个任务中处理多个异步操作。接着,我们等待 5 秒并打印输出,以便在调用 block_on 函数之前能够感知系统的运行方式。

运行我们的程序,会在终端打印出以下冗长但至关重要的输出:

Here is the queue count: 1
Here is the queue count: 2
Here is the queue count: 3
runnable accepted
polling with result: 1
runnable accepted
polling with result: 1
runnable accepted
async fn
async fn
before the block
async fn
async fn
runnable accepted
polling with result: 2
runnable accepted
polling with result: 2
runnable accepted
polling with result: 3
runnable accepted
polling with result: 3

我们的打印输出给出了异步运行时的时间线。我们可以看到,队列中填充了我们生成的三个任务,并且在调用 block_on 函数之前,运行时按顺序异步地处理它们。即使在调用第一个 block_on 函数之后,阻塞在我们生成的第一个任务上,另外两个计数任务也在同时被处理。

注意,我们在第三个任务中构建并调用了四次的异步函数本质上是阻塞的。即使我们使用了 await 语法,像这样:

async {
    async_fn().await;
    async_fn().await;
    async_fn().await;
    async_fn().await;
}

async_fn 中的调用堆栈会阻塞处理任务队列的线程,直到整个任务完成。当轮询结果是 Pending 时,任务会被放回队列等待再次轮询。

我们的异步运行时可以通过图 3-1 来总结。

image.png

让我们用一个类比来描述发生的事情。假设我们有一件需要清洗的脏外套。外套内标签上的清洗说明和材质内容就像是 future。我们走进干洗店,交给工作人员外套和说明。工作人员通过给外套套上塑料袋并分配一个编号来让外套变得“可运行”。工作人员还给了我们一张带编号的票,这就像是 main 函数得到的任务。

接着,我们可以去做其他事情,外套被清洗。如果外套第一次没有清洗干净,它会继续经过清洗周期,直到清洁完成。然后,我们带着票回到干洗店,将票交给工作人员。这就像是 block_on 函数的阶段。如果我们在很长时间后才回来,外套可能已经清洗干净,我们可以拿走它,继续我们的日常活动。如果我们太早去干洗店,外套还没有清洗干净,我们必须等到它清洗完成才能带走。干净的外套就是结果。

目前,我们的异步运行时只有一个线程在处理队列。这就像我们坚持只让干洗店有一个工作人员。这样的资源使用效率不是最高的,因为大多数 CPU 都有多个核心。考虑到这一点,增加工作线程和队列的数量来提高处理更多任务的能力是很有用的。

增加工作线程和队列

为了增加处理队列的线程数量,我们可以通过克隆接收器来添加另一个线程从队列中消费任务:

let (tx, rx) = flume::unbounded::<Runnable>();

let queue_one = rx.clone();
let queue_two = rx.clone();

thread::spawn(move || {
    while let Ok(runnable) = queue_one.recv() {
        let _ = catch_unwind(|| runnable.run());
    }
});
thread::spawn(move || {
    while let Ok(runnable) = queue_two.recv() {
        let _ = catch_unwind(|| runnable.run());
    }
});

如果我们通过通道发送任务,任务将通常在两个线程之间分配。如果一个线程被一个 CPU 密集型任务阻塞,另一个线程将继续处理其他任务。回想一下,第 1 章通过斐波那契数的例子证明了 CPU 密集型任务可以通过使用线程并行执行。我们还可以用以下代码以更符合人体工程学的方式构建线程池:

for _ in 0..3 {
    let receiver = rx.clone();
    thread::spawn(move || {
        while let Ok(runnable) = receiver.recv() {
            let _ = catch_unwind(|| runnable.run());
        }
    });
}

我们可以将 CPU 密集型任务卸载到线程池中,并继续处理程序的其他部分,在需要任务结果时再进行阻塞。虽然这不完全符合异步编程的精神(因为我们使用异步编程来优化 I/O 操作的调度),但这个方法可以提醒我们,某些问题可以通过在程序早期卸载 CPU 密集型任务来解决。

警告

你可能遇到过类似“异步编程不适用于计算密集型任务”的警告。异步编程只是一个机制,只要使用得当,你可以用它来做任何你需要做的事情。然而,这个警告并非毫无道理。例如,如果你使用异步运行时处理传入的请求,就像大多数 Web 框架一样,那么将计算密集型任务丢到异步运行时的队列中,可能会阻塞你处理传入请求的能力,直到这些计算任务完成。

现在我们已经探讨了多个工作线程,我们还需要进一步研究多个队列。

将任务传递到不同的队列

我们希望拥有多个队列的原因之一是,我们可能希望对任务进行不同的优先级划分。在本节中,我们将构建一个具有两个消费线程的高优先级队列和一个具有一个消费线程的低优先级队列。为了支持多个队列,我们需要以下 enum 来分类任务所要进入的队列类型:

#[derive(Debug, Clone, Copy)]
enum FutureType {
    High,
    Low
}

我们还需要让我们的 futures 在传递到 spawn 函数时能够返回未来的类型,这可以通过利用 trait 来实现:

trait FutureOrderLabel: Future {
    fn get_order(&self) -> FutureType;
}

接着,我们需要通过添加一个额外的字段来加入未来类型:

struct CounterFuture {
    count: u32,
    order: FutureType
}

我们的 poll 函数保持不变,因此不需要再次讨论。然而,我们需要为我们的 future 实现 FutureOrderLabel trait:

impl FutureOrderLabel for CounterFuture {
    fn get_order(&self) -> FutureType {
        self.order
    }
}

现在,我们的 future 已经准备好处理了,我们需要重新格式化我们的异步运行时以使用未来类型。我们可以通过在 spawn_task 函数签名中添加额外的 trait 来实现这一点:

fn spawn_task<F, T>(future: F) -> Task<T>
where
    F: Future<Output = T> + Send + 'static + FutureOrderLabel,
    T: Send + 'static,
{
    . . .
}

现在我们可以定义我们的两个队列。到此为止,你可以尝试自己编写它们,然后再继续前进,因为我们已经讨论了构建这两个队列所需的一切。如果你尝试构建队列,应该会得到类似于以下的形式:

static HIGH_QUEUE: LazyLock<flume::Sender<Runnable>> = LazyLock::new(|| {
    let (tx, rx) = flume::unbounded::<Runnable>();
    for _ in 0..2 {
        let receiver = rx.clone();
        thread::spawn(move || {
            while let Ok(runnable) = receiver.recv() {
                let _ = catch_unwind(|| runnable.run());
            }
        });
    }
    tx
});

static LOW_QUEUE: LazyLock<flume::Sender<Runnable>> = LazyLock::new(|| {
    let (tx, rx) = flume::unbounded::<Runnable>();
    for _ in 0..1 {
        let receiver = rx.clone();
        thread::spawn(move || {
            while let Ok(runnable) = receiver.recv() {
                let _ = catch_unwind(|| runnable.run());
            }
        });
    }
    tx
});

低优先级队列有一个消费线程,高优先级队列有两个消费线程。接下来,我们需要将 futures 路由到正确的队列。这可以通过为每个队列定义一个单独的 runner 闭包,然后根据 future 类型传递正确的闭包来实现:

let schedule_high = |runnable| HIGH_QUEUE.send(runnable).unwrap();
let schedule_low = |runnable| LOW_QUEUE.send(runnable).unwrap();

let schedule = match future.get_order() {
    FutureType::High => schedule_high,
    FutureType::Low => schedule_low
};
let (runnable, task) = async_task::spawn(future, schedule);
runnable.schedule();
return task;

现在我们可以创建一个可以插入选定队列的 future:

let one = CounterFuture { count: 0 , order: FutureType::High};

然而,我们遇到了一些问题。假设大量低优先级任务被创建,而没有高优先级任务。在这种情况下,我们将只有一个消费线程在处理所有任务,而另外两个消费线程将闲置。这意味着我们只能以三分之一的容量工作。此时,任务偷取(task stealing)就派上用场了。

注意

我们不需要编写自己的异步运行时队列来控制任务的分配。例如,Tokio 通过使用 LocalSet 允许你控制任务的分配。我们将在第 7 章中介绍这一点。

任务偷取

在任务偷取中,消费线程会在自己的队列为空时,从其他队列中偷取任务。图 3-2 展示了在当前异步系统中任务偷取的情况。

image.png

我们还需要理解,任务偷取也可以反过来进行。如果低优先级队列为空,我们希望低优先级的消费线程从高优先级队列中偷取任务。

为了实现任务偷取,我们需要将高优先级和低优先级队列的通道传递给两个队列。在定义我们的通道之前,我们需要导入以下内容:

use flume::{Sender, Receiver};

如果我们使用标准库中的 SenderReceiver,我们将无法将它们传递到其他线程。通过使用 flume,我们将两个通道都设为静态,并在 spawn_task 函数内延迟评估:

static HIGH_CHANNEL: LazyLock<(Sender<Runnable>, Receiver<Runnable>)> =
    LazyLock::new(|| flume::unbounded::<Runnable>());
static LOW_CHANNEL: LazyLock<(Sender<Runnable>, Receiver<Runnable>)> =
    LazyLock::new(|| flume::unbounded::<Runnable>());

现在我们有了两个通道,我们需要定义高优先级队列的消费线程,这些线程将执行以下步骤,每次循环时重复执行:

  1. 检查 HIGH_CHANNEL 是否有消息。
  2. 如果 HIGH_CHANNEL 没有消息,检查 LOW_CHANNEL 是否有消息。
  3. 如果 LOW_CHANNEL 没有消息,等待 100 毫秒再进行下一次循环。

注意

我们可以在线程空闲时将其挂起,并在需要处理传入任务时唤醒它们。这可以节省过多的循环和休眠操作,当没有任务需要处理时,避免浪费资源。然而,在生产代码中依赖线程休眠可能会增加响应延迟,这是不可取的。在生产环境中,应该避免在线程中使用休眠,而是使用更响应式的机制,如线程停车或条件变量。我们将在第 10 章中讨论与异步队列相关的线程停车。

我们的高优先级队列可以通过以下代码执行这些步骤:

static HIGH_QUEUE: LazyLock<flume::Sender<Runnable>> = LazyLock::new(|| {
    for _ in 0..2 {
        let high_receiver = HIGH_CHANNEL.1.clone();
        let low_receiver = LOW_CHANNEL.1.clone();
        thread::spawn(move || {
            loop {
                match high_receiver.try_recv() {
                    Ok(runnable) => {
                        let _ = catch_unwind(|| runnable.run());
                    },
                    Err(_) => {
                        match low_receiver.try_recv() {
                            Ok(runnable) => {
                                let _ = catch_unwind(|| runnable.run());
                            },
                            Err(_) => {
                                thread::sleep(Duration::from_millis(100));
                            }
                        }
                    }
                };
            }
        });
    }
    HIGH_CHANNEL.0.clone()
});

我们的低优先级队列仅需要将步骤 1 和步骤 2 对调,并返回 LOW_CHANNEL.0.clone。现在,我们的两个队列都首先从各自的队列中拉取任务,在自己的队列没有任务时,才会从其他队列中拉取任务。当没有任务时,我们的消费线程会放慢处理速度。

警告

记住,队列和通道的评估是延迟的。需要向队列发送任务,队列才能开始运行。如果你只向低优先级队列发送任务,而从不向高优先级队列发送任务,那么高优先级队列永远不会启动,因此也不会从低优先级队列偷取任务。

当前进度

到目前为止,我们已经创建了自己的异步运行时队列,并定义了不同的队列!我们现在可以细粒度地控制我们的异步任务如何运行。需要注意的是,我们可能不希望进行任务偷取。例如,如果我们将 CPU 密集型任务放到高优先级队列,将轻量级的网络任务放到低优先级队列,我们就不希望低优先级队列从高优先级队列偷取任务。否则,我们可能会因低优先级队列的消费线程被 CPU 密集型任务阻塞,而关闭网络处理。

虽然实现 trait 约束并看到它如何应用于我们的 future 很有意思,但我们现在面临一个问题。我们不能传递简单的异步块或异步函数,因为它们没有实现 FutureOrderLabel trait。其他开发者可能只想要一个简单的接口来运行他们的任务。如果我们必须为每个异步任务实现 Future trait,并在所有任务上实现 FutureOrderLabel,你能想象我们的代码会有多臃肿吗?为了更好的开发者体验,我们需要重构 spawn_task 函数。

重构我们的 spawn_task 函数

为了允许将异步块和异步函数传递到 spawn_task 函数中,我们需要移除 FutureOrderLabel trait,并移除 CounterFuture 结构体中的 order 字段。接着,我们必须在 spawn_task 函数中移除 FutureOrderLabel trait 的约束,并为任务的 order 添加另一个参数,修改后的函数签名如下:

fn spawn_task<F, T>(future: F, order: FutureType) -> Task<T>
where
    F: Future<Output = T> + Send + 'static,
    T: Send + 'static,
{
    . . .
}

我们还需要更新 spawn_task 函数中正确调度闭包的选择:

let schedule = match order {
    FutureType::High => schedule_high,
    FutureType::Low => schedule_low
};

然而,我们仍然不希望开发者过于担心任务的顺序,因此我们可以为 spawn_task 函数创建一个宏:

macro_rules! spawn_task {
    ($future:expr) => {
        spawn_task!($future, FutureType::Low)
    };
    ($future:expr, $order:expr) => {
        spawn_task($future, $order)
    };
}

这个宏允许我们仅传入 future。如果我们只传入 future,宏将默认为低优先级类型(FutureType::Low)。如果显式传入了 order,则会将其传递给 spawn_task 函数。通过这个宏,Rust 会自动推断你至少需要传递 future 表达式,并且如果没有提供 future,代码将无法编译。我们现在拥有了更符合人机工程学的任务生成方式,示例如下:

fn main() {
    let one = CounterFuture { count: 0 };
    let two = CounterFuture { count: 0 };

    let t_one = spawn_task!(one, FutureType::High);
    let t_two = spawn_task!(two);
    let t_three = spawn_task!(async_fn());
    let t_four = spawn_task!(async {
        async_fn().await;
        async_fn().await;
    }, FutureType::High);

    future::block_on(t_one);
    future::block_on(t_two);
    future::block_on(t_three);
    future::block_on(t_four);
}

这个宏非常灵活。使用它的开发者可以轻松地生成任务而不必过多考虑任务的优先级,但如果需要的话,也能明确声明任务的优先级是高的。我们还可以传入异步块和异步函数,因为它们实际上是 futures 的语法糖。不过,在阻塞主函数以等待多个任务时,我们会遇到重复代码的问题。为了避免这种重复,我们需要创建一个自己的 join 宏。

创建我们自己的 join

为了创建我们自己的 join 宏,我们需要接受一组任务并调用 block_on 函数。我们可以用以下代码来定义我们的 join 宏:

macro_rules! join {
    ($($future:expr),*) => {
        {
            let mut results = Vec::new();
            $(
                results.push(future::block_on($future));
            )*
            results
        }
    };
}

确保我们保持结果的顺序与传入的 futures 顺序一致是非常重要的。否则,用户将无法知道哪个结果属于哪个任务。还要注意,我们的 join 宏将只返回一个类型,因此我们可以像这样使用 join 宏:

let outcome: Vec<u32> = join!(t_one, t_two);
let outcome_two: Vec<()> = join!(t_four, t_three);

outcome 是计数器的输出向量,而 outcome_two 是未返回任何内容的异步函数的输出向量。只要我们有相同的返回类型,这段代码就会有效。

我们必须记住,任务是直接执行的,任务执行过程中可能会发生错误。为了返回结果的向量,我们可以创建一个 try_join 宏,如下所示:

macro_rules! try_join {
    ($($future:expr),*) => {
        {
            let mut results = Vec::new();
            $(
                let result = catch_unwind(|| future::block_on($future));
                results.push(result);
            )*
            results
        }
    };
}

这个宏与我们的 join! 宏类似,但它会返回任务的执行结果。

到目前为止,我们已经具备了几乎所有需要的功能,可以在我们的运行时环境中以符合人机工程学的方式运行异步任务,包括任务偷取和不同的队列。尽管生成任务本身并不完全符合人机工程学,但我们仍然需要一个方便的接口来配置我们的运行时环境。

配置我们的运行时

你可能还记得队列是懒加载的:它在被调用之前不会启动。这直接影响到我们的任务偷取。举例来说,如果没有任务被发送到高优先级队列,该队列就不会启动,因此也无法在低优先级队列为空时从中偷取任务,反之亦然。为了让运行时开始工作,并优化消费循环的数量,配置一个运行时并不是解决这个问题的罕见方法。例如,我们可以参考下面的 Tokio 运行时启动示例:

use tokio::runtime::Runtime;

// 创建运行时
let rt = Runtime::new().unwrap();

// 在运行时中生成一个 future
rt.spawn(async {
    println!("现在运行在一个工作线程上");
});

在写作本文时,上面的示例出现在 Tokio 的运行时结构体文档中。Tokio 库也使用过程宏来设置运行时,但它们超出了本书的范围。你可以在 Rust 文档中找到有关过程宏的更多信息。对于我们的运行时,我们可以构建一个基本的运行时构建器来定义高优先级队列和低优先级队列的消费线程数量。

首先,我们从定义一个运行时结构体开始:

struct Runtime {
    high_num: usize,
    low_num: usize,
}

其中,high_num 是高优先级队列的消费线程数,low_num 是低优先级队列的消费线程数。我们为运行时实现以下函数:

impl Runtime {
    pub fn new() -> Self {
        let num_cores = std::thread::available_parallelism().unwrap().get();
        Self {
            high_num: num_cores - 2,
            low_num: 1,
        }
    }

    pub fn with_high_num(mut self, num: usize) -> Self {
        self.high_num = num;
        self
    }

    pub fn with_low_num(mut self, num: usize) -> Self {
        self.low_num = num;
        self
    }

    pub fn run(&self) {
        . . .
    }
}

在这里,我们使用计算机上可用核心数来标准化定义队列的消费线程数。我们也可以根据需要自己定义高低优先级队列的消费线程数。run 函数用来设置环境变量,然后启动两个任务到两个队列,设置好队列:

pub fn run(&self) {
    std::env::set_var("HIGH_NUM", self.high_num.to_string());
    std::env::set_var("LOW_NUM", self.low_num.to_string());

    let high = spawn_task!(async {}, FutureType::High);
    let low = spawn_task!(async {}, FutureType::Low);
    join!(high, low);
}

我们使用 join! 来确保 run 函数执行完后,两个队列都已经准备好开始任务偷取。

在尝试我们的运行时之前,我们需要使用这些环境变量来确定每个队列的消费线程数量。在我们的 spawn_task 函数中,我们引用每个队列定义中的环境变量:

static HIGH_QUEUE: LazyLock<flume::Sender<Runnable>> = LazyLock::new(|| {
    let high_num = std::env::var("HIGH_NUM").unwrap().parse::<usize>().unwrap();
    for _ in 0..high_num {
        . . .

对于低优先级队列同样适用。然后,我们可以在主函数中定义我们的运行时,并使用默认的线程数量:

Runtime::new().run();

或者,我们可以使用自定义的线程数量:

Runtime::new().with_low_num(2).with_high_num(4).run();

现在,我们可以在程序的其他部分随时运行 spawn_task 函数和 join 宏。我们有了一个可以配置的运行时,支持两种类型的队列和任务偷取!

我们现在已经基本完成了所有工作。不过,在本章结束之前,我们需要讨论最后一个概念:后台进程。

运行后台进程

后台进程是在程序生命周期内定期在后台执行的任务。这些进程可用于监控和维护任务,如数据库清理、日志轮换和数据更新,确保程序始终可以访问最新的信息。将一个基本的后台进程实现为异步运行时中的任务,可以展示如何处理长时间运行的任务。

在处理后台任务之前,我们需要创建一个永远不会停止轮询的 future。在本章的这一阶段,你应该能够自己构建这个任务,尝试构建它,然后再继续。

如果你尝试自己构建这个 future,它的结构应该是这样的,假设正在执行的任务是阻塞的:

#[derive(Debug, Clone, Copy)]
struct BackgroundProcess;

impl Future for BackgroundProcess {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        println!("background process firing");
        std::thread::sleep(Duration::from_secs(1));
        cx.waker().wake_by_ref();
        Poll::Pending
    }
}

你的实现可能会有所不同,但关键是我们始终返回 Pending

我们需要注意的是,如果我们在主函数中丢弃了一个任务,正在异步运行时中执行的任务将被取消,并且不会被执行,因此我们的后台任务必须在程序的整个生命周期中存在。我们需要在主函数的一开始,定义完运行时之后立即发送后台任务:

Runtime::new().with_low_num(2).with_high_num(4).run();
let _background = spawn_task!(BackgroundProcess{});

这样,我们的后台进程将会在程序的整个生命周期内定期运行。

然而,这种方式并不符合人机工程学。举例来说,假设一个结构体或函数可以创建一个后台运行的任务。我们不需要在程序中来回调整任务,以避免任务被丢弃,从而取消后台任务。我们可以通过使用 detach 方法,消除调整任务的需要,使后台任务持续运行:

Runtime::new().with_low_num(2).with_high_num(4).run();
spawn_task!(BackgroundProcess{}).detach();

此方法将任务的指针移到一个不安全的循环中,这个循环将轮询任务并调度它,直到任务完成。与任务关联的指针在主函数中被丢弃,消除了在主函数中保持任务引用的需要。

总结

在本章中,我们实现了自己的运行时,并在此过程中学到了很多东西。我们最初构建了一个基本的异步运行时环境,它接受 futures,创建任务和 runnable。runnable 被放入队列,由消费线程处理,任务则返回给主函数,我们可以阻塞主函数来等待任务的结果。我们花了一些时间巩固了 futures 和任务在异步运行时中经历的步骤。最后,我们实现了具有不同消费线程数量的队列,并使用这种模式实现了任务偷取,用于队列为空的情况。然后,我们为用户创建了自己的宏,使他们能够轻松生成任务并将其合并。

任务偷取引入的细微差别突显了异步编程的真正性质。异步运行时仅仅是你用来解决问题的工具。如果你有一个流量较小的程序,但流量触发了长时间运行的任务,那么你完全可以将一个线程用于接收网络流量,五个线程用于处理长时间、CPU 密集型的任务。在这种情况下,你不希望你的网络队列从 CPU 密集型队列中偷取任务。当然,你的解决方案应该追求合理性。然而,深入了解你使用的异步运行时后,你就能以有趣的方式解决复杂问题。

我们构建的异步运行时显然不是最好的。成熟的异步运行时有一支非常聪明的团队在解决问题和边缘案例。然而,既然你已经完成了本章的学习,你应该明白需要深入阅读你选择的运行时的具体细节,这样你才能将其特性应用到你正在尝试解决的问题上。还要注意,使用 flume 通道实现的简单异步任务队列是可以在生产环境中使用的。

在第 4 章,我们将介绍如何将 HTTP 集成到我们自己的异步运行时中。