std_modrpc/role_impls/
multi_stream_receiver.rs

1use core::{
2    cell::RefCell,
3    marker::PhantomData,
4};
5use crate::{
6    proto::{
7        MultiStreamInitState,
8        MultiStreamItem,
9        MultiStreamItemLazy,
10        MultiStreamId,
11        MultiStreamReceiverConfig,
12    },
13    receive_stream::{ReceiveStream, StreamState},
14};
15use std::collections::HashMap;
16use std::rc::Rc;
17use modrpc::RoleSetup;
18
19pub enum ReceiveMultiStreamNextError {
20    Shutdown,
21    DecodeItem(mproto::DecodeError),
22}
23
24pub struct ReceiveMultiStream<T> {
25    stream_id: MultiStreamId,
26    receive_stream: ReceiveStream,
27    phantom: PhantomData<T>,
28}
29
30struct BrokerState {
31    streams: RefCell<HashMap<u32, Rc<StreamState>>>,
32}
33
34pub struct MultiStreamReceiver<T> {
35    hooks: crate::MultiStreamReceiverHooks<T>,
36    broker_state: Rc<BrokerState>,
37}
38
39pub struct MultiStreamReceiverBuilder<T> {
40    name: &'static str,
41    hooks: crate::MultiStreamReceiverHooks<T>,
42    stubs: crate::MultiStreamReceiverStubs<T>,
43
44    broker_state: Rc<BrokerState>,
45}
46
47impl<T: mproto::Owned> MultiStreamReceiver<T> {
48    pub fn new_stream(&self, stream_id: MultiStreamId, next_seq: Option<u64>) -> ReceiveMultiStream<T> {
49        let receive_stream = ReceiveStream::new(next_seq);
50        self.broker_state.streams.borrow_mut()
51            .insert(stream_id.id, receive_stream.stream_state().clone());
52
53        ReceiveMultiStream {
54            stream_id,
55            receive_stream,
56            phantom: PhantomData,
57        }
58    }
59}
60
61impl<T> Clone for MultiStreamReceiver<T> {
62    fn clone(&self) -> Self {
63        Self {
64            hooks: self.hooks.clone(),
65            broker_state: self.broker_state.clone(),
66        }
67    }
68}
69
70impl<T: mproto::Owned> ReceiveMultiStream<T> {
71    pub fn id(&self) -> MultiStreamId {
72        self.stream_id
73    }
74
75    pub async fn next(&mut self) -> Result<Option<T>, ReceiveMultiStreamNextError> {
76        use mproto::BaseLen;
77
78        let packet = self.receive_stream.next_packet().await;
79
80        let stream_item: MultiStreamItemLazy<T> = mproto::decode_value(
81            &packet.as_ref()[modrpc::TransmitPacket::BASE_LEN..]
82        )
83        .map_err(|e| ReceiveMultiStreamNextError::DecodeItem(e))?;
84
85        let owned_result = stream_item.payload()
86            .map_err(|e| ReceiveMultiStreamNextError::DecodeItem(e))?
87            .map(|i| T::lazy_to_owned(i))
88            .transpose()
89            .map_err(|e| ReceiveMultiStreamNextError::DecodeItem(e))?;
90
91        Ok(owned_result)
92    }
93
94    pub async fn next_lazy(&mut self)
95        -> Result<mproto::LazyBuf<Option<T>, modrpc::Packet>, ReceiveMultiStreamNextError>
96    {
97        use mproto::BaseLen;
98
99        let packet = self.receive_stream.next_packet().await;
100        packet.advance(modrpc::TransmitPacket::BASE_LEN);
101
102        let stream_item: mproto::LazyBuf<MultiStreamItem<T>, _> = mproto::LazyBuf::new(packet);
103
104        Ok(stream_item.map(|s| s.payload().unwrap()))
105    }
106
107    pub fn with_try_next<R>(
108        &mut self,
109        f: impl FnOnce(Option<mproto::DecodeResult<Option<T::Lazy<'_>>>>) -> R,
110    ) -> R {
111        use mproto::BaseLen;
112
113        let Some(packet) = self.receive_stream.try_next_packet() else {
114            return f(None);
115        };
116
117        let stream_item =
118            match mproto::decode_value::<MultiStreamItemLazy<T>>(
119                &packet.as_ref()[modrpc::TransmitPacket::BASE_LEN..]
120            ) {
121                Ok(x) => x,
122                Err(e) => {
123                    return f(Some(Err(e)));
124                },
125            };
126
127        let payload =
128            match stream_item.payload() {
129                Ok(x) => x,
130                Err(e) => {
131                    return f(Some(Err(e)));
132                },
133            };
134
135        f(Some(Ok(payload)))
136    }
137
138    pub async fn with_next<'a, Fut, R>(
139        &mut self,
140        f: impl FnOnce(mproto::DecodeResult<Option<T::Lazy<'_>>>) -> Fut,
141    ) -> Option<R>
142        where Fut: std::future::Future<Output = R>
143    {
144        use mproto::BaseLen;
145
146        let packet = self.receive_stream.next_packet().await;
147
148        let stream_item =
149            match mproto::decode_value::<MultiStreamItemLazy<T>>(
150                &packet.as_ref()[modrpc::TransmitPacket::BASE_LEN..]
151            ) {
152                Ok(x) => x,
153                Err(e) => {
154                    return Some(f(Err(e)).await);
155                },
156            };
157
158        let payload =
159            match stream_item.payload() {
160                Ok(x) => x,
161                Err(e) => {
162                    return Some(f(Err(e)).await);
163                },
164            };
165
166        Some(f(Ok(payload)).await)
167    }
168
169    pub async fn with_next_sync<'a, R>(
170        &mut self,
171        f: impl FnOnce(mproto::DecodeResult<T::Lazy<'_>>) -> R,
172    ) -> Option<R> {
173        use mproto::BaseLen;
174
175        let packet = self.receive_stream.next_packet().await;
176
177        let stream_item =
178            match mproto::decode_value::<MultiStreamItemLazy<T>>(
179                &packet.as_ref()[modrpc::TransmitPacket::BASE_LEN..]
180            ) {
181                Ok(x) => x,
182                Err(e) => {
183                    return Some(f(Err(e)));
184                }
185            };
186
187        let payload =
188            match stream_item.payload() {
189                Ok(Some(x)) => x,
190                Ok(None) => {
191                    // End of stream
192                    return None;
193                }
194                Err(e) => {
195                    return Some(f(Err(e)));
196                }
197            };
198
199        Some(f(Ok(payload)))
200    }
201
202    pub async fn collect(&mut self) -> Result<Vec<T>, ReceiveMultiStreamNextError> {
203        let mut collected = Vec::new();
204        while let Some(item) = self.next().await? {
205            collected.push(item);
206        }
207        Ok(collected)
208    }
209}
210
211#[derive(thiserror::Error, Debug)]
212pub enum MultiStreamTryCollectError<E: std::fmt::Debug> {
213    #[error("failed to decode MultiStream item")]
214    DecodeError(#[from] mproto::DecodeError),
215    #[error("stream sender failed: {0:?}")]
216    SenderError(E),
217    #[error("plane is shutting down")]
218    Shutdown,
219}
220
221impl<E: std::fmt::Debug> From<ReceiveMultiStreamNextError> for MultiStreamTryCollectError<E> {
222    fn from(other: ReceiveMultiStreamNextError) -> Self {
223        match other {
224            ReceiveMultiStreamNextError::DecodeItem(e) =>
225                MultiStreamTryCollectError::DecodeError(e),
226            ReceiveMultiStreamNextError::Shutdown =>
227                MultiStreamTryCollectError::Shutdown,
228        }
229    }
230}
231
232impl<T: mproto::Owned, E: mproto::Owned + std::fmt::Debug> ReceiveMultiStream<Result<T, E>> {
233    pub async fn try_collect(&mut self) -> Result<Vec<T>, MultiStreamTryCollectError<E>> {
234        let mut collected = Vec::new();
235        while let Some(item) =
236            self.next().await?
237                .transpose()
238                .map_err(|e| MultiStreamTryCollectError::SenderError(e))?
239        {
240            collected.push(item);
241        }
242        Ok(collected)
243    }
244}
245
246impl<T: mproto::Owned> MultiStreamReceiverBuilder<T> {
247    pub fn new(
248        name: &'static str,
249        hooks: crate::MultiStreamReceiverHooks<T>,
250        stubs: crate::MultiStreamReceiverStubs<T>,
251        _config: &MultiStreamReceiverConfig,
252        _init: MultiStreamInitState,
253    ) -> Self {
254        Self {
255            name, hooks, stubs,
256            broker_state: Rc::new(BrokerState {
257                streams: RefCell::new(HashMap::new()),
258            }),
259        }
260    }
261
262    pub fn create_handle(
263        &self,
264        _setup: &RoleSetup,
265    ) -> crate::MultiStreamReceiver<T> {
266        crate::MultiStreamReceiver {
267            hooks: self.hooks.clone(),
268            broker_state: self.broker_state.clone(),
269        }
270    }
271
272    pub fn build(
273        self,
274        setup: &RoleSetup,
275    ) {
276        use mproto::BaseLen;
277
278        let broker_state = self.broker_state;
279        self.stubs.item.inline_untyped(setup, move |_source, packet| {
280            let stream_item_bytes = &packet[modrpc::TransmitPacket::BASE_LEN..];
281            let (seq, stream_id, shutdown) = {
282                let Ok(stream_item) =
283                    mproto::decode_value::<MultiStreamItemLazy<T>>(stream_item_bytes)
284                else {
285                    return;
286                };
287                let Ok(seq) = stream_item.seq() else {
288                    return;
289                };
290                let Ok(stream_id) = stream_item.stream_id().and_then(|r| r.id()) else {
291                    return;
292                };
293                let Ok(payload) = stream_item.payload() else {
294                    return;
295                };
296                (seq, stream_id, payload.is_none())
297            };
298
299            let Some(stream_state) = broker_state.streams.borrow().get(&stream_id).cloned() else {
300                log::warn!("Unknown stream_id name={} stream_id={stream_id}", self.name);
301                return;
302            };
303
304            let stream_is_done = stream_state.handle_item(seq, shutdown, packet.clone());
305            if stream_is_done {
306                log::debug!("MultiStreamReciever shutdown stream stream_id={stream_id} seq={seq}");
307                broker_state.streams.borrow_mut().remove(&stream_id);
308            }
309        })
310        .subscribe();
311    }
312}
313
std_modrpc/role_impls/multi_stream_receiver.rs

std_modrpc/role_impls/
multi_stream_receiver.rs
