wrpc_transport/frame/conn/
server.rs

1use core::fmt::{Debug, Display};
2use core::marker::PhantomData;
3
4use std::collections::{HashMap, hash_map};
5use std::sync::Arc;
6
7use anyhow::bail;
8use futures::{Stream, StreamExt as _};
9use tokio::io::{AsyncRead, AsyncWrite};
10use tokio::sync::{Mutex, mpsc};
11use tokio_stream::wrappers::ReceiverStream;
12use tracing::{instrument, trace};
13
14use crate::Serve;
15use crate::frame::{ConnHandler, Header, HeaderReadError, Incoming, Outgoing};
16
17/// wRPC server for framed transports
18#[derive(Debug)]
19pub struct Server<C, I, O, H = ()> {
20    handlers: Mutex<HashMap<String, HashMap<String, mpsc::Sender<(C, I, O)>>>>,
21    conn_handler: PhantomData<H>,
22}
23
24impl<C, I, O, H> Server<C, I, O, H> {
25    /// Constructs a new [Server]
26    #[must_use]
27    pub fn new() -> Self {
28        Self {
29            handlers: Mutex::default(),
30            conn_handler: PhantomData,
31        }
32    }
33}
34
35impl<C, I, O> Default for Server<C, I, O> {
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41/// Error returned by [`Server::accept`]
42pub enum AcceptError<C, I, O> {
43    /// Header read error
44    HeaderRead(HeaderReadError),
45    /// Function was not handled
46    UnhandledFunction {
47        /// Instance
48        instance: String,
49        /// Function name
50        name: String,
51    },
52    /// Message sending failed
53    Send(mpsc::error::SendError<(C, I, O)>),
54}
55
56impl<C, I, O> From<HeaderReadError> for AcceptError<C, I, O> {
57    fn from(err: HeaderReadError) -> Self {
58        Self::HeaderRead(err)
59    }
60}
61
62impl<C, I, O> From<mpsc::error::SendError<(C, I, O)>> for AcceptError<C, I, O> {
63    fn from(err: mpsc::error::SendError<(C, I, O)>) -> Self {
64        Self::Send(err)
65    }
66}
67
68impl<C, I, O> Debug for AcceptError<C, I, O> {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        match self {
71            Self::HeaderRead(err) => Debug::fmt(err, f),
72            Self::UnhandledFunction { instance, name } => {
73                write!(f, "`{instance}#{name}` does not have a handler registered")
74            }
75            Self::Send(err) => Debug::fmt(err, f),
76        }
77    }
78}
79
80impl<C, I, O> Display for AcceptError<C, I, O> {
81    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82        match self {
83            Self::HeaderRead(err) => Display::fmt(err, f),
84            Self::UnhandledFunction { instance, name } => {
85                write!(f, "`{instance}#{name}` does not have a handler registered")
86            }
87            Self::Send(err) => Display::fmt(err, f),
88        }
89    }
90}
91
92impl<C, I, O> core::error::Error for AcceptError<C, I, O> {}
93
94impl<C, I, O, H> Server<C, I, O, H>
95where
96    I: AsyncRead + Unpin,
97    H: ConnHandler<I, O>,
98{
99    /// Accept an already-established connection.
100    ///
101    /// # Errors
102    ///
103    /// Returns an error if handling the invocation fails
104    #[instrument(level = "trace", skip_all, ret(level = "trace"))]
105    pub async fn accept(&self, cx: C, tx: O, mut rx: I) -> Result<(), AcceptError<C, I, O>> {
106        let Header { instance, name } = Header::read(&mut rx).await?;
107        let h = self.handlers.lock().await;
108        let h = h
109            .get(&instance)
110            .and_then(|h| h.get(&name))
111            .ok_or_else(|| AcceptError::UnhandledFunction { instance, name })?;
112        h.send((cx, rx, tx)).await?;
113        Ok(())
114    }
115}
116
117#[instrument(level = "trace", skip(srv, paths))]
118async fn serve<C, I, O, H>(
119    srv: &Server<C, I, O, H>,
120    instance: &str,
121    func: &str,
122    paths: Arc<[Box<[Option<usize>]>]>,
123) -> anyhow::Result<
124    impl Stream<Item = anyhow::Result<(C, Outgoing, Incoming)>> + 'static + use<C, I, O, H>,
125>
126where
127    C: Send + Sync + 'static,
128    I: AsyncRead + Send + Sync + Unpin + 'static,
129    O: AsyncWrite + Send + Sync + Unpin + 'static,
130    H: ConnHandler<I, O>,
131{
132    let (tx, rx) = mpsc::channel(1024);
133    let mut handlers = srv.handlers.lock().await;
134    match handlers
135        .entry(instance.to_string())
136        .or_default()
137        .entry(func.to_string())
138    {
139        hash_map::Entry::Occupied(_) => {
140            bail!("handler for `{instance}#{func}` already exists")
141        }
142        hash_map::Entry::Vacant(entry) => {
143            entry.insert(tx);
144        }
145    }
146    Ok(ReceiverStream::new(rx).map(move |(cx, rx, tx)| {
147        trace!("received invocation");
148        let rx = Incoming::new(rx, paths.as_ref(), |rx, res| H::on_ingress(rx, res));
149        let tx = Outgoing::new(tx, |tx, res| H::on_egress(tx, res));
150        Ok((cx, tx, rx))
151    }))
152}
153
154impl<C, I, O, H> Serve for Server<C, I, O, H>
155where
156    C: Send + Sync + 'static,
157    I: AsyncRead + Send + Sync + Unpin + 'static,
158    O: AsyncWrite + Send + Sync + Unpin + 'static,
159    H: ConnHandler<I, O> + Send + Sync,
160{
161    type Context = C;
162
163    async fn serve(
164        &self,
165        instance: &str,
166        func: &str,
167        paths: Arc<[Box<[Option<usize>]>]>,
168    ) -> anyhow::Result<
169        impl Stream<Item = anyhow::Result<(Self::Context, Outgoing, Incoming)>>
170        + 'static
171        + use<C, I, O, H>,
172    > {
173        serve(self, instance, func, paths).await
174    }
175}
176
177impl<'a, C, I, O, H> Serve for &'a Server<C, I, O, H>
178where
179    C: Send + Sync + 'static,
180    I: AsyncRead + Send + Sync + Unpin + 'static,
181    O: AsyncWrite + Send + Sync + Unpin + 'static,
182    H: ConnHandler<I, O> + Send + Sync,
183{
184    type Context = C;
185
186    async fn serve(
187        &self,
188        instance: &str,
189        func: &str,
190        paths: Arc<[Box<[Option<usize>]>]>,
191    ) -> anyhow::Result<
192        impl Stream<Item = anyhow::Result<(Self::Context, Outgoing, Incoming)>>
193        + 'static
194        + use<'a, C, I, O, H>,
195    > {
196        serve(self, instance, func, paths).await
197    }
198}
wrpc_transport/frame/conn/server.rs

wrpc_transport/frame/conn/
server.rs
