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;
#[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",
¶ms,
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",
¶ms,
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,
))
}