Rust原子和锁——构建自己的通道

160 阅读27分钟

一些通道只能在一个发送端和一个接收端之间使用,而另一些通道可以允许任意数量的线程发送数据,甚至支持多个接收端。一些通道是阻塞的,这意味着接收(有时是发送)操作是一个阻塞操作,会让线程休眠直到操作可以完成。有些通道经过优化以提高吞吐量,而另一些则优化为低延迟。

通道的变体是无穷无尽的,没有一种可以适用于所有场景的通用版本。

在这一章中,我们将实现一些相对简单的通道,探索原子操作的更多应用,同时学习如何使用 Rust 的类型系统来捕捉我们的需求和假设。

基于 Mutex 的简单通道

实现一个基本的通道并不需要对原子操作有太多的了解。我们可以使用 VecDeque,它基本上是一个允许高效在两端添加和移除元素的 Vec,并通过 Mutex 来保护它,以允许多个线程访问。我们将使用 VecDeque 作为数据队列,通常称为消息队列,其中存放已发送但尚未接收的消息。任何想要发送消息的线程都可以简单地将消息添加到队列的末尾,而任何想要接收消息的线程只需从队列的开头取出消息即可。

我们还需要添加一个用于使接收操作阻塞的组件:条件变量(Condvar,见“条件变量”),用于通知等待的接收端有新消息到来。

这样的实现可以非常简短且相对直接,如下所示:

use std::sync::{Condvar, Mutex};
use std::collections::VecDeque;

pub struct Channel<T> {
    queue: Mutex<VecDeque<T>>,
    item_ready: Condvar,
}

impl<T> Channel<T> {
    pub fn new() -> Self {
        Self {
            queue: Mutex::new(VecDeque::new()),
            item_ready: Condvar::new(),
        }
    }

    pub fn send(&self, message: T) {
        self.queue.lock().unwrap().push_back(message);
        self.item_ready.notify_one();
    }

    pub fn receive(&self) -> T {
        let mut b = self.queue.lock().unwrap();
        loop {
            if let Some(message) = b.pop_front() {
                return message;
            }
            b = self.item_ready.wait(b).unwrap();
        }
    }
}

注意,我们并不需要使用任何原子操作或不安全代码,也不需要考虑 SendSync 特性。编译器了解 Mutex 的接口以及该类型提供的保证,并会隐式理解如果 Mutex<T>Condvar 都可以安全地在线程之间共享,那么我们的 Channel<T> 也可以。

send 方法锁住了互斥锁,将新消息添加到队列的末尾,并在解锁队列后直接使用条件变量通知一个可能在等待的接收端。

receive 方法也会锁住互斥锁以从队列的开头弹出下一个消息,如果没有消息可用,则会使用条件变量进行等待。

提示

请记住,Condvar::wait 方法在等待期间会解锁 Mutex,并在返回之前重新锁住它。因此,我们的 receive 函数在等待时不会一直保持互斥锁锁定状态。

虽然这个通道在使用上非常灵活,因为它允许任意数量的发送和接收线程,但其实现可能在许多情况下并不理想。即使有足够多的消息可以接收,任何发送或接收操作都会短暂地阻塞其他发送或接收操作,因为它们都必须锁定同一个互斥锁。如果 VecDeque::push 需要扩展 VecDeque 的容量,那么所有的发送和接收线程都必须等待进行重新分配的线程完成,这在某些场景中可能是无法接受的。

另一个可能不理想的特性是,这个通道的队列可能会无限增长。没有任何机制阻止发送端以比接收端处理消息更快的速率不断发送新消息。

不安全的单次传输通道

通道的使用场景可以说是无穷无尽的。然而,在本章的剩余部分中,我们将专注于一种特定类型的使用场景:在一个线程中发送一个消息到另一个线程。为这种场景设计的通道通常被称为单次传输通道(one-shot channel)。

我们可以使用上面基于 Mutex<VecDeque> 的实现,将 VecDeque 替换为 Option,从而将队列的容量有效地减少到只允许一个消息。这样做可以避免分配内存,但仍然存在使用 Mutex 带来的一些问题。我们可以通过从头开始使用原子操作构建自己的单次传输通道来避免这些问题。

