staircase 0.0.7

//! Typed reconcile plans.
//!
//! This module implements typed step chaining with ordinary Rust types instead
//! of runtime type checks. Each [`Step`] has an `InData` and `OutData`;
//! [`Plan`] stores a nested type-level chain where every `.then(...)` requires
//! the next step's `InData` to match the previous stage's `OutData`.
//!
//! The main public pieces are:
//! - [`RunPlan`]: execution trait implemented by every plan node.
//! - [`Plan`]: small builder wrapper exposing `from`, `then`, `join`,
//!   `join_map`, and `boxed`.
//! - [`JoinSteps`]: trait implemented for tuples of immutable observation
//!   steps.
//!
//! Parallel joins only accept immutable steps, so modifying branches are
//! rejected by the type system before a plan can be built.
//!
//! The concrete node types used in `Plan<...>` return values are public only
//! because Rust currently requires public functions to expose nameable return
//! types. Treat them as implementation details and build plans through
//! [`Plan`].

// The public builder API deliberately hides a nested type tree. `Plan::then`
// creates `Then<PreviousPlan, Leaf<NextStep>>`, so the compiler can enforce the
// data edge with a normal associated-type equality bound:
// `NextStep::InData == PreviousPlan::OutData`.
//
// `Plan::join` works the same way, but inserts `Join<TupleOfSteps>`.
// `Plan::join_map` wraps that join node in a `Map`, so all parallel outcome
// handling remains in one place. Rust has no variadic generics, so
// `impl_join_steps!` generates `JoinSteps` impls for tuple arities. Each impl
// uses `futures_util::join!`, which polls all branch futures concurrently and
// returns their results in tuple order.
//
// `BoxPlan` is the escape hatch for branchy reconcilers. It erases the nested
// concrete plan type while preserving the typed input/output boundary through
// the `RunPlan` trait object.
//
// `PlanFuture` is pinned because boxed `dyn Future`s are not `Unpin` by
// default. Returning one alias from `RunPlan` keeps recursive composition
// object safe enough for `BoxPlan` without exposing concrete async block types.

use std::{borrow::Cow, error::Error as StdError, fmt, future::Future, marker::PhantomData, pin::Pin};

#[cfg(feature = "controller_metrics")]
use std::time::Instant;

use super::{ImmutableStep, MutableStepOutcome, RunContext, RunStepConfig, Step, StepMode};
use k8s_openapi::{NamespaceResourceScope, serde::Serialize};
use kube::{Resource, core::object::HasStatus};
use thiserror::Error;

/// Failure returned by one step.
#[derive(Debug, Error)]
#[error("step `{step}` failed")]
pub struct StepFailure<E> {
    /// Trace/display name of the failed step.
    pub step:   Cow<'static, str>,
    /// Error returned by the failed step.
    #[source]
    pub source: E,
}

/// Error returned by a reconcile run.
#[derive(Debug)]
pub enum RunError<E> {
    /// An error that happens within this crate, e.g. network issues with
    /// communication with kubernetes api server.
    Reconciler { source: E },
    /// One sequential step failed.
    Step { step: Cow<'static, str>, source: E },
    /// One or more joined immutable steps failed.
    Parallel { errors: Vec<StepFailure<E>> },
}

impl<E> fmt::Display for RunError<E> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Reconciler { .. } => f.write_str("reconciler failed"),
            Self::Step { step, .. } => write!(f, "step `{step}` failed"),
            Self::Parallel { errors } => write!(f, "one or more parallel steps failed ({} failures)", errors.len()),
        }
    }
}

