use std::error::Error;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::time::Instant;
use beamr::native::native_process::{NativeContext, NativeHandler, NativeOutcome};
use beamr::process::ExitReason;
use super::tests::fixtures::ControlledConsumer;
use super::{DispatchConversation, DispatchError};
use crate::conversation::ParticipantPid;
use crate::conversation::handoff::{ConsumptionAcceptance, ReceiptEvent, ReceiptMetrics};
use crate::routing::group::{ConsumerGroup, ConsumerRegistration};
use crate::routing::{ConsumerId, ConsumerStateView, RoutingMessage};
#[derive(Debug)]
pub(super) struct AttemptProbe {
pub(super) waiter_resumes: AtomicU64,
pub(super) attempts: AtomicU64,
pub(super) reroutes: AtomicU64,
pub(super) cell: Option<(usize, usize)>,
pub(super) receipts: Arc<ReceiptMetrics>,
armed: mpsc::Sender<ParticipantPid>,
only: Option<ParticipantPid>,
pub(super) accepted_signal: Mutex<Option<mpsc::Sender<ConsumptionAcceptance>>>,
pub(super) observe_permission: Mutex<Option<mpsc::Receiver<()>>>,
pub(super) captured_exit: Mutex<Option<mpsc::Receiver<Instant>>>,
}
impl AttemptProbe {
pub(super) fn new() -> (Arc<Self>, mpsc::Receiver<ParticipantPid>) {
Self::filtered(None)
}
fn filtered(only: Option<ParticipantPid>) -> (Arc<Self>, mpsc::Receiver<ParticipantPid>) {
let (armed, receiver) = mpsc::channel();
(
Arc::new(Self {
waiter_resumes: AtomicU64::new(0),
attempts: AtomicU64::new(0),
reroutes: AtomicU64::new(0),
cell: None,
receipts: Arc::new(ReceiptMetrics::default()),
armed,
only,
accepted_signal: Mutex::new(None),
observe_permission: Mutex::new(None),
captured_exit: Mutex::new(None),
}),
receiver,
)
}
pub(super) fn arm_waiter(&self, pid: ParticipantPid) -> Result<(), DispatchError> {
if self.only.is_some_and(|selected| selected != pid) {
return Ok(());
}
self.armed.send(pid).map_err(|error| {
DispatchError::Conversation(format!(
"dispatch waiter-armed signal disconnected: {error}"
))
})?;
let permission = self
.observe_permission
.lock()
.map_err(|error| {
DispatchError::Conversation(format!("observe permission lock poisoned: {error}"))
})?
.take();
if let Some(permission) = permission {
permission.recv().map_err(|error| {
DispatchError::Conversation(format!(
"observe permission signal disconnected: {error}"
))
})?;
}
Ok(())
}
pub(super) fn record_event(&self, event: ReceiptEvent) {
let ReceiptEvent::Acceptance(acceptance) = event else {
return;
};
let signal = match self.accepted_signal.lock() {
Ok(signal) => signal.clone(),
Err(error) => {
tracing::error!(%error, "controlled acceptance signal lock poisoned");
return;
}
};
if let Some(signal) = signal {
if let Err(error) = signal.send(acceptance) {
tracing::error!(%error, "controlled acceptance signal disconnected");
}
}
}
fn snapshot(&self, taken: u64, completed: u64) -> CostVector {
CostVector {
attempts: self.attempts.load(Ordering::Acquire),
taken,
acceptances: self.receipts.acceptances.load(Ordering::Acquire),
exits: self.receipts.exits.load(Ordering::Acquire),
terminals: self.receipts.terminals.load(Ordering::Acquire),
reroutes: self.reroutes.load(Ordering::Acquire),
resumes: self.waiter_resumes.load(Ordering::Acquire),
timer_arms: self.receipts.rearms.load(Ordering::Acquire),
timer_wakeups: self.receipts.idle_wakeups.load(Ordering::Acquire),
completed,
registrations: self.receipts.registrations.load(Ordering::Acquire),
releases: self.receipts.releases.load(Ordering::Acquire),
}
}
}
pub(super) fn register(dispatch: &DispatchConversation, consumer: &ControlledConsumer, name: &str) {
let _ = dispatch.group().add_consumer(ConsumerRegistration::new(
consumer.pid,
ConsumerStateView::new(ConsumerId::new(name), 0, 4, 0, Vec::new()),
));
}
#[test]
fn none_reply_consumption_records_acceptance_for_the_selected_attempt() -> Result<(), Box<dyn Error>>
{
let mut dispatch = DispatchConversation::new(ConsumerGroup::new(super::select_first()))?;
let consumer = ControlledConsumer::new(dispatch.supervisor())?;
register(&dispatch, &consumer, "selected");
let (probe, armed) = AttemptProbe::new();
dispatch.probe = Some(Arc::clone(&probe));
let worker_dispatch = dispatch.clone();
let worker = std::thread::spawn(move || worker_dispatch.dispatch(&RoutingMessage::new()));
assert_eq!(armed.recv()?, consumer.pid);
assert_eq!(consumer.permit.messages_taken(), 0);
assert_eq!(probe.receipts.acceptances.load(Ordering::Acquire), 0);
consumer.permit.release()?;
let result = worker.join().map_err(|_| "dispatch worker panicked")??;
assert_eq!(result.delivered_to(), &ConsumerId::new("selected"));
assert_eq!(consumer.permit.messages_taken(), 1);
assert_eq!(probe.receipts.acceptances.load(Ordering::Acquire), 1);
consumer.owner.finalize();
dispatch.supervisor().shutdown();
Ok(())
}
struct CompletionWork {
completed: Arc<AtomicU64>,
signal: mpsc::Sender<()>,
}
impl NativeHandler for CompletionWork {
fn handle(&mut self, _: &mut NativeContext<'_>) -> NativeOutcome {
self.completed.fetch_add(1, Ordering::Relaxed);
if let Err(error) = self.signal.send(()) {
tracing::error!(%error, "unrelated work completion signal disconnected");
}
NativeOutcome::Stop(ExitReason::Normal)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct CostVector {
attempts: u64,
taken: u64,
acceptances: u64,
exits: u64,
terminals: u64,
reroutes: u64,
resumes: u64,
timer_arms: u64,
timer_wakeups: u64,
completed: u64,
registrations: u64,
releases: u64,
}
struct IdleVector {
pending: usize,
unrelated: usize,
before: CostVector,
after: CostVector,
cleaned: CostVector,
outcome_pending: bool,
}
fn idle_cell(pending: usize, unrelated: usize) -> Result<IdleVector, Box<dyn Error>> {
let mut dispatch = DispatchConversation::new(ConsumerGroup::new(super::select_first()))?;
let consumer = ControlledConsumer::new(dispatch.supervisor())?;
register(&dispatch, &consumer, "held");
let (mut probe, armed) = AttemptProbe::new();
Arc::get_mut(&mut probe)
.ok_or("idle probe unexpectedly shared")?
.cell = Some((pending, unrelated));
dispatch.probe = Some(Arc::clone(&probe));
let mut workers = Vec::new();
let (settled, outcomes) = mpsc::channel();
for _ in 0..pending {
let worker_dispatch = dispatch.clone();
let settled = settled.clone();
workers.push(std::thread::spawn(move || {
let result = worker_dispatch.dispatch(&RoutingMessage::new());
settled.send(()).map_err(|error| error.to_string())?;
Ok::<_, String>(result)
}));
}
for _ in 0..pending {
assert_eq!(armed.recv()?, consumer.pid);
}
let completed = Arc::new(AtomicU64::new(0));
let before = probe.snapshot(
consumer.permit.messages_taken(),
completed.load(Ordering::Acquire),
);
let (done, completed_work) = mpsc::channel();
for _ in 0..unrelated {
let completed = Arc::clone(&completed);
let done = done.clone();
dispatch
.supervisor()
.scheduler()
.spawn_native(Box::new(move || {
Box::new(CompletionWork {
completed: Arc::clone(&completed),
signal: done.clone(),
})
}))
.map_err(|error| format!("unrelated work spawn failed: {error}"))?;
}
for _ in 0..unrelated {
completed_work.recv()?;
}
let after = probe.snapshot(
consumer.permit.messages_taken(),
completed.load(Ordering::Acquire),
);
let pending_outcome = outcomes.try_recv();
dispatch
.supervisor()
.scheduler()
.terminate_process(consumer.pid.get(), ExitReason::Error);
for worker in workers {
let result = worker.join().map_err(|_| "idle matrix worker panicked")??;
assert_eq!(result, Err(DispatchError::NoConsumerAvailable));
}
let cleaned = probe.snapshot(
consumer.permit.messages_taken(),
completed.load(Ordering::Acquire),
);
dispatch.supervisor().shutdown();
Ok(IdleVector {
pending,
unrelated,
before,
after,
cleaned,
outcome_pending: matches!(pending_outcome, Err(mpsc::TryRecvError::Empty)),
})
}
#[test]
fn idle_cost_matrix_has_no_timer_arms_wakeups_or_waiter_resumes() -> Result<(), Box<dyn Error>> {
let mut vectors = Vec::new();
for pending in [1, 128] {
for unrelated in [0, 256] {
vectors.push(idle_cell(pending, unrelated)?);
}
}
for vector in &vectors {
eprintln!(
"idle pending={} unrelated={} before={:?} after={:?} cleaned={:?} pending_outcome={}",
vector.pending,
vector.unrelated,
vector.before,
vector.after,
vector.cleaned,
vector.outcome_pending
);
}
for vector in vectors {
let pending = u64::try_from(vector.pending)?;
assert_eq!(
vector.before,
CostVector {
attempts: pending,
taken: 0,
acceptances: 0,
exits: 0,
terminals: 0,
reroutes: 0,
resumes: 0,
timer_arms: 0,
timer_wakeups: 0,
completed: 0,
registrations: pending,
releases: 0,
}
);
let expected_after = CostVector {
completed: u64::try_from(vector.unrelated)?,
..vector.before
};
assert_eq!(
vector.after, expected_after,
"unrelated work must not wake pending dispatches"
);
assert!(vector.outcome_pending);
assert_eq!(
vector.cleaned,
CostVector {
exits: pending,
terminals: pending,
reroutes: pending,
resumes: pending,
releases: pending,
..expected_after
}
);
}
Ok(())
}
pub(super) struct WaiterReady {
ready: mpsc::Receiver<ParticipantPid>,
selected: ParticipantPid,
}
impl WaiterReady {
pub(super) fn wait(&self) -> Result<(), String> {
let selected = self
.ready
.recv()
.map_err(|error| format!("dispatch waiter-armed signal disconnected: {error}"))?;
if selected != self.selected {
return Err("dispatch waiter armed for a different consumer".to_owned());
}
Ok(())
}
}
pub(super) fn arm_waiter_signal(
dispatch: &mut DispatchConversation,
selected: ParticipantPid,
) -> WaiterReady {
let (probe, ready) = AttemptProbe::filtered(Some(selected));
dispatch.probe = Some(probe);
WaiterReady { ready, selected }
}