dbos 0.6.0

//! The ambient workflow context.
//!
//! A workflow function is an ordinary `async fn` taking its own arguments and nothing else. What
//! makes it durable — which workflow it is, and which step comes next — travels in a
//! [`tokio::task_local!`], set for the duration of the body and unset outside it. Threading a
//! context parameter instead would put it in every workflow and step signature, and through every
//! helper between them.

use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicI32, AtomicU64, AtomicUsize, Ordering};

use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;

use crate::instance::Executor;
use crate::serialization::PayloadEncoding;
use crate::sysdb::types::Timestamp;

tokio::task_local! {
    /// Set while a workflow body runs, and read by everything the body calls.
    static CURRENT: Ctx;
}

/// What a running workflow knows about itself.
///
/// Cheap to clone: a few `Arc`s. Cloning it does not make a second workflow — the clone shares the
/// same step counter, which is the point, because a step allocated through either must not reuse
/// an id allocated through the other.
///
/// **Internal.** Everything a workflow may ask about itself is a free function — [`workflow_id`],
/// [`step_id`], [`step_status`] and [`cancellation_token`] — so user code never names this type or
/// reaches through it. That keeps the ambient context an implementation detail: the questions are
/// stable API, the thing that answers them is not.
#[derive(Clone)]
pub(crate) struct Ctx {
    executor: Arc<Executor>,
    workflow: Arc<WorkflowState>,
    /// Which step body this context is inside, or `None` in the workflow proper.
    ///
    /// A step is a leaf: the checkpoint it writes stands for everything the body did, so a step
    /// inside a step is a plain call. Without this the inner call would allocate a step id of its
    /// own and every step after it would replay against the wrong slot — a correctness trap rather
    /// than a policy question.
    ///
    /// **Here rather than on [`WorkflowState`], because the question is about one call stack and
    /// not about the workflow.** [`in_step_scope`](Self::in_step_scope) binds it on the context it
    /// rebinds the task-local with, so it is in scope only while that body is being polled: two
    /// steps in flight cannot see each other's, and one that finishes cannot answer for another
    /// still inside its own. A *count* of live steps would belong on the shared state instead,
    /// since refusing concurrency outright is a question about the workflow.
    step: Option<StepScope>,
}

/// What a context inside a step body knows about that body.
///
/// **One `Option` for both, because a context has both or neither.** The marker and the status are
/// bound together by [`in_step_scope`](Ctx::in_step_scope) and go out of scope together, and that
/// invariant is worth more as a type than as a convention two fields keep by agreement: nothing can
/// leave a context holding a status it is no longer inside, or a marker with nothing to report.
///
/// They stay *distinct* inside it, because they answer different questions on different clocks —
/// see [`StepMarker`] for the one that is per attempt, and [`StepStatus`] for what the body may
/// read about itself.
#[derive(Clone, Debug)]
pub(crate) struct StepScope {
    /// Which body, for the leaf rule. Opaque, and never compared across runs.
    marker: StepMarker,
    /// What [`step_status`] reports.
    ///
    /// **The id inside it is per step, where the marker is per attempt.** A retry enters a new
    /// scope with a fresh marker and the *same* id: the attempts are different bodies, and they are
    /// attempts at one step, competing to record under one row. `current_attempt` is the field that
    /// moves between them.
    status: StepStatus,
    /// Fires when the body running under this scope should stop.
    ///
    /// **Per attempt, like the marker beside it.** A retried step gets a fresh token, and shared
    /// state is exactly what must not carry one — a token cancelled by attempt one would arrive
    /// already-cancelled at attempt two.
    ///
    /// **Not an `Option`, because a step body always has one.** The engine mints a token for every
    /// attempt whether or not anything watches it, so "in a step" and "has a token" are the same
    /// condition, and one `Option` on the `Ctx` says it once.
    ///
    /// **This is the receiving end, and only that.** The engine raises a cancellation on the token
    /// it holds itself; what lands here is a clone, handed to the body so it can watch — see
    /// [`cancellation_token`]. Nothing a body does with this cancels anything.
    cancellation: CancellationToken,
}

