use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};
use tokio::sync::oneshot;
use crate::actor::{Actor, ActorContext, AsyncHandler, Handler, Message};
use crate::instrument;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemoteInvocation {
pub call_id: u64,
pub actor_label: String,
pub message_type: String,
pub payload: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemoteResponse {
pub call_id: u64,
pub result: Result<Vec<u8>, String>,
}
pub struct DispatchRequest {
pub invocation: RemoteInvocation,
pub respond_to: oneshot::Sender<RemoteResponse>,
}
#[derive(Debug, thiserror::Error)]
pub enum SendError {
#[error("actor mailbox closed")]
MailboxClosed,
#[error("response channel dropped")]
ResponseDropped,
#[error("wire connection closed")]
WireClosed,
#[error("serialization failed: {0}")]
SerializationFailed(String),
#[error("deserialization failed: {0}")]
DeserializationFailed(String),
#[error("remote error: {0}")]
RemoteError(String),
}
pub struct ReplySender<R: Send + 'static> {
inner: Arc<Mutex<Option<oneshot::Sender<R>>>>,
}
impl<R: Send + 'static> Clone for ReplySender<R> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl<R: Send + 'static> ReplySender<R> {
pub fn new(sender: oneshot::Sender<R>) -> Self {
Self {
inner: Arc::new(Mutex::new(Some(sender))),
}
}
pub fn send(&self, value: R) -> bool {
if let Some(sender) = self.inner.lock().unwrap().take() {
sender.send(value).is_ok()
} else {
false
}
}
pub fn is_consumed(&self) -> bool {
self.inner.lock().unwrap().is_none()
}
}
impl<R: Send + 'static> Drop for ReplySender<R> {
fn drop(&mut self) {
if Arc::strong_count(&self.inner) == 1 && !self.is_consumed() {
tracing::warn!("ReplySender dropped without sending a reply (forgotten reply?)");
}
}
}
pub trait EnvelopeProxy<A: Actor>: Send {
fn handle<'a>(
self: Box<Self>,
ctx: &'a ActorContext<A>,
actor: &'a mut A,
state: &'a mut A::State,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
}
pub(crate) struct TypedEnvelope<M: Message> {
pub(crate) message: M,
pub(crate) respond_to: oneshot::Sender<M::Result>,
}
impl<A, M> EnvelopeProxy<A> for TypedEnvelope<M>
where
A: Handler<M>,
M: Message,
{
fn handle<'a>(
self: Box<Self>,
ctx: &'a ActorContext<A>,
actor: &'a mut A,
state: &'a mut A::State,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
Box::pin(async move {
let reply = ReplySender::new(self.respond_to);
*ctx.reply_token.lock().unwrap() = Some(Box::new(reply.clone()));
let result = actor.handle(ctx, state, self.message);
let token_was_claimed = ctx.reply_token.lock().unwrap().take().is_none();
if !token_was_claimed {
reply.send(result);
}
})
}
}
pub(crate) struct AsyncTypedEnvelope<M: Message> {
pub(crate) message: M,
pub(crate) respond_to: oneshot::Sender<M::Result>,
}
impl<A, M> EnvelopeProxy<A> for AsyncTypedEnvelope<M>
where
A: AsyncHandler<M>,
M: Message,
{
fn handle<'a>(
self: Box<Self>,
ctx: &'a ActorContext<A>,
actor: &'a mut A,
state: &'a mut A::State,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
Box::pin(async move {
let reply = ReplySender::new(self.respond_to);
*ctx.reply_token.lock().unwrap() = Some(Box::new(reply.clone()));
let result = actor.handle(ctx, state, self.message).await;
let token_was_claimed = ctx.reply_token.lock().unwrap().take().is_none();
if !token_was_claimed {
reply.send(result);
}
})
}
}
pub(crate) struct ForwardedEnvelope<M: Message> {
pub(crate) message: M,
pub(crate) reply_sender: ReplySender<M::Result>,
}
impl<A, M> EnvelopeProxy<A> for ForwardedEnvelope<M>
where
A: Handler<M>,
M: Message,
{
fn handle<'a>(
self: Box<Self>,
ctx: &'a ActorContext<A>,
actor: &'a mut A,
state: &'a mut A::State,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
Box::pin(async move {
*ctx.reply_token.lock().unwrap() = Some(Box::new(self.reply_sender.clone()));
let result = actor.handle(ctx, state, self.message);
let token_was_claimed = ctx.reply_token.lock().unwrap().take().is_none();
if !token_was_claimed {
self.reply_sender.send(result);
}
})
}
}
#[derive(Clone)]
pub struct ResponseRegistry {
inner: Arc<Mutex<HashMap<u64, oneshot::Sender<RemoteResponse>>>>,
}
impl ResponseRegistry {
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn register(&self, call_id: u64, tx: oneshot::Sender<RemoteResponse>) {
let mut map = self.inner.lock().unwrap();
map.insert(call_id, tx);
instrument::inflight_calls_set(map.len() as f64);
}
pub fn complete(&self, response: RemoteResponse) {
let mut map = self.inner.lock().unwrap();
if let Some(tx) = map.remove(&response.call_id) {
instrument::inflight_calls_set(map.len() as f64);
drop(map);
let _ = tx.send(response);
}
}
pub fn fail_all(&self, error_msg: &str) {
let mut map = self.inner.lock().unwrap();
let count = map.len() as u64;
for (call_id, tx) in map.drain() {
let _ = tx.send(RemoteResponse {
call_id,
result: Err(error_msg.to_string()),
});
}
instrument::inflight_calls_set(0.0);
instrument::dead_letters(count);
}
}
impl Default for ResponseRegistry {
fn default() -> Self {
Self::new()
}
}
pub(crate) fn next_call_id() -> u64 {
static COUNTER: AtomicU64 = AtomicU64::new(1);
COUNTER.fetch_add(1, Ordering::Relaxed)
}