08cefb1
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use futures::future::BoxFuture;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct Mailbox<M> {
tx: mpsc::UnboundedSender<M>,
}
impl<M> Mailbox<M> {
pub fn send(&self, msg: M) -> bool {
self.tx.send(msg).is_ok()
}
pub fn is_closed(&self) -> bool {
self.tx.is_closed()
}
}
type SpawnFn<M> = dyn Fn(Arc<ActorRegistry<M>>, Uuid, mpsc::UnboundedReceiver<M>) -> BoxFuture<'static, ()>
+ Send
+ Sync;
pub struct ActorRegistry<M> {
name: &'static str,
actors: Mutex<HashMap<Uuid, mpsc::UnboundedSender<M>>>,
spawn: Arc<SpawnFn<M>>,
shutdown: CancellationToken,
}
impl<M> std::fmt::Debug for ActorRegistry<M> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let live = self.actors.lock().unwrap().len();
f.debug_struct("ActorRegistry")
.field("name", &self.name)
.field("live_actors", &live)
.finish()
}
}
impl<M> ActorRegistry<M>
where
M: Send + 'static,
{
pub fn new<F, Fut>(name: &'static str, spawn: F, shutdown: CancellationToken) -> Arc<Self>
where
F: Fn(Arc<Self>, Uuid, mpsc::UnboundedReceiver<M>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let spawn: Arc<SpawnFn<M>> = Arc::new(move |reg, key, rx| Box::pin(spawn(reg, key, rx)));
Arc::new(Self {
name,
actors: Mutex::new(HashMap::new()),
spawn,
shutdown,
})
}
pub fn shutdown_token(&self) -> CancellationToken {
self.shutdown.clone()
}
pub fn live_count(&self) -> usize {
self.actors.lock().unwrap().len()
}
pub fn dispatch(self: &Arc<Self>, key: Uuid, mut msg: M) {
{
let guard = self.actors.lock().unwrap();
if let Some(tx) = guard.get(&key) {
match tx.send(msg) {
Ok(()) => return,
Err(e) => msg = e.0,
}
}
}
let (tx, rx) = mpsc::unbounded_channel();
let mailbox = Mailbox { tx: tx.clone() };
{
let mut guard = self.actors.lock().unwrap();
if let Some(existing) = guard.get(&key) {
match existing.send(msg) {
Ok(()) => return,
Err(e) => msg = e.0,
}
}
guard.insert(key, tx);
}
let registry = self.clone();
let spawn = self.spawn.clone();
tokio::spawn(async move {
let fut = (spawn)(registry.clone(), key, rx);
fut.await;
registry.actors.lock().unwrap().remove(&key);
});
let _ = mailbox.send(msg);
crate::metrics::actor_spawned(self.name);
}
}
pub async fn run_actor<M>(
mut rx: mpsc::UnboundedReceiver<M>,
shutdown: CancellationToken,
mut apply: impl FnMut(M) -> BoxFuture<'static, ()>,
) {
loop {
tokio::select! {
_ = shutdown.cancelled() => break,
maybe = rx.recv() => {
match maybe {
Some(msg) => apply(msg).await,
None => break,
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn dispatch_spawns_and_orders() {
let shutdown = CancellationToken::new();
let seen = Arc::new(Mutex::new(Vec::new()));
let reg = ActorRegistry::new(
"test",
{
let seen = seen.clone();
move |reg: Arc<ActorRegistry<u32>>, _key, rx| {
let seen = seen.clone();
async move {
run_actor(rx, reg.shutdown_token(), move |msg| {
let seen = seen.clone();
Box::pin(async move {
seen.lock().unwrap().push(msg);
})
})
.await;
}
}
},
shutdown,
);
let key = Uuid::new_v4();
reg.dispatch(key, 1);
reg.dispatch(key, 2);
reg.dispatch(key, 3);
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let got = seen.lock().unwrap().clone();
assert_eq!(got, vec![1, 2, 3]);
}
}