/// What a step body can learn about the attempt it is running as.
///
/// Read with [`step_status`].
///
/// The token a body watches for an abandoned attempt is not a field here but
/// [`cancellation_token`] — a free function, because it is useful to a body that has no interest
/// in which attempt it is.
///
/// **A read-only snapshot, and read-only by construction.** The fields are behind accessors rather
/// than public, so there is no way to write one. That is not a guard against a body changing its
/// own retry policy — it could not anyway, since this is `Copy` and the engine rebuilds its own
/// copy from its own counters on every attempt — but against the *appearance* of one: a settable
/// `max_attempts` that silently changed nothing would be worse than no field at all. A step decides
/// how many attempts it gets through [`StepOptions`](crate::StepOptions), before it runs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StepStatus {
    /// The step's ordinal position in its workflow, counting from zero — the same number
    /// [`step_id`] reports, and half a checkpoint row's key.
    pub(crate) step_id: i32,
    /// Which attempt is running, counting from **one**.
    ///
    /// **One-based**, so a plain step is "attempt 1 of 1" rather than "attempt 0 of 1". A step's
    /// tracing span spells the same number one-based, and a body that logs "attempt 2 of 3"
    /// should not have to add one to say so.
    pub(crate) current_attempt: u32,
    /// How many attempts the policy allows in total, so a body can tell it is on its last.
    ///
    /// **Always a number, even for a step that does not retry.** Every step runs the same loop
    /// with `max_attempts` defaulting to 1, so a plain step is honestly attempt 1 of 1 —
    /// and "does this step retry?" is `max_attempts() > 1` rather than a second `Option` to unwrap
    /// inside one.
    ///
    /// **A ceiling, not a promise.** The loop can also end early: a failure that
    /// [`should_retry`](crate::StepOptions::should_retry) declines is the step's last whatever the
    /// count says, so `current_attempt() == max_attempts()` means "no attempts left *by count*",
    /// not "certainly the last". A body cannot know it is on its last attempt before it fails,
    /// because the predicate is only asked once there is a failure to show it.
    pub(crate) max_attempts: u32,
}

impl StepStatus {
    /// The step's ordinal position in its workflow, counting from zero.
    pub fn step_id(&self) -> i32 {
        self.step_id
    }

    /// Which attempt is running, counting from one.
    pub fn current_attempt(&self) -> u32 {
        self.current_attempt
    }

    /// How many attempts the policy allows in total — a ceiling the
    /// [`should_retry`](crate::StepOptions::should_retry) predicate may stop short of.
    pub fn max_attempts(&self) -> u32 {
        self.max_attempts
    }
}

#[cfg(test)]
impl StepStatus {
    /// A first attempt at `step_id`, for tests that only care which step they are inside.
    pub(crate) fn first(step_id: i32) -> Self {
        Self {
            step_id,
            current_attempt: 1,
            max_attempts: 1,
        }
    }
}

/// Which step body a context is inside.
///
/// **A newtype rather than a bare integer**, because the only other integer in reach is a step id
/// and the two mean nothing alike: a step id is an ordinal position, restarts from zero on every
/// replay, and is half the primary key of a checkpoint row, where a marker is opaque, never
/// persisted, and never compared across runs. Distinct types are what keep a later edit from
/// passing one where the other belongs.
///
/// **One per attempt, not one per step.** [`Ctx::in_step_scope`] is entered again for every retry,
/// and each entry is a body of its own — work the first attempt handed out is not inside the
/// second.
///
/// Named `marker` rather than `body` because in [`step`](crate::step()) a step's *body* is already
/// its closure, and the field would shadow it wherever both are in scope.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct StepMarker(u64);

/// Hands out a fresh [`StepMarker`] for each step body.
///
/// Process-wide rather than per-workflow, because a marker is only ever compared for equality and
/// a value distinct across the process is distinct within any one workflow. `Relaxed` is enough:
/// the counter orders nothing, it only has to stop handing out the same value twice.
static NEXT_STEP_MARKER: AtomicU64 = AtomicU64::new(0);

/// The parts of a workflow that outlive any one call within it.
struct WorkflowState {
    workflow_id: String,
    /// The token this execution owns the workflow by, which every write it makes presents.
    ///
    /// A direct start's creator token, or the token of the dequeue claim that dispatched it. When
    /// the workflow is handed to another execution (a cancel, a resume, a recovery), the row's
    /// `owner_xid` stops matching this, and this execution's next write is refused.
    owner_xid: String,
    /// When this workflow must stop, if it was given a budget.
    ///
    /// **The instant the database holds, not the timeout its caller offered.** A recovered
    /// workflow reads the deadline stored on its row, so a workflow given an hour that crashed
    /// after fifty minutes has ten left — and a child that inherits this inherits what is left
    /// rather than a fresh hour.
    ///
    /// Here rather than passed down to the one place that watches it, because a **child workflow**
    /// has to read it at launch: an inherited deadline is how a parent's budget reaches work the
    /// parent is not itself running, and the child's `select!` fires on the same instant with no
    /// signal passing between them.
    deadline: Option<Timestamp>,
    /// How this workflow's input, output and error are encoded, as its row's tag says.
    ///
    /// Here because a send, an event or a stream write made inside the workflow takes the same
    /// encoding unless it asks for another — so a portable workflow's writes are readable by
    /// whoever can read the workflow.
    encoding: PayloadEncoding,
    /// The next step id to hand out, so the first step in a workflow is step 0.
    ///
    /// **Zero-based, and the base is part of the contract.** The id is public: it is stored in
    /// `operation_outputs.function_id`, shown in step listings, and passed to `fork` as the step
    /// to resume at, so other SDKs sharing the database rely on the same base.
    next_step_id: AtomicI32,
    /// The child starts this workflow has in flight, which it waits for before recording its
    /// outcome.
    ///
    /// **A start writes under this execution's owner token, and recording the outcome gives the
    /// token up.** A start the body dropped part-way goes on in a task of its own, and if the
    /// outcome landed first that task's write would be refused, with nobody left awaiting it to
    /// hear so: the child would never exist, and nothing would record why. Waiting here is what
    /// makes a dropped start keep its promise to start the child.
    starts: Arc<Starts>,
}