首先,我们从构建一个最小实现的单次传输通道开始,而不对接口进行过多的设计。在本章的后面,我们将探索如何改进接口以及如何利用 Rust 的类型系统,为用户提供更好的体验。

我们需要的工具与实现 SpinLock<T>(参见第 4 章)时基本相同:用于存储的 UnsafeCell 以及用于表示状态的 AtomicBool。在这种情况下,我们使用原子布尔值来表示消息是否已准备好供消费。

在消息发送之前,通道是“空”的,还没有包含类型 T 的消息。我们可以在 UnsafeCell 中使用 Option<T> 来表示 T 的缺失。然而,由于我们的原子布尔值已经告诉我们是否有消息,这么做会浪费宝贵的内存空间。因此,我们可以使用 std::mem::MaybeUninit<T>,它本质上是 Option<T> 的不安全精简版:它要求用户手动跟踪是否已初始化,并且它几乎所有的接口都是不安全的,因为它无法自行进行检查。

综合以上,我们可以使用如下结构体定义来开始我们的第一个尝试:

use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicBool, Ordering};
use std::cell::UnsafeCell;

pub struct Channel<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    ready: AtomicBool,
}

SpinLock<T> 一样,我们需要告诉编译器,我们的通道在某些情况下可以安全地在线程之间共享,至少在 T 实现了 Send 时:

unsafe impl<T> Sync for Channel<T> where T: Send {}

一个新的通道是空的,其中 ready 设置为 falsemessage 保持未初始化状态:

impl<T> Channel<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            ready: AtomicBool::new(false),
        }
    }

    …
}

要发送消息,我们首先需要将它存储到 UnsafeCell 中,然后通过将 ready 标志设置为 true 来将其释放给接收方。尝试多次发送将会很危险,因为在设置 ready 标志之后,接收方可能在任何时候读取消息,这与第二次发送消息的尝试可能发生竞态条件。为此,我们将这个责任交给用户,通过将方法标记为 unsafe 并留下一些说明:

/// Safety: 只能调用一次!
pub unsafe fn send(&self, message: T) {
    (*self.message.get()).write(message);
    self.ready.store(true, Ordering::Release);
}

在上面的代码片段中,我们使用 UnsafeCell::get 方法获取指向 MaybeUninit<T> 的指针,然后不安全地解引用该指针,调用 MaybeUninit::write 来初始化它。当误用时,这可能导致未定义行为,但我们将这个责任转移给调用者。

对于内存排序,我们需要使用发布顺序(release ordering),因为该原子存储操作实际上是将消息发布给接收方。这确保了从接收线程的角度来看,如果它使用获取顺序(acquire ordering)从 self.ready 加载 true,则消息的初始化已完成。

对于接收操作,我们暂时不会提供阻塞接口,而是提供两个方法:一个检查消息是否可用,另一个用于接收消息。我们将允许用户使用类似线程停车的方式(见“线程停车”)来实现阻塞。

下面是完成此版本通道的最后两个方法:

pub fn is_ready(&self) -> bool {
    self.ready.load(Ordering::Acquire)
}

/// Safety: 只能调用一次,且只能在 `is_ready()` 返回 `true` 之后调用!
pub unsafe fn receive(&self) -> T {
    (*self.message.get()).assume_init_read()
}

虽然 is_ready 方法可以始终安全地调用,但 receive 方法使用了 MaybeUninit::assume_init_read(),该方法不安全地假设内容已初始化,并且不会用于生成非 Copy 对象的多个副本。与 send 一样,我们也将这个责任交给用户,通过将函数本身标记为 unsafe

这样得到的结果是一个技术上可用的通道,但它并不易于使用且总体上令人失望。如果正确使用,它可以完成它的任务,但存在许多微妙的误用方式。

多次调用 send 可能导致数据竞态条件,因为第二次发送的数据会覆盖第一次发送的数据,而接收方可能正试图读取第一条消息。即使接收操作得到了适当的同步,从多个线程中调用 send 也可能导致多个线程尝试同时写入消息,从而再次导致数据竞态。此外,多次调用 receive 会导致多次复制消息,即使 T 没有实现 Copy,因此不能安全地复制。

