aion_server/worker/

lease_record.rs

//! The lease-record seam (WA-010 R3): the one call that turns a transport's
//! "the worker took it" into the durable [`aion_core::Event::ActivityLeased`]
//! fact, on every transport and on both dispatch paths.
//!
//! # Where the call sits
//!
//! A transport arm calls [`DeliveryAccepted::accepted`] at the instant it knows
//! the worker HOLDS the attempt — the gRPC stream accepted the frame; the
//! liminal push was acknowledged — and BEFORE it waits for any reply. That
//! placement is the whole of the ordering guarantee on the liminal outbox arm,
//! where the reply IS the completion: recorded after the wait, the lease would
//! land after the terminal it is supposed to precede.
//!
//! On the gRPC arm the completion arrives on another task, so the placement
//! alone cannot order the two appends. [`LeaseHandoff`] therefore arms the
//! completion fences' lease gate before the push and settles it after the
//! record lands (or fails, or the delivery is abandoned — the guard settles on
//! drop), and the completion entry waits on that gate before the fence accepts
//! it. See [`CompletionFences::lease_settled`].
//!
//! # What a failed record means
//!
//! An append that fails cannot write its own failure into the history. The
//! ERROR line and [`LeaseRecordLedger`] (exported as
//! `aion_activity_lease_record_failures_total` and read back on the API
//! provenance, R4) are the ONLY record that this install knew the worker and
//! lost the fact — as opposed to a history written before the event existed,
//! which never knew. A failed record never fails or repeats the dispatch: the
//! worker holds the work, and the history simply reads that attempt as
//! unattributed.
//!
//! An UNINSTALLED seam is the same loss by a different route — a boot wired a
//! dispatcher without handing it a recorder — so it is counted on the ledger
//! and said out loud, never silently skipped.

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};

use aion_core::{ActivityId, RunId, WorkerAttribution, WorkflowId};
use async_trait::async_trait;

use super::envelope::{CompletionFences, CompletionToken};
use super::registry::WorkerHandle;
use super::task_delivery::DeliveryAccepted;
use crate::error::ServerError;
use crate::observability::Metrics;

/// The attempt a lease is recorded for.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LeaseKey {
    /// Owning workflow.
    pub workflow_id: WorkflowId,
    /// The run whose history carries the lease.
    pub run_id: RunId,
    /// The activity the attempt belongs to.
    pub activity_id: ActivityId,
    /// One-based delivery attempt, as stamped on the wire.
    pub attempt: u32,
}

impl std::fmt::Display for LeaseKey {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            formatter,
            "workflow {} run {} activity {} attempt {}",
            self.workflow_id, self.run_id, self.activity_id, self.attempt
        )
    }
}

/// Something that can append the lease fact for one attempt.
///
/// Production implements this over the engine ([`EngineLeaseRecorder`]) so the
/// append goes through the workflow's one `Recorder`; tests implement it over a
/// vector or a fault.
#[async_trait]
pub trait ActivityLeaseRecorder: Send + Sync + 'static {
    /// Append `ActivityLeased` for `key`, attributed to `worker`.
    ///
    /// # Errors
    ///
    /// Returns the recorder's own reason when the append did not land. The
    /// caller counts and logs it; it never retries the dispatch.
    async fn record(&self, key: &LeaseKey, worker: WorkerAttribution) -> Result<(), String>;
}

/// The production recorder: the engine's single-writer lease API.
pub struct EngineLeaseRecorder {
    engine: Arc<aion::Engine>,
}

impl EngineLeaseRecorder {
    /// Record leases through `engine`.
    #[must_use]
    pub fn new(engine: Arc<aion::Engine>) -> Self {
        Self { engine }
    }
}

impl std::fmt::Debug for EngineLeaseRecorder {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.debug_struct("EngineLeaseRecorder").finish()
    }
}

#[async_trait]
impl ActivityLeaseRecorder for EngineLeaseRecorder {
    async fn record(&self, key: &LeaseKey, worker: WorkerAttribution) -> Result<(), String> {
        self.engine
            .record_activity_lease(
                &key.workflow_id,
                &key.run_id,
                key.activity_id.clone(),
                key.attempt,
                worker,
            )
            .await
            .map_err(|error| error.to_string())
    }
}

/// The install-lifetime count of leases this server knew and failed to record.
///
/// Shared by clone: every dispatcher on a boot increments the same count, and
/// the API provenance (R4) reads it back so a reader can tell "never recorded
/// here" from "recorded here with N losses".
#[derive(Clone, Debug, Default)]
pub struct LeaseRecordLedger {
    failures: Arc<AtomicU64>,
}

impl LeaseRecordLedger {
    /// Leases lost on this install since boot.
    #[must_use]
    pub fn failures(&self) -> u64 {
        self.failures.load(Ordering::SeqCst)
    }

