下一站 - Ihcblog!

远方的风景与脚下的路 | 子站点:ihc.im

0%

在 Rust 中构建可靠的上下文传递组件

发表于

This article also has an English version.

本篇将介绍我在 Rust 中构建可靠的上下文传递组件的一些思考、设计与实现。我实现的 certain-map 已经开源（一年多以前开源的，近期做了更多改进，本文后续会介绍），欢迎使用！

项目地址：https://github.com/ihciah/certain-map

它解决了什么问题：

  1. 在跨组件传递上下文时，它可以借助编译器保证字段的存在性（即当某组件对 Context 中某字段存在读依赖时，前置组件必须写入过该字段，否则无法通过编译）
  2. 通用组件实现依赖的上下文可被定义为泛型参数并加以约束，这使组件实现更为通用，不耦合 Context 具体类型

注：虽然项目名字看起来是一种 map 实现，但其实这是一个基于过程宏生成的 struct。之所以命名为 certain-map 是因为其设计初衷是为了替换 TypeMap，并保证字段的存在性。

Service 抽象

如果你对此已经很熟悉，可以跳过这一节。

在 Rust 中，利用类型系统，我们很容易构建分层抽象，将复杂过程拆分为独立的组件，并在使用时将其组合。一个典型的例子是 tower Service（我也定义了一个面向纯 async 的 Service 及其周边工具：service-async）。

一个典型的 Service 定义如下：

1
2
3
4
5
6
trait Service<Request> {
    type Response;
    type Error;

    fn call(&self, req: Request) -> impl Future<Output = Result<Self::Response, Self::Error>>;
}

以我正在开发的通用网关框架为例，L7 能力可能构建在 L5 能力之上，而 L5 又是基于 L4 实现的。为了使逻辑充分解耦并可插拔，利用 Service 抽象我们可以实现下面的 Service：

  1. L4Svc: Service<(SocketAddr, T)>
  2. TLSSvc: Service<T> where T: AsyncRead + AsyncWrite
  3. H1Dispatcher: Service<T> where T: AsyncRead + AsyncWrite
  4. H1Svc: Service<http::Request<Bytes>, Response = http::Response<Bytes>>

在构建组件时，L4Svc 作为最外层组件，接收对端地址和连接，并将连接传递给下一层组件。而 H1Dispatcher 则需要实现一个循环，持续解析 HTTP 请求，并将请求交给 H1Svc 处理，最后将 Response 写回。

所以我们可以有类似下面的代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
struct L4Svc<T> {
    inner: T,
}
impl<T, R> Service<(SocketAddr, R)> for L4Svc<T>
where T: Service<R> {
    type Response = T::Response;
    type Error = T::Error;

    async fn call(&self, (addr, r): (SocketAddr, R)) -> Result<Self::Response, Self::Error> {
        println!("L4Svc handle: {:?}", addr);
        self.inner.call(r).await
    }
}

struct TLSSvc<T> {
    inner: T,
    tls: TlsAcceptor,
}
impl<T, R> Service<R> for TLSSvc<T>
where
    R: AsyncRead + AsyncWrite,
    T: Service<TlsStream>,
    T::Error: Into<TlsError>,
{
    type Response = T::Response;
    type Error = TlsError;

    async fn call(&self, r: R) -> Result<Self::Response, Self::Error> {
        // handshake and wrap the connection with TLS
        let tls_stream = self.tls.handshake(r).await?;
        // call inner service and map error
        self.inner.call(tls_stream).await.map_err(Into::into)
    }
}

struct H1Dispatcher<T> {
    inner: T,
}
impl<T, R> Service<R> for H1Dispatcher<T>
where
    R: AsyncRead + AsyncWrite,
    T: Service<http::Request<Bytes>, Response = http::Response<Bytes>>,
    T::Error: Into<HttpError>,
{
    // maybe the transmitted size, here is just a demo
    type Response = (u64, u64);
    type Error = HttpError;

    async fn call(&self, r: R) -> Result<Self::Response, Self::Error> {
        let (mut rb, mut wb) = (0, 0);
        loop {
            // read request
            let req = read_request(r).await?;
            // call inner service
            let resp = self.inner.call(req).await?;
            // write response
            write_response(r, resp).await?;
        }
        Ok((rb, wb))
    }
}