impl WorkflowState {
    fn next_step_id(&self) -> i32 {
        self.next_step_id.fetch_add(1, Ordering::Relaxed)
    }

    fn peek_step_id(&self) -> i32 {
        self.next_step_id.load(Ordering::Relaxed)
    }

    fn claim_step_id(&self, step_id: i32) -> bool {
        self.next_step_id
            .compare_exchange(step_id, step_id + 1, Ordering::Relaxed, Ordering::Relaxed)
            .is_ok()
    }
}

/// A count of the child starts a workflow has in flight, and a way to wait until there are none.
///
/// A start holds a [`StartToken`] for as long as its task exists, and the execution waits on
/// [`drained`](Self::drained) once its body has ended. Built from an atomic and a
/// [`tokio::sync::Notify`] alone, which need no particular runtime to run on.
#[derive(Debug, Default)]
pub(crate) struct Starts {
    in_flight: AtomicUsize,
    none_left: Notify,
}

impl Starts {
    /// Counts one more start in flight, until the token is dropped.
    pub(crate) fn token(self: &Arc<Self>) -> StartToken {
        self.in_flight.fetch_add(1, Ordering::AcqRel);
        StartToken(Arc::clone(self))
    }

    /// Returns once no start is in flight: at once if none is.
    ///
    /// Registers for the wake-up before reading the count, so a last token dropped between the
    /// read and the await still wakes it.
    pub(crate) async fn drained(&self) {
        loop {
            let mut notified = std::pin::pin!(self.none_left.notified());
            notified.as_mut().enable();
            if self.in_flight.load(Ordering::Acquire) == 0 {
                return;
            }
            notified.await;
        }
    }
}

/// One start in flight, counted in its workflow's [`Starts`] until this is dropped — however the
/// task holding it ends.
#[derive(Debug)]
pub(crate) struct StartToken(Arc<Starts>);

impl Drop for StartToken {
    fn drop(&mut self) {
        if self.0.in_flight.fetch_sub(1, Ordering::AcqRel) == 1 {
            self.0.none_left.notify_waiters();
        }
    }
}

impl Ctx {
    /// A context for a workflow about to run.
    pub(crate) fn new(
        executor: Arc<Executor>,
        workflow_id: impl Into<String>,
        owner_xid: impl Into<String>,
        deadline: Option<Timestamp>,
        encoding: PayloadEncoding,
    ) -> Self {
        Self {
            executor,
            workflow: Arc::new(WorkflowState {
                workflow_id: workflow_id.into(),
                owner_xid: owner_xid.into(),
                deadline,
                encoding,
                next_step_id: AtomicI32::new(0),
                starts: Arc::default(),
            }),
            step: None,
        }
    }

    /// The context of the workflow this code is running inside, or `None` outside one.
    ///
    /// The sole accessor. Crate-internal: user code asks its questions through the free functions
    /// that wrap this, which is what keeps the context itself off the public surface.
    ///
    /// `None` is not an error. A step called outside a workflow runs plainly and undurably, which
    /// is what makes a function built from steps ordinarily testable.
    pub(crate) fn current() -> Option<Ctx> {
        CURRENT.try_with(Ctx::clone).ok()
    }

    /// [`current`](Self::current) without the clone: runs `f` on the ambient context, or on
    /// `None`.
    ///
    /// **For callers that only want to look.** `current` hands back an owned `Ctx`, which costs
    /// three atomic increments and three decrements — the executor, the workflow state, and the
    /// attempt's cancellation token are each behind an `Arc`. Those refcounts are shared by every
    /// concurrent step of the workflow, so they are exactly the words under contention when steps
    /// run together, and paying for them to answer a question that borrows is waste. A step's
    /// `poll` asks where it stands on every poll, which is what makes that waste worth a second
    /// accessor.
    pub(crate) fn with_current<R>(f: impl FnOnce(Option<&Ctx>) -> R) -> R {
        // The `Option` is what lets one `FnOnce` serve both arms: `try_with` runs the closure
        // exactly when it returns `Ok`, so precisely one of these two takes finds a value.
        let mut f = Some(f);
        match CURRENT.try_with(|ctx| f.take().expect("the closure runs once")(Some(ctx))) {
            Ok(answer) => answer,
            Err(_) => f.take().expect("the closure runs once")(None),
        }
    }

    /// The id of the workflow this context belongs to.
    pub(crate) fn workflow_id(&self) -> &str {
        &self.workflow.workflow_id
    }

    /// The token this execution owns its workflow by, which every write it makes presents.
    pub(crate) fn owner_xid(&self) -> &str {
        &self.workflow.owner_xid
    }

