datum-agent 0.10.9

Embeddable Datum job registry and lifecycle supervisor
Documentation
//! Built-in zero-infrastructure job factories for demos and end-to-end cluster tests.

use std::{collections::HashMap, time::Duration};

use datum::{Keep, RestartSettings, Sink, Source, ThrottleMode};

use crate::{AgentError, AgentResult, JobMat, JobRestartPolicy, JobSpec, dcp::DcpJobFactories};

const DEFAULT_TICKER_INTERVAL_MS: u64 = 20;
const DEFAULT_COUNTER_INTERVAL_MS: u64 = 20;

/// Return the daemon's built-in DCP job factories.
#[must_use]
pub fn registered_factories() -> DcpJobFactories {
    let factories = DcpJobFactories::new();
    register_empty(&factories);
    register_ticker(&factories);
    register_counter(&factories);
    factories
}

fn register_empty(factories: &DcpJobFactories) {
    factories
        .register("empty", |instance_name, _params| {
            Ok(JobSpec::new(instance_name, |context| {
                let control = context.control();
                Ok(Source::<u64>::empty()
                    .instrumented(
                        format!("{}:{}", context.name(), context.generation()),
                        context.instrumentation_registry(),
                    )
                    .via_mat(context.drain_flow(), Keep::right)
                    .to_mat(Sink::ignore(), move |_switch, completion| {
                        JobMat::new(completion, control.clone())
                    }))
            }))
        })
        .expect("built-in empty factory registers");
}

fn register_ticker(factories: &DcpJobFactories) {
    factories
        .register("ticker", |instance_name, params| {
            let interval = interval_param(
                "ticker",
                &params,
                Duration::from_millis(DEFAULT_TICKER_INTERVAL_MS),
            )?;
            Ok(JobSpec::new(instance_name, move |context| {
                let control = context.control();
                Ok(Source::tick(Duration::ZERO, interval, 1_u64)
                    .instrumented(
                        format!("{}:{}", context.name(), context.generation()),
                        context.instrumentation_registry(),
                    )
                    .via_mat(context.drain_flow(), Keep::right)
                    .to_mat(Sink::ignore(), move |_switch, completion| {
                        JobMat::new(completion, control.clone())
                    }))
            })
            .with_restart_policy(demo_restart_policy()))
        })
        .expect("built-in ticker factory registers");
}

fn register_counter(factories: &DcpJobFactories) {
    factories
        .register("counter", |instance_name, params| {
            let interval = interval_param(
                "counter",
                &params,
                Duration::from_millis(DEFAULT_COUNTER_INTERVAL_MS),
            )?;
            Ok(JobSpec::new(instance_name, move |context| {
                let control = context.control();
                Ok(
                    Source::unfold(0_u64, |current| Some((current.saturating_add(1), current)))
                        .throttle(1, interval, 0, ThrottleMode::Shaping)
                        .instrumented(
                            format!("{}:{}", context.name(), context.generation()),
                            context.instrumentation_registry(),
                        )
                        .via_mat(context.drain_flow(), Keep::right)
                        .to_mat(Sink::ignore(), move |_switch, completion| {
                            JobMat::new(completion, control.clone())
                        }),
                )
            })
            .with_restart_policy(demo_restart_policy()))
        })
        .expect("built-in counter factory registers");
}

fn interval_param(
    factory: &'static str,
    params: &HashMap<String, String>,
    default: Duration,
) -> AgentResult<Duration> {
    let Some(raw) = params.get("interval_ms") else {
        return Ok(default);
    };
    let millis = raw.parse::<u64>().map_err(|error| {
        invalid_factory_parameter(
            factory,
            "interval_ms",
            raw,
            format!("expected a positive integer millisecond value ({error})"),
        )
    })?;
    if millis == 0 {
        return Err(invalid_factory_parameter(
            factory,
            "interval_ms",
            raw,
            "must be greater than zero",
        ));
    }
    Ok(Duration::from_millis(millis))
}

fn invalid_factory_parameter(
    factory: &str,
    parameter: &str,
    value: &str,
    reason: impl Into<String>,
) -> AgentError {
    AgentError::InvalidFactoryParameter {
        factory: factory.to_owned(),
        parameter: parameter.to_owned(),
        value: value.to_owned(),
        reason: reason.into(),
    }
}

fn demo_restart_policy() -> JobRestartPolicy {
    JobRestartPolicy::always(RestartSettings::new(
        Duration::from_millis(100),
        Duration::from_secs(1),
        0.0,
    ))
}