是不是非常容易？我们可以通过组合不同的 Service 来构建复杂的逻辑，这些 Service 可以是不同的人实现的，我们只需要明确每个 Service 接收什么样的 Request，并约束其 inner 实现什么样的 Service。

Context 传递

从前面的例子中我们可以看出，Service 的本质是对异步函数的抽象，其借助类型系统将异步函数的输入和输出约束在了一起。外层 Service 产生的信息都可以丢给内层 Service，并由内层 Service 丢给下一层 Service。

这种显式传递可以很好地标注 Request 和 Response 的类型变换，反映 Service 的逻辑和约束（例如 H1Dispatcher 接收 impl Read+Write 并约束其 inner svc 能够处理 http::Request<Bytes>，那么显然其内部主要实现 http 协议编解码）。

但有时候这种显式传递也会导致代码冗余，例如，我们期望每层都可以访问到请求者的 IP 地址（前面例子中的 SocketAddr），那么我们就需要在每一层都传递这个数据。如果我们有非常多需要跨层传递的信息，那么这种显式传递就会变得非常繁琐，并使组件不再通用，耦合度增加。

要解决这个问题，一个方法是将这类信息全部收集到一个结构内，并在每一层传递这个结构。这个结构可以是预定义的 struct，也可以是一个 type map。这样我们就可以在每一层通过这个结构访问到所有信息，并写入后续 Service 可能会用到的信息。

两种 Context 存储方式

Struct-based Context

1
2
3
4
5
6
7
8
9
10
11
12
13
14
struct MyContext {
    addr: Option<SocketAddr>,
    // more fields
}

impl<T, R> Service<(MyContext, SocketAddr, R)> for L4Svc<T>
where ...
{ ... }
impl<T, R> Service<(MyContext, R)> for TLSSvc<T>
where ...
{ ... }
impl<T, R> Service<(MyContext, R)> for H1Dispatcher<T>
where ...
{ ... }

MyContext 由用户定义，其中包含了所有需要传递的信息。此处“用户”指开发了一部分 Service 的开发者，它可能会使用其他开发者提供的 Service。

此时所有 Service 实现均需要接收 concreate type MyContext，并将其传递给下一层 Service。显然，基于这种方式实现的 Service 不具有任何可复用性，因为每个用户的 Context 都是不同的类型。

TypeMap-based Context

TypeMap 并非是一个特殊结构，其本质是一个 HashMap，key 是一个 type id，value 是一个 trait object。它可以存放任意的不同类型的数据。

1
2
3
4
5
6
7
8
9
impl<T, R> Service<(TypeMap, SocketAddr, R)> for L4Svc<T>
where ...
{ ... }
impl<T, R> Service<(TypeMap, R)> for TLSSvc<T>
where ...
{ ... }
impl<T, R> Service<(TypeMap, R)> for H1Dispatcher<T>
where ...
{ ... }

这个方式解决了 struct-based context 的问题，因为 TypeMap 足够通用，被所有 Service 强耦合是可以接受的。

hyper 使用这种形式传递一些内部使用的上下文信息，http::Request 和 http::Response 中的 Extensions 就是一个 TypeMap。

我们可以定义一个 NewType 来包装字段，以避免 key 冲突：

1
2
struct PeerAddr(pub SocketAddr);
struct LocalAddr(pub SocketAddr);

TypeMap 的缺陷

由前面的分析，我们可以看出 TypeMap 是一个非常通用的解决方案，但它也有一些缺陷：

  1. 无法保证 Value 的存在性，即在取值时无法保证 Value 的存在，这可能导致 panic 或书写多余的错误处理逻辑。
  2. 存在堆分配开销

这两个问题中，第一个问题尤为严重。我在字节跳动参与开发了内部的 RPC 框架（现在它开源为 Volo），在初期它发生了多次意外的 panic，这些 panic 是由于 TypeMap 的 Value 不存在导致的。外层 Service 由于疏漏导致在某些分支忘记插入某个 Context 值，而内层 Service 又强依赖并假设这个值一定存在，这就导致了 panic。

另外，Golang 中的 Context 也存在该问题。

可靠的 Context 传递

让我们反思一下，为什么 TypeMap 存在这种缺陷？这是因为 TypeMap 是运行时读写的 map，所以无法充分利用编译期检查能力。