    fn increment(&self) {
        self.failures.fetch_add(1, Ordering::SeqCst);
    }
}

struct Installed {
    recorder: Arc<dyn ActivityLeaseRecorder>,
    metrics: Option<Metrics>,
}

/// The shared, late-installed recorder every dispatch path records through.
///
/// Built with the worker seams before the engine exists and installed once the
/// engine has booted — the same shape as the outbox delivery callback. Cloning
/// shares the installation and the ledger.
#[derive(Clone, Default)]
pub struct LeaseRecorderSeam {
    installed: Arc<OnceLock<Installed>>,
    ledger: LeaseRecordLedger,
}

impl std::fmt::Debug for LeaseRecorderSeam {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("LeaseRecorderSeam")
            .field("installed", &self.installed.get().is_some())
            .field("failures", &self.ledger.failures())
            .finish()
    }
}

impl LeaseRecorderSeam {
    /// Install the recorder (and the metrics the failure counter exports on).
    ///
    /// Returns `false`, and keeps the first installation, when one is already
    /// in place — a second install is a wiring bug and is said so.
    pub fn install(
        &self,
        recorder: Arc<dyn ActivityLeaseRecorder>,
        metrics: Option<Metrics>,
    ) -> bool {
        if self.installed.set(Installed { recorder, metrics }).is_err() {
            tracing::warn!("activity lease recorder already installed; ignoring duplicate install");
            return false;
        }
        true
    }

    /// Whether a recorder has been installed.
    #[must_use]
    pub fn is_installed(&self) -> bool {
        self.installed.get().is_some()
    }

    /// The shared loss count.
    #[must_use]
    pub const fn ledger(&self) -> &LeaseRecordLedger {
        &self.ledger
    }

    /// Record the lease for `key`, counting and logging a loss instead of
    /// failing: the worker already holds the work.
    pub async fn record(&self, key: &LeaseKey, worker: WorkerAttribution) {
        let Some(installed) = self.installed.get() else {
            self.lost(
                key,
                &worker,
                None,
                "no activity lease recorder is installed on this dispatcher",
            );
            return;
        };
        if let Err(reason) = installed.recorder.record(key, worker.clone()).await {
            self.lost(key, &worker, installed.metrics.as_ref(), &reason);
        }
    }

    /// [`Self::record`] from a thread that cannot await — the engine-seam
    /// bridge, which dispatches synchronously.
    ///
    /// Runs the record on `handle` (or the ambient runtime). Inside a
    /// multi-thread runtime the wait is hosted with `block_in_place`; a
    /// current-thread runtime cannot host it without starving the recorder,
    /// and a thread with no runtime in reach cannot run it at all — both are
    /// counted as losses by name rather than deadlocked or skipped.
    pub fn record_blocking(
        &self,
        handle: Option<&tokio::runtime::Handle>,
        key: &LeaseKey,
        worker: WorkerAttribution,
    ) {
        match tokio::runtime::Handle::try_current() {
            Ok(ambient) => match ambient.runtime_flavor() {
                tokio::runtime::RuntimeFlavor::MultiThread => {
                    let runner = handle.cloned().unwrap_or(ambient);
                    tokio::task::block_in_place(|| runner.block_on(self.record(key, worker)));
                }
                flavor => self.lost(
                    key,
                    &worker,
                    self.installed.get().and_then(|i| i.metrics.as_ref()),
                    &format!(
                        "the lease cannot be recorded from inside a {flavor:?} tokio runtime: \
                         blocking here would starve the recorder of its only thread"
                    ),
                ),
            },
            Err(_) => match handle {
                Some(runner) => runner.block_on(self.record(key, worker)),
                None => self.lost(
                    key,
                    &worker,
                    self.installed.get().and_then(|i| i.metrics.as_ref()),
                    "no tokio runtime is reachable from this thread to record the lease on",
                ),
            },
        }
    }

    fn lost(
        &self,
        key: &LeaseKey,
        worker: &WorkerAttribution,
        metrics: Option<&Metrics>,
        why: &str,
    ) {
        self.ledger.increment();
        if let Some(metrics) = metrics {
            metrics.activity_lease_record_failed();
        }
        tracing::error!(
            workflow_id = %key.workflow_id,
            run_id = %key.run_id,
            activity_id = %key.activity_id,
            attempt = key.attempt,
            worker_identity = %worker.identity,
            task_queue = %worker.task_queue,
            transport = worker.transport.name(),
            lease_record_failures_total = self.ledger.failures(),
            reason = why,
            "activity lease was not recorded; the worker holds the attempt and the history reads \
             it as unattributed"
        );
    }
}