一个更微妙的问题是我们的通道缺少 Drop 实现。MaybeUninit 类型并不跟踪它是否已经被初始化,因此在丢弃时不会自动丢弃其内容。这意味着如果消息已经发送但未接收,消息将永远不会被丢弃。这虽然不是未定义行为,但仍然应该避免。尽管在 Rust 中内存泄漏被普遍认为是安全的,但它通常只有在其他内存泄漏的情况下才是可以接受的。例如,泄漏一个 Vec 也会泄漏其内容,但正常使用 Vec 并不会导致任何泄漏。

由于我们让用户对一切负责,这不可避免地会导致意外事故的发生。

通过运行时检查确保安全性

为了提供一个更安全的接口,我们可以添加一些检查,使误用导致明确的错误信息,而不是未定义行为。

首先处理在消息准备好之前调用 receive 的问题。这个很容易解决,因为我们只需在 receive 方法中验证 ready 标志,然后再尝试读取消息:

/// 如果消息尚不可用,则会触发 panic。
///
/// 提示:可以先使用 `is_ready` 进行检查。
///
/// 安全性:只能调用一次!
pub unsafe fn receive(&self) -> T {
    if !self.ready.load(Ordering::Acquire) {
        panic!("no message available!");
    }
    (*self.message.get()).assume_init_read()
}

该函数仍然是 unsafe 的,因为用户仍需负责不调用超过一次,但如果没有先检查 is_ready(),则不会再导致未定义行为。

由于现在 receive 方法中包含了获取(acquire)顺序的 ready 标志加载操作,从而提供了必要的同步,我们可以将 is_ready 中的加载顺序降低为 Relaxed,因为现在它仅用于指示状态:

pub fn is_ready(&self) -> bool {
    self.ready.load(Ordering::Relaxed)
}

注意:总的修改顺序(见“Relaxed Ordering”)确保在 is_ready 加载 true 后,receive 也会看到 true。无论 is_ready 中使用哪种内存顺序,都不会出现 is_ready 返回 truereceive() 仍然 panic 的情况。

接下来处理的问题是多次调用 receive 的情况。通过在 receive 方法中将 ready 标志重新设置为 false,我们可以轻松地使其也导致 panic,如下所示:

/// 如果消息尚不可用或已被消费,则会触发 panic。
///
/// 提示:可以先使用 `is_ready` 进行检查。
pub fn receive(&self) -> T {
    if !self.ready.swap(false, Ordering::Acquire) {
        panic!("no message available!");
    }
    // 安全性:我们刚刚检查(并重置了)ready 标志。
    unsafe { (*self.message.get()).assume_init_read() }
}

我们只需将加载操作更改为交换操作,将 ready 标志设置为 false,这样 receive 方法在任何条件下都是完全安全的。该函数不再标记为 unsafe。我们不再让用户负责所有事情,而是承担了不安全代码的责任,从而减少了用户的负担。

对于 send,情况稍微复杂一些。为了防止多个 send 调用同时访问存储单元,我们需要知道是否已经有其他 send 调用启动。而 ready 标志只能告诉我们是否已有其他 send 调用完成,这还不够。

我们添加第二个标志 in_use 来表示通道是否正在使用:

pub struct Channel<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    in_use: AtomicBool, // 新增
    ready: AtomicBool,
}

impl<T> Channel<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            in_use: AtomicBool::new(false), // 新增
            ready: AtomicBool::new(false),
        }
    }

    …
}

现在,我们只需在 send 方法中在访问存储单元之前将 in_use 设置为 true,并在已被其他调用设置的情况下触发 panic:

/// 当尝试发送多个消息时会触发 panic。
pub fn send(&self, message: T) {
    if self.in_use.swap(true, Ordering::Relaxed) {
        panic!("can't send more than one message!");
    }
    unsafe { (*self.message.get()).write(message) };
    self.ready.store(true, Ordering::Release);
}

我们可以对原子交换操作使用 Relaxed 内存顺序,因为 in_use 的总修改顺序(见“Relaxed Ordering”)确保 in_use 只会有一个交换操作返回 false,而这是 send 访问存储单元的唯一情况。

现在我们有了一个完全安全的接口,但仍有一个问题。如果发送的消息从未被接收,它将永远不会被丢弃。虽然这不会导致未定义行为,在安全代码中也是允许的,但这绝对是需要避免的事情。

