sunfishcode's blog sunfishcode 的博客

Rust 中写入未初始化 Buffer 的方法探讨

发布于 2025 年 3 月 11 日

在 Rust 中使用未初始化的 buffer 是一个长期存在的问题,例如:

最近,John Nunley 和 Alex Saveau 提出了一个新方法,使用 Buffer trait。该 trait 现在已经在 rustix 1.0 中实现,我将在本文中介绍它。 更新:这个想法现在已经作为一个独立的库发布:buffer-trait。

Buffer trait 介绍

POSIX 的 read 函数从文件描述符读取字节到 buffer 中,并且可以读取比请求更少的字节。 使用 Buffer,rustix 中的 read 函数看起来像这样:

pub fn read<Fd: AsFd, Buf: Buffer<u8>>(fd: Fd, buf: Buf) -> Result<Buf::Output>

这里使用 Buffer trait 来描述 buffer 参数。Buffer trait 如下所示:

pub trait Buffer<T> {
/// 带有 `Buffer` 参数的函数返回值的类型。
type Output;
/// 返回底层 buffer 的原始指针和长度。
fn parts_mut(&mut self) -> (*mut T, usize);
/// 断言 `len` 个元素已被写入,并提供返回值。
unsafe fn assume_init(self, len: usize) -> Self::Output;
}

(感谢 Yoshua Wuyts 对这个 trait 的反馈以及对整个想法的鼓励!)

(Rustix 自己的 Buffer trait 是被封闭的,并且它的函数是私有的,但这只是 rustix 目前选择保留在不破坏兼容性的前提下发展该 trait 的能力,但代价是暂时不允许用户使用 Buffer 来定义他们自己的 I/O 函数。)

Buffer 为 &mut [T] 实现了,所以用户可以传递一个 &mut [u8] buffer 给 read 函数进行写入,成功时它会返回一个 Result<usize>,其中 usize 指示实际读取的字节数。这与 rustix 中 read 之前的工作方式相匹配。 使用方法如下:

let mut buf = [0_u8; 16];
let num_read = read(fd, &mut buf)?;
use(&buf[..num_read]);

Buffer 也为 &mut [MaybeUninit<T>] 实现了,所以用户可以传递 read 一个 &mut [MaybeUninit<u8>],在这种情况下,他们会得到一个 Result<(&mut [u8], &mut [MaybeUninit<u8>])>。 成功时,它会提供一对 slices,它们是原始 buffer 的子 slices,包含已读入数据的字节范围,以及剩余的未初始化的字节。 Rustix 之前有一个名为 read_uninit 的函数 以这种方式工作,并且在 rustix 1.0 中它被这个新的启用 Buffer 的 read 函数取代。 使用方法如下:

let mut buf = [MaybeUninit::<u8>::uninit(); 16];
let (init, uninit) = read(fd, &mut buf)?;
use(init);

这允许使用安全的 API 写入未初始化的 buffers。

此外,Buffer 也支持读取到 Vec 的备用容量中。spare_capacity 函数 接受一个 &mut Vec<T> 并返回一个 SpareCapacity newtype,该 newtype 实现了 Buffer,并在 read 之后自动设置 vector 的长度以包含已初始化的元素数量,从而封装了 Vec::set_len 的不安全性。 使用方法如下:

let mut buf = Vec::<u8>::with_capacity(1024);
let num_read = read(fd, spare_capacity(&mut buf))?;
use(&buf);

在 rustix 中,所有之前接受 &mut [u8] buffers 用于写入的函数现在都接受 impl Buffer<u8> buffers,所以它们支持写入到未初始化的 buffers。

底层实现

read 的实现如下:

let len = unsafe { backend::io::syscalls::read(fd.as_fd(), buf.parts_mut())? };
unsafe { Ok(buf.assume_init(len)) }

首先,我们调用底层的系统调用,它返回读取的字节数。 然后我们将其传递给 assume_init,它计算要返回的 Buffer::Output。 输出可能只是该数字,或者可能是反映该数字的一对 slices。

如果 T 不是 u8 怎么办?