/// The attribution a lease carries for `worker`, read off its registration.
#[must_use]
pub fn attribution_for(worker: &WorkerHandle) -> WorkerAttribution {
    WorkerAttribution {
        identity: worker.identity().to_owned(),
        task_queue: worker.task_queue().to_owned(),
        node: worker.node().map(str::to_owned),
        deployment: worker
            .instance()
            .map(|instance| instance.deployment.clone()),
        instance_id: worker
            .instance()
            .map(|instance| instance.instance_id.clone()),
        transport: worker.delivery().transport(),
    }
}

/// One delivery's lease: armed on the fences before the push, recorded at the
/// transport's accept point, settled on the gate either way.
///
/// Dropping an unaccepted handoff (the push failed, the candidate was skipped)
/// settles the gate without recording, so a completion that can never come is
/// never waited for and a later candidate arms its own.
pub struct LeaseHandoff {
    seam: LeaseRecorderSeam,
    key: LeaseKey,
    worker: WorkerAttribution,
    fences: CompletionFences,
    token: CompletionToken,
    settled: AtomicBool,
}

impl LeaseHandoff {
    /// The attempt this handoff leases.
    #[must_use]
    pub fn key(&self) -> &LeaseKey {
        &self.key
    }

    /// The completion token the lease gate is armed under.
    #[must_use]
    pub fn token(&self) -> &CompletionToken {
        &self.token
    }
}

impl std::fmt::Debug for LeaseHandoff {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("LeaseHandoff")
            .field("key", &self.key)
            .field("settled", &self.settled.load(Ordering::SeqCst))
            .finish_non_exhaustive()
    }
}

impl LeaseHandoff {
    /// Arm the fences' lease gate for `token` and hold the lease to record.
    ///
    /// # Errors
    ///
    /// Returns lock poison from the fences.
    pub fn arm(
        seam: LeaseRecorderSeam,
        key: LeaseKey,
        worker: WorkerAttribution,
        fences: CompletionFences,
        token: CompletionToken,
    ) -> Result<Self, ServerError> {
        fences.arm_lease(&token)?;
        Ok(Self {
            seam,
            key,
            worker,
            fences,
            token,
            settled: AtomicBool::new(false),
        })
    }

    /// Record the lease from a thread that cannot await, then settle the gate.
    pub fn accepted_blocking(&self, handle: Option<&tokio::runtime::Handle>) {
        if self.settled.load(Ordering::SeqCst) {
            tracing::warn!(key = %self.key, "lease handoff accepted twice; the second is ignored");
            return;
        }
        self.seam
            .record_blocking(handle, &self.key, self.worker.clone());
        self.settle();
    }

    fn settle(&self) {
        if self.settled.swap(true, Ordering::SeqCst) {
            return;
        }
        if let Err(error) = self.fences.settle_lease(&self.token) {
            tracing::error!(
                key = %self.key,
                %error,
                "lease gate could not be settled; completions for this attempt may wait on a \
                 poisoned fence"
            );
        }
    }
}

impl Drop for LeaseHandoff {
    fn drop(&mut self) {
        self.settle();
    }
}

