staircase 0.0.7

Kubernetes Step-based Operator
Documentation
use std::{future::Future, sync::Arc};

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

use k8s_openapi::{
    NamespaceResourceScope,
    jiff::Timestamp,
    serde::{Serialize, de::DeserializeOwned},
};
use kube::{Resource, core::object::HasStatus};

#[cfg(feature = "controller_trace")]
use tracing::{Instrument, Level};

#[cfg(feature = "controller_metrics")]
use opentelemetry::{
    KeyValue, global,
    metrics::{Histogram, Meter},
};

use super::{
    KubeStatusUpdater, RunContext, RunError, RunOutcome, RunPlan, StatusUpdater,
    run::{RunStepConfig, run_plan},
};

/// High-level reconciler interface for running a step-based plan.
///
/// Implement this trait when a controller should resolve an incoming
/// evaluation resource into the concrete resource to reconcile, initial step
/// data, and the [`RunPlan`] to execute.
///
/// Use [`run`] to execute a reconciler. Use [`run_plan`] directly when you need
/// lower-level control over context construction, status updates, or execution
/// configuration.
///
/// [`run`]: crate::controller::run
/// [`run_plan`]: crate::controller::run_plan
pub trait Reconciler {
    type Resource: Resource
        + HasStatus<Status: Serialize + Send>
        + Resource<Scope = NamespaceResourceScope, DynamicType: Default>
        + DeserializeOwned
        + Send
        + Sync;

    /// Input received from the surrounding controller.
    ///
    /// Evaluation resolves this into the concrete [`Self::Resource`] for the
    /// run. For example, this may be an `Arc<Self::Resource>` or a finalizer
    /// event.
    type EvaluationResource;

    /// Additional value returned to the caller after the run completes.
    type EvaluationOutcome;

    type Error: Send + 'static;

    /// Data passed to the first step in the plan.
    type InitialData: Send + 'static;

    /// Plan executed for this run.
    type Plan: RunPlan<Resource = Self::Resource, Error = Self::Error, InData = Self::InitialData> + Send;

    /// Builds the plan and initial data for one run.
    #[allow(clippy::type_complexity)]
    fn evaluate_plan(
        &self,
        resource: Self::EvaluationResource,
    ) -> impl Future<
        Output = (
            Self::InitialData,
            Self::EvaluationOutcome,
            Arc<Self::Resource>,
            Self::Plan,
        ),
    > + Send;

    fn build_config(&self) -> RunReconcilerConfig { RunReconcilerConfig::default() }

    fn build_status_updater(&self) -> impl StatusUpdater<Self::Resource> + Send + Sync { KubeStatusUpdater {} }

    /// Returns the Kubernetes client associated with this reconciler.
    fn client(&self) -> impl Future<Output = Result<kube::Client, Self::Error>> + Send;
}

#[derive(Clone, Debug)]
pub struct RunReconcilerConfig {
    #[cfg(feature = "controller_metrics")]
    run_duration:    Histogram<u64>,
    run_step_config: RunStepConfig,
}

#[allow(clippy::derivable_impls)]
impl Default for RunReconcilerConfig {
    fn default() -> Self {
        #[cfg(feature = "controller_metrics")]
        return Self::with_meter(global::meter("reconciler"), Default::default());
        #[cfg(not(feature = "controller_metrics"))]
        return Self {
            run_step_config: Default::default(),
        };
    }
}

impl RunReconcilerConfig {
    #[cfg(feature = "controller_metrics")]
    pub fn with_meter(meter: Meter, run_step_config: RunStepConfig) -> Self {
        Self {
            run_duration: meter
                .u64_histogram("run.duration")
                .with_unit("ms")
                .with_description("runtime of the reconcile run")
                .with_boundaries(vec![
                    0.0, 2.0, 5.0, 10.0, 20.0, 35.0, 75.0, 125.0, 250.0, 500.0, 1000.0, 2500.0, 10000.0,
                ])
                .build(),
            run_step_config,
        }
    }
}

/// Executes one run of a high-level [`Reconciler`].
///
/// Features:
///  - `controller_trace` adds tracing spans around evaluation and step
///    execution.
///  - `controller_metrics` records run and step duration metrics.
///
/// [`Reconciler`]: crate::controller::Reconciler
pub async fn run<R: Reconciler>(
    reconciler: &R,
    resource: <R as Reconciler>::EvaluationResource,
) -> Result<(RunOutcome, R::EvaluationOutcome), RunError<R::Error>> {
    let run_start = Timestamp::now();

    #[cfg(feature = "controller_metrics")]
    let run_start_instant = Instant::now();
    let client = reconciler
        .client()
        .await
        .map_err(|source| RunError::Reconciler { source })?;

    #[cfg(feature = "controller_trace")]
    let id = format!("{:X}", fastrand::u32(..));

    let future = reconciler.evaluate_plan(resource);
    #[cfg(feature = "controller_trace")]
    let future = future.instrument(tracing::span!(Level::INFO, "evaluate_plan", id));

    let (data, eval_outcome, resource, steps) = future.await;

    let context = RunContext {
        client,
        resource,
        run_start,
    };

    let status_updater = reconciler.build_status_updater();
    let config = reconciler.build_config();

    let future = run_plan(context, data, steps, &status_updater, &config.run_step_config);
    #[cfg(feature = "controller_trace")]
    let future = future.instrument(tracing::span!(
        Level::INFO,
        "reconcile",
        id,
        step = tracing::field::Empty
    )); // step will be overwritten

    let result = future.await;

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

    result.map(|outcome| (outcome, eval_outcome))
}

