reifydb_runtime/actor/system/host/

mod.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

#![allow(clippy::disallowed_types)]

mod pool;

use std::{
	any::Any,
	error, fmt,
	fmt::{Debug, Formatter},
	mem,
	sync::{Arc, OnceLock, Weak},
	time,
	time::Duration,
};

use crossbeam_channel::{Receiver, RecvTimeoutError as CcRecvTimeoutError};

use crate::{
	actor::{
		context::CancellationToken, system::host::pool::PoolActorHandle, timers::scheduler::SchedulerHandle,
		traits::Actor,
	},
	context::clock::Clock,
	pool::{
		PoolConfig, Pools,
		actor_pool::{EPHEMERAL_BATCH_SIZE, Schedule},
	},
	sync::mutex::Mutex,
};

static TESTING_ROOT: OnceLock<ActorSystem> = OnceLock::new();

struct ActorSystemInner {
	cancel: CancellationToken,
	scheduler: SchedulerHandle,
	clock: Clock,
	pools: Pools,
	wakers: Mutex<Vec<Arc<dyn Fn() + Send + Sync>>>,
	keepalive: Mutex<Vec<Box<dyn Any + Send + Sync>>>,
	done_rxs: Mutex<Vec<Receiver<()>>>,
	children: Mutex<Vec<ActorSystem>>,
}

#[derive(Clone)]
pub struct ActorSystem {
	inner: Arc<ActorSystemInner>,
}

impl ActorSystem {
	pub fn new(pools: Pools, clock: Clock) -> Self {
		let scheduler = SchedulerHandle::new();

		Self {
			inner: Arc::new(ActorSystemInner {
				cancel: CancellationToken::new(),
				scheduler,
				clock,
				pools,
				wakers: Mutex::new(Vec::new()),
				keepalive: Mutex::new(Vec::new()),
				done_rxs: Mutex::new(Vec::new()),
				children: Mutex::new(Vec::new()),
			}),
		}
	}

	pub fn testing(clock: Clock) -> Self {
		TESTING_ROOT.get_or_init(|| Self::new(Pools::new(PoolConfig::default()), Clock::Real)).scope_with(clock)
	}

	pub fn scope(&self) -> Self {
		self.scope_with(self.inner.clock.clone())
	}

	fn scope_with(&self, clock: Clock) -> Self {
		let child = Self {
			inner: Arc::new(ActorSystemInner {
				cancel: self.inner.cancel.child_token(),
				scheduler: self.inner.scheduler.shared(),
				clock,
				pools: self.inner.pools.clone(),
				wakers: Mutex::new(Vec::new()),
				keepalive: Mutex::new(Vec::new()),
				done_rxs: Mutex::new(Vec::new()),
				children: Mutex::new(Vec::new()),
			}),
		};
		self.inner.children.lock().push(child.clone());
		child
	}

	pub fn pools(&self) -> Pools {
		self.inner.pools.clone()
	}

	pub fn spawner(&self) -> ActorSpawner {
		ActorSpawner {
			inner: Arc::downgrade(&self.inner),
			clock: self.inner.clock.clone(),
		}
	}

	pub fn cancellation_token(&self) -> CancellationToken {
		self.inner.cancel.clone()
	}

	pub fn is_cancelled(&self) -> bool {
		self.inner.cancel.is_cancelled()
	}

	pub fn shutdown(&self) {
		self.inner.cancel.cancel();

		{
			let mut children = self.inner.children.lock();
			for child in children.iter() {
				child.shutdown();
			}
			children.clear();
		}

		let wakers = mem::take(&mut *self.inner.wakers.lock());
		for waker in &wakers {
			waker();
		}
		drop(wakers);

		self.inner.scheduler.shutdown();

		self.inner.keepalive.lock().clear();
	}

	pub(crate) fn register_waker(&self, f: Arc<dyn Fn() + Send + Sync>) {
		self.inner.wakers.lock().push(f);
	}

	pub(crate) fn register_keepalive(&self, cell: Box<dyn Any + Send + Sync>) {
		self.inner.keepalive.lock().push(cell);
	}

	pub(crate) fn register_done_rx(&self, rx: Receiver<()>) {
		self.inner.done_rxs.lock().push(rx);
	}

	pub fn join(&self) -> Result<(), JoinError> {
		self.join_timeout(Duration::from_secs(5))
	}

	#[allow(clippy::disallowed_methods)]
	pub fn join_timeout(&self, timeout: Duration) -> Result<(), JoinError> {
		let deadline = time::Instant::now() + timeout;
		let rxs: Vec<_> = mem::take(&mut *self.inner.done_rxs.lock());
		for rx in rxs {
			let remaining = deadline.saturating_duration_since(time::Instant::now());
			match rx.recv_timeout(remaining) {
				Ok(()) => {}
				Err(CcRecvTimeoutError::Disconnected) => {}
				Err(CcRecvTimeoutError::Timeout) => {
					return Err(JoinError::new("timed out waiting for actors to stop"));
				}
			}
		}
		Ok(())
	}

	pub fn scheduler(&self) -> &SchedulerHandle {
		&self.inner.scheduler
	}

	pub fn clock(&self) -> &Clock {
		&self.inner.clock
	}