#[async_trait]
impl DeliveryAccepted for LeaseHandoff {
    async fn accepted(&self) {
        if self.settled.load(Ordering::SeqCst) {
            tracing::warn!(key = %self.key, "lease handoff accepted twice; the second is ignored");
            return;
        }
        self.seam.record(&self.key, self.worker.clone()).await;
        self.settle();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aion_core::WorkerTransport;
    use std::sync::Mutex;
    use uuid::Uuid;

    struct Recording(Mutex<Vec<(LeaseKey, WorkerAttribution)>>);

    #[async_trait]
    impl ActivityLeaseRecorder for Recording {
        async fn record(&self, key: &LeaseKey, worker: WorkerAttribution) -> Result<(), String> {
            self.0
                .lock()
                .map_err(|_| "recording lock poisoned".to_owned())?
                .push((key.clone(), worker));
            Ok(())
        }
    }

    struct Faulting;

    #[async_trait]
    impl ActivityLeaseRecorder for Faulting {
        async fn record(&self, _key: &LeaseKey, _worker: WorkerAttribution) -> Result<(), String> {
            Err("store fault injected".to_owned())
        }
    }

    fn key() -> LeaseKey {
        LeaseKey {
            workflow_id: WorkflowId::new(Uuid::new_v4()),
            run_id: RunId::new(Uuid::new_v4()),
            activity_id: ActivityId::from_sequence_position(3),
            attempt: 1,
        }
    }

    fn worker() -> WorkerAttribution {
        WorkerAttribution {
            identity: "w-1".to_owned(),
            task_queue: "q".to_owned(),
            node: None,
            deployment: None,
            instance_id: None,
            transport: WorkerTransport::Grpc,
        }
    }

    #[tokio::test]
    async fn an_installed_recorder_receives_the_lease_and_the_ledger_stays_at_zero() {
        let seam = LeaseRecorderSeam::default();
        let recording = Arc::new(Recording(Mutex::new(Vec::new())));
        assert!(seam.install(recording.clone(), None));
        let key = key();
        seam.record(&key, worker()).await;
        let seen = recording
            .0
            .lock()
            .map(|seen| seen.clone())
            .unwrap_or_default();
        assert_eq!(seen.len(), 1);
        assert_eq!(seen[0].0, key);
        assert_eq!(seen[0].1.identity, "w-1");
        assert_eq!(seam.ledger().failures(), 0);
    }

    #[tokio::test]
    async fn a_failed_record_counts_on_the_ledger_and_on_the_metrics_surface() {
        let seam = LeaseRecorderSeam::default();
        let metrics = Metrics::new().ok();
        assert!(seam.install(Arc::new(Faulting), metrics.clone()));
        seam.record(&key(), worker()).await;
        seam.record(&key(), worker()).await;
        assert_eq!(seam.ledger().failures(), 2);
        if let Some(metrics) = metrics {
            let text = String::from_utf8(metrics.encode().unwrap_or_default()).unwrap_or_default();
            assert!(
                text.contains("aion_activity_lease_record_failures_total 2"),
                "the metrics surface must carry the same count; got:\n{text}"
            );
        }
    }

    #[tokio::test]
    async fn an_uninstalled_seam_counts_the_loss_instead_of_skipping_it() {
        let seam = LeaseRecorderSeam::default();
        assert!(!seam.is_installed());
        seam.record(&key(), worker()).await;
        assert_eq!(seam.ledger().failures(), 1);
    }

    #[tokio::test]
    async fn a_second_install_keeps_the_first() {
        let seam = LeaseRecorderSeam::default();
        let recording = Arc::new(Recording(Mutex::new(Vec::new())));
        assert!(seam.install(recording.clone(), None));
        assert!(!seam.install(Arc::new(Faulting), None));
        seam.record(&key(), worker()).await;
        assert_eq!(seam.ledger().failures(), 0);
        assert_eq!(recording.0.lock().map_or(0, |s| s.len()), 1);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_handoff_arms_the_gate_and_settles_it_after_recording() -> Result<(), ServerError> {
        let seam = LeaseRecorderSeam::default();
        let recording = Arc::new(Recording(Mutex::new(Vec::new())));
        seam.install(recording.clone(), None);
        let fences = CompletionFences::default();
        let k = key();
        let token = fences.issue(&k.workflow_id, &k.run_id, &k.activity_id, k.attempt)?;
        let handoff = LeaseHandoff::arm(seam.clone(), k, worker(), fences.clone(), token.clone())?;
        assert!(fences.lease_pending(&token)?);
        handoff.accepted().await;
        assert!(!fences.lease_pending(&token)?);
        assert_eq!(recording.0.lock().map_or(0, |s| s.len()), 1);
        Ok(())
    }

    #[tokio::test]
    async fn dropping_an_unaccepted_handoff_settles_without_recording() -> Result<(), ServerError> {
        let seam = LeaseRecorderSeam::default();
        let recording = Arc::new(Recording(Mutex::new(Vec::new())));
        seam.install(recording.clone(), None);
        let fences = CompletionFences::default();
        let k = key();
        let token = fences.issue(&k.workflow_id, &k.run_id, &k.activity_id, k.attempt)?;
        let handoff = LeaseHandoff::arm(seam, k, worker(), fences.clone(), token.clone())?;
        assert!(fences.lease_pending(&token)?);
        drop(handoff);
        assert!(!fences.lease_pending(&token)?);
        assert_eq!(recording.0.lock().map_or(0, |s| s.len()), 0);
        Ok(())
    }

    #[tokio::test]
    async fn recording_from_a_current_thread_runtime_is_a_named_loss_not_a_deadlock() {
        let seam = LeaseRecorderSeam::default();
        seam.install(Arc::new(Recording(Mutex::new(Vec::new()))), None);
        seam.record_blocking(None, &key(), worker());
        assert_eq!(seam.ledger().failures(), 1);
    }

    #[test]
    fn recording_from_a_plain_thread_runs_on_the_given_handle() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(1)
            .enable_all()
            .build();
        let Ok(runtime) = runtime else {
            return;
        };
        let seam = LeaseRecorderSeam::default();
        let recording = Arc::new(Recording(Mutex::new(Vec::new())));
        seam.install(recording.clone(), None);
        seam.record_blocking(Some(runtime.handle()), &key(), worker());
        assert_eq!(seam.ledger().failures(), 0);
        assert_eq!(recording.0.lock().map_or(0, |s| s.len()), 1);
    }
}