    /// This execution, as the writes it makes present it.
    pub(crate) fn owner(&self) -> crate::sysdb::types::WorkflowOwner<'_> {
        crate::sysdb::types::WorkflowOwner {
            workflow_id: self.workflow_id(),
            owner_xid: self.owner_xid(),
        }
    }

    /// The id of the step body this context is inside, or `None` in the workflow proper.
    pub(crate) fn step_id(&self) -> Option<i32> {
        self.step.as_ref().map(|step| step.status.step_id)
    }

    /// What the body this context is inside may read about its own attempt, or `None` outside one.
    pub(crate) fn step_status(&self) -> Option<StepStatus> {
        self.step.as_ref().map(|step| step.status)
    }

    /// Whether this and `other` are the same *execution* of the same workflow.
    ///
    /// **Identity, not equality of ids.** A workflow id names a row; this asks whether the two
    /// contexts share one [`WorkflowState`], and therefore one step counter. Two things a matching
    /// id would wave through do not share one: a second `DBOS` in the process serving the same id,
    /// which is what [`Error::WrongInstance`](crate::Error::WrongInstance) exists to catch, and a
    /// second *execution* of one id — a recovery re-run — whose counter restarts from zero. A
    /// step id means nothing across either, so anything holding one has to ask this rather than
    /// compare strings.
    ///
    /// Cheap: one pointer comparison. Clones of a `Ctx` share the state, so the ordinary case —
    /// a step built and polled inside one workflow body — answers true without touching memory
    /// the caller did not already have.
    pub(crate) fn is_same_execution(&self, other: &Ctx) -> bool {
        Arc::ptr_eq(&self.workflow, &other.workflow)
    }

    /// The child starts this workflow has in flight. A start holds a token from it for as long as
    /// its task exists, and the execution waits for them all before recording its outcome.
    pub(crate) fn starts(&self) -> &Arc<Starts> {
        &self.workflow.starts
    }

    /// When this workflow must stop, if it has a deadline at all.
    pub(crate) fn deadline(&self) -> Option<Timestamp> {
        self.workflow.deadline
    }

    /// How this workflow's own payloads are encoded, which its writes inherit.
    pub(crate) fn encoding(&self) -> PayloadEncoding {
        self.workflow.encoding.clone()
    }

    /// Runs `future` with `ctx` ambient.
    ///
    /// The context is restored on every poll and removed on every yield, so it is present for the
    /// whole body including across `.await`. It does **not** cross a spawn, on any runtime: a
    /// spawned task is a new task with its own empty task-local map. That is the correct default —
    /// the spawned work is not part of the durable workflow, and silently adopting the parent's
    /// step counter would let two tasks allocate the same step id — but it is surprising enough
    /// that `the_context_does_not_cross_a_spawn` pins it as behaviour rather than leaving it to be
    /// discovered.
    pub(crate) async fn scope<F: Future>(ctx: Ctx, future: F) -> F::Output {
        CURRENT.scope(ctx, future).await
    }

    /// The next step id in this workflow, zero-based and never reused.
    pub(crate) fn next_step_id(&self) -> i32 {
        self.workflow.next_step_id()
    }

    /// The id the next call will take, without taking it.
    ///
    /// For a call that has to look at a position before deciding whether to occupy it — a
    /// [`patch`](crate::patch()) — and that pairs this with [`claim_step_id`](Self::claim_step_id)
    /// to learn whether anything else took the position while it looked.
    pub(crate) fn peek_step_id(&self) -> i32 {
        self.workflow.peek_step_id()
    }

    /// Takes `step_id` if it is still the next id, and reports whether it was.
    ///
    /// One compare-and-swap, so the check and the take are a single step: `false` means some other
    /// call of this execution took an id since `step_id` was [peeked](Self::peek_step_id), and the
    /// counter is left as that call left it.
    pub(crate) fn claim_step_id(&self, step_id: i32) -> bool {
        self.workflow.claim_step_id(step_id)
    }

    /// The executor running this workflow.
    pub(crate) fn executor(&self) -> &Arc<Executor> {
        &self.executor
    }

    /// Whether this call is inside a step body.
    ///
    /// **Per call stack, not per workflow**: it reads the [`StepMarker`] that
    /// [`in_step_scope`](Self::in_step_scope) binds on the context it rebinds for the body, so a
    /// sibling step running concurrently has no bearing on the answer, and neither has one that
    /// has just finished.
    pub(crate) fn in_step(&self) -> bool {
        self.step.is_some()
    }

    /// Which step body this context is inside, if any.
    ///
    /// [`in_step`](Self::in_step) asks whether there is one; this asks *which*, and the difference
    /// is what lets a durable call built inside a step body be refused when it is polled somewhere
    /// else. Both places have the same workflow id, so comparing workflow identity cannot tell them
    /// apart. See [`StepMarker`].
    pub(crate) fn step_marker(&self) -> Option<StepMarker> {
        self.step.as_ref().map(|step| step.marker)
    }

    /// Runs `body` under a context that is [`in_step`](Self::in_step).
    ///
    /// **Nothing to unset afterwards.** The marker lives on the `Ctx` bound for this body alone, so
    /// it goes out of scope with the body however the body ends — an early return, an error, a
    /// panic — and no other call stack ever saw it. A flag on the shared state would need a guard
    /// to clear it, and a guard could not make it correct: restoring rather than clearing still
    /// hands one step's answer to another.
    pub(crate) async fn in_step_scope<F: Future>(
        &self,
        cancellation: CancellationToken,
        status: StepStatus,
        body: F,
    ) -> F::Output {
        // Rebinding rather than mutating: the body must see this attempt's token, its own marker
        // and its own id, and the `Ctx` the workflow body holds must acquire none of them.
        let scoped = Ctx {
            executor: Arc::clone(&self.executor),
            workflow: Arc::clone(&self.workflow),
            step: Some(StepScope {
                marker: StepMarker(NEXT_STEP_MARKER.fetch_add(1, Ordering::Relaxed)),
                status,
                cancellation,
            }),
        };
        CURRENT.scope(scoped, body).await
    }

    /// Fires when the running step should stop: the attempt's token, or one that never fires when
    /// there is no step running.
    ///
    /// What [`cancellation_token`] hands out; its documentation says
    /// when the token fires and who should watch it.
    pub(crate) fn cancellation(&self) -> CancellationToken {
        // `unwrap_or_default` means exactly one thing: there is no step here, and a token that
        // never fires is the honest answer.
        self.step
            .as_ref()
            .map(|step| step.cancellation.clone())
            .unwrap_or_default()
    }
}

