xrpc/
server.rs

1use async_trait::async_trait;
2use bytes::Bytes;
3use futures::Stream;
4use parking_lot::RwLock;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::future::Future;
8use std::sync::Arc;
9use tokio::sync::mpsc;
10
11use crate::channel::message::MessageChannel;
12use crate::codec::{BincodeCodec, Codec};
13use crate::error::{Result, RpcError, TransportError};
14use crate::message::Message;
15use crate::message::metadata::MessageMetadata;
16use crate::message::types::{MessageId, MessageType};
17use crate::streaming::{StreamId, next_stream_id};
18
19#[async_trait]
20pub trait Handler<C: Codec>: Send + Sync {
21    async fn handle(&self, request: Message<C>, codec: &C) -> Result<Message<C>>;
22    fn method_name(&self) -> &str;
23}
24
25#[async_trait]
26pub trait StreamHandler<C: Codec>: Send + Sync {
27    async fn handle(
28        &self,
29        request: Message<C>,
30        sender: ServerStreamSender<C>,
31        codec: &C,
32    ) -> Result<()>;
33    fn method_name(&self) -> &str;
34}
35
36type HandlerMap<C> = Arc<RwLock<HashMap<String, Arc<dyn Handler<C>>>>>;
37type StreamHandlerMap<C> = Arc<RwLock<HashMap<String, Arc<dyn StreamHandler<C>>>>>;
38
39fn is_connected_idle_timeout<T, C>(transport: &T, error: &TransportError) -> bool
40where
41    T: MessageChannel<C> + ?Sized,
42    C: Codec,
43{
44    matches!(error, TransportError::Timeout { .. }) && transport.is_connected()
45}
46
47pub struct FnHandler<F, C> {
48    method: String,
49    func: Arc<F>,
50    _codec: std::marker::PhantomData<C>,
51}
52
53impl<F, Fut, C> FnHandler<F, C>
54where
55    F: Fn(Message<C>) -> Fut + Send + Sync + 'static,
56    Fut: Future<Output = Result<Message<C>>> + Send + 'static,
57    C: Codec,
58{
59    pub fn new(method: impl Into<String>, func: F) -> Self {
60        Self {
61            method: method.into(),
62            func: Arc::new(func),
63            _codec: std::marker::PhantomData,
64        }
65    }
66}
67
68#[async_trait]
69impl<F, Fut, C: Codec + Default> Handler<C> for FnHandler<F, C>
70where
71    F: Fn(Message<C>) -> Fut + Send + Sync + 'static,
72    Fut: Future<Output = Result<Message<C>>> + Send + 'static,
73{
74    async fn handle(&self, request: Message<C>, _codec: &C) -> Result<Message<C>> {
75        (self.func)(request).await
76    }
77
78    fn method_name(&self) -> &str {
79        &self.method
80    }
81}
82
83pub struct TypedHandler<Req, Resp, F, C> {
84    method: String,
85    func: Arc<F>,
86    _phantom: std::marker::PhantomData<(Req, Resp, C)>,
87}
88
89impl<Req, Resp, F, Fut, C> TypedHandler<Req, Resp, F, C>
90where
91    Req: for<'de> Deserialize<'de> + Send + 'static,
92    Resp: Serialize + Send + 'static,
93    F: Fn(Req) -> Fut + Send + Sync + 'static,
94    Fut: Future<Output = Result<Resp>> + Send + 'static,
95    C: Codec,
96{
97    pub fn new(method: impl Into<String>, func: F) -> Self {
98        Self {
99            method: method.into(),
100            func: Arc::new(func),
101            _phantom: std::marker::PhantomData,
102        }
103    }
104}
105
106#[async_trait]
107impl<Req, Resp, F, Fut, C> Handler<C> for TypedHandler<Req, Resp, F, C>
108where
109    Req: for<'de> Deserialize<'de> + Send + Sync + 'static,
110    Resp: Serialize + Send + Sync + 'static,
111    F: Fn(Req) -> Fut + Send + Sync + 'static,
112    Fut: Future<Output = Result<Resp>> + Send + 'static,
113    C: Codec + Default,
114{
115    async fn handle(&self, request: Message<C>, codec: &C) -> Result<Message<C>> {
116        let req: Req = codec.decode(&request.payload)?;
117        let resp = (self.func)(req).await?;
118        let payload = codec.encode(&resp)?;
119        Ok(Message::new(
120            request.id,
121            MessageType::Reply,
122            "",
123            Bytes::from(payload),
124            MessageMetadata::new(),
125        ))
126    }
127
128    fn method_name(&self) -> &str {
129        &self.method
130    }
131}
132
133pub struct ServerStreamSender<C: Codec> {
134    stream_id: StreamId,
135    tx: mpsc::UnboundedSender<Bytes>,
136    sequence: std::sync::atomic::AtomicU64,
137    codec: C,
138}
139
140impl<C: Codec> ServerStreamSender<C> {
141    fn new(stream_id: StreamId, tx: mpsc::UnboundedSender<Bytes>, codec: C) -> Self {
142        Self {
143            stream_id,
144            tx,
145            sequence: std::sync::atomic::AtomicU64::new(0),
146            codec,
147        }
148    }
149
150    pub fn stream_id(&self) -> StreamId {
151        self.stream_id
152    }
153
154    pub fn send<T: Serialize>(&self, data: T) -> Result<()> {
155        let seq = self
156            .sequence
157            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
158        let payload = self.codec.encode(&data)?;
159        let chunk: Message = Message::new(
160            MessageId::new(),
161            MessageType::StreamChunk,
162            "",
163            Bytes::from(payload),
164            MessageMetadata::new().with_stream(self.stream_id, seq),
165        );
166        let encoded = chunk.encode().map_err(RpcError::Transport)?;
167
168        self.tx
169            .send(encoded.freeze())
170            .map_err(|_| RpcError::StreamError("Stream closed".to_string()))
171    }
172
173    pub fn end(&self) -> Result<()> {
174        let end_msg: Message = Message::stream_end(self.stream_id);
175        let encoded = end_msg.encode().map_err(RpcError::Transport)?;
176
177        self.tx
178            .send(encoded.freeze())
179            .map_err(|_| RpcError::StreamError("Stream closed".to_string()))
180    }
181}
182
183pub struct TypedStreamHandler<Req, Item, F, C> {
184    method: String,
185    func: Arc<F>,
186    _phantom: std::marker::PhantomData<(Req, Item, C)>,
187}
188
189impl<Req, Item, F, S, C> TypedStreamHandler<Req, Item, F, C>
190where
191    Req: for<'de> Deserialize<'de> + Send + 'static,
192    Item: Serialize + Send + 'static,
193    S: Stream<Item = Result<Item>> + Send + 'static,
194    F: Fn(Req) -> S + Send + Sync + 'static,
195    C: Codec,
196{
197    pub fn new(method: impl Into<String>, func: F) -> Self {
198        Self {
199            method: method.into(),
200            func: Arc::new(func),
201            _phantom: std::marker::PhantomData,
202        }
203    }
204}
205
206#[async_trait]
207impl<Req, Item, F, S, C> StreamHandler<C> for TypedStreamHandler<Req, Item, F, C>
208where
209    Req: for<'de> Deserialize<'de> + Send + Sync + 'static,
210    Item: Serialize + Send + Sync + 'static,
211    S: Stream<Item = Result<Item>> + Send + 'static,
212    F: Fn(Req) -> S + Send + Sync + 'static,
213    C: Codec + Default,
214{
215    async fn handle(
216        &self,
217        request: Message<C>,
218        sender: ServerStreamSender<C>,
219        codec: &C,
220    ) -> Result<()> {
221        use futures::StreamExt;
222
223        let req: Req = codec.decode(&request.payload)?;
224        let mut stream = Box::pin((self.func)(req));
225
226        while let Some(result) = stream.next().await {
227            match result {
228                Ok(item) => sender.send(item)?,
229                Err(e) => return Err(e),
230            }
231        }
232
233        sender.end()?;
234        Ok(())
235    }
236
237    fn method_name(&self) -> &str {
238        &self.method
239    }
240}
241
242pub struct FnStreamHandler<F, C> {
243    method: String,
244    func: Arc<F>,
245    _codec: std::marker::PhantomData<C>,
246}
247
248impl<F, Fut, C> FnStreamHandler<F, C>
249where
250    F: Fn(Message<C>, ServerStreamSender<C>) -> Fut + Send + Sync + 'static,
251    Fut: Future<Output = Result<()>> + Send + 'static,
252    C: Codec,
253{
254    pub fn new(method: impl Into<String>, func: F) -> Self {
255        Self {
256            method: method.into(),
257            func: Arc::new(func),
258            _codec: std::marker::PhantomData,
259        }
260    }
261}
262
263#[async_trait]
264impl<F, Fut, C> StreamHandler<C> for FnStreamHandler<F, C>
265where
266    F: Fn(Message<C>, ServerStreamSender<C>) -> Fut + Send + Sync + 'static,
267    Fut: Future<Output = Result<()>> + Send + 'static,
268    C: Codec + Default,
269{
270    async fn handle(
271        &self,
272        request: Message<C>,
273        sender: ServerStreamSender<C>,
274        _codec: &C,
275    ) -> Result<()> {
276        (self.func)(request, sender).await
277    }
278
279    fn method_name(&self) -> &str {
280        &self.method
281    }
282}
283
284pub struct RpcServer<C: Codec = BincodeCodec> {
285    handlers: HandlerMap<C>,
286    stream_handlers: StreamHandlerMap<C>,
287    codec: C,
288}
289
290impl RpcServer<BincodeCodec> {
291    pub fn new() -> Self {
292        Self {
293            handlers: Arc::new(RwLock::new(HashMap::new())),
294            stream_handlers: Arc::new(RwLock::new(HashMap::new())),
295            codec: BincodeCodec,
296        }
297    }
298}
299
300impl<C: Codec + Clone + Default + 'static> RpcServer<C> {
301    pub fn with_codec(codec: C) -> Self {
302        Self {
303            handlers: Arc::new(RwLock::new(HashMap::new())),
304            stream_handlers: Arc::new(RwLock::new(HashMap::new())),
305            codec,
306        }
307    }
308
309    pub fn register(&self, handler: Arc<dyn Handler<C>>) {
310        let method = handler.method_name().to_string();
311        self.handlers.write().insert(method, handler);
312    }
313
314    pub fn register_fn<F, Fut>(&self, method: impl Into<String>, func: F)
315    where
316        F: Fn(Message<C>) -> Fut + Send + Sync + 'static,
317        Fut: Future<Output = Result<Message<C>>> + Send + 'static,
318    {
319        let handler: Arc<FnHandler<F, C>> = Arc::new(FnHandler::new(method, func));
320        self.register(handler);
321    }
322
323    pub fn register_typed<Req, Resp, F, Fut>(&self, method: impl Into<String>, func: F)
324    where
325        Req: for<'de> Deserialize<'de> + Send + Sync + 'static,
326        Resp: Serialize + Send + Sync + 'static,
327        F: Fn(Req) -> Fut + Send + Sync + 'static,
328        Fut: Future<Output = Result<Resp>> + Send + 'static,
329    {
330        let handler: Arc<TypedHandler<Req, Resp, F, C>> = Arc::new(TypedHandler::new(method, func));
331        self.register(handler);
332    }
333
334    pub fn register_stream<Req, Item, F, S>(&self, method: impl Into<String>, func: F)
335    where
336        Req: for<'de> Deserialize<'de> + Send + Sync + 'static,
337        Item: Serialize + Send + Sync + 'static,
338        S: Stream<Item = Result<Item>> + Send + 'static,
339        F: Fn(Req) -> S + Send + Sync + 'static,
340    {
341        let method = method.into();
342        let handler: Arc<TypedStreamHandler<Req, Item, F, C>> =
343            Arc::new(TypedStreamHandler::new(method.clone(), func));
344        self.stream_handlers.write().insert(method, handler);
345    }
346
347    pub fn register_stream_fn<F, Fut>(&self, method: impl Into<String>, func: F)
348    where
349        F: Fn(Message<C>, ServerStreamSender<C>) -> Fut + Send + Sync + 'static,
350        Fut: Future<Output = Result<()>> + Send + 'static,
351    {
352        let method = method.into();
353        let handler: Arc<FnStreamHandler<F, C>> =
354            Arc::new(FnStreamHandler::new(method.clone(), func));
355        self.stream_handlers.write().insert(method, handler);
356    }
357
358    /// Dispatch a single message, sending responses via `send_tx`.
359    async fn dispatch_message(
360        handlers: &HandlerMap<C>,
361        stream_handlers: &StreamHandlerMap<C>,
362        codec: &C,
363        message: Message<C>,
364        send_tx: &mpsc::UnboundedSender<Message<C>>,
365    ) {
366        match message.msg_type {
367            MessageType::Call => {
368                if message.metadata.stream_id.is_some() {
369                    Self::dispatch_stream_call(stream_handlers, codec, message, send_tx).await;
370                    return;
371                }
372
373                let handler = handlers.read().get(&message.method).cloned();
374                let response = match handler {
375                    Some(h) => match h.handle(message.clone(), codec).await {
376                        Ok(response) => response,
377                        Err(e) => {
378                            Message::error_with_codec(message.id, e.to_string(), codec.clone())
379                        }
380                    },
381                    None => Message::error_with_codec(
382                        message.id,
383                        format!("Method not found: {}", message.method),
384                        codec.clone(),
385                    ),
386                };
387                let _ = send_tx.send(response);
388            }
389            MessageType::Notification => {
390                let handler = handlers.read().get(&message.method).cloned();
391                if let Some(h) = handler {
392                    let _ = h.handle(message, codec).await;
393                }
394            }
395            _ => {}
396        }
397    }
398
399    /// Dispatch a streaming call, forwarding all chunks via `send_tx`.
400    async fn dispatch_stream_call(
401        stream_handlers: &StreamHandlerMap<C>,
402        codec: &C,
403        message: Message<C>,
404        send_tx: &mpsc::UnboundedSender<Message<C>>,
405    ) {
406        let stream_id = message.metadata.stream_id.unwrap_or_else(next_stream_id);
407        let handler = stream_handlers.read().get(&message.method).cloned();
408
409        let Some(h) = handler else {
410            let error = Message::stream_error_with_codec(
411                message.id,
412                stream_id,
413                format!("Stream method not found: {}", message.method),
414                codec.clone(),
415            );
416            let _ = send_tx.send(error);
417            return;
418        };
419
420        let (tx, mut rx) = mpsc::unbounded_channel::<Bytes>();
421        let sender = ServerStreamSender::new(stream_id, tx, codec.clone());
422
423        let send_tx_clone = send_tx.clone();
424        let forward_task = async move {
425            while let Some(data) = rx.recv().await {
426                if let Ok(msg) = Message::<C>::decode(&data[..]) {
427                    let _ = send_tx_clone.send(msg);
428                }
429            }
430        };
431
432        let codec = codec.clone();
433        let handler_task = async {
434            if let Err(e) = h.handle(message.clone(), sender, &codec).await {
435                let error = Message::stream_error_with_codec(
436                    message.id,
437                    stream_id,
438                    e.to_string(),
439                    codec.clone(),
440                );
441                let _ = send_tx.send(error);
442            }
443        };
444
445        tokio::join!(handler_task, forward_task);
446    }
447
448    /// Serve requests concurrently, dispatching each message in its own tokio task.
449    pub async fn serve<T: MessageChannel<C> + 'static>(&self, transport: Arc<T>) -> Result<()> {
450        let (send_tx, mut send_rx) = mpsc::unbounded_channel::<Message<C>>();
451
452        let transport_writer = transport.clone();
453        let writer_handle = tokio::spawn(async move {
454            while let Some(msg) = send_rx.recv().await {
455                if transport_writer.send(&msg).await.is_err() {
456                    break;
457                }
458            }
459        });
460
461        let result = loop {
462            let message = match transport.recv().await {
463                Ok(msg) => msg,
464                Err(error) if is_connected_idle_timeout(transport.as_ref(), &error) => continue,
465                Err(error) => break Err(RpcError::Transport(error)),
466            };
467
468            let handlers = self.handlers.clone();
469            let stream_handlers = self.stream_handlers.clone();
470            let codec = self.codec.clone();
471            let send_tx = send_tx.clone();
472
473            tokio::spawn(async move {
474                Self::dispatch_message(&handlers, &stream_handlers, &codec, message, &send_tx)
475                    .await;
476            });
477        };
478
479        drop(send_tx);
480        let _ = writer_handle.await;
481
482        result
483    }
484
485    /// Serve requests sequentially, processing one message at a time.
486    ///
487    /// This preserves strict one-at-a-time message processing for cases where ordering guarantees or single-threaded handler execution is required.
488    pub async fn serve_sequential<T: MessageChannel<C> + 'static>(
489        &self,
490        transport: Arc<T>,
491    ) -> Result<()> {
492        let (send_tx, mut send_rx) = mpsc::unbounded_channel::<Message<C>>();
493
494        let transport_writer = transport.clone();
495        let writer_handle = tokio::spawn(async move {
496            while let Some(msg) = send_rx.recv().await {
497                if transport_writer.send(&msg).await.is_err() {
498                    break;
499                }
500            }
501        });
502
503        let result = loop {
504            let message = match transport.recv().await {
505                Ok(msg) => msg,
506                Err(error) if is_connected_idle_timeout(transport.as_ref(), &error) => continue,
507                Err(error) => break Err(RpcError::Transport(error)),
508            };
509
510            Self::dispatch_message(
511                &self.handlers,
512                &self.stream_handlers,
513                &self.codec,
514                message,
515                &send_tx,
516            )
517            .await;
518        };
519
520        drop(send_tx);
521        let _ = writer_handle.await;
522
523        result
524    }
525
526    pub fn spawn_handler<T: MessageChannel<C> + 'static>(&self, transport: T) -> ServerHandle {
527        let handlers = self.handlers.clone();
528        let stream_handlers = self.stream_handlers.clone();
529        let codec = self.codec.clone();
530        let transport = Arc::new(transport);
531
532        let handle = tokio::spawn(async move {
533            let server = RpcServer {
534                handlers,
535                stream_handlers,
536                codec,
537            };
538            let _ = server.serve(transport).await;
539        });
540
541        ServerHandle { handle }
542    }
543
544    pub fn handler_count(&self) -> usize {
545        self.handlers.read().len() + self.stream_handlers.read().len()
546    }
547}
548
549impl Default for RpcServer<BincodeCodec> {
550    fn default() -> Self {
551        Self::new()
552    }
553}
554
555pub struct ServerHandle {
556    handle: tokio::task::JoinHandle<()>,
557}
558
559impl ServerHandle {
560    pub async fn shutdown(self) {
561        self.handle.abort();
562        let _ = self.handle.await;
563    }
564
565    pub fn is_finished(&self) -> bool {
566        self.handle.is_finished()
567    }
568}
569
570#[cfg(test)]
571mod tests {
572    use super::*;
573    use crate::channel::message::{MessageChannel, MessageChannelAdapter};
574    use crate::error::TransportResult;
575    use crate::streaming::StreamReceiver;
576    use crate::transport::channel::{ChannelConfig, ChannelFrameTransport};
577    use parking_lot::Mutex;
578    use std::collections::VecDeque;
579    use std::sync::atomic::{AtomicBool, Ordering};
580    use std::time::Duration;
581    use tokio::sync::Notify;
582
583    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
584    struct AddRequest {
585        a: i32,
586        b: i32,
587    }
588
589    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
590    struct AddResponse {
591        result: i32,
592    }
593
594    #[derive(Debug)]
595    enum ScriptedReceive {
596        IdleTimeout,
597        Message(Message),
598        DisconnectAfterReply,
599        DisconnectedTimeout,
600        Error(TransportError),
601    }
602
603    #[derive(Debug)]
604    struct ScriptedServerChannel {
605        connected: AtomicBool,
606        received: Mutex<VecDeque<ScriptedReceive>>,
607        sent: Mutex<Vec<Message>>,
608        sent_notify: Notify,
609    }
610
611    impl ScriptedServerChannel {
612        fn new(received: impl IntoIterator<Item = ScriptedReceive>) -> Self {
613            Self {
614                connected: AtomicBool::new(true),
615                received: Mutex::new(received.into_iter().collect()),
616                sent: Mutex::new(Vec::new()),
617                sent_notify: Notify::new(),
618            }
619        }
620
621        async fn wait_for_reply(&self) {
622            loop {
623                let notified = self.sent_notify.notified();
624                if !self.sent.lock().is_empty() {
625                    return;
626                }
627                notified.await;
628            }
629        }
630    }
631
632    #[async_trait::async_trait]
633    impl MessageChannel for ScriptedServerChannel {
634        async fn send(&self, message: &Message) -> TransportResult<()> {
635            self.sent.lock().push(message.clone());
636            self.sent_notify.notify_waiters();
637            Ok(())
638        }
639
640        async fn recv(&self) -> TransportResult<Message> {
641            let step = self
642                .received
643                .lock()
644                .pop_front()
645                .expect("scripted receive sequence exhausted");
646            match step {
647                ScriptedReceive::IdleTimeout => {
648                    tokio::time::sleep(Duration::from_millis(5)).await;
649                    Err(TransportError::Timeout {
650                        duration_ms: 5,
651                        operation: "receive".to_string(),
652                    })
653                }
654                ScriptedReceive::Message(message) => Ok(message),
655                ScriptedReceive::DisconnectAfterReply => {
656                    self.wait_for_reply().await;
657                    self.connected.store(false, Ordering::Release);
658                    Err(TransportError::ConnectionClosed)
659                }
660                ScriptedReceive::DisconnectedTimeout => {
661                    self.connected.store(false, Ordering::Release);
662                    Err(TransportError::Timeout {
663                        duration_ms: 5,
664                        operation: "receive".to_string(),
665                    })
666                }
667                ScriptedReceive::Error(error) => Err(error),
668            }
669        }
670
671        fn is_connected(&self) -> bool {
672            self.connected.load(Ordering::Acquire)
673        }
674
675        async fn close(&self) -> TransportResult<()> {
676            self.connected.store(false, Ordering::Release);
677            self.sent_notify.notify_waiters();
678            Ok(())
679        }
680    }
681
682    async fn run_connected_timeout_test(sequential: bool) {
683        let request = Message::call("add", AddRequest { a: 20, b: 22 }).unwrap();
684        let channel = Arc::new(ScriptedServerChannel::new([
685            ScriptedReceive::IdleTimeout,
686            ScriptedReceive::Message(request),
687            ScriptedReceive::DisconnectAfterReply,
688        ]));
689        let server = RpcServer::new();
690        server.register_typed("add", |request: AddRequest| async move {
691            Ok(AddResponse {
692                result: request.a + request.b,
693            })
694        });
695
696        let result = tokio::time::timeout(Duration::from_secs(1), async {
697            if sequential {
698                server.serve_sequential(Arc::clone(&channel)).await
699            } else {
700                server.serve(Arc::clone(&channel)).await
701            }
702        })
703        .await
704        .expect("server did not finish after the scripted disconnect");
705
706        assert!(matches!(
707            result,
708            Err(RpcError::Transport(TransportError::ConnectionClosed))
709        ));
710        let sent = channel.sent.lock();
711        assert_eq!(sent.len(), 1);
712        assert_eq!(sent[0].msg_type, MessageType::Reply);
713        let response: AddResponse = sent[0].deserialize_payload().unwrap();
714        assert_eq!(response.result, 42);
715    }
716
717    async fn run_terminal_receive_test(sequential: bool, step: ScriptedReceive) -> RpcError {
718        let channel = Arc::new(ScriptedServerChannel::new([step]));
719        let server = RpcServer::new();
720        tokio::time::timeout(Duration::from_secs(1), async {
721            if sequential {
722                server.serve_sequential(channel).await
723            } else {
724                server.serve(channel).await
725            }
726        })
727        .await
728        .expect("terminal receive error did not stop the server")
729        .unwrap_err()
730    }
731
732    #[derive(Debug, Clone)]
733    struct StatefulCodec {
734        tag: &'static str,
735    }
736
737    impl StatefulCodec {
738        fn configured() -> Self {
739            Self { tag: "configured" }
740        }
741    }
742
743    impl Default for StatefulCodec {
744        fn default() -> Self {
745            Self { tag: "default" }
746        }
747    }
748
749    impl Codec for StatefulCodec {
750        fn encode<T: Serialize>(&self, data: &T) -> Result<Vec<u8>> {
751            let mut encoded = format!("{}:", self.tag).into_bytes();
752            encoded.extend(serde_json::to_vec(data)?);
753            Ok(encoded)
754        }
755
756        fn decode<T: for<'de> Deserialize<'de>>(&self, data: &[u8]) -> Result<T> {
757            let prefix = format!("{}:", self.tag);
758            let payload = data.strip_prefix(prefix.as_bytes()).ok_or_else(|| {
759                RpcError::Serialization(format!("missing {} codec prefix", self.tag))
760            })?;
761            Ok(serde_json::from_slice(payload)?)
762        }
763    }
764
765    #[tokio::test]
766    async fn test_concurrent_server_survives_connected_idle_timeout() {
767        run_connected_timeout_test(false).await;
768    }
769
770    #[tokio::test]
771    async fn test_sequential_server_survives_connected_idle_timeout() {
772        run_connected_timeout_test(true).await;
773    }
774
775    #[tokio::test]
776    async fn test_disconnected_timeout_remains_terminal_in_both_server_modes() {
777        for sequential in [false, true] {
778            let error =
779                run_terminal_receive_test(sequential, ScriptedReceive::DisconnectedTimeout).await;
780            assert!(matches!(
781                error,
782                RpcError::Transport(TransportError::Timeout { .. })
783            ));
784        }
785    }
786
787    #[tokio::test]
788    async fn test_non_timeout_receive_error_remains_terminal_in_both_server_modes() {
789        for sequential in [false, true] {
790            let error = run_terminal_receive_test(
791                sequential,
792                ScriptedReceive::Error(TransportError::Protocol("invalid frame".to_string())),
793            )
794            .await;
795            assert!(matches!(
796                error,
797                RpcError::Transport(TransportError::Protocol(message))
798                    if message == "invalid frame"
799            ));
800        }
801    }
802
803    #[tokio::test]
804    async fn test_long_unary_handler_survives_connected_receive_timeouts() {
805        let config = ChannelConfig::default().with_read_timeout(Duration::from_millis(20));
806        let (client_transport, server_transport) =
807            ChannelFrameTransport::create_pair("long-unary-timeout", config).unwrap();
808        let client_channel = MessageChannelAdapter::new(client_transport);
809        let server_channel = MessageChannelAdapter::new(server_transport);
810        let server = RpcServer::new();
811        server.register_typed("slow_add", |request: AddRequest| async move {
812            tokio::time::sleep(Duration::from_millis(100)).await;
813            Ok(AddResponse {
814                result: request.a + request.b,
815            })
816        });
817        let server_handle = server.spawn_handler(server_channel);
818        let client = crate::client::RpcClient::with_timeout(client_channel, Duration::from_secs(1));
819        let client_handle = client.start();
820
821        let first: AddResponse = client
822            .call("slow_add", &AddRequest { a: 20, b: 22 })
823            .await
824            .unwrap();
825        assert_eq!(first.result, 42);
826        let second: AddResponse = client
827            .call("slow_add", &AddRequest { a: 1, b: 2 })
828            .await
829            .unwrap();
830        assert_eq!(second.result, 3);
831
832        client.close().await.unwrap();
833        client_handle.join().await.unwrap();
834        server_handle.shutdown().await;
835    }
836
837    #[tokio::test]
838    async fn test_delayed_server_stream_survives_connected_receive_timeouts() {
839        let config = ChannelConfig::default().with_read_timeout(Duration::from_millis(20));
840        let (client_transport, server_transport) =
841            ChannelFrameTransport::create_pair("long-stream-timeout", config).unwrap();
842        let client_channel = MessageChannelAdapter::new(client_transport);
843        let server_channel = MessageChannelAdapter::new(server_transport);
844        let server = RpcServer::new();
845        server.register_stream_fn(
846            "delayed_stream",
847            |_request: Message, sender: ServerStreamSender<BincodeCodec>| async move {
848                tokio::time::sleep(Duration::from_millis(100)).await;
849                sender.send("ready".to_string())?;
850                sender.end()?;
851                Ok(())
852            },
853        );
854        let server_handle = server.spawn_handler(server_channel);
855        let client = crate::client::RpcClient::new(client_channel);
856        let client_handle = client.start();
857
858        let mut stream: StreamReceiver<String> = client
859            .call_server_stream("delayed_stream", &())
860            .await
861            .unwrap();
862        assert_eq!(stream.recv().await.unwrap().unwrap(), "ready");
863        assert!(stream.recv().await.is_none());
864
865        client.close().await.unwrap();
866        client_handle.join().await.unwrap();
867        server_handle.shutdown().await;
868    }
869
870    #[tokio::test]
871    async fn test_explicit_peer_close_interrupts_long_server_receive() {
872        let config = ChannelConfig::default().with_read_timeout(Duration::from_secs(300));
873        let (client_transport, server_transport) =
874            ChannelFrameTransport::create_pair("server-explicit-close", config).unwrap();
875        let client_channel = MessageChannelAdapter::new(client_transport);
876        let server_channel = Arc::new(MessageChannelAdapter::new(server_transport));
877        let server = RpcServer::new();
878        let serve = tokio::spawn(async move { server.serve(server_channel).await });
879
880        client_channel.close().await.unwrap();
881        let result = tokio::time::timeout(Duration::from_secs(1), serve)
882            .await
883            .expect("explicit peer close did not wake the server receive")
884            .unwrap();
885        assert!(matches!(
886            result,
887            Err(RpcError::Transport(TransportError::ConnectionClosed))
888        ));
889    }
890
891    #[tokio::test]
892    async fn test_server_typed_handler() {
893        let config = ChannelConfig::default();
894        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
895
896        let client_channel = MessageChannelAdapter::new(t1);
897        let server_channel = MessageChannelAdapter::new(t2);
898
899        let server = RpcServer::new();
900        server.register_typed("add", |req: AddRequest| async move {
901            Ok(AddResponse {
902                result: req.a + req.b,
903            })
904        });
905
906        let _handle = server.spawn_handler(server_channel);
907
908        let request: Message = Message::call("add", AddRequest { a: 10, b: 32 }).unwrap();
909        client_channel.send(&request).await.unwrap();
910
911        let response = client_channel.recv().await.unwrap();
912        assert_eq!(response.msg_type, MessageType::Reply);
913
914        let resp: AddResponse = response.deserialize_payload().unwrap();
915        assert_eq!(resp.result, 42);
916    }
917
918    #[tokio::test]
919    async fn test_server_errors_use_configured_codec_instance() {
920        let (client_transport, server_transport) =
921            ChannelFrameTransport::create_pair("stateful-error", ChannelConfig::default()).unwrap();
922        let client_channel =
923            MessageChannelAdapter::<_, StatefulCodec>::with_codec(client_transport);
924        let server_channel =
925            MessageChannelAdapter::<_, StatefulCodec>::with_codec(server_transport);
926        let codec = StatefulCodec::configured();
927        let server = RpcServer::with_codec(codec.clone());
928        let server_handle = server.spawn_handler(server_channel);
929        let client = crate::client::RpcClient::with_codec(client_channel, codec);
930        let client_handle = client.start();
931
932        let result: Result<AddResponse> = client.call("missing", &AddRequest { a: 1, b: 2 }).await;
933
934        assert!(matches!(
935            result,
936            Err(RpcError::ServerError(error)) if error.contains("Method not found: missing")
937        ));
938        client.close().await.unwrap();
939        client_handle.join().await.unwrap();
940        server_handle.shutdown().await;
941    }
942
943    #[tokio::test]
944    async fn test_server_stream_handler() {
945        let config = ChannelConfig::default();
946        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
947
948        let client_channel = Arc::new(MessageChannelAdapter::new(t1));
949        let server_channel = MessageChannelAdapter::new(t2);
950
951        let server = RpcServer::new();
952        server.register_stream("range", |count: i32| {
953            futures::stream::iter((1..=count).map(Ok))
954        });
955
956        let _handle = server.spawn_handler(server_channel);
957
958        let stream_id = next_stream_id();
959        let mut request: Message = Message::call("range", 5i32).unwrap();
960        request.metadata = request.metadata.with_stream(stream_id, 0);
961
962        let manager = crate::streaming::StreamManager::new();
963        let mut receiver: StreamReceiver<i32> = manager.create_receiver(stream_id);
964
965        client_channel.send(&request).await.unwrap();
966
967        let client_channel_clone = client_channel.clone();
968        let recv_task = tokio::spawn(async move {
969            while let Ok(msg) = client_channel_clone.recv().await {
970                if msg.msg_type == MessageType::StreamEnd {
971                    manager.handle_message(&msg);
972                    break;
973                }
974                manager.handle_message(&msg);
975            }
976        });
977
978        let mut items = Vec::new();
979        while let Some(result) = receiver.recv().await {
980            items.push(result.unwrap());
981        }
982
983        recv_task.await.unwrap();
984        assert_eq!(items, vec![1, 2, 3, 4, 5]);
985    }
986
987    #[tokio::test]
988    async fn test_server_notification() {
989        use std::sync::atomic::{AtomicBool, Ordering};
990
991        let config = ChannelConfig::default();
992        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
993
994        let client_channel = MessageChannelAdapter::new(t1);
995        let server_channel = MessageChannelAdapter::new(t2);
996
997        let called = Arc::new(AtomicBool::new(false));
998        let called_clone = called.clone();
999
1000        let server = RpcServer::new();
1001        server.register_fn("log", move |_msg: Message| {
1002            let called = called_clone.clone();
1003            async move {
1004                called.store(true, Ordering::Release);
1005                Message::reply(MessageId::new(), ())
1006            }
1007        });
1008
1009        let _handle = server.spawn_handler(server_channel);
1010
1011        let notification: Message = Message::notification("log", "test").unwrap();
1012        client_channel.send(&notification).await.unwrap();
1013
1014        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1015        assert!(called.load(Ordering::Acquire));
1016    }
1017
1018    #[tokio::test]
1019    async fn test_unary_during_stream() {
1020        let config = ChannelConfig::default();
1021        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1022        let client_channel = MessageChannelAdapter::new(t1);
1023        let server_channel = MessageChannelAdapter::new(t2);
1024
1025        let server = RpcServer::new();
1026        server.register_stream_fn(
1027            "slow_stream",
1028            |_msg: Message, sender: ServerStreamSender<BincodeCodec>| async move {
1029                for i in 0..5 {
1030                    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1031                    sender.send(i)?;
1032                }
1033                sender.end()?;
1034                Ok(())
1035            },
1036        );
1037        server.register_typed("add", |req: AddRequest| async move {
1038            Ok(AddResponse {
1039                result: req.a + req.b,
1040            })
1041        });
1042
1043        let _handle = server.spawn_handler(server_channel);
1044
1045        let client = crate::client::RpcClient::new(client_channel);
1046        let _client_handle = client.start();
1047
1048        let mut stream: StreamReceiver<i32> =
1049            client.call_server_stream("slow_stream", &()).await.unwrap();
1050
1051        let response: AddResponse = tokio::time::timeout(
1052            std::time::Duration::from_secs(1),
1053            client.call("add", &AddRequest { a: 10, b: 32 }),
1054        )
1055        .await
1056        .expect("Unary call should not be blocked by active stream")
1057        .unwrap();
1058        assert_eq!(response.result, 42);
1059
1060        let mut items = Vec::new();
1061        while let Some(result) = stream.recv().await {
1062            items.push(result.unwrap());
1063        }
1064        assert_eq!(items, vec![0, 1, 2, 3, 4]);
1065    }
1066
1067    #[tokio::test]
1068    async fn test_multiple_concurrent_streams() {
1069        let config = ChannelConfig::default();
1070        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1071        let client_channel = MessageChannelAdapter::new(t1);
1072        let server_channel = MessageChannelAdapter::new(t2);
1073
1074        let server = RpcServer::new();
1075        server.register_stream_fn(
1076            "countdown",
1077            |msg: Message, sender: ServerStreamSender<BincodeCodec>| async move {
1078                let count: i32 = msg.deserialize_payload().unwrap_or(3);
1079                for i in (0..count).rev() {
1080                    tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1081                    sender.send(i)?;
1082                }
1083                sender.end()?;
1084                Ok(())
1085            },
1086        );
1087
1088        let _handle = server.spawn_handler(server_channel);
1089
1090        let client = crate::client::RpcClient::new(client_channel);
1091        let _client_handle = client.start();
1092
1093        let mut stream_a: StreamReceiver<i32> =
1094            client.call_server_stream("countdown", &3i32).await.unwrap();
1095        let mut stream_b: StreamReceiver<i32> =
1096            client.call_server_stream("countdown", &4i32).await.unwrap();
1097
1098        let (items_a, items_b) = tokio::join!(
1099            async {
1100                let mut items = Vec::new();
1101                while let Some(result) = stream_a.recv().await {
1102                    items.push(result.unwrap());
1103                }
1104                items
1105            },
1106            async {
1107                let mut items = Vec::new();
1108                while let Some(result) = stream_b.recv().await {
1109                    items.push(result.unwrap());
1110                }
1111                items
1112            }
1113        );
1114
1115        assert_eq!(items_a, vec![2, 1, 0]);
1116        assert_eq!(items_b, vec![3, 2, 1, 0]);
1117    }
1118
1119    #[tokio::test]
1120    async fn test_concurrent_unary_calls() {
1121        let config = ChannelConfig::default();
1122        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
1123        let client_channel = MessageChannelAdapter::new(t1);
1124        let server_channel = MessageChannelAdapter::new(t2);
1125
1126        let server = RpcServer::new();
1127        server.register_typed("slow_add", |req: AddRequest| async move {
1128            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1129            Ok(AddResponse {
1130                result: req.a + req.b,
1131            })
1132        });
1133
1134        let _handle = server.spawn_handler(server_channel);
1135
1136        let client = crate::client::RpcClient::new(client_channel);
1137        let _client_handle = client.start();
1138
1139        let start = std::time::Instant::now();
1140        let handles: Vec<_> = (0..5)
1141            .map(|i| {
1142                let client = &client;
1143                async move {
1144                    let resp: AddResponse = client
1145                        .call("slow_add", &AddRequest { a: i, b: i * 10 })
1146                        .await
1147                        .unwrap();
1148                    resp.result
1149                }
1150            })
1151            .collect();
1152
1153        let results = futures::future::join_all(handles).await;
1154        let elapsed = start.elapsed();
1155
1156        assert_eq!(results, vec![0, 11, 22, 33, 44]);
1157        assert!(
1158            elapsed < std::time::Duration::from_millis(200),
1159            "Concurrent calls took {:?}, expected < 200ms",
1160            elapsed
1161        );
1162    }
1163}
xrpc/server.rs

xrpc/
server.rs