impl<E> StdError for RunError<E>
where
    E: StdError + 'static,
{
    fn source(&self) -> Option<&(dyn StdError + 'static)> {
        match self {
            Self::Reconciler { source } | Self::Step { source, .. } => Some(source),
            Self::Parallel { errors } => errors.first().map(|error| error as &(dyn StdError + 'static)),
        }
    }
}

impl<E> From<StepFailure<E>> for RunError<E> {
    fn from(value: StepFailure<E>) -> Self {
        Self::Step {
            step:   value.step,
            source: value.source,
        }
    }
}

pub type BoxPlan<R, E, InData, OutData> =
    Box<dyn RunPlan<Resource = R, Error = E, InData = InData, OutData = OutData> + Send + Sync>;

pub type PlanFuture<'a, R, OutData, E> = Pin<
    Box<dyn Future<Output = Result<MutableStepOutcome<<R as HasStatus>::Status, OutData>, RunError<E>>> + Send + 'a>,
>;

pub trait RunPlan {
    type Resource: HasStatus<Status: Send> + Send + Sync;
    type Error: Send + 'static;
    type InData: Send + 'static;
    type OutData: Send + 'static;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error>;
}

impl<P> RunPlan for Box<P>
where
    P: RunPlan + ?Sized,
{
    type Error = P::Error;
    type InData = P::InData;
    type OutData = P::OutData;
    type Resource = P::Resource;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        self.as_ref().run_plan(context, data, config)
    }
}

/// Typed execution plan for a reconcile run.
///
/// Use [`Plan::from`] to start with one step, [`Plan::then`] to add sequential
/// steps, and [`Plan::join`] or [`Plan::join_map`] to run observation steps in
/// parallel.
pub struct Plan<P> {
    inner: P,
}

impl Plan<()> {
    /// Starts a plan from a single step.
    pub fn from<S>(step: S) -> Plan<Leaf<S>>
    where
        S: Step,
    {
        Plan { inner: Leaf { step } }
    }

    /// Builds a plan that does no work and passes its input data through.
    pub fn empty<R, E, D>() -> Plan<Empty<R, E, D>>
    where
        R: HasStatus<Status: Send> + Send + Sync,
        E: Send + 'static,
        D: Send + Sync + 'static,
    {
        Plan {
            inner: Empty { _marker: PhantomData },
        }
    }
}

impl<P> Plan<P> {
    /// Adds a sequential step.
    pub fn then<S>(self, step: S) -> Plan<Then<P, Leaf<S>>>
    where
        P: RunPlan,
        S: Step<Resource = P::Resource, Error = P::Error, InData = P::OutData>,
    {
        Plan {
            inner: Then {
                first:  self.inner,
                second: Leaf { step },
            },
        }
    }

    /// Adds a parallel stage.
    ///
    /// Each joined step receives a clone of the previous stage's output. Joined
    /// steps need `type Mode = Immutable` so they implement [`ImmutableStep`].
    pub fn join<J>(self, steps: J) -> Plan<Then<P, Join<J>>>
    where
        P: RunPlan,
        J: JoinSteps<Resource = P::Resource, Error = P::Error, InData = P::OutData>,
    {
        Plan {
            inner: Then {
                first:  self.inner,
                second: Join { steps },
            },
        }
    }

    /// Adds a parallel stage and maps its tuple output.
    ///
    /// This has the same execution and error semantics as [`Plan::join`], but
    /// lets callers immediately turn the tuple of branch outputs into a domain
    /// type.
    ///
    /// Like [`Plan::join`], every branch step needs `type Mode = Immutable` so
    /// it implements [`ImmutableStep`].
    #[expect(clippy::type_complexity)]
    pub fn join_map<J, F, OutData>(self, steps: J, f: F) -> Plan<Then<P, Map<Join<J>, F, OutData>>>
    where
        P: RunPlan,
        J: JoinSteps<Resource = P::Resource, Error = P::Error, InData = P::OutData>,
        F: Fn(J::OutData) -> OutData + Send + Sync,
        OutData: Send + Sync + 'static,
    {
        Plan {
            inner: Then {
                first:  self.inner,
                second: Map {
                    plan: Join { steps },
                    f,
                    _marker: PhantomData,
                },
            },
        }
    }

    /// Erases the concrete plan type.
    ///
    /// This is useful for reconcilers with branches that return different plan
    /// shapes.
    pub fn boxed(self) -> BoxPlan<P::Resource, P::Error, P::InData, P::OutData>
    where
        P: RunPlan + Send + Sync + 'static,
    {
        Box::new(self)
    }
}

impl<P> RunPlan for Plan<P>
where
    P: RunPlan,
{
    type Error = P::Error;
    type InData = P::InData;
    type OutData = P::OutData;
    type Resource = P::Resource;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        self.inner.run_plan(context, data, config)
    }
}

#[doc(hidden)]
pub struct Empty<R, E, D> {
    _marker: PhantomData<(R, E, D)>,
}

impl<R, E, D> RunPlan for Empty<R, E, D>
where
    R: HasStatus<Status: Send> + Send + Sync,
    E: Send + Sync + 'static,
    D: Send + Sync + 'static,
{
    type Error = E;
    type InData = D;
    type OutData = D;
    type Resource = R;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        _: &'b RunContext<Self::Resource>,
        data: Self::InData,
        _: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        Box::pin(async move { Ok(MutableStepOutcome::NoModification { data }) })
    }
}