/// The id of the workflow this code is running inside, or `None` outside one.
///
/// A workflow that wants to log its own id, publish it, or hand it to something that will later
/// address the workflow by it — the id is an idempotency key, so this is how a workflow tells a
/// caller what to send to — has no parameter to read it from: a DBOS workflow is an ordinary
/// `async fn` taking its own arguments and nothing else.
///
/// `None` is not an error, and code that runs both inside and outside a workflow is the reason it
/// is an `Option` rather than a panic: a helper called from a workflow and from a plain handler
/// gets an answer in both places.
///
/// Answers from inside a step as well as from the workflow body — a step is part of its workflow,
/// and asking which workflow it belongs to is not the same question as
/// [`step_id`].
///
/// ```no_run
/// # async fn f() -> dbos::Result<()> {
/// // The payment id a caller pays against *is* this workflow's id.
/// let id = dbos::workflow_id().expect("inside a workflow");
/// dbos::set_event("payment_id", &id).await?;
/// # Ok(())
/// # }
/// ```
pub fn workflow_id() -> Option<String> {
    Ctx::with_current(|ctx| ctx.map(|ctx| ctx.workflow_id().to_owned()))
}

/// The id of the step this code is running inside, or `None` when it is not inside one.
///
/// The ordinal position of the step within its workflow, counting from zero — the number that
/// addresses a checkpoint row, labels a step in a listing, and names a fork point. It restarts
/// from zero on every replay, which is what makes it an address rather than a serial number: the
/// same step in the same workflow has the same id on every execution.
///
/// **`None` in the workflow body itself**, not just outside a workflow. A step id belongs to a
/// step, and between two steps a workflow is inside neither. So this is `Some` only while a step
/// body is executing, which includes the retries of a step — every attempt of one step reports
/// that step's id.
///
/// **A step called from inside another step reports the enclosing step's id**, because a nested
/// call is not a step of its own: it takes no id, writes no checkpoint, and is replayed as part of
/// the body that called it. Reporting the durable step the work belongs to is the only answer with
/// a checkpoint row behind it.
///
/// ```no_run
/// async fn charge() -> dbos::Result<()> {
///     // Inside a step body, so this is the id of the step being run.
///     tracing::info!(step = dbos::step_id(), "charging");
///     Ok(())
/// }
/// ```
pub fn step_id() -> Option<i32> {
    Ctx::with_current(|ctx| ctx.and_then(Ctx::step_id))
}

/// What the step this code is running inside knows about its own attempt, or `None` outside one.
///
/// [`step_id`] is the common case and stays its own function; this is the rest of what a body may
/// ask — chiefly **which attempt it is**, so a step can behave differently on its last one: log the
/// failure loudly, fall back to a cheaper path, or stop paying for a cache it is about to give up
/// on.
///
/// `None` in the workflow body proper and outside a workflow, exactly as
/// [`step_id`] is — a status belongs to a step, and between two steps a workflow is
/// inside neither.
///
/// **Inside a step called from another step, this describes the enclosing step** — its id, its
/// attempt, its cap — for the reason [`step_id`] gives: the nested call is part of that step and
/// has no attempt of its own. Worth knowing before a shared helper branches on it, since "the last
/// attempt" it reads is its caller's, and can be the last attempt of a step the helper itself has
/// never failed.
///
/// ```no_run
/// # #[derive(Debug, thiserror::Error, serde::Serialize, serde::Deserialize)]
/// # #[error("upstream is down")]
/// # struct Upstream;
/// async fn charge() -> dbos::Result<(), Upstream> {
///     // Matched rather than unwrapped: called outside a workflow this body still runs, plainly,
///     // and there is no attempt to describe — which is what keeps it testable on its own.
///     if let Some(status) = dbos::step_status()
///         && status.current_attempt() == status.max_attempts()
///     {
///         tracing::warn!(step = status.step_id(), "last attempt; the step is about to fail");
///     }
///     Err(Upstream)?
/// }
/// ```
pub fn step_status() -> Option<StepStatus> {
    Ctx::with_current(|ctx| ctx.and_then(Ctx::step_status))
}