由于我们在 receive 方法中重置了 ready 标志,因此修复这个问题很简单:ready 标志表示存储单元中是否有尚未接收的消息需要丢弃。

ChannelDrop 实现中,我们不需要使用原子操作来检查 ready 标志,因为对象只能在它被完全拥有且没有未完成的借用时被丢弃。这意味着我们可以使用 AtomicBool::get_mut 方法,它接受一个独占引用(&mut self),证明不需要原子访问。对于 UnsafeCell,也是通过 UnsafeCell::get_mut 实现的。

使用这些方法,这里是我们完全安全且无泄漏的通道的最后一部分:

impl<T> Drop for Channel<T> {
    fn drop(&mut self) {
        if *self.ready.get_mut() {
            unsafe { self.message.get_mut().assume_init_drop() }
        }
    }
}

试用

由于我们的 Channel 目前尚未提供阻塞接口,因此我们手动使用线程停车来等待消息。接收线程在没有消息准备好的情况下会调用 park() 使其自己进入等待状态,而发送线程在发送完消息后会调用 unpark() 使接收线程解除等待。

以下是一个完整的测试程序,它通过我们的 Channel 从第二个线程向主线程发送字符串 "hello world!"

fn main() {
    let channel = Channel::new();
    let t = thread::current();
    thread::scope(|s| {
        s.spawn(|| {
            channel.send("hello world!");
            t.unpark();
        });
        while !channel.is_ready() {
            thread::park();
        }
        assert_eq!(channel.receive(), "hello world!");
    });
}

该程序可以正常编译、运行并干净退出,证明我们的 Channel 正常工作。

如果我们重复 send 行,还可以看到我们的安全检查产生的效果,程序运行时会产生如下 panic 消息:

thread '<unnamed>' panicked at 'can't send more than one message!', src/main.rs

虽然程序 panic 不是一个好的现象,但让程序可靠地 panic 总比接近潜在的未定义行为要好得多。

使用单个原子变量表示通道状态

如果你喜欢实现通道,这里有一个微妙的变化可以节省一个字节的内存。

我们可以使用一个 AtomicU8 来表示所有四种状态,而不是用两个原子布尔值分别表示通道的状态。我们将使用 compare_exchange 来原子性地检查通道是否处于预期状态,并将其更改为另一状态。

const EMPTY: u8 = 0;
const WRITING: u8 = 1;
const READY: u8 = 2;
const READING: u8 = 3;

pub struct Channel<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    state: AtomicU8,
}

unsafe impl<T: Send> Sync for Channel<T> {}

impl<T> Channel<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            state: AtomicU8::new(EMPTY),
        }
    }

    pub fn send(&self, message: T) {
        if self.state.compare_exchange(
            EMPTY, WRITING, Ordering::Relaxed, Ordering::Relaxed
        ).is_err() {
            panic!("can't send more than one message!");
        }
        unsafe { (*self.message.get()).write(message) };
        self.state.store(READY, Ordering::Release);
    }

    pub fn is_ready(&self) -> bool {
        self.state.load(Ordering::Relaxed) == READY
    }

    pub fn receive(&self) -> T {
        if self.state.compare_exchange(
            READY, READING, Ordering::Acquire, Ordering::Relaxed
        ).is_err() {
            panic!("no message available!");
        }
        unsafe { (*self.message.get()).assume_init_read() }
    }
}

impl<T> Drop for Channel<T> {
    fn drop(&mut self) {
        if *self.state.get_mut() == READY {
            unsafe { self.message.get_mut().assume_init_drop() }
        }
    }
}

这种实现使用了单个原子变量来跟踪所有可能的状态,使代码更加简洁并节省了一些内存。

通过类型确保安全性

虽然我们已经成功保护了 Channel 的用户不遇到未定义行为,但他们在错误使用时仍然可能导致程序 panic。理想情况下,编译器应检查代码的正确使用,并在程序运行之前指出误用。

让我们来看看多次调用 sendreceive 的问题。

为了防止一个函数被调用多次,我们可以让它通过值来接收一个参数,这样对于非 Copy 类型来说,它会消耗该对象。对象一旦被消耗(或移动),它就不再存在于调用者中,从而防止它再次被使用。