要引入编译期检查解决该问题，一个 idea 是，我们利用不同的类型来描述不同存在性的状态，并有条件的为它们实现某些 trait。当字段存在性变化时，我们也需要变更类型。

我们需要解决三个问题：

  1. 如何设计 Trait 来约束字段的存在性并对其进行操作？
  2. 如何定义类型并随字段存在性变更类型？
  3. 如何为字段存在性不同的状态可选地实现 Trait？

设计 Trait 来操作字段

在我的这篇文章里也介绍了类似的设计：Rust HTTP 框架设计 - 以 Axum 0.6 为例

Rust 标准库提供了 AsRef, AsMut 等抽象，Linkerd 中也设计了 Param 描述类似抽象，axum 中设计了反向描述的 FromRef。

我们可以将这种设计发扬光大，提供更多的操作 Trait：

我发布了一个包来定义这些 trait：param；它的源代码在这里。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
pub trait Param<T> {
    fn param(&self) -> T;
}
pub trait ParamRef<T> {
    fn param_ref(&self) -> &T;
}
pub trait ParamMut<T> {
    fn param_mut(&mut self) -> &mut T;
}
pub trait ParamSet<T> {
    type Transformed;
    fn param_set(self, item: T) -> Self::Transformed;
}
pub trait ParamRemove<T> {
    type Transformed;
    fn param_remove(self) -> Self::Transformed;
}
...

当我们需要约束某个字段的存在性时，我们就可以声明 CX: Param<PeerAddr> 来约束它。

当涉及对字段的插入、删除等变更字段存在性的操作时，我们需要获得当前类型的所有权并返回新的类型，新的类型定义为该 trait 的关联类型。

定义类型并在操作时变更类型

struct 显然是最佳的存储方式，它足够高效。由于 field 个数和类型都取决于用户定义，所以我们需要基于用户的定义生成代码，过程宏是我们的最佳选择。

在开发过程宏前，我们需要手动写出一个目标结构，用于证明这种设计的可行性，并将其作为模板。过程宏的实现并不复杂，在此我只展示目标结构。

对于下面这种用户定义的结构：

1
2
3
4
struct MyContext {
    peer: PeerAddr,
    local: LocalAddr,
}

我们可以生成下面的结构：

1
2
3
4
struct MyContext<T1, T2> {
    peer: T1,
    local: T2,
}

我们可以定义两个辅助结构放在我们的 lib 中：

1
2
pub struct Occupied<T>(T);
pub struct Vacant;

这两个结构即可去填充 T1、T2 的位置。

例如，初始状态下其类型展开为：

1
2
3
4
struct MyContext {
    peer: Vacant,
    local: Vacant,
}

在写入 peer addr 后，其类型变为：

1
2
3
4
struct MyContext {
    peer: Occupied<PeerAddr>,
    local: Vacant,
}

在写操作时（写操作指对字段存在性有变更的操作），我们需要拿到当前类型的所有权，并拆分字段重组到新类型中。例如：

1
2
3
4
5
6
7
8
9
impl<T1, T2> ParamSet<PeerAddr> for MyContext<T1, T2> {
    type Transformed = MyContext<Occupied<PeerAddr>, T2>;
    fn param_set(self, item: T) -> Self::Transformed {
        MyContext {
            peer: Occupied(item),
            local: self.local,
        }
    }
}

为不同状态有条件地实现 Trait

有了前面的设计，我们就可以为不同状态有条件地实现 Trait 了。例如：

1
2
3
4
5
impl<T2> ParamRef<PeerAddr> for MyContext<Occupied<PeerAddr>, T2> {
    fn param(&self) -> PeerAddr {
        &self.peer.0
    }
}

MyContext<Vacant, T> 则不实现 ParamRef<PeerAddr>。

更进一步，我们可以定义 Available 和 MaybeAvailable 两个 trait 来表示字段可操作能力，这样可以减少生成代码的复杂度：

1
2
3
4
5
6
7
8
pub trait Available {
    fn get_ref<T>(&self) -> &T;
    ...
}
pub trait MaybeAvailable {
    fn set<T>(data: T) -> Occupied<T>;
    ...
}

其中，Available 表示字段一定存在（只为 Occupied<T> 实现），主要涉及读方法、take 方法等；MaybeAvailable 表示字段可能存在（所以可以为 Occupied<T> 和 Vacant 两个类型实现），含有覆盖写方法、remove 方法等。为了减少误用可能，我们还可以引入 sealed trait 来限制 trait 的实现。

