dnet_rpc/parts/
producer.rs

1//! Producer parts.
2
3use std::future::Future;
4
5use futures::{
6    channel::{mpsc::UnboundedSender, oneshot},
7    select, FutureExt, Sink, SinkExt, Stream, StreamExt,
8};
9
10use dportable::{create_non_sync_send_variant_for_wasm, spawn};
11use futures::pin_mut;
12
13use crate::{
14    consumer::RequestId,
15    producer::{self, abortable::AborterToken, StreamResponse},
16    ShutdownType,
17};
18
19macro_rules! wrap_task {
20    ($task:expr, $task_aborter:expr) => {{
21        #[cfg(target_arch = "wasm32")]
22        {
23            if $task_aborter.is_some() {
24                spawn(async move { $task.await })
25                    .map(|join_handle| join_handle.unwrap())
26                    .boxed_local()
27            } else {
28                $task.boxed_local()
29            }
30            .fuse()
31        }
32
33        #[cfg(not(target_arch = "wasm32"))]
34        {
35            if $task_aborter.is_some() {
36                spawn(async move { $task.await })
37                    .map(|join_handle| join_handle.unwrap())
38                    .boxed()
39            } else {
40                $task.boxed()
41            }
42            .fuse()
43        }
44    }};
45}
46
47/// Handle response to customer.
48pub async fn handle_response<Response, S>(
49    message: producer::Message<Response>,
50    sender: &mut S,
51    stop_sender: &mut Option<oneshot::Sender<ShutdownType>>,
52) where
53    S: Sink<producer::Message<Response>> + Unpin,
54{
55    match message {
56        producer::Message::Response { .. } => {
57            if sender.send(message).await.is_err() {
58                if let Some(stop_sender) = stop_sender.take() {
59                    let _ = stop_sender.send(ShutdownType::Closed);
60                }
61            }
62        }
63        _ => {
64            let _ = sender.send(message).await;
65        }
66    }
67}
68
69/// Handle request from consumer.
70pub fn handle_request<F, O, E, R, Response>(
71    id: RequestId,
72    task: F,
73    response_factory: R,
74    reply_sender: UnboundedSender<producer::Message<Response>>,
75    remove_aborter_sender: UnboundedSender<RequestId>,
76    mut abort_receiver: oneshot::Receiver<()>,
77    task_aborter: Option<AborterToken>,
78) where
79    F: Future<Output = Result<O, E>> + SendUnlessWasm + 'static,
80    O: Send + 'static,
81    E: Send + 'static,
82    R: Fn(O) -> Response + Send + 'static,
83    Response: Send + 'static,
84{
85    spawn(async move {
86        let task = wrap_task!(task, task_aborter);
87        pin_mut!(task);
88        select! {
89            result = task => {
90                if let Ok(result) = result {
91                    let response = response_factory(result);
92                    let message = producer::Message::Response { id, response };
93                    let _ = reply_sender.unbounded_send(message);
94                }
95                let _ = remove_aborter_sender.unbounded_send(id);
96            },
97            _ = abort_receiver => {
98                if let Some(task_aborter) = task_aborter {
99                    task_aborter.cancel();
100                }
101            },
102        };
103    });
104}
105
106/// Handle no-ack request from consumer.
107pub fn handle_no_ack_request<F, E>(
108    id: RequestId,
109    task: F,
110    remove_aborter_sender: UnboundedSender<RequestId>,
111    mut abort_receiver: oneshot::Receiver<()>,
112    task_aborter: Option<AborterToken>,
113) where
114    F: Future<Output = Result<(), E>> + SendUnlessWasm + 'static,
115    E: Send + 'static,
116{
117    spawn(async move {
118        let task = wrap_task!(task, task_aborter);
119        pin_mut!(task);
120        select! {
121            _result = task => {
122                let _ = remove_aborter_sender.unbounded_send(id);
123            },
124            _ = abort_receiver => {
125                if let Some(task_aborter) = task_aborter {
126                    task_aborter.cancel();
127                }
128            },
129        };
130    });
131}
132
133/// Handle stream request from consumer.
134pub fn handle_stream_request<F, S, E, O, R, Response>(
135    id: RequestId,
136    task: F,
137    response_factory: R,
138    reply_sender: UnboundedSender<producer::Message<Response>>,
139    remove_aborter_sender: UnboundedSender<RequestId>,
140    mut abort_receiver: oneshot::Receiver<()>,
141    task_aborter: Option<AborterToken>,
142) where
143    F: Future<Output = Result<S, E>> + SendUnlessWasm + 'static,
144    S: Stream<Item = O> + Send + Unpin + 'static,
145    E: Send + 'static,
146    R: Fn(StreamResponse<O>) -> Response + Send + 'static,
147    Response: Send + 'static,
148{
149    spawn(async move {
150        let task = wrap_task!(task, task_aborter);
151
152        let mut stream = match task.await {
153            Ok(stream) => stream.fuse(),
154            Err(_) => {
155                let _ = remove_aborter_sender.unbounded_send(id);
156                return;
157            }
158        };
159
160        let response = StreamResponse::Open;
161        let response = response_factory(response);
162        let message = producer::Message::Response { id, response };
163        let _ = reply_sender.unbounded_send(message);
164        loop {
165            select! {
166                result = stream.next() => {
167                    if let Some(result) = result {
168                        let response = StreamResponse::Item(result);
169                        let response = response_factory(response);
170                        let message = producer::Message::Response { id, response };
171                        let _ = reply_sender.unbounded_send(message);
172                    } else {
173                        let response = StreamResponse::Closed;
174                        let response = response_factory(response);
175                        let message = producer::Message::Response { id, response };
176                        let _ = reply_sender.unbounded_send(message);
177
178                        let _ = remove_aborter_sender.unbounded_send(id);
179                        break;
180                    }
181                },
182                _ = abort_receiver => {
183                    if let Some(task_aborter) = task_aborter {
184                        task_aborter.cancel();
185                    }
186                    break;
187                },
188            };
189        }
190    });
191}
192
193create_non_sync_send_variant_for_wasm! {
194    /// Trait for types implementing [Send] unless running under
195    /// WASM targets - then it does nothing.
196    pub trait SendUnlessWasm: Send {}
197
198    impl<T> SendUnlessWasm for T where T: Send {}
199}
dnet_rpc/parts/producer.rs

dnet_rpc/parts/
producer.rs