通过将 sendreceive 的能力分别表示为单独的(非 Copy)类型,并在执行操作时消耗该对象,我们可以确保每个操作只能执行一次。

这样就得到了以下接口设计,其中一个通道由一个 Sender 和一个 Receiver 表示,每个都有一个通过值接收 self 的方法:

pub fn channel<T>() -> (Sender<T>, Receiver<T>) { … }

pub struct Sender<T> { … }
pub struct Receiver<T> { … }

impl<T> Sender<T> {
    pub fn send(self, message: T) { … }
}

impl<T> Receiver<T> {
    pub fn is_ready(&self) -> bool { … }
    pub fn receive(self) -> T { … }
}

用户可以通过调用 channel() 来创建一个通道,这将返回一个 Sender 和一个 Receiver。他们可以自由传递这些对象,将其移动到另一个线程等。然而,他们不能拥有其中任何一个的多个副本,从而保证了 sendreceive 只能各调用一次。

为了实现这一点,我们需要为 UnsafeCellAtomicBool 找一个位置。之前,我们只有一个结构体包含这些字段,但现在我们有两个独立的结构体,每个结构体可能会比另一个结构体更久地存在。

由于发送者和接收者需要共享这些变量的所有权,我们将使用一个 Arc(引用计数)来为我们提供一个引用计数的共享分配,其中存储了共享的 Channel 对象。如下所示,Channel 类型不必是公共的,因为它的存在只是实现细节,与用户无关。

pub struct Sender<T> {
    channel: Arc<Channel<T>>,
}

pub struct Receiver<T> {
    channel: Arc<Channel<T>>,
}

struct Channel<T> { // 不再是 `pub`
    message: UnsafeCell<MaybeUninit<T>>,
    ready: AtomicBool,
}

unsafe impl<T> Sync for Channel<T> where T: Send {}

就像之前一样,我们在 Channel<T> 上实现 Sync,条件是 TSend,以便它可以跨线程使用。

注意,我们不再需要之前通道实现中的 in_use 原子布尔变量。它之前只被 send 用来检查它是否已被多次调用,现在通过类型系统已经静态地保证了这一点。

创建一个通道和发送接收器对的 channel 函数看起来与之前的 Channel::new 函数类似,只是它将 Channel 包装在一个 Arc 中,并将该 Arc 及其副本包装在 SenderReceiver 类型中:

pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let a = Arc::new(Channel {
        message: UnsafeCell::new(MaybeUninit::uninit()),
        ready: AtomicBool::new(false),
    });
    (Sender { channel: a.clone() }, Receiver { channel: a })
}

sendis_readyreceive 方法基本上与之前实现的类似,有一些差异:

  • 它们现在被移动到了各自的类型中,这样只有(唯一的)发送者可以发送,只有(唯一的)接收者可以接收。
  • sendreceive 现在通过值接收 self,以确保它们只能调用一次。
  • send 不再会 panic,因为它的前置条件(只能调用一次)现在通过类型系统静态保证。

所以它们现在看起来如下所示:

impl<T> Sender<T> {
    /// 不会 panic :)
    pub fn send(self, message: T) {
        unsafe { (*self.channel.message.get()).write(message) };
        self.channel.ready.store(true, Ordering::Release);
    }
}

impl<T> Receiver<T> {
    pub fn is_ready(&self) -> bool {
        self.channel.ready.load(Ordering::Relaxed)
    }

    pub fn receive(self) -> T {
        if !self.channel.ready.swap(false, Ordering::Acquire) {
            panic!("no message available!");
        }
        unsafe { (*self.channel.message.get()).assume_init_read() }
    }
}

receive 函数仍然可能会 panic,因为用户仍然可能在 is_ready() 返回 true 之前调用它。它仍然使用交换操作将 ready 标志重新设置为 false(而不是仅加载),这样 ChannelDrop 实现就知道是否有未读取的消息需要被丢弃。

ChannelDrop 实现与我们之前实现的完全相同:

impl<T> Drop for Channel<T> {
    fn drop(&mut self) {
        if *self.ready.get_mut() {
            unsafe { self.message.get_mut().assume_init_drop() }
        }
    }
}

