use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, LazyLock, Mutex, RwLock};
use nmbrs_metrics::component::Component;
static SESSION_ROOT: LazyLock<Mutex<Option<Arc<RwLock<Component>>>>> =
LazyLock::new(|| Mutex::new(None));
#[cfg(test)]
pub(crate) fn session_root_test_guard() -> std::sync::MutexGuard<'static, ()> {
static TEST_LOCK: Mutex<()> = Mutex::new(());
TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn set_session_root(root: Arc<RwLock<Component>>) {
*SESSION_ROOT.lock().unwrap_or_else(|e| e.into_inner()) = Some(root);
}
fn session_root() -> Option<Arc<RwLock<Component>>> {
SESSION_ROOT
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
}
pub fn session_root_handle() -> Option<Arc<RwLock<Component>>> {
session_root()
}
pub struct FiberContext {
pub phase: Arc<str>,
pub cycle: AtomicU64,
pub controls: ControlMap,
}
pub type ControlMap =
Arc<std::collections::HashMap<String, Arc<dyn nmbrs_metrics::controls::ErasedControl>>>;
pub fn snapshot_controls(component: &Arc<RwLock<Component>>) -> ControlMap {
Arc::new(nmbrs_metrics::component::Component::control_snapshot(
component,
))
}
pub fn empty_controls() -> ControlMap {
Arc::new(std::collections::HashMap::new())
}
tokio::task_local! {
static FIBER_CTX: FiberContext;
}
pub async fn with_fiber_context<F>(phase: Arc<str>, controls: ControlMap, fut: F) -> F::Output
where
F: Future,
{
FIBER_CTX
.scope(
FiberContext {
phase,
cycle: AtomicU64::new(0),
controls,
},
fut,
)
.await
}
fn current_phase_control(name: &str) -> Option<Arc<dyn nmbrs_metrics::controls::ErasedControl>> {
FIBER_CTX
.try_with(|ctx| ctx.controls.get(name).cloned())
.ok()
.flatten()
}
pub fn set_task_cycle(cycle: u64) {
let _ = FIBER_CTX.try_with(|ctx| ctx.cycle.store(cycle, Ordering::Relaxed));
}
fn task_phase() -> Option<Arc<str>> {
FIBER_CTX.try_with(|ctx| ctx.phase.clone()).ok()
}
fn task_cycle() -> u64 {
FIBER_CTX
.try_with(|ctx| ctx.cycle.load(Ordering::Relaxed))
.unwrap_or(0)
}
fn capture_binding(ctx: &polydat::dsl::factory::BuildContext, name: &str) -> String {
ctx.binding()
.map(str::to_string)
.unwrap_or_else(|| name.to_string())
}
#[polydat::polydat_node(
category = Context,
purity = SideChannel(Other),
)]
fn control_set(
name: Const<&str>,
#[poly_const(capture_binding, from = (ctx, name))] binding: &String,
value: f64,
) -> u64 {
let Some(erased) = resolve_control(name.0) else {
return 0;
};
if erased.gauge_f64() == Some(value) {
return 0;
}
let name = name.0.to_string();
let binding = binding.clone();
tokio::spawn(async move {
let origin = nmbrs_metrics::controls::ControlOrigin::Polydat { binding };
if let Err(e) = erased.set_f64(value, origin).await {
polydat::audit::warn(&format!("control_set({name}, {value}) failed: {e}"));
}
});
1
}
fn resolve_control(name: &str) -> Option<Arc<dyn nmbrs_metrics::controls::ErasedControl>> {
if let Some(handle) = current_phase_control(name) {
return Some(handle);
}
let root = session_root()?;
let guard = root.read().ok()?;
guard.find_control_erased_up(name)
}
fn control_gauge_f64(name: &str) -> f64 {
resolve_control(name)
.and_then(|c| c.gauge_f64())
.unwrap_or(0.0)
}
fn control_value_string(name: &str) -> String {
resolve_control(name)
.map(|c| c.value_string())
.unwrap_or_default()
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads a live dynamic control value; changes over the run"),
)]
fn control(name: Const<&str>) -> f64 {
control_gauge_f64(name.0)
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads a live dynamic control value; changes over the run"),
)]
fn control_u64(name: Const<&str>) -> u64 {
let v = control_gauge_f64(name.0);
if v < 0.0 { 0 } else { v as u64 }
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads a live dynamic control value; changes over the run"),
)]
fn control_bool(name: Const<&str>) -> bool {
control_gauge_f64(name.0) != 0.0
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads a live dynamic control value; changes over the run"),
)]
fn control_str(name: Const<&str>) -> String {
control_value_string(name.0)
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads the live rate control; changes over the run"),
)]
fn rate() -> f64 {
control_gauge_f64("rate")
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads the live concurrency control; changes over the run"),
)]
fn concurrency() -> f64 {
control_gauge_f64("concurrency")
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads the per-fiber phase name; varies across fibers and over the run"),
)]
fn phase() -> String {
task_phase().map(|s| s.to_string()).unwrap_or_default()
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads the current phase's start time; differs per phase and per run"),
)]
fn phase_start_millis() -> u64 {
crate::execution_context::current_phase_start_ms().unwrap_or(0)
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("monotonic elapsed time within the current phase"),
)]
fn phase_elapsed_millis() -> u64 {
let Some(start) = crate::execution_context::current_phase_start_ms() else {
return 0;
};
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0);
now.saturating_sub(start)
}
#[polydat::polydat_node(
category = Context,
purity = Nondeterministic("reads the per-fiber cycle ordinal; changes every cycle"),
)]
fn cycle() -> u64 {
task_cycle()
}
#[cfg(test)]
mod tests {
#![allow(clippy::await_holding_lock)]
use super::*;
use nmbrs_metrics::controls::{BranchScope, ControlBuilder};
use nmbrs_metrics::labels::Labels;
use polydat::ast::Value;
use std::collections::HashMap;
use std::sync::MutexGuard;
fn serial_test() -> MutexGuard<'static, ()> {
super::session_root_test_guard()
}
fn install_session_with_control(name: &str, initial: u32) -> Arc<RwLock<Component>> {
let root = Component::root(Labels::empty().with("session", "t"), HashMap::new());
root.read().unwrap().controls().declare(
ControlBuilder::new(name, initial)
.reify_as_gauge(|v| Some(*v as f64))
.branch_scope(BranchScope::Subtree)
.build(),
);
set_session_root(root.clone());
root
}
#[test]
fn control_reads_current_value() {
let _g = serial_test();
install_session_with_control("rate", 500);
let mut k =
polydat::dsl::compile_polydat_interpreter("x := control(\"rate\")").expect("compile");
assert_eq!(k.pull_ref("x").as_f64(), 500.0);
}
#[test]
fn control_missing_name_returns_zero() {
let _g = serial_test();
install_session_with_control("rate", 500);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control(\"not_declared\")")
.expect("compile");
assert_eq!(k.pull_ref("x").as_f64(), 0.0);
}
#[test]
fn rate_node_is_alias_of_control_rate() {
let _g = serial_test();
install_session_with_control("rate", 750);
let mut k = polydat::dsl::compile_polydat_interpreter("x := rate()").expect("compile");
assert_eq!(k.pull_ref("x").as_f64(), 750.0);
}
#[test]
fn concurrency_node_reads_concurrency_control() {
let _g = serial_test();
install_session_with_control("concurrency", 32);
let mut k =
polydat::dsl::compile_polydat_interpreter("x := concurrency()").expect("compile");
assert_eq!(k.pull_ref("x").as_f64(), 32.0);
}
#[tokio::test]
async fn phase_and_cycle_read_from_task_locals() {
let phase_arc: Arc<str> = Arc::from("rampup");
with_fiber_context(phase_arc.clone(), empty_controls(), async {
set_task_cycle(4242);
let mut k = polydat::dsl::compile_polydat_interpreter("p := phase()\nc := cycle()")
.expect("compile phase/cycle");
assert_eq!(k.pull_ref("p").as_str(), "rampup");
assert_eq!(k.pull_ref("c").as_u64(), 4242);
})
.await;
}
#[test]
fn phase_is_empty_outside_fiber_scope() {
let mut k = polydat::dsl::compile_polydat_interpreter("p := phase()").expect("compile");
assert_eq!(k.pull_ref("p").as_str(), "");
}
#[test]
fn cycle_is_zero_outside_fiber_scope() {
let mut k = polydat::dsl::compile_polydat_interpreter("c := cycle()").expect("compile");
assert_eq!(k.pull_ref("c").as_u64(), 0);
}
#[tokio::test]
async fn set_task_cycle_is_noop_outside_scope() {
set_task_cycle(99);
let mut k = polydat::dsl::compile_polydat_interpreter("c := cycle()").expect("compile");
assert_eq!(k.pull_ref("c").as_u64(), 0);
}
#[tokio::test]
async fn control_set_writes_through_converter_and_reaches_committed() {
let _g = serial_test();
let root = Component::root(
Labels::empty().with("session", "s_cs"),
std::collections::HashMap::new(),
);
let c: nmbrs_metrics::controls::Control<u32> =
nmbrs_metrics::controls::ControlBuilder::new("concurrency", 4u32)
.reify_as_gauge(|v| Some(*v as f64))
.from_f64(|v| {
if v < 0.0 || v > u32::MAX as f64 {
Err(format!("out of range: {v}"))
} else {
Ok(v as u32)
}
})
.branch_scope(nmbrs_metrics::controls::BranchScope::Subtree)
.build();
root.read().unwrap().controls().declare(c.clone());
set_session_root(root);
let ctx = polydat::dsl::factory::BuildContext::with_binding("feedback_loop");
let consts = [polydat::dsl::factory::ConstArg::Str("concurrency".into())];
let node = polydat::dsl::factory::build_node(&ctx, "control_set", &[], &[], &consts)
.expect("control_set should build");
let mut out = [Value::None];
node.eval(&[Value::F64(64.0)], &mut out);
assert_eq!(out[0].as_u64(), 1, "write should report submitted");
for _ in 0..10 {
tokio::task::yield_now().await;
if c.value() == 64u32 {
break;
}
}
assert_eq!(c.value(), 64u32);
let committed = c.get();
assert!(matches!(
committed.origin,
nmbrs_metrics::controls::ControlOrigin::Polydat { .. }
));
let rev_before = c.get().rev;
node.eval(&[Value::F64(64.0)], &mut out);
assert_eq!(
out[0].as_u64(),
0,
"write of the committed value must be elided"
);
assert_eq!(
c.get().rev,
rev_before,
"an elided write must not touch the control"
);
node.eval(&[Value::F64(32.0)], &mut out);
assert_eq!(out[0].as_u64(), 1, "a real change dispatches");
for _ in 0..10 {
tokio::task::yield_now().await;
if c.value() == 32u32 {
break;
}
}
assert_eq!(c.value(), 32u32);
}
#[test]
fn control_u64_casts_gauge_to_integer() {
let _g = serial_test();
install_session_with_control("concurrency", 64);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_u64(\"concurrency\")")
.expect("compile control_u64");
assert_eq!(k.pull_ref("x").as_u64(), 64);
}
#[test]
fn control_u64_missing_name_returns_zero() {
let _g = serial_test();
install_session_with_control("concurrency", 5);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_u64(\"not_there\")")
.expect("compile control_u64");
assert_eq!(k.pull_ref("x").as_u64(), 0);
}
#[test]
fn control_u64_is_volatile_not_const_folded() {
let _g = serial_test();
install_session_with_control("concurrency", 32);
let k = polydat::dsl::compile_polydat_interpreter("x := control_u64(\"concurrency\")")
.expect("compile control_u64");
assert!(
k.get_constant("x").is_none(),
"control_u64 must be volatile — its wire must NOT be const-folded",
);
}
#[test]
fn control_bool_projects_gauge_to_boolean() {
let _g = serial_test();
install_session_with_control("enabled", 1);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_bool(\"enabled\")")
.expect("compile");
assert!(k.pull_ref("x").as_bool());
}
#[test]
fn control_bool_zero_is_false() {
let _g = serial_test();
install_session_with_control("enabled", 0);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_bool(\"enabled\")")
.expect("compile");
assert!(!k.pull_ref("x").as_bool());
}
#[test]
fn control_bool_missing_name_is_false() {
let _g = serial_test();
install_session_with_control("enabled", 1);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_bool(\"absent\")")
.expect("compile");
assert!(!k.pull_ref("x").as_bool());
}
#[test]
fn control_str_renders_value_string() {
let _g = serial_test();
install_session_with_control("concurrency", 42);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_str(\"concurrency\")")
.expect("compile");
assert_eq!(k.pull_ref("x").as_str(), "42");
}
#[test]
fn control_str_missing_name_returns_empty() {
let _g = serial_test();
install_session_with_control("concurrency", 42);
let mut k = polydat::dsl::compile_polydat_interpreter("x := control_str(\"log_level\")")
.expect("compile");
assert_eq!(k.pull_ref("x").as_str(), "");
}
#[tokio::test]
async fn control_set_records_compile_time_binding_attribution() {
let _g = serial_test();
let root = Component::root(
Labels::empty().with("session", "attr"),
std::collections::HashMap::new(),
);
let c: nmbrs_metrics::controls::Control<f64> =
nmbrs_metrics::controls::ControlBuilder::new("rate", 100.0)
.reify_as_gauge(|v| Some(*v))
.from_f64(Ok)
.branch_scope(nmbrs_metrics::controls::BranchScope::Subtree)
.build();
root.read().unwrap().controls().declare(c.clone());
set_session_root(root.clone());
let ctx = polydat::dsl::factory::BuildContext::with_binding("rate_adj");
let consts = [polydat::dsl::factory::ConstArg::Str("rate".into())];
let node = polydat::dsl::factory::build_node(&ctx, "control_set", &[], &[], &consts)
.expect("control_set should build");
let mut out = [Value::None];
node.eval(&[Value::F64(4242.0)], &mut out);
for _ in 0..40 {
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
if c.value() == 4242.0 {
break;
}
}
assert_eq!(c.value(), 4242.0);
match c.get().origin {
nmbrs_metrics::controls::ControlOrigin::Polydat { ref binding } => {
assert_eq!(
binding, "rate_adj",
"attribution should be the DSL binding name, not the control name"
);
}
other => panic!("expected Polydat origin, got {other:?}"),
}
}
#[test]
fn control_set_returns_zero_without_session_root() {
let _g = serial_test();
*SESSION_ROOT.lock().unwrap_or_else(|e| e.into_inner()) = None;
let consts = [polydat::dsl::factory::ConstArg::Str("anything".into())];
let node = polydat::dsl::factory::build_node(
&polydat::dsl::factory::BuildContext::default(),
"control_set",
&[],
&[],
&consts,
)
.expect("control_set should build");
let mut out = [Value::None];
node.eval(&[Value::F64(1.0)], &mut out);
assert_eq!(out[0].as_u64(), 0);
}
fn component_with_concurrency(val: u32) -> Arc<RwLock<Component>> {
let comp = Component::root(Labels::empty().with("phase", "p"), HashMap::new());
comp.read()
.unwrap_or_else(|e| e.into_inner())
.controls()
.declare(
ControlBuilder::new("concurrency", val)
.reify_as_gauge(|v| Some(*v as f64))
.branch_scope(BranchScope::Subtree)
.build(),
);
comp
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn per_execution_control_map_isolates_concurrent_reads() {
let _g = serial_test();
*SESSION_ROOT.lock().unwrap_or_else(|e| e.into_inner()) = None;
let comp_a = component_with_concurrency(2);
let comp_b = component_with_concurrency(32);
let phase: Arc<str> = Arc::from("p");
let a_val = with_fiber_context(phase.clone(), snapshot_controls(&comp_a), async {
control_gauge_f64("concurrency")
})
.await;
let b_val = with_fiber_context(phase.clone(), snapshot_controls(&comp_b), async {
control_gauge_f64("concurrency")
})
.await;
assert_eq!(a_val, 2.0, "execution A must read its OWN concurrency (2)");
assert_eq!(
b_val, 32.0,
"execution B must read its OWN concurrency (32), not A's",
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn control_read_falls_back_to_session_root_without_exec_context() {
let _g = serial_test();
install_session_with_control("concurrency", 7);
assert_eq!(control_gauge_f64("concurrency"), 7.0);
}
}