|
|
# 40｜异步处理：如何处理异步IO？
|
|
|
|
|
|
你好，我是陈天。
|
|
|
|
|
|
前面两讲我们学习了异步处理基本的功能和原理（Future/async/await），但是还没有正式介绍在具体场合下该用哪些工具来处理异步 IO。不过之前讲 trait 的时候，已经了解和使用过一些处理同步 IO 的结构和 trait。
|
|
|
|
|
|
今天我们就对比同步 IO 来学习异步 IO。毕竟在学习某个新知识的时候，如果能够和头脑中已有的知识联系起来，大脑神经元之间的连接就会被激活，学习的效果会事半功倍。
|
|
|
|
|
|
回忆一下同步环境都有哪些结构和 trait呢？首先，单个的值可以用类型 T 表述，一组值可以用 Iterator trait 表述；同步 IO，我们有标准的 Read/Write/Seek trait。顾名思义，Read/Write 是进行 IO 的读写，而 Seek 是在 IO 中前后移动当前的位置。
|
|
|
|
|
|
那么异步呢？我们已经学习到，对于单个的、在未来某个时刻会得到的值，可以用 Future 来表示：
|
|
|
|
|
|
|
|
|
|
|
|
但还不知道一组未来才能得到的值该用什么 trait 来表述，也不知道异步的 Read/Write 该是什么样子。今天，我们就来聊聊这些重要的异步数据类型。
|
|
|
|
|
|
## Stream trait
|
|
|
|
|
|
首先来了解一下 Iterator 在异步环境下的表兄弟：Stream。
|
|
|
|
|
|
我们知道，对于 Iterator，可以不断调用其 next() 方法，获得新的值，直到 Iterator 返回 None。Iterator 是阻塞式返回数据的，每次调用 next()，必然独占 CPU 直到得到一个结果，**而异步的 Stream 是非阻塞的，在等待的过程中会空出 CPU 做其他事情**。
|
|
|
|
|
|
不过和 Future 已经在标准库稳定下来不同，Stream trait 目前还只能在 nightly 版本使用。一般跟 Stream 打交道，我们会使用 futures 库。来对比 Iterator 和 Stream的源码定义：
|
|
|
|
|
|
```rust
|
|
|
pub trait Iterator {
|
|
|
type Item;
|
|
|
fn next(&mut self) -> Option<Self::Item>;
|
|
|
|
|
|
fn size_hint(&self) -> (usize, Option<usize>) { ... }
|
|
|
fn map<B, F>(self, f: F) -> Map<Self, F> where F: FnMut(Self::Item) -> B { ... }
|
|
|
... // 还有 67 个方法
|
|
|
}
|
|
|
|
|
|
pub trait Stream {
|
|
|
type Item;
|
|
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>;
|
|
|
|
|
|
fn size_hint(&self) -> (usize, Option<usize>) { ... }
|
|
|
}
|
|
|
|
|
|
pub trait StreamExt: Stream {
|
|
|
fn next(&mut self) -> Next<'_, Self> where Self: Unpin { ... }
|
|
|
fn map<T, F>(self, f: F) -> Map<Self, F> where F: FnMut(Self::Item) -> T { ... }
|
|
|
... // 还有 41 个方法
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
可以看到，Iterator 把所有方法都放在 Iterator trait 里，而Stream 把需要开发者实现的基本方法和有缺省实现的衍生方法区别开，放在不同的 trait 里。比如 map。
|
|
|
|
|
|
实现 Stream 的时候，和 Iterator 类似，你需要提供 Item 类型，这是每次拿出一个值时，值的类型；此外，还有 poll\_next() 方法，它长得和 Future 的 poll() 方法很像，和 Iterator 版本的 next() 的作用类似。
|
|
|
|
|
|
然而，poll\_next() 调用起来不方便，我们需要自己处理 Poll 状态，所以，StreamExt 提供了 next() 方法，返回一个实现了 Future trait 的 Next 结构，这样，我们就可以直接通过 stream.next().await 来获取下一个值了。来看 next() 方法以及 Next 结构的实现（[源码](https://docs.rs/futures-util/0.3.17/src/futures_util/stream/stream/next.rs.html#10-34)）：
|
|
|
|
|
|
```rust
|
|
|
pub trait StreamExt: Stream {
|
|
|
fn next(&mut self) -> Next<'_, Self> where Self: Unpin {
|
|
|
assert_future::<Option<Self::Item>, _>(Next::new(self))
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// next 返回了 Next 结构
|
|
|
pub struct Next<'a, St: ?Sized> {
|
|
|
stream: &'a mut St,
|
|
|
}
|
|
|
|
|
|
// 如果 Stream Unpin 那么 Next 也是 Unpin
|
|
|
impl<St: ?Sized + Unpin> Unpin for Next<'_, St> {}
|
|
|
|
|
|
impl<'a, St: ?Sized + Stream + Unpin> Next<'a, St> {
|
|
|
pub(super) fn new(stream: &'a mut St) -> Self {
|
|
|
Self { stream }
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// Next 实现了 Future，每次 poll() 实际上就是从 stream 中 poll_next()
|
|
|
impl<St: ?Sized + Stream + Unpin> Future for Next<'_, St> {
|
|
|
type Output = Option<St::Item>;
|
|
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
|
self.stream.poll_next_unpin(cx)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
看个小例子（[代码](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=4374c885eaf1a386fe6a67d8a54dc37b)）：
|
|
|
|
|
|
```rust
|
|
|
use futures::prelude::*;
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() {
|
|
|
let mut st = stream::iter(1..10)
|
|
|
.filter(|x| future::ready(x % 2 == 0))
|
|
|
.map(|x| x * x);
|
|
|
|
|
|
while let Some(x) = st.next().await {
|
|
|
println!("Got item: {}", x);
|
|
|
}
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
我们使用 stream::iter 生成了一个 Stream，并对其进行 filter / map 的操作。最后，遍历整个 stream，把获得的数据打印出来。从使用的感受来看，Stream 和 Iterator 也很相似，可以对比着来用。
|
|
|
|
|
|
### 生成 Stream
|
|
|
|
|
|
futures 库提供了一些基本的生成 Stream 的方法，除了上面用到的 iter 方法，还有：
|
|
|
|
|
|
* empty()：生成一个空的 Stream
|
|
|
* once()：生成一个只包含单个值的 Stream
|
|
|
* pending()：生成一个不包含任何值，只返回 Poll::Pending 的 Stream
|
|
|
* repeat()：生成一个一直返回相同值的 Stream
|
|
|
* repeat\_with()：通过闭包函数无穷尽地返回数据的 Stream
|
|
|
* poll\_fn()：通过一个返回 Poll<Option<T>> 的闭包来产生 Stream
|
|
|
* unfold()：通过初始值和返回 Future 的闭包来产生 Stream
|
|
|
|
|
|
前几种产生 Stream 的方法都很好理解，最后三种引入了闭包复杂一点，我们分别使用它们来实现斐波那契数列，对比一下差异（[代码](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=d83d2877d953b381c4f412b5768288ff)）：
|
|
|
|
|
|
```rust
|
|
|
use futures::{prelude::*, stream::poll_fn};
|
|
|
use std::task::Poll;
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() {
|
|
|
consume(fib().take(10)).await;
|
|
|
consume(fib1(10)).await;
|
|
|
// unfold 产生的 Unfold stream 没有实现 Unpin，
|
|
|
// 所以我们将其 Pin<Box<T>> 一下，使其满足 consume 的接口
|
|
|
consume(fib2(10).boxed()).await;
|
|
|
}
|
|
|
|
|
|
async fn consume(mut st: impl Stream<Item = i32> + Unpin) {
|
|
|
while let Some(v) = st.next().await {
|
|
|
print!("{} ", v);
|
|
|
}
|
|
|
print!("\\n");
|
|
|
}
|
|
|
|
|
|
// 使用 repeat_with 创建 stream，无法控制何时结束
|
|
|
fn fib() -> impl Stream<Item = i32> {
|
|
|
let mut a = 1;
|
|
|
let mut b = 1;
|
|
|
stream::repeat_with(move || {
|
|
|
let c = a + b;
|
|
|
a = b;
|
|
|
b = c;
|
|
|
b
|
|
|
})
|
|
|
}
|
|
|
|
|
|
// 使用 poll_fn 创建 stream，可以通过返回 Poll::Ready(None) 来结束
|
|
|
fn fib1(mut n: usize) -> impl Stream<Item = i32> {
|
|
|
let mut a = 1;
|
|
|
let mut b = 1;
|
|
|
poll_fn(move |_cx| -> Poll<Option<i32>> {
|
|
|
if n == 0 {
|
|
|
return Poll::Ready(None);
|
|
|
}
|
|
|
n -= 1;
|
|
|
let c = a + b;
|
|
|
a = b;
|
|
|
b = c;
|
|
|
Poll::Ready(Some(b))
|
|
|
})
|
|
|
}
|
|
|
|
|
|
fn fib2(n: usize) -> impl Stream<Item = i32> {
|
|
|
stream::unfold((n, (1, 1)), |(mut n, (a, b))| async move {
|
|
|
if n == 0 {
|
|
|
None
|
|
|
} else {
|
|
|
n -= 1;
|
|
|
let c = a + b;
|
|
|
// c 作为 poll_next() 的返回值，(n, (a, b)) 作为 state
|
|
|
Some((c, (n, (b, c))))
|
|
|
}
|
|
|
})
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
值得注意的是，**使用 unfold 的时候，同时使用了局部变量和 Future，所以生成的 Stream 没有实现 Unpin**，我们在使用的时候，需要将其 pin 住。怎么做呢？
|
|
|
|
|
|
Pin<Box<T>> 是一种很简单的方法，能将数据 Pin 在堆上，我们可以使用 StreamExt 的 boxed() 方法来生成一个 Pin<Box<T>>。
|
|
|
|
|
|
除了上面讲的方法，我们还可以为一个数据结构实现 Stream trait，从而使其支持 Stream。看一个例子（[代码](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=94797b51e031dbde9d81ff04c5ff9f83)）：
|
|
|
|
|
|
```rust
|
|
|
use futures::prelude::*;
|
|
|
use pin_project::pin_project;
|
|
|
use std::{
|
|
|
pin::Pin,
|
|
|
task::{Context, Poll},
|
|
|
};
|
|
|
use tokio::{
|
|
|
fs,
|
|
|
io::{AsyncBufReadExt, AsyncRead, BufReader, Lines},
|
|
|
};
|
|
|
|
|
|
/// LineStream 内部使用 tokio::io::Lines
|
|
|
#[pin_project]
|
|
|
struct LineStream<R> {
|
|
|
#[pin]
|
|
|
lines: Lines<BufReader<R>>,
|
|
|
}
|
|
|
|
|
|
impl<R: AsyncRead> LineStream<R> {
|
|
|
/// 从 BufReader 创建一个 LineStream
|
|
|
pub fn new(reader: BufReader<R>) -> Self {
|
|
|
Self {
|
|
|
lines: reader.lines(),
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
/// 为 LineStream 实现 Stream trait
|
|
|
impl<R: AsyncRead> Stream for LineStream<R> {
|
|
|
type Item = std::io::Result<String>;
|
|
|
|
|
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
|
self.project()
|
|
|
.lines
|
|
|
.poll_next_line(cx)
|
|
|
.map(Result::transpose)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() -> std::io::Result<()> {
|
|
|
let file = fs::File::open("Cargo.toml").await?;
|
|
|
let reader = BufReader::new(file);
|
|
|
let mut st = LineStream::new(reader);
|
|
|
while let Some(Ok(line)) = st.next().await {
|
|
|
println!("Got: {}", line);
|
|
|
}
|
|
|
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
这段代码封装了 [Lines](https://docs.rs/tokio/1.14.0/tokio/io/struct.Lines.html) 结构，我们可以通过 [AsyncBufReadExt](https://docs.rs/tokio/1.14.0/tokio/io/trait.AsyncBufReadExt.html#) 的 [lines()](https://docs.rs/tokio/1.14.0/tokio/io/trait.AsyncBufReadExt.html#method.lines) 方法，把一个实现了 [AsyncBufRead](https://docs.rs/tokio/1.14.0/tokio/io/trait.AsyncBufRead.html) trait 的 reader 转换成 Lines。
|
|
|
|
|
|
你也许注意到代码中引入的 [pin\_project](https://docs.rs/pin-project) 库，它提供了一些便利的宏，方便我们操作数据结构里需要被 pin 住的字段。在数据结构中，**可以使用 #\[pin\] 来声明某个字段在使用的时候需要被封装为 Pin<T>**。这样，调用时，我们就可以使用 self.project().lines 得到一个 Pin<&mut Lines>，以便调用其 [poll\_next\_line()](https://docs.rs/tokio/1.14.0/tokio/io/struct.Lines.html#method.poll_next_line) 方法（这个方法的第一个参数是 Pin<&mut Self>）。
|
|
|
|
|
|
在Lines这个结构内部，异步的 [next\_line()](https://docs.rs/tokio/1.14.0/tokio/io/struct.Lines.html#method.next_line) 方法可以读取下一行，它实际上就是比较低阶的 [poll\_next\_line()](https://docs.rs/tokio/1.14.0/src/tokio/io/util/lines.rs.html#112-134) 接口的一个封装。
|
|
|
|
|
|
**虽然 Lines 结构提供了 next\_line()，但并没有实现 Stream**，所以我们无法像其他 Stream 那样统一用 next() 方法获取下一行。于是，我们将其包裹在自己的 LineStream 下，并且为 LineStream 实现了 Stream 方法。
|
|
|
|
|
|
注意，由于 poll\_next\_line() 的结果是 Result<Option<T>>，而 Stream 的 poll\_next() 的结果是 Option<Result<T>>，所以我们需要使用 Result 方法的 [transpose](https://doc.rust-lang.org/std/result/enum.Result.html#method.transpose) 来将二者对调。这个transpose 方法是一个很基础的方法，非常实用。
|
|
|
|
|
|
## 异步 IO 接口
|
|
|
|
|
|
在实现 LineStream 时，我们遇到了两个异步 I/O 接口：AsyncRead 以及 AsyncBufRead。回到开头的那张表，相信你现在已经有大致答案了吧：**所有同步的 Read / Write / Seek trait，前面加一个 Async，就构成了对应的异步 IO 接口**。
|
|
|
|
|
|
不过，和 Stream 不同的是，如果你对比 futures 下定义的 IO trait 以及 tokio 下定义的 IO trait，会发现它们都有各自的定义，双方并未统一，有些许的差别：
|
|
|
|
|
|
|
|
|
|
|
|
比如 futures 下 [AsyncRead](https://docs.rs/futures/0.3.17/futures/io/trait.AsyncRead.html) 的定义：
|
|
|
|
|
|
```rust
|
|
|
pub trait AsyncRead {
|
|
|
fn poll_read(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
buf: &mut [u8]
|
|
|
) -> Poll<Result<usize, Error>>;
|
|
|
|
|
|
unsafe fn initializer(&self) -> Initializer { ... }
|
|
|
fn poll_read_vectored(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
bufs: &mut [IoSliceMut<'_>]
|
|
|
) -> Poll<Result<usize, Error>> { ... }
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
而 tokio 下 [AsyncRead](https://docs.rs/tokio/1.14.0/tokio/io/trait.AsyncRead.html) 的定义：
|
|
|
|
|
|
```rust
|
|
|
pub trait AsyncRead {
|
|
|
fn poll_read(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
buf: &mut ReadBuf<'_>
|
|
|
) -> Poll<Result<()>>;
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
我们看不同之处：tokio 的 poll\_read() 方法需要 [ReadBuf](https://docs.rs/tokio/1.14.0/src/tokio/io/read_buf.rs.html#27-31)，而 futures 的 poll\_read() 方法需要 &mut \[u8\]。此外，futures 的 AsyncRead 还多了两个缺省方法。
|
|
|
|
|
|
再看 AsyncWrite。futures 下的 [AsyncWrite](https://docs.rs/futures/0.3.17/futures/io/trait.AsyncWrite.html) 接口如下：
|
|
|
|
|
|
```rust
|
|
|
pub trait AsyncWrite {
|
|
|
fn poll_write(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
buf: &[u8]
|
|
|
) -> Poll<Result<usize, Error>>;
|
|
|
fn poll_flush(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Error>>;
|
|
|
fn poll_close(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Error>>;
|
|
|
|
|
|
fn poll_write_vectored(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
bufs: &[IoSlice<'_>]
|
|
|
) -> Poll<Result<usize, Error>> { ... }
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
而 tokio 下的 [AsyncWrite](https://docs.rs/tokio/1.14.0/tokio/io/trait.AsyncWrite.html) 的定义：
|
|
|
|
|
|
```rust
|
|
|
pub trait AsyncWrite {
|
|
|
fn poll_write(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
buf: &[u8]
|
|
|
) -> Poll<Result<usize, Error>>;
|
|
|
fn poll_flush(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Error>>;
|
|
|
fn poll_shutdown(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Error>>;
|
|
|
|
|
|
fn poll_write_vectored(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
bufs: &[IoSlice<'_>]
|
|
|
) -> Poll<Result<usize, Error>> { ... }
|
|
|
fn is_write_vectored(&self) -> bool { ... }
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
可以看到，AsyncWrite 二者的差距就只有 poll\_close() 和 poll\_shutdown() 命名上的分别。其它的异步 IO 接口我就不一一举例了，你可以自己去看代码对比。
|
|
|
|
|
|
### 异步 IO 接口的兼容性处理
|
|
|
|
|
|
为什么 Rust 的异步 IO trait 会有这样的分裂？这是因为在 tokio / futures 库实现的早期，社区还没有形成比较统一的异步 IO trait，不同的接口背后也有各自不同的考虑，这种分裂就沿袭下来。
|
|
|
|
|
|
所以，如果我们使用 tokio 进行异步开发，那么，代码需要使用 tokio::io 下的异步 IO trait。也许，未来等 Async IO trait 稳定并进入标准库后，tokio 会更新自己的 trait。
|
|
|
|
|
|
虽然 Rust 的异步 IO trait 有这样的分裂，你也不必过分担心。**tokio-util 提供了相应的**[Compat](https://docs.rs/tokio-util/0.6.9/tokio_util/compat/index.html)**功能，可以让你的数据结构在二者之间自如切换**。看一个使用 [yamux](https://docs.rs/yamux) 做多路复用的例子，重点位置详细注释了：
|
|
|
|
|
|
```rust
|
|
|
use anyhow::Result;
|
|
|
use futures::prelude::*;
|
|
|
use tokio::net::TcpListener;
|
|
|
use tokio_util::{
|
|
|
codec::{Framed, LinesCodec},
|
|
|
compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt},
|
|
|
};
|
|
|
use tracing::info;
|
|
|
use yamux::{Config, Connection, Mode, WindowUpdateMode};
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() -> Result<()> {
|
|
|
tracing_subscriber::fmt::init();
|
|
|
let addr = "0.0.0.0:8080";
|
|
|
let listener = TcpListener::bind(addr).await?;
|
|
|
info!("Listening on: {:?}", addr);
|
|
|
loop {
|
|
|
let (stream, addr) = listener.accept().await?;
|
|
|
info!("Accepted: {:?}", addr);
|
|
|
let mut config = Config::default();
|
|
|
config.set_window_update_mode(WindowUpdateMode::OnRead);
|
|
|
// 使用 compat() 方法把 tokio AsyncRead/AsyncWrite 转换成 futures 对应的 trait
|
|
|
let conn = Connection::new(stream.compat(), config, Mode::Server);
|
|
|
// Yamux ctrl stream 可以用来打开新的 stream
|
|
|
let _ctrl = conn.control();
|
|
|
tokio::spawn(
|
|
|
yamux::into_stream(conn).try_for_each_concurrent(None, move |s| async move {
|
|
|
// 使用 compat() 方法把 futures AsyncRead/AsyncWrite 转换成 tokio 对应的 trait
|
|
|
let mut framed = Framed::new(s.compat(), LinesCodec::new());
|
|
|
while let Some(Ok(line)) = framed.next().await {
|
|
|
println!("Got: {}", line);
|
|
|
framed
|
|
|
.send(format!("Hello! I got '{}'", line))
|
|
|
.await
|
|
|
.unwrap();
|
|
|
}
|
|
|
|
|
|
Ok(())
|
|
|
}),
|
|
|
);
|
|
|
}
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
yamux 是一个类似 HTTP/2 内部多路复用机制的协议，可以让你在一个 TCP 连接上打开多个逻辑 yamux stream，而yamux stream 之间并行工作，互不干扰。
|
|
|
|
|
|
yamux crate 在实现的时候，使用了 futures 下的异步 IO 接口。但是当我们使用 tokio Listener 接受一个客户端，得到对应的 TcpStream 时，这个 TcpStream 使用的是 tokio 下的异步 IO 接口。所以我们需要 tokio\_util::compat 来协助接口的兼容。
|
|
|
|
|
|
在代码中，首先我用 stream.compat() 生成一个 Compat 结构，供 yamux Connection 使用：
|
|
|
|
|
|
```rust
|
|
|
let conn = Connection::new(stream.compat(), config, Mode::Server);
|
|
|
|
|
|
```
|
|
|
|
|
|
之后，拿到 yamux connection 下所有 stream 进行处理时，我们想用 tokio 的 Frame 和 Codec 一行行读取和写入，也就需要把使用 futures 异步接口的 yamux stream，转换成使用 tokio 接口的数据结构，这样就可以用在 Framed::new() 中：
|
|
|
|
|
|
```rust
|
|
|
let mut framed = Framed::new(s.compat(), LinesCodec::new());
|
|
|
|
|
|
```
|
|
|
|
|
|
如果你想运行这段代码，可以看这门课的 [GitHub repo](https://github.com/tyrchen/geektime-rust) 下的完整版，包括依赖以及客户端的代码。
|
|
|
|
|
|
### 实现异步 IO 接口
|
|
|
|
|
|
异步 IO 主要应用在文件处理、网络处理等场合，而这些场合的数据结构都已经实现了对应的接口，比如 File 或者 TcpStream，它们也已经实现了 AsyncRead / AsyncWrite。所以基本上，我们不用自己实现异步 IO 接口，只需要会用就可以了。
|
|
|
|
|
|
不过有些情况，我们可能会把已有的数据结构封装在自己的数据结构中，此时，也应该实现相应的异步 IO 接口（[代码](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=b53409b4ad26aaf078d4799bfe95f65c)）：
|
|
|
|
|
|
```rust
|
|
|
use anyhow::Result;
|
|
|
use pin_project::pin_project;
|
|
|
use std::{
|
|
|
pin::Pin,
|
|
|
task::{Context, Poll},
|
|
|
};
|
|
|
use tokio::{
|
|
|
fs::File,
|
|
|
io::{AsyncRead, AsyncReadExt, ReadBuf},
|
|
|
};
|
|
|
|
|
|
#[pin_project]
|
|
|
struct FileWrapper {
|
|
|
#[pin]
|
|
|
file: File,
|
|
|
}
|
|
|
|
|
|
impl FileWrapper {
|
|
|
pub async fn try_new(name: &str) -> Result<Self> {
|
|
|
let file = File::open(name).await?;
|
|
|
Ok(Self { file })
|
|
|
}
|
|
|
}
|
|
|
|
|
|
impl AsyncRead for FileWrapper {
|
|
|
fn poll_read(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>,
|
|
|
buf: &mut ReadBuf<'_>,
|
|
|
) -> Poll<std::io::Result<()>> {
|
|
|
self.project().file.poll_read(cx, buf)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() -> Result<()> {
|
|
|
let mut file = FileWrapper::try_new("./Cargo.toml").await?;
|
|
|
let mut buffer = String::new();
|
|
|
file.read_to_string(&mut buffer).await?;
|
|
|
println!("{}", buffer);
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
这段代码封装了 tokio::fs::File 结构，我们想读取内部的 file 字段，但又不想把 File 暴露出来，因此实现了 AsyncRead trait。
|
|
|
|
|
|
## Sink trait
|
|
|
|
|
|
在同步环境下往 IO 中发送连续的数据，可以一次性发送，也可以使用 Write trait 多次发送，使用起来并没有什么麻烦；但在异步 IO 下，做同样的事情，我们需要更方便的接口。因此异步IO还有一个比较独特的 Sink trait，它是一个用于发送一系列异步值的接口。
|
|
|
|
|
|
看 Sink trait 的定义：
|
|
|
|
|
|
```rust
|
|
|
pub trait Sink<Item> {
|
|
|
type Error;
|
|
|
fn poll_ready(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Self::Error>>;
|
|
|
fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error>;
|
|
|
fn poll_flush(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Self::Error>>;
|
|
|
fn poll_close(
|
|
|
self: Pin<&mut Self>,
|
|
|
cx: &mut Context<'_>
|
|
|
) -> Poll<Result<(), Self::Error>>;
|
|
|
}
|
|
|
|
|
|
pub trait SinkExt<Item>: Sink<Item> {
|
|
|
...
|
|
|
fn send(&mut self, item: Item) -> Send<'_, Self, Item> where Self: Unpin { ... }
|
|
|
...
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
和 Stream trait 不同的是，Sink trait 的 Item 是 trait 的泛型参数，而不是关联类型。**一般而言，当 trait 接受某个 input，应该使用泛型参数，比如 Add<Rhs>；当它输出某个 output，那么应该使用关联类型，比如 Future、Stream、Iterator 等**。
|
|
|
|
|
|
Item 对于 Sink 来说是输入，所以使用泛型参数是正确的选择。因为这也意味着，在发送端，可以发送不同类型的数据结构。
|
|
|
|
|
|
看上面的定义源码，Sink trait 有四个方法：
|
|
|
|
|
|
* poll\_ready()：用来准备 Sink 使其可以发送数据。只有 poll\_ready() 返回 Poll::Ready(Ok(())) 后，Sink 才会开展后续的动作。poll\_ready() 可以用来控制背压。
|
|
|
* start\_send()：开始发送数据到 Sink。但是start\_send() 并不保证数据被发送完毕，所以调用者要调用 poll\_flush() 或者 poll\_close() 来保证完整发送。
|
|
|
* poll\_flush()：将任何尚未发送的数据 flush 到这个 Sink。
|
|
|
* poll\_close()：将任何尚未发送的数据 flush 到这个 Sink，并关闭这个 Sink。
|
|
|
|
|
|
其中三个方法和 Item 是无关的，这会导致，如果不同的输入类型有多个实现，Sink的poll\_ready、poll\_flush 和 poll\_close 可能会有重复的代码。所以一般我们在使用 Sink 时，如果确实需要处理不同的数据类型，可以用 enum 将它们统一（感兴趣的话，可以进一步阅读这个[讨论](https://github.com/rust-lang/futures-rs/issues/623)）。
|
|
|
|
|
|
我们就用一个简单的 FileSink 的例子，看看如何实现这些方法。tokio::fs 下的 File 结构已经实现了 AsyncRead / AsyncWrite，我们只需要在 Sink 的几个方法中调用 AsyncWrite 的方法即可（[代码](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=b3aab166023b7478ccd947703f8f53cd)）：
|
|
|
|
|
|
```rust
|
|
|
use anyhow::Result;
|
|
|
use bytes::{BufMut, BytesMut};
|
|
|
use futures::{Sink, SinkExt};
|
|
|
use pin_project::pin_project;
|
|
|
use std::{
|
|
|
pin::Pin,
|
|
|
task::{Context, Poll},
|
|
|
};
|
|
|
use tokio::{fs::File, io::AsyncWrite};
|
|
|
|
|
|
#[pin_project]
|
|
|
struct FileSink {
|
|
|
#[pin]
|
|
|
file: File,
|
|
|
buf: BytesMut,
|
|
|
}
|
|
|
|
|
|
impl FileSink {
|
|
|
pub fn new(file: File) -> Self {
|
|
|
Self {
|
|
|
file,
|
|
|
buf: BytesMut::new(),
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
impl Sink<&str> for FileSink {
|
|
|
type Error = std::io::Error;
|
|
|
|
|
|
fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
|
Poll::Ready(Ok(()))
|
|
|
}
|
|
|
|
|
|
fn start_send(self: Pin<&mut Self>, item: &str) -> Result<(), Self::Error> {
|
|
|
let this = self.project();
|
|
|
eprint!("{}", item);
|
|
|
this.buf.put(item.as_bytes());
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
|
// 如果想 project() 多次，需要先把 self reborrow 一下
|
|
|
let this = self.as_mut().project();
|
|
|
let buf = this.buf.split_to(this.buf.len());
|
|
|
if buf.is_empty() {
|
|
|
return Poll::Ready(Ok(()));
|
|
|
}
|
|
|
|
|
|
// 写入文件
|
|
|
if let Err(e) = futures::ready!(this.file.poll_write(cx, &buf[..])) {
|
|
|
return Poll::Ready(Err(e));
|
|
|
}
|
|
|
// 刷新文件
|
|
|
self.project().file.poll_flush(cx)
|
|
|
}
|
|
|
|
|
|
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
|
let this = self.project();
|
|
|
// 结束写入
|
|
|
this.file.poll_shutdown(cx)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() -> Result<()> {
|
|
|
let file_sink = FileSink::new(File::create("/tmp/hello").await?);
|
|
|
// pin_mut 可以把变量 pin 住
|
|
|
futures::pin_mut!(file_sink);
|
|
|
file_sink.send("hello\\n").await?;
|
|
|
file_sink.send("world!\\n").await?;
|
|
|
file_sink.send("Tyr!\\n").await?;
|
|
|
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
对于 poll\_ready() 方法，直接返回 Poll::Ready(Ok(()))。
|
|
|
|
|
|
在 start\_send() 方法中，我们把传入的 item，写入 FileSink 的 BytesMut 中。然后在 poll\_flush() 时，我们拿到 buf，把已有的内容调用 [split\_to()](https://docs.rs/bytes/1.1.0/bytes/struct.BytesMut.html#method.split_to)，得到一个包含所有未写入文件的新 buffer。这个 buffer 和 self 无关，所以传入 poll\_write() 时，不会有对 self 的引用问题。
|
|
|
|
|
|
在写入文件后，我们再次调用 poll\_flush() ，确保写入的内容刷新到磁盘上。最后，在 poll\_close() 时调用 poll\_shutdown() 关闭文件。
|
|
|
|
|
|
这段代码虽然实现了 Sink trait，也展示了如何实现 Sink 的几个方法，但是这么简单的一个问题，处理起来还是颇为费劲。有没有更简单的方法呢？
|
|
|
|
|
|
有的。futures 里提供了 sink::unfold 方法，类似 stream::unfold，我们来重写上面的 File Sink 的例子（[代码](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=2582cfe2473f615b4e2f021893d738a3)）：
|
|
|
|
|
|
```rust
|
|
|
use anyhow::Result;
|
|
|
use futures::prelude::*;
|
|
|
use tokio::{fs::File, io::AsyncWriteExt};
|
|
|
|
|
|
#[tokio::main]
|
|
|
async fn main() -> Result<()> {
|
|
|
let file_sink = writer(File::create("/tmp/hello").await?);
|
|
|
// pin_mut 可以把变量 pin 住
|
|
|
futures::pin_mut!(file_sink);
|
|
|
if let Err(_) = file_sink.send("hello\\n").await {
|
|
|
println!("Error on send");
|
|
|
}
|
|
|
if let Err(_) = file_sink.send("world!\\n").await {
|
|
|
println!("Error on send");
|
|
|
}
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
/// 使用 unfold 生成一个 Sink 数据结构
|
|
|
fn writer<'a>(file: File) -> impl Sink<&'a str> {
|
|
|
sink::unfold(file, |mut file, line: &'a str| async move {
|
|
|
file.write_all(line.as_bytes()).await?;
|
|
|
eprint!("Received: {}", line);
|
|
|
Ok::<_, std::io::Error>(file)
|
|
|
})
|
|
|
}
|
|
|
|
|
|
```
|
|
|
|
|
|
可以看到，通过 unfold 方法，我们不需要撰写 Sink 的几个方法了，而且可以在一个返回 Future 的闭包中来提供处理逻辑，这就意味着我们可以不使用 poll\_xxx 这样的方法，直接在闭包中使用这样的异步函数：
|
|
|
|
|
|
```rust
|
|
|
file.write_all(line.as_bytes()).await?
|
|
|
|
|
|
```
|
|
|
|
|
|
你看，短短 5 行代码，就实现了刚才五十多行代码要表达的逻辑。
|
|
|
|
|
|
## 小结
|
|
|
|
|
|
今天我们学习了和异步 IO 相关的 Stream / Sink trait，以及和异步读写相关的 AsyncRead / AsyncWrite 等 trait。在学习异步 IO 时，很多内容都可以和同步 IO 的处理对比着学，这样事半功倍。
|
|
|
|
|
|
|
|
|
|
|
|
在处理异步 IO 时，底层的 poll\_xxx() 函数很难写，因为它的约束很多。好在有 pin\_project 这个项目，用宏帮我们解决了很多关于 Pin/Unpin 的问题。
|
|
|
|
|
|
一般情况下，我们不太需要直接实现 Stream / Sink / AsyncRead / AsyncWrite trait，如果的确需要，先看看有没有可以使用的辅助函数，比如通过 poll\_fn / unfold 创建 Stream、通过 unfold 创建 Sink。
|
|
|
|
|
|
### 思考题
|
|
|
|
|
|
我们知道 tokio:sync::mpsc 下有支持异步的 MPSC channel，生产者可以通过 send() 发送消息，消费者可以通过 recv() 来接收消息。你能不能为其封装 Sink 和 Stream 的实现，让 MPSC channel 可以像 Stream / Sink 一样使用？（提示：tokio-stream 有 ReceiverStream 的实现）。
|
|
|
|
|
|
欢迎在留言区分享你的思考和学习收获，感谢收听，恭喜你已经完成了rust学习的40次打卡，如果觉得有收获，也欢迎分享给你身边的朋友，邀他一起讨论。我们下节课见。
|
|
|
|