Arc<Channel<T>>Drop 实现将在 Sender<T>Receiver<T> 之一被丢弃时减少分配的引用计数。在丢弃第二个时,该计数将变为零,Channel<T> 本身将被丢弃。这将调用我们上面的 Drop 实现,如果有消息发送但未接收,我们可以在其中丢弃该消息。

试用

fn main() {
    thread::scope(|s| {
        let (sender, receiver) = channel();
        let t = thread::current();
        s.spawn(move || {
            sender.send("hello world!");
            t.unpark();
        });
        while !receiver.is_ready() {
            thread::park();
        }
        assert_eq!(receiver.receive(), "hello world!");
    });
}

尽管我们仍然需要手动使用线程停车来等待消息,这有点不便,但稍后我们会处理这个问题。

我们现在的目标是至少让某些形式的误用在编译时就不可能发生。不像之前,尝试发送两次不会导致程序 panic,而是根本不会生成有效的程序。如果我们在上述工作程序中添加另一个 send 调用,编译器会捕捉到问题并耐心地告知我们错误:

error[E0382]: use of moved value: `sender`
  --> src/main.rs
   |
   |             sender.send("hello world!");
   |                    --------------------
   |                     `sender` moved due to this method call
   |
   |             sender.send("second message");
   |             ^^^^^^ value used here after move
   |
note: this function takes ownership of the receiver `self`, which moves `sender`
  --> src/lib.rs
   |
   |     pub fn send(self, message: T) {
   |                 ^^^^
   = note: move occurs because `sender` has type `Sender<&str>`,
           which does not implement the `Copy` trait

根据情况,设计一个在编译时捕捉错误的接口可能非常棘手。如果情况适合这样的接口,不仅可以为用户提供更多的便利,还可以减少运行时检查,因为某些事情现在已经静态保证。例如,我们不再需要 in_use 标志,并从 send 方法中移除了交换和检查操作。

不幸的是,可能会出现新的问题,导致更多的运行时开销。在这种情况下,问题在于分离的所有权,我们不得不使用 Arc 并付出分配的成本。

在安全性、便利性、灵活性、简单性和性能之间做权衡是不幸但有时不可避免的。Rust 通常努力使这些方面都做到出色,但有时你需要在其中一个方面做出些许妥协,以最大化另一个方面。

通过借用避免分配

我们刚刚设计的基于 Arc 的通道实现在使用上非常方便,但代价是性能下降,因为它需要分配内存。如果我们希望优化效率,可以牺牲一些便利性来提高性能,将共享 Channel 对象的管理交给用户。与其在后台处理 Channel 的分配和所有权,我们可以强制用户创建一个可以被 SenderReceiver 借用的 Channel,这样他们可以简单地将 Channel 放在本地变量中,从而避免分配内存的开销。

我们还必须在简单性上做出一些让步,因为我们现在必须处理借用和生命周期。

因此,三个类型现在如下所示,其中 Channel 再次变为公共的,而 SenderReceiver 以某个生命周期借用它:

pub struct Channel<T> {
    message: UnsafeCell<MaybeUninit<T>>,
    ready: AtomicBool,
}

unsafe impl<T> Sync for Channel<T> where T: Send {}

pub struct Sender<'a, T> {
    channel: &'a Channel<T>,
}

pub struct Receiver<'a, T> {
    channel: &'a Channel<T>,
}

我们不再使用创建 (Sender, Receiver) 对的 channel() 函数,而是回到了本章早些时候的 Channel::new,允许用户将这样的对象创建为本地变量。

此外,我们需要一种方法让用户创建借用 ChannelSenderReceiver 对象。这需要是独占借用(&mut Channel),以确保不能为同一个通道创建多个发送者或接收者。通过同时提供 SenderReceiver,我们可以将独占借用拆分为两个共享借用,从而发送者和接收者都可以引用通道,同时防止其他任何东西访问通道。

这导致了以下实现:

impl<T> Channel<T> {
    pub const fn new() -> Self {
        Self {
            message: UnsafeCell::new(MaybeUninit::uninit()),
            ready: AtomicBool::new(false),
        }
    }

    pub fn split<'a>(&'a mut self) -> (Sender<'a, T>, Receiver<'a, T>) {
        *self = Self::new();
        (Sender { channel: self }, Receiver { channel: self })
    }
}