/// A token that fires when the step this code is running inside is abandoned.
///
/// **This receives a cancellation; it does not raise one.** The engine holds the token and fires
/// it — on a step's timeout, on the preemption of a
/// [`preemptible`](crate::StepOptions::preemptible) step whose workflow was cancelled elsewhere,
/// and on any other path that abandons an attempt. What this hands back is a clone to watch.
///
/// **A clone is not a read-only handle, and cancelling one reaches no further than the body's own
/// watchers.** `cancel()` and `drop_guard()` are public on [`CancellationToken`], so a body *can*
/// fire the token it was handed, and every other clone of it — including ones the body passed to
/// its own detached work — will see cancelled. Nothing else changes: the engine never waits on this
/// token, it only fires it, so a body cannot end its own attempt, fail its step or cancel its
/// workflow this way. A step ends by returning, and a workflow is cancelled through
/// [`DBOS::cancel`](crate::DBOS::cancel). Firing it yourself only tells your own watchers that an
/// attempt was abandoned when it was not, which is a lie worth not telling; the type stays
/// `CancellationToken` rather than a wrapper because hiding `cancel()` would mean reimplementing
/// the half of it a body actually waits on.
///
/// Watch it from work the runtime cannot stop by dropping the step's future — a
/// [`spawn_blocking`](tokio::task::spawn_blocking) thread, or a client holding its own cancel
/// handle. **Ordinary `async` code needs nothing:** a step that times out has its future dropped,
/// which stops it at its next suspension point and runs its destructors on the way out, so a
/// connection is returned and a guard released without the body containing a line about it. This
/// token covers the remainder that dropping cannot reach.
///
/// It fires whenever the engine **abandons** an attempt. An attempt that reaches an outcome does
/// *not* fire it: the body has had its chance to clean up, and work it deliberately left running is
/// not the engine's to stop.
///
/// **A cancellation from outside only reaches a step that asked to notice it.** Cancelling a
/// workflow marks its status row; a step reads that row only under
/// [`preemptible`](crate::StepOptions::preemptible), so a default step runs to its outcome and its
/// token stays quiet, even though the workflow is already `CANCELLED`. A body on a blocking thread
/// that must learn about an out-of-process cancellation needs that option set.
///
/// Returns a token that is never cancelled when there is no step running, so a body that is also
/// called outside a workflow needs no second path — which is why this is a [`CancellationToken`]
/// rather than an `Option` of one, unlike [`workflow_id`] and [`step_id`].
///
/// ```no_run
/// async fn hashes() -> dbos::Result<u64> {
///     let token = dbos::cancellation_token();
///     // A blocking thread: dropping this step's future cannot reach it, so it watches instead.
///     let hashed = tokio::task::spawn_blocking(move || {
///         let mut total = 0;
///         while !token.is_cancelled() {
///             total += 1;
///         }
///         total
///     });
///     Ok(hashed.await.expect("the hashing thread panicked"))
/// }
/// ```
pub fn cancellation_token() -> CancellationToken {
    Ctx::with_current(|ctx| ctx.map(Ctx::cancellation).unwrap_or_default())
}

impl std::fmt::Debug for Ctx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Ctx")
            .field("workflow_id", &self.workflow.workflow_id)
            .field(
                "steps_taken",
                &self.workflow.next_step_id.load(Ordering::Relaxed),
            )
            .field("step", &self.step)
            .finish_non_exhaustive()
    }
}

/// The owner token the crate's unit tests run their contexts under. A test that writes against a
/// workflow row makes the row's `owner_xid` this, as a start or a claim would.
#[cfg(test)]
pub(crate) const TEST_OWNER: &str = "test-owner";

