Next Stop - Ihcblog!

Some creations and thoughts sharing | sub site:ihc.im

0%

Building Reliable Context-Passing Components in Rust

Posted on

This article also has an Chinese version.

This article discusses my thoughts, designs, and implementations on creating a reliable context-passing component in Rust. My project, certain-map, has been open-sourced (it was first released over a year ago, with more improvements made recently, which I will discuss later in this article). Feel free to use it!

Project URL: https://github.com/ihciah/certain-map

What problem does it solve:

  1. When passing context across components, it can leverage the compiler to ensure the existence of fields (i.e., when a component has a read dependency on a field in the Context, the preceding component must have written to that field, otherwise it will not compile).
  2. The context required by generic components can be defined as a generic parameter and constrained, which makes the component implementation more generic and not coupled to a specific type of Context.

Note: Although the project name seems to suggest a map implementation, it is actually a struct generated using procedural macros. The reason it is named certain-map is that it was originally designed to replace TypeMap and ensure the existence of fields.

Service Abstraction

If you are already familiar with this, you may skip this section.

In Rust, thanks to the type system, it is easy to build layered abstractions, break down complex processes into independent components, and combine them when needed. A typical example is the tower Service (I have also defined an async-oriented Service and its auxiliary tools: service-async).

A typical Service is defined as follows:

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>>;
}

Using the general-purpose gateway framework I am currently developing as an example, L7 capabilities may be built on top of L5 capabilities, which in turn are based on L4 implementations. To ensure that the logic is fully decoupled and pluggable, using the Service abstraction allows us to implement the following Services:

  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>>

When building components, L4Svc acts as the outermost layer, receiving peer addresses and connections and passing these connections to the next layer. The H1Dispatcher, on the other hand, needs to implement a loop that continuously parses HTTP requests and passes them to H1Svc for processing, finally writing the Response back.

Thus, we can have code similar to the following:

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))
    }
}

Isn’t it very straightforward? We can build complex logic by composing different Services, which can be implemented by different people. We just need to clearly define what kind of Request each Service accepts and constrain its inner to implement a certain type of Service.

Context Passing

As can be seen from the previous examples, the essence of a Service is the abstraction of asynchronous functions, where the type system constrains the inputs and outputs of these functions together. Information generated by an outer Service can be passed on to an inner Service, and from the inner Service to the next layer.

This explicit passing can effectively show the types of transformations between Requests and Responses, reflecting the logic and constraints of the Service (for example, H1Dispatcher receives impl Read+Write and constrains its inner svc to handle http::Request<Bytes>, obviously focusing on HTTP protocol encoding and decoding internally).

However, sometimes this explicit passing can lead to code redundancy. For instance, if each layer needs access to the requester’s IP address (the SocketAddr in the previous example), then this data must be passed at every layer. If there is a lot of information that needs to be passed across layers, this explicit passing can become cumbersome and make components less general and more coupled.

One solution to this problem is to collect all such information into a single structure and pass this structure through each layer. This structure could be a predefined struct or a type map. This way, at each layer, all information can be accessed and information that the subsequent Services might need can be inputted.

Two Methods of Context Storage

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 is defined by the user and contains all necessary information for passing. Here, the “user” refers to the developer who has implemented a part of the Service, and they may use Services provided by other developers.

At this point, all Service implementations need to receive the concrete type MyContext and pass it on to the next layer of Service. Clearly, Services implemented in this way lack reusability, as each user’s Context is a different type.

TypeMap-based Context

TypeMap is not a special structure; fundamentally, it is a HashMap where the key is a type id and the value is a trait object. It can store data of any different type.

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 ...
{ ... }

This method addresses the problem of struct-based context, as TypeMap is sufficiently general and can be acceptably tightly coupled with all services.

The hyper library uses this form to pass certain internally used context information, with Extensions in http::Request and http::Response functioning as a TypeMap.

We can define new types to wrap fields, in order to avoid key conflicts:

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

Flaws of TypeMap

From the analysis above, we can see that TypeMap is a very versatile solution; however, it has several shortcomings:

  1. Inability to Guarantee the Presence of Values: When retrieving values, there is no guarantee that the value will exist, which can lead to panic or necessitate additional error handling logic.
  2. Heap Allocation Overhead

Among these issues, the first one is particularly severe. I participated in the development of an internal RPC framework at ByteDance, which has now been open-sourced as Volo. In its early stages, it experienced several unexpected panics, which were due to the absence of values in the TypeMap. Outer layer Services, due to oversight, forgot to insert certain context values in some branches, while inner layer Services strongly depended on and assumed the presence of these values, leading to panic.

Reliable Context Passing

Let’s reflect on why TypeMap holds these defects? The reason lies in its nature as a runtime read-write general map, making it incapable of fully leveraging compile-time check capabilities.