实现效果

到目前为止，我们解决了前面三个问题，定义好了操作 trait、生成了目标结构、实现了字段存在性变更的操作。我们可以在编译期检查字段的存在性，避免 panic。

下面贴一个 certain-map 的 example(link)：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
use certain_map::{certain_map, Param, ParamRef, ParamRemove, ParamSet, ParamTake};

#[derive(Clone)]
pub struct UserName(String);

#[derive(Copy, Clone)]
pub struct UserAge(u8);

certain_map! {
    #[empty(MyCertainMapEmpty)]
    #[full(MyCertainMapFull)]
    #[style = "unfilled"]
    #[derive(Clone)]
    pub struct MyCertainMap {
        name: UserName,
        #[ensure(Clone)]
        age: UserAge,
    }
}

fn main() {
    let meta = MyCertainMap::new();

    // With #[default(MyCertainMapEmpty)] we can get an empty type.
    assert_type::<MyCertainMapEmpty>(&meta);

    // The following line compiles fail since there's no UserName in the map.
    // log_username(&meta);

    let meta = meta.param_set(UserName("ihciah".to_string()));
    // Now we can get it with certainty.
    log_username(&meta);

    let (meta, removed) = ParamTake::<UserName>::param_take(meta);
    assert_eq!(removed.0, "ihciah");
    // The following line compiles fail since the UserName is removed.
    // log_username(&meta);

    // We can also remove a type no matter if it exist.
    let meta = ParamRemove::<UserName>::param_remove(meta);

    let meta = meta.param_set(UserAge(24));
    // we can get ownership of fields with #[ensure(Clone)]
    log_age(&meta);
}

fn log_username<T: ParamRef<UserName>>(meta: &T) {
    println!("username: {}", meta.param_ref().0);
}

fn log_age<T: Param<UserAge>>(meta: &T) {
    println!("user age: {}", meta.param().0);
}

fn assert_type<T>(_: &T) {}

如果我们在插入 UserName 前，或删除 UserName 后读取它，则编译器会报错，无法通过编译。

该 example 中，context 结构的类型做了如下变化：

  1. 初始状态下，MyCertainMap 是一个空的 struct，没有任何字段，所以对应 MyCertainMap<Vacant, Vacant>，其 size 是 0。
  2. 插入 UserName 后，MyCertainMap 变为 MyCertainMap<Occupied<UserName>, Vacant>，其 size 等同于 UserName。
  3. Take UserName 后，MyCertainMap 变为 MyCertainMap<Vacant, Vacant>，其 size 又变为 0。
  4. 插入 UserAge 后，MyCertainMap 变为 MyCertainMap<Vacant, Occupied<UserAge>>，其 size 等同于 UserAge。

到此为止，该设计已经完美地解决了我们的问题，它可以保证只要可以编译，则不可能有字段不存在的情况。另外，由于引入了 param 系列 trait，我们可以在定义 Service 时将 Context 作为泛型参数，解耦其 concreate 类型，让组件更通用。

更高效的 Context 传递

我将 certain-map（v0.2）集成在 MonoLake 中，这是我在字节跳动开发的一个基于 Monoio 的通用网关框架（尚未开源，但预计会在近期开源）。

在对其做压力测试和 perf 时，我观察到一些相关的内存拷贝开销。这些开销有两个来源：

  1. 由于字段存在性变更时 struct 类型和 size 都存在变更，所以需要拆分字段并重组到新类型中，这会导致一定的栈拷贝。
  2. 当外层 Service 传递 Context 到内层时，也可能存在一定的拷贝开销（该开销是否实际存在取决于 async 生成器是否做了该优化）。

为了解决这个问题，我提出了一个新的设计：将存储与状态分离，存储部分预分配全部字段空间，并传递状态和存储的引用构成的一个复合结构。这样，当状态变更时，只需要修改状态，而存储不移动，这样就避免了上述的栈拷贝开销。

这个设计的实现在 certain-map 的 v0.3 版本中。

存储结构与状态结构

存储结构是一个预分配的 struct，其不感知状态（即字段的存在性）。例如，对于下面这种用户定义的结构：

