use std::sync::Mutex;
use futures_util::{StreamExt, stream::FuturesUnordered};
use tokio::{sync::mpsc, task::JoinHandle};
use super::{
actor::ScheduledActor,
reaper::{AttemptReaper, ReapFuture, ReaperCommand},
};
pub(in crate::core::registry) struct ActorRuntime {
attempt_reaper: AttemptReaper,
reaper_rx: Mutex<Option<mpsc::UnboundedReceiver<ReaperCommand>>>,
reaper_handle: Mutex<Option<JoinHandle<()>>>,
}
impl ActorRuntime {
pub(in crate::core::registry) fn new() -> Self {
let (reaper_tx, reaper_rx) = mpsc::unbounded_channel();
Self {
attempt_reaper: AttemptReaper::new(reaper_tx),
reaper_rx: Mutex::new(Some(reaper_rx)),
reaper_handle: Mutex::new(None),
}
}
pub(in crate::core::registry) fn attempt_reaper(&self) -> AttemptReaper {
self.attempt_reaper.clone()
}
pub(in crate::core::registry) fn reaping_attempts(&self) -> usize {
self.attempt_reaper.active()
}
pub(in crate::core::registry) fn spawn(&self) {
let mut reaper_rx = self
.reaper_rx
.lock()
.unwrap_or_else(|error| error.into_inner())
.take()
.expect("attempt reaper starts exactly once");
let handle = tokio::spawn(async move {
let mut active = FuturesUnordered::<ReapFuture>::new();
let mut closing = false;
loop {
if closing && active.is_empty() {
break;
}
tokio::select! {
command = reaper_rx.recv(), if !closing => match command {
Some(ReaperCommand::Reap(future)) => active.push(future),
Some(ReaperCommand::Close) | None => {
closing = true;
reaper_rx.close();
while let Ok(command) = reaper_rx.try_recv() {
if let ReaperCommand::Reap(future) = command {
active.push(future);
}
}
}
},
completed = active.next(), if !active.is_empty() => {
debug_assert!(completed.is_some());
}
}
}
});
*self
.reaper_handle
.lock()
.unwrap_or_else(|error| error.into_inner()) = Some(handle);
}
pub(in crate::core::registry) fn schedule(&self, actor: ScheduledActor) {
actor.spawn();
}
pub(in crate::core::registry) fn schedule_batch(
&self,
actors: impl IntoIterator<Item = ScheduledActor>,
) {
for actor in actors {
actor.spawn();
}
}
pub(in crate::core::registry) async fn join(&self) -> bool {
self.attempt_reaper.close();
if self.attempt_reaper.active() != 0 {
return true;
}
let handle = self
.reaper_handle
.lock()
.unwrap_or_else(|error| error.into_inner())
.take();
match handle {
Some(handle) => handle.await.is_ok(),
None => true,
}
}
}