xrpc/
streaming.rs

1use bytes::Bytes;
2use futures::Stream;
3use parking_lot::Mutex;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::pin::Pin;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::task::{Context, Poll};
10use tokio::sync::{Mutex as AsyncMutex, mpsc};
11
12use crate::channel::message::MessageChannel;
13use crate::codec::{BincodeCodec, Codec};
14use crate::error::{Result, RpcError};
15use crate::message::Message;
16use crate::message::types::{MessageId, MessageType};
17
18pub type StreamId = u64;
19
20type StreamSenders = Arc<Mutex<HashMap<StreamId, mpsc::UnboundedSender<Result<Bytes>>>>>;
21
22struct StreamRegistration {
23    stream_id: StreamId,
24    streams: StreamSenders,
25}
26
27impl Drop for StreamRegistration {
28    fn drop(&mut self) {
29        self.streams.lock().remove(&self.stream_id);
30    }
31}
32
33static STREAM_ID_COUNTER: AtomicU64 = AtomicU64::new(1);
34
35pub fn next_stream_id() -> StreamId {
36    STREAM_ID_COUNTER.fetch_add(1, Ordering::Relaxed)
37}
38
39pub struct StreamSender<T: MessageChannel, C: Codec = BincodeCodec> {
40    stream_id: StreamId,
41    sequence: AtomicU64,
42    transport: Arc<T>,
43    codec: C,
44    operation: AsyncMutex<()>,
45    ended: std::sync::atomic::AtomicBool,
46}
47
48impl<T: MessageChannel> StreamSender<T, BincodeCodec> {
49    pub fn new(stream_id: StreamId, transport: Arc<T>) -> Self {
50        Self {
51            stream_id,
52            sequence: AtomicU64::new(0),
53            transport,
54            codec: BincodeCodec,
55            operation: AsyncMutex::new(()),
56            ended: std::sync::atomic::AtomicBool::new(false),
57        }
58    }
59}
60
61impl<T: MessageChannel, C: Codec> StreamSender<T, C> {
62    pub fn with_codec(stream_id: StreamId, transport: Arc<T>, codec: C) -> Self {
63        Self {
64            stream_id,
65            sequence: AtomicU64::new(0),
66            transport,
67            codec,
68            operation: AsyncMutex::new(()),
69            ended: std::sync::atomic::AtomicBool::new(false),
70        }
71    }
72
73    pub async fn send<D: Serialize>(&self, data: D) -> Result<()> {
74        let _operation = self.operation.lock().await;
75        if self.ended.load(Ordering::Acquire) {
76            return Err(RpcError::StreamError("Stream already ended".to_string()));
77        }
78
79        let seq = self.sequence.load(Ordering::Relaxed);
80        let payload = self.codec.encode(&data)?;
81        let chunk = Message::new(
82            MessageId::new(),
83            MessageType::StreamChunk,
84            "",
85            Bytes::from(payload),
86            crate::message::metadata::MessageMetadata::new().with_stream(self.stream_id, seq),
87        );
88        self.transport
89            .send(&chunk)
90            .await
91            .map_err(RpcError::Transport)?;
92        self.sequence.fetch_add(1, Ordering::Relaxed);
93        Ok(())
94    }
95
96    pub async fn end(&self) -> Result<()> {
97        let _operation = self.operation.lock().await;
98        if self.ended.load(Ordering::Acquire) {
99            return Ok(());
100        }
101
102        let end_msg: Message = Message::stream_end(self.stream_id);
103        self.transport
104            .send(&end_msg)
105            .await
106            .map_err(RpcError::Transport)?;
107        self.ended.store(true, Ordering::Release);
108        Ok(())
109    }
110
111    pub fn stream_id(&self) -> StreamId {
112        self.stream_id
113    }
114
115    pub fn is_ended(&self) -> bool {
116        self.ended.load(Ordering::Acquire)
117    }
118}
119
120/// Receives stream chunks. Assumes transport delivers messages in order.
121pub struct StreamReceiver<D, C: Codec = BincodeCodec> {
122    stream_id: StreamId,
123    rx: mpsc::UnboundedReceiver<Result<Bytes>>,
124    ended: bool,
125    codec: C,
126    registration: Option<StreamRegistration>,
127    on_terminate: Option<Box<dyn FnOnce(StreamId) + Send + Sync>>,
128    _phantom: std::marker::PhantomData<D>,
129}
130
131impl<D, C: Codec> StreamReceiver<D, C> {
132    fn finish(&mut self) {
133        self.ended = true;
134        self.registration.take();
135        if let Some(on_terminate) = self.on_terminate.take() {
136            on_terminate(self.stream_id);
137        }
138    }
139
140    pub(crate) fn set_termination_callback(
141        &mut self,
142        callback: impl FnOnce(StreamId) + Send + Sync + 'static,
143    ) {
144        self.on_terminate = Some(Box::new(callback));
145    }
146}
147
148impl<D, C> StreamReceiver<D, C>
149where
150    D: for<'de> Deserialize<'de>,
151    C: Codec,
152{
153    fn new(
154        stream_id: StreamId,
155        rx: mpsc::UnboundedReceiver<Result<Bytes>>,
156        codec: C,
157        registration: StreamRegistration,
158    ) -> Self {
159        Self {
160            stream_id,
161            rx,
162            ended: false,
163            codec,
164            registration: Some(registration),
165            on_terminate: None,
166            _phantom: std::marker::PhantomData,
167        }
168    }
169
170    pub fn stream_id(&self) -> StreamId {
171        self.stream_id
172    }
173
174    pub fn is_ended(&self) -> bool {
175        self.ended
176    }
177
178    pub async fn recv(&mut self) -> Option<Result<D>> {
179        if self.ended {
180            return None;
181        }
182
183        match self.rx.recv().await {
184            Some(Ok(data)) => Some(self.codec.decode(&data)),
185            Some(Err(e)) => {
186                self.finish();
187                Some(Err(e))
188            }
189            None => {
190                self.finish();
191                None
192            }
193        }
194    }
195
196    pub async fn collect(mut self) -> Result<Vec<D>> {
197        let mut items = Vec::new();
198        while let Some(result) = self.recv().await {
199            items.push(result?);
200        }
201        Ok(items)
202    }
203
204    pub fn cancel(&mut self) {
205        self.finish();
206        self.rx.close();
207    }
208}
209
210impl<D, C> Stream for StreamReceiver<D, C>
211where
212    D: for<'de> Deserialize<'de> + Unpin,
213    C: Codec + Unpin,
214{
215    type Item = Result<D>;
216
217    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
218        if self.ended {
219            return Poll::Ready(None);
220        }
221
222        match self.rx.poll_recv(cx) {
223            Poll::Ready(Some(Ok(data))) => {
224                let result = self.codec.decode(&data);
225                Poll::Ready(Some(result))
226            }
227            Poll::Ready(Some(Err(e))) => {
228                self.finish();
229                Poll::Ready(Some(Err(e)))
230            }
231            Poll::Ready(None) => {
232                self.finish();
233                Poll::Ready(None)
234            }
235            Poll::Pending => Poll::Pending,
236        }
237    }
238}
239
240impl<D, C: Codec> Drop for StreamReceiver<D, C> {
241    fn drop(&mut self) {
242        self.finish();
243    }
244}
245
246pub struct StreamManager<C: Codec = BincodeCodec> {
247    streams: StreamSenders,
248    codec: C,
249}
250
251impl StreamManager<BincodeCodec> {
252    pub fn new() -> Self {
253        Self {
254            streams: Arc::new(Mutex::new(HashMap::new())),
255            codec: BincodeCodec,
256        }
257    }
258}
259
260impl<C: Codec> StreamManager<C> {
261    pub fn with_codec(codec: C) -> Self {
262        Self {
263            streams: Arc::new(Mutex::new(HashMap::new())),
264            codec,
265        }
266    }
267
268    pub fn create_receiver<D>(&self, stream_id: StreamId) -> StreamReceiver<D, C>
269    where
270        D: for<'de> Deserialize<'de>,
271        C: Clone,
272    {
273        let (tx, rx) = mpsc::unbounded_channel();
274        self.streams.lock().insert(stream_id, tx);
275        let registration = StreamRegistration {
276            stream_id,
277            streams: Arc::clone(&self.streams),
278        };
279        StreamReceiver::new(stream_id, rx, self.codec.clone(), registration)
280    }
281
282    pub fn handle_message(&self, message: &Message<C>) -> bool {
283        let stream_id = match message.metadata.stream_id {
284            Some(id) => id,
285            None => return false,
286        };
287
288        let streams = self.streams.lock();
289        let sender = match streams.get(&stream_id) {
290            Some(tx) => tx,
291            None => return false,
292        };
293
294        match message.msg_type {
295            MessageType::StreamChunk => {
296                let _ = sender.send(Ok(message.payload.clone()));
297                true
298            }
299            MessageType::StreamEnd => {
300                drop(streams);
301                self.remove_stream(stream_id);
302                true
303            }
304            _ => false,
305        }
306    }
307
308    /// Forward an error to the stream receiver and close the stream.
309    pub fn send_error(&self, stream_id: StreamId, error_msg: String) {
310        let sender = self.streams.lock().remove(&stream_id);
311        if let Some(sender) = sender {
312            let _ = sender.send(Err(RpcError::ServerError(error_msg)));
313        }
314    }
315
316    pub fn remove_stream(&self, stream_id: StreamId) {
317        self.streams.lock().remove(&stream_id);
318    }
319
320    pub(crate) fn close_all(&self, error: RpcError) {
321        let streams = {
322            let mut streams = self.streams.lock();
323            std::mem::take(&mut *streams)
324        };
325        for sender in streams.into_values() {
326            let _ = sender.send(Err(error.clone()));
327        }
328    }
329
330    pub fn active_stream_count(&self) -> usize {
331        self.streams.lock().len()
332    }
333}
334
335impl Default for StreamManager<BincodeCodec> {
336    fn default() -> Self {
337        Self::new()
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use crate::channel::message::MessageChannelAdapter;
345    use crate::error::{TransportError, TransportResult};
346    use crate::transport::channel::{ChannelConfig, ChannelFrameTransport};
347    use std::sync::atomic::AtomicUsize;
348
349    #[derive(Debug, Default)]
350    struct SerializedSendChannel {
351        chunk_release: tokio::sync::Notify,
352        chunk_started: tokio::sync::Notify,
353        committed: Mutex<Vec<MessageType>>,
354    }
355
356    #[async_trait::async_trait]
357    impl MessageChannel for SerializedSendChannel {
358        async fn send(&self, message: &Message) -> TransportResult<()> {
359            if message.msg_type == MessageType::StreamChunk {
360                self.chunk_started.notify_one();
361                self.chunk_release.notified().await;
362            }
363            self.committed.lock().push(message.msg_type);
364            Ok(())
365        }
366
367        async fn recv(&self) -> TransportResult<Message> {
368            std::future::pending().await
369        }
370
371        fn is_connected(&self) -> bool {
372            true
373        }
374
375        async fn close(&self) -> TransportResult<()> {
376            Ok(())
377        }
378    }
379
380    #[derive(Debug, Default)]
381    struct RetryEndChannel {
382        attempts: AtomicUsize,
383        first_started: tokio::sync::Notify,
384    }
385
386    #[async_trait::async_trait]
387    impl MessageChannel for RetryEndChannel {
388        async fn send(&self, message: &Message) -> TransportResult<()> {
389            assert_eq!(message.msg_type, MessageType::StreamEnd);
390            match self.attempts.fetch_add(1, Ordering::SeqCst) {
391                0 => {
392                    self.first_started.notify_one();
393                    std::future::pending().await
394                }
395                1 => Err(TransportError::Timeout {
396                    duration_ms: 1,
397                    operation: "stream end".to_string(),
398                }),
399                _ => Ok(()),
400            }
401        }
402
403        async fn recv(&self) -> TransportResult<Message> {
404            std::future::pending().await
405        }
406
407        fn is_connected(&self) -> bool {
408            true
409        }
410
411        async fn close(&self) -> TransportResult<()> {
412            Ok(())
413        }
414    }
415
416    #[tokio::test]
417    async fn test_stream_send_and_end_commit_in_serial_order() {
418        let transport = Arc::new(SerializedSendChannel::default());
419        let sender = Arc::new(StreamSender::new(next_stream_id(), Arc::clone(&transport)));
420        let chunk_sender = Arc::clone(&sender);
421        let chunk = tokio::spawn(async move { chunk_sender.send(1).await });
422
423        tokio::time::timeout(
424            std::time::Duration::from_secs(1),
425            transport.chunk_started.notified(),
426        )
427        .await
428        .unwrap();
429        let end_sender = Arc::clone(&sender);
430        let mut end = tokio::spawn(async move { end_sender.end().await });
431        assert!(
432            tokio::time::timeout(std::time::Duration::from_millis(50), &mut end)
433                .await
434                .is_err()
435        );
436
437        transport.chunk_release.notify_one();
438        chunk.await.unwrap().unwrap();
439        end.await.unwrap().unwrap();
440
441        assert_eq!(
442            *transport.committed.lock(),
443            vec![MessageType::StreamChunk, MessageType::StreamEnd]
444        );
445        assert!(matches!(
446            sender.send(2).await,
447            Err(RpcError::StreamError(_))
448        ));
449    }
450
451    #[tokio::test]
452    async fn test_stream_end_can_retry_after_cancellation_and_failure() {
453        let transport = Arc::new(RetryEndChannel::default());
454        let sender = Arc::new(StreamSender::new(next_stream_id(), Arc::clone(&transport)));
455        let first_sender = Arc::clone(&sender);
456        let first = tokio::spawn(async move { first_sender.end().await });
457
458        tokio::time::timeout(
459            std::time::Duration::from_secs(1),
460            transport.first_started.notified(),
461        )
462        .await
463        .unwrap();
464        first.abort();
465        assert!(first.await.unwrap_err().is_cancelled());
466        assert!(!sender.is_ended());
467
468        assert!(matches!(
469            sender.end().await,
470            Err(RpcError::Transport(TransportError::Timeout { .. }))
471        ));
472        assert!(!sender.is_ended());
473
474        sender.end().await.unwrap();
475        assert!(sender.is_ended());
476        sender.end().await.unwrap();
477        assert_eq!(transport.attempts.load(Ordering::SeqCst), 3);
478    }
479
480    #[tokio::test]
481    async fn test_stream_sender_receiver() {
482        let config = ChannelConfig::default();
483        let (t1, t2) = ChannelFrameTransport::create_pair("test", config).unwrap();
484
485        let sender_channel = Arc::new(MessageChannelAdapter::new(t1));
486        let receiver_channel = MessageChannelAdapter::new(t2);
487
488        let stream_id = next_stream_id();
489        let sender = StreamSender::new(stream_id, sender_channel);
490
491        let manager = StreamManager::new();
492        let receiver: StreamReceiver<i32> = manager.create_receiver(stream_id);
493
494        let recv_handle = tokio::spawn(async move {
495            loop {
496                let msg = receiver_channel.recv().await.unwrap();
497                if !manager.handle_message(&msg) {
498                    break;
499                }
500                if msg.msg_type == MessageType::StreamEnd {
501                    break;
502                }
503            }
504            receiver
505        });
506
507        sender.send(1i32).await.unwrap();
508        sender.send(2i32).await.unwrap();
509        sender.send(3i32).await.unwrap();
510        sender.end().await.unwrap();
511
512        let mut receiver = recv_handle.await.unwrap();
513
514        let items: Vec<i32> = vec![
515            receiver.recv().await.unwrap().unwrap(),
516            receiver.recv().await.unwrap().unwrap(),
517            receiver.recv().await.unwrap().unwrap(),
518        ];
519        assert_eq!(items, vec![1, 2, 3]);
520        assert!(receiver.recv().await.is_none());
521    }
522
523    #[tokio::test]
524    async fn test_stream_end_is_normal_completion() {
525        let manager = StreamManager::new();
526        let stream_id = next_stream_id();
527        let mut receiver: StreamReceiver<i32> = manager.create_receiver(stream_id);
528        let end: Message = Message::stream_end(stream_id);
529
530        assert!(manager.handle_message(&end));
531        assert!(receiver.recv().await.is_none());
532        assert_eq!(manager.active_stream_count(), 0);
533    }
534
535    #[tokio::test]
536    async fn test_close_all_reports_connection_closed() {
537        let manager = StreamManager::new();
538        let mut first: StreamReceiver<i32> = manager.create_receiver(next_stream_id());
539        let mut second: StreamReceiver<i32> = manager.create_receiver(next_stream_id());
540
541        manager.close_all(RpcError::ConnectionClosed);
542
543        assert!(matches!(
544            first.recv().await,
545            Some(Err(RpcError::ConnectionClosed))
546        ));
547        assert!(matches!(
548            second.recv().await,
549            Some(Err(RpcError::ConnectionClosed))
550        ));
551        assert_eq!(manager.active_stream_count(), 0);
552    }
553
554    #[test]
555    fn test_dropped_receiver_removes_stream_registration() {
556        let manager = StreamManager::new();
557        let receiver: StreamReceiver<i32> = manager.create_receiver(next_stream_id());
558
559        assert_eq!(manager.active_stream_count(), 1);
560        drop(receiver);
561        assert_eq!(manager.active_stream_count(), 0);
562    }
563
564    #[tokio::test]
565    async fn test_close_all_preserves_supplied_error() {
566        let manager = StreamManager::new();
567        let mut receiver: StreamReceiver<i32> = manager.create_receiver(next_stream_id());
568
569        manager.close_all(RpcError::Transport(crate::error::TransportError::Protocol(
570            "invalid frame".to_string(),
571        )));
572
573        assert!(matches!(
574            receiver.recv().await,
575            Some(Err(RpcError::Transport(crate::error::TransportError::Protocol(error))))
576                if error == "invalid frame"
577        ));
578    }
579
580    #[test]
581    fn test_stream_id_generation() {
582        let id1 = next_stream_id();
583        let id2 = next_stream_id();
584        assert_ne!(id1, id2);
585    }
586}
xrpc/streaming.rs

xrpc/
streaming.rs