To introduce compile-time checks to solve this issue, one idea is to use different types to describe states of existence, and conditionally implement certain traits for them. When the existence of fields changes, we would need to change the type accordingly.

We need to address three questions:

  1. How to design Traits to constrain the presence of fields and manipulate them?
  2. How to define types that change as the presence of fields changes?
  3. How to optionally implement Traits for different states of field presence?

Designing Traits to Manipulate Fields

This design is similarly discussed in my article: Rust HTTP Framework Design - Taking Axum 0.6 as an Example

The Rust standard library provides abstractions like AsRef and AsMut. Linkerd has designed the Param to describe similar abstractions, and Axum has designed the reverse trait FromRef.

We can further develop this idea by providing more operational Traits:

I have published a package to define these traits: param; its source code can be found here.

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;
}
...

When we need to constrain the presence of a specific field, we can declare CX: Param<PeerAddr> to enforce this constraint.

When it comes to operations such as inserting or deleting fields, which change the field’s existence, we need to obtain ownership of the current type and return a new type. This new type is defined as the associated type of that trait.

Defining Types and Changing Types During Operations

Struct is clearly the best storage method as it is highly efficient. Since the number and types of fields depend on user definitions, we need to generate code based on user specifications. A procedural macro is our best choice for this.

Before developing a procedural macro, we need to manually write out a target structure to validate the feasibility of this design and use it as a template. The implementation of the procedural macro is not complicated, so here I will just show the target code.

Consider the following user-defined structure:

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

We can generate the following structure:

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

We can define two helper structures to be included in our library:

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

These two structures can be used to fill the positions of T1 and T2.

For example, initially, their types expand to:

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

After writing the peer addr, the type becomes:

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

During write operations (which refer to operations that change the existence of fields), we need to take ownership of the current type and reassemble the fields into the new type. For example:

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,
        }
    }
}

Conditionally Implementing Traits for Different States

With the design previously discussed, we can now conditionally implement traits for different states. For example:

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

MyContext<Vacant, T> does not implement ParamRef<PeerAddr>.

Furthermore, we can define two traits, Available and MaybeAvailable, to represent the operability of fields, which can reduce the complexity of the generated code:

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 indicates that a field definitely exists (implemented only for Occupied<T>), primarily involving read methods, take methods, etc.; MaybeAvailable indicates that a field might exist (thus it can be implemented for both Occupied<T> and Vacant), including overwrite write methods, remove methods, etc. To reduce the possibility of misuse, we can also introduce sealed traits to restrict the implementation of traits.

Implementation Results

So far, we have addressed the first three issues, defined the operational traits, generated the target structures, and implemented operations for changing the existence of fields. We can check the existence of fields at compile time to avoid panic.

Below is an example from certain-map (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) {}

If we attempt to read UserName before inserting it, or after deleting it, the compiler will throw an error, preventing the code from compiling.

In the given example, the type of the context structure undergoes the following changes:

  1. Initially, MyCertainMap is an empty struct without any fields, corresponding to MyCertainMap<Vacant, Vacant>, and its size is 0.
  2. After inserting UserName, MyCertainMap changes to MyCertainMap<Occupied<UserName>, Vacant>, and its size becomes equivalent to that of UserName.
  3. After taking UserName, MyCertainMap reverts to MyCertainMap<Vacant, Vacant>, and its size returns to 0.
  4. After inserting UserAge, MyCertainMap changes to MyCertainMap<Vacant, Occupied<UserAge>, and its size becomes equivalent to that of UserAge.

With this design, we have perfectly solved our problem; it ensures that if the code compiles, it is impossible for the fields to not exist. Furthermore, by introducing the param series of traits, we can use Context as a generic parameter when defining Service, decoupling its concrete type and making the components more versatile.

More Efficient Context Passing

I integrated certain-map (v0.2) into MonoLake, which is a generic gateway framework based on Monoio that I developed at ByteDance. It has not yet been open-sourced but is expected to be in the near future.

During stress testing and performance profiling, I observed some related memory copying overhead. There are two sources of this overhead:

  1. As the field existence changes, the struct type and size also change, necessitating the splitting of fields and recombination into the new type, which leads to a certain amount of stack copying.
  2. When an outer layer Service passes the Context to an inner layer, there can also be some copying overhead (whether this overhead actually exists depends on whether the async generator has optimized for this).

To address this issue, I proposed a new design: separate storage from state. The storage part pre-allocates space for all fields, and a composite structure made up of references to the state and storage is passed. Thus, when the state changes, only the state is modified, and the storage does not move, thereby avoiding the aforementioned stack copying overhead.

This design has been implemented in version 0.3 of certain-map.

Storage Structure and State Structure

The storage structure is a pre-allocated struct that is unaware of the state (i.e., the presence of fields). For example, for a user-defined structure like:

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