#[doc(hidden)]
pub struct Leaf<S> {
    step: S,
}

impl<S> RunPlan for Leaf<S>
where
    S: Step + Send + Sync,
    S::Resource: Resource<Scope = NamespaceResourceScope, DynamicType: Default> + HasStatus + Send + Sync,
    <S::Resource as HasStatus>::Status: Serialize + Send,
{
    type Error = S::Error;
    type InData = S::InData;
    type OutData = S::OutData;
    type Resource = S::Resource;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        Box::pin(async move {
            execute_step(&self.step, context, data, config)
                .await
                .map_err(RunError::from)
        })
    }
}

#[doc(hidden)]
pub struct Then<A, B> {
    first:  A,
    second: B,
}

impl<A, B> RunPlan for Then<A, B>
where
    A: RunPlan + Sync,
    B: RunPlan<Resource = A::Resource, Error = A::Error, InData = A::OutData> + Sync,
{
    type Error = A::Error;
    type InData = A::InData;
    type OutData = B::OutData;
    type Resource = A::Resource;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        Box::pin(async move {
            match self.first.run_plan(context, data, config).await? {
                MutableStepOutcome::Stop => Ok(MutableStepOutcome::Stop),
                MutableStepOutcome::Modified { status } => Ok(MutableStepOutcome::Modified { status }),
                MutableStepOutcome::NoModification { data } => self.second.run_plan(context, data, config).await,
            }
        })
    }
}

#[doc(hidden)]
pub struct Join<J> {
    steps: J,
}

impl<J> RunPlan for Join<J>
where
    J: JoinSteps + Sync,
{
    type Error = J::Error;
    type InData = J::InData;
    type OutData = J::OutData;
    type Resource = J::Resource;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        self.steps.run_join(context, data, config)
    }
}

#[doc(hidden)]
pub struct Map<P, F, OutData> {
    plan:    P,
    f:       F,
    _marker: PhantomData<OutData>,
}

impl<P, F, OutData> RunPlan for Map<P, F, OutData>
where
    P: RunPlan + Sync,
    F: Fn(P::OutData) -> OutData + Send + Sync,
    OutData: Send + Sync + 'static,
{
    type Error = P::Error;
    type InData = P::InData;
    type OutData = OutData;
    type Resource = P::Resource;

    fn run_plan<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
        Box::pin(async move {
            match self.plan.run_plan(context, data, config).await? {
                MutableStepOutcome::Stop => Ok(MutableStepOutcome::Stop),
                MutableStepOutcome::Modified { status } => Ok(MutableStepOutcome::Modified { status }),
                MutableStepOutcome::NoModification { data } => {
                    Ok(MutableStepOutcome::NoModification { data: (self.f)(data) })
                },
            }
        })
    }
}

pub trait JoinSteps {
    type Resource: HasStatus<Status: Send> + Send + Sync;
    type Error: Send + 'static;
    type InData: Clone + Send + 'static;
    type OutData: Send + 'static;

    fn run_join<'a, 'b: 'a>(
        &'a self,
        context: &'b RunContext<Self::Resource>,
        data: Self::InData,
        config: &'a RunStepConfig,
    ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error>;
}

async fn execute_step<'a, S>(
    step: &'a S,
    context: &'a RunContext<S::Resource>,
    data: S::InData,
    #[allow(unused_variables)] config: &'a RunStepConfig,
) -> Result<MutableStepOutcome<<S::Resource as HasStatus>::Status, S::OutData>, StepFailure<S::Error>>
where
    S: Step + Sync + ?Sized,
    S::Resource: Resource<Scope = NamespaceResourceScope, DynamicType: Default> + HasStatus + Send + Sync,
    <S::Resource as HasStatus>::Status: Serialize + Send,
{
    #[cfg(any(feature = "controller_trace", feature = "controller_metrics"))]
    let step_name = step.traced_name();
    #[cfg(not(any(feature = "controller_trace", feature = "controller_metrics")))]
    let step_name = step.traced_name();

    let func = step.run(context, data);

    #[cfg(feature = "controller_trace")]
    use tracing::Instrument;

    #[cfg(feature = "controller_trace")]
    let func = func.instrument(tracing::span!(
        tracing::Level::INFO,
        "step",
        name = match step_name {
            Cow::Borrowed(step_name) => step_name,
            Cow::Owned(ref step_name) => step_name.as_str(),
        },
        comp = tracing::field::Empty,
    ));

    #[cfg(feature = "controller_metrics")]
    let step_key = match &step_name {
        Cow::Borrowed(step_name) => opentelemetry::KeyValue::new("step", *step_name),
        Cow::Owned(step_name) => opentelemetry::KeyValue::new("step", step_name.clone()),
    };

    #[cfg(feature = "controller_metrics")]
    let step_start_instant = Instant::now();

    let result = func.await.map(S::Mode::into_mutable_outcome);

    #[cfg(feature = "controller_metrics")]
    {
        let ms = step_start_instant.elapsed().as_millis().min(u64::MAX as u128) as u64;
        match &result {
            Ok(outcome) => config.step_duration.record(ms, &[
                opentelemetry::KeyValue::new("ok", true),
                step_key.clone(),
                key_value_of_outcome(outcome),
            ]),
            Err(_) => config
                .step_duration
                .record(ms, &[opentelemetry::KeyValue::new("ok", false), step_key.clone()]),
        };
    }

    result.map_err(|source| StepFailure {
        step: step_name,
        source,
    })
}