1
2
3
4
struct MyContext {
    peer: PeerAddr,
    local: LocalAddr,
}

其实际的存储结构是（下文称其为 Store 结构）：

1
2
3
4
pub struct MyContextStorage {
    peer: MaybeUninit<PeerAddr>,
    local: MaybeUninit<LocalAddr>,
}

而状态结构是一个零大小的 struct，泛型参数表示存储结构中对应字段的存在性（下文称其为 State 结构）：

1
2
3
4
pub struct MyContextState<T1, T2> {
    peer: PhantomData<T1>,
    local: PhantomData<T2>,
}

最终传递的结构是一个复合结构（下文称其为 Handler 结构；该结构实质上等价于单个引用，所以传递成本非常低）：

1
2
3
4
pub struct MyContextHandler<'a, T1, T2> {
    storage: &'a mut MyContextStorage,
    state: MyContextState<T1, T2>,
}

在 trait 层面我们依旧可以采用前面的设计，我们需要为 MyContextHandler 实现 Param 等 trait；Available/MaybeAvailable 设计也可以继续沿用（函数签名需要修改）。

我们还需要一些新的 trait 来提供生成 Handler 结构等能力：

1
2
3
4
pub trait Handler {
    type Hdr<'a>;
    fn handler<'a>(&'a mut self) -> Self::Hdr<'a>;
}

最后，我们还需要为 Handler 结构实现 Drop 方法，以保证存储结构的字段被正确 drop。

优势与问题

对比前面的设计，这种设计有以下优势：

  1. 存储结构预分配在栈上，在可以快速访问的同时避免了栈拷贝开销（对比 TypeMap 的方案则节约了堆开销），对 CPU 缓存也更友好。
  2. 传递的结构是一个引用，所以传递成本非常低。可以避免不必要的栈拷贝开销。
  3. 当状态变更时，只需要修改状态，而存储不移动，这样就避免了栈拷贝开销。

但是，这个设计也引入了新的问题：

  1. 用户需要感知 Store 结构和 Handler 结构
  2. 生命周期管理更加复杂（下文中会介绍遇到的问题与解法）
  3. clone 需要特殊实现

第一个问题并不是一个大问题，用户只需要创建这个存储，并生成其 Handler，之后就可以像之前使用 Context 结构一样使用 Handler。

下面我们会展开讨论后面两个问题。

生命周期与约束定义

由于 Service 抽象下调用往往是嵌套的，所以理论上讲，用户将 Store 结构放在外层 Service 的栈上，并将其 Handler 传递给内层 Service，这样是没有问题的。

但实际上，在实现 Service 时，由于不绑定 Context 的 concreate type，所以基于 trait 的操作方式需要关心生命周期。

在这个例子中，我们实现了几个简单的 Service 并将其组合：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
struct Add1<T>(T);

impl<T, CX> Service<(u8, CX)> for Add1<T>
where
    T: Service<(u8, CX::Transformed)>,
    CX: ParamSet<RawBeforeAdd>,
{
    type Response = T::Response;
    type Error = T::Error;

    fn call(
        &self,
        (num, cx): (u8, CX),
    ) -> impl Future<Output = Result<Self::Response, Self::Error>> {
        self.0.call((num + 1, cx.param_set(RawBeforeAdd(num))))
    }
}

struct Mul2<T>(T);

impl<T, CX> Service<(u8, CX)> for Mul2<T>
where
    T: Service<(u8, CX::Transformed)>,
    CX: ParamSet<RawBeforeMul>,
{
    type Response = T::Response;
    type Error = T::Error;

    fn call(
        &self,
        (num, cx): (u8, CX),
    ) -> impl Future<Output = Result<Self::Response, Self::Error>> {
        self.0.call((num * 2, cx.param_set(RawBeforeMul(num))))
    }
}

struct Identical;

impl<CX> Service<(u8, CX)> for Identical
where
    CX: ParamRef<RawBeforeAdd> + ParamRef<RawBeforeMul>,
{
    type Response = u8;
    type Error = Infallible;

    async fn call(&self, (num, cx): (u8, CX)) -> Result<Self::Response, Self::Error> {
        println!(
            "num before add: {}",
            ParamRef::<RawBeforeAdd>::param_ref(&cx).0
        );
        println!(
            "num before mul: {}",
            ParamRef::<RawBeforeMul>::param_ref(&cx).0
        );
        println!("num: {num}");
        Ok(num)
    }
}

