use serde::Serialize;
use serde::de::DeserializeOwned;
use tokio::sync::{mpsc, oneshot};
use crate::actor::{Actor, AsyncHandler, Handler, Message, RemoteMessage};
use crate::instrument;
use crate::wire::{
AsyncTypedEnvelope, EnvelopeProxy, ForwardedEnvelope, RemoteInvocation, ReplySender,
ResponseRegistry, SendError, TypedEnvelope, next_call_id,
};
pub struct Endpoint<A: Actor> {
inner: EndpointInner<A>,
}
impl<A: Actor> Clone for Endpoint<A> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
enum EndpointInner<A: Actor> {
Local {
mailbox_tx: mpsc::UnboundedSender<Box<dyn EnvelopeProxy<A>>>,
},
Remote {
label: String,
wire_tx: mpsc::UnboundedSender<RemoteInvocation>,
response_registry: ResponseRegistry,
},
}
impl<A: Actor> Clone for EndpointInner<A> {
fn clone(&self) -> Self {
match self {
Self::Local { mailbox_tx } => Self::Local {
mailbox_tx: mailbox_tx.clone(),
},
Self::Remote {
label,
wire_tx,
response_registry,
} => Self::Remote {
label: label.clone(),
wire_tx: wire_tx.clone(),
response_registry: response_registry.clone(),
},
}
}
}
impl<A: Actor + 'static> Endpoint<A> {
pub(crate) fn local(mailbox_tx: mpsc::UnboundedSender<Box<dyn EnvelopeProxy<A>>>) -> Self {
Self {
inner: EndpointInner::Local { mailbox_tx },
}
}
pub(crate) fn remote(
label: String,
wire_tx: mpsc::UnboundedSender<RemoteInvocation>,
response_registry: ResponseRegistry,
) -> Self {
Self {
inner: EndpointInner::Remote {
label,
wire_tx,
response_registry,
},
}
}
pub async fn send<M>(&self, message: M) -> Result<M::Result, SendError>
where
A: Handler<M>,
M: RemoteMessage,
M::Result: Serialize + DeserializeOwned,
{
let actor_type = std::any::type_name::<A>();
match &self.inner {
EndpointInner::Local { mailbox_tx } => {
instrument::send_local(actor_type);
let (tx, rx) = oneshot::channel();
let envelope = TypedEnvelope {
message,
respond_to: tx,
};
mailbox_tx.send(Box::new(envelope)).map_err(|_| {
instrument::send_error(actor_type, "mailbox_closed");
SendError::MailboxClosed
})?;
rx.await.map_err(|_| {
instrument::send_error(actor_type, "response_dropped");
SendError::ResponseDropped
})
}
EndpointInner::Remote {
label,
wire_tx,
response_registry,
} => {
instrument::send_remote(actor_type);
#[cfg(feature = "monitor")]
let start = std::time::Instant::now();
let call_id = next_call_id();
let payload = bincode::serde::encode_to_vec(&message, bincode::config::standard())
.map_err(|e| {
instrument::send_error(actor_type, "serialization_failed");
SendError::SerializationFailed(e.to_string())
})?;
let (resp_tx, resp_rx) = oneshot::channel();
response_registry.register(call_id, resp_tx);
let invocation = RemoteInvocation {
call_id,
actor_label: label.clone(),
message_type: M::TYPE_ID.to_string(),
payload,
};
wire_tx.send(invocation).map_err(|_| {
instrument::send_error(actor_type, "wire_closed");
SendError::WireClosed
})?;
let response = resp_rx.await.map_err(|_| {
instrument::send_error(actor_type, "response_dropped");
SendError::ResponseDropped
})?;
let result_bytes = response.result.map_err(|e| {
instrument::send_error(actor_type, "remote_error");
SendError::RemoteError(e)
})?;
let (result, _): (M::Result, _) =
bincode::serde::decode_from_slice(&result_bytes, bincode::config::standard())
.map_err(|e| {
instrument::send_error(actor_type, "deserialization_failed");
SendError::DeserializationFailed(e.to_string())
})?;
#[cfg(feature = "monitor")]
instrument::remote_roundtrip_duration(actor_type, start.elapsed());
Ok(result)
}
}
}
pub async fn send_async<M>(&self, message: M) -> Result<M::Result, SendError>
where
A: AsyncHandler<M>,
M: RemoteMessage,
M::Result: Serialize + DeserializeOwned,
{
let actor_type = std::any::type_name::<A>();
match &self.inner {
EndpointInner::Local { mailbox_tx } => {
instrument::send_local(actor_type);
let (tx, rx) = oneshot::channel();
let envelope = AsyncTypedEnvelope {
message,
respond_to: tx,
};
mailbox_tx.send(Box::new(envelope)).map_err(|_| {
instrument::send_error(actor_type, "mailbox_closed");
SendError::MailboxClosed
})?;
rx.await.map_err(|_| {
instrument::send_error(actor_type, "response_dropped");
SendError::ResponseDropped
})
}
EndpointInner::Remote {
label,
wire_tx,
response_registry,
} => {
instrument::send_remote(actor_type);
#[cfg(feature = "monitor")]
let start = std::time::Instant::now();
let call_id = next_call_id();
let payload = bincode::serde::encode_to_vec(&message, bincode::config::standard())
.map_err(|e| {
instrument::send_error(actor_type, "serialization_failed");
SendError::SerializationFailed(e.to_string())
})?;
let (resp_tx, resp_rx) = oneshot::channel();
response_registry.register(call_id, resp_tx);
let invocation = RemoteInvocation {
call_id,
actor_label: label.clone(),
message_type: M::TYPE_ID.to_string(),
payload,
};
wire_tx.send(invocation).map_err(|_| {
instrument::send_error(actor_type, "wire_closed");
SendError::WireClosed
})?;
let response = resp_rx.await.map_err(|_| {
instrument::send_error(actor_type, "response_dropped");
SendError::ResponseDropped
})?;
let result_bytes = response.result.map_err(|e| {
instrument::send_error(actor_type, "remote_error");
SendError::RemoteError(e)
})?;
let (result, _): (M::Result, _) =
bincode::serde::decode_from_slice(&result_bytes, bincode::config::standard())
.map_err(|e| {
instrument::send_error(actor_type, "deserialization_failed");
SendError::DeserializationFailed(e.to_string())
})?;
#[cfg(feature = "monitor")]
instrument::remote_roundtrip_duration(actor_type, start.elapsed());
Ok(result)
}
}
}
pub(crate) fn send_with_reply_sender<M>(
&self,
message: M,
reply_sender: ReplySender<M::Result>,
) -> bool
where
A: Handler<M>,
M: Message + Send + 'static,
M::Result: Send + 'static,
{
match &self.inner {
EndpointInner::Local { mailbox_tx } => {
let envelope = ForwardedEnvelope {
message,
reply_sender,
};
mailbox_tx.send(Box::new(envelope)).is_ok()
}
EndpointInner::Remote { .. } => {
tracing::warn!("ctx.forward() to remote endpoint is not yet supported");
false
}
}
}
}