The actual storage structure (referred to below as the Store structure) is:

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

The state structure, on the other hand, is a zero-sized struct, with generic parameters indicating the presence of corresponding fields in the storage structure (referred to below as the State structure):

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

The final structure passed is a composite structure (referred to below as the Handler structure; this structure is essentially equivalent to a single reference, thus very low in passing cost):

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

At the trait level, we can still use the preceding design, where we need to implement traits like Param for MyContextHandler; the Available/MaybeAvailable design can also be continued (though function signatures need modification).

We also need some new traits to provide capabilities like generating the Handler structure:

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

Finally, we need to implement the Drop method for the Handler structure to ensure that the fields within the storage structure are properly dropped.

Advantages and Issues

Compared to previous designs, this design offers the following advantages:

  1. The storage structure is pre-allocated on the stack, which allows for quick access while avoiding stack copy overhead (compared to the TypeMap approach, it also saves on heap expenses), and is more friendly to CPU caches.
  2. The passed structure is a reference, so the passing cost is very low. This avoids unnecessary stack copy overhead.
  3. When the state changes, only the state needs to be modified, while the storage remains unmoved, thus avoiding stack copying overhead.

However, this design also introduces new problems:

  1. Users need to be aware of the Store structure and the Handler structure.
  2. Lifecycle management becomes more complex (the problems and solutions encountered will be discussed later).
  3. Cloning requires a special implementation.

The first issue is not a major problem; users only need to create the storage and generate its Handler, after which they can use the Handler just like they previously used the Context structure.

We will now proceed to discuss the last two issues in more detail.

Lifecycle and Constraint Definitions

Since service abstraction typically involves nested calling, in theory, placing the Store structure on the stack of the outer layer service and passing its Handler to the inner layer service should pose no issues.

However, in practical terms, when implementing services, due to not binding to the concrete type of Context, operations based on traits need to be mindful of lifetimes.

In this example, we implemented a few simple services and combined them:

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)
    }
}

To verify its correctness:

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

When invoking, it is necessary to create the Store structure, and after generating the Handler, pass it as part of the Request using CX.

However, I am not satisfied with just this; the automatic injection of Context should also be implemented as a Service too!

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!

In this case, the Handler structure needs to specify the lifetime, while there is no lifetime defined on the structure and the Service itself. At this point, it is necessary to use HRTB (Higher-Ranked Trait Bounds) to constrain the inner service (type T) to accept Handlers with any lifetime.

Does everything look normal? Not quite!

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

The reason is that the compiler cannot prove that the returned Response or Error does not contain a reference to the store. If it does contain such a reference, since the store is created inside the function, its reference will become invalid after the function ends, and therefore cannot be returned.

To solve this problem, we need to employ some tricks:

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
    }
}

This implementation allows for correct compilation. But why is that?

By making the Response and Error generic parameters, we decouple their types from the lifetime of the Handler. This means that for any lifetime of the Handler, the same Response and Error are returned, allowing the compiler to prove that the Response and Error do not contain references to the store.

Now, we can verify that the current implementation meets the expectations:

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

Implementing Fork

In version 0.2, we could directly derive Clone to implement cloning. However, in version 0.3, with storage and state separated, we need to implement a special trait to convey the semantics of Clone, which we refer to as Fork.

I have divided the Clone semantics into the following two behaviors: one is using the Handler to copy the Store and State (the Handler must be used, otherwise the existence of fields in State is unknown), and the other is constructing the Handler using the Store and State:

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<'_>;
}

Building on the previous example, we can implement a new Service to test Fork. This Service will constrain Req: Copy (actual calls use numbers) and Resp: Add<Output = Resp> (actual also use numbers), and implement it by invoking the inner service twice with the Req, and then adding the results:

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)
    }
}

Just writing this Svc definition might seem to compile, but will it perform as expected? No, it will not. During invocation, the compiler will complain about unsatisfied constraints!

We need to carefully understand how to properly use HRTB. We need to modify the code as follows:

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)
    }
}

The difference lies here:

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>,

In the erroneous case, we treat HDR as a generic parameter. T: Service<(R, HDR)> effectively poses an existential constraint (intersection), but for parameters, we should use HRTB to constrain their lifetimes, that is, to constrain over all possible Hdr<'a> (union).

After correctly using HRTB, we can compile and execute this Svc correctly:

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

Summary

In this article, I introduced how to design a reliable context passing scheme and provided two implementations along with their design rationale.

In version 0.3 of certain-map, I separated storage and state by passing them through a Handler structure, avoiding stack copy overhead and enhancing performance. However, this increased the complexity of lifetime management, so I presented the correct way to use HRTB to solve this issue.

Finally, I welcome everyone to use this component in your projects and invite suggestions and improvements.

Welcome to my other publishing channels