volo_grpc/client/
meta.rs

1use std::{net::SocketAddr, str::FromStr};
2
3use metainfo::{Backward, Forward};
4use volo::{Service, context::Context};
5
6use crate::{
7    Request, Response, Status,
8    context::ClientContext,
9    metadata::{
10        DESTINATION_METHOD, DESTINATION_SERVICE, HEADER_TRANS_REMOTE_ADDR, KeyAndValueRef,
11        MetadataKey, SOURCE_SERVICE,
12    },
13};
14
15#[derive(Clone)]
16pub struct MetaService<S> {
17    inner: S,
18}
19
20impl<S> MetaService<S> {
21    pub fn new(inner: S) -> Self {
22        Self { inner }
23    }
24}
25
26impl<T, U, S> Service<ClientContext, Request<T>> for MetaService<S>
27where
28    S: Service<ClientContext, Request<T>, Response = Response<U>, Error = Status>
29        + Send
30        + 'static
31        + Sync,
32    T: Send + 'static,
33{
34    type Response = S::Response;
35
36    type Error = S::Error;
37
38    async fn call(
39        &self,
40        cx: &mut ClientContext,
41        mut volo_req: Request<T>,
42    ) -> Result<Self::Response, Self::Error> {
43        let metadata = volo_req.metadata_mut();
44        _ = metainfo::METAINFO.with(|metainfo| {
45            let metainfo = metainfo.borrow_mut();
46
47            // persistents for multi-hops
48            if let Some(ap) = metainfo.get_all_persistents() {
49                for (key, value) in ap {
50                    let key = metainfo::HTTP_PREFIX_PERSISTENT.to_owned() + key;
51                    metadata.insert(MetadataKey::from_str(key.as_str())?, value.parse()?);
52                }
53            }
54
55            // transients for one-hop
56            if let Some(at) = metainfo.get_all_transients() {
57                for (key, value) in at {
58                    let key = metainfo::HTTP_PREFIX_TRANSIENT.to_owned() + key;
59                    metadata.insert(MetadataKey::from_str(key.as_str())?, value.parse()?);
60                }
61            }
62
63            // caller
64            metadata.insert(SOURCE_SERVICE, cx.rpc_info.caller().service_name().parse()?);
65
66            // callee
67            metadata.insert(
68                DESTINATION_SERVICE,
69                cx.rpc_info.callee().service_name().parse()?,
70            );
71            metadata.insert(DESTINATION_METHOD, cx.rpc_info.method().parse()?);
72
73            Ok::<(), Status>(())
74        });
75
76        let mut volo_resp = self.inner.call(cx, volo_req).await?;
77
78        let metadata = volo_resp.metadata_mut();
79        _ = metainfo::METAINFO.with(|metainfo| {
80            let mut metainfo = metainfo.borrow_mut();
81
82            // callee
83            if let Some(ad) = metadata.remove(HEADER_TRANS_REMOTE_ADDR) {
84                let maybe_addr = ad.to_str()?.parse::<SocketAddr>();
85                if let Ok(addr) = maybe_addr {
86                    cx.rpc_info_mut()
87                        .callee_mut()
88                        .set_address(volo::net::Address::from(addr));
89                }
90            }
91
92            // backward
93            let mut vec = Vec::with_capacity(metadata.len());
94            for key_and_value in metadata.iter() {
95                match key_and_value {
96                    KeyAndValueRef::Ascii(k, v) => {
97                        let k = k.as_str();
98                        let v = v.to_str()?;
99                        if k.starts_with(metainfo::HTTP_PREFIX_BACKWARD) {
100                            vec.push(k.to_owned());
101                            metainfo.strip_http_prefix_and_set_backward_downstream(k, v.to_owned());
102                        }
103                    }
104                    _ => unreachable!(),
105                }
106            }
107            for k in vec {
108                metadata.remove(k);
109            }
110
111            Ok::<(), Status>(())
112        });
113
114        Ok(volo_resp)
115    }
116}
volo_grpc/client/meta.rs

volo_grpc/client/
meta.rs
