use std::{
panic::{AssertUnwindSafe, catch_unwind},
sync::{
Arc,
atomic::{AtomicU8, Ordering, fence},
},
};
use crossbeam_channel::{Receiver, Sender, TryRecvError as CcTryRecvError, bounded};
use reifydb_value::reifydb_assertions;
use tracing::{debug, error};
use super::{ActorSystem, JoinError};
use crate::{
actor::{
context::{CancellationToken, Context},
mailbox::{ActorRef, create_mailbox},
traits::{Actor, Directive},
},
pool::actor_pool::{Runnable, Schedule},
sync::mutex::{Mutex, MutexGuard},
};
const IDLE: u8 = 0;
const SCHEDULED: u8 = 1;
const NOTIFIED: u8 = 2;
enum CellState<S> {
Uninit,
Running(S),
Stopped,
}
struct ActorCell<A: Actor> {
actor: A,
name: String,
state: Mutex<CellState<A::State>>,
rx: Receiver<A::Message>,
ctx: Context<A::Message>,
cancel: CancellationToken,
schedule_state: AtomicU8,
completion_tx: Sender<()>,
done_tx: Sender<()>,
schedule: Schedule,
batch_size: usize,
}
impl<A: Actor> Runnable for ActorCell<A>
where
A::State: Send,
{
fn run(self: Arc<Self>) {
let result = catch_unwind(AssertUnwindSafe(|| run_batch(Arc::clone(&self))));
if result.is_err() {
error!(actor = %self.name, "actor batch panicked; stopping actor");
*self.state.lock() = CellState::Stopped;
self.schedule_state.store(IDLE, Ordering::Release);
let _ = self.completion_tx.send(());
let _ = self.done_tx.send(());
}
}
}
fn notify<A: Actor>(cell: &Arc<ActorCell<A>>)
where
A::State: Send,
{
let prev = cell.schedule_state.fetch_max(SCHEDULED, Ordering::AcqRel);
if prev == IDLE {
cell.schedule.enqueue(Arc::clone(cell) as Arc<dyn Runnable>);
}
}
fn run_batch<A: Actor>(cell: Arc<ActorCell<A>>)
where
A::State: Send,
{
reifydb_assertions! {
let s = cell.schedule_state.load(Ordering::Acquire);
assert!(
s == SCHEDULED || s == NOTIFIED,
"actor run_batch entered while schedule_state={s} (expected SCHEDULED={SCHEDULED} or NOTIFIED={NOTIFIED}); this is a spurious wakeup that bypassed the notify guard"
);
}
let mut guard = match lock_state_or_bail(&cell) {
Some(guard) => guard,
None => return,
};
let state = match &mut *guard {
CellState::Running(state) => state,
CellState::Uninit | CellState::Stopped => {
unreachable!("lock_state_or_bail returned a non-running state")
}
};
let flow = process_message_batch(&cell, state);
dispatch_directive(&cell, guard, flow);
}
#[inline]
fn lock_state_or_bail<A: Actor>(cell: &Arc<ActorCell<A>>) -> Option<MutexGuard<'_, CellState<A::State>>>
where
A::State: Send,
{
let mut guard = cell.state.lock();
match &*guard {
CellState::Running(_) => Some(guard),
CellState::Stopped => {
cell.schedule_state.store(IDLE, Ordering::Release);
None
}
CellState::Uninit => {
debug!(actor = %cell.name, "Pool actor starting");
*guard = CellState::Running(cell.actor.init(&cell.ctx));
Some(guard)
}
}
}
#[inline]
fn process_message_batch<A: Actor>(cell: &Arc<ActorCell<A>>, state: &mut A::State) -> Directive
where
A::State: Send,
{
let mut processed = 0;
let mut flow = Directive::Continue;
while processed < cell.batch_size {
if cell.cancel.is_cancelled() {
flow = Directive::Stop;
break;
}
match cell.rx.try_recv() {
Ok(msg) => {
processed += 1;
flow = cell.actor.handle(state, msg, &cell.ctx);
match flow {
Directive::Continue => continue,
Directive::Yield | Directive::Park | Directive::Stop => break,
}
}
Err(CcTryRecvError::Empty) => {
flow = cell.actor.idle(&cell.ctx);
break;
}
Err(CcTryRecvError::Disconnected) => {
debug!("Pool actor mailbox closed, stopping");
flow = Directive::Stop;
break;
}
}
}
flow
}
#[inline]
fn dispatch_directive<A: Actor>(
cell: &Arc<ActorCell<A>>,
mut guard: MutexGuard<'_, CellState<A::State>>,
flow: Directive,
) where
A::State: Send,
{
match flow {
Directive::Stop => {
cell.actor.post_stop();
*guard = CellState::Stopped;
cell.schedule_state.store(IDLE, Ordering::Release);
let _ = cell.completion_tx.send(());
let _ = cell.done_tx.send(());
}
Directive::Park => {
cell.schedule_state.store(IDLE, Ordering::Release);
drop(guard);
fence(Ordering::SeqCst);
let has_msgs = !cell.rx.is_empty();
let cancelled = cell.cancel.is_cancelled();
if has_msgs || cancelled {
notify(cell);
}
}
Directive::Yield | Directive::Continue => {
drop(guard);
let prev = cell.schedule_state.compare_exchange(
NOTIFIED,
SCHEDULED,
Ordering::AcqRel,
Ordering::Acquire,
);
match prev {
Ok(_) => {
cell.schedule.enqueue(Arc::clone(cell) as Arc<dyn Runnable>);
}
Err(SCHEDULED) => {
if !cell.rx.is_empty() || cell.cancel.is_cancelled() {
cell.schedule.enqueue(Arc::clone(cell) as Arc<dyn Runnable>);
} else {
cell.schedule_state.store(IDLE, Ordering::Release);
fence(Ordering::SeqCst);
if !cell.rx.is_empty() || cell.cancel.is_cancelled() {
notify(cell);
}
}
}
Err(_) => {}
}
}
}
}
pub struct PoolActorHandle<M> {
pub actor_ref: ActorRef<M>,
completion_rx: Receiver<()>,
}
impl<M> PoolActorHandle<M> {
pub fn actor_ref(&self) -> &ActorRef<M> {
&self.actor_ref
}
pub fn join(self) -> Result<(), JoinError> {
self.completion_rx.recv().map_err(|_| JoinError::new("actor completion channel disconnected"))
}
}
pub(super) fn spawn_on_schedule<A: Actor>(
system: &ActorSystem,
name: &str,
actor: A,
schedule: Schedule,
batch_size: usize,
) -> PoolActorHandle<A::Message>
where
A::State: Send,
{
let config = actor.config();
let (actor_ref, mailbox) = create_mailbox(config.mailbox_capacity);
let ctx = Context::new(actor_ref.clone(), system.clone(), system.cancellation_token());
let cancel = system.cancellation_token();
let (completion_tx, completion_rx) = bounded(1);
let (done_tx, done_rx) = bounded(1);
system.register_done_rx(done_rx);
let cell = Arc::new(ActorCell {
actor,
name: name.to_string(),
state: Mutex::new(CellState::Uninit),
rx: mailbox.rx,
ctx,
cancel,
schedule_state: AtomicU8::new(SCHEDULED),
completion_tx,
done_tx,
schedule,
batch_size,
});
register_actor_hooks(&cell, &actor_ref, system);
cell.schedule.enqueue(Arc::clone(&cell) as Arc<dyn Runnable>);
PoolActorHandle {
actor_ref,
completion_rx,
}
}
#[inline]
fn register_actor_hooks<A: Actor>(cell: &Arc<ActorCell<A>>, actor_ref: &ActorRef<A::Message>, system: &ActorSystem)
where
A::State: Send,
{
let cell_weak = Arc::downgrade(cell);
let notify_fn: Arc<dyn Fn() + Send + Sync> = Arc::new(move || {
if let Some(cell) = cell_weak.upgrade() {
notify(&cell);
}
});
actor_ref.set_notify(notify_fn.clone());
system.register_waker(notify_fn);
system.register_keepalive(Box::new(Arc::clone(cell)));
}