1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use bytes::Bytes;
6use serde_json::Value;
7use tokio::sync::{broadcast, mpsc, oneshot, watch};
8use unb_core::{ClientOperationId, CoreError, Envelope, ErrorCode, Kind};
9
10use crate::cancellation::CancellationToken;
11use crate::client::{ClientDelivery, ClientSession};
12use crate::core_runtime::{ProtocolCoreHandle, SessionOutcome};
13use crate::error::WsError;
14use crate::transport::Pipe;
15
16pub type OnOpened = Box<dyn FnOnce(String) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send>;
17
18pub enum Directive {
19 StartClientOperation {
20 target_path: String,
21 kind: Kind,
22 payload: Bytes,
23 hops: Option<u8>,
24 headers: serde_json::Map<String, Value>,
25 body: Option<crate::BodyStream>,
26 timeout: Option<std::time::Duration>,
27 sender: mpsc::Sender<ClientDelivery>,
28 reply: oneshot::Sender<Result<ClientOperationId, CoreError>>,
29 },
30 OpenStream {
31 target_path: String,
32 kind: Kind,
33 payload: Bytes,
34 hops: Option<u8>,
35 headers: serde_json::Map<String, Value>,
36 body: Option<crate::BodyStream>,
37 opened: Option<OnOpened>,
38 reply: oneshot::Sender<Result<String, CoreError>>,
39 },
40 Send {
41 corr: String,
42 payload: Bytes,
43 },
44 Respond {
45 corr: String,
46 payload: Bytes,
47 headers: serde_json::Map<String, Value>,
48 body: Option<crate::BodyStream>,
49 },
50 Fail {
51 corr: String,
52 code: ErrorCode,
53 message: String,
54 },
55 Cancel {
56 corr: String,
57 },
58 CancelClientOperation {
59 operation: ClientOperationId,
60 },
61 Control {
62 kind: Kind,
63 payload: Bytes,
64 },
65}
66
67pub struct Wire {
68 directives: mpsc::Sender<Directive>,
69 outcome: watch::Receiver<Option<SessionOutcome>>,
70 routes: watch::Receiver<bool>,
71 cancellation: CancellationToken,
72 client: ClientSession,
73 observations: broadcast::Sender<Envelope>,
74 _core: Option<ProtocolCoreHandle>,
75}
76
77impl Wire {
78 pub fn open(transport: Pipe) -> Wire {
79 let (wire, core) =
80 ProtocolCoreHandle::open(transport, &crate::core_runtime::RuntimeHandle::current());
81 let mut wire = Arc::into_inner(wire).expect("standalone wire ownership");
82 wire._core = Some(core);
83 wire
84 }
85
86 pub fn channel(
87 cancellation: CancellationToken,
88 ) -> (
89 Arc<Wire>,
90 mpsc::Receiver<Directive>,
91 watch::Sender<Option<SessionOutcome>>,
92 watch::Sender<bool>,
93 ) {
94 let (directives_tx, directives_rx) = mpsc::channel(64);
95 let (outcome_tx, outcome) = watch::channel(None);
96 let (routes_tx, routes) = watch::channel(false);
97 let (observations, _) = broadcast::channel(256);
98 let wire = Arc::new(Wire {
99 directives: directives_tx.clone(),
100 outcome,
101 routes,
102 cancellation,
103 client: ClientSession::connected(directives_tx.clone()),
104 observations,
105 _core: None,
106 });
107 (wire, directives_rx, outcome_tx, routes_tx)
108 }
109
110 pub(crate) fn standalone_channel(
111 cancellation: CancellationToken,
112 ) -> (
113 Arc<Wire>,
114 mpsc::Receiver<Directive>,
115 watch::Sender<Option<SessionOutcome>>,
116 watch::Sender<bool>,
117 ) {
118 let (directives_tx, directives_rx) = mpsc::channel(64);
119 let (outcome_tx, outcome) = watch::channel(None);
120 let (routes_tx, routes) = watch::channel(false);
121 let (observations, _) = broadcast::channel(256);
122 let wire = Arc::new(Wire {
123 directives: directives_tx.clone(),
124 outcome,
125 routes,
126 cancellation,
127 client: ClientSession::connected(directives_tx.clone()),
128 observations,
129 _core: None,
130 });
131 (wire, directives_rx, outcome_tx, routes_tx)
132 }
133
134 pub async fn open_stream(
135 &self,
136 target_path: &str,
137 kind: Kind,
138 payload: Value,
139 ) -> Result<String, WsError> {
140 self.open_forward(target_path, kind, Envelope::encode_payload(&payload), None)
141 .await
142 }
143
144 pub async fn open_stream_with(
145 &self,
146 target_path: &str,
147 kind: Kind,
148 payload: Value,
149 headers: serde_json::Map<String, Value>,
150 ) -> Result<String, WsError> {
151 let (reply, response) = oneshot::channel();
152 self.directives
153 .send(Directive::OpenStream {
154 target_path: target_path.into(),
155 kind,
156 payload: Envelope::encode_payload(&payload),
157 hops: None,
158 headers,
159 body: None,
160 opened: None,
161 reply,
162 })
163 .await
164 .map_err(|_| WsError::Gone)?;
165 Ok(response.await.map_err(|_| WsError::Gone)??)
166 }
167
168 pub async fn open_stream_streaming(
169 &self,
170 target_path: &str,
171 kind: Kind,
172 body: crate::BodyStream,
173 ) -> Result<String, WsError> {
174 let (reply, response) = oneshot::channel();
175 self.directives
176 .send(Directive::OpenStream {
177 target_path: target_path.into(),
178 kind,
179 payload: Bytes::new(),
180 hops: None,
181 headers: Default::default(),
182 body: Some(body),
183 opened: None,
184 reply,
185 })
186 .await
187 .map_err(|_| WsError::Gone)?;
188 Ok(response.await.map_err(|_| WsError::Gone)??)
189 }
190
191 pub async fn respond_streaming(
192 &self,
193 corr: &str,
194 body: crate::BodyStream,
195 ) -> Result<(), WsError> {
196 self.push(Directive::Respond {
197 corr: corr.into(),
198 payload: Bytes::new(),
199 headers: serde_json::Map::new(),
200 body: Some(body),
201 })
202 .await
203 }
204
205 pub async fn open_forward(
206 &self,
207 target_path: &str,
208 kind: Kind,
209 payload: Bytes,
210 hops: Option<u8>,
211 ) -> Result<String, WsError> {
212 let (reply, response) = oneshot::channel();
213 self.directives
214 .send(Directive::OpenStream {
215 target_path: target_path.into(),
216 kind,
217 payload,
218 hops,
219 headers: Default::default(),
220 body: None,
221 opened: None,
222 reply,
223 })
224 .await
225 .map_err(|_| WsError::Gone)?;
226 Ok(response.await.map_err(|_| WsError::Gone)??)
227 }
228
229 #[allow(clippy::too_many_arguments)]
230 pub async fn open_forward_with<F, Fut>(
231 &self,
232 target_path: &str,
233 kind: Kind,
234 payload: Bytes,
235 hops: Option<u8>,
236 headers: serde_json::Map<String, Value>,
237 body: Option<crate::BodyStream>,
238 opened: F,
239 ) -> Result<String, WsError>
240 where
241 F: FnOnce(String) -> Fut + Send + 'static,
242 Fut: Future<Output = ()> + Send + 'static,
243 {
244 let (reply, response) = oneshot::channel();
245 self.directives
246 .send(Directive::OpenStream {
247 target_path: target_path.into(),
248 kind,
249 payload,
250 hops,
251 headers,
252 body,
253 opened: Some(Box::new(move |corr| Box::pin(opened(corr)))),
254 reply,
255 })
256 .await
257 .map_err(|_| WsError::Gone)?;
258 Ok(response.await.map_err(|_| WsError::Gone)??)
259 }
260
261 pub async fn send(&self, corr: &str, payload: Value) -> Result<(), WsError> {
262 self.send_raw(corr, Envelope::encode_payload(&payload))
263 .await
264 }
265
266 #[inline]
267 pub async fn send_bytes(&self, corr: &str, payload: Bytes) -> Result<(), WsError> {
268 self.send_raw(corr, payload).await
269 }
270
271 #[inline]
272 pub(crate) async fn send_raw(&self, corr: &str, payload: Bytes) -> Result<(), WsError> {
273 self.push(Directive::Send {
274 corr: corr.into(),
275 payload,
276 })
277 .await
278 }
279
280 pub async fn respond(&self, corr: &str, payload: Value) -> Result<(), WsError> {
281 self.push(Directive::Respond {
282 corr: corr.into(),
283 payload: Envelope::encode_payload(&payload),
284 headers: serde_json::Map::new(),
285 body: None,
286 })
287 .await
288 }
289
290 pub async fn respond_with(
291 &self,
292 corr: &str,
293 payload: Bytes,
294 headers: serde_json::Map<String, Value>,
295 ) -> Result<(), WsError> {
296 self.push(Directive::Respond {
297 corr: corr.into(),
298 payload,
299 headers,
300 body: None,
301 })
302 .await
303 }
304
305 pub async fn fail(&self, corr: &str, code: ErrorCode, message: &str) -> Result<(), WsError> {
306 self.push(Directive::Fail {
307 corr: corr.into(),
308 code,
309 message: message.into(),
310 })
311 .await
312 }
313
314 pub async fn cancel(&self, corr: &str) -> Result<(), WsError> {
315 self.push(Directive::Cancel { corr: corr.into() }).await
316 }
317
318 pub async fn control(&self, kind: Kind, payload: Bytes) -> Result<(), WsError> {
319 self.push(Directive::Control { kind, payload }).await
320 }
321
322 pub async fn session_outcome(&self) -> Result<SessionOutcome, WsError> {
323 let mut outcome = self.outcome.clone();
324 loop {
325 if let Some(result) = outcome.borrow().clone() {
326 return Ok(result);
327 }
328 outcome.changed().await.map_err(|_| WsError::Gone)?;
329 }
330 }
331
332 pub async fn routes_acked(&self) -> Result<(), WsError> {
333 let mut routes = self.routes.clone();
334 loop {
335 if *routes.borrow() {
336 return Ok(());
337 }
338 routes.changed().await.map_err(|_| WsError::Gone)?;
339 }
340 }
341
342 pub fn shutdown(&self) {
343 self.cancellation.cancel();
344 }
345
346 pub fn is_closed(&self) -> bool {
347 self.cancellation.is_cancelled()
348 }
349
350 pub fn client_session(&self) -> ClientSession {
351 self.client.clone()
352 }
353
354 pub fn observe(&self) -> broadcast::Receiver<Envelope> {
355 self.observations.subscribe()
356 }
357
358 pub(crate) fn observation_sender(&self) -> broadcast::Sender<Envelope> {
359 self.observations.clone()
360 }
361
362 pub async fn closed(&self) {
363 self.cancellation.cancelled().await;
364 }
365
366 async fn push(&self, directive: Directive) -> Result<(), WsError> {
367 self.directives
368 .send(directive)
369 .await
370 .map_err(|_| WsError::Gone)
371 }
372}