sunfishcode's blog sunfishcode 的博客

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

发布于 2025 年 3 月 11 日

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

最近,John NunleyAlex 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 中它被这个新的启用 Bufferread 函数取代。 使用方法如下:

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