split 方法的签名比较复杂,需要仔细看一下。它通过独占引用来借用 self,但将其拆分为两个共享引用,分别包装在 SenderReceiver 类型中。生命周期 'a 清楚地表明,这两个对象都借用了具有有限生命周期的内容,在本例中就是 Channel 本身。由于 Channel 被独占借用,只要 SenderReceiver 对象存在,调用者就无法借用或移动它。

然而,一旦这两个对象都不再存在,独占借用就会过期,编译器会允许再次借用 Channel 对象并调用 split()。虽然我们可以假设在 SenderReceiver 仍然存在时不能再次调用 split(),但我们无法防止在这些对象被丢弃或忘记后进行第二次调用。我们需要确保不会意外地为已经设置 ready 标志的通道创建新的 SenderReceiver 对象,因为这会破坏防止未定义行为的假设。

通过在 split() 中用一个新的空通道覆盖 *self,我们可以确保在创建 SenderReceiver 状态时通道处于预期状态。这也会调用旧 *selfDrop 实现,该实现将负责丢弃之前发送但未接收的消息。

注意

由于 split 方法的签名中的生命周期来自 self,因此可以省略。上面代码片段中的 split 签名与以下更简洁的版本相同:

pub fn split(&mut self) -> (Sender<T>, Receiver<T>) { … }

虽然这个版本没有明确显示返回的对象借用了 self,但编译器仍会像处理更详细的版本一样检查生命周期的正确使用。

其余的方法和 Drop 实现与基于 Arc 的实现相同,只是对 SenderReceiver 类型添加了 '_' 生命周期参数(如果您忘记添加,编译器会建议您添加)。

为了完整起见,以下是剩余的代码:

impl<T> Sender<'_, T> {
    pub fn send(self, message: T) {
        unsafe { (*self.channel.message.get()).write(message) };
        self.channel.ready.store(true, Release);
    }
}

impl<T> Receiver<'_, T> {
    pub fn is_ready(&self) -> bool {
        self.channel.ready.load(Relaxed)
    }

    pub fn receive(self) -> T {
        if !self.channel.ready.swap(false, Acquire) {
            panic!("no message available!");
        }
        unsafe { (*self.channel.message.get()).assume_init_read() }
    }
}

impl<T> Drop for Channel<T> {
    fn drop(&mut self) {
        if *self.ready.get_mut() {
            unsafe { self.message.get_mut().assume_init_drop() }
        }
    }
}

测试

fn main() {
    let mut channel = Channel::new();
    thread::scope(|s| {
        let (sender, receiver) = channel.split();
        let t = thread::current();
        s.spawn(move || {
            sender.send("hello world!");
            t.unpark();
        });
        while !receiver.is_ready() {
            thread::park();
        }
        assert_eq!(receiver.receive(), "hello world!");
    });
}

与基于 Arc 的版本相比,便利性的减少非常小:我们只需要多写一行手动创建一个 Channel 对象的代码。然而,请注意,通道必须在作用域(scope)之前创建,以向编译器证明它的存在会比发送者和接收者都长。

要查看编译器的借用检查器的作用,可以在不同位置添加第二次调用 channel.split()。你会发现,如果在线程作用域内第二次调用它会导致错误,而在作用域之后调用则是可以接受的。即使在作用域开始之前调用 split() 也可以,只要在作用域开始之前停止使用返回的 SenderReceiver

阻塞

最后,我们来解决 Channel 的最后一个主要不便:缺少阻塞接口。每次测试通道的新变体时,我们都使用了线程挂起。这种模式不难集成到通道本身中。

为了能够唤醒接收线程,发送者需要知道要唤醒哪个线程。std::thread::Thread 类型表示对线程的句柄,正是我们需要用来调用 unpark() 的东西。我们将在 Sender 对象中存储指向接收线程的句柄,如下所示:

use std::thread::Thread;

pub struct Sender<'a, T> {
    channel: &'a Channel<T>,
    receiving_thread: Thread, // 新增!
}

但是,如果 Receiver 对象在不同线程之间传递,这个句柄就会引用错误的线程。发送者对此一无所知,仍然引用最初持有接收者的线程。