/// Makes [`TEST_OWNER`] the owner of `workflow_id`'s row, as a start or a claim would make its own
/// token.
#[cfg(test)]
pub(crate) async fn own_for_test(db: &dbos_test_support::TestDatabase, workflow_id: &str) {
    sqlx::query(r#"UPDATE "dbos"."workflow_status" SET owner_xid = $1 WHERE workflow_uuid = $2"#)
        .bind(TEST_OWNER)
        .bind(workflow_id)
        .execute(&db.pool().await)
        .await
        .expect("could not take ownership of the workflow row");
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    /// The step counter is the half of the context that needs no executor, so it is tested on
    /// `WorkflowState` directly. The scoping tests below need a real `Ctx`, so they launch against
    /// a leased database rather than faking an executor into existence.
    fn state() -> WorkflowState {
        WorkflowState {
            workflow_id: "wf-1".to_owned(),
            owner_xid: TEST_OWNER.to_owned(),
            deadline: None,
            encoding: PayloadEncoding::Native,
            next_step_id: AtomicI32::new(0),
            starts: Arc::default(),
        }
    }

    /// With no start in flight, the wait is over at once.
    #[tokio::test]
    async fn no_starts_in_flight_drains_at_once() {
        let starts = Arc::new(Starts::default());
        crate::rt::timeout(Duration::from_secs(1), starts.drained())
            .await
            .expect("nothing was in flight");
    }

    /// The wait lasts until the last token is dropped, and not past it.
    #[tokio::test]
    async fn drained_waits_for_the_last_start() {
        let starts = Arc::new(Starts::default());
        let first = starts.token();
        let second = starts.token();

        let mut drained = std::pin::pin!(starts.drained());
        assert!(
            crate::rt::timeout(Duration::from_millis(20), drained.as_mut())
                .await
                .is_err(),
            "drained with two starts in flight"
        );
        drop(first);
        assert!(
            crate::rt::timeout(Duration::from_millis(20), drained.as_mut())
                .await
                .is_err(),
            "drained with one start still in flight"
        );
        drop(second);
        crate::rt::timeout(Duration::from_secs(1), drained)
            .await
            .expect("the last start was dropped");
    }

    /// A token dropped by another task while the wait is parked wakes it.
    #[tokio::test]
    async fn a_token_dropped_elsewhere_wakes_the_wait() {
        let starts = Arc::new(Starts::default());
        let token = starts.token();
        let dropper = crate::rt::spawn(async move {
            crate::rt::sleep(Duration::from_millis(20)).await;
            drop(token);
        });
        crate::rt::timeout(Duration::from_secs(1), starts.drained())
            .await
            .expect("the token was dropped");
        dropper.await.unwrap();
    }

    #[test]
    fn step_ids_are_zero_based() {
        let state = state();
        assert_eq!(state.next_step_id(), 0, "step ids are zero-based");
        assert_eq!(state.next_step_id(), 1);
        assert_eq!(state.next_step_id(), 2);
    }

    #[test]
    fn step_ids_are_never_reused_across_threads() {
        // Two `Ctx` clones share one counter, so ids must be unique even when steps are allocated
        // from different threads. Steps may run concurrently, so a counter that could hand out a
        // duplicate would be a replay bug rather than a policy question.
        let state = Arc::new(state());
        let handles: Vec<_> = (0..8)
            .map(|_| {
                let state = Arc::clone(&state);
                std::thread::spawn(move || {
                    (0..100).map(|_| state.next_step_id()).collect::<Vec<_>>()
                })
            })
            .collect();
        let mut ids: Vec<i32> = handles
            .into_iter()
            .flat_map(|h| h.join().unwrap())
            .collect();
        ids.sort_unstable();
        assert_eq!(
            ids,
            (0..800).collect::<Vec<_>>(),
            "800 allocations, 800 distinct ids"
        );
    }

    #[tokio::test]
    async fn there_is_no_context_outside_a_workflow() {
        assert!(Ctx::current().is_none());
    }

    /// A real `Ctx` needs a real `Executor`, which needs a database. These live here rather than
    /// in `tests/` because `Ctx::new` and `Ctx::scope` are crate-internal — entering a context is
    /// something the engine does, never something a caller does.
    async fn ctx(
        workflow_id: &str,
        deadline: Option<Timestamp>,
    ) -> (Ctx, crate::DBOS, dbos_test_support::TestDatabase) {
        let db = dbos_test_support::test_database().await;
        let dbos = crate::DBOS::new(crate::Config {
            migrate: false,
            app_version: Some("1.0.0".to_owned()),
            ..crate::Config::new("ctx-test", db.url())
        });
        dbos.launch().await.expect("launch failed");
        let ctx = Ctx::new(
            dbos.executor("test").expect("launched"),
            workflow_id,
            crate::context::TEST_OWNER,
            deadline,
            PayloadEncoding::Native,
        );
        (ctx, dbos, db)
    }

    #[tokio::test]
    async fn a_context_is_ambient_within_its_scope_and_gone_outside_it() {
        let (ctx, dbos, _db) = ctx("wf-42", None).await;

        assert!(Ctx::current().is_none(), "nothing before");
        Ctx::scope(ctx, async {
            assert_eq!(Ctx::current().expect("inside").workflow_id(), "wf-42");
            crate::rt::yield_now().await;
            assert_eq!(
                Ctx::current().expect("still inside").workflow_id(),
                "wf-42",
                "the context survives an await, which is the whole point of a task-local"
            );
        })
        .await;
        assert!(Ctx::current().is_none(), "nothing after");

        dbos.shutdown().await;
    }

    /// The deadline reaches everything the workflow calls, including from inside a step.
    ///
    /// It lives on the shared `WorkflowState` rather than on the `Ctx`, and the difference shows up
    /// exactly here: `in_step_scope` builds a *new* `Ctx` to carry the attempt's cancellation
    /// token, so anything held on the `Ctx` itself would be silently dropped at every step
    /// boundary. A child workflow launched after a step has run must still inherit the budget its
    /// parent has left.
    #[tokio::test]
    async fn the_deadline_travels_with_the_context_and_survives_a_step_scope() {
        let deadline = Timestamp::from_epoch_ms(1_700_000_000_000);
        let (ctx, dbos, _db) = ctx("wf-deadline", Some(deadline)).await;

        let inner = ctx.clone();
        Ctx::scope(ctx, async move {
            assert_eq!(
                Ctx::current().expect("inside").deadline(),
                Some(deadline),
                "the body reads the deadline its row carries"
            );
            inner
                .in_step_scope(CancellationToken::new(), StepStatus::first(0), async {
                    assert_eq!(
                        Ctx::current().expect("inside a step").deadline(),
                        Some(deadline),
                        "a step's rebinding shares the workflow state, so the deadline comes with it"
                    );
                })
                .await;
        })
        .await;

        dbos.shutdown().await;
    }

    /// A sibling that has finished does not answer for a step still inside its own body.
    ///
    /// One half of what a workflow-wide flag would get wrong, and the half that corrupts a step
    /// that is still *running*: whichever scope ended first would clear the answer for the other,
    /// so a body still inside itself would read as no longer in a step, and the next call it made
    /// would allocate a step id instead of being the plain call a nested one has to be.
    #[tokio::test]
    async fn a_sibling_that_finishes_does_not_end_this_step() {
        let (ctx, dbos, _db) = ctx("wf-siblings", None).await;
        let (inside, wait_for_inside) = tokio::sync::oneshot::channel();
        let (done, wait_for_done) = tokio::sync::oneshot::channel();

        let long = ctx.in_step_scope(CancellationToken::new(), StepStatus::first(0), async move {
            inside.send(()).expect("the sibling is waiting on this");
            wait_for_done.await.expect("the sibling ran to completion");
            Ctx::current().expect("inside a step").in_step()
        });
        let short = async {
            wait_for_inside
                .await
                .expect("the long step reached its body");
            // A whole step scope opens and closes while the other one is suspended inside its own.
            ctx.in_step_scope(
                CancellationToken::new(),
                StepStatus::first(0),
                std::future::ready(()),
            )
            .await;
            done.send(()).expect("the long step is waiting on this");
        };

        let (still_inside, ()) = tokio::join!(long, short);
        assert!(
            still_inside,
            "a sibling finishing must not tell a step it has left its own body"
        );

        dbos.shutdown().await;
    }

    /// A step in flight does not make the workflow body around it look nested.
    ///
    /// The other half of the same hazard: with the answer shared, anything the workflow proper did
    /// while a step was running would take the plain path meant for a nested call, and so never be
    /// checkpointed at all. `a_step_built_while_a_sibling_runs_is_still_checkpointed` in
    /// [`step`](crate::step()) is what that would cost a caller.
    #[tokio::test]
    async fn a_running_step_does_not_make_the_workflow_proper_look_nested() {
        let (ctx, dbos, _db) = ctx("wf-proper", None).await;
        let (inside, wait_for_inside) = tokio::sync::oneshot::channel();
        let (looked, wait_for_look) = tokio::sync::oneshot::channel();

        let stepping = ctx.clone();
        Ctx::scope(ctx, async move {
            let held = stepping.in_step_scope(
                CancellationToken::new(),
                StepStatus::first(0),
                async move {
                    inside
                        .send(())
                        .expect("the workflow body is waiting on this");
                    wait_for_look.await.expect("the workflow body looked");
                },
            );
            let proper = async {
                wait_for_inside.await.expect("the step reached its body");
                let seen = Ctx::current().expect("inside the workflow").in_step();
                looked.send(()).expect("the step is waiting on this");
                seen
            };

            let ((), seen) = tokio::join!(held, proper);
            assert!(
                !seen,
                "the workflow body is not inside the step it is waiting on"
            );
        })
        .await;

        dbos.shutdown().await;
    }

    #[tokio::test]
    async fn the_context_does_not_cross_a_spawn() {
        // Documented rather than discovered: a spawned task is a new task with an empty
        // task-local map. That is the correct default — the spawned work is not part of the
        // durable workflow — and adopting the parent's counter would let two tasks allocate the
        // same step id.
        let (ctx, dbos, _db) = ctx("wf-42", None).await;

        let spawned_saw = Ctx::scope(ctx, async {
            crate::rt::spawn(async { Ctx::current().map(|c| c.workflow_id().to_owned()) })
                .await
                .expect("the spawned task panicked")
        })
        .await;
        assert_eq!(spawned_saw, None);

        dbos.shutdown().await;
    }

    #[tokio::test]
    async fn a_clone_shares_one_step_counter() {
        let (ctx, dbos, _db) = ctx("wf-42", None).await;
        let clone = ctx.clone();

        assert_eq!(ctx.next_step_id(), 0);
        assert_eq!(
            clone.next_step_id(),
            1,
            "a clone is the same workflow, not a second one"
        );
        assert_eq!(ctx.next_step_id(), 2);

        dbos.shutdown().await;
    }
}