	pub fn spawn_coordination<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		let group = self.inner.pools.actor_pool().coordination();
		pool::spawn_on_schedule(self, name, actor, Schedule::Pinned(group.assign()), group.batch_size())
	}

	pub fn spawn_flow<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		let group = self.inner.pools.actor_pool().flow();
		pool::spawn_on_schedule(self, name, actor, Schedule::Pinned(group.assign()), group.batch_size())
	}

	pub fn spawn_maintenance<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		let group = self.inner.pools.actor_pool().maintenance();
		pool::spawn_on_schedule(self, name, actor, Schedule::Pinned(group.assign()), group.batch_size())
	}

	pub fn spawn_ephemeral<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		pool::spawn_on_schedule(self, name, actor, self.inner.pools.task_injector(), EPHEMERAL_BATCH_SIZE)
	}
}

impl Debug for ActorSystem {
	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
		f.debug_struct("ActorSystem").field("cancelled", &self.is_cancelled()).finish_non_exhaustive()
	}
}

#[derive(Clone)]
pub struct ActorSpawner {
	inner: Weak<ActorSystemInner>,
	clock: Clock,
}

impl ActorSpawner {
	fn system(&self) -> ActorSystem {
		ActorSystem {
			inner: self.inner.upgrade().expect("runtime already shut down: cannot spawn actor"),
		}
	}

	pub fn clock(&self) -> &Clock {
		&self.clock
	}

	pub fn pools(&self) -> Pools {
		self.system().pools()
	}

	pub fn is_alive(&self) -> bool {
		self.inner.strong_count() > 0
	}

	pub fn cancellation_token(&self) -> Option<CancellationToken> {
		self.inner.upgrade().map(|inner| inner.cancel.clone())
	}

	pub fn scope(&self) -> ActorSpawner {
		self.system().scope().spawner()
	}

	pub fn shutdown(&self) {
		if let Some(inner) = self.inner.upgrade() {
			ActorSystem {
				inner,
			}
			.shutdown();
		}
	}

	pub fn spawn_coordination<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		self.system().spawn_coordination(name, actor)
	}

	pub fn spawn_flow<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		self.system().spawn_flow(name, actor)
	}

	pub fn spawn_maintenance<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		self.system().spawn_maintenance(name, actor)
	}

	pub fn spawn_ephemeral<A: Actor>(&self, name: &str, actor: A) -> ActorHandle<A::Message>
	where
		A::State: Send,
	{
		self.system().spawn_ephemeral(name, actor)
	}
}

impl Debug for ActorSpawner {
	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
		f.debug_struct("ActorSpawner").field("alive", &self.is_alive()).finish_non_exhaustive()
	}
}

pub type ActorHandle<M> = PoolActorHandle<M>;

#[derive(Debug)]
pub struct JoinError {
	message: String,
}

impl JoinError {
	pub fn new(message: impl Into<String>) -> Self {
		Self {
			message: message.into(),
		}
	}
}

impl fmt::Display for JoinError {
	fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
		write!(f, "actor join failed: {}", self.message)
	}
}

impl error::Error for JoinError {}

#[cfg(test)]
mod tests {
	use std::sync;

	use super::*;
	use crate::{
		actor::{context::Context, traits::Directive},
		pool::{PoolConfig, Pools},
	};

	fn test_system() -> ActorSystem {
		let pools = Pools::new(PoolConfig::default());
		ActorSystem::new(pools, Clock::Real)
	}

	struct CounterActor;

	#[derive(Debug)]
	enum CounterMessage {
		Inc,
		Get(sync::mpsc::Sender<i64>),
		Stop,
	}

	impl Actor for CounterActor {
		type State = i64;
		type Message = CounterMessage;

		fn init(&self, _ctx: &Context<Self::Message>) -> Self::State {
			0
		}

		fn handle(
			&self,
			state: &mut Self::State,
			msg: Self::Message,
			_ctx: &Context<Self::Message>,
		) -> Directive {
			match msg {
				CounterMessage::Inc => *state += 1,
				CounterMessage::Get(tx) => {
					let _ = tx.send(*state);
				}
				CounterMessage::Stop => return Directive::Stop,
			}
			Directive::Continue
		}
	}

	#[test]
	fn test_spawn_and_send() {
		let system = test_system();
		let handle = system.spawn_coordination("counter", CounterActor);

		let actor_ref = handle.actor_ref().clone();
		actor_ref.send(CounterMessage::Inc).unwrap();
		actor_ref.send(CounterMessage::Inc).unwrap();
		actor_ref.send(CounterMessage::Inc).unwrap();

		let (tx, rx) = sync::mpsc::channel();
		actor_ref.send(CounterMessage::Get(tx)).unwrap();

		let value = rx.recv().unwrap();
		assert_eq!(value, 3);

		actor_ref.send(CounterMessage::Stop).unwrap();
		handle.join().unwrap();
	}

	#[test]
	fn test_shutdown_join() {
		let system = test_system();

		for i in 0..5 {
			system.spawn_coordination(&format!("counter-{i}"), CounterActor);
		}

		// join() must not return before every actor has finished, or shutdown races teardown.
		system.shutdown();
		system.join().unwrap();
	}

	#[test]
	fn test_shutdown_stops_the_timer_scheduler() {
		// Timers must be dead before the cells are released, otherwise a callback owns the last system ref.
		let system = test_system();
		system.shutdown();

		let (tx, rx) = sync::mpsc::channel();
		system.scheduler().schedule_once(Duration::from_millis(5), move || {
			let _ = tx.send(());
		});

		assert!(rx.recv_timeout(Duration::from_millis(200)).is_err());
	}
}