#[cfg(feature = "controller_metrics")]
fn key_value_of_outcome(outcome: &RunOutcome) -> KeyValue {
    match outcome {
        RunOutcome::Stopped => KeyValue::new("outcome", "stopped"),
        RunOutcome::Modified => KeyValue::new("outcome", "modified"),
        RunOutcome::NoOp => KeyValue::new("outcome", "noop"),
    }
}

#[cfg(test)]
mod tests {
    use std::{
        convert::Infallible,
        error::Error,
        fmt,
        sync::{Arc, Mutex},
    };

    use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
    use kube::{CustomResource, client::Body};
    use schemars::JsonSchema;
    use serde::{Deserialize, Serialize};

    use crate::controller::{BoxPlan, Mutable, MutableStepOutcome, MutableStepResult, Plan, Step};

    use super::*;

    #[derive(CustomResource, Deserialize, Serialize, Clone, Debug, JsonSchema)]
    #[kube(
        group = "tests.staircase.dev",
        version = "v1",
        kind = "RunWrapperTest",
        status = "RunWrapperTestStatus",
        namespaced
    )]
    struct RunWrapperTestSpec {}

    #[derive(Deserialize, Serialize, Clone, Debug, Default, JsonSchema)]
    struct RunWrapperTestStatus {}

    #[derive(Debug, PartialEq, Eq)]
    enum TestError {
        Client,
    }

    impl fmt::Display for TestError {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                TestError::Client => f.write_str("client failed"),
            }
        }
    }

    impl Error for TestError {}

    struct ModifiedStep {
        log: Arc<Mutex<Vec<&'static str>>>,
    }

    impl Step for ModifiedStep {
        type Error = TestError;
        type InData = ();
        type Mode = Mutable;
        type OutData = ();
        type Resource = RunWrapperTest;

        async fn run(
            &self,
            _context: &RunContext<Self::Resource>,
            _data: Self::InData,
        ) -> MutableStepResult<Self::Resource, Self::OutData, Self::Error> {
            self.log.lock().unwrap().push("step");
            Ok(MutableStepOutcome::Modified { status: None })
        }
    }

    struct FakeReconciler {
        client: kube::Client,
        log:    Arc<Mutex<Vec<&'static str>>>,
    }

    impl Reconciler for FakeReconciler {
        type Error = TestError;
        type EvaluationOutcome = &'static str;
        type EvaluationResource = Arc<RunWrapperTest>;
        type InitialData = ();
        type Plan = BoxPlan<Self::Resource, Self::Error, Self::InitialData, ()>;
        type Resource = RunWrapperTest;

        async fn evaluate_plan(
            &self,
            resource: Self::EvaluationResource,
        ) -> ((), Self::EvaluationOutcome, Arc<Self::Resource>, Self::Plan) {
            self.log.lock().unwrap().push("evaluate_plan");
            let plan = Plan::from(ModifiedStep { log: self.log.clone() }).boxed();
            ((), "evaluated", resource, plan)
        }

        async fn client(&self) -> Result<kube::Client, Self::Error> {
            self.log.lock().unwrap().push("client");
            Ok(self.client.clone())
        }
    }

    struct FailingClientReconciler {
        log: Arc<Mutex<Vec<&'static str>>>,
    }

    impl Reconciler for FailingClientReconciler {
        type Error = TestError;
        type EvaluationOutcome = &'static str;
        type EvaluationResource = Arc<RunWrapperTest>;
        type InitialData = ();
        type Plan = BoxPlan<Self::Resource, Self::Error, Self::InitialData, ()>;
        type Resource = RunWrapperTest;

        async fn evaluate_plan(
            &self,
            resource: Self::EvaluationResource,
        ) -> ((), Self::EvaluationOutcome, Arc<Self::Resource>, Self::Plan) {
            self.log.lock().unwrap().push("evaluate_plan");
            (
                (),
                "evaluated",
                resource,
                Plan::empty::<RunWrapperTest, TestError, ()>().boxed(),
            )
        }

        async fn client(&self) -> Result<kube::Client, Self::Error> {
            self.log.lock().unwrap().push("client");
            Err(TestError::Client)
        }
    }

    fn dummy_client() -> kube::Client {
        let service = tower::service_fn(|_request: http::Request<Body>| async {
            Ok::<_, Infallible>(http::Response::new(Body::empty()))
        });
        kube::Client::new(service, "default")
    }

    fn resource() -> Arc<RunWrapperTest> {
        Arc::new(RunWrapperTest {
            metadata: ObjectMeta {
                name: Some("sample".to_string()),
                namespace: Some("default".to_string()),
                ..Default::default()
            },
            spec:     RunWrapperTestSpec {},
            status:   None,
        })
    }

    #[test]
    fn run_calls_client_evaluates_steps_and_returns_evaluation_outcome() {
        let log = Arc::new(Mutex::new(Vec::new()));
        let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();
        let (outcome, evaluation_outcome) = rt
            .block_on(async {
                let reconciler = FakeReconciler {
                    client: dummy_client(),
                    log:    log.clone(),
                };
                run(&reconciler, resource()).await
            })
            .unwrap();

        assert!(matches!(outcome, RunOutcome::Modified));
        assert_eq!(evaluation_outcome, "evaluated");
        assert_eq!(*log.lock().unwrap(), vec!["client", "evaluate_plan", "step"]);
    }

    #[test]
    fn run_returns_client_error_without_evaluating_steps() {
        let log = Arc::new(Mutex::new(Vec::new()));
        let reconciler = FailingClientReconciler { log: log.clone() };
        let rt = tokio::runtime::Builder::new_current_thread().build().unwrap();

        let error = rt.block_on(run(&reconciler, resource())).unwrap_err();

        assert!(matches!(error, RunError::Reconciler {
            source: TestError::Client,
        }));
        assert_eq!(*log.lock().unwrap(), vec!["client"]);
    }
}