sova_grpc/
router.rs

1//! In-process unary method registry.
2
3use crate::error::GrpcError;
4use bytes::Bytes;
5use serde::{de::DeserializeOwned, Serialize};
6use sova_core::extend::BoxFuture;
7use sova_core::Request;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11type BodyHandler = Arc<dyn Fn(Bytes) -> BoxFuture<Result<Bytes, GrpcError>> + Send + Sync>;
12type CtxHandler = Arc<dyn Fn(Request, Bytes) -> BoxFuture<Result<Bytes, GrpcError>> + Send + Sync>;
13
14#[derive(Clone)]
15enum Handler {
16    Body(BodyHandler),
17    Ctx(CtxHandler),
18}
19
20#[derive(Clone, Default)]
21pub struct MethodRouter {
22    handlers: Arc<std::sync::RwLock<HashMap<String, Handler>>>,
23}
24
25impl MethodRouter {
26    pub fn new() -> Self {
27        Self::default()
28    }
29
30    pub fn unary<Req, Res, F, Fut>(&self, method: impl Into<String>, f: F) -> &Self
31    where
32        Req: DeserializeOwned + Send + 'static,
33        Res: Serialize + Send + 'static,
34        F: Fn(Req) -> Fut + Send + Sync + 'static,
35        Fut: std::future::Future<Output = Result<Res, GrpcError>> + Send + 'static,
36    {
37        let method = method.into();
38        let f = Arc::new(f);
39        let handler: BodyHandler = Arc::new(move |body| {
40            let f = Arc::clone(&f);
41            Box::pin(async move {
42                let req: Req =
43                    serde_json::from_slice(&body).map_err(|e| GrpcError::Decode(e.to_string()))?;
44                let res = f(req).await?;
45                let bytes =
46                    serde_json::to_vec(&res).map_err(|e| GrpcError::Decode(e.to_string()))?;
47                Ok(Bytes::from(bytes))
48            })
49        });
50        self.handlers
51            .write()
52            .unwrap()
53            .insert(method, Handler::Body(handler));
54        self
55    }
56
57    /// Unary handler with access to the incoming HTTP [`Request`] (auth, state, headers).
58    pub fn unary_with_request<Req, Res, F, Fut>(&self, method: impl Into<String>, f: F) -> &Self
59    where
60        Req: DeserializeOwned + Send + 'static,
61        Res: Serialize + Send + 'static,
62        F: Fn(Request, Req) -> Fut + Send + Sync + 'static,
63        Fut: std::future::Future<Output = Result<Res, GrpcError>> + Send + 'static,
64    {
65        let method = method.into();
66        let f = Arc::new(f);
67        let handler: CtxHandler = Arc::new(move |http_req, body| {
68            let f = Arc::clone(&f);
69            Box::pin(async move {
70                let req: Req =
71                    serde_json::from_slice(&body).map_err(|e| GrpcError::Decode(e.to_string()))?;
72                let res = f(http_req, req).await?;
73                let bytes =
74                    serde_json::to_vec(&res).map_err(|e| GrpcError::Decode(e.to_string()))?;
75                Ok(Bytes::from(bytes))
76            })
77        });
78        self.handlers
79            .write()
80            .unwrap()
81            .insert(method, Handler::Ctx(handler));
82        self
83    }
84
85    pub async fn invoke_raw(&self, method: &str, body: Bytes) -> Result<Bytes, GrpcError> {
86        let handler = self
87            .handlers
88            .read()
89            .unwrap()
90            .get(method)
91            .cloned()
92            .ok_or_else(|| GrpcError::NotFound(method.to_string()))?;
93        match handler {
94            Handler::Body(h) => h(body).await,
95            Handler::Ctx(h) => {
96                let req = Request::new(http::Method::POST, format!("/{method}"));
97                h(req, body).await
98            }
99        }
100    }
101
102    pub async fn invoke_with_request(
103        &self,
104        method: &str,
105        http_req: Request,
106        body: Bytes,
107    ) -> Result<Bytes, GrpcError> {
108        let handler = self
109            .handlers
110            .read()
111            .unwrap()
112            .get(method)
113            .cloned()
114            .ok_or_else(|| GrpcError::NotFound(method.to_string()))?;
115        match handler {
116            Handler::Body(h) => h(body).await,
117            Handler::Ctx(h) => h(http_req, body).await,
118        }
119    }
120
121    pub async fn invoke<Req, Res>(&self, method: &str, req: &Req) -> Result<Res, GrpcError>
122    where
123        Req: Serialize,
124        Res: DeserializeOwned,
125    {
126        let body =
127            Bytes::from(serde_json::to_vec(req).map_err(|e| GrpcError::Decode(e.to_string()))?);
128        let out = self.invoke_raw(method, body).await?;
129        serde_json::from_slice(&out).map_err(|e| GrpcError::Decode(e.to_string()))
130    }
131
132    pub fn methods(&self) -> Vec<String> {
133        self.handlers.read().unwrap().keys().cloned().collect()
134    }
135}
sova_grpc/router.rs

sova_grpc/
router.rs