#[cfg(feature = "controller_metrics")]
fn key_value_of_outcome<S, D>(outcome: &MutableStepOutcome<S, D>) -> opentelemetry::KeyValue {
    match outcome {
        MutableStepOutcome::Stop => opentelemetry::KeyValue::new("outcome", "stop"),
        MutableStepOutcome::Modified { .. } => opentelemetry::KeyValue::new("outcome", "modified"),
        MutableStepOutcome::NoModification { .. } => opentelemetry::KeyValue::new("outcome", "nomodification"),
    }
}

macro_rules! impl_join_steps {
    ($($type_name:ident:$var_name:ident:$idx:tt),+) => {
        impl<InData, Resource, Error, $($type_name),+> JoinSteps for ($($type_name,)+)
        where
            InData: Clone + Send + 'static,
            Resource: kube::Resource<Scope = NamespaceResourceScope, DynamicType: Default> + HasStatus + Send + Sync,
            <Resource as HasStatus>::Status: Serialize + Send,
            Error: Send + 'static,
            $(
                $type_name: ImmutableStep<Resource = Resource, Error = Error, InData = InData> + Send + Sync,
            )+
        {
            type Error = Error;
            type InData = InData;
            type OutData = ($($type_name::OutData,)+);
            type Resource = Resource;

            fn run_join<'a, 'b: 'a>(
                &'a self,
                context: &'b RunContext<Self::Resource>,
                data: Self::InData,
                config: &'a RunStepConfig,
            ) -> PlanFuture<'a, Self::Resource, Self::OutData, Self::Error> {
                Box::pin(async move {
                    $(
                        let $var_name = execute_step(&self.$idx, context, data.clone(), config);
                    )+

                    let ($($var_name,)+) = futures_util::join!($($var_name,)+);

                    let mut errors = Vec::new();
                    let mut stopped = false;

                    $(
                        let $var_name = match $var_name {
                            Err(error) => {
                                errors.push(error);
                                None
                            },
                            Ok(MutableStepOutcome::Stop) => {
                                stopped = true;
                                None
                            },
                            Ok(MutableStepOutcome::Modified { .. }) => unreachable!("immutable steps cannot modify"),
                            Ok(MutableStepOutcome::NoModification { data }) => Some(data),
                        };
                    )+

                    if !errors.is_empty() {
                        return Err(RunError::Parallel { errors });
                    }
                    if stopped {
                        return Ok(MutableStepOutcome::Stop);
                    }

                    Ok(MutableStepOutcome::NoModification {
                        data: ($($var_name.expect("join output is present when branch did not stop, modify or fail"),)+),
                    })
                })
            }
        }
    };
}

impl_join_steps!(A:a:0, B:b:1);
impl_join_steps!(A:a:0, B:b:1, C:c:2);
impl_join_steps!(A:a:0, B:b:1, C:c:2, D:d:3);
impl_join_steps!(A:a:0, B:b:1, C:c:2, D:d:3, E:e:4);
impl_join_steps!(A:a:0, B:b:1, C:c:2, D:d:3, E:e:4, F:f:5);
impl_join_steps!(A:a:0, B:b:1, C:c:2, D:d:3, E:e:4, F:f:5, G:g:6);
impl_join_steps!(A:a:0, B:b:1, C:c:2, D:d:3, E:e:4, F:f:5, G:g:6, H:h:7);