Buffer 使用类型参数 T 而不是硬编码 u8,以便它可以被诸如 epoll::wait 函数、kevent 函数 和 port::get 函数 使用,以返回事件记录而不是字节。 使用方法如下:

let mut event_list = Vec::<epoll::Event>::with_capacity(16);
loop {
let _num = epoll::wait(&epoll, spare_capacity(&mut event_list), None)?;
for event in event_list.drain(..) {
handle(event);
  }
}

这里使用 drain 耗尽 Vec,以便在每次 wait 之前它是空的,因为 spare_capacity 将附加到 Vec 而不是覆盖任何元素。

循环内部没有动态分配;SpareCapacity 仅使用现有的备用容量,并且仅调用 set_len,而不调用 resize。

或者,由于 Buffer 也适用于 slices,因此可以使用以下代码编写代码而不使用 Vec:

let mut event_list = [MaybeUninit::<epoll::Event>; 16];
loop {
let (init, _uninit) = epoll::wait(&epoll, &mut event_list, None)?;
for event in init {
handle(event);
  }
}

错误信息

Buffer trait 方法的一个缺点是,它有时会引发来自 rustc 的不太明显的错误消息。这种情况发生的次数足够多,以至于我们现在在 rustix 的文档中有一个部分 专门介绍这些错误,并且还有一个示例 展示了它们出现的情况。

安全地使用 Buffer

Rust 的 std 目前包含一个基于 BorrowedBuf 的实验性 API,它具有允许用户在不使用 unsafe 的情况下使用它,并且不进行任何非常低效的操作(例如初始化整个 buffer)的优点。为了实现这一点,BorrowedBuf 使用 “双游标” 设计来避免重新初始化已经初始化的内存。

这里描述的 Buffer trait 更简单,避免了对“双游标”的需求,但是它确实有一个需要 unsafe 的方法。 我们是否可以修改它以支持安全使用?

像 BorrowedCursor 一样的 Cursor API 可以做到这一点。 它支持安全地且增量地写入未初始化的 buffer。 BorrowedCursor 的一个关键特性是它从不需要完全初始化 buffer。

这样,Buffer trait 可能如下所示:

pub trait Buffer<T> {
  // ... 现有的内容
  /// `parts_mut` 的替代方案,用于 `init`。
  ///
  /// 返回一个 `Cursor`。
  fn cursor(&mut self) -> Cursor<T> {
    Cursor::new(self)
  }
}
impl<T, B: Buffer<T>> Cursor<T, B> {
  /// ... cursor API
  fn finish(self) -> B::Output {
    // SAFETY: `Cursor` 确保已写入了 `pos` 个字节。
    unsafe { self.b.assume_init(selff.pos) }
  }
}

这样,用户可以编写自己的接受 Buffer 参数的函数,并使用 cursor 实现它们,而无需使用 unsafe。

为什么是 parts_mut 和原始指针?

Buffer trait 中的 parts_mut 函数如下所示:

fn parts_mut(&mut self) -> (*mut T, usize);

为什么返回一个原始指针和长度,而不是 &mut [MaybeUninit<T>]? 因为 &mut [MaybeUninit<T>] 在一种微妙的方式下是不健全的。 我们为 &mut [T] 实现了 Buffer,它不能包含任何未初始化的元素,并且将其公开为 &mut [MaybeUninit<T>] 将允许将未初始化的元素写入其中。

使用原始指针,我们将保证 buffer 已正确写入的负担放在 assume_init 调用上。

展望未来

该 Buffer trait 的一个有限版本现在已在 rustix 1.0 中,所以我们将看看它在实践中如何。

这个想法现在也可以在独立的库中找到:buffer-trait。

如果它运行良好,我认为这种 Buffer 设计值得考虑用于 Rust 的 std,以替代 BorrowedBuf 函数(目前不稳定)。 它更简单,因为它避免了 “双游标” 模式,并且它具有支持 Vec 备用容量用例并封装不安全的 Vec::set_len 调用的有趣功能。

RSS · Mastodon · Github