可以验证其正确性：

1
2
3
let svc = Add1(Mul2(Identical));
let mut store = MyCertainMap::new();
assert_eq!(svc.call((2, store.handler())).await.unwrap(), 6);

调用时需要创建 Store 结构，并生成 Handler 后作为 CX 作为 Request 的一部分传递。

但我并不满足于此，Context 自动注入也应当实现为一个 Service！

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
struct CXSvc<CXStore, T> {
    inner: T,
    cx: PhantomData<CXStore>,
}

impl<CXStore, T> CXSvc<CXStore, T> {
    fn new(inner: T) -> Self {
        Self {
            inner,
            cx: PhantomData,
        }
    }
}

impl<CXStore, T> CXSvc<CXStore, T> {
    async fn call<R>(
        &self,
        num: R,
    ) -> Result<
        <T as Service<(R, <CXStore as Handler>::Hdr<'_>)>>::Response,
        <T as Service<(R, <CXStore as Handler>::Hdr<'_>)>>::Error,
    >
    where
        CXStore: Handler + Default + 'static,
        for<'a> T: Service<(R, CXStore::Hdr<'a>)>,
    {
        let mut store = CXStore::default();
        let hdr = store.handler();
        self.inner.call((num, hdr)).await
    }
}

impl<CXStore, T, R> Service<R> for CXSvc<CXStore, T>
where
    CXStore: Handler + Default + 'static,
    for<'a> T: Service<(R, CXStore::Hdr<'a>)>,
{
    type Response = <T as Service<(R, <CXStore as Handler>::Hdr<'_>)>>::Response;
    type Error = <T as Service<(R, <CXStore as Handler>::Hdr<'_>)>>::Error;

    async fn call(&self, num: R) -> Result<Self::Response, Self::Error> {
        let mut store = CXStore::default();
        let hdr = store.handler();
        self.inner.call((num, hdr)).await
    }
}

HRTB is your friend!

此处 Handler 结构需要写明 lifetime，而结构定义与 Service 上均无 lifetime。此时需要利用 HRTB 来约束 inner svc（类型 T）可以接受任意 lifetime 的 Handler。

看起来一切正常？并没有！

cannot return value referencing local variable store returns a value referencing data owned by the current function

其原因是，编译器无法证明返回的 Response 或 Error 中不包含对 store 的引用。如果包含，那么因为 store 是在函数内部创建的，其引用将会在函数结束后失效，则不能返回。

为了解决这个问题，我们需要一些小技巧：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
impl<CXStore, T, R, RESP, ERR> Service<R> for CXSvc<CXStore, T>
where
    CXStore: Handler + Default + 'static,
    for<'a> T: Service<(R, CXStore::Hdr<'a>), Response = RESP, Error = ERR>,
{
    type Response = RESP;
    type Error = ERR;

    async fn call(&self, num: R) -> Result<Self::Response, Self::Error> {
        let mut store = CXStore::default();
        let hdr = store.handler();
        self.inner.call((num, hdr)).await
    }
}

这种实现即可正确编译。但是为什么呢？

通过将 Response 和 Error 作为泛型参数，我们可以将其类型与 Handler 的生命周期解除绑定，即对于任意 lifetime 的 Handler 均返回相同的 Response 和 Error，这样编译器就可以证明 Response 和 Error 不会包含对 store 的引用。

可以验证现在的实现符合预期：

1
2
let svc = CXSvc::<MyCertainMap, _>::new(svc);
assert_eq!(svc.call(2).await.unwrap(), 6);

实现 Fork

在 0.2 版本中，我们直接 derive Clone 即可实现 Clone。但在 0.3 版本中，存储与状态分离，我们需要实现特殊的 trait 来表达 Clone 语义，此处称之为 Fork。

我将 Clone 语义拆分为下面两个行为，一个是使用 Handler 拷贝 Store 和 State（必须使用 Handler，否则不知道 State 中字段的存在性），另一个是使用 Store 和 State 构造 Handler：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
pub trait Fork {
    type Store;
    type State;
    fn fork(&self) -> (Self::Store, Self::State);
}

pub trait Attach<Store> {
    type Hdr<'a>
    where
        Store: 'a;
    /// # Safety
    /// The caller must make sure the attached map has the data of current state.
    unsafe fn attach(self, store: &mut Store) -> Self::Hdr<'_>;
}

在前面例子的基础上，我们可以新实现一个 Service 来测试 Fork。这个 Svc 会约束 Req: Copy（实际调用时是数字）以及 Resp: Add<Output = Resp>（实际调用时也是数字），并将其实现为使用 Req 调用两次 inner svc，并将结果相加：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
struct DupSvc<T>(T);
impl<T, R, CXIn, CXStore, CXState, Resp, Err, HDR> Service<(R, CXIn)> for DupSvc<T>
where
    R: Copy,
    Resp: Add<Output = Resp>,
    CXIn: Fork<Store = CXStore, State = CXState>,
    CXStore: 'static,
    for<'a> CXState: Attach<CXStore, Hdr<'a> = HDR>,
    T: Service<(R, HDR), Response = Resp, Error = Err>,
{
    type Response = Resp;
    type Error = Err;

    async fn call(&self, (req, ctx): (R, CXIn)) -> Result<Self::Response, Self::Error> {
        // fork ctx
        let (mut store, state) = ctx.fork();
        let forked_ctx = unsafe { state.attach(&mut store) };
        let r1 = self.0.call((req, forked_ctx)).await?;

        // fork ctx
        let (mut store, state) = ctx.fork();
        let forked_ctx = unsafe { state.attach(&mut store) };
        let r2 = self.0.call((req, forked_ctx)).await?;

        Ok(r1 + r2)
    }
}

只写这个 Svc 定义看起来编译可以通过，但预期正常执行吗？并不，在调用时编译器会 complain 约束不满足！

我们需要认真搞清楚 HRTB 的使用姿势。我们需要将代码修改为下面这样：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
impl<T, R, CXIn, CXStore, CXState, Resp, Err> Service<(R, CXIn)> for DupSvc<T>
where
    R: Copy,
    Resp: Add<Output = Resp>,
    CXIn: Fork<Store = CXStore, State = CXState>,
    CXStore: 'static,
    for<'a> CXState: Attach<CXStore>,
    for<'a> T: Service<(R, <CXState as Attach<CXStore>>::Hdr<'a>), Response = Resp, Error = Err>,
{
    type Response = Resp;
    type Error = Err;

    async fn call(&self, (req, ctx): (R, CXIn)) -> Result<Self::Response, Self::Error> {
        // fork ctx
        let (mut store, state) = ctx.fork();
        let forked_ctx = unsafe { state.attach(&mut store) };
        let r1 = self.0.call((req, forked_ctx)).await?;

        // fork ctx
        let (mut store, state) = ctx.fork();
        let forked_ctx = unsafe { state.attach(&mut store) };
        let r2 = self.0.call((req, forked_ctx)).await?;

        Ok(r1 + r2)
    }
}