我们可以通过对 Receiver 进行一些限制,来处理这个问题:不允许它在不同线程之间传递。正如在“线程安全:Send 和 Sync”中讨论的那样,我们可以使用特殊的 PhantomData 标记类型将此限制添加到我们的结构体中。PhantomData<*const ()> 就能胜任此工作,因为原始指针(如 *const ())并不实现 Send

pub struct Receiver<'a, T> {
    channel: &'a Channel<T>,
    _no_send: PhantomData<*const ()>, // 新增!
}

接下来,我们必须修改 Channel::split 方法以填充新的字段,如下所示:

pub fn split<'a>(&'a mut self) -> (Sender<'a, T>, Receiver<'a, T>) {
    *self = Self::new();
    (
        Sender {
            channel: self,
            receiving_thread: thread::current(), // 新增!
        },
        Receiver {
            channel: self,
            _no_send: PhantomData, // 新增!
        }
    )
}

我们将当前线程的句柄用于 receiving_thread 字段,因为我们返回的 Receiver 对象将保留在当前线程中。

send 方法的变化不大,如下所示。我们只需在接收线程可能正在等待时调用 unpark() 来唤醒它:

impl<T> Sender<'_, T> {
    pub fn send(self, message: T) {
        unsafe { (*self.channel.message.get()).write(message) };
        self.channel.ready.store(true, Release);
        self.receiving_thread.unpark(); // 新增!
    }
}

receive 函数进行了较大的修改。新版本在没有消息时不会引发恐慌,而是使用 thread::park() 耐心等待消息,并根据需要多次尝试:

impl<T> Receiver<'_, T> {
    pub fn receive(self) -> T {
        while !self.channel.ready.swap(false, Acquire) {
            thread::park();
        }
        unsafe { (*self.channel.message.get()).assume_init_read() }
    }
}

注意

请记住,thread::park() 可能会虚假返回(或者是因为除了我们的 send 方法之外的其他东西调用了 unpark())。这意味着我们不能假设当 park() 返回时 ready 标志已经设置。因此,我们需要在解除挂起后使用循环再次检查标志。

Channel<T> 结构体、其 Sync 实现、其 new 函数和其 Drop 实现保持不变。

测试

fn main() {
    let mut channel = Channel::new();
    thread::scope(|s| {
        let (sender, receiver) = channel.split();
        s.spawn(move || {
            sender.send("hello world!");
        });
        assert_eq!(receiver.receive(), "hello world!");
    });
}

显然,在这个简单的测试程序中,这个 Channel 比上一个更加方便使用。我们为此便利性付出的代价是减少了一些灵活性:只有调用 split() 的线程才能调用 receive()。如果你交换 sendreceive 的顺序,程序将无法编译。根据使用场景,这可能完全没问题、有帮助,或者非常不方便。

有很多方法可以解决这个问题,其中许多方法会使我们付出更多的复杂性并影响性能。一般来说,我们可以继续探索的变化和权衡几乎是无穷无尽的。

我们可以轻松花费大量时间实现二十种不同变体的单次通道,每种变体的属性略有不同,以适应各种想象得到的使用场景。虽然这听起来很有趣,但我们最好避免陷入这个深坑,结束本章,以免事情失控。

总结

  1. 通道用于在线程之间传递消息。
  2. 使用 MutexCondvar 实现一个简单且灵活的通道相对容易,但可能效率低下。
  3. 单次通道(one-shot channel)是一种专为发送一次消息而设计的通道。
  4. MaybeUninit<T> 类型可以用来表示一个可能尚未初始化的 T。它的接口大部分是不安全的,要求其用户负责跟踪是否已初始化,避免重复使用非 Copy 数据,并在必要时释放其内容。
  5. 不释放对象(也称为泄漏或忘记)是安全的,但在没有正当理由的情况下不被鼓励。
  6. 使用 panic 是创建安全接口的重要工具。
  7. 通过值传递一个非 Copy 对象,可以防止某些操作被执行多次。
  8. 独占借用和拆分借用是一种强制确保正确性的有力工具。
  9. 可以通过确保其类型不实现 Send 来确保对象留在同一线程中,这可以通过 PhantomData 标记类型来实现。
  10. 每个设计和实现决策都有取舍,最好根据特定的用例来做决定。
  11. 没有用例的设计虽然有趣且具有教育意义,但可能会变成一个无止境的任务。