use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aion_core::WorkerTransport;
use serde::{Deserialize, Serialize};
use tokio::sync::watch;
use tracing::{info, warn};
use super::grpc_liveness::{GrpcLivenessTarget, GrpcLivenessWaiters, ping_grpc_worker};
use super::heartbeat::{
DISPATCH_PROBATION_PINGS, DispatchExclusion, ExcludedWorker, HeartbeatTracker, sweep_interval,
};
use super::liminal_transport::{LiminalConnectionNotifier, LiminalWorkerDelivery};
use super::liveness::{PingFailure, ProbedTransport};
use super::registry::{ConnectedWorkerRegistry, WorkerId};
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LivenessPing {
pub liveness_ping: u64,
pub silence_window_ms: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LivenessPong {
pub liveness_pong: u64,
}
#[derive(Clone, Debug)]
pub struct LivenessTarget {
pub pid: u64,
pub worker_id: WorkerId,
pub delivery: LiminalWorkerDelivery,
}
pub struct LivenessProbe {
notifier: Option<Arc<LiminalConnectionNotifier>>,
grpc: Option<GrpcLivenessWaiters>,
tracker: HeartbeatTracker,
registry: ConnectedWorkerRegistry,
cadence: Duration,
silence_window: Duration,
}
impl std::fmt::Debug for LivenessProbe {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("LivenessProbe")
.field("cadence", &self.cadence)
.field("silence_window", &self.silence_window)
.finish_non_exhaustive()
}
}
impl LivenessProbe {
#[must_use]
pub fn new(
notifier: Arc<LiminalConnectionNotifier>,
tracker: HeartbeatTracker,
registry: ConnectedWorkerRegistry,
heartbeat_window: Duration,
) -> Self {
Self::across_transports(Some(notifier), None, tracker, registry, heartbeat_window)
}
#[must_use]
pub fn across_transports(
notifier: Option<Arc<LiminalConnectionNotifier>>,
grpc: Option<GrpcLivenessWaiters>,
tracker: HeartbeatTracker,
registry: ConnectedWorkerRegistry,
heartbeat_window: Duration,
) -> Self {
Self {
notifier,
grpc,
tracker,
registry,
cadence: sweep_interval(heartbeat_window),
silence_window: heartbeat_window,
}
}
#[must_use]
pub const fn cadence(&self) -> Duration {
self.cadence
}
#[must_use]
pub const fn silence_window(&self) -> Duration {
self.silence_window
}
pub async fn run(self, mut shutdown: watch::Receiver<bool>) {
info!(
cadence_ms = self.cadence.as_millis(),
silence_window_ms = self.silence_window.as_millis(),
"worker liveness probe started (liminal + grpc)"
);
let mut sequence = 0_u64;
let mut ticks = tokio::time::interval(self.cadence);
ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = ticks.tick() => {
if *shutdown.borrow() {
break;
}
sequence = sequence.saturating_add(1);
self.probe_once(sequence).await;
}
changed = shutdown.changed() => {
if changed.is_err() || *shutdown.borrow() {
break;
}
}
}
}
info!("worker liveness probe stopped");
}
pub async fn probe_once(&self, sequence: u64) {
let liminal = self
.notifier
.as_ref()
.map_or_else(Vec::new, |notifier| notifier.liveness_targets());
let grpc = self.grpc_targets();
let silence_window_ms = u64::try_from(self.silence_window.as_millis()).unwrap_or(u64::MAX);
tokio::join!(
self.probe_liminal(sequence, silence_window_ms, liminal),
self.probe_grpc(sequence, silence_window_ms, grpc),
);
self.publish_reachability_verdict(sequence);
}
fn grpc_targets(&self) -> Vec<GrpcLivenessTarget> {
let Some(_) = self.grpc.as_ref() else {
return Vec::new();
};
match self.registry.grpc_liveness_targets() {
Ok(targets) => targets,
Err(error) => {
warn!(
%error,
"could not enumerate gRPC-delivered workers for this liveness round; they \
go unprobed and will therefore not clear their dispatch probation"
);
Vec::new()
}
}
}
async fn probe_liminal(
&self,
sequence: u64,
silence_window_ms: u64,
targets: Vec<LivenessTarget>,
) {
if targets.is_empty() {
return;
}
let ping = LivenessPing {
liveness_ping: sequence,
silence_window_ms,
};
let payload = match serde_json::to_vec(&ping) {
Ok(payload) => payload,
Err(error) => {
warn!(%error, "liminal liveness probe could not encode its ping; skipping round");
return;
}
};
let deadline = self.cadence;
let answers = targets.into_iter().map(|target| {
let payload = payload.clone();
async move {
let outcome = tokio::task::spawn_blocking(move || {
ping_one(&target.delivery, payload, deadline)
})
.await;
(target.pid, target.worker_id, outcome)
}
});
for (pid, worker_id, outcome) in futures::future::join_all(answers).await {
self.apply_answer(pid, worker_id, sequence, outcome);
}
}
async fn probe_grpc(
&self,
sequence: u64,
silence_window_ms: u64,
targets: Vec<GrpcLivenessTarget>,
) {
let (Some(waiters), false) = (self.grpc.as_ref(), targets.is_empty()) else {
return;
};
let deadline = self.cadence;
let answers = targets.iter().map(|target| async move {
let outcome = ping_grpc_worker(
waiters,
target,
aion_proto::ProtoLivenessPing {
liveness_ping: sequence,
silence_window_ms,
},
deadline,
)
.await;
(target.worker_id, outcome)
});
for (worker_id, outcome) in futures::future::join_all(answers).await {
self.apply_grpc_answer(worker_id, sequence, outcome);
}
}
fn publish_reachability_verdict(&self, sequence: u64) {
let now = Instant::now();
let measured = match self.tracker.unreachable_workers(now) {
Ok(workers) => workers,
Err(error) => {
warn!(
%error,
liveness_ping = sequence,
"could not read worker reachability; leaving the previous dispatch \
eligibility verdict in place rather than guessing"
);
return;
}
};
let Some(unreachable) = self.within_reach(&measured) else {
warn!(
liveness_ping = sequence,
"could not read which transport each unreachable worker is delivered over, so \
this round cannot tell an exclusion it is entitled to publish from one it has \
no wire to test; the previous verdict stands unchanged"
);
return;
};
let excluded_now: BTreeMap<WorkerId, DispatchExclusion> = unreachable
.iter()
.map(|excluded| (excluded.worker_id, excluded.exclusion))
.collect();
match self.registry.dispatch_ineligible() {
Ok(previously_excluded) => {
self.announce_transitions(
sequence,
&previously_excluded,
&unreachable,
&excluded_now,
);
}
Err(error) => warn!(
%error,
liveness_ping = sequence,
"could not read the published liminal dispatch eligibility set; this round's \
eligibility changes go UNANNOUNCED, though the verdict itself is still published"
),
}
if let Err(error) = self.registry.set_dispatch_ineligible(excluded_now) {
warn!(
%error,
liveness_ping = sequence,
"could not publish worker dispatch eligibility; selection keeps the \
previous verdict"
);
}
}
fn within_reach(&self, measured: &[ExcludedWorker]) -> Option<Vec<ExcludedWorker>> {
let transports = self
.registry
.transports_of(measured.iter().map(|excluded| excluded.worker_id))
.map_err(|error| warn!(%error, "could not read worker transports"))
.ok()?;
Some(
measured
.iter()
.filter(|excluded| {
transports
.get(&excluded.worker_id)
.is_some_and(|transport| self.covers(*transport))
})
.copied()
.collect(),
)
}
const fn covers(&self, transport: WorkerTransport) -> bool {
match transport {
WorkerTransport::Grpc => self.grpc.is_some(),
WorkerTransport::Liminal => self.notifier.is_some(),
}
}
fn announce_transitions(
&self,
sequence: u64,
previously_excluded: &BTreeMap<WorkerId, DispatchExclusion>,
unreachable: &[ExcludedWorker],
excluded_now: &BTreeMap<WorkerId, DispatchExclusion>,
) {
for excluded in unreachable {
let Some(announcement) = transition(
previously_excluded.contains_key(&excluded.worker_id),
Some(excluded.exclusion),
) else {
continue;
};
self.say(sequence, excluded.worker_id, announcement);
}
for worker_id in previously_excluded
.keys()
.filter(|worker_id| !excluded_now.contains_key(worker_id))
{
let Some(task_queue) = self.task_queue_of(*worker_id) else {
continue;
};
let Some(announcement) = transition(true, None) else {
continue;
};
self.say_with_queue(sequence, *worker_id, &task_queue, announcement);
}
}
fn say(&self, sequence: u64, worker_id: WorkerId, announcement: Announcement) {
let task_queue = self
.task_queue_of(worker_id)
.unwrap_or_else(|| "<unregistered>".to_owned());
self.say_with_queue(sequence, worker_id, &task_queue, announcement);
}
fn say_with_queue(
&self,
sequence: u64,
worker_id: WorkerId,
task_queue: &str,
announcement: Announcement,
) {
let transport = self.transport_of(worker_id);
match announcement {
Announcement::ProbationOpened { answers } => info!(
worker_id = worker_id.value(),
task_queue,
transport,
liveness_ping = sequence,
answers_banked = answers,
answers_required = DISPATCH_PROBATION_PINGS,
"worker is SERVING ITS DISPATCH PROBATION: it has answered {answers} of \
{DISPATCH_PROBATION_PINGS} consecutive liveness pings since it connected, and is \
not selected for dispatch until the run is complete. This is the ordinary cost of \
connecting — every healthy worker start passes through it, on every transport — \
not a fault, and not a statement that anything is unreachable"
),
Announcement::EligibilityWithdrawn => warn!(
worker_id = worker_id.value(),
task_queue,
transport,
liveness_ping = sequence,
answers_required = DISPATCH_PROBATION_PINGS,
silence_window_ms = self.silence_window.as_millis(),
"WITHDRAWING DISPATCH ELIGIBILITY from worker: it had PROVED its dispatch path \
reachable on this connection and the server can no longer prove it — either a \
liveness ping failed or the last proof aged out of the window. It stays \
registered and keeps its in-flight work, and becomes eligible again after \
{DISPATCH_PROBATION_PINGS} consecutive answered pings"
),
Announcement::EligibilityRestored => info!(
worker_id = worker_id.value(),
task_queue,
transport,
liveness_ping = sequence,
"DISPATCH ELIGIBILITY RESTORED to worker: it has answered a full run of \
consecutive liveness pings, so the server can again prove it reaches this \
worker's dispatch path. Dispatch selection includes it from now"
),
}
}
fn transport_of(&self, worker_id: WorkerId) -> &'static str {
match self.registry.worker_by_id(worker_id) {
Ok(Some(handle)) => handle.delivery().transport().name(),
Ok(None) | Err(_) => "<unregistered>",
}
}
fn apply_answer(
&self,
pid: u64,
worker_id: WorkerId,
sequence: u64,
outcome: Result<Result<LivenessPong, PingFailure>, tokio::task::JoinError>,
) {
let failure = match outcome {
Err(join_error) => {
PingFailure::Unanswered(format!("ping task failed to run: {join_error}"))
}
Ok(Err(failure)) => failure,
Ok(Ok(pong)) if pong.liveness_pong != sequence => PingFailure::Unanswered(format!(
"worker answered with mismatched sequence {}",
pong.liveness_pong
)),
Ok(Ok(_)) => {
self.record_reachable(ProbedTransport::Liminal { pid }, worker_id);
return;
}
};
self.record_unreachable(
ProbedTransport::Liminal { pid },
worker_id,
sequence,
failure,
);
}
fn apply_grpc_answer(
&self,
worker_id: WorkerId,
sequence: u64,
outcome: Result<u64, PingFailure>,
) {
match outcome {
Ok(echoed) if echoed == sequence => {
self.record_reachable(ProbedTransport::Grpc, worker_id);
}
Ok(echoed) => self.record_unreachable(
ProbedTransport::Grpc,
worker_id,
sequence,
PingFailure::Unanswered(format!(
"worker answered with mismatched sequence {echoed}"
)),
),
Err(failure) => {
self.record_unreachable(ProbedTransport::Grpc, worker_id, sequence, failure);
}
}
}
fn record_reachable(&self, transport: ProbedTransport, worker_id: WorkerId) {
if let Err(error) = self
.tracker
.record_dispatch_reachability(worker_id, Instant::now())
{
warn!(
%error,
connection_pid = transport.pid(),
transport = transport.name(),
worker_id = worker_id.value(),
"failed to record worker dispatch reachability from a liveness answer"
);
}
}
fn record_unreachable(
&self,
transport: ProbedTransport,
worker_id: WorkerId,
sequence: u64,
failure: PingFailure,
) {
if let Err(error) = self.tracker.record_dispatch_unreachable(worker_id) {
warn!(
%error,
connection_pid = transport.pid(),
transport = transport.name(),
worker_id = worker_id.value(),
"failed to reset worker dispatch probation after a failed liveness ping; \
its eligibility may outlive the proof that earned it"
);
}
let task_queue = self
.task_queue_of(worker_id)
.unwrap_or_else(|| "<unregistered>".to_owned());
let transport_name = transport.name();
match failure {
PingFailure::Unaskable(reason) => warn!(
connection_pid = transport.pid(),
transport = transport_name,
worker_id = worker_id.value(),
task_queue = %task_queue,
liveness_ping = sequence,
reason = %reason,
silence_window_ms = self.silence_window.as_millis(),
"THE SERVER COULD NOT ASK this {transport_name} worker for liveness — the push \
itself was refused, so nothing was sent and the worker has no idea it was probed. \
A dispatch would be refused by the same channel for the same reason. This says \
nothing about whether the worker is healthy; it says this server cannot currently \
reach it. Its dispatch eligibility is withdrawn NOW and it must answer a full run \
of consecutive pings to earn it back"
),
PingFailure::Unanswered(reason) => warn!(
connection_pid = transport.pid(),
transport = transport_name,
worker_id = worker_id.value(),
task_queue = %task_queue,
liveness_ping = sequence,
reason = %reason,
silence_window_ms = self.silence_window.as_millis(),
"{transport_name} worker did not answer its liveness ping; the server could not \
prove it can reach this worker's dispatch path, so its dispatch eligibility is \
withdrawn NOW and it must answer a full run of consecutive pings to earn it back \
— an unanswered probe is direct evidence about the push leg, not mere silence. \
NOTE: the worker's connection lease may still be fresh — its liveness pump beats \
from a background task and keeps proving the process is alive — so do NOT expect \
an expiry sweep to reap it"
),
}
}
fn task_queue_of(&self, worker_id: WorkerId) -> Option<String> {
self.registry
.worker_by_id(worker_id)
.ok()
.flatten()
.map(|handle| handle.task_queue().to_owned())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Announcement {
ProbationOpened {
answers: u32,
},
EligibilityWithdrawn,
EligibilityRestored,
}
const fn transition(
was_excluded: bool,
now_excluded: Option<DispatchExclusion>,
) -> Option<Announcement> {
match (was_excluded, now_excluded) {
(false, Some(DispatchExclusion::OpeningProbation { answers })) => {
Some(Announcement::ProbationOpened { answers })
}
(false, Some(DispatchExclusion::ReachabilityLost)) => {
Some(Announcement::EligibilityWithdrawn)
}
(true, None) => Some(Announcement::EligibilityRestored),
(true, Some(_)) | (false, None) => None,
}
}
fn ping_one(
delivery: &LiminalWorkerDelivery,
payload: Vec<u8>,
deadline: Duration,
) -> Result<LivenessPong, PingFailure> {
let awaiter = delivery
.push_payload_with_deadline(payload, deadline)
.map_err(|error| PingFailure::Unaskable(error.to_string()))?;
let reply = awaiter
.receive(deadline)
.map_err(|error| PingFailure::Unanswered(format!("no answer arrived: {error}")))?;
serde_json::from_slice(&reply)
.map_err(|error| PingFailure::Unanswered(format!("answer could not be decoded: {error}")))
}
#[cfg(test)]
mod tests {
use crate::namespace::NamespaceGuard;
use std::sync::Arc;
use std::time::{Duration, Instant};
use super::super::heartbeat::{DISPATCH_PROBATION_PINGS, HeartbeatTracker};
use super::super::liminal_transport::LiminalConnectionNotifier;
use super::super::registry::{ConnectedWorkerRegistry, WorkerId};
use super::{
Announcement, DispatchExclusion, LivenessPing, LivenessPong, LivenessProbe, PingFailure,
WorkerTransport, transition,
};
fn test_admission_runtime() -> std::io::Result<tokio::runtime::Runtime> {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
}
#[test]
fn every_standing_change_has_exactly_one_announcement() {
assert_eq!(
transition(
false,
Some(DispatchExclusion::OpeningProbation { answers: 1 })
),
Some(Announcement::ProbationOpened { answers: 1 }),
"a fresh connection serving its probation is ORDINARY and must not be announced as a \
worker the server cannot reach"
);
assert_eq!(
transition(false, Some(DispatchExclusion::ReachabilityLost)),
Some(Announcement::EligibilityWithdrawn),
"losing eligibility that was actually held is the incident the WARN exists for"
);
assert_eq!(
transition(true, None),
Some(Announcement::EligibilityRestored),
"the all-clear must be said out loud — an alarm nobody cancels is read as ongoing"
);
assert_eq!(
transition(true, Some(DispatchExclusion::ReachabilityLost)),
None,
"a standing exclusion must not re-log every cadence"
);
assert_eq!(
transition(
true,
Some(DispatchExclusion::OpeningProbation { answers: 0 })
),
None,
"including while a still-excluded worker is still serving its probation"
);
assert_eq!(
transition(false, None),
None,
"an eligible worker that stayed eligible is not news"
);
}
#[test]
fn a_probe_covers_exactly_the_transports_it_holds_a_wire_for() -> std::io::Result<()> {
let admission_runtime = test_admission_runtime()?;
let window = Duration::from_secs(30);
let registry = ConnectedWorkerRegistry::default();
let tracker = HeartbeatTracker::new(window);
let notifier = || {
Arc::new(
LiminalConnectionNotifier::new(registry.clone()).with_admission(
NamespaceGuard::shared_engine(),
false,
admission_runtime.handle().clone(),
),
)
};
let build = |liminal: bool, grpc: bool| {
LivenessProbe::across_transports(
liminal.then(notifier),
grpc.then(super::GrpcLivenessWaiters::new),
tracker.clone(),
registry.clone(),
window,
)
};
let both = build(true, true);
assert!(both.covers(WorkerTransport::Liminal));
assert!(both.covers(WorkerTransport::Grpc));
let liminal_only = build(true, false);
assert!(liminal_only.covers(WorkerTransport::Liminal));
assert!(
!liminal_only.covers(WorkerTransport::Grpc),
"a probe holding no gRPC answer registry cannot hear a gRPC answer, so it must not \
judge a gRPC worker: the exclusion it would publish could never be lifted"
);
let grpc_only = build(false, true);
assert!(
!grpc_only.covers(WorkerTransport::Liminal),
"and the same in the other direction — the reach follows the wire, not the fleet"
);
assert!(grpc_only.covers(WorkerTransport::Grpc));
let neither = build(false, false);
assert!(!neither.covers(WorkerTransport::Liminal));
assert!(!neither.covers(WorkerTransport::Grpc));
Ok(())
}
#[test]
fn the_liveness_pair_round_trips_through_json() -> Result<(), serde_json::Error> {
let ping = LivenessPing {
liveness_ping: 7,
silence_window_ms: 30_000,
};
let encoded = serde_json::to_string(&ping)?;
assert_eq!(encoded, r#"{"liveness_ping":7,"silence_window_ms":30000}"#);
assert_eq!(serde_json::from_str::<LivenessPing>(&encoded)?, ping);
let answer = LivenessPong { liveness_pong: 7 };
let encoded = serde_json::to_string(&answer)?;
assert_eq!(encoded, r#"{"liveness_pong":7}"#);
assert_eq!(serde_json::from_str::<LivenessPong>(&encoded)?, answer);
Ok(())
}
#[test]
fn a_liveness_ping_is_disjoint_from_the_other_pushed_frames() -> Result<(), serde_json::Error> {
let ping = serde_json::to_vec(&LivenessPing {
liveness_ping: 1,
silence_window_ms: 1_000,
})?;
assert!(
serde_json::from_slice::<super::super::liminal_transport::DispatchRequest>(&ping)
.is_err(),
"a liveness ping must never decode as a dispatch"
);
assert!(
serde_json::from_slice::<super::super::liminal_transport::InterventionRequest>(&ping)
.is_err(),
"a liveness ping must never decode as an intervention"
);
Ok(())
}
#[test]
fn probe_timings_derive_from_the_heartbeat_window() {
let window = Duration::from_secs(30);
assert_eq!(super::sweep_interval(window), Duration::from_millis(7_500));
assert!(super::sweep_interval(window) * 4 <= window);
}
#[test]
fn every_failed_probe_class_withdraws_dispatch_eligibility() -> std::io::Result<()> {
let admission_runtime = test_admission_runtime()?;
let window = Duration::from_secs(30);
let worker = WorkerId::from_value(1);
let start = Instant::now();
let failures = [
(
"the push was refused, so nothing was even asked",
Ok(Err(PingFailure::Unaskable("push refused".to_owned()))),
),
(
"the ping was sent and no answer came back",
Ok(Err(PingFailure::Unanswered("no answer arrived".to_owned()))),
),
(
"an answer came back carrying the wrong sequence",
Ok(Ok(LivenessPong { liveness_pong: 99 })),
),
];
for (class, outcome) in failures {
let registry = ConnectedWorkerRegistry::default();
let tracker = HeartbeatTracker::new(window);
let probe = LivenessProbe::new(
Arc::new(
LiminalConnectionNotifier::new(registry.clone()).with_admission(
NamespaceGuard::shared_engine(),
false,
admission_runtime.handle().clone(),
),
),
tracker.clone(),
registry,
window,
);
assert!(
tracker.register_connection(worker, start).is_ok(),
"tracker registration must succeed for {class}"
);
for _ in 0..DISPATCH_PROBATION_PINGS {
assert!(
tracker
.record_dispatch_reachability(worker, start)
.is_ok_and(|tracked| tracked),
"the worker serves its probation before {class}"
);
}
assert!(
tracker
.is_dispatch_reachable(worker, start)
.is_ok_and(|reachable| reachable),
"precondition: the worker is eligible before {class}"
);
probe.apply_answer(7, worker, 1, outcome);
assert!(
tracker
.is_dispatch_reachable(worker, start)
.is_ok_and(|reachable| !reachable),
"the probe must withdraw dispatch eligibility when {class} — otherwise the \
probation never resets and a one-way link keeps its eligibility forever"
);
}
Ok(())
}
#[test]
fn an_answered_probe_keeps_dispatch_eligibility_through_the_same_seam() -> std::io::Result<()> {
let admission_runtime = test_admission_runtime()?;
let window = Duration::from_secs(30);
let worker = WorkerId::from_value(1);
let start = Instant::now();
let registry = ConnectedWorkerRegistry::default();
let tracker = HeartbeatTracker::new(window);
let probe = LivenessProbe::new(
Arc::new(
LiminalConnectionNotifier::new(registry.clone()).with_admission(
NamespaceGuard::shared_engine(),
false,
admission_runtime.handle().clone(),
),
),
tracker.clone(),
registry,
window,
);
assert!(tracker.register_connection(worker, start).is_ok());
for sequence in 1..=u64::from(DISPATCH_PROBATION_PINGS) {
probe.apply_answer(
7,
worker,
sequence,
Ok(Ok(LivenessPong {
liveness_pong: sequence,
})),
);
}
assert!(
tracker
.is_dispatch_reachable(worker, start)
.is_ok_and(|reachable| reachable),
"answered probes must EARN eligibility through the probe seam, not merely fail to \
withdraw it"
);
Ok(())
}
}