它们的差别在于这里：

1
2
3
4
5
6
// fail
    for<'a> CXState: Attach<CXStore, Hdr<'a> = HDR>,
    T: Service<(R, HDR), Response = Resp, Error = Err>,
// pass
    for<'a> CXState: Attach<CXStore>,
    for<'a> T: Service<(R, <CXState as Attach<CXStore>>::Hdr<'a>), Response = Resp, Error = Err>,

在错误的 case 中，我们将 HDR 作为泛型参数，T: Service<(R, HDR)> 实质上是一个存在性约束（取交集），但对于入参，我们应当使用 HRTB 来约束其生命周期，即约束所有可能的 Hdr<'a>（取并集）。

在正确使用 HRTB 后，我们可以正确编译并执行这个 Svc：

1
2
let svc = CXSvc::<MyCertainMap, _>::new(DupSvc(Add1(Mul2(Identical))));
assert_eq!(svc.call(2).await.unwrap(), 12);

总结

在这篇文章中，我介绍了如何设计一个可靠的 Context 传递方案，并给出了两种实现以及实现时的设计思路。

在 certain-map 0.3 版本中，我将存储和状态分离，通过 Handler 结构传递，避免了栈拷贝开销，提高了性能。但这提高了生命周期管理的复杂度，所以我给出了使用 HRTB 解决该问题的正确方式。

最后欢迎大家在项目中使用这个组件，也欢迎提出建议和改进。

欢迎关注我的其它发布渠道