use std::sync::OnceLock;
use ikigai_scheduler::{Scheduler, SchedulerSpec};
pub const CONFIG_KEY: &str = "scheduler";
pub const ENV_VAR: &str = "IKIGAI_SCHEDULER";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SchedulerSource {
Flag,
Config,
Env,
Default,
}
impl SchedulerSource {
pub fn as_str(self) -> &'static str {
match self {
SchedulerSource::Flag => "flag",
SchedulerSource::Config => "config",
SchedulerSource::Env => "env",
SchedulerSource::Default => "default",
}
}
}
static FLAG_SPEC: OnceLock<String> = OnceLock::new();
pub fn set_scheduler_spec(spec: impl Into<String>) -> Result<(), String> {
let spec = spec.into();
spec.parse::<SchedulerSpec>()?;
let _ = FLAG_SPEC.set(spec);
Ok(())
}
fn resolved() -> &'static (Scheduler, SchedulerSource) {
static RESOLVED: OnceLock<(Scheduler, SchedulerSource)> = OnceLock::new();
RESOLVED.get_or_init(|| {
let decision = decide(
FLAG_SPEC.get().map(String::as_str),
config_spec().as_deref(),
std::env::var(ENV_VAR).ok().as_deref(),
);
for warning in &decision.warnings {
eprintln!("ikigai: {warning}");
}
(decision.spec.build(), decision.source)
})
}
pub fn scheduler() -> Scheduler {
resolved().0.clone()
}
pub fn scheduler_source() -> SchedulerSource {
resolved().1
}
fn config_spec() -> Option<String> {
crate::config::get(&format!("{}.{CONFIG_KEY}", crate::instance_name()))
.or_else(|| crate::config::get(CONFIG_KEY))
}
pub const WIDTH_ROUTING_CONFIG_KEY: &str = "width-routing";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RoutingSource {
Flag,
Config,
Default,
}
impl RoutingSource {
pub fn as_str(self) -> &'static str {
match self {
RoutingSource::Flag => "flag",
RoutingSource::Config => "config",
RoutingSource::Default => "default",
}
}
}
static WIDTH_ROUTING_FLAG: OnceLock<bool> = OnceLock::new();
pub fn set_width_routing(value: &str) -> Result<(), String> {
let on = parse_switch(value)?;
let _ = WIDTH_ROUTING_FLAG.set(on);
Ok(())
}
pub fn width_routing() -> bool {
resolved_routing().0
}
pub fn width_routing_source() -> RoutingSource {
resolved_routing().1
}
fn resolved_routing() -> &'static (bool, RoutingSource) {
static RESOLVED: OnceLock<(bool, RoutingSource)> = OnceLock::new();
RESOLVED.get_or_init(|| {
let (on, source, warnings) = decide_routing(
WIDTH_ROUTING_FLAG.get().copied(),
routing_config_spec().as_deref(),
);
for warning in &warnings {
eprintln!("ikigai: {warning}");
}
(on, source)
})
}
fn routing_config_spec() -> Option<String> {
crate::config::get(&format!(
"{}.{WIDTH_ROUTING_CONFIG_KEY}",
crate::instance_name()
))
.or_else(|| crate::config::get(WIDTH_ROUTING_CONFIG_KEY))
}
fn parse_switch(value: &str) -> Result<bool, String> {
match value.trim().to_ascii_lowercase().as_str() {
"on" | "true" => Ok(true),
"off" | "false" => Ok(false),
other => Err(format!("`{other}` is not a width-routing setting (on|off)")),
}
}
fn decide_routing(flag: Option<bool>, config: Option<&str>) -> (bool, RoutingSource, Vec<String>) {
let mut warnings = Vec::new();
if let Some(on) = flag {
return (on, RoutingSource::Flag, warnings);
}
if let Some(value) = config {
match parse_switch(value) {
Ok(on) => return (on, RoutingSource::Config, warnings),
Err(e) => warnings.push(format!(
"{e} (from {WIDTH_ROUTING_CONFIG_KEY}); ignoring it"
)),
}
}
(false, RoutingSource::Default, warnings)
}
struct Decision {
spec: SchedulerSpec,
source: SchedulerSource,
warnings: Vec<String>,
}
fn decide(flag: Option<&str>, config: Option<&str>, env: Option<&str>) -> Decision {
let mut warnings = Vec::new();
let channels = [
(SchedulerSource::Flag, flag, "--scheduler"),
(SchedulerSource::Config, config, CONFIG_KEY),
(SchedulerSource::Env, env, ENV_VAR),
];
for (source, value, label) in channels {
let Some(value) = value else { continue };
match value.parse::<SchedulerSpec>() {
Ok(spec) => {
if source == SchedulerSource::Env {
warnings.push(format!(
"{ENV_VAR} is deprecated; write `{CONFIG_KEY} = \"{spec}\"` in \
{} or pass `--scheduler {spec}`",
crate::config::config_path().display()
));
}
return Decision {
spec,
source,
warnings,
};
}
Err(e) => warnings.push(format!("{e} (from {label}); ignoring it")),
}
}
Decision {
spec: SchedulerSpec::Single,
source: SchedulerSource::Default,
warnings,
}
}
pub struct ConfiguredScheduler {
scheduler: Scheduler,
source: SchedulerSource,
width_routing: bool,
routing_source: RoutingSource,
}
impl ikigai_core::SchedulerReporter for ConfiguredScheduler {
fn rows(&self) -> Vec<(String, String)> {
let mut rows = ikigai_core::SchedulerReporter::rows(&self.scheduler);
rows.push(("source".to_string(), self.source.as_str().to_string()));
rows.push((
"routing".to_string(),
if self.width_routing {
"by-width"
} else {
"off"
}
.to_string(),
));
rows.push((
"routing.by".to_string(),
self.routing_source.as_str().to_string(),
));
rows
}
}
pub fn reporter() -> ConfiguredScheduler {
let (scheduler, source) = resolved();
ConfiguredScheduler {
scheduler: scheduler.clone(),
source: *source,
width_routing: width_routing(),
routing_source: width_routing_source(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use ikigai_core::SchedulerReporter;
#[test]
fn flag_beats_config_beats_env_beats_default() {
let all = decide(Some("pool:2"), Some("pool:3"), Some("pool:4"));
assert_eq!(all.spec, SchedulerSpec::Pool(2));
assert_eq!(all.source, SchedulerSource::Flag);
let no_flag = decide(None, Some("pool:3"), Some("pool:4"));
assert_eq!(no_flag.spec, SchedulerSpec::Pool(3));
assert_eq!(no_flag.source, SchedulerSource::Config);
let env_only = decide(None, None, Some("pool:4"));
assert_eq!(env_only.spec, SchedulerSpec::Pool(4));
assert_eq!(env_only.source, SchedulerSource::Env);
let nothing = decide(None, None, None);
assert_eq!(nothing.spec, SchedulerSpec::Single);
assert_eq!(nothing.source, SchedulerSource::Default);
assert!(nothing.warnings.is_empty(), "silence when nothing is set");
}
#[test]
fn the_env_channel_warns_only_when_it_decides() {
let decided = decide(None, None, Some("pool:4"));
assert_eq!(decided.warnings.len(), 1);
assert!(
decided.warnings[0].contains(ENV_VAR) && decided.warnings[0].contains("deprecated"),
"{:?}",
decided.warnings
);
assert!(
decided.warnings[0].contains("pool:4"),
"{:?}",
decided.warnings
);
let overridden = decide(Some("single"), None, Some("pool:4"));
assert!(overridden.warnings.is_empty(), "{:?}", overridden.warnings);
}
#[test]
fn an_invalid_env_value_warns_and_falls_back() {
let decision = decide(None, None, Some("pool:xyz"));
assert_eq!(decision.spec, SchedulerSpec::Single);
assert_eq!(decision.source, SchedulerSource::Default);
assert_eq!(decision.warnings.len(), 1);
assert!(
decision.warnings[0].contains("pool:xyz"),
"{:?}",
decision.warnings
);
}
#[test]
fn an_invalid_config_value_falls_through_to_the_env() {
let decision = decide(None, Some("nonsense"), Some("pool:4"));
assert_eq!(decision.spec, SchedulerSpec::Pool(4));
assert_eq!(decision.source, SchedulerSource::Env);
assert_eq!(decision.warnings.len(), 2, "{:?}", decision.warnings);
}
#[test]
fn an_invalid_flag_value_is_an_error() {
let e = set_scheduler_spec("pool:xyz").expect_err("a typo'd flag must not be accepted");
assert!(e.contains("pool:xyz"), "{e}");
assert!(set_scheduler_spec("nonsense").is_err());
}
#[test]
fn the_reporter_adds_the_deciding_channel_as_a_row() {
let reporter = ConfiguredScheduler {
scheduler: SchedulerSpec::Pool(3).build(),
source: SchedulerSource::Config,
width_routing: false,
routing_source: RoutingSource::Default,
};
let rows = reporter.rows();
assert!(rows.contains(&("backend".to_string(), "pool:3".to_string())));
assert!(rows.contains(&("threads".to_string(), "3".to_string())));
assert!(rows.contains(&("source".to_string(), "config".to_string())));
}
#[test]
fn the_reporter_states_whether_width_routing_is_on_and_who_said_so() {
let off = ConfiguredScheduler {
scheduler: SchedulerSpec::Single.build(),
source: SchedulerSource::Default,
width_routing: false,
routing_source: RoutingSource::Default,
};
assert!(off
.rows()
.contains(&("routing".to_string(), "off".to_string())));
assert!(off
.rows()
.contains(&("routing.by".to_string(), "default".to_string())));
let on = ConfiguredScheduler {
scheduler: SchedulerSpec::Pool(8).build(),
source: SchedulerSource::Flag,
width_routing: true,
routing_source: RoutingSource::Config,
};
assert!(on
.rows()
.contains(&("routing".to_string(), "by-width".to_string())));
assert!(on
.rows()
.contains(&("routing.by".to_string(), "config".to_string())));
for (label, _) in on.rows() {
assert!(label.len() <= 10, "`{label}` overflows the row column");
}
}
#[test]
fn width_routing_is_off_by_default_and_the_flag_beats_the_config() {
assert!(
!decide_routing(None, None).0,
"off unless something says so"
);
assert_eq!(decide_routing(None, None).1, RoutingSource::Default);
let from_config = decide_routing(None, Some("on"));
assert!(from_config.0);
assert_eq!(from_config.1, RoutingSource::Config);
let flag_wins = decide_routing(Some(false), Some("on"));
assert!(!flag_wins.0);
assert_eq!(flag_wins.1, RoutingSource::Flag);
}
#[test]
fn an_invalid_width_routing_config_value_warns_and_stays_off() {
let (on, source, warnings) = decide_routing(None, Some("maybe"));
assert!(!on);
assert_eq!(source, RoutingSource::Default);
assert_eq!(warnings.len(), 1);
assert!(warnings[0].contains("maybe"), "{warnings:?}");
}
#[test]
fn the_switch_accepts_on_off_and_the_toml_booleans() {
assert_eq!(parse_switch("on"), Ok(true));
assert_eq!(parse_switch("ON"), Ok(true));
assert_eq!(parse_switch("true"), Ok(true));
assert_eq!(parse_switch("off"), Ok(false));
assert_eq!(parse_switch("false"), Ok(false));
assert!(parse_switch("yes").is_err());
assert!(set_width_routing("yes").is_err());
}
}