use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use bytes::Bytes;
use serde_json::Value;
use tokio::sync::{broadcast, mpsc, oneshot, watch};
use unb_core::{ClientOperationId, CoreError, Envelope, ErrorCode, Kind};
use crate::cancellation::CancellationToken;
use crate::client::{ClientDelivery, ClientSession};
use crate::core_runtime::{ProtocolCoreHandle, SessionOutcome};
use crate::error::WsError;
use crate::transport::Pipe;
pub type OnOpened = Box<dyn FnOnce(String) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send>;
pub enum Directive {
StartClientOperation {
target_path: String,
kind: Kind,
payload: Bytes,
hops: Option<u8>,
headers: serde_json::Map<String, Value>,
body: Option<crate::BodyStream>,
timeout: Option<std::time::Duration>,
sender: mpsc::Sender<ClientDelivery>,
reply: oneshot::Sender<Result<ClientOperationId, CoreError>>,
},
OpenStream {
target_path: String,
kind: Kind,
payload: Bytes,
hops: Option<u8>,
headers: serde_json::Map<String, Value>,
body: Option<crate::BodyStream>,
opened: Option<OnOpened>,
reply: oneshot::Sender<Result<String, CoreError>>,
},
Send {
corr: String,
payload: Bytes,
},
Respond {
corr: String,
payload: Bytes,
headers: serde_json::Map<String, Value>,
body: Option<crate::BodyStream>,
},
Fail {
corr: String,
code: ErrorCode,
message: String,
},
Cancel {
corr: String,
},
CancelClientOperation {
operation: ClientOperationId,
},
Control {
kind: Kind,
payload: Bytes,
},
}
pub struct Wire {
directives: mpsc::Sender<Directive>,
outcome: watch::Receiver<Option<SessionOutcome>>,
routes: watch::Receiver<bool>,
cancellation: CancellationToken,
client: ClientSession,
observations: broadcast::Sender<Envelope>,
_core: Option<ProtocolCoreHandle>,
}
impl Wire {
pub fn open(transport: Pipe) -> Wire {
let (wire, core) =
ProtocolCoreHandle::open(transport, &crate::core_runtime::RuntimeHandle::current());
let mut wire = Arc::into_inner(wire).expect("standalone wire ownership");
wire._core = Some(core);
wire
}
pub fn channel(
cancellation: CancellationToken,
) -> (
Arc<Wire>,
mpsc::Receiver<Directive>,
watch::Sender<Option<SessionOutcome>>,
watch::Sender<bool>,
) {
let (directives_tx, directives_rx) = mpsc::channel(64);
let (outcome_tx, outcome) = watch::channel(None);
let (routes_tx, routes) = watch::channel(false);
let (observations, _) = broadcast::channel(256);
let wire = Arc::new(Wire {
directives: directives_tx.clone(),
outcome,
routes,
cancellation,
client: ClientSession::connected(directives_tx.clone()),
observations,
_core: None,
});
(wire, directives_rx, outcome_tx, routes_tx)
}
pub(crate) fn standalone_channel(
cancellation: CancellationToken,
) -> (
Arc<Wire>,
mpsc::Receiver<Directive>,
watch::Sender<Option<SessionOutcome>>,
watch::Sender<bool>,
) {
let (directives_tx, directives_rx) = mpsc::channel(64);
let (outcome_tx, outcome) = watch::channel(None);
let (routes_tx, routes) = watch::channel(false);
let (observations, _) = broadcast::channel(256);
let wire = Arc::new(Wire {
directives: directives_tx.clone(),
outcome,
routes,
cancellation,
client: ClientSession::connected(directives_tx.clone()),
observations,
_core: None,
});
(wire, directives_rx, outcome_tx, routes_tx)
}
pub async fn open_stream(
&self,
target_path: &str,
kind: Kind,
payload: Value,
) -> Result<String, WsError> {
self.open_forward(target_path, kind, Envelope::encode_payload(&payload), None)
.await
}
pub async fn open_stream_with(
&self,
target_path: &str,
kind: Kind,
payload: Value,
headers: serde_json::Map<String, Value>,
) -> Result<String, WsError> {
let (reply, response) = oneshot::channel();
self.directives
.send(Directive::OpenStream {
target_path: target_path.into(),
kind,
payload: Envelope::encode_payload(&payload),
hops: None,
headers,
body: None,
opened: None,
reply,
})
.await
.map_err(|_| WsError::Gone)?;
Ok(response.await.map_err(|_| WsError::Gone)??)
}
pub async fn open_stream_streaming(
&self,
target_path: &str,
kind: Kind,
body: crate::BodyStream,
) -> Result<String, WsError> {
let (reply, response) = oneshot::channel();
self.directives
.send(Directive::OpenStream {
target_path: target_path.into(),
kind,
payload: Bytes::new(),
hops: None,
headers: Default::default(),
body: Some(body),
opened: None,
reply,
})
.await
.map_err(|_| WsError::Gone)?;
Ok(response.await.map_err(|_| WsError::Gone)??)
}
pub async fn respond_streaming(
&self,
corr: &str,
body: crate::BodyStream,
) -> Result<(), WsError> {
self.push(Directive::Respond {
corr: corr.into(),
payload: Bytes::new(),
headers: serde_json::Map::new(),
body: Some(body),
})
.await
}
pub async fn open_forward(
&self,
target_path: &str,
kind: Kind,
payload: Bytes,
hops: Option<u8>,
) -> Result<String, WsError> {
let (reply, response) = oneshot::channel();
self.directives
.send(Directive::OpenStream {
target_path: target_path.into(),
kind,
payload,
hops,
headers: Default::default(),
body: None,
opened: None,
reply,
})
.await
.map_err(|_| WsError::Gone)?;
Ok(response.await.map_err(|_| WsError::Gone)??)
}
#[allow(clippy::too_many_arguments)]
pub async fn open_forward_with<F, Fut>(
&self,
target_path: &str,
kind: Kind,
payload: Bytes,
hops: Option<u8>,
headers: serde_json::Map<String, Value>,
body: Option<crate::BodyStream>,
opened: F,
) -> Result<String, WsError>
where
F: FnOnce(String) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let (reply, response) = oneshot::channel();
self.directives
.send(Directive::OpenStream {
target_path: target_path.into(),
kind,
payload,
hops,
headers,
body,
opened: Some(Box::new(move |corr| Box::pin(opened(corr)))),
reply,
})
.await
.map_err(|_| WsError::Gone)?;
Ok(response.await.map_err(|_| WsError::Gone)??)
}
pub async fn send(&self, corr: &str, payload: Value) -> Result<(), WsError> {
self.send_raw(corr, Envelope::encode_payload(&payload))
.await
}
#[inline]
pub async fn send_bytes(&self, corr: &str, payload: Bytes) -> Result<(), WsError> {
self.send_raw(corr, payload).await
}
#[inline]
pub(crate) async fn send_raw(&self, corr: &str, payload: Bytes) -> Result<(), WsError> {
self.push(Directive::Send {
corr: corr.into(),
payload,
})
.await
}
pub async fn respond(&self, corr: &str, payload: Value) -> Result<(), WsError> {
self.push(Directive::Respond {
corr: corr.into(),
payload: Envelope::encode_payload(&payload),
headers: serde_json::Map::new(),
body: None,
})
.await
}
pub async fn respond_with(
&self,
corr: &str,
payload: Bytes,
headers: serde_json::Map<String, Value>,
) -> Result<(), WsError> {
self.push(Directive::Respond {
corr: corr.into(),
payload,
headers,
body: None,
})
.await
}
pub async fn fail(&self, corr: &str, code: ErrorCode, message: &str) -> Result<(), WsError> {
self.push(Directive::Fail {
corr: corr.into(),
code,
message: message.into(),
})
.await
}
pub async fn cancel(&self, corr: &str) -> Result<(), WsError> {
self.push(Directive::Cancel { corr: corr.into() }).await
}
pub async fn control(&self, kind: Kind, payload: Bytes) -> Result<(), WsError> {
self.push(Directive::Control { kind, payload }).await
}
pub async fn session_outcome(&self) -> Result<SessionOutcome, WsError> {
let mut outcome = self.outcome.clone();
loop {
if let Some(result) = outcome.borrow().clone() {
return Ok(result);
}
outcome.changed().await.map_err(|_| WsError::Gone)?;
}
}
pub async fn routes_acked(&self) -> Result<(), WsError> {
let mut routes = self.routes.clone();
loop {
if *routes.borrow() {
return Ok(());
}
routes.changed().await.map_err(|_| WsError::Gone)?;
}
}
pub fn shutdown(&self) {
self.cancellation.cancel();
}
pub fn is_closed(&self) -> bool {
self.cancellation.is_cancelled()
}
pub fn client_session(&self) -> ClientSession {
self.client.clone()
}
pub fn observe(&self) -> broadcast::Receiver<Envelope> {
self.observations.subscribe()
}
pub(crate) fn observation_sender(&self) -> broadcast::Sender<Envelope> {
self.observations.clone()
}
pub async fn closed(&self) {
self.cancellation.cancelled().await;
}
async fn push(&self, directive: Directive) -> Result<(), WsError> {
self.directives
.send(directive)
.await
.map_err(|_| WsError::Gone)
}
}