use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use nmbrs_metrics::controls::{Control, ControlBuilder, ControlOrigin, ControlRegistry, SetError};
use nmbrs_metrics::labels::Labels;
use nmbrs_rate::{RateLimiter, RateLimiterApplier, RateSpec};
fn build_rate_control(initial: f64) -> Control<RateSpec> {
ControlBuilder::new("rate", RateSpec::new(initial))
.validator(|spec: &RateSpec| {
if spec.ops_per_sec <= 0.0 {
return Err(format!("rate must be > 0 (got {})", spec.ops_per_sec,));
}
if spec.ops_per_sec > 1_000_000.0 {
return Err(format!(
"rate must be <= 1_000_000 (got {})",
spec.ops_per_sec,
));
}
if spec.burst_ratio < 1.0 {
return Err(format!(
"burst_ratio must be >= 1.0 (got {})",
spec.burst_ratio,
));
}
Ok(())
})
.reify_as_gauge(|spec: &RateSpec| Some(spec.ops_per_sec))
.build()
}
#[tokio::test]
async fn full_flow_workload_like_rate_control() {
let registry = ControlRegistry::new();
let control = build_rate_control(500.0);
registry.declare(control.clone());
let limiter = Arc::new(RateLimiter::start(RateSpec::new(500.0)));
control.register_applier(RateLimiterApplier::new(limiter.clone()));
assert_eq!(limiter.rate(), 500.0);
tokio::time::sleep(Duration::from_millis(50)).await;
let ops = Arc::new(AtomicU64::new(0));
let stop = Arc::new(AtomicU64::new(0));
let mut workers = Vec::new();
for _ in 0..4 {
let l = limiter.clone();
let ops = ops.clone();
let stop = stop.clone();
workers.push(tokio::spawn(async move {
while stop.load(Ordering::Relaxed) == 0 {
l.acquire().await;
ops.fetch_add(1, Ordering::Relaxed);
}
}));
}
tokio::time::sleep(Duration::from_millis(120)).await;
let ops_after_slow = ops.load(Ordering::Relaxed);
control
.set(RateSpec::new(20_000.0), ControlOrigin::Test)
.await
.expect("reconfigure to higher rate should succeed");
assert_eq!(limiter.rate(), 20_000.0);
tokio::time::sleep(Duration::from_millis(120)).await;
let ops_after_fast = ops.load(Ordering::Relaxed);
stop.store(1, Ordering::Relaxed);
for w in workers {
let _ = tokio::time::timeout(Duration::from_secs(1), w).await;
}
let slow_rate = ops_after_slow as f64 / 0.120;
let fast_rate = (ops_after_fast - ops_after_slow) as f64 / 0.120;
assert!(
fast_rate > slow_rate * 3.0,
"expected fast-phase throughput to exceed slow by 3x: \
slow={slow_rate:.0} ops/s fast={fast_rate:.0} ops/s",
);
let snap = registry.snapshot_gauges(&Labels::of("phase", "rampup"), Instant::now());
let family = snap
.family("control_rate")
.expect("control_rate gauge family should exist");
let metric = family.metrics().next().unwrap();
assert_eq!(metric.labels().get("control"), Some("rate"));
assert_eq!(metric.labels().get("phase"), Some("rampup"));
match control.set(RateSpec::new(-1.0), ControlOrigin::Test).await {
Err(SetError::ValidationFailed(msg)) => assert!(msg.contains("> 0")),
other => panic!("expected validation rejection, got {other:?}"),
}
assert_eq!(limiter.rate(), 20_000.0);
}
#[tokio::test]
async fn failing_applier_reports_through_control() {
let limiter = Arc::new(RateLimiter::start(RateSpec::new(1_000.0)));
let control = build_rate_control(1_000.0);
control.register_applier(RateLimiterApplier::new(limiter.clone()));
control.register_applier(nmbrs_metrics::controls::SyncApplier::new(|_: RateSpec| {
Err("subscriber B refused".into())
}));
match control
.set(RateSpec::new(5_000.0), ControlOrigin::Test)
.await
{
Err(SetError::ApplyFailed(failures)) => {
assert!(
failures
.iter()
.any(|f| f.message.contains("subscriber B refused"))
);
}
other => panic!("expected ApplyFailed, got {other:?}"),
}
assert_eq!(
control.value().ops_per_sec,
1_000.0,
"control committed value is the source of truth"
);
}