// Copyright 2026 PingCAP, Inc.

use std::{
    sync::{
        Arc, Mutex, RwLock,
        atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};

use dashmap::{DashMap, DashSet};
use prometheus::IntGauge;
use tikv_util::{info, warn};

use crate::{
    cpu_config::CpuThrottleConfig,
    metrics::{
        CPU_THROTTLE_ALLOCATIONS, CPU_THROTTLE_BURST_COUNT, CPU_THROTTLE_BURST_TOKENS_USED,
        CPU_THROTTLE_EFFECTIVE_MAX_READ_CPU_RATIO, CPU_THROTTLE_FAIR_ALLOCATIONS,
        CPU_THROTTLE_FAIR_CHARGE_SCALE, CPU_THROTTLE_GLOBAL_BUCKET_AVAILABLE,
        CPU_THROTTLE_GLOBAL_BUCKET_CAPACITY, CPU_THROTTLE_GLOBAL_BUCKET_DEBT_US,
        CPU_THROTTLE_GLOBAL_REFILL_RATE, CPU_THROTTLE_KEYSPACE_BUCKET_AVAILABLE,
        CPU_THROTTLE_KEYSPACE_BUCKET_CAPACITY, CPU_THROTTLE_KEYSPACE_DEBT_US,
        CPU_THROTTLE_KEYSPACE_INFLIGHT_US, CPU_THROTTLE_KEYSPACE_REFILL_RATE,
        CPU_THROTTLE_LOW_PRIORITY_DAG_ALLOCATIONS, CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE,
        CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_CAPACITY, CPU_THROTTLE_LOW_PRIORITY_DAG_REFILL_RATE,
        CPU_THROTTLE_REFILL_RATE_ADJUSTMENTS, CPU_THROTTLE_REQUESTED_MAX_READ_CPU_RATIO,
        CPU_THROTTLE_TENANT_CAP_QUOTA_US, CPU_THROTTLE_TENANT_CAP_RATIO,
        CPU_TOKEN_RUNTIME_ALLOCATIONS, CPU_TOKEN_RUNTIME_WAIT_DURATION_HISTOGRAM, keyspace_label,
        remove_keyspace_label,
    },
};

/// Error types for CPU throttling
#[derive(Debug, Clone, PartialEq)]
pub enum ThrottleError {
    /// Global CPU token bucket exhausted
    GlobalCpuExhausted,
    /// Low-priority DAG global CPU token bucket exhausted
    LowPriorityDagGlobalCpuExhausted,
    /// Keyspace CPU token bucket exhausted
    KeyspaceCpuExhausted,
    /// Keyspace exceeded fair share allocation.
    KeyspaceFairShareExceeded {
        /// Earliest time either the keyspace or fair-share path can recover.
        retry_at: Instant,
    },
    /// Keyspace CPU exhausted and burst was not currently available.
    KeyspaceCpuExhaustedNoBurst {
        /// Earliest time either the keyspace or relevant burst path can
        /// recover.
        retry_at: Instant,
    },
    /// Request timeout while waiting for token refill
    RequestTimeout,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BurstDenyReason {
    BurstDisabled,
    KeyspaceBurstDisabled,
    GlobalRatioExceeded,
    GlobalBucketEmpty,
    BurstBudgetExhausted,
    IntervalLimitExceeded,
}

impl BurstDenyReason {
    fn as_str(self) -> &'static str {
        match self {
            BurstDenyReason::BurstDisabled => "burst_disabled",
            BurstDenyReason::KeyspaceBurstDisabled => "keyspace_burst_disabled",
            BurstDenyReason::GlobalRatioExceeded => "global_ratio_exceeded",
            BurstDenyReason::GlobalBucketEmpty => "global_bucket_empty",
            BurstDenyReason::BurstBudgetExhausted => "burst_budget_exhausted",
            BurstDenyReason::IntervalLimitExceeded => "interval_limit_exceeded",
        }
    }
}

fn clamp_gauge_value(value: u64) -> i64 {
    if value > i64::MAX as u64 {
        i64::MAX
    } else {
        value as i64
    }
}

fn token_amount_to_i64(tokens: u64) -> i64 {
    tokens.min(i64::MAX as u64) as i64
}

fn token_amount_to_i128(tokens: u64) -> i128 {
    i128::from(token_amount_to_i64(tokens))
}

fn clamp_i128_to_i64(value: i128) -> i64 {
    value.clamp(i128::from(i64::MIN), i128::from(i64::MAX)) as i64
}

fn token_balance_to_available(token_balance: i64) -> u64 {
    token_balance.max(0) as u64
}

fn token_balance_to_debt(token_balance: i64) -> u64 {
    if token_balance < 0 {
        (-(i128::from(token_balance))) as u64
    } else {
        0
    }
}

fn atomic_saturating_add(value: &AtomicU64, delta: u64) -> u64 {
    let mut current = value.load(Ordering::Relaxed);
    loop {
        let new_value = current.saturating_add(delta);
        match value.compare_exchange_weak(current, new_value, Ordering::Relaxed, Ordering::Relaxed)
        {
            Ok(_) => return new_value,
            Err(actual) => current = actual,
        }
    }
}

fn burst_cap_us(capacity_us: u64, burst_threshold: f64) -> u64 {
    if capacity_us == u64::MAX {
        return u64::MAX;
    }
    if !(0.0..=1.0).contains(&burst_threshold) {
        return capacity_us;
    }
    (capacity_us as f64 * burst_threshold) as u64
}

fn reserve_burst_budget(
    tracker: &BurstUsageTracker,
    capacity_us: u64,
    burst_threshold: f64,
    amount_us: u64,
) -> Result<Option<u64>, ()> {
    let cap_us = burst_cap_us(capacity_us, burst_threshold);
    if cap_us == u64::MAX {
        return Ok(None);
    }
    let window_id = tracker.current_window_id();
    tracker
        .try_reserve_global(window_id, amount_us, cap_us)
        .map(Some)
        .ok_or(())
}

fn throttle_error_label(err: &ThrottleError) -> &'static str {
    match err {
        ThrottleError::GlobalCpuExhausted => "global_exhausted",
        ThrottleError::LowPriorityDagGlobalCpuExhausted => "low_priority_global_exhausted",
        ThrottleError::KeyspaceCpuExhausted => "keyspace_exhausted",
        ThrottleError::KeyspaceFairShareExceeded { .. } => "fair_share_exceeded",
        ThrottleError::KeyspaceCpuExhaustedNoBurst { .. } => "keyspace_exhausted_no_burst",
        ThrottleError::RequestTimeout => "request_timeout",
    }
}

/// Fixed-point unit for the fair-gate charge scale. Using fixed point keeps
/// the per-poll charge path free of floating-point work.
pub(crate) const FAIR_CHARGE_SCALE_ONE: u64 = 1 << 16;

/// Convert a base/effective ratio pair into a fixed-point charge scale.
///
/// `effective_ratio <= 0` means the override is disabled and returns
/// [`FAIR_CHARGE_SCALE_ONE`]. Otherwise the scale is
/// `base_ratio / effective_ratio`: a lower effective ratio yields a scale above
/// 1.0, so the same measured CPU consumes proportionally more tokens.
pub(crate) fn fair_charge_scale_from_ratios(base_ratio: f64, effective_ratio: f64) -> u64 {
    if base_ratio <= 0.0 || effective_ratio <= 0.0 || base_ratio == effective_ratio {
        return FAIR_CHARGE_SCALE_ONE;
    }
    let scale = (base_ratio / effective_ratio) * FAIR_CHARGE_SCALE_ONE as f64;
    if !scale.is_finite() || scale <= 0.0 {
        return FAIR_CHARGE_SCALE_ONE;
    }
    (scale.round() as u64).max(1)
}

/// Bits reserved for the carried remainder in the packed charge state. The
/// remainder is always less than [`FAIR_CHARGE_SCALE_ONE`], so 16 bits suffice.
const FAIR_CHARGE_RESIDUE_BITS: u32 = 16;
const FAIR_CHARGE_RESIDUE_MASK: u64 = (1 << FAIR_CHARGE_RESIDUE_BITS) - 1;

/// Pack a scale and its carried remainder into one word.
///
/// Keeping them in a single atomic makes an exchange-rate change and the
/// remainder that belongs to it a single transition. A charge cannot then read
/// one scale, be preempted, and commit a remainder that belongs to another.
const fn pack_fair_charge_state(scale: u64, residue_q: u64) -> u64 {
    (scale << FAIR_CHARGE_RESIDUE_BITS) | (residue_q & FAIR_CHARGE_RESIDUE_MASK)
}

const fn unpack_fair_charge_state(state: u64) -> (u64, u64) {
    (
        state >> FAIR_CHARGE_RESIDUE_BITS,
        state & FAIR_CHARGE_RESIDUE_MASK,
    )
}

/// Apply the current charge scale to a measured CPU amount, carrying the
/// fractional part forward.
///
/// Rounding each delta independently is wrong here. The fair gate charges at
/// every poll checkpoint, so deltas are often a handful of microseconds, and a
/// per-delta rounding error is a fixed fraction of a tiny number:
///
/// - scale 0.667, delta 1us: rounds to 1 token, so a stream of 1us deltas is
///   charged 1:1 and the override is silently ignored.
/// - scale below 0.5, delta 1us: rounds to 0 tokens, so a stream of short
///   dispatches consumes nothing at all.
///
/// Instead accumulate the exact scaled amount in fixed point, charge only whole
/// tokens, and keep the sub-token remainder for the next call. The total
/// charged then tracks the exact product to within one token regardless of how
/// the CPU time is split across calls.
///
/// The scale is read from the same word as the remainder it is combined with,
/// so a concurrent scale change either loses this CAS (and the charge retries
/// under the new scale) or happens entirely after it. A remainder is therefore
/// never applied to a scale it was not produced under.
fn accumulate_fair_charge(cpu_delta_us: u64, state: &AtomicU64) -> u64 {
    let one = u128::from(FAIR_CHARGE_SCALE_ONE);

    loop {
        let packed = state.load(Ordering::Acquire);
        let (scale, residue) = unpack_fair_charge_state(packed);
        if scale == FAIR_CHARGE_SCALE_ONE {
            // No override in force. There is no remainder to keep, so skip the
            // CAS entirely and leave the state untouched.
            return cpu_delta_us;
        }

        let total_q = u128::from(cpu_delta_us) * u128::from(scale) + u128::from(residue);
        let tokens = total_q / one;
        // The remainder is always < FAIR_CHARGE_SCALE_ONE, so it fits the field.
        let new_residue = (total_q % one) as u64;
        if state
            .compare_exchange(
                packed,
                pack_fair_charge_state(scale, new_residue),
                Ordering::AcqRel,
                Ordering::Acquire,
            )
            .is_ok()
        {
            return tokens.min(u128::from(u64::MAX)) as u64;
        }
    }
}

fn low_priority_dag_quota_us(global_cpu_us: u64, ratio: f64, enabled: bool) -> u64 {
    if !enabled || ratio <= 0.0 {
        return 0;
    }
    (global_cpu_us as f64 * ratio).floor() as u64
}

fn refill_rate_to_capacity(
    refill_rate_us: u64,
    refill_interval_ms: u64,
    burst_window_ms: u64,
) -> u64 {
    if refill_rate_us == u64::MAX {
        return u64::MAX;
    }
    let capacity = u128::from(refill_rate_us).saturating_mul(u128::from(burst_window_ms))
        / u128::from(refill_interval_ms.max(1));
    capacity.min(u128::from(u64::MAX)) as u64
}

const TENANT_CAP_KEYSPACE_CLEANUP_INTERVAL_SECS: u64 = 60;

fn scale_refill_rate_by_reference(
    base_rate_us: u64,
    reference_current_rate_us: u64,
    reference_base_rate_us: u64,
) -> u64 {
    if reference_base_rate_us == 0 {
        return base_rate_us;
    }

    let scaled = (base_rate_us as u128).saturating_mul(reference_current_rate_us as u128)
        / reference_base_rate_us as u128;
    scaled.min(u64::MAX as u128) as u64
}

#[derive(Debug)]
pub(crate) struct KeyspaceInflightState {
    allocated_us: AtomicU64,
    bucket_available: IntGauge,
    inflight: IntGauge,
    debt: IntGauge,
}

fn update_global_bucket_metrics(bucket: &CpuTokenBucket) {
    CPU_THROTTLE_GLOBAL_BUCKET_AVAILABLE
        .set(clamp_gauge_value(bucket.metric_cpu_us(bucket.available())));
    CPU_THROTTLE_GLOBAL_BUCKET_DEBT_US.set(clamp_gauge_value(bucket.metric_cpu_us(bucket.debt())));
}

fn update_global_bucket_limit_metrics(bucket: &CpuTokenBucket) {
    CPU_THROTTLE_GLOBAL_BUCKET_CAPACITY
        .set(clamp_gauge_value(bucket.metric_cpu_us(bucket.capacity())));
    CPU_THROTTLE_GLOBAL_REFILL_RATE.set(clamp_gauge_value(
        bucket.metric_cpu_us(bucket.current_refill_rate()),
    ));
}

impl KeyspaceInflightState {
    fn new(keyspace_id: u32) -> Self {
        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_label = keyspace_label.as_ref();
        Self {
            allocated_us: AtomicU64::new(0),
            bucket_available: CPU_THROTTLE_KEYSPACE_BUCKET_AVAILABLE
                .with_label_values(&[keyspace_label]),
            inflight: CPU_THROTTLE_KEYSPACE_INFLIGHT_US.with_label_values(&[keyspace_label]),
            debt: CPU_THROTTLE_KEYSPACE_DEBT_US.with_label_values(&[keyspace_label]),
        }
    }

    fn add(&self, delta_us: u64) -> u64 {
        self.allocated_us.fetch_add(delta_us, Ordering::AcqRel) + delta_us
    }

    fn sub(&self, delta_us: u64) -> u64 {
        let mut current = self.allocated_us.load(Ordering::Acquire);
        loop {
            let new_value = current.saturating_sub(delta_us);
            match self.allocated_us.compare_exchange_weak(
                current,
                new_value,
                Ordering::AcqRel,
                Ordering::Relaxed,
            ) {
                Ok(_) => return new_value,
                Err(actual) => current = actual,
            }
        }
    }

    fn update_inflight_metrics(&self, inflight_us: u64, capacity_us: u64) {
        self.inflight.set(clamp_gauge_value(inflight_us));
        self.debt
            .set(clamp_gauge_value(inflight_us.saturating_sub(capacity_us)));
    }
}

fn atomic_saturating_sub(value: &AtomicU64, delta_us: u64) -> u64 {
    loop {
        let current = value.load(Ordering::Acquire);
        let new_value = current.saturating_sub(delta_us);
        if value
            .compare_exchange(current, new_value, Ordering::AcqRel, Ordering::Relaxed)
            .is_ok()
        {
            return new_value;
        }
    }
}

const ADAPTIVE_ESTIMATE_ALPHA: f64 = 0.2;
const ADAPTIVE_ESTIMATE_LOG_INTERVAL_MS: u64 = 5000;

#[derive(Debug)]
struct AdaptiveEstimateEntry {
    estimate_us: AtomicU64,
    window_id: AtomicU64,
    sum_us: AtomicU64,
    count: AtomicU64,
    last_log_ms: AtomicU64,
}

impl AdaptiveEstimateEntry {
    fn new(estimate_us: u64, window_id: u64) -> Self {
        Self {
            estimate_us: AtomicU64::new(estimate_us.max(1)),
            window_id: AtomicU64::new(window_id),
            sum_us: AtomicU64::new(0),
            count: AtomicU64::new(0),
            last_log_ms: AtomicU64::new(0),
        }
    }
}

#[derive(Debug)]
pub(crate) struct AdaptiveEstimator {
    alpha: f64,
    default_estimate_us: u64,
    interval_ms: u64,
    start_instant: Instant,
    estimates: DashMap<u32, Arc<AdaptiveEstimateEntry>>,
}

impl AdaptiveEstimator {
    fn new(
        default_estimate_us: u64,
        interval_ms: u64,
        initial: std::collections::HashMap<u32, u64>,
    ) -> Self {
        let estimates = DashMap::new();
        let interval_ms = interval_ms.max(1);
        let start_instant = Instant::now();
        for (keyspace_id, estimate_us) in initial {
            estimates.insert(
                keyspace_id,
                Arc::new(AdaptiveEstimateEntry::new(estimate_us, 0)),
            );
        }
        AdaptiveEstimator {
            alpha: ADAPTIVE_ESTIMATE_ALPHA,
            default_estimate_us: default_estimate_us.max(1),
            interval_ms,
            start_instant,
            estimates,
        }
    }

    fn get(&self, keyspace_id: u32) -> u64 {
        self.estimates
            .get(&keyspace_id)
            .map(|v| v.value().estimate_us.load(Ordering::Relaxed))
            .unwrap_or(self.default_estimate_us)
    }

    fn current_window_id(&self) -> u64 {
        let elapsed_ms = self.start_instant.elapsed().as_millis() as u64;
        elapsed_ms / self.interval_ms
    }

    fn observe(&self, keyspace_id: u32, actual_cpu_us: u64) {
        if actual_cpu_us == 0 {
            return;
        }
        let window_id = self.current_window_id();
        let entry = if let Some(entry) = self.estimates.get(&keyspace_id) {
            Arc::clone(entry.value())
        } else {
            let entry = self.estimates.entry(keyspace_id).or_insert_with(|| {
                Arc::new(AdaptiveEstimateEntry::new(
                    self.default_estimate_us,
                    window_id,
                ))
            });
            Arc::clone(entry.value())
        };
        loop {
            let observed_window_id = entry.window_id.load(Ordering::Acquire);
            if observed_window_id >= window_id {
                break;
            }
            if entry
                .window_id
                .compare_exchange(
                    observed_window_id,
                    window_id,
                    Ordering::AcqRel,
                    Ordering::Acquire,
                )
                .is_err()
            {
                continue;
            }

            let sum_us = entry.sum_us.swap(0, Ordering::AcqRel);
            let count = entry.count.swap(0, Ordering::AcqRel);
            if let Some(avg) = sum_us.checked_div(count) {
                let old_estimate_us = entry.estimate_us.load(Ordering::Relaxed);
                let new_estimate_us =
                    (old_estimate_us as f64 * (1.0 - self.alpha) + avg as f64 * self.alpha) as u64;
                let new_estimate_us = new_estimate_us.max(1);
                entry.estimate_us.store(new_estimate_us, Ordering::Relaxed);

                let elapsed_ms = self.start_instant.elapsed().as_millis() as u64;
                let last_log_ms = entry.last_log_ms.load(Ordering::Relaxed);
                if elapsed_ms.saturating_sub(last_log_ms) >= ADAPTIVE_ESTIMATE_LOG_INTERVAL_MS
                    && entry
                        .last_log_ms
                        .compare_exchange(
                            last_log_ms,
                            elapsed_ms,
                            Ordering::Relaxed,
                            Ordering::Relaxed,
                        )
                        .is_ok()
                {
                    info!(
                        "[CPU throttle] Adaptive estimate updated";
                        "keyspace_id" => keyspace_id,
                        "window_id" => observed_window_id,
                        "avg_cpu_us" => avg,
                        "old_estimate_us" => old_estimate_us,
                        "new_estimate_us" => new_estimate_us,
                    );
                }
            }
            break;
        }

        // These statistics are advisory. Saturating CAS loops preserve the old
        // overflow behavior without serializing every request on a DashMap write
        // lock.
        atomic_saturating_add(&entry.sum_us, actual_cpu_us);
        atomic_saturating_add(&entry.count, 1);
    }
}

/// CPU usage monitor data
#[derive(Debug)]
pub struct CpuUsageMonitor {
    /// Global CPU usage ratio (0.0 - 1.0)
    pub global_usage_ratio: f64,
    /// Keyspace CPU usage ratios
    pub keyspace_usage_ratios: std::collections::HashMap<u32, f64>,
    /// Last update time
    pub last_update: Instant,
    /// Last log time for rate limiting
    pub last_log_time: Instant,
}

impl Default for CpuUsageMonitor {
    fn default() -> Self {
        let now = Instant::now();
        CpuUsageMonitor {
            global_usage_ratio: 0.0,
            keyspace_usage_ratios: std::collections::HashMap::new(),
            last_update: now,
            last_log_time: now,
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct BurstUsage {
    window_id: u64,
    used_us: u64,
}

#[derive(Debug)]
pub(crate) struct BurstUsageTracker {
    start_instant: Instant,
    interval_ms: u64,
    usages: DashMap<u32, BurstUsage>,
    global_window: Mutex<GlobalBurstWindow>,
}

#[derive(Debug)]
struct GlobalBurstWindow {
    window_id: u64,
    used_us: u64,
}

impl BurstUsageTracker {
    fn new(interval_ms: u64) -> Self {
        BurstUsageTracker {
            start_instant: Instant::now(),
            interval_ms: interval_ms.max(1),
            usages: DashMap::new(),
            global_window: Mutex::new(GlobalBurstWindow {
                window_id: 0,
                used_us: 0,
            }),
        }
    }

    fn current_window_id(&self) -> u64 {
        let elapsed_ms = self.start_instant.elapsed().as_millis() as u64;
        elapsed_ms / self.interval_ms
    }

    fn next_window_deadline(&self) -> Instant {
        let next_window = self.current_window_id().saturating_add(1);
        self.start_instant
            .checked_add(Duration::from_millis(
                next_window.saturating_mul(self.interval_ms),
            ))
            .unwrap_or_else(Instant::now)
    }

    fn usage_in_window(&self, keyspace_id: u32, window_id: u64) -> u64 {
        let mut entry = self.usages.entry(keyspace_id).or_insert(BurstUsage {
            window_id,
            used_us: 0,
        });
        if entry.window_id < window_id {
            entry.window_id = window_id;
            entry.used_us = 0;
        } else if entry.window_id > window_id {
            // A delayed request from an older window must not reset usage that
            // has already been recorded in a newer window.
            return 0;
        }
        entry.used_us
    }

    fn add_usage_in_window(&self, keyspace_id: u32, window_id: u64, amount_us: u64) {
        let mut entry = self.usages.entry(keyspace_id).or_insert(BurstUsage {
            window_id,
            used_us: 0,
        });
        if entry.window_id < window_id {
            entry.window_id = window_id;
            entry.used_us = 0;
        }
        if entry.window_id == window_id {
            entry.used_us = entry.used_us.saturating_add(amount_us);
        }
    }

    fn release_usage_in_window(&self, keyspace_id: u32, window_id: u64, amount_us: u64) {
        if let Some(mut entry) = self.usages.get_mut(&keyspace_id) {
            if entry.window_id == window_id {
                entry.used_us = entry.used_us.saturating_sub(amount_us);
            }
        }
    }

    fn try_reserve_global(&self, window_id: u64, amount_us: u64, cap_us: u64) -> Option<u64> {
        if cap_us == u64::MAX {
            return Some(window_id);
        }
        let mut global = self.global_window.lock().unwrap();
        // A request can sample its window ID, get descheduled, and acquire the
        // lock after another request has advanced the tracker. Returning and
        // charging the effective (newer) window fixes the race where a stale
        // request could reset current usage or later release tokens from the
        // wrong window.
        let effective_window_id = window_id
            .max(self.current_window_id())
            .max(global.window_id);
        if global.window_id < effective_window_id {
            global.window_id = effective_window_id;
            global.used_us = 0;
        }
        if global.used_us.saturating_add(amount_us) > cap_us {
            return None;
        }
        global.used_us = global.used_us.saturating_add(amount_us);
        Some(effective_window_id)
    }

    fn release_global(&self, window_id: u64, amount_us: u64) {
        let mut global = self.global_window.lock().unwrap();
        if global.window_id != window_id {
            return;
        }
        global.used_us = global.used_us.saturating_sub(amount_us);
    }

    fn remove_keyspace(&self, keyspace_id: u32) {
        self.usages.remove(&keyspace_id);
    }
}

#[derive(Debug, Clone)]
pub(crate) struct BurstUsageToken {
    keyspace_id: u32,
    window_id: u64,
    allocated_us: u64,
    tracker: Arc<BurstUsageTracker>,
}

#[derive(Debug)]
struct IntervalAllocationWindow {
    window_id: u64,
    allocated_us: u64,
}

/// Token bucket for CPU throttling
#[derive(Debug)]
pub struct CpuTokenBucket {
    /// Bucket capacity in microseconds
    capacity_us: AtomicU64,
    /// Signed token balance. Positive values are available tokens and negative
    /// values are debt.
    token_balance: AtomicI64,
    /// Base refill rate (tokens per refill interval)
    base_refill_rate_us: AtomicU64,
    /// Current refill rate (tokens per refill interval), can be adjusted
    /// dynamically
    current_refill_rate_us: AtomicU64,
    /// CPU microseconds represented by one token, in fixed point, used only to
    /// preserve the `_us` contract of global bucket metrics. Non-global
    /// buckets leave this at 1.0.
    metric_cpu_per_token_q: AtomicU64,
    /// Refill interval
    refill_interval: Duration,
    /// Last refill timestamp, as elapsed milliseconds since `start_instant`.
    ///
    /// `refill` is also called by fair read-pool admission before every pop
    /// attempt, so keep the no-refill fast path lock-free.
    last_refill_elapsed_ms: AtomicU64,
    /// Whether this bucket is enabled
    enabled: AtomicBool,
    /// Window ID and allocated tokens for per-interval burst limiting. These
    /// values must transition atomically at a window boundary.
    interval_allocation: Mutex<IntervalAllocationWindow>,
    /// Allocation limit per interval (0 = disabled)
    interval_allocation_limit: AtomicU64,
    /// Creation instant for window ID calculation
    start_instant: Instant,
}

impl CpuTokenBucket {
    /// Create a new token bucket
    ///
    /// # Arguments
    /// * `capacity_us` - Bucket capacity in microseconds
    /// * `refill_rate_us` - Refill rate in microseconds per interval
    /// * `refill_interval` - Refill interval duration
    /// * `interval_allocation_limit` - Per-interval allocation limit (0 =
    ///   disabled)
    pub fn new(
        capacity_us: u64,
        refill_rate_us: u64,
        refill_interval: Duration,
        interval_allocation_limit: u64,
    ) -> Self {
        let start_instant = Instant::now();
        CpuTokenBucket {
            capacity_us: AtomicU64::new(capacity_us),
            token_balance: AtomicI64::new(token_amount_to_i64(capacity_us)),
            base_refill_rate_us: AtomicU64::new(refill_rate_us),
            current_refill_rate_us: AtomicU64::new(refill_rate_us),
            metric_cpu_per_token_q: AtomicU64::new(FAIR_CHARGE_SCALE_ONE),
            refill_interval,
            last_refill_elapsed_ms: AtomicU64::new(0),
            enabled: AtomicBool::new(true),
            interval_allocation: Mutex::new(IntervalAllocationWindow {
                window_id: 0,
                allocated_us: 0,
            }),
            interval_allocation_limit: AtomicU64::new(interval_allocation_limit),
            start_instant,
        }
    }

    fn elapsed_since_start_ms(&self) -> u64 {
        let elapsed_ms = self.start_instant.elapsed().as_millis();
        elapsed_ms.min(u128::from(u64::MAX)) as u64
    }

    fn capacity_as_i64(&self) -> i64 {
        token_amount_to_i64(self.capacity_us.load(Ordering::Acquire))
    }

    fn set_metric_charge_scale(&self, scale_q: u64) {
        let scale_q = scale_q.max(1);
        let numerator =
            u128::from(FAIR_CHARGE_SCALE_ONE).saturating_mul(u128::from(FAIR_CHARGE_SCALE_ONE));
        let cpu_per_token_q = numerator / u128::from(scale_q);
        self.metric_cpu_per_token_q.store(
            cpu_per_token_q.min(u128::from(u64::MAX)) as u64,
            Ordering::Release,
        );
    }

    fn metric_cpu_us(&self, tokens: u64) -> u64 {
        let cpu_per_token_q = self.metric_cpu_per_token_q.load(Ordering::Acquire);
        let cpu_us = u128::from(tokens).saturating_mul(u128::from(cpu_per_token_q))
            >> FAIR_CHARGE_RESIDUE_BITS;
        cpu_us.min(u128::from(u64::MAX)) as u64
    }

    /// Try to allocate tokens from the bucket
    ///
    /// # Arguments
    /// * `tokens` - Number of tokens to allocate
    ///
    /// # Returns
    /// `true` if allocation succeeded, `false` otherwise
    pub fn try_allocate(&self, tokens: u64) -> bool {
        if !self.enabled.load(Ordering::Relaxed) {
            return true;
        }
        if tokens == 0 {
            return true;
        }

        let tokens = token_amount_to_i64(tokens);

        loop {
            let current = self.token_balance.load(Ordering::Acquire);
            if current < tokens {
                return false;
            }

            if self
                .token_balance
                .compare_exchange(
                    current,
                    current - tokens,
                    Ordering::AcqRel,
                    Ordering::Acquire,
                )
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Allocate tokens and allow the bucket to go into debt.
    ///
    /// This is intentionally separate from [`CpuTokenBucket::try_allocate`].
    /// Existing CPU throttle paths keep their hard-fail semantics, while the
    /// fair read-pool path charges measured CPU at executor checkpoints and
    /// allows bounded debt before yielding.
    pub fn allocate_with_debt(&self, tokens: u64) {
        if tokens == 0 || !self.enabled.load(Ordering::Relaxed) {
            return;
        }

        let tokens = token_amount_to_i64(tokens);

        let _ = self
            .token_balance
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
                Some(current.saturating_sub(tokens))
            });
    }

    /// Clear fair-gate debt after a dry-run sample.
    ///
    /// Production admission never resets a bucket this way. The manager only
    /// calls it from `reset_fair_global_cpu_after_dry_run` after the fair gate
    /// has explicitly registered its dry-run reset hook.
    fn reset_to_capacity(&self) {
        let now_ms = self.elapsed_since_start_ms();
        let interval_ms = self.refill_interval.as_millis().max(1) as u64;
        self.last_refill_elapsed_ms.store(now_ms, Ordering::Release);
        let mut interval = self.interval_allocation.lock().unwrap();
        interval.window_id = now_ms / interval_ms;
        interval.allocated_us = 0;
        self.token_balance
            .store(self.capacity_as_i64(), Ordering::Release);
    }

    /// Release tokens back to the bucket
    ///
    /// # Arguments
    /// * `tokens` - Number of tokens to release
    pub fn release(&self, tokens: u64) {
        if tokens == 0 || !self.enabled.load(Ordering::Relaxed) {
            return;
        }

        let tokens = token_amount_to_i64(tokens);

        loop {
            let current = self.token_balance.load(Ordering::Acquire);
            let capacity = self.capacity_as_i64();
            let new_value = current.saturating_add(tokens).min(capacity);

            if self
                .token_balance
                .compare_exchange(current, new_value, Ordering::AcqRel, Ordering::Acquire)
                .is_ok()
            {
                break;
            }
        }
    }

    /// Refill tokens based on elapsed time.
    ///
    /// Returns `true` when this call advances the refill timestamp and performs
    /// the interval's refill work.
    pub fn refill(&self) -> bool {
        if !self.enabled.load(Ordering::Relaxed) {
            return false;
        }

        let interval_ms = self.refill_interval.as_millis() as u64;
        if interval_ms == 0 {
            return false;
        }

        loop {
            let now_ms = self.elapsed_since_start_ms();
            let last_ms = self.last_refill_elapsed_ms.load(Ordering::Acquire);
            let elapsed_ms = now_ms.saturating_sub(last_ms);
            if elapsed_ms < interval_ms {
                return false;
            }

            let intervals = elapsed_ms / interval_ms;
            // Advance by complete intervals instead of storing `now_ms`.
            // Keeping `elapsed_ms % interval_ms` as remainder makes the next
            // refill occur at the original boundary; discarding it would push
            // every later refill back and accumulate timestamp drift.
            let new_last_ms = last_ms.saturating_add(interval_ms.saturating_mul(intervals));
            if self
                .last_refill_elapsed_ms
                .compare_exchange(last_ms, new_last_ms, Ordering::AcqRel, Ordering::Acquire)
                .is_err()
            {
                continue;
            }

            let refill_rate = self.current_refill_rate_us.load(Ordering::Relaxed);
            // Skip refill for unlimited quota (represented by u64::MAX)
            // to avoid overflow issues in saturating_mul. The timestamp is
            // already advanced by the CAS, matching the previous behavior.
            if refill_rate == u64::MAX {
                return true;
            }

            let tokens_to_add = refill_rate.saturating_mul(intervals);
            self.release(tokens_to_add);
            return true;
        }
    }

    /// Returns the absolute deadline of the next refill boundary.
    pub fn next_refill_deadline(&self) -> Instant {
        let interval_ms = self.refill_interval.as_millis().max(1) as u64;
        let next_elapsed_ms = self
            .last_refill_elapsed_ms
            .load(Ordering::Acquire)
            .saturating_add(interval_ms);
        self.start_instant
            .checked_add(Duration::from_millis(next_elapsed_ms))
            .unwrap_or_else(Instant::now)
    }

    /// Get current available tokens
    pub fn available(&self) -> u64 {
        token_balance_to_available(self.token_balance.load(Ordering::Relaxed))
    }

    /// Get current debt tokens.
    pub fn debt(&self) -> u64 {
        token_balance_to_debt(self.token_balance.load(Ordering::Relaxed))
    }

    /// Returns whether the bucket has positive token balance after debt.
    pub fn has_positive_available_with_debt(&self) -> bool {
        if !self.enabled.load(Ordering::Relaxed) {
            return true;
        }
        self.token_balance.load(Ordering::Acquire) > 0
    }

    /// Check if allocation can proceed within interval limit
    /// Returns Some(window_id) if allocation succeeds, None if limit is
    /// exceeded. When limit is disabled (0), returns Some(0).
    /// Self-resetting: detects a window change and atomically resets the window
    /// ID and allocation counter before reserving from the new interval.
    pub fn try_allocate_in_interval(&self, amount: u64) -> Option<u64> {
        let interval_limit = self.interval_allocation_limit.load(Ordering::Acquire);
        if interval_limit == 0 {
            return Some(0); // Limit disabled, return dummy window_id
        }

        let interval_ms = self.refill_interval.as_millis() as u64;
        if interval_ms == 0 {
            return Some(0);
        }

        let mut interval = self.interval_allocation.lock().unwrap();
        let elapsed_total_ms = self.start_instant.elapsed().as_millis() as u64;
        let current_window = elapsed_total_ms / interval_ms;
        if current_window > interval.window_id {
            interval.window_id = current_window;
            interval.allocated_us = 0;
        }

        if interval.allocated_us.saturating_add(amount) > interval_limit {
            return None;
        }
        interval.allocated_us = interval.allocated_us.saturating_add(amount);
        Some(interval.window_id)
    }

    /// Release previously allocated amount from interval counter (rollback)
    /// Used when try_allocate_in_interval succeeds but subsequent try_allocate
    /// fails. Validates window_id to prevent underflow if window was reset
    /// between allocation and rollback.
    pub fn release_in_interval(&self, amount: u64, window_id: u64) {
        if self.interval_allocation_limit.load(Ordering::Acquire) == 0 {
            return; // Limit disabled, nothing to rollback
        }

        let mut interval = self.interval_allocation.lock().unwrap();
        if interval.window_id == window_id {
            interval.allocated_us = interval.allocated_us.saturating_sub(amount);
        }
    }

    /// Update per-interval allocation limit (0 = disabled)
    pub fn set_interval_allocation_limit(&self, new_limit: u64) {
        self.interval_allocation_limit
            .store(new_limit, Ordering::Release);
    }

    fn interval_allocation_limit(&self) -> u64 {
        self.interval_allocation_limit.load(Ordering::Acquire)
    }

    #[cfg(test)]
    fn allocated_in_interval(&self) -> u64 {
        self.interval_allocation.lock().unwrap().allocated_us
    }

    /// Set refill rate dynamically
    ///
    /// # Arguments
    /// * `rate_us` - New refill rate in microseconds per interval
    pub fn set_refill_rate(&self, rate_us: u64) {
        self.current_refill_rate_us
            .store(rate_us, Ordering::Release);
    }

    /// Restore to base refill rate
    pub fn restore_base_rate(&self) {
        let base_rate = self.base_refill_rate_us.load(Ordering::Acquire);
        self.current_refill_rate_us
            .store(base_rate, Ordering::Release);
    }

    /// Get base refill rate
    pub fn base_refill_rate(&self) -> u64 {
        self.base_refill_rate_us.load(Ordering::Acquire)
    }

    /// Get current refill rate
    pub fn current_refill_rate(&self) -> u64 {
        self.current_refill_rate_us.load(Ordering::Acquire)
    }

    /// Get bucket capacity
    pub fn capacity(&self) -> u64 {
        self.capacity_us.load(Ordering::Acquire)
    }

    /// Update bucket capacity and refill rate while preserving consumed tokens.
    pub fn update_quota(&self, new_capacity_us: u64, new_refill_rate_us: u64) {
        let old_capacity = token_amount_to_i128(self.capacity_us.load(Ordering::Acquire));
        let current_balance = i128::from(self.token_balance.load(Ordering::Acquire));
        let consumed = old_capacity.saturating_sub(current_balance);
        let new_capacity = token_amount_to_i128(new_capacity_us);
        let new_balance = clamp_i128_to_i64(new_capacity.saturating_sub(consumed));

        self.capacity_us.store(new_capacity_us, Ordering::Release);
        self.base_refill_rate_us
            .store(new_refill_rate_us, Ordering::Release);
        self.current_refill_rate_us
            .store(new_refill_rate_us, Ordering::Release);
        self.token_balance.store(new_balance, Ordering::Release);
    }

    fn update_quota_reserving_inflight(
        &self,
        new_capacity_us: u64,
        new_refill_rate_us: u64,
        inflight_allocated_us: u64,
    ) {
        let old_capacity = token_amount_to_i128(self.capacity_us.load(Ordering::Acquire));
        let current_available = token_amount_to_i128(self.available());
        let old_consumed = old_capacity.saturating_sub(current_available);
        let reserved = old_consumed.max(token_amount_to_i128(inflight_allocated_us));
        let new_capacity = token_amount_to_i128(new_capacity_us);
        let new_available = (new_capacity - reserved).max(0);

        self.capacity_us.store(new_capacity_us, Ordering::Release);
        self.base_refill_rate_us
            .store(new_refill_rate_us, Ordering::Release);
        self.current_refill_rate_us
            .store(new_refill_rate_us, Ordering::Release);
        self.token_balance
            .store(clamp_i128_to_i64(new_available), Ordering::Release);
    }

    /// Set enabled state
    pub fn set_enabled(&self, enabled: bool) {
        self.enabled.store(enabled, Ordering::Release);
    }
}

/// Token handle for CPU resource allocation
/// Automatically returns unused tokens when dropped
#[derive(Debug)]
pub struct CpuTokenHandle {
    /// Keyspace ID
    pub(crate) keyspace_id: u32,
    /// Total allocated tokens in microseconds (atomic for runtime updates)
    pub(crate) allocated_us: AtomicU64,
    /// Tokens allocated from keyspace bucket (for accurate release on drop)
    pub(crate) keyspace_allocated_us: AtomicU64,
    /// Tokens allocated from global bucket (for tracking/release)
    pub(crate) global_allocated_us: AtomicU64,
    /// Actually used CPU time in microseconds
    pub(crate) actual_used_us: AtomicU64,
    /// Reference to global token bucket
    pub(crate) global_bucket: Arc<CpuTokenBucket>,
    /// Reference to keyspace token bucket (optional)
    pub(crate) keyspace_bucket: Option<Arc<CpuTokenBucket>>,
    /// Reference to the low-priority DAG global bucket (optional)
    pub(crate) low_priority_bucket: Option<Arc<CpuTokenBucket>>,
    /// Cached keyspace inflight state and metric handles (optional)
    pub(crate) keyspace_inflight_state: Option<Arc<KeyspaceInflightState>>,
    /// Tokens allocated from the low-priority DAG global bucket
    pub(crate) low_priority_allocated_us: AtomicU64,
    /// Low-priority DAG inflight allocation tracking (optional)
    pub(crate) low_priority_inflight_allocated_us: Option<Arc<AtomicU64>>,
    /// Whether this allocation is in burst mode
    pub(crate) is_burst: AtomicBool,
    /// Request deadline for timeout checks
    pub(crate) request_deadline: Instant,
    /// Whether to skip throttle for this keyspace (no quota configured)
    pub(crate) skip_throttle: bool,
    /// Whether burst is allowed for this keyspace
    pub(crate) keyspace_burst_allowed: bool,
    /// CPU usage monitor for dynamic burst decision
    pub(crate) cpu_usage_monitor: Option<Arc<RwLock<CpuUsageMonitor>>>,
    /// CPU throttle config for burst decision
    pub(crate) config: Option<Arc<CpuThrottleConfig>>,
    /// Adaptive estimator for per-keyspace estimated CPU
    pub(crate) adaptive_estimator: Option<Arc<AdaptiveEstimator>>,
    /// Burst usage tracking token (only for burst allocations)
    pub(crate) burst_usage_token: Option<BurstUsageToken>,
    /// Burst usage tracker for global burst budget
    pub(crate) burst_usage_tracker: Option<Arc<BurstUsageTracker>>,
    /// Last time "Switched to dynamic burst mode" log was printed
    pub(crate) last_dynamic_burst_log: Option<Arc<RwLock<Instant>>>,
    /// Last time "CPU token periodic allocation succeeded" log was printed
    pub(crate) last_periodic_allocation_log: Option<Arc<RwLock<Instant>>>,
}

impl CpuTokenHandle {
    /// Sleep helper function compatible with yatp runtime
    /// Uses global timer handle for truly async sleep without thread blocking
    async fn yatp_sleep(duration: Duration) {
        use tikv_util::retry::sleep_async;
        sleep_async(duration).await;
    }

    fn try_allocate_low_priority_tokens(&self, amount_us: u64) -> bool {
        self.low_priority_bucket
            .as_ref()
            .map(|bucket| bucket.try_allocate(amount_us))
            .unwrap_or(true)
    }

    fn commit_low_priority_allocation(&self, amount_us: u64) {
        if let Some(ref bucket) = self.low_priority_bucket {
            self.low_priority_allocated_us
                .fetch_add(amount_us, Ordering::AcqRel);
            if let Some(ref inflight) = self.low_priority_inflight_allocated_us {
                inflight.fetch_add(amount_us, Ordering::AcqRel);
            }
            CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE
                .set(clamp_gauge_value(bucket.available()));
        }
    }

    fn rollback_low_priority_allocation(&self, amount_us: u64) {
        if let Some(ref bucket) = self.low_priority_bucket {
            bucket.release(amount_us);
            CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE
                .set(clamp_gauge_value(bucket.available()));
        }
    }

    /// Rollback low-priority tokens, refresh metrics, and signal allocation
    /// failure. Every early-return-false path in `allocate_more` needs
    /// exactly this sequence, so it lives here to avoid repetition.
    fn fail_allocation(&self, additional_us: u64) -> bool {
        self.rollback_low_priority_allocation(additional_us);
        self.update_available_metrics();
        false
    }

    /// Allocate more tokens during execution (non-blocking)
    ///
    /// # Arguments
    /// * `additional_us` - Additional tokens needed in microseconds
    ///
    /// # Returns
    /// `true` if allocation succeeded, `false` otherwise
    pub fn allocate_more(&self, additional_us: u64) -> bool {
        if !self.try_allocate_low_priority_tokens(additional_us) {
            return false;
        }

        // Check burst state
        let is_burst = self.is_burst.load(Ordering::Acquire);

        if is_burst {
            if !self.can_dynamic_burst() {
                return self.fail_allocation(additional_us);
            }
            let mut budget_window = None;
            if let (Some(ref tracker), Some(ref config)) = (&self.burst_usage_tracker, &self.config)
            {
                match reserve_burst_budget(
                    tracker,
                    self.global_bucket.base_refill_rate(),
                    config.burst_threshold,
                    additional_us,
                ) {
                    Ok(window_id) => budget_window = window_id,
                    Err(()) => {
                        return self.fail_allocation(additional_us);
                    }
                }
            }
            // Burst mode: check per-interval limit then allocate from global
            let interval_window_id =
                match self.global_bucket.try_allocate_in_interval(additional_us) {
                    Some(window_id) => window_id,
                    None => {
                        if let (Some(window_id), Some(ref tracker)) =
                            (budget_window, &self.burst_usage_tracker)
                        {
                            tracker.release_global(window_id, additional_us);
                        }
                        return self.fail_allocation(additional_us);
                    }
                };
            if !self.global_bucket.try_allocate(additional_us) {
                // Rollback interval allocation since token allocation failed
                self.global_bucket
                    .release_in_interval(additional_us, interval_window_id);
                if let (Some(window_id), Some(ref tracker)) =
                    (budget_window, &self.burst_usage_tracker)
                {
                    tracker.release_global(window_id, additional_us);
                }
                return self.fail_allocation(additional_us);
            }
            self.global_allocated_us
                .fetch_add(additional_us, Ordering::AcqRel);
            self.allocated_us.fetch_add(additional_us, Ordering::AcqRel);
            self.commit_low_priority_allocation(additional_us);
            self.update_available_metrics();
            return true;
        }

        // Normal mode: allocate from keyspace bucket (hard isolation by default)
        if let Some(ref bucket) = self.keyspace_bucket {
            // Pre-check: if keyspace doesn't have enough tokens and dynamic burst not
            // allowed, fail fast without touching global bucket (avoid
            // contention)
            if bucket.available() < additional_us && !self.can_dynamic_burst() {
                return self.fail_allocation(additional_us);
            }

            if bucket.try_allocate(additional_us) {
                // Keyspace allocation succeeded
                let enforce_global = self.config.as_ref().is_some_and(|c| c.enforce_global_limit);
                let mut global_tracked = false;
                if enforce_global {
                    if !self.global_bucket.try_allocate(additional_us) {
                        // Global exhausted, rollback keyspace allocation
                        bucket.release(additional_us);
                        return self.fail_allocation(additional_us);
                    }
                    global_tracked = true;
                } else if self.global_bucket.try_allocate(additional_us) {
                    global_tracked = true;
                }

                // Update tracking counters
                self.keyspace_allocated_us
                    .fetch_add(additional_us, Ordering::AcqRel);
                if global_tracked {
                    self.global_allocated_us
                        .fetch_add(additional_us, Ordering::AcqRel);
                }
                self.allocated_us.fetch_add(additional_us, Ordering::AcqRel);
                self.inflight_add(additional_us);
                self.commit_low_priority_allocation(additional_us);
                self.update_available_metrics();
                return true;
            } else {
                // Keyspace bucket exhausted, try dynamic burst.
                //
                // Note: fair-share enforcement (calculate_fair_share) is not
                // applied here. The handle lacks access to the per-keyspace
                // quota map required to compute a keyspace's fair share, and
                // per-keyspace burst usage is not recorded in the
                // BurstUsageTracker for this runtime allocation. As a result,
                // in the overlap zone where fair_allocation_threshold <=
                // global_usage_ratio < burst_threshold, a long-running request
                // that exhausts its keyspace quota mid-flight may consume more
                // burst tokens than its proportional fair share for that
                // interval. This is a known limitation; stricter enforcement
                // would require passing keyspace quota context into the handle.
                if self.can_dynamic_burst() {
                    let mut budget_window = None;
                    if let (Some(ref tracker), Some(ref config)) =
                        (&self.burst_usage_tracker, &self.config)
                    {
                        match reserve_burst_budget(
                            tracker,
                            self.global_bucket.base_refill_rate(),
                            config.burst_threshold,
                            additional_us,
                        ) {
                            Ok(window_id) => budget_window = window_id,
                            Err(()) => {
                                return self.fail_allocation(additional_us);
                            }
                        }
                    }
                    // Check per-interval limit before burst allocation
                    let interval_window_id =
                        match self.global_bucket.try_allocate_in_interval(additional_us) {
                            Some(window_id) => window_id,
                            None => {
                                if let (Some(window_id), Some(ref tracker)) =
                                    (budget_window, &self.burst_usage_tracker)
                                {
                                    tracker.release_global(window_id, additional_us);
                                }
                                return self.fail_allocation(additional_us);
                            }
                        };
                    if !self.global_bucket.try_allocate(additional_us) {
                        // Rollback interval allocation since token allocation failed
                        self.global_bucket
                            .release_in_interval(additional_us, interval_window_id);
                        if let (Some(window_id), Some(ref tracker)) =
                            (budget_window, &self.burst_usage_tracker)
                        {
                            tracker.release_global(window_id, additional_us);
                        }
                        return self.fail_allocation(additional_us);
                    }
                    // Switch to burst mode
                    self.is_burst.store(true, Ordering::Release);
                    self.global_allocated_us
                        .fetch_add(additional_us, Ordering::AcqRel);

                    // Log based on debug mode
                    let should_log = if self.config.as_ref().is_some_and(|c| c.debug) {
                        true
                    } else {
                        self.last_dynamic_burst_log
                            .as_ref()
                            .is_some_and(|last_log| {
                                let mut last_log = last_log.write().unwrap();
                                if last_log.elapsed() >= Duration::from_secs(5) {
                                    *last_log = Instant::now();
                                    true
                                } else {
                                    false
                                }
                            })
                    };

                    if should_log {
                        info!(
                            "[CPU throttle] Switched to dynamic burst mode during execution";
                            "keyspace_id" => self.keyspace_id,
                            "additional_us" => additional_us,
                        );
                    }

                    self.allocated_us.fetch_add(additional_us, Ordering::AcqRel);
                    self.commit_low_priority_allocation(additional_us);
                    self.update_available_metrics();
                    return true;
                }
                return self.fail_allocation(additional_us);
            }
        }

        // No keyspace bucket, allocate from global only
        if !self.global_bucket.try_allocate(additional_us) {
            return self.fail_allocation(additional_us);
        }
        self.global_allocated_us
            .fetch_add(additional_us, Ordering::AcqRel);
        self.allocated_us.fetch_add(additional_us, Ordering::AcqRel);
        self.commit_low_priority_allocation(additional_us);
        self.update_available_metrics();
        true
    }

    fn update_available_metrics(&self) {
        update_global_bucket_metrics(&self.global_bucket);
        if let Some(ref bucket) = self.low_priority_bucket {
            CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE
                .set(clamp_gauge_value(bucket.available()));
        }
        if let Some(ref bucket) = self.keyspace_bucket {
            if let Some(ref state) = self.keyspace_inflight_state {
                state
                    .bucket_available
                    .set(clamp_gauge_value(bucket.available()));
            }
        }
    }

    /// Check if dynamic burst is allowed during execution
    fn can_dynamic_burst(&self) -> bool {
        self.dynamic_burst_policy_allows() && self.global_bucket.available() > 0
    }

    /// Check the policy conditions that can make dynamic burst viable after a
    /// future global refill. Unlike `can_dynamic_burst`, this intentionally
    /// ignores the current global token balance.
    fn dynamic_burst_policy_allows(&self) -> bool {
        let (monitor, config) = match (&self.cpu_usage_monitor, &self.config) {
            (Some(m), Some(c)) => (m, c),
            _ => return false,
        };

        if !self.keyspace_burst_allowed {
            return false;
        }

        if !config.enable_burst {
            return false;
        }

        let monitor_guard = monitor.read().unwrap();
        let global_ratio = monitor_guard.global_usage_ratio;

        global_ratio < config.burst_threshold
    }

    fn runtime_allocation_retry_at(&self, additional_us: u64) -> Instant {
        let mut retry_at = None;
        let mut include = |deadline| {
            retry_at = Some(retry_at.map_or(deadline, |current: Instant| current.max(deadline)));
        };
        if let Some(ref bucket) = self.low_priority_bucket {
            if bucket.available() < additional_us {
                include(bucket.next_refill_deadline());
            }
        }
        if self.is_burst.load(Ordering::Acquire) {
            if self.global_bucket.available() < additional_us {
                include(self.global_bucket.next_refill_deadline());
            }
        } else {
            if let Some(ref bucket) = self.keyspace_bucket {
                if bucket.available() < additional_us {
                    let keyspace_retry_at = bucket.next_refill_deadline();
                    if self.dynamic_burst_policy_allows() {
                        // Keyspace refill and global burst are alternative
                        // recovery paths. The global path remains viable after
                        // refill even when the bucket is currently empty.
                        include(keyspace_retry_at.min(self.global_bucket.next_refill_deadline()));
                    } else {
                        include(keyspace_retry_at);
                    }
                }
                if self
                    .config
                    .as_ref()
                    .is_some_and(|config| config.enforce_global_limit)
                    && self.global_bucket.available() < additional_us
                {
                    include(self.global_bucket.next_refill_deadline());
                }
            } else if !self.skip_throttle && self.global_bucket.available() < additional_us {
                include(self.global_bucket.next_refill_deadline());
            }
        }
        retry_at.unwrap_or_else(|| self.global_bucket.next_refill_deadline())
    }

    /// Allocate more tokens with wait on refill
    ///
    /// # Arguments
    /// * `additional_us` - Additional tokens needed in microseconds
    ///
    /// # Returns
    /// `Ok(u64)` with allocated amount on success, `Err(ThrottleError)` on
    /// failure/timeout
    pub async fn allocate_more_with_wait(&self, additional_us: u64) -> Result<u64, ThrottleError> {
        let start = Instant::now();
        let keyspace_label = keyspace_label(self.keyspace_id);
        let keyspace_label = keyspace_label.as_ref();

        // Refill immediately first to maximize available tokens
        // This ensures we account for time elapsed since last refill before starting
        // wait logic
        self.global_bucket.refill();
        if let Some(ref bucket) = self.low_priority_bucket {
            bucket.refill();
        }
        if !self.is_burst.load(Ordering::Acquire) {
            if let Some(ref bucket) = self.keyspace_bucket {
                bucket.refill();
            }
        }

        let now = Instant::now();
        if now >= self.request_deadline {
            CPU_TOKEN_RUNTIME_WAIT_DURATION_HISTOGRAM
                .with_label_values(&[keyspace_label, "timeout"])
                .observe(start.elapsed().as_secs_f64());
            CPU_TOKEN_RUNTIME_ALLOCATIONS
                .with_label_values(&["timeout"])
                .inc();
            return Err(ThrottleError::RequestTimeout);
        }
        // 1. Try immediate allocation after refilling
        if self.allocate_more(additional_us) {
            CPU_TOKEN_RUNTIME_WAIT_DURATION_HISTOGRAM
                .with_label_values(&[keyspace_label, "success"])
                .observe(start.elapsed().as_secs_f64());
            CPU_TOKEN_RUNTIME_ALLOCATIONS
                .with_label_values(&["success"])
                .inc();
            return Ok(additional_us);
        }

        // 2. Immediate allocation failed, wait for refill
        loop {
            let now = Instant::now();
            if now >= self.request_deadline {
                CPU_TOKEN_RUNTIME_WAIT_DURATION_HISTOGRAM
                    .with_label_values(&[keyspace_label, "timeout"])
                    .observe(start.elapsed().as_secs_f64());
                CPU_TOKEN_RUNTIME_ALLOCATIONS
                    .with_label_values(&["timeout"])
                    .inc();
                return Err(ThrottleError::RequestTimeout);
            }

            let retry_at = self
                .runtime_allocation_retry_at(additional_us)
                .min(self.request_deadline);
            let wait_duration = retry_at.saturating_duration_since(now);
            Self::yatp_sleep(wait_duration).await;

            // Refill tokens before retrying allocation
            self.global_bucket.refill();
            if let Some(ref bucket) = self.low_priority_bucket {
                bucket.refill();
            }
            if let Some(ref bucket) = self.keyspace_bucket {
                bucket.refill();
            }

            // Retry allocation
            if self.allocate_more(additional_us) {
                CPU_TOKEN_RUNTIME_WAIT_DURATION_HISTOGRAM
                    .with_label_values(&[keyspace_label, "success"])
                    .observe(start.elapsed().as_secs_f64());
                CPU_TOKEN_RUNTIME_ALLOCATIONS
                    .with_label_values(&["success"])
                    .inc();
                return Ok(additional_us);
            }
        }
    }

    /// Record actual CPU usage
    ///
    /// # Arguments
    /// * `actual_cpu_us` - Actual CPU time used in microseconds
    pub fn record_actual_usage(&self, actual_cpu_us: u64) {
        self.actual_used_us
            .fetch_max(actual_cpu_us, Ordering::AcqRel);
    }

    /// Get allocated tokens
    pub fn allocated(&self) -> u64 {
        self.allocated_us.load(Ordering::Acquire)
    }

    /// Get the largest additional runtime allocation that can succeed in one
    /// token-bucket attempt.
    pub fn max_single_runtime_allocation_us(&self) -> u64 {
        let mut max_allocation_us = self.global_bucket.capacity();
        if self.is_burst.load(Ordering::Acquire) {
            let interval_limit = self.global_bucket.interval_allocation_limit();
            if interval_limit > 0 {
                max_allocation_us = max_allocation_us.min(interval_limit);
            }
        } else {
            if let Some(ref bucket) = self.keyspace_bucket {
                max_allocation_us = max_allocation_us.min(bucket.capacity());
            }
        }
        if let Some(ref bucket) = self.low_priority_bucket {
            max_allocation_us = max_allocation_us.min(bucket.capacity());
        }
        max_allocation_us.max(1)
    }

    /// Get the keyspace ID associated with this handle.
    pub fn keyspace_id(&self) -> u32 {
        self.keyspace_id
    }

    /// Get actual used CPU time
    pub fn actual_used(&self) -> u64 {
        self.actual_used_us.load(Ordering::Relaxed)
    }

    /// Get runtime token management configuration if enabled
    ///
    /// # Returns
    /// `Some((check_interval, threshold, per_allocation_us))` if runtime token
    /// management is enabled `None` if disabled or config not available
    pub fn get_runtime_config(&self) -> Option<(Duration, f64, u64)> {
        self.config.as_ref().and_then(|cfg| {
            if cfg.enable_runtime_token_management {
                Some((
                    Duration::from_micros(cfg.runtime_check_interval_us),
                    cfg.additional_allocation_threshold,
                    cfg.per_allocation_us,
                ))
            } else {
                None
            }
        })
    }

    fn inflight_add(&self, delta_us: u64) {
        if let (Some(ref state), Some(ref bucket)) =
            (&self.keyspace_inflight_state, &self.keyspace_bucket)
        {
            let new_value = state.add(delta_us);
            state.update_inflight_metrics(new_value, bucket.capacity());
        }
    }

    fn inflight_sub(&self, delta_us: u64) {
        if let (Some(ref state), Some(ref bucket)) =
            (&self.keyspace_inflight_state, &self.keyspace_bucket)
        {
            let new_value = state.sub(delta_us);
            state.update_inflight_metrics(new_value, bucket.capacity());
        }
    }

    fn low_priority_inflight_sub(&self, delta_us: u64) {
        if let Some(ref inflight) = self.low_priority_inflight_allocated_us {
            atomic_saturating_sub(inflight, delta_us);
        }
    }
}

impl Drop for CpuTokenHandle {
    fn drop(&mut self) {
        // `skip_throttle` only means the regular keyspace/global path was
        // skipped because no quota was configured. A low-priority DAG request
        // may still hold low-priority tokens in that case, so we can only
        // return early when neither path allocated anything.
        if self.skip_throttle && self.low_priority_bucket.is_none() {
            return;
        }

        let actual = self.actual_used_us.load(Ordering::Acquire);
        if !self.skip_throttle {
            if let Some(ref estimator) = self.adaptive_estimator {
                estimator.observe(self.keyspace_id, actual);
            }
        }
        let keyspace_allocated = self.keyspace_allocated_us.load(Ordering::Acquire);
        let global_allocated = self.global_allocated_us.load(Ordering::Acquire);
        let low_priority_allocated = self.low_priority_allocated_us.load(Ordering::Acquire);

        if !self.skip_throttle {
            // Handle burst usage tracking
            if let Some(ref token) = self.burst_usage_token {
                let used_for_initial = actual.min(token.allocated_us);
                let unused = token.allocated_us - used_for_initial;
                if unused > 0 {
                    token.tracker.release_usage_in_window(
                        token.keyspace_id,
                        token.window_id,
                        unused,
                    );
                }
            }

            // A handle may start in normal mode, allocate from the keyspace bucket,
            // and later switch to burst mode at runtime. Refund every tracked bucket
            // independently based on the total actual CPU used.
            if actual < global_allocated {
                let global_unused = global_allocated - actual;
                self.global_bucket.release(global_unused);
            }
            if let Some(ref bucket) = self.keyspace_bucket {
                if actual < keyspace_allocated {
                    let keyspace_unused = keyspace_allocated - actual;
                    bucket.release(keyspace_unused);
                }
            }

            self.inflight_sub(keyspace_allocated);
        }

        if let Some(ref bucket) = self.low_priority_bucket {
            if actual < low_priority_allocated {
                let low_priority_unused = low_priority_allocated - actual;
                bucket.release(low_priority_unused);
            }
            self.low_priority_inflight_sub(low_priority_allocated);
        }

        self.update_available_metrics();
    }
}

/// CPU throttle manager
/// Manages global and per-keyspace token buckets
pub struct CpuThrottleManager {
    /// Global token bucket
    global_bucket: Arc<CpuTokenBucket>,
    /// Whether the keyspace-fair read pool owns global CPU accounting.
    fair_global_cpu_gate_registered: AtomicBool,
    /// Packed fair-gate charge state: the fixed-point scale in the high bits
    /// and the carried sub-token remainder in the low
    /// [`FAIR_CHARGE_RESIDUE_BITS`].
    ///
    /// A scale of [`FAIR_CHARGE_SCALE_ONE`] means 1.0 (no override). Larger
    /// values make each microsecond of measured CPU cost more tokens, which
    /// lowers the effective CPU ratio without resizing the bucket.
    ///
    /// The two live in one word so that changing the exchange rate and
    /// discarding the remainder that belonged to the old rate is a single
    /// atomic transition. See
    /// [`CpuThrottleManager::update_fair_max_read_cpu_ratio`].
    fair_charge_state: AtomicU64,
    /// Read CPU ratio config is asking for (f64 bits), whether or not the fair
    /// gate can enforce it. Metrics only.
    requested_max_read_cpu_ratio_bits: AtomicU64,
    /// Read CPU ratio the global bucket was sized from at startup (f64 bits).
    /// The enforced ratio is derived from this and the charge scale, so it is
    /// the anchor for the effective gauge. Metrics only.
    startup_max_read_cpu_ratio_bits: AtomicU64,
    /// Serializes the one-time startup anchor initialization with online fair
    /// ratio updates. These transitions publish several related atomics and
    /// metrics, while the hot charge path remains lock-free.
    fair_charge_update_mutex: Mutex<()>,
    /// Low-priority DAG global token bucket
    low_priority_dag_bucket: Arc<CpuTokenBucket>,
    /// Tokens currently held by in-flight low-priority DAG requests.
    low_priority_dag_inflight_allocated_us: Arc<AtomicU64>,
    /// Per-keyspace token buckets
    keyspace_buckets: Arc<DashMap<u32, Arc<CpuTokenBucket>>>,
    /// CPU usage monitor
    cpu_usage_monitor: Arc<RwLock<CpuUsageMonitor>>,
    /// Burst usage tracker for fair allocation
    burst_usage_tracker: Arc<BurstUsageTracker>,
    /// Keyspace quotas (keyspace_id -> quota in millicores)
    keyspace_quotas: Arc<DashMap<u32, u64>>,
    /// Keyspaces configured via config keyspace_cpu_quotas
    config_keyspaces: DashSet<u32>,
    /// Keyspaces auto-created by uniform tenant cap mode.
    tenant_cap_keyspaces: DashSet<u32>,
    /// Last request time for auto-created tenant cap keyspaces.
    tenant_cap_keyspace_last_access: DashMap<u32, Instant>,
    /// Last time idle tenant cap keyspaces were scanned.
    last_tenant_cap_cleanup: Arc<RwLock<Instant>>,
    /// Idle TTL for auto-created tenant cap keyspaces, in milliseconds.
    tenant_cap_idle_ttl_ms: AtomicU64,
    /// Serializes tenant cap bucket (re)creation against disable/teardown.
    /// An in-flight allocation that snapshotted a stale quota must not be able
    /// to resurrect a keyspace bucket after a concurrent disable has torn it
    /// down: recreation re-loads the quota inside this lock, and disable holds
    /// the same lock while clearing quota + removing buckets, so the two paths
    /// cannot interleave.
    tenant_cap_lifecycle: Mutex<()>,
    /// Per-keyspace inflight allocated tokens and cached metric handles
    keyspace_inflight_states: Arc<DashMap<u32, Arc<KeyspaceInflightState>>>,
    /// Whether the low-priority DAG global bucket is enabled
    low_priority_dag_enabled: AtomicBool,
    /// Low-priority DAG global ratio (f64 bits)
    low_priority_dag_ratio_bits: AtomicU64,
    /// The immutable global CPU budget used to derive the low-priority DAG
    /// bucket
    low_priority_dag_global_cpu_us: u64,
    /// Token refill interval in milliseconds
    refill_interval_ms: u64,
    /// Maximum duration of tokens retained by each bucket.
    burst_window_ms: u64,
    /// Configuration
    config: Arc<CpuThrottleConfig>,
    /// Dynamic adjustment enabled flag
    dynamic_adjustment_enabled: AtomicBool,
    /// Serializes dynamic refill updates against fair-gate registration and
    /// online configuration. This runs on periodic usage/config paths, not on
    /// the per-poll fair CPU charge path.
    dynamic_adjustment_mutex: Mutex<()>,
    /// High watermark for dynamic adjustment (f64 bits)
    high_watermark_bits: AtomicU64,
    /// Low watermark for dynamic adjustment (f64 bits)
    low_watermark_bits: AtomicU64,
    /// Fair allocation enabled flag
    fair_allocation_enabled: AtomicBool,
    /// Fair allocation threshold (f64 bits)
    fair_allocation_threshold_bits: AtomicU64,
    /// Burst enabled flag
    burst_enabled: AtomicBool,
    /// Burst threshold (f64 bits)
    burst_threshold_bits: AtomicU64,
    /// Enforce global limit flag
    enforce_global_limit_enabled: AtomicBool,
    /// Whether PD service_limit updates may manage keyspace quotas.
    pd_service_limit_enabled: AtomicBool,
    /// Uniform tenant max CPU ratio (f64 bits). 0 disables cap mode.
    tenant_max_cpu_ratio_bits: AtomicU64,
    /// Pre-computed tenant cap quota in microseconds. 0 disables cap mode.
    tenant_cap_quota_us: AtomicU64,
    /// Debug logging enabled flag
    debug_enabled: AtomicBool,
    /// Keyspace overrides for estimated CPU time per request
    keyspace_estimated_cpu_per_request_us: Arc<std::collections::HashMap<u32, u64>>,
    /// Keyspace-specific burst enable flags
    keyspace_burst_enabled: DashMap<u32, bool>,
    /// Adaptive estimator for per-keyspace estimated CPU time per request
    adaptive_estimator: Option<Arc<AdaptiveEstimator>>,
    /// Last time "Global CPU exhausted" warning was logged
    last_global_exhausted_log: Arc<RwLock<Instant>>,
    /// Last time "Keyspace CPU exhausted" warning was logged
    last_keyspace_exhausted_log: Arc<RwLock<Instant>>,
    /// Last time "Burst mode activated" log was printed
    last_burst_activated_log: Arc<RwLock<Instant>>,
    /// Last time "Keyspace exceeded fair share" warning was logged
    last_fair_share_exceeded_log: Arc<RwLock<Instant>>,
    /// Last time "CPU token periodic allocation succeeded" log was printed
    last_periodic_allocation_log: Arc<RwLock<Instant>>,
    /// Last time "Switched to dynamic burst mode during execution" log was
    /// printed
    last_dynamic_burst_log: Arc<RwLock<Instant>>,
    /// Last time global refill rate adjustment log was printed
    last_global_refill_rate_log: Arc<RwLock<Instant>>,
}

impl CpuThrottleManager {
    fn low_priority_dag_active(&self) -> bool {
        self.low_priority_dag_enabled.load(Ordering::Acquire)
            && self.low_priority_dag_bucket.capacity() > 0
    }

    fn current_low_priority_dag_refill_rate(&self) -> u64 {
        scale_refill_rate_by_reference(
            self.low_priority_dag_bucket.base_refill_rate(),
            self.global_bucket.current_refill_rate(),
            self.global_bucket.base_refill_rate(),
        )
    }

    fn update_low_priority_dag_bucket_metrics(&self) {
        CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_CAPACITY
            .set(clamp_gauge_value(self.low_priority_dag_bucket.capacity()));
        CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE
            .set(clamp_gauge_value(self.low_priority_dag_bucket.available()));
        CPU_THROTTLE_LOW_PRIORITY_DAG_REFILL_RATE.set(clamp_gauge_value(
            self.low_priority_dag_bucket.current_refill_rate(),
        ));
    }

    /// Remove every refill-rate change made by dynamic adjustment.
    ///
    /// The fair gate owns the global bucket directly, so leaving an earlier
    /// dynamic reduction in place would silently stack two independent CPU
    /// controls. Keyspace and low-priority rates are restored as well to keep
    /// manager state and metrics internally consistent.
    fn restore_dynamic_refill_rates(&self) {
        self.global_bucket.restore_base_rate();
        update_global_bucket_limit_metrics(&self.global_bucket);

        self.low_priority_dag_bucket.restore_base_rate();
        self.update_low_priority_dag_bucket_metrics();

        for entry in self.keyspace_buckets.iter() {
            let keyspace_id = *entry.key();
            let bucket = entry.value();
            bucket.restore_base_rate();
            let keyspace_label = keyspace_label(keyspace_id);
            CPU_THROTTLE_KEYSPACE_REFILL_RATE
                .with_label_values(&[keyspace_label.as_ref()])
                .set(clamp_gauge_value(bucket.current_refill_rate()));
        }
    }
}

impl CpuThrottleManager {
    /// Create a new CPU throttle manager
    ///
    /// # Arguments
    /// * `config` - CPU throttle configuration
    /// * `global_cpu_us` - Global CPU quota in microseconds per interval
    /// * `keyspace_cpu_us` - Per-keyspace CPU quotas (keyspace_id -> us per
    ///   interval)
    pub fn new(
        config: CpuThrottleConfig,
        global_cpu_us: u64,
        keyspace_cpu_us: std::collections::HashMap<u32, u64>,
    ) -> Self {
        info!(
            "[CPU throttle] Manager initialized";
            "config" => ?config,
            "global_refill_rate_us_per_interval" => global_cpu_us,
            "keyspace_refill_rates_us_per_interval" => ?keyspace_cpu_us,
            "refill_interval_ms" => config.refill_interval_ms,
            "burst_window_ms" => config.burst_window_ms,
            "keyspace_count" => keyspace_cpu_us.len(),
        );

        let refill_interval_ms = config.refill_interval_ms.max(1);
        let burst_window_ms = config.burst_window_ms.max(refill_interval_ms);
        let refill_interval = Duration::from_millis(refill_interval_ms);
        let burst_usage_tracker = Arc::new(BurstUsageTracker::new(refill_interval_ms));

        // Keep the interval allocation limit tied to refill rate. Burst
        // capacity must not increase the sustainable allocation rate.
        let global_interval_limit = burst_cap_us(global_cpu_us, config.burst_threshold);

        let global_capacity_us =
            refill_rate_to_capacity(global_cpu_us, refill_interval_ms, burst_window_ms);
        let global_bucket = Arc::new(CpuTokenBucket::new(
            global_capacity_us,
            global_cpu_us,
            refill_interval,
            global_interval_limit,
        ));
        let low_priority_dag_refill_rate = low_priority_dag_quota_us(
            global_cpu_us,
            config.low_priority_dag_global_ratio,
            config.low_priority_dag_global_enabled,
        );
        let low_priority_dag_capacity = refill_rate_to_capacity(
            low_priority_dag_refill_rate,
            refill_interval_ms,
            burst_window_ms,
        );
        let low_priority_dag_bucket = Arc::new(CpuTokenBucket::new(
            low_priority_dag_capacity,
            low_priority_dag_refill_rate,
            refill_interval,
            0,
        ));
        let low_priority_dag_inflight_allocated_us = Arc::new(AtomicU64::new(0));
        let keyspace_buckets = Arc::new(DashMap::new());

        // Initialize keyspace buckets (no per-interval limit for keyspace buckets)
        for (keyspace_id, quota_us) in &keyspace_cpu_us {
            let capacity_us =
                refill_rate_to_capacity(*quota_us, refill_interval_ms, burst_window_ms);
            keyspace_buckets.insert(
                *keyspace_id,
                Arc::new(CpuTokenBucket::new(
                    capacity_us,
                    *quota_us,
                    refill_interval,
                    0,
                )),
            );
        }

        // Initialize keyspace quotas (convert us to millicores for fair allocation)
        let keyspace_quotas = Arc::new(DashMap::new());
        let config_keyspaces = DashSet::new();
        let tenant_cap_keyspaces = DashSet::new();
        let tenant_cap_keyspace_last_access = DashMap::new();
        for (keyspace_id, quota_us) in &keyspace_cpu_us {
            // quota_us per interval -> millicores
            let millicores = quota_us / refill_interval_ms;
            keyspace_quotas.insert(*keyspace_id, millicores);
            config_keyspaces.insert(*keyspace_id);
        }
        let keyspace_inflight_states = Arc::new(DashMap::new());
        let low_priority_dag_enabled = AtomicBool::new(config.low_priority_dag_global_enabled);
        let low_priority_dag_ratio_bits =
            AtomicU64::new(config.low_priority_dag_global_ratio.to_bits());
        let dynamic_adjustment_enabled = AtomicBool::new(config.enable_dynamic_adjustment);
        let high_watermark_bits = AtomicU64::new(config.high_watermark.to_bits());
        let low_watermark_bits = AtomicU64::new(config.low_watermark.to_bits());
        let fair_allocation_enabled = AtomicBool::new(config.enable_fair_allocation);
        let fair_allocation_threshold_bits =
            AtomicU64::new(config.fair_allocation_threshold.to_bits());
        let burst_enabled = AtomicBool::new(config.enable_burst);
        let burst_threshold_bits = AtomicU64::new(config.burst_threshold.to_bits());
        let enforce_global_limit_enabled = AtomicBool::new(config.enforce_global_limit);
        let pd_service_limit_enabled = AtomicBool::new(config.enable_pd_service_limit);
        let tenant_max_cpu_ratio_bits = AtomicU64::new(config.tenant_max_cpu_ratio.to_bits());
        let tenant_cap_quota_us = AtomicU64::new(config.compute_tenant_cap_quota_us().unwrap_or(0));
        let tenant_cap_idle_ttl_ms = AtomicU64::new(config.tenant_cap_idle_ttl_ms.max(1));
        let tenant_cap_lifecycle = Mutex::new(());
        let debug_enabled = AtomicBool::new(config.debug);

        let keyspace_estimated_cpu_per_request_us_map =
            CpuThrottleConfig::parse_keyspace_estimated_cpu_per_request_us(
                &config.keyspace_estimated_cpu_per_request_us,
            );
        let adaptive_estimator = if config.enable_adaptive_estimated_cpu_per_request_us {
            Some(Arc::new(AdaptiveEstimator::new(
                config.estimated_cpu_per_request_us,
                config.stats_interval_ms,
                keyspace_estimated_cpu_per_request_us_map.clone(),
            )))
        } else {
            None
        };
        let keyspace_estimated_cpu_per_request_us =
            Arc::new(keyspace_estimated_cpu_per_request_us_map);
        let keyspace_burst_enabled = DashMap::new();
        for (keyspace_id, enabled) in
            CpuThrottleConfig::parse_keyspace_burst_enabled(&config.keyspace_burst_enabled)
        {
            keyspace_burst_enabled.insert(keyspace_id, enabled);
        }

        // Check for quota oversubscription
        let sum_quotas_us: u64 = keyspace_cpu_us.values().sum();
        if sum_quotas_us > global_cpu_us {
            let oversubscription_ratio = sum_quotas_us as f64 / global_cpu_us as f64;
            if config.enforce_global_limit {
                tikv_util::warn!(
                    "[CPU throttle] Sum of keyspace quotas exceeds global quota";
                    "sum_quotas_us" => sum_quotas_us,
                    "global_cpu_us" => global_cpu_us,
                    "oversubscription_ratio" => format!("{:.2}x", oversubscription_ratio),
                    "enforce_global_limit" => true,
                    "note" => "Global limit will be enforced - requests may be throttled when aggregate usage exceeds global quota",
                );
            } else {
                info!(
                    "[CPU throttle] Sum of keyspace quotas exceeds global quota (hard isolation mode)";
                    "sum_quotas_us" => sum_quotas_us,
                    "global_cpu_us" => global_cpu_us,
                    "oversubscription_ratio" => format!("{:.2}x", oversubscription_ratio),
                    "enforce_global_limit" => false,
                    "note" => "Global limit bypassed - keyspaces can exceed global quota in aggregate. Consider enabling enforce_global_limit if strict global limit is required.",
                );
            }
        }

        update_global_bucket_metrics(&global_bucket);
        update_global_bucket_limit_metrics(&global_bucket);
        CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_CAPACITY
            .set(clamp_gauge_value(low_priority_dag_bucket.capacity()));
        CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE
            .set(clamp_gauge_value(low_priority_dag_bucket.available()));
        CPU_THROTTLE_LOW_PRIORITY_DAG_REFILL_RATE.set(clamp_gauge_value(
            low_priority_dag_bucket.current_refill_rate(),
        ));
        CPU_THROTTLE_TENANT_CAP_RATIO.set(config.tenant_max_cpu_ratio);
        CPU_THROTTLE_TENANT_CAP_QUOTA_US.set(clamp_gauge_value(
            tenant_cap_quota_us.load(Ordering::Acquire),
        ));

        for keyspace_id in keyspace_cpu_us.keys() {
            let keyspace_label = keyspace_label(*keyspace_id);
            let keyspace_label = keyspace_label.as_ref();
            if let Some(bucket) = keyspace_buckets.get(keyspace_id) {
                CPU_THROTTLE_KEYSPACE_BUCKET_CAPACITY
                    .with_label_values(&[keyspace_label])
                    .set(clamp_gauge_value(bucket.capacity()));
                CPU_THROTTLE_KEYSPACE_BUCKET_AVAILABLE
                    .with_label_values(&[keyspace_label])
                    .set(clamp_gauge_value(bucket.available()));
                CPU_THROTTLE_KEYSPACE_REFILL_RATE
                    .with_label_values(&[keyspace_label])
                    .set(clamp_gauge_value(bucket.current_refill_rate()));
            }
        }

        CpuThrottleManager {
            global_bucket,
            fair_global_cpu_gate_registered: AtomicBool::new(false),
            fair_charge_state: AtomicU64::new(pack_fair_charge_state(FAIR_CHARGE_SCALE_ONE, 0)),
            requested_max_read_cpu_ratio_bits: AtomicU64::new(0.0f64.to_bits()),
            startup_max_read_cpu_ratio_bits: AtomicU64::new(0.0f64.to_bits()),
            fair_charge_update_mutex: Mutex::new(()),
            low_priority_dag_bucket,
            low_priority_dag_inflight_allocated_us,
            keyspace_buckets,
            cpu_usage_monitor: Arc::new(RwLock::new(CpuUsageMonitor::default())),
            burst_usage_tracker,
            keyspace_quotas,
            config_keyspaces,
            tenant_cap_keyspaces,
            tenant_cap_keyspace_last_access,
            last_tenant_cap_cleanup: Arc::new(RwLock::new(Instant::now())),
            tenant_cap_idle_ttl_ms,
            tenant_cap_lifecycle,
            keyspace_inflight_states,
            low_priority_dag_enabled,
            low_priority_dag_ratio_bits,
            low_priority_dag_global_cpu_us: global_cpu_us,
            refill_interval_ms,
            burst_window_ms,
            config: Arc::new(config),
            dynamic_adjustment_enabled,
            dynamic_adjustment_mutex: Mutex::new(()),
            high_watermark_bits,
            low_watermark_bits,
            fair_allocation_enabled,
            fair_allocation_threshold_bits,
            burst_enabled,
            burst_threshold_bits,
            enforce_global_limit_enabled,
            pd_service_limit_enabled,
            tenant_max_cpu_ratio_bits,
            tenant_cap_quota_us,
            debug_enabled,
            keyspace_estimated_cpu_per_request_us,
            keyspace_burst_enabled,
            adaptive_estimator,
            last_global_exhausted_log: Arc::new(RwLock::new(Instant::now())),
            last_keyspace_exhausted_log: Arc::new(RwLock::new(Instant::now())),
            last_burst_activated_log: Arc::new(RwLock::new(Instant::now())),
            last_fair_share_exceeded_log: Arc::new(RwLock::new(Instant::now())),
            last_periodic_allocation_log: Arc::new(RwLock::new(Instant::now())),
            last_dynamic_burst_log: Arc::new(RwLock::new(Instant::now())),
            last_global_refill_rate_log: Arc::new(RwLock::new(Instant::now())),
        }
    }

    /// Calculate a keyspace's share of the legacy throttle burst budget for
    /// one refill interval.
    ///
    /// This is used only when the legacy `enable_fair_allocation` guard is
    /// enabled. It is unrelated to read-pool keyspace fairness, whose fair
    /// queue and global admission controller perform dispatch-time fairness.
    ///
    /// # Arguments
    /// * `keyspace_id` - Keyspace ID
    /// * `interval_budget_us` - Global burst budget for one refill interval
    ///
    /// # Returns
    /// Fair share in microseconds
    fn calculate_fair_share(&self, keyspace_id: u32, interval_budget_us: u64) -> Option<u64> {
        // Get keyspace quota
        let keyspace_quota_millicores = self
            .keyspace_quotas
            .get(&keyspace_id)
            .map(|q| *q.value())
            .unwrap_or(0);
        if keyspace_quota_millicores == 0 {
            return None;
        }

        // Get total quotas
        let total_quota_millicores: u64 = self
            .keyspace_quotas
            .iter()
            .map(|entry| *entry.value())
            .sum();

        if total_quota_millicores == 0 {
            return None;
        }

        // Skip fair allocation when global quota is unlimited (u64::MAX)
        if interval_budget_us == u64::MAX {
            return None;
        }

        // Both the per-keyspace usage tracker and this budget cover the same
        // refill interval. Using the bucket's available balance here would mix
        // this accounting window with the longer burst-capacity window.
        let fair_share_us =
            (interval_budget_us.saturating_mul(keyspace_quota_millicores)) / total_quota_millicores;

        Some(fair_share_us)
    }

    fn is_keyspace_burst_allowed(&self, keyspace_id: u32) -> bool {
        self.keyspace_burst_enabled
            .get(&keyspace_id)
            .map(|entry| *entry.value())
            .unwrap_or(true)
    }

    fn tenant_max_cpu_ratio(&self) -> f64 {
        f64::from_bits(self.tenant_max_cpu_ratio_bits.load(Ordering::Acquire))
    }

    fn tenant_cap_quota_us(&self) -> Option<u64> {
        let quota_us = self.tenant_cap_quota_us.load(Ordering::Acquire);
        if quota_us == 0 { None } else { Some(quota_us) }
    }

    pub fn tenant_cap_enabled(&self) -> bool {
        self.tenant_cap_quota_us().is_some()
    }

    pub fn pd_service_limit_enabled(&self) -> bool {
        self.pd_service_limit_enabled.load(Ordering::Acquire)
    }

    pub fn is_dynamic_adjustment_enabled(&self) -> bool {
        self.dynamic_adjustment_enabled.load(Ordering::Acquire)
    }

    pub fn is_fair_allocation_enabled(&self) -> bool {
        self.fair_allocation_enabled.load(Ordering::Acquire)
    }

    pub fn is_enforce_global_limit_enabled(&self) -> bool {
        self.enforce_global_limit_enabled.load(Ordering::Acquire)
    }

    pub fn is_burst_enabled(&self) -> bool {
        self.burst_enabled.load(Ordering::Acquire)
    }

    fn ensure_tenant_cap_bucket(&self, keyspace_id: u32, now: Instant) -> bool {
        // Hold the lifecycle lock so a concurrent disable cannot tear down the
        // bucket between the quota re-load and the insert below. Re-load the
        // quota *inside* the lock: a request that snapshotted a stale non-zero
        // quota before a disable ran must observe quota == 0 here and refuse to
        // resurrect the bucket.
        let _guard = self.tenant_cap_lifecycle.lock().unwrap();
        let quota_us = self.tenant_cap_quota_us.load(Ordering::Acquire);
        if quota_us == 0 {
            return false;
        }

        self.update_keyspace_quota_inner(keyspace_id, quota_us);
        self.tenant_cap_keyspaces.insert(keyspace_id);
        self.tenant_cap_keyspace_last_access
            .insert(keyspace_id, now);
        true
    }

    fn remove_tenant_cap_keyspace(&self, keyspace_id: u32) {
        self.tenant_cap_keyspaces.remove(&keyspace_id);
        self.tenant_cap_keyspace_last_access.remove(&keyspace_id);
        self.remove_keyspace(keyspace_id);
    }

    fn maybe_cleanup_tenant_cap_keyspaces(&self, now: Instant) {
        let cleanup_interval = Duration::from_secs(TENANT_CAP_KEYSPACE_CLEANUP_INTERVAL_SECS);
        {
            let last_cleanup = self.last_tenant_cap_cleanup.read().unwrap();
            if now.saturating_duration_since(*last_cleanup) < cleanup_interval {
                return;
            }
        }
        {
            let mut last_cleanup = self.last_tenant_cap_cleanup.write().unwrap();
            if now.saturating_duration_since(*last_cleanup) < cleanup_interval {
                return;
            }
            *last_cleanup = now;
        }

        let idle_ttl =
            Duration::from_millis(self.tenant_cap_idle_ttl_ms.load(Ordering::Acquire).max(1));
        let stale_keyspaces: Vec<u32> = self
            .tenant_cap_keyspace_last_access
            .iter()
            .filter_map(|entry| {
                let keyspace_id = *entry.key();
                if now.saturating_duration_since(*entry.value()) < idle_ttl {
                    return None;
                }
                if self.keyspace_inflight_us(keyspace_id) == 0 {
                    Some(keyspace_id)
                } else {
                    None
                }
            })
            .collect();

        for keyspace_id in stale_keyspaces {
            // Serialize the removal against concurrent recreation / disable so
            // teardown and (re)creation of the same keyspace cannot interleave.
            let _guard = self.tenant_cap_lifecycle.lock().unwrap();
            let still_stale = self
                .tenant_cap_keyspace_last_access
                .get(&keyspace_id)
                .map(|last_access| now.saturating_duration_since(*last_access.value()) >= idle_ttl)
                .unwrap_or(false);
            if !still_stale {
                continue;
            }
            if self.keyspace_inflight_us(keyspace_id) == 0 {
                self.remove_tenant_cap_keyspace(keyspace_id);
            }
        }
    }

    /// Check if burst is allowed for a keyspace.
    ///
    /// # Arguments
    /// * `keyspace_id` - Keyspace ID
    ///
    /// # Returns
    /// `Ok(())` if burst is allowed, otherwise the condition that denied it.
    fn check_burst_allowed(&self, keyspace_id: u32) -> Result<(), BurstDenyReason> {
        if !self.burst_enabled.load(Ordering::Acquire) {
            return Err(BurstDenyReason::BurstDisabled);
        }

        if !self.is_keyspace_burst_allowed(keyspace_id) {
            return Err(BurstDenyReason::KeyspaceBurstDisabled);
        }

        let monitor = self.cpu_usage_monitor.read().unwrap();
        let global_ratio = monitor.global_usage_ratio;

        let burst_threshold = f64::from_bits(self.burst_threshold_bits.load(Ordering::Acquire));

        // Burst conditions:
        // 1. Global CPU usage < burst_threshold (default 90%)
        // 2. Global bucket has available tokens
        if global_ratio.partial_cmp(&burst_threshold) != Some(std::cmp::Ordering::Less) {
            return Err(BurstDenyReason::GlobalRatioExceeded);
        }

        if self.global_bucket.available() == 0 {
            return Err(BurstDenyReason::GlobalBucketEmpty);
        }

        Ok(())
    }

    fn try_allocate_global_bucket_only(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
    ) -> Result<(), ThrottleError> {
        if !self.global_bucket.try_allocate(estimated_cpu_us) {
            CPU_THROTTLE_ALLOCATIONS
                .with_label_values(&["global_exhausted"])
                .inc();

            let debug_enabled = self.debug_enabled.load(Ordering::Acquire);
            let should_log = if debug_enabled {
                true
            } else {
                let mut last_log = self.last_global_exhausted_log.write().unwrap();
                if last_log.elapsed() >= Duration::from_secs(5) {
                    *last_log = Instant::now();
                    true
                } else {
                    false
                }
            };

            if should_log {
                tikv_util::warn!(
                    "[CPU throttle] Global CPU exhausted";
                    "keyspace_id" => keyspace_id,
                    "estimated_cpu_us" => estimated_cpu_us,
                    "global_available" => self.global_bucket.available(),
                    "global_capacity" => self.global_bucket.capacity(),
                );
            }
            return Err(ThrottleError::GlobalCpuExhausted);
        }

        CPU_THROTTLE_ALLOCATIONS
            .with_label_values(&["success"])
            .inc();
        update_global_bucket_metrics(&self.global_bucket);
        Ok(())
    }

    fn try_allocate_global_only(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<(CpuTokenHandle, bool), ThrottleError> {
        self.try_allocate_global_bucket_only(keyspace_id, estimated_cpu_us)?;
        let handle = CpuTokenHandle {
            keyspace_id,
            allocated_us: AtomicU64::new(estimated_cpu_us),
            keyspace_allocated_us: AtomicU64::new(0),
            global_allocated_us: AtomicU64::new(estimated_cpu_us),
            actual_used_us: AtomicU64::new(0),
            global_bucket: self.global_bucket.clone(),
            keyspace_bucket: None,
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: AtomicBool::new(false),
            request_deadline,
            skip_throttle: false,
            keyspace_burst_allowed: self.is_keyspace_burst_allowed(keyspace_id),
            cpu_usage_monitor: Some(self.cpu_usage_monitor.clone()),
            config: Some(self.build_handle_config()),
            adaptive_estimator: self.adaptive_estimator.clone(),
            burst_usage_token: None,
            burst_usage_tracker: Some(self.burst_usage_tracker.clone()),
            last_dynamic_burst_log: Some(self.last_dynamic_burst_log.clone()),
            last_periodic_allocation_log: Some(self.last_periodic_allocation_log.clone()),
        };
        handle.update_available_metrics();
        Ok((handle, false))
    }

    /// Returns the next retry deadline when fair read-pool dispatch is blocked.
    /// `None` means the global bucket currently has positive tokens.
    pub fn fair_global_cpu_dispatch_retry_at(&self) -> Option<Instant> {
        if self.global_bucket.refill() {
            update_global_bucket_metrics(&self.global_bucket);
        }
        if self.global_bucket.has_positive_available_with_debt() {
            None
        } else {
            Some(self.global_bucket.next_refill_deadline())
        }
    }

    /// Returns the exact bucket refill deadline relevant to an allocation
    /// failure.
    pub fn allocation_retry_at(&self, keyspace_id: u32, error: &ThrottleError) -> Instant {
        match error {
            ThrottleError::LowPriorityDagGlobalCpuExhausted => {
                self.low_priority_dag_bucket.next_refill_deadline()
            }
            ThrottleError::KeyspaceCpuExhausted => self
                .keyspace_buckets
                .get(&keyspace_id)
                .map(|bucket| bucket.next_refill_deadline())
                .unwrap_or_else(|| self.global_bucket.next_refill_deadline()),
            ThrottleError::KeyspaceFairShareExceeded { retry_at }
            | ThrottleError::KeyspaceCpuExhaustedNoBurst { retry_at } => *retry_at,
            ThrottleError::GlobalCpuExhausted | ThrottleError::RequestTimeout => {
                self.global_bucket.next_refill_deadline()
            }
        }
    }

    fn keyspace_or_earlier_retry_at(&self, keyspace_id: u32, alternative: Instant) -> Instant {
        self.keyspace_buckets
            .get(&keyspace_id)
            .map(|bucket| bucket.next_refill_deadline().min(alternative))
            .unwrap_or(alternative)
    }

    fn burst_denied_retry_at(&self, keyspace_id: u32, reason: BurstDenyReason) -> Instant {
        let alternative = match reason {
            BurstDenyReason::GlobalBucketEmpty | BurstDenyReason::IntervalLimitExceeded => {
                Some(self.global_bucket.next_refill_deadline())
            }
            BurstDenyReason::BurstBudgetExhausted => {
                Some(self.burst_usage_tracker.next_window_deadline())
            }
            BurstDenyReason::BurstDisabled
            | BurstDenyReason::KeyspaceBurstDisabled
            | BurstDenyReason::GlobalRatioExceeded => None,
        };
        alternative.map_or_else(
            || {
                self.keyspace_buckets
                    .get(&keyspace_id)
                    .map(|bucket| bucket.next_refill_deadline())
                    .unwrap_or_else(|| self.global_bucket.next_refill_deadline())
            },
            |deadline| self.keyspace_or_earlier_retry_at(keyspace_id, deadline),
        )
    }

    /// Mark this manager as owned by a keyspace-fair global CPU gate.
    pub fn register_fair_global_cpu_gate(&self) {
        {
            let _guard = self.dynamic_adjustment_mutex.lock().unwrap();
            self.fair_global_cpu_gate_registered
                .store(true, Ordering::Release);
            let dynamic_adjustment_was_enabled = self
                .dynamic_adjustment_enabled
                .swap(false, Ordering::AcqRel);
            self.restore_dynamic_refill_rates();
            if dynamic_adjustment_was_enabled {
                warn!(
                    "[CPU throttle] Dynamic adjustment disabled for fair global CPU gate";
                );
            }
        }
        let (scale, _) = unpack_fair_charge_state(self.fair_charge_state.load(Ordering::Acquire));
        self.global_bucket.set_metric_charge_scale(scale);
        update_global_bucket_metrics(&self.global_bucket);
        update_global_bucket_limit_metrics(&self.global_bucket);
        // The effective-ratio gauge depends on whether a gate exists, and
        // registration can happen after the ratios were published.
        self.publish_read_cpu_ratio_metrics();
    }

    /// Initialize the startup ratio used as the fair gate's token-space anchor.
    ///
    /// Both the read limiter and the usage subscriber can observe the manager
    /// during startup. Only the first initialization wins, so a later observer
    /// cannot reset an already-applied fairness override.
    pub(crate) fn initialize_fair_max_read_cpu_ratio(&self, startup_ratio: f64) {
        let _guard = self.fair_charge_update_mutex.lock().unwrap();
        if self.startup_max_read_cpu_ratio_bits.load(Ordering::Acquire) == 0.0f64.to_bits() {
            self.update_fair_charge_scale(startup_ratio, startup_ratio);
        }
    }

    /// Override the fair global CPU gate's effective read CPU ratio.
    ///
    /// `0` restores the startup ratio. The caller validates non-zero values;
    /// this method only translates the configured ratio into the charge scale
    /// anchored at the startup token-bucket ratio.
    pub fn update_fair_max_read_cpu_ratio(&self, max_read_cpu_ratio: f64) {
        let _guard = self.fair_charge_update_mutex.lock().unwrap();
        let startup_ratio =
            f64::from_bits(self.startup_max_read_cpu_ratio_bits.load(Ordering::Acquire));
        if startup_ratio <= 0.0 {
            warn!(
                "[CPU throttle] Fair max read CPU ratio update ignored before startup ratio initialization";
                "max_read_cpu_ratio" => max_read_cpu_ratio,
            );
            return;
        }
        let effective_ratio = if max_read_cpu_ratio == 0.0 {
            startup_ratio
        } else {
            max_read_cpu_ratio
        };
        self.update_fair_charge_scale(startup_ratio, effective_ratio);
    }

    /// Charge actual fair read-pool dispatch CPU usage to the global bucket.
    ///
    /// TODO(#5480): intermediate fair-gate state. This dispatch-level path
    /// replaces the legacy DAG admission while keyspace fairness is enabled:
    /// tenant caps, keyspace buckets, PD service-limit quotas, low-priority DAG
    /// sub-bucket, burst fair allocation, keyspace fair-share logic, and
    /// request-level adaptive estimation are not consulted here.
    ///
    /// Returns `true` when the charge creates global debt and the caller should
    /// yield.
    ///
    /// The charge is scaled by the current fair-gate charge scale, so the
    /// fairness read-CPU-ratio override takes effect here rather than by
    /// resizing the bucket. See
    /// [`CpuThrottleManager::update_fair_max_read_cpu_ratio`].
    pub fn consume_fair_global_cpu(&self, cpu_delta_us: u64) -> bool {
        self.global_bucket.refill();
        if cpu_delta_us > 0 {
            let tokens = accumulate_fair_charge(cpu_delta_us, &self.fair_charge_state);
            if tokens > 0 {
                self.global_bucket.allocate_with_debt(tokens);
            }
        }
        update_global_bucket_metrics(&self.global_bucket);
        self.global_bucket.debt() > 0
    }

    /// Hot-update the effective read CPU ratio for the fair global CPU gate.
    ///
    /// Instead of resizing the global bucket, this changes the exchange rate
    /// between measured CPU and tokens. Because the bucket's capacity and
    /// refill rate stay in token space, scaling the charge divides both the
    /// sustained rate and the burst depth by the same factor:
    ///
    /// ```text
    /// effective sustained CPU = refill_rate / scale
    /// effective burst depth   = capacity    / scale
    /// ```
    ///
    /// Both land exactly on `effective_ratio`, so this is equivalent to a full
    /// resize while remaining a single atomic store with no reconciliation of
    /// in-flight state.
    ///
    /// Tokens already in the bucket keep their token value, so an outstanding
    /// balance or debt is retroactively reinterpreted at the new rate. Existing
    /// debt is not bounded by one bucket capacity: `allocate_with_debt` lets
    /// the balance go arbitrarily negative, so a deeply indebted bucket
    /// needs proportionally more refill intervals to recover.
    /// Reinterpretation does not create debt, it only rescales what is
    /// already owed.
    ///
    /// This only affects the fair gate. The legacy DAG admission path compares
    /// token counts against measured CPU in several places and still requires a
    /// restart to change its ratio.
    fn update_fair_charge_scale(&self, base_ratio: f64, effective_ratio: f64) {
        let scale = fair_charge_scale_from_ratios(base_ratio, effective_ratio);
        // One store publishes the new rate and discards the remainder that
        // belonged to the old one. A remainder accumulated under the old scale
        // has no meaning under the new one, and splitting this into two writes
        // would let a concurrent charge commit a remainder across the boundary.
        let previous_state = self
            .fair_charge_state
            .swap(pack_fair_charge_state(scale, 0), Ordering::AcqRel);
        let (previous_scale, _) = unpack_fair_charge_state(previous_state);

        self.startup_max_read_cpu_ratio_bits
            .store(base_ratio.to_bits(), Ordering::Release);
        self.requested_max_read_cpu_ratio_bits
            .store(effective_ratio.to_bits(), Ordering::Release);
        let gate_registered = self.fair_global_cpu_gate_registered.load(Ordering::Acquire);
        self.global_bucket
            .set_metric_charge_scale(if gate_registered {
                scale
            } else {
                FAIR_CHARGE_SCALE_ONE
            });
        // Republish all global bucket gauges so their `_us` values immediately
        // reflect the new token-to-CPU exchange rate.
        update_global_bucket_metrics(&self.global_bucket);
        update_global_bucket_limit_metrics(&self.global_bucket);
        self.publish_read_cpu_ratio_metrics();
        CPU_THROTTLE_FAIR_CHARGE_SCALE.set(scale as f64 / FAIR_CHARGE_SCALE_ONE as f64);

        if previous_scale != scale {
            info!(
                "[CPU throttle] Fair gate charge scale updated";
                "base_max_read_cpu_ratio" => base_ratio,
                "effective_max_read_cpu_ratio" => effective_ratio,
                "scale" => scale as f64 / FAIR_CHARGE_SCALE_ONE as f64,
                "fair_gate_registered" => gate_registered,
                "in_force" => gate_registered,
            );
        }
    }

    /// Current fair-gate charge scale as a float. Test and metrics helper.
    pub fn fair_charge_scale(&self) -> f64 {
        let (scale, _) = unpack_fair_charge_state(self.fair_charge_state.load(Ordering::Acquire));
        scale as f64 / FAIR_CHARGE_SCALE_ONE as f64
    }

    /// Publish both ratio gauges from stored state.
    ///
    /// The effective gauge reports what is actually in force: without a
    /// registered fair gate the override enforces nothing, so publishing the
    /// requested value there would claim a throttle that does not exist.
    fn publish_read_cpu_ratio_metrics(&self) {
        CPU_THROTTLE_REQUESTED_MAX_READ_CPU_RATIO.set(f64::from_bits(
            self.requested_max_read_cpu_ratio_bits
                .load(Ordering::Acquire),
        ));
        CPU_THROTTLE_EFFECTIVE_MAX_READ_CPU_RATIO.set(self.enforced_max_read_cpu_ratio());
    }

    /// The read CPU ratio the fair gate actually enforces.
    ///
    /// Derived from the startup ratio and the live charge scale rather than
    /// from the requested gauge. Without a registered fair gate nothing is
    /// enforced through the scale, so this reports the startup ratio.
    pub fn enforced_max_read_cpu_ratio(&self) -> f64 {
        let startup = f64::from_bits(self.startup_max_read_cpu_ratio_bits.load(Ordering::Acquire));
        if !self.fair_global_cpu_gate_registered.load(Ordering::Acquire) {
            return startup;
        }
        let scale = self.fair_charge_scale();
        if scale <= 0.0 {
            return startup;
        }
        startup / scale
    }

    #[cfg(test)]
    fn fair_charge_residue_q(&self) -> u64 {
        let (_, residue) = unpack_fair_charge_state(self.fair_charge_state.load(Ordering::Acquire));
        residue
    }

    /// Discard simulated fair-gate debt before leaving dry-run mode.
    pub(crate) fn reset_fair_global_cpu_after_dry_run(&self) {
        if !self.fair_global_cpu_gate_registered.load(Ordering::Acquire) {
            return;
        }
        self.global_bucket.reset_to_capacity();
        update_global_bucket_metrics(&self.global_bucket);
    }

    /// Try to allocate CPU tokens with burst support
    ///
    /// # Arguments
    /// * `keyspace_id` - Keyspace ID
    /// * `estimated_cpu_us` - Estimated CPU time in microseconds
    /// * `request_deadline` - Request deadline for timeout checks
    ///
    /// # Returns
    /// `Ok((CpuTokenHandle, is_burst))` if allocation succeeded,
    /// `Err(ThrottleError)` otherwise
    pub fn try_allocate_with_burst(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<(CpuTokenHandle, bool), ThrottleError> {
        // Refill tokens based on elapsed time before attempting allocation
        self.global_bucket.refill();
        if let Some(bucket) = self.keyspace_buckets.get(&keyspace_id) {
            bucket.refill();
        }

        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_label = keyspace_label.as_ref();

        let keyspace_burst_allowed = self.is_keyspace_burst_allowed(keyspace_id);
        let debug_enabled = self.debug_enabled.load(Ordering::Acquire);
        let enforce_global_limit = self.enforce_global_limit_enabled.load(Ordering::Acquire);
        let burst_threshold = f64::from_bits(self.burst_threshold_bits.load(Ordering::Acquire));
        let fair_allocation_enabled = self.fair_allocation_enabled.load(Ordering::Acquire);
        let fair_allocation_threshold =
            f64::from_bits(self.fair_allocation_threshold_bits.load(Ordering::Acquire));
        let tenant_cap_active = if self.tenant_cap_quota_us().is_some() {
            let now = Instant::now();
            self.maybe_cleanup_tenant_cap_keyspaces(now);
            if !self.tenant_cap_keyspaces.contains(&keyspace_id)
                || !self.keyspace_buckets.contains_key(&keyspace_id)
            {
                // (Re)create through the lifecycle-locked helper, which
                // re-loads the quota inside the lock. If a concurrent disable
                // has already cleared the cap, this returns false and the
                // request is treated as cap-inactive instead of resurrecting a
                // bucket on the stale quota.
                self.ensure_tenant_cap_bucket(keyspace_id, now)
            } else {
                self.tenant_cap_keyspace_last_access
                    .insert(keyspace_id, now);
                true
            }
        } else {
            false
        };

        // 0. Skip throttle for keyspaces without configured quotas. In tenant
        // cap mode, a concurrent idle cleanup can remove the auto-created quota
        // after we marked the cap active; recreate it instead of falling back to
        // skip_throttle/global-only.
        if !self.keyspace_quotas.contains_key(&keyspace_id) {
            let tenant_cap_recreated =
                tenant_cap_active && self.ensure_tenant_cap_bucket(keyspace_id, Instant::now());
            if !tenant_cap_recreated {
                CPU_THROTTLE_ALLOCATIONS
                    .with_label_values(&["skip_no_quota"])
                    .inc();
                return Ok((
                    CpuTokenHandle {
                        keyspace_id,
                        allocated_us: AtomicU64::new(0),
                        keyspace_allocated_us: AtomicU64::new(0),
                        global_allocated_us: AtomicU64::new(0),
                        actual_used_us: AtomicU64::new(0),
                        global_bucket: self.global_bucket.clone(),
                        keyspace_bucket: None,
                        low_priority_bucket: None,
                        keyspace_inflight_state: None,
                        low_priority_allocated_us: AtomicU64::new(0),
                        low_priority_inflight_allocated_us: None,
                        is_burst: AtomicBool::new(false),
                        request_deadline,
                        skip_throttle: true,
                        keyspace_burst_allowed,
                        cpu_usage_monitor: None,
                        config: None,
                        adaptive_estimator: None,
                        burst_usage_token: None,
                        burst_usage_tracker: None,
                        last_dynamic_burst_log: Some(self.last_dynamic_burst_log.clone()),
                        last_periodic_allocation_log: Some(
                            self.last_periodic_allocation_log.clone(),
                        ),
                    },
                    false,
                ));
            }
        }

        // 1. Try to allocate from keyspace bucket first (hard isolation by default)
        let mut keyspace_bucket = self
            .keyspace_buckets
            .get(&keyspace_id)
            .map(|bucket| bucket.value().clone());
        if keyspace_bucket.is_none() && tenant_cap_active {
            self.ensure_tenant_cap_bucket(keyspace_id, Instant::now());
            keyspace_bucket = self
                .keyspace_buckets
                .get(&keyspace_id)
                .map(|bucket| bucket.value().clone());
        }

        if let Some(bucket) = keyspace_bucket {
            if bucket.try_allocate(estimated_cpu_us) {
                // Normal mode: keyspace allocation succeeded
                let mut global_tracked_us = 0;
                if enforce_global_limit {
                    if !self.global_bucket.try_allocate(estimated_cpu_us) {
                        // Global exhausted, rollback keyspace allocation
                        bucket.release(estimated_cpu_us);
                        CPU_THROTTLE_ALLOCATIONS
                            .with_label_values(&["global_exhausted"])
                            .inc();

                        let should_log = if debug_enabled {
                            true
                        } else {
                            let mut last_log = self.last_global_exhausted_log.write().unwrap();
                            if last_log.elapsed() >= Duration::from_secs(5) {
                                *last_log = Instant::now();
                                true
                            } else {
                                false
                            }
                        };

                        if should_log {
                            tikv_util::warn!(
                                "[CPU throttle] Global CPU exhausted";
                                "keyspace_id" => keyspace_id,
                                "estimated_cpu_us" => estimated_cpu_us,
                                "global_available" => self.global_bucket.available(),
                                "global_capacity" => self.global_bucket.capacity(),
                            );
                        }
                        return Err(ThrottleError::GlobalCpuExhausted);
                    }
                    global_tracked_us = estimated_cpu_us;
                } else if self.global_bucket.try_allocate(estimated_cpu_us) {
                    global_tracked_us = estimated_cpu_us;
                }

                // Normal mode: both allocations succeeded (or global check skipped)
                CPU_THROTTLE_ALLOCATIONS
                    .with_label_values(&["success"])
                    .inc();
                let handle = CpuTokenHandle {
                    keyspace_id,
                    allocated_us: AtomicU64::new(estimated_cpu_us),
                    keyspace_allocated_us: AtomicU64::new(estimated_cpu_us),
                    global_allocated_us: AtomicU64::new(global_tracked_us),
                    actual_used_us: AtomicU64::new(0),
                    global_bucket: self.global_bucket.clone(),
                    keyspace_bucket: Some(bucket.clone()),
                    low_priority_bucket: None,
                    keyspace_inflight_state: Some(self.keyspace_inflight_state(keyspace_id)),
                    low_priority_allocated_us: AtomicU64::new(0),
                    low_priority_inflight_allocated_us: None,
                    is_burst: AtomicBool::new(false),
                    request_deadline,
                    skip_throttle: false,
                    keyspace_burst_allowed,
                    cpu_usage_monitor: Some(self.cpu_usage_monitor.clone()),
                    config: Some(self.build_handle_config()),
                    adaptive_estimator: self.adaptive_estimator.clone(),
                    burst_usage_token: None,
                    burst_usage_tracker: Some(self.burst_usage_tracker.clone()),
                    last_dynamic_burst_log: Some(self.last_dynamic_burst_log.clone()),
                    last_periodic_allocation_log: Some(self.last_periodic_allocation_log.clone()),
                };
                handle.inflight_add(estimated_cpu_us);
                handle.update_available_metrics();
                return Ok((handle, false));
            }

            // Keyspace exhausted, check if burst is allowed
            let mut burst_budget_window = None;
            let mut burst_denied_reason = None;
            let mut burst_allowed = match self.check_burst_allowed(keyspace_id) {
                Ok(()) => true,
                Err(reason) => {
                    burst_denied_reason = Some(reason);
                    false
                }
            };
            if burst_allowed {
                match reserve_burst_budget(
                    &self.burst_usage_tracker,
                    self.global_bucket.base_refill_rate(),
                    burst_threshold,
                    estimated_cpu_us,
                ) {
                    Ok(window_id) => burst_budget_window = window_id,
                    Err(()) => {
                        burst_denied_reason = Some(BurstDenyReason::BurstBudgetExhausted);
                        burst_allowed = false;
                    }
                }
            }

            let mut interval_window_id = None;
            if burst_allowed {
                // Check per-interval allocation limit for global bucket
                match self
                    .global_bucket
                    .try_allocate_in_interval(estimated_cpu_us)
                {
                    Some(window_id) => interval_window_id = Some(window_id),
                    None => {
                        if let Some(window_id) = burst_budget_window {
                            self.burst_usage_tracker
                                .release_global(window_id, estimated_cpu_us);
                        }
                        burst_denied_reason = Some(BurstDenyReason::IntervalLimitExceeded);
                        burst_allowed = false;
                    }
                }
            }

            if burst_allowed {
                // Try to allocate from global bucket for burst mode
                if !self.global_bucket.try_allocate(estimated_cpu_us) {
                    // Rollback interval allocation since token allocation failed
                    if let Some(window_id) = interval_window_id {
                        self.global_bucket
                            .release_in_interval(estimated_cpu_us, window_id);
                    }
                    if let Some(window_id) = burst_budget_window {
                        self.burst_usage_tracker
                            .release_global(window_id, estimated_cpu_us);
                    }
                    CPU_THROTTLE_ALLOCATIONS
                        .with_label_values(&["global_exhausted"])
                        .inc();

                    let should_log = if debug_enabled {
                        true
                    } else {
                        let mut last_log = self.last_global_exhausted_log.write().unwrap();
                        if last_log.elapsed() >= Duration::from_secs(5) {
                            *last_log = Instant::now();
                            true
                        } else {
                            false
                        }
                    };

                    if should_log {
                        tikv_util::warn!(
                            "[CPU throttle] Global CPU exhausted during burst";
                            "keyspace_id" => keyspace_id,
                            "estimated_cpu_us" => estimated_cpu_us,
                            "global_available" => self.global_bucket.available(),
                            "global_capacity" => self.global_bucket.capacity(),
                        );
                    }
                    return Err(ThrottleError::KeyspaceCpuExhaustedNoBurst {
                        retry_at: self.keyspace_or_earlier_retry_at(
                            keyspace_id,
                            self.global_bucket.next_refill_deadline(),
                        ),
                    });
                }

                // Global allocation succeeded, setup burst with fair allocation check
                let burst_usage_token = if fair_allocation_enabled {
                    let window_id = burst_budget_window
                        .unwrap_or_else(|| self.burst_usage_tracker.current_window_id());
                    let current_burst_us = self
                        .burst_usage_tracker
                        .usage_in_window(keyspace_id, window_id);
                    let enforce_fair_allocation = {
                        let monitor = self.cpu_usage_monitor.read().unwrap();
                        monitor.global_usage_ratio >= fair_allocation_threshold
                    };

                    if enforce_fair_allocation {
                        let interval_budget_us = self.global_bucket.interval_allocation_limit();
                        if let Some(fair_share_us) =
                            self.calculate_fair_share(keyspace_id, interval_budget_us)
                        {
                            if current_burst_us.saturating_add(estimated_cpu_us) > fair_share_us {
                                self.global_bucket.release(estimated_cpu_us);
                                if let Some(window_id) = interval_window_id {
                                    self.global_bucket
                                        .release_in_interval(estimated_cpu_us, window_id);
                                }
                                if let Some(window_id) = burst_budget_window {
                                    self.burst_usage_tracker
                                        .release_global(window_id, estimated_cpu_us);
                                }
                                update_global_bucket_metrics(&self.global_bucket);
                                CPU_THROTTLE_ALLOCATIONS
                                    .with_label_values(&["fair_share_exceeded"])
                                    .inc();
                                CPU_THROTTLE_FAIR_ALLOCATIONS
                                    .with_label_values(&[keyspace_label])
                                    .inc();

                                let should_log = if debug_enabled {
                                    true
                                } else {
                                    let mut last_log =
                                        self.last_fair_share_exceeded_log.write().unwrap();
                                    if last_log.elapsed() >= Duration::from_secs(5) {
                                        *last_log = Instant::now();
                                        true
                                    } else {
                                        false
                                    }
                                };

                                if should_log {
                                    tikv_util::warn!(
                                        "[CPU throttle] Keyspace exceeded fair share";
                                        "keyspace_id" => keyspace_id,
                                        "fair_share_us" => fair_share_us,
                                        "current_burst_us" => current_burst_us,
                                        "estimated_cpu_us" => estimated_cpu_us,
                                        "interval_budget_us" => interval_budget_us,
                                    );
                                }
                                let retry_at = self.keyspace_or_earlier_retry_at(
                                    keyspace_id,
                                    self.burst_usage_tracker.next_window_deadline(),
                                );
                                return Err(ThrottleError::KeyspaceFairShareExceeded { retry_at });
                            }
                        }
                    }

                    self.burst_usage_tracker.add_usage_in_window(
                        keyspace_id,
                        window_id,
                        estimated_cpu_us,
                    );
                    Some(BurstUsageToken {
                        keyspace_id,
                        window_id,
                        allocated_us: estimated_cpu_us,
                        tracker: self.burst_usage_tracker.clone(),
                    })
                } else {
                    None
                };

                // Burst mode: keyspace exhausted but burst is allowed
                CPU_THROTTLE_ALLOCATIONS
                    .with_label_values(&["success"])
                    .inc();
                CPU_THROTTLE_BURST_COUNT
                    .with_label_values(&[keyspace_label])
                    .inc();
                CPU_THROTTLE_BURST_TOKENS_USED
                    .with_label_values(&[keyspace_label])
                    .inc_by(estimated_cpu_us);

                let should_log = if debug_enabled {
                    true
                } else {
                    let mut last_log = self.last_burst_activated_log.write().unwrap();
                    if last_log.elapsed() >= Duration::from_secs(5) {
                        *last_log = Instant::now();
                        true
                    } else {
                        false
                    }
                };

                if should_log {
                    info!(
                        "[CPU throttle] Burst mode activated";
                        "keyspace_id" => keyspace_id,
                        "estimated_cpu_us" => estimated_cpu_us,
                    );
                }
                let handle = CpuTokenHandle {
                    keyspace_id,
                    allocated_us: AtomicU64::new(estimated_cpu_us),
                    keyspace_allocated_us: AtomicU64::new(0),
                    global_allocated_us: AtomicU64::new(estimated_cpu_us),
                    actual_used_us: AtomicU64::new(0),
                    global_bucket: self.global_bucket.clone(),
                    keyspace_bucket: None,
                    low_priority_bucket: None,
                    keyspace_inflight_state: None,
                    low_priority_allocated_us: AtomicU64::new(0),
                    low_priority_inflight_allocated_us: None,
                    is_burst: AtomicBool::new(true),
                    request_deadline,
                    skip_throttle: false,
                    keyspace_burst_allowed,
                    cpu_usage_monitor: Some(self.cpu_usage_monitor.clone()),
                    config: Some(self.build_handle_config()),
                    adaptive_estimator: self.adaptive_estimator.clone(),
                    burst_usage_token,
                    burst_usage_tracker: Some(self.burst_usage_tracker.clone()),
                    last_dynamic_burst_log: Some(self.last_dynamic_burst_log.clone()),
                    last_periodic_allocation_log: Some(self.last_periodic_allocation_log.clone()),
                };
                handle.update_available_metrics();
                Ok((handle, true))
            } else {
                // Keyspace exhausted and burst not allowed
                CPU_THROTTLE_ALLOCATIONS
                    .with_label_values(&[if tenant_cap_active {
                        "tenant_cap_exhausted"
                    } else {
                        "keyspace_exhausted"
                    }])
                    .inc();

                let should_log = if debug_enabled {
                    true
                } else {
                    let mut last_log = self.last_keyspace_exhausted_log.write().unwrap();
                    if last_log.elapsed() >= Duration::from_secs(5) {
                        *last_log = Instant::now();
                        true
                    } else {
                        false
                    }
                };

                if should_log {
                    tikv_util::warn!(
                        "[CPU throttle] Keyspace CPU exhausted, burst not allowed";
                        "keyspace_id" => keyspace_id,
                        "estimated_cpu_us" => estimated_cpu_us,
                        "keyspace_available" => bucket.available(),
                        "keyspace_capacity" => bucket.capacity(),
                        "global_available" => self.global_bucket.available(),
                        "global_capacity" => self.global_bucket.capacity(),
                        "global_ratio" => self.cpu_usage_monitor.read().unwrap().global_usage_ratio,
                        "burst_denied_reason" => burst_denied_reason.map_or("unknown", BurstDenyReason::as_str),
                    );
                }
                let retry_at = self.burst_denied_retry_at(
                    keyspace_id,
                    burst_denied_reason.expect("denied burst must have a reason"),
                );
                Err(ThrottleError::KeyspaceCpuExhaustedNoBurst { retry_at })
            }
        } else if tenant_cap_active && self.tenant_cap_quota_us().is_some() {
            CPU_THROTTLE_ALLOCATIONS
                .with_label_values(&["tenant_cap_exhausted"])
                .inc();
            Err(ThrottleError::KeyspaceCpuExhaustedNoBurst {
                retry_at: self.global_bucket.next_refill_deadline(),
            })
        } else {
            self.try_allocate_global_only(keyspace_id, estimated_cpu_us, request_deadline)
        }
    }

    /// Try to allocate CPU tokens (non-burst version for backward
    /// compatibility)
    ///
    /// # Arguments
    /// * `keyspace_id` - Keyspace ID
    /// * `estimated_cpu_us` - Estimated CPU time in microseconds
    /// * `request_deadline` - Request deadline for timeout checks
    ///
    /// # Returns
    /// `Ok(CpuTokenHandle)` if allocation succeeded, `Err(ThrottleError)`
    /// otherwise
    pub fn try_allocate(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<CpuTokenHandle, ThrottleError> {
        self.try_allocate_with_burst(keyspace_id, estimated_cpu_us, request_deadline)
            .map(|(handle, _is_burst)| handle)
    }

    /// Try to allocate CPU tokens for a low-priority DAG request.
    ///
    /// The request must first pass the low-priority DAG global bucket, then
    /// pass the existing keyspace/global allocation path. Keyspaces without a
    /// configured quota still charge the shared global bucket in stage 2, so
    /// low-priority usage remains part of the underlying CPU budget. If the
    /// second stage fails, the low-priority allocation is rolled back
    /// immediately.
    pub fn try_allocate_low_priority_dag(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<CpuTokenHandle, ThrottleError> {
        self.try_allocate_low_priority_dag_with_burst(
            keyspace_id,
            estimated_cpu_us,
            request_deadline,
        )
        .map(|(handle, _is_burst)| handle)
    }

    pub fn try_allocate_low_priority_dag_with_burst(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<(CpuTokenHandle, bool), ThrottleError> {
        if !self.low_priority_dag_active() {
            return self.try_allocate_with_burst(keyspace_id, estimated_cpu_us, request_deadline);
        }

        self.low_priority_dag_bucket.refill();
        if !self.low_priority_dag_bucket.try_allocate(estimated_cpu_us) {
            CPU_THROTTLE_LOW_PRIORITY_DAG_ALLOCATIONS
                .with_label_values(&["low_priority_global_exhausted"])
                .inc();
            return Err(ThrottleError::LowPriorityDagGlobalCpuExhausted);
        }

        let allocation_result = if self.tenant_cap_quota_us().is_some()
            || self.keyspace_quotas.contains_key(&keyspace_id)
        {
            self.try_allocate_with_burst(keyspace_id, estimated_cpu_us, request_deadline)
        } else {
            self.global_bucket.refill();
            self.try_allocate_global_only(keyspace_id, estimated_cpu_us, request_deadline)
        };

        match allocation_result {
            Ok((mut handle, is_burst)) => {
                handle.low_priority_bucket = Some(self.low_priority_dag_bucket.clone());
                handle.low_priority_inflight_allocated_us =
                    Some(self.low_priority_dag_inflight_allocated_us.clone());
                handle.commit_low_priority_allocation(estimated_cpu_us);
                handle.update_available_metrics();
                CPU_THROTTLE_LOW_PRIORITY_DAG_ALLOCATIONS
                    .with_label_values(&["success"])
                    .inc();
                Ok((handle, is_burst))
            }
            Err(err) => {
                self.low_priority_dag_bucket.release(estimated_cpu_us);
                self.update_low_priority_dag_bucket_metrics();
                CPU_THROTTLE_LOW_PRIORITY_DAG_ALLOCATIONS
                    .with_label_values(&[throttle_error_label(&err)])
                    .inc();
                Err(err)
            }
        }
    }

    /// Update CPU usage from resource monitor
    ///
    /// # Arguments
    /// * `global_ratio` - Global CPU usage ratio (0.0 - 1.0)
    /// * `keyspace_ratios` - Per-keyspace CPU usage ratios
    pub fn update_usage(
        &self,
        global_ratio: f64,
        keyspace_ratios: std::collections::HashMap<u32, f64>,
    ) {
        let mut monitor = self.cpu_usage_monitor.write().unwrap();
        monitor.global_usage_ratio = global_ratio;
        monitor.keyspace_usage_ratios = keyspace_ratios;
        monitor.last_update = Instant::now();

        // Rate limit logging: only log once every 5 seconds
        const LOG_INTERVAL_SECS: u64 = 5;
        if monitor
            .last_update
            .duration_since(monitor.last_log_time)
            .as_secs()
            >= LOG_INTERVAL_SECS
        {
            info!(
                "[CPU throttle] Manager updated usage";
                "global_ratio" => global_ratio,
                "keyspace_count" => monitor.keyspace_usage_ratios.len(),
            );
            monitor.last_log_time = monitor.last_update;
        }
    }

    /// Update only global CPU usage ratio without clearing keyspace ratios
    ///
    /// # Arguments
    /// * `global_ratio` - Global CPU usage ratio (0.0 - 1.0)
    pub fn update_global_ratio(&self, global_ratio: f64) {
        let mut monitor = self.cpu_usage_monitor.write().unwrap();
        monitor.global_usage_ratio = global_ratio;
        monitor.last_update = Instant::now();

        // Rate limit logging: only log once every 5 seconds
        const LOG_INTERVAL_SECS: u64 = 5;
        if monitor
            .last_update
            .duration_since(monitor.last_log_time)
            .as_secs()
            >= LOG_INTERVAL_SECS
        {
            info!(
                "[CPU throttle] Manager updated global ratio";
                "global_ratio" => global_ratio,
                "keyspace_count" => monitor.keyspace_usage_ratios.len(),
            );
            monitor.last_log_time = monitor.last_update;
        }
    }

    /// Adjust refill rates based on CPU usage (for dynamic adjustment feature)
    pub fn adjust_refill_rates(&self) {
        let _guard = self.dynamic_adjustment_mutex.lock().unwrap();
        if !self.dynamic_adjustment_enabled.load(Ordering::Acquire)
            || self.fair_global_cpu_gate_registered.load(Ordering::Acquire)
        {
            return;
        }

        let high_watermark = f64::from_bits(self.high_watermark_bits.load(Ordering::Acquire));
        let low_watermark = f64::from_bits(self.low_watermark_bits.load(Ordering::Acquire));
        let monitor = self.cpu_usage_monitor.read().unwrap();
        let global_ratio = monitor.global_usage_ratio;

        // Adjust global bucket refill rate
        if global_ratio > high_watermark {
            // High watermark: reduce refill rate
            let reduction = (global_ratio - high_watermark) / (1.0 - high_watermark);
            let new_rate =
                (self.global_bucket.base_refill_rate() as f64 * (1.0 - reduction * 0.5)) as u64;
            self.global_bucket.set_refill_rate(new_rate);
            update_global_bucket_limit_metrics(&self.global_bucket);

            CPU_THROTTLE_REFILL_RATE_ADJUSTMENTS
                .with_label_values(&["global", "decrease"])
                .inc();

            // Rate limit logging: only log once every 5 seconds
            let should_log = {
                let mut last_log = self.last_global_refill_rate_log.write().unwrap();
                if last_log.elapsed() >= Duration::from_secs(5) {
                    *last_log = Instant::now();
                    true
                } else {
                    false
                }
            };

            if should_log {
                info!(
                    "[CPU throttle] Reducing global refill rate";
                    "global_ratio" => global_ratio,
                    "new_rate" => new_rate,
                );
            }
        } else if global_ratio < low_watermark {
            // Low watermark: restore base rate
            self.global_bucket.restore_base_rate();
            update_global_bucket_limit_metrics(&self.global_bucket);
            CPU_THROTTLE_REFILL_RATE_ADJUSTMENTS
                .with_label_values(&["global", "increase"])
                .inc();

            // Rate limit logging: only log once every 5 seconds
            let should_log = {
                let mut last_log = self.last_global_refill_rate_log.write().unwrap();
                if last_log.elapsed() >= Duration::from_secs(5) {
                    *last_log = Instant::now();
                    true
                } else {
                    false
                }
            };

            if should_log {
                info!(
                    "[CPU throttle] Restoring global refill rate";
                    "global_ratio" => global_ratio,
                );
            }
        }

        if self.low_priority_dag_bucket.capacity() > 0 {
            let scaled_rate = self.current_low_priority_dag_refill_rate();
            self.low_priority_dag_bucket.set_refill_rate(scaled_rate);
            CPU_THROTTLE_LOW_PRIORITY_DAG_REFILL_RATE.set(clamp_gauge_value(scaled_rate));
        } else {
            CPU_THROTTLE_LOW_PRIORITY_DAG_REFILL_RATE.set(0);
        }

        let refill_interval_ms = self.refill_interval_ms.max(1);
        let sum_quotas_millicores: u64 = self
            .keyspace_quotas
            .iter()
            .fold(0, |acc, entry| acc.saturating_add(*entry.value()));
        let sum_quotas_us = sum_quotas_millicores.saturating_mul(refill_interval_ms);
        let global_refill_rate_us = self.global_bucket.base_refill_rate();
        let effective_refill_budget_us = std::cmp::max(global_refill_rate_us, sum_quotas_us);

        // Adjust keyspace bucket refill rates
        for entry in self.keyspace_buckets.iter() {
            let keyspace_id = *entry.key();
            let bucket = entry.value();
            let keyspace_label = keyspace_label(keyspace_id);
            let keyspace_label = keyspace_label.as_ref();

            if let Some(&keyspace_ratio) = monitor.keyspace_usage_ratios.get(&keyspace_id) {
                let quota_millicores = self
                    .keyspace_quotas
                    .get(&keyspace_id)
                    .map(|q| *q.value())
                    .unwrap_or(0);

                // Align keyspace adjustment conditions with global watermarks.
                if quota_millicores > 0 {
                    let quota_us = quota_millicores.saturating_mul(refill_interval_ms);

                    // Skip adjustment for unlimited quota or an invalid refill budget.
                    if quota_us == u64::MAX || effective_refill_budget_us == 0 {
                        continue;
                    }

                    let quota_ratio = quota_us as f64 / effective_refill_budget_us as f64;
                    if quota_ratio <= 0.0 {
                        continue;
                    }

                    let normalized_ratio = keyspace_ratio / quota_ratio;

                    if normalized_ratio > high_watermark {
                        // Keyspace usage exceeds high watermark, reduce refill rate.
                        let reduction =
                            (normalized_ratio - high_watermark) / (1.0 - high_watermark);
                        let new_rate =
                            (bucket.base_refill_rate() as f64 * (1.0 - reduction * 0.5)) as u64;
                        bucket.set_refill_rate(new_rate);
                        CPU_THROTTLE_KEYSPACE_REFILL_RATE
                            .with_label_values(&[keyspace_label])
                            .set(clamp_gauge_value(new_rate));
                        CPU_THROTTLE_REFILL_RATE_ADJUSTMENTS
                            .with_label_values(&["keyspace", "decrease"])
                            .inc();
                    } else if normalized_ratio < low_watermark {
                        // Restore base rate when below low watermark.
                        bucket.restore_base_rate();
                        CPU_THROTTLE_KEYSPACE_REFILL_RATE
                            .with_label_values(&[keyspace_label])
                            .set(clamp_gauge_value(bucket.base_refill_rate()));
                        CPU_THROTTLE_REFILL_RATE_ADJUSTMENTS
                            .with_label_values(&["keyspace", "increase"])
                            .inc();
                    }
                }
            }
        }
    }

    /// Update dynamic adjustment settings (hot update for watermarks)
    pub fn update_dynamic_adjustment(
        &self,
        enabled: bool,
        high_watermark: f64,
        low_watermark: f64,
    ) {
        let _guard = self.dynamic_adjustment_mutex.lock().unwrap();
        self.high_watermark_bits
            .store(high_watermark.to_bits(), Ordering::Release);
        self.low_watermark_bits
            .store(low_watermark.to_bits(), Ordering::Release);
        if self.fair_global_cpu_gate_registered.load(Ordering::Acquire) {
            self.dynamic_adjustment_enabled
                .store(false, Ordering::Release);
            self.restore_dynamic_refill_rates();
            if enabled {
                warn!(
                    "[CPU throttle] Dynamic adjustment ignored for fair global CPU gate";
                );
            }
        } else {
            self.dynamic_adjustment_enabled
                .store(enabled, Ordering::Release);
        }
    }

    /// Update low-priority DAG global throttle settings (hot update).
    pub fn update_low_priority_dag_settings(&self, enabled: bool, ratio: f64) {
        self.low_priority_dag_enabled
            .store(enabled, Ordering::Release);
        self.low_priority_dag_ratio_bits
            .store(ratio.to_bits(), Ordering::Release);

        let refill_rate_us =
            low_priority_dag_quota_us(self.low_priority_dag_global_cpu_us, ratio, enabled);
        let capacity_us = refill_rate_to_capacity(
            refill_rate_us,
            self.refill_interval_ms,
            self.burst_window_ms,
        );
        let inflight_us = self
            .low_priority_dag_inflight_allocated_us
            .load(Ordering::Acquire);
        self.low_priority_dag_bucket
            .update_quota_reserving_inflight(capacity_us, refill_rate_us, inflight_us);
        let scaled_rate = self.current_low_priority_dag_refill_rate();
        self.low_priority_dag_bucket.set_refill_rate(scaled_rate);
        self.update_low_priority_dag_bucket_metrics();
    }

    /// Update fair allocation settings (hot update)
    pub fn update_fair_allocation(&self, enabled: bool, threshold: f64) {
        self.fair_allocation_enabled
            .store(enabled, Ordering::Release);
        self.fair_allocation_threshold_bits
            .store(threshold.to_bits(), Ordering::Release);
    }

    /// Update burst settings (hot update)
    pub fn update_burst_settings(
        &self,
        enabled: bool,
        threshold: f64,
        keyspace_burst_enabled: std::collections::HashMap<u32, bool>,
    ) {
        self.burst_enabled.store(enabled, Ordering::Release);
        self.burst_threshold_bits
            .store(threshold.to_bits(), Ordering::Release);
        let interval_limit = burst_cap_us(self.global_bucket.base_refill_rate(), threshold);
        self.global_bucket
            .set_interval_allocation_limit(interval_limit);
        self.update_keyspace_burst_enabled(keyspace_burst_enabled);
    }

    /// Update enforce global limit setting (hot update)
    pub fn update_enforce_global_limit(&self, enabled: bool) {
        self.enforce_global_limit_enabled
            .store(enabled, Ordering::Release);
    }

    pub fn update_pd_service_limit_enabled(&self, enabled: bool) {
        self.pd_service_limit_enabled
            .store(enabled, Ordering::Release);
    }

    fn service_limit_can_mutate_keyspace(&self) -> bool {
        self.pd_service_limit_enabled() && self.tenant_cap_quota_us().is_none()
    }

    /// Update uniform tenant cap setting (hot update).
    pub fn update_tenant_max_cpu_ratio(&self, ratio: f64) {
        let quota_us =
            CpuThrottleConfig::compute_tenant_cap_quota_us_for(ratio, self.refill_interval_ms)
                .unwrap_or(0);

        // Serialize the whole transition against concurrent bucket
        // (re)creation in the allocation path. Without this, an in-flight
        // request that already snapshotted a non-zero quota could reinsert a
        // keyspace bucket after the disable branch below removed it, leaving
        // the node throttling on a stale quota after the cap was turned off.
        let _guard = self.tenant_cap_lifecycle.lock().unwrap();

        if quota_us > 0 {
            let now = Instant::now();
            // Publish the new quota first so any concurrent recreation that
            // wins the lock after us observes the up-to-date value.
            self.tenant_cap_quota_us.store(quota_us, Ordering::Release);
            self.tenant_max_cpu_ratio_bits
                .store(ratio.to_bits(), Ordering::Release);
            let mut keyspaces_to_update: std::collections::HashSet<u32> = self
                .keyspace_quotas
                .iter()
                .map(|entry| *entry.key())
                .collect();
            keyspaces_to_update.extend(self.tenant_cap_keyspaces.iter().map(|entry| *entry.key()));
            for keyspace_id in keyspaces_to_update {
                self.update_keyspace_quota_inner(keyspace_id, quota_us);
                self.tenant_cap_keyspaces.insert(keyspace_id);
                self.tenant_cap_keyspace_last_access
                    .insert(keyspace_id, now);
            }
            CPU_THROTTLE_TENANT_CAP_QUOTA_US.set(clamp_gauge_value(quota_us));
            CPU_THROTTLE_TENANT_CAP_RATIO.set(ratio);
        } else {
            // Clear the quota first: once this store is visible, any recreation
            // that acquires the lifecycle lock after us re-loads quota == 0 and
            // refuses to resurrect a bucket.
            self.tenant_cap_quota_us.store(0, Ordering::Release);
            self.tenant_max_cpu_ratio_bits
                .store(ratio.to_bits(), Ordering::Release);
            let keyspaces: Vec<u32> = self
                .tenant_cap_keyspaces
                .iter()
                .map(|entry| *entry.key())
                .collect();
            for keyspace_id in keyspaces {
                self.remove_tenant_cap_keyspace(keyspace_id);
            }
            self.tenant_cap_keyspaces.clear();
            self.tenant_cap_keyspace_last_access.clear();
            CPU_THROTTLE_TENANT_CAP_QUOTA_US.set(0);
            CPU_THROTTLE_TENANT_CAP_RATIO.set(ratio);
        }
    }

    /// Update idle TTL for auto-created tenant cap keyspaces (hot update).
    pub fn update_tenant_cap_idle_ttl(&self, idle_ttl_ms: u64) {
        self.tenant_cap_idle_ttl_ms
            .store(idle_ttl_ms.max(1), Ordering::Release);
        let now = Instant::now();
        let next_cleanup_due = now
            .checked_sub(Duration::from_secs(
                TENANT_CAP_KEYSPACE_CLEANUP_INTERVAL_SECS + 1,
            ))
            .unwrap_or(now);
        *self.last_tenant_cap_cleanup.write().unwrap() = next_cleanup_due;
    }

    /// Update debug logging setting (hot update)
    pub fn update_debug(&self, enabled: bool) {
        self.debug_enabled.store(enabled, Ordering::Release);
    }

    fn update_keyspace_burst_enabled(&self, new_values: std::collections::HashMap<u32, bool>) {
        let new_keyspaces: std::collections::HashSet<u32> = new_values.keys().copied().collect();

        for (keyspace_id, enabled) in new_values {
            self.keyspace_burst_enabled.insert(keyspace_id, enabled);
        }

        let mut to_remove = Vec::new();
        for entry in self.keyspace_burst_enabled.iter() {
            let keyspace_id = *entry.key();
            if !new_keyspaces.contains(&keyspace_id) {
                to_remove.push(keyspace_id);
            }
        }
        for keyspace_id in to_remove {
            self.keyspace_burst_enabled.remove(&keyspace_id);
        }
    }

    fn build_handle_config_value(&self) -> CpuThrottleConfig {
        let mut config = (*self.config).clone();
        config.enable_fair_allocation = self.fair_allocation_enabled.load(Ordering::Acquire);
        config.fair_allocation_threshold =
            f64::from_bits(self.fair_allocation_threshold_bits.load(Ordering::Acquire));
        config.enable_burst = self.burst_enabled.load(Ordering::Acquire);
        config.burst_threshold = f64::from_bits(self.burst_threshold_bits.load(Ordering::Acquire));
        config.enforce_global_limit = self.enforce_global_limit_enabled.load(Ordering::Acquire);
        config.tenant_max_cpu_ratio = self.tenant_max_cpu_ratio();
        config.debug = self.debug_enabled.load(Ordering::Acquire);
        config
    }

    fn build_handle_config(&self) -> Arc<CpuThrottleConfig> {
        Arc::new(self.build_handle_config_value())
    }

    /// Perform refill on all buckets
    pub fn refill_all(&self) {
        self.global_bucket.refill();
        update_global_bucket_metrics(&self.global_bucket);
        self.low_priority_dag_bucket.refill();
        CPU_THROTTLE_LOW_PRIORITY_DAG_BUCKET_AVAILABLE
            .set(clamp_gauge_value(self.low_priority_dag_bucket.available()));
        for entry in self.keyspace_buckets.iter() {
            entry.value().refill();
            let keyspace_label = keyspace_label(*entry.key());
            CPU_THROTTLE_KEYSPACE_BUCKET_AVAILABLE
                .with_label_values(&[keyspace_label.as_ref()])
                .set(clamp_gauge_value(entry.value().available()));
        }
    }

    /// Get global bucket available tokens
    pub fn global_available(&self) -> u64 {
        self.global_bucket.available()
    }

    /// Get global bucket debt tokens.
    pub fn global_debt(&self) -> u64 {
        self.global_bucket.debt()
    }

    /// Get global bucket burst capacity in microseconds.
    pub fn global_capacity(&self) -> u64 {
        self.global_bucket.capacity()
    }

    /// Get the global refill budget in microseconds per refill interval.
    pub fn global_refill_rate(&self) -> u64 {
        self.global_bucket.base_refill_rate()
    }

    /// Get low-priority DAG bucket available tokens.
    pub fn low_priority_dag_available(&self) -> u64 {
        self.low_priority_dag_bucket.available()
    }

    /// Get low-priority DAG bucket burst capacity in microseconds.
    pub fn low_priority_dag_capacity(&self) -> u64 {
        self.low_priority_dag_bucket.capacity()
    }

    /// Whether the low-priority DAG global throttle is currently active.
    pub fn is_low_priority_dag_cpu_throttle_enabled(&self) -> bool {
        self.low_priority_dag_active()
    }

    /// Get refill interval in milliseconds
    pub fn refill_interval_ms(&self) -> u64 {
        self.refill_interval_ms
    }

    /// Get the token bucket burst window in milliseconds.
    pub fn burst_window_ms(&self) -> u64 {
        self.burst_window_ms
    }

    /// Get keyspace bucket available tokens
    pub fn keyspace_available(&self, keyspace_id: u32) -> Option<u64> {
        self.keyspace_buckets
            .get(&keyspace_id)
            .map(|b| b.available())
    }

    fn keyspace_inflight_state(&self, keyspace_id: u32) -> Arc<KeyspaceInflightState> {
        if let Some(state) = self.keyspace_inflight_states.get(&keyspace_id) {
            return Arc::clone(state.value());
        }
        let state = self
            .keyspace_inflight_states
            .entry(keyspace_id)
            .or_insert_with(|| Arc::new(KeyspaceInflightState::new(keyspace_id)));
        Arc::clone(state.value())
    }

    /// Reads the inflight allocation of a keyspace without registering state or
    /// Prometheus label values for keyspaces that have never allocated tokens.
    fn keyspace_inflight_us(&self, keyspace_id: u32) -> u64 {
        self.keyspace_inflight_states
            .get(&keyspace_id)
            .map(|state| state.allocated_us.load(Ordering::Acquire))
            .unwrap_or(0)
    }

    fn update_keyspace_quota_inner(&self, keyspace_id: u32, quota_us: u64) {
        let refill_interval = Duration::from_millis(self.refill_interval_ms.max(1));
        let capacity_us =
            refill_rate_to_capacity(quota_us, self.refill_interval_ms, self.burst_window_ms);
        let inflight_state = self.keyspace_inflight_state(keyspace_id);
        let inflight_us = inflight_state.allocated_us.load(Ordering::Acquire);
        let bucket = if let Some(entry) = self.keyspace_buckets.get(&keyspace_id) {
            let bucket = entry.value().clone();
            bucket.update_quota_reserving_inflight(capacity_us, quota_us, inflight_us);
            bucket
        } else {
            let bucket = Arc::new(CpuTokenBucket::new(
                capacity_us,
                quota_us,
                refill_interval,
                0,
            ));
            self.keyspace_buckets.insert(keyspace_id, bucket.clone());
            bucket
        };
        if inflight_us > capacity_us {
            bucket.token_balance.store(0, Ordering::Release);
        }

        // Update quota map
        let millicores = quota_us / self.refill_interval_ms.max(1);
        self.keyspace_quotas.insert(keyspace_id, millicores);

        // Check for quota oversubscription after update
        let sum_quotas_millicores: u64 = self.keyspace_quotas.iter().map(|e| *e.value()).sum();
        let sum_quotas_us = sum_quotas_millicores * self.refill_interval_ms.max(1);
        let global_refill_rate_us = self.global_bucket.base_refill_rate();

        if sum_quotas_us > global_refill_rate_us {
            let oversubscription_ratio = sum_quotas_us as f64 / global_refill_rate_us as f64;
            if self.enforce_global_limit_enabled.load(Ordering::Acquire) {
                tikv_util::warn!(
                    "[CPU throttle] Quota update: sum exceeds global (global limit enforced)";
                    "keyspace_id" => keyspace_id,
                    "new_quota_us" => quota_us,
                    "sum_quotas_us" => sum_quotas_us,
                    "global_refill_rate_us" => global_refill_rate_us,
                    "oversubscription_ratio" => format!("{:.2}x", oversubscription_ratio),
                );
            } else {
                info!(
                    "[CPU throttle] Quota update: sum exceeds global (hard isolation mode)";
                    "keyspace_id" => keyspace_id,
                    "new_quota_us" => quota_us,
                    "sum_quotas_us" => sum_quotas_us,
                    "global_refill_rate_us" => global_refill_rate_us,
                    "oversubscription_ratio" => format!("{:.2}x", oversubscription_ratio),
                );
            }
        }

        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_label = keyspace_label.as_ref();
        CPU_THROTTLE_KEYSPACE_BUCKET_CAPACITY
            .with_label_values(&[keyspace_label])
            .set(clamp_gauge_value(capacity_us));
        inflight_state
            .bucket_available
            .set(clamp_gauge_value(bucket.available()));
        CPU_THROTTLE_KEYSPACE_REFILL_RATE
            .with_label_values(&[keyspace_label])
            .set(clamp_gauge_value(bucket.current_refill_rate()));
        inflight_state.update_inflight_metrics(inflight_us, capacity_us);
    }

    /// Add or update keyspace bucket.
    pub fn update_keyspace_quota(&self, keyspace_id: u32, quota_us: u64) {
        if self.tenant_cap_quota_us().is_some() {
            tikv_util::warn!(
                "[CPU throttle] Ignore non-tenant quota update while tenant cap is enabled";
                "keyspace_id" => keyspace_id,
                "requested_quota_us" => quota_us,
            );
            return;
        }
        self.update_keyspace_quota_inner(keyspace_id, quota_us);
    }

    /// Update config-managed keyspace quotas (hot update)
    pub fn update_config_keyspace_quotas(&self, new_quotas: std::collections::HashMap<u32, u64>) {
        let new_keyspaces: std::collections::HashSet<u32> = new_quotas.keys().copied().collect();
        if self.tenant_cap_quota_us().is_some() {
            tikv_util::warn!(
                "[CPU throttle] Ignore config keyspace quotas while tenant cap is enabled";
                "keyspace_count" => new_quotas.len(),
            );
            self.replace_config_keyspaces(&new_keyspaces);
            return;
        }

        let mut to_remove = Vec::new();

        for entry in self.config_keyspaces.iter() {
            let keyspace_id = *entry.key();
            if !new_keyspaces.contains(&keyspace_id) {
                to_remove.push(keyspace_id);
            }
        }

        for keyspace_id in to_remove {
            self.config_keyspaces.remove(&keyspace_id);
            self.remove_keyspace(keyspace_id);
        }

        for (keyspace_id, quota_us) in new_quotas {
            self.update_keyspace_quota(keyspace_id, quota_us);
            self.config_keyspaces.insert(keyspace_id);
        }
    }

    fn replace_config_keyspaces(&self, new_keyspaces: &std::collections::HashSet<u32>) {
        let mut to_remove = Vec::new();
        for entry in self.config_keyspaces.iter() {
            let keyspace_id = *entry.key();
            if !new_keyspaces.contains(&keyspace_id) {
                to_remove.push(keyspace_id);
            }
        }
        for keyspace_id in to_remove {
            self.config_keyspaces.remove(&keyspace_id);
        }
        for keyspace_id in new_keyspaces {
            self.config_keyspaces.insert(*keyspace_id);
        }
    }

    pub fn update_keyspace_quota_for_pd_service_limit(
        &self,
        keyspace_id: u32,
        quota_us: u64,
    ) -> bool {
        if !self.service_limit_can_mutate_keyspace() {
            tikv_util::warn!(
                "[CPU throttle] Ignore PD service_limit quota update while disabled or tenant cap is enabled";
                "keyspace_id" => keyspace_id,
                "requested_quota_us" => quota_us,
            );
            return false;
        }
        self.update_keyspace_quota_inner(keyspace_id, quota_us);
        true
    }

    pub fn remove_keyspace_for_pd_service_limit(&self, keyspace_id: u32) -> bool {
        if !self.service_limit_can_mutate_keyspace() {
            tikv_util::warn!(
                "[CPU throttle] Ignore PD service_limit keyspace removal while disabled or tenant cap is enabled";
                "keyspace_id" => keyspace_id,
            );
            return false;
        }
        self.remove_keyspace(keyspace_id);
        true
    }

    /// Remove keyspace bucket
    pub fn remove_keyspace(&self, keyspace_id: u32) {
        self.keyspace_buckets.remove(&keyspace_id);
        self.keyspace_quotas.remove(&keyspace_id);
        self.burst_usage_tracker.remove_keyspace(keyspace_id);
        self.keyspace_inflight_states.remove(&keyspace_id);

        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_label = keyspace_label.as_ref();
        let _ = CPU_THROTTLE_KEYSPACE_BUCKET_AVAILABLE.remove_label_values(&[keyspace_label]);
        let _ = CPU_THROTTLE_KEYSPACE_BUCKET_CAPACITY.remove_label_values(&[keyspace_label]);
        let _ = CPU_THROTTLE_KEYSPACE_REFILL_RATE.remove_label_values(&[keyspace_label]);
        let _ = CPU_THROTTLE_KEYSPACE_INFLIGHT_US.remove_label_values(&[keyspace_label]);
        let _ = CPU_THROTTLE_KEYSPACE_DEBT_US.remove_label_values(&[keyspace_label]);
        remove_keyspace_label(keyspace_id);
    }

    /// Get the capacity cap for estimated CPU
    /// When runtime token management is enabled, cap to capacity / 10
    /// to leave room for runtime allocation.
    /// When disabled, cap to full capacity to avoid starving requests.
    fn get_capacity_cap(&self, keyspace_id: u32, include_low_priority_dag: bool) -> u64 {
        let divisor = if self.config.enable_runtime_token_management {
            10
        } else {
            1
        };

        let tenant_cap = self.tenant_cap_quota_us().map(|quota_us| {
            refill_rate_to_capacity(quota_us, self.refill_interval_ms, self.burst_window_ms)
                / divisor
        });
        let keyspace_cap = tenant_cap
            .or_else(|| {
                self.keyspace_buckets
                    .get(&keyspace_id)
                    .map(|b| b.capacity() / divisor)
            })
            .unwrap_or(u64::MAX);

        // Also consider global capacity
        let global_cap = self.global_bucket.capacity() / divisor;
        let low_priority_cap = if include_low_priority_dag && self.low_priority_dag_active() {
            self.low_priority_dag_bucket.capacity() / divisor
        } else {
            u64::MAX
        };

        // Return min of both, ensure at least 1
        keyspace_cap.min(global_cap).min(low_priority_cap).max(1)
    }

    fn estimated_cpu_per_request_us(
        &self,
        keyspace_id: u32,
        include_low_priority_dag: bool,
    ) -> u64 {
        let estimated = if let Some(ref estimator) = self.adaptive_estimator {
            estimator.get(keyspace_id)
        } else {
            self.keyspace_estimated_cpu_per_request_us
                .get(&keyspace_id)
                .copied()
                .unwrap_or(self.config.estimated_cpu_per_request_us)
        };

        // Cap to prevent exceeding bucket capacity
        let cap = self.get_capacity_cap(keyspace_id, include_low_priority_dag);
        estimated.min(cap).max(1) // Ensure > 0
    }

    /// Get the estimated CPU time per request for a keyspace
    /// Applies cap to prevent estimates from exceeding bucket capacity
    pub fn get_estimated_cpu_per_request_us(&self, keyspace_id: u32) -> u64 {
        self.estimated_cpu_per_request_us(keyspace_id, false)
    }

    /// Get the estimated CPU time per low-priority DAG request for a keyspace.
    /// Applies the extra low-priority global bucket cap in addition to the
    /// regular keyspace/global capacity caps.
    pub fn get_estimated_cpu_per_low_priority_dag_request_us(&self, keyspace_id: u32) -> u64 {
        self.estimated_cpu_per_request_us(keyspace_id, true)
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::Barrier, thread, time::Duration};

    use tikv_util::sys::SysQuota;

    use super::*;

    const REFILL_INTERVAL_MATRIX_MS: [u64; 6] = [100, 50, 20, 10, 5, 1];

    /// Preserve the explicit 100ms interval used by the fixed microsecond
    /// quotas in legacy behavior tests. Production defaults and their capacity
    /// compatibility are covered separately below.
    fn test_cpu_throttle_config() -> CpuThrottleConfig {
        let mut config = CpuThrottleConfig::default();
        config.refill_interval_ms = 100;
        config
    }

    fn manager_for_refill_interval(refill_interval_ms: u64) -> CpuThrottleManager {
        let mut config = test_cpu_throttle_config();
        config.refill_interval_ms = refill_interval_ms;
        config.burst_window_ms = 100;
        config.enforce_global_limit = true;
        let refill_rate_us = refill_interval_ms * 1_000;
        CpuThrottleManager::new(config, refill_rate_us, [(1, refill_rate_us)].into())
    }

    #[test]
    fn test_refill_deadline_matrix_fairness_enabled() {
        for refill_interval_ms in REFILL_INTERVAL_MATRIX_MS {
            let manager = manager_for_refill_interval(refill_interval_ms);
            assert!(
                manager
                    .global_bucket
                    .try_allocate(manager.global_bucket.capacity())
            );
            let retry_at = manager
                .fair_global_cpu_dispatch_retry_at()
                .expect("empty fair global bucket must wait");
            assert_eq!(retry_at, manager.global_bucket.next_refill_deadline());
        }
    }

    #[test]
    fn test_refill_deadline_matrix_fairness_disabled() {
        for refill_interval_ms in REFILL_INTERVAL_MATRIX_MS {
            let manager = manager_for_refill_interval(refill_interval_ms);
            let keyspace_bucket = manager.keyspace_buckets.get(&1).unwrap();
            keyspace_bucket.refill();
            assert!(keyspace_bucket.try_allocate(keyspace_bucket.capacity()));
            let retry_at = manager.allocation_retry_at(1, &ThrottleError::KeyspaceCpuExhausted);
            assert_eq!(retry_at, keyspace_bucket.next_refill_deadline());
        }
    }

    #[test]
    fn test_allocation_retry_uses_earliest_viable_path() {
        for (burst_enabled, expect_global_deadline) in [(true, true), (false, false)] {
            let mut config = test_cpu_throttle_config();
            config.refill_interval_ms = 60_000;
            config.burst_window_ms = 60_000;
            config.enable_burst = burst_enabled;
            let manager = CpuThrottleManager::new(config, 1_000, [(1, 1_000)].into());
            let keyspace_bucket = manager.keyspace_buckets.get(&1).unwrap();

            // Put the keyspace bucket one interval behind the global bucket so
            // its next boundary is later. If burst is enabled, an empty global
            // bucket can recover first; otherwise only keyspace refill helps.
            keyspace_bucket
                .last_refill_elapsed_ms
                .store(60_000, Ordering::Release);
            assert!(
                manager
                    .global_bucket
                    .try_allocate(manager.global_bucket.capacity())
            );
            assert!(keyspace_bucket.try_allocate(keyspace_bucket.capacity()));
            let global_deadline = manager.global_bucket.next_refill_deadline();
            let keyspace_deadline = keyspace_bucket.next_refill_deadline();
            assert!(global_deadline < keyspace_deadline);

            let error = manager
                .try_allocate(1, 1, Instant::now() + Duration::from_secs(5))
                .unwrap_err();
            assert!(matches!(
                error,
                ThrottleError::KeyspaceCpuExhaustedNoBurst { .. }
            ));
            let expected = if expect_global_deadline {
                global_deadline
            } else {
                keyspace_deadline
            };
            assert_eq!(manager.allocation_retry_at(1, &error), expected);
        }
    }

    #[test]
    fn test_burst_allocation_failure_uses_earliest_viable_path() {
        let mut config = test_cpu_throttle_config();
        config.refill_interval_ms = 60_000;
        config.burst_window_ms = 60_000;
        config.enable_burst = true;
        let manager = CpuThrottleManager::new(config, 100, [(1, 100)].into());
        let keyspace_bucket = manager.keyspace_buckets.get(&1).unwrap();

        // Keep one global token so the burst precheck passes, but not enough
        // for this request. The actual global allocation failure must still
        // preserve the earlier keyspace-or-global recovery deadline.
        keyspace_bucket
            .last_refill_elapsed_ms
            .store(60_000, Ordering::Release);
        assert!(manager.global_bucket.try_allocate(99));
        assert!(keyspace_bucket.try_allocate(100));
        let global_deadline = manager.global_bucket.next_refill_deadline();
        let keyspace_deadline = keyspace_bucket.next_refill_deadline();
        assert!(global_deadline < keyspace_deadline);

        let error = manager
            .try_allocate(1, 2, Instant::now() + Duration::from_secs(5))
            .unwrap_err();
        assert!(matches!(
            error,
            ThrottleError::KeyspaceCpuExhaustedNoBurst { .. }
        ));
        assert_eq!(manager.allocation_retry_at(1, &error), global_deadline);
    }

    #[test]
    fn test_token_bucket_allocation() {
        let bucket = CpuTokenBucket::new(10000, 10000, Duration::from_millis(100), 0);

        // Test successful allocation
        assert!(bucket.try_allocate(5000));
        assert_eq!(bucket.available(), 5000);

        // Test allocation that would exceed capacity
        assert!(!bucket.try_allocate(6000));
        assert_eq!(bucket.available(), 5000);

        // Test successful second allocation
        assert!(bucket.try_allocate(3000));
        assert_eq!(bucket.available(), 2000);
    }

    #[test]
    fn test_token_bucket_debt_allocation_and_release() {
        let bucket = CpuTokenBucket::new(100, 100, Duration::from_millis(100), 0);

        assert!(!bucket.try_allocate(150));
        assert_eq!(bucket.available(), 100);
        assert_eq!(bucket.debt(), 0);

        bucket.allocate_with_debt(150);
        assert_eq!(bucket.available(), 0);
        assert_eq!(bucket.debt(), 50);
        assert!(!bucket.has_positive_available_with_debt());

        bucket.release(30);
        assert_eq!(bucket.available(), 0);
        assert_eq!(bucket.debt(), 20);
        assert!(!bucket.try_allocate(1));

        bucket.release(50);
        assert_eq!(bucket.available(), 30);
        assert_eq!(bucket.debt(), 0);
        assert!(bucket.has_positive_available_with_debt());
    }

    #[test]
    fn test_token_bucket_release() {
        let bucket = CpuTokenBucket::new(10000, 10000, Duration::from_millis(100), 0);

        assert!(bucket.try_allocate(8000));
        assert_eq!(bucket.available(), 2000);

        // Release tokens
        bucket.release(3000);
        assert_eq!(bucket.available(), 5000);

        // Release more than capacity should cap at capacity
        bucket.release(10000);
        assert_eq!(bucket.available(), 10000);
    }

    #[test]
    fn test_token_bucket_refill() {
        let bucket = CpuTokenBucket::new(10000, 10000, Duration::from_millis(100), 0);

        // Allocate some tokens
        assert!(bucket.try_allocate(8000));
        assert_eq!(bucket.available(), 2000);

        // Wait for more than the refill interval
        thread::sleep(Duration::from_millis(150));

        // Refill should restore to capacity
        bucket.refill();
        assert_eq!(bucket.available(), 10000);
    }

    #[test]
    fn test_token_bucket_refill_interval() {
        let interval = Duration::from_millis(200);
        let mut bucket = CpuTokenBucket::new(20000, 5000, interval, 0);

        assert!(bucket.try_allocate(20000));
        assert_eq!(bucket.available(), 0);

        bucket.start_instant = Instant::now() - Duration::from_millis(450);
        bucket.last_refill_elapsed_ms.store(0, Ordering::Release);

        bucket.refill();
        assert_eq!(bucket.available(), 10000);
        assert_eq!(
            bucket.last_refill_elapsed_ms.load(Ordering::Acquire),
            400,
            "the incomplete 50ms interval must be preserved"
        );
        let retry_in = bucket
            .next_refill_deadline()
            .saturating_duration_since(Instant::now());
        assert!(retry_in <= Duration::from_millis(155));
        assert!(retry_in >= Duration::from_millis(140));
    }

    #[test]
    fn test_refill_interval_does_not_change_burst_capacity() {
        const BURST_WINDOW_MS: u64 = 100;
        const CPU_CORES: u64 = 2;
        const EXPECTED_CAPACITY_US: u64 = CPU_CORES * BURST_WINDOW_MS * 1_000;

        for refill_interval_ms in [100, 50, 20, 10, 5, 1] {
            let refill_rate_us = CPU_CORES * refill_interval_ms * 1_000;
            assert_eq!(
                refill_rate_to_capacity(refill_rate_us, refill_interval_ms, BURST_WINDOW_MS,),
                EXPECTED_CAPACITY_US,
                "refill interval {refill_interval_ms}ms"
            );

            let mut bucket = CpuTokenBucket::new(
                EXPECTED_CAPACITY_US,
                refill_rate_us,
                Duration::from_millis(refill_interval_ms),
                0,
            );
            assert!(bucket.try_allocate(EXPECTED_CAPACITY_US));
            bucket.start_instant = Instant::now() - Duration::from_millis(refill_interval_ms + 1);
            assert!(bucket.refill());
            // A 1ms test interval can cross a second boundary while the test
            // thread is descheduled. Every credited interval must still add
            // exactly one refill quantum.
            assert!(bucket.available() >= refill_rate_us);
            assert_eq!(bucket.available() % refill_rate_us, 0);
            assert!(bucket.available() <= EXPECTED_CAPACITY_US);
        }

        assert_eq!(
            refill_rate_to_capacity(u64::MAX, 1, BURST_WINDOW_MS),
            u64::MAX
        );
        assert_eq!(
            refill_rate_to_capacity(u64::MAX - 1, 1, BURST_WINDOW_MS),
            u64::MAX
        );
    }

    #[test]
    fn test_default_burst_window_preserves_pre_decoupling_capacity() {
        let config = CpuThrottleConfig::default();
        assert_eq!(config.refill_interval_ms, 10);
        assert_eq!(config.burst_window_ms, 100);

        // One CPU core contributes 10ms of tokens per default refill. The
        // 100ms burst window retains 100ms of tokens, matching the bucket
        // capacity before refill cadence and capacity were decoupled.
        let refill_rate_us = config.refill_interval_ms * 1_000;
        let manager =
            CpuThrottleManager::new(config, refill_rate_us, std::collections::HashMap::new());
        assert_eq!(manager.global_bucket.current_refill_rate(), 10_000);
        assert_eq!(manager.global_capacity(), 100_000);
    }

    #[test]
    fn test_concurrent_refill_advances_one_boundary_once() {
        let mut bucket = CpuTokenBucket::new(10_000, 1_000, Duration::from_secs(1), 0);
        assert!(bucket.try_allocate(10_000));
        bucket.start_instant = Instant::now() - Duration::from_millis(2_500);
        let bucket = Arc::new(bucket);

        let advanced = (0..16)
            .map(|_| {
                let bucket = bucket.clone();
                thread::spawn(move || bucket.refill())
            })
            .map(|thread| thread.join().unwrap())
            .filter(|advanced| *advanced)
            .count();

        assert_eq!(advanced, 1);
        assert_eq!(bucket.available(), 2_000);
        assert_eq!(bucket.last_refill_elapsed_ms.load(Ordering::Acquire), 2_000);
    }

    #[test]
    fn test_token_bucket_unlimited_quota() {
        // Test that unlimited quota (u64::MAX) doesn't cause overflow in refill
        let mut bucket = CpuTokenBucket::new(u64::MAX, u64::MAX, Duration::from_millis(100), 0);

        // Allocate some tokens
        assert!(bucket.try_allocate(10000));
        let available_before = bucket.available();

        // Simulate elapsed time and refill
        bucket.start_instant = Instant::now() - Duration::from_millis(250);
        bucket.last_refill_elapsed_ms.store(0, Ordering::Release);

        // Refill should skip the refill logic for u64::MAX
        bucket.refill();

        // Available tokens should remain unchanged after refill
        // (no overflow, no erratic behavior)
        let available_after = bucket.available();
        assert_eq!(available_before, available_after);
    }

    #[test]
    fn test_token_handle_drop() {
        let global_bucket = Arc::new(CpuTokenBucket::new(
            10000,
            10000,
            Duration::from_millis(100),
            0,
        ));
        let keyspace_bucket = Arc::new(CpuTokenBucket::new(
            5000,
            5000,
            Duration::from_millis(100),
            0,
        ));

        // Allocate tokens
        assert!(global_bucket.try_allocate(3000));
        assert!(keyspace_bucket.try_allocate(3000));

        {
            let mut config = test_cpu_throttle_config();
            config.enforce_global_limit = true; // Use legacy mode for this test
            let _handle = CpuTokenHandle {
                keyspace_id: 1,
                allocated_us: AtomicU64::new(3000),
                keyspace_allocated_us: AtomicU64::new(3000),
                global_allocated_us: AtomicU64::new(3000),
                actual_used_us: AtomicU64::new(1000), // Only used 1000
                global_bucket: global_bucket.clone(),
                keyspace_bucket: Some(keyspace_bucket.clone()),
                low_priority_bucket: None,
                keyspace_inflight_state: None,
                low_priority_allocated_us: AtomicU64::new(0),
                low_priority_inflight_allocated_us: None,
                is_burst: AtomicBool::new(false),
                request_deadline: Instant::now() + std::time::Duration::from_secs(5),
                skip_throttle: false,
                keyspace_burst_allowed: true,
                cpu_usage_monitor: None,
                config: Some(Arc::new(config)),
                adaptive_estimator: None,
                burst_usage_token: None,
                burst_usage_tracker: None,
                last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
                last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
            };

            // Handle is in scope
            assert_eq!(global_bucket.available(), 7000);
            assert_eq!(keyspace_bucket.available(), 2000);
        } // Handle dropped here

        // After drop, unused tokens (2000) should be returned
        assert_eq!(global_bucket.available(), 9000);
        assert_eq!(keyspace_bucket.available(), 4000);
    }

    #[test]
    fn test_runtime_retry_waits_for_both_keyspace_and_global_refill() {
        let global_bucket = Arc::new(CpuTokenBucket::new(1, 1, Duration::from_millis(100), 0));
        let keyspace_bucket = Arc::new(CpuTokenBucket::new(1, 1, Duration::from_millis(10), 0));
        assert!(global_bucket.try_allocate(1));
        assert!(keyspace_bucket.try_allocate(1));

        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        let handle = CpuTokenHandle {
            keyspace_id: 1,
            allocated_us: AtomicU64::new(0),
            keyspace_allocated_us: AtomicU64::new(0),
            global_allocated_us: AtomicU64::new(0),
            actual_used_us: AtomicU64::new(0),
            global_bucket: global_bucket.clone(),
            keyspace_bucket: Some(keyspace_bucket.clone()),
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: AtomicBool::new(false),
            request_deadline: Instant::now() + Duration::from_secs(5),
            skip_throttle: false,
            keyspace_burst_allowed: true,
            cpu_usage_monitor: None,
            config: Some(Arc::new(config)),
            adaptive_estimator: None,
            burst_usage_token: None,
            burst_usage_tracker: None,
            last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
            last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
        };

        let global_deadline = global_bucket.next_refill_deadline();
        let keyspace_deadline = keyspace_bucket.next_refill_deadline();
        assert!(global_deadline > keyspace_deadline);
        assert_eq!(handle.runtime_allocation_retry_at(1), global_deadline);
    }

    #[test]
    fn test_runtime_retry_considers_burst_after_global_refill() {
        let global_bucket = Arc::new(CpuTokenBucket::new(1, 1, Duration::from_millis(10), 1));
        let keyspace_bucket = Arc::new(CpuTokenBucket::new(1, 1, Duration::from_millis(100), 0));
        assert!(global_bucket.try_allocate(1));
        assert!(keyspace_bucket.try_allocate(1));

        let mut config = test_cpu_throttle_config();
        config.enable_burst = true;
        config.burst_threshold = 0.9;
        let handle = CpuTokenHandle {
            keyspace_id: 1,
            allocated_us: AtomicU64::new(0),
            keyspace_allocated_us: AtomicU64::new(0),
            global_allocated_us: AtomicU64::new(0),
            actual_used_us: AtomicU64::new(0),
            global_bucket: global_bucket.clone(),
            keyspace_bucket: Some(keyspace_bucket.clone()),
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: AtomicBool::new(false),
            request_deadline: Instant::now() + Duration::from_secs(5),
            skip_throttle: false,
            keyspace_burst_allowed: true,
            cpu_usage_monitor: Some(Arc::new(RwLock::new(CpuUsageMonitor::default()))),
            config: Some(Arc::new(config)),
            adaptive_estimator: None,
            burst_usage_token: None,
            burst_usage_tracker: None,
            last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
            last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
        };

        let global_deadline = global_bucket.next_refill_deadline();
        let keyspace_deadline = keyspace_bucket.next_refill_deadline();
        assert!(global_deadline < keyspace_deadline);
        assert_eq!(handle.runtime_allocation_retry_at(1), global_deadline);
    }

    #[test]
    fn test_token_handle_drop_after_runtime_burst_refunds_keyspace_tokens() {
        let global_bucket = Arc::new(CpuTokenBucket::new(
            10000,
            10000,
            Duration::from_millis(100),
            0,
        ));
        let keyspace_bucket = Arc::new(CpuTokenBucket::new(
            3000,
            3000,
            Duration::from_millis(100),
            0,
        ));
        let cpu_usage_monitor = Arc::new(RwLock::new(CpuUsageMonitor {
            global_usage_ratio: 0.5,
            ..CpuUsageMonitor::default()
        }));

        // Simulate an initial normal allocation that later needs runtime burst.
        assert!(global_bucket.try_allocate(3000));
        assert!(keyspace_bucket.try_allocate(3000));

        {
            let mut config = test_cpu_throttle_config();
            config.enable_burst = true;
            config.burst_threshold = 0.9;

            let handle = CpuTokenHandle {
                keyspace_id: 1,
                allocated_us: AtomicU64::new(3000),
                keyspace_allocated_us: AtomicU64::new(3000),
                global_allocated_us: AtomicU64::new(3000),
                actual_used_us: AtomicU64::new(0),
                global_bucket: global_bucket.clone(),
                keyspace_bucket: Some(keyspace_bucket.clone()),
                low_priority_bucket: None,
                keyspace_inflight_state: None,
                low_priority_allocated_us: AtomicU64::new(0),
                low_priority_inflight_allocated_us: None,
                is_burst: AtomicBool::new(false),
                request_deadline: Instant::now() + std::time::Duration::from_secs(5),
                skip_throttle: false,
                keyspace_burst_allowed: true,
                cpu_usage_monitor: Some(cpu_usage_monitor),
                config: Some(Arc::new(config)),
                adaptive_estimator: None,
                burst_usage_token: None,
                burst_usage_tracker: None,
                last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
                last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
            };

            assert!(handle.allocate_more(1000));
            assert!(handle.is_burst.load(Ordering::Acquire));

            handle.record_actual_usage(2000);

            assert_eq!(global_bucket.available(), 6000);
            assert_eq!(keyspace_bucket.available(), 0);
        }

        // The request only used part of the pre-burst keyspace allocation.
        // Drop should refund both the tracked global tokens and the unused
        // keyspace tokens allocated before switching to burst mode.
        assert_eq!(global_bucket.available(), 8000);
        assert_eq!(keyspace_bucket.available(), 1000);
    }

    #[test]
    fn test_throttle_manager_basic() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true; // Use legacy mode for this test
        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 50000); // 50ms per interval
        keyspace_quotas.insert(2, 30000); // 30ms per interval

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // Test allocation
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle1 = manager.try_allocate(1, 10000, deadline).unwrap();
        let _handle2 = manager.try_allocate(2, 10000, deadline).unwrap();

        // Check available tokens
        assert_eq!(manager.global_available(), 80000);
        assert_eq!(manager.keyspace_available(1), Some(40000));
        assert_eq!(manager.keyspace_available(2), Some(20000));
    }

    #[test]
    fn test_global_only_allocation_ignores_keyspace_bucket() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 100);
        let manager = CpuThrottleManager::new(config, 1_000, keyspace_quotas);
        let deadline = Instant::now() + Duration::from_secs(5);

        let _keyspace_handle = manager.try_allocate(1, 100, deadline).unwrap();
        assert_eq!(manager.keyspace_available(1), Some(0));

        let (_global_handle, _) = manager.try_allocate_global_only(1, 500, deadline).unwrap();
        assert_eq!(manager.global_available(), 400);
        assert_eq!(manager.keyspace_available(1), Some(0));
    }

    #[test]
    fn test_global_only_allocation_refunds_unused_tokens_on_drop() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        {
            let (handle, _) = manager.try_allocate_global_only(42, 400, deadline).unwrap();
            assert_eq!(manager.global_available(), 600);
            handle.record_actual_usage(150);
        }

        assert_eq!(manager.global_available(), 850);
    }

    #[test]
    fn test_global_only_runtime_allocation_uses_global_bucket_only() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 100);
        let manager = CpuThrottleManager::new(config, 1_000, keyspace_quotas);
        let deadline = Instant::now() + Duration::from_secs(5);

        let (handle, _) = manager.try_allocate_global_only(1, 100, deadline).unwrap();
        assert!(handle.allocate_more(200));
        assert_eq!(manager.global_available(), 700);
        assert_eq!(manager.keyspace_available(1), Some(100));
    }

    #[test]
    fn test_fair_charge_scale_from_ratios() {
        // Override disabled or equal to base: no scaling.
        assert_eq!(
            fair_charge_scale_from_ratios(0.6, 0.6),
            FAIR_CHARGE_SCALE_ONE
        );
        assert_eq!(
            fair_charge_scale_from_ratios(0.6, 0.0),
            FAIR_CHARGE_SCALE_ONE
        );

        // Lower effective ratio costs proportionally more tokens.
        assert_eq!(
            fair_charge_scale_from_ratios(0.6, 0.3),
            2 * FAIR_CHARGE_SCALE_ONE
        );
        // Raising the ratio above the startup value discounts each charge.
        assert_eq!(
            fair_charge_scale_from_ratios(0.6, 0.9),
            (FAIR_CHARGE_SCALE_ONE as f64 * 2.0 / 3.0).round() as u64
        );

        // 1000us of CPU at ratio 0.6 -> 0.9 costs ~667 tokens.
        let state = fair_charge_state_for(0.6, 0.9);
        assert_eq!(accumulate_fair_charge(1_000, &state), 666);
        // And the inverse direction doubles the cost.
        let state = fair_charge_state_for(0.6, 0.3);
        assert_eq!(accumulate_fair_charge(1_000, &state), 2_000);
    }

    #[test]
    fn test_global_bucket_metric_values_follow_registered_fair_scale() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());

        // An override is inert without a fair gate, so token values still map
        // one-to-one to CPU microseconds.
        manager.update_fair_charge_scale(0.6, 0.3);
        assert_eq!(manager.global_bucket.metric_cpu_us(1_000), 1_000);

        // Registering the gate makes each CPU microsecond cost two tokens. The
        // existing `_us` metrics therefore convert 1,000 tokens back to 500us.
        manager.register_fair_global_cpu_gate();
        assert_eq!(manager.global_bucket.metric_cpu_us(1_000), 500);

        manager.update_fair_charge_scale(0.6, 0.6);
        assert_eq!(manager.global_bucket.metric_cpu_us(1_000), 1_000);
    }

    #[test]
    fn test_fair_max_read_cpu_ratio_override_uses_startup_anchor() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());
        manager.initialize_fair_max_read_cpu_ratio(0.6);
        manager.register_fair_global_cpu_gate();

        manager.update_fair_max_read_cpu_ratio(0.3);
        assert_eq!(manager.fair_charge_scale(), 2.0);
        assert!((manager.enforced_max_read_cpu_ratio() - 0.3).abs() < f64::EPSILON);

        // A second startup observer must not reset the live override.
        manager.initialize_fair_max_read_cpu_ratio(0.8);
        assert_eq!(manager.fair_charge_scale(), 2.0);
        assert!((manager.enforced_max_read_cpu_ratio() - 0.3).abs() < f64::EPSILON);

        manager.update_fair_max_read_cpu_ratio(0.0);
        assert_eq!(manager.fair_charge_scale(), 1.0);
        assert!((manager.enforced_max_read_cpu_ratio() - 0.6).abs() < f64::EPSILON);
    }

    /// A packed charge state at the scale implied by the given ratios, with no
    /// carried remainder.
    fn fair_charge_state_for(base_ratio: f64, effective_ratio: f64) -> AtomicU64 {
        let scale = fair_charge_scale_from_ratios(base_ratio, effective_ratio);
        AtomicU64::new(pack_fair_charge_state(scale, 0))
    }

    #[test]
    fn test_fair_charge_state_packing_roundtrip() {
        for scale in [1, FAIR_CHARGE_SCALE_ONE, 99 * FAIR_CHARGE_SCALE_ONE] {
            for residue in [0, 1, FAIR_CHARGE_SCALE_ONE - 1] {
                let (got_scale, got_residue) =
                    unpack_fair_charge_state(pack_fair_charge_state(scale, residue));
                assert_eq!((got_scale, got_residue), (scale, residue));
            }
        }
    }

    #[test]
    fn test_fair_charge_accumulates_fraction_across_small_deltas() {
        // Rounding each delta independently would charge 1 token per 1us delta
        // here, ignoring the override entirely. Carrying the remainder keeps the
        // total on the exact product.
        let state = fair_charge_state_for(0.6, 0.9);
        let charged: u64 = (0..1_000).map(|_| accumulate_fair_charge(1, &state)).sum();
        assert!(
            (666..=667).contains(&charged),
            "1000 x 1us at scale 2/3 should total ~667 tokens, got {charged}"
        );

        // A scale below 0.5 would round every 1us delta to zero, letting short
        // dispatches run free forever.
        let state = fair_charge_state_for(0.3, 0.9);
        let (scale, _) = unpack_fair_charge_state(state.load(Ordering::Acquire));
        assert!(scale < FAIR_CHARGE_SCALE_ONE / 2);
        let charged: u64 = (0..300).map(|_| accumulate_fair_charge(1, &state)).sum();
        assert!(
            (99..=101).contains(&charged),
            "300 x 1us at scale 1/3 should total ~100 tokens, got {charged}"
        );
    }

    #[test]
    fn test_fair_charge_split_matches_single_charge() {
        // Splitting the same CPU time across many calls must not change the
        // total charged, beyond a single token of rounding.
        let state = fair_charge_state_for(0.6, 0.7);
        let single = accumulate_fair_charge(10_000, &state);

        let state = fair_charge_state_for(0.6, 0.7);
        let split: u64 = (0..10_000).map(|_| accumulate_fair_charge(1, &state)).sum();

        assert!(
            single.abs_diff(split) <= 1,
            "single={single} split={split} must agree within one token"
        );
    }

    #[test]
    fn test_fair_charge_state_never_mixes_scale_and_residue() {
        // Hammer the charge path while flipping the scale underneath it. Because
        // scale and remainder share one word, every observed state must be a
        // scale that was actually published together with a remainder produced
        // under it -- never a remainder carried across a scale boundary.
        const WORKERS: usize = 4;
        const ITERATIONS: usize = 20_000;

        let scale_a = fair_charge_scale_from_ratios(0.6, 0.9);
        let scale_b = fair_charge_scale_from_ratios(0.6, 0.3);
        let state = Arc::new(AtomicU64::new(pack_fair_charge_state(scale_a, 0)));
        let start = Arc::new(Barrier::new(WORKERS + 1));

        let workers: Vec<_> = (0..WORKERS)
            .map(|_| {
                let state = state.clone();
                let start = start.clone();
                thread::spawn(move || {
                    start.wait();
                    for _ in 0..ITERATIONS {
                        accumulate_fair_charge(1, &state);
                        let (scale, residue) =
                            unpack_fair_charge_state(state.load(Ordering::Acquire));
                        assert!(
                            scale == scale_a || scale == scale_b,
                            "observed unpublished scale {scale}"
                        );
                        assert!(
                            residue < FAIR_CHARGE_SCALE_ONE,
                            "remainder {residue} escaped its field"
                        );
                    }
                })
            })
            .collect();

        start.wait();
        for i in 0..ITERATIONS {
            let scale = if i % 2 == 0 { scale_b } else { scale_a };
            // Mirrors update_fair_charge_scale: one store publishes the new rate
            // and drops the remainder that belonged to the old one.
            state.store(pack_fair_charge_state(scale, 0), Ordering::Release);
            thread::yield_now();
        }

        for worker in workers {
            worker.join().unwrap();
        }
    }

    #[test]
    fn test_fair_charge_scale_change_clears_stale_residue() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());
        manager.register_fair_global_cpu_gate();

        manager.update_fair_charge_scale(0.6, 0.9);
        // Leave a partial remainder behind.
        assert!(!manager.consume_fair_global_cpu(1));

        manager.update_fair_charge_scale(0.6, 0.3);
        assert_eq!(manager.fair_charge_residue_q(), 0);
    }

    #[test]
    fn test_fair_charge_scale_changes_effective_rate_and_burst_depth() {
        let config = test_cpu_throttle_config();
        // refill_interval_ms = 100, burst_window default gives a 10x deeper
        // bucket than one interval's refill rate.
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());
        manager.register_fair_global_cpu_gate();

        let capacity = manager.global_capacity();
        let refill_rate = manager.global_refill_rate();

        // Halving the effective ratio doubles the cost of each charge, so both
        // the burst depth and the sustained rate halve in real CPU terms.
        manager.update_fair_charge_scale(0.6, 0.3);
        assert_eq!(manager.fair_charge_scale(), 2.0);

        // Burst depth: draining the whole bucket now takes half the real CPU.
        assert!(!manager.consume_fair_global_cpu(capacity / 2));
        assert_eq!(manager.global_available(), 0);
        assert_eq!(manager.global_debt(), 0);

        // The bucket itself was never resized.
        assert_eq!(manager.global_capacity(), capacity);
        assert_eq!(manager.global_refill_rate(), refill_rate);

        // Sustained rate: one interval's refill covers half the real CPU it
        // used to.
        assert!(manager.consume_fair_global_cpu(1));
    }

    #[test]
    fn test_fair_charge_scale_disable_restores_unscaled_charges() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());
        manager.register_fair_global_cpu_gate();

        manager.update_fair_charge_scale(0.6, 0.3);
        assert_eq!(manager.fair_charge_scale(), 2.0);
        assert!(!manager.consume_fair_global_cpu(100));
        assert_eq!(manager.global_available(), 800);

        // Restoring the startup ratio disables the override exactly, with no
        // accumulated drift.
        manager.update_fair_charge_scale(0.6, 0.6);
        assert_eq!(manager.fair_charge_scale(), 1.0);
        assert!(!manager.consume_fair_global_cpu(100));
        assert_eq!(manager.global_available(), 700);
    }

    #[test]
    fn test_fair_charge_scale_does_not_affect_legacy_admission() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        let manager = CpuThrottleManager::new(config, 1_000, [(1, 1_000)].into());
        let deadline = Instant::now() + Duration::from_secs(5);

        manager.update_fair_charge_scale(0.6, 0.3);

        // Legacy admission still charges estimated CPU 1:1. Its unit mixing
        // between tokens and measured CPU is why the override is fair-gate
        // only.
        let handle = manager.try_allocate(1, 400, deadline).unwrap();
        assert_eq!(manager.global_available(), 600);
        assert_eq!(manager.keyspace_available(1), Some(600));
        drop(handle);
    }

    #[test]
    fn test_consume_fair_global_cpu_positive_balance_continues() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());

        assert!(!manager.consume_fair_global_cpu(400));
        assert_eq!(manager.global_available(), 600);
        assert_eq!(manager.global_debt(), 0);

        assert!(!manager.consume_fair_global_cpu(600));
        assert_eq!(manager.global_available(), 0);
        assert_eq!(manager.global_debt(), 0);
    }

    #[test]
    fn test_consume_fair_global_cpu_debt_yields() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 100, std::collections::HashMap::new());

        assert!(manager.consume_fair_global_cpu(150));
        assert_eq!(manager.global_available(), 0);
        assert_eq!(manager.global_debt(), 50);
    }

    #[test]
    fn test_consume_fair_global_cpu_refill_recovers() {
        let mut config = test_cpu_throttle_config();
        config.refill_interval_ms = 1;
        config.burst_window_ms = 1;
        let manager = CpuThrottleManager::new(config, 100, std::collections::HashMap::new());

        assert!(manager.consume_fair_global_cpu(150));
        thread::sleep(Duration::from_millis(2));

        assert!(!manager.consume_fair_global_cpu(0));
        assert!(manager.global_available() > 0);
        assert_eq!(manager.global_debt(), 0);
    }

    #[test]
    fn test_token_bucket_reset_to_capacity_clears_debt_and_interval_usage() {
        let bucket = CpuTokenBucket::new(100, 100, Duration::from_secs(60), 100);

        bucket.allocate_with_debt(150);
        assert_eq!(bucket.available(), 0);
        assert_eq!(bucket.debt(), 50);
        assert!(bucket.try_allocate_in_interval(100).is_some());
        assert!(bucket.try_allocate_in_interval(1).is_none());

        bucket.reset_to_capacity();

        assert_eq!(bucket.available(), 100);
        assert_eq!(bucket.debt(), 0);
        assert!(bucket.try_allocate_in_interval(100).is_some());
    }

    #[test]
    fn test_interval_allocation_stale_window_does_not_reset_newer_usage() {
        let bucket = CpuTokenBucket::new(100, 100, Duration::from_secs(60), 100);
        let old_window = bucket.try_allocate_in_interval(10).unwrap();
        {
            // Simulate another thread advancing the shared window after this
            // caller sampled the old ID but before it acquired the lock.
            let mut interval = bucket.interval_allocation.lock().unwrap();
            interval.window_id = old_window + 1;
            interval.allocated_us = 40;
        }

        assert_eq!(bucket.try_allocate_in_interval(10), Some(old_window + 1));
        assert_eq!(bucket.allocated_in_interval(), 50);
        bucket.release_in_interval(10, old_window);
        assert_eq!(bucket.allocated_in_interval(), 50);
    }

    #[test]
    fn test_dry_run_reset_requires_registered_fair_global_cpu_gate() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 100, std::collections::HashMap::new());

        assert!(manager.consume_fair_global_cpu(150));
        manager.reset_fair_global_cpu_after_dry_run();
        assert_eq!(manager.global_debt(), 50);

        manager.register_fair_global_cpu_gate();
        manager.reset_fair_global_cpu_after_dry_run();
        assert_eq!(manager.global_available(), 100);
        assert_eq!(manager.global_debt(), 0);
    }

    #[test]
    fn test_throttle_manager_exhaustion() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = false;
        config.enforce_global_limit = true; // Use legacy mode for this test
        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 10000); // 10ms per interval

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // Exhaust keyspace quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle = manager.try_allocate(1, 8000, deadline).unwrap();
        manager.try_allocate(1, 3000, deadline).unwrap_err(); // Should fail

        // Global still has capacity
        assert_eq!(manager.global_available(), 92000);
    }

    #[test]
    fn test_dynamic_refill_adjustment() {
        let mut config = test_cpu_throttle_config();
        config.enable_dynamic_adjustment = true;
        config.high_watermark = 0.8;
        config.low_watermark = 0.5;

        let manager = CpuThrottleManager::new(config, 100000, std::collections::HashMap::new());

        // Update with high CPU usage
        manager.update_usage(0.85, std::collections::HashMap::new());
        manager.adjust_refill_rates();

        // Refill rate should be reduced (can't directly verify internal state,
        // but this tests that the method runs without panic)

        // Update with low CPU usage
        manager.update_usage(0.4, std::collections::HashMap::new());
        manager.adjust_refill_rates();

        // Refill rate should be restored
    }

    #[test]
    fn test_low_priority_dag_bucket_initialization() {
        let mut config = test_cpu_throttle_config();
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.6;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());

        assert_eq!(manager.low_priority_dag_capacity(), 60_000);
        assert_eq!(manager.low_priority_dag_available(), 60_000);
        assert_eq!(manager.low_priority_dag_bucket.base_refill_rate(), 60_000);
    }

    #[test]
    fn test_adaptive_estimator_records_concurrent_observations() {
        const KEYSPACE_ID: u32 = 42;
        const THREADS: u64 = 8;
        const OBSERVATIONS_PER_THREAD: u64 = 10_000;
        const ACTUAL_CPU_US: u64 = 200;

        // Keep every observation in the same window so this test isolates the
        // lock-free aggregation path used by concurrent coprocessor requests.
        let estimator = Arc::new(AdaptiveEstimator::new(
            1_000,
            60_000,
            std::collections::HashMap::new(),
        ));
        let mut threads = Vec::new();
        for _ in 0..THREADS {
            let estimator = Arc::clone(&estimator);
            threads.push(thread::spawn(move || {
                for _ in 0..OBSERVATIONS_PER_THREAD {
                    estimator.observe(KEYSPACE_ID, ACTUAL_CPU_US);
                }
            }));
        }
        for thread in threads {
            thread.join().unwrap();
        }

        let entry = estimator.estimates.get(&KEYSPACE_ID).unwrap();
        assert_eq!(
            entry.count.load(Ordering::Relaxed),
            THREADS * OBSERVATIONS_PER_THREAD
        );
        assert_eq!(
            entry.sum_us.load(Ordering::Relaxed),
            THREADS * OBSERVATIONS_PER_THREAD * ACTUAL_CPU_US
        );
    }

    #[test]
    fn test_low_priority_dag_refill_rate_tracks_global_scale() {
        let mut config = test_cpu_throttle_config();
        config.enable_dynamic_adjustment = true;
        config.high_watermark = 0.8;
        config.low_watermark = 0.5;
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.6;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());

        manager.update_usage(0.85, std::collections::HashMap::new());
        manager.adjust_refill_rates();
        assert_eq!(manager.global_bucket.current_refill_rate(), 87_500);
        assert_eq!(
            manager.low_priority_dag_bucket.current_refill_rate(),
            52_500
        );

        manager.update_usage(0.4, std::collections::HashMap::new());
        manager.adjust_refill_rates();
        assert_eq!(manager.global_bucket.current_refill_rate(), 100_000);
        assert_eq!(
            manager.low_priority_dag_bucket.current_refill_rate(),
            60_000
        );
    }

    #[test]
    fn test_fair_global_cpu_gate_disables_dynamic_refill_adjustment() {
        let mut config = test_cpu_throttle_config();
        config.enable_dynamic_adjustment = true;
        config.high_watermark = 0.8;
        config.low_watermark = 0.5;
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.6;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());

        // Demonstrate the failure mode: a usage ratio evaluated against a
        // smaller denominator can drive the legacy feedback controller to a
        // zero refill rate.
        manager.update_usage(1.4, std::collections::HashMap::new());
        manager.adjust_refill_rates();
        assert_eq!(manager.global_bucket.current_refill_rate(), 0);
        assert_eq!(manager.low_priority_dag_bucket.current_refill_rate(), 0);

        // Registering the fair gate removes the stale reduction and disables
        // the legacy feedback controller.
        manager.register_fair_global_cpu_gate();
        assert!(!manager.is_dynamic_adjustment_enabled());
        assert_eq!(manager.global_bucket.current_refill_rate(), 100_000);
        assert_eq!(
            manager.low_priority_dag_bucket.current_refill_rate(),
            60_000
        );

        // An accidental online attempt to re-enable it remains inert.
        manager.update_dynamic_adjustment(true, 0.8, 0.5);
        assert!(!manager.is_dynamic_adjustment_enabled());
        manager.adjust_refill_rates();
        assert_eq!(manager.global_bucket.current_refill_rate(), 100_000);
        assert_eq!(
            manager.low_priority_dag_bucket.current_refill_rate(),
            60_000
        );
    }

    #[test]
    fn test_low_priority_dag_allocation_rolls_back_on_keyspace_failure() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = false;
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.6;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 1_000);

        let manager = CpuThrottleManager::new(config, 100_000, keyspace_quotas);
        let deadline = Instant::now() + Duration::from_secs(5);

        let _first = manager
            .try_allocate_low_priority_dag(1, 900, deadline)
            .unwrap();
        let low_priority_available = manager.low_priority_dag_available();
        assert_eq!(low_priority_available, 60_000 - 900);

        let result = manager.try_allocate_low_priority_dag(1, 200, deadline);
        assert!(matches!(
            result,
            Err(ThrottleError::KeyspaceCpuExhaustedNoBurst { .. })
        ));
        assert_eq!(manager.low_priority_dag_available(), low_priority_available);
    }

    #[test]
    fn test_low_priority_dag_without_keyspace_quota_charges_global_bucket() {
        let mut config = test_cpu_throttle_config();
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.5;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let handle = manager
            .try_allocate_low_priority_dag(42, 1_000, deadline)
            .unwrap();
        assert!(!handle.skip_throttle);
        assert_eq!(manager.global_available(), 99_000);
        assert_eq!(manager.low_priority_dag_available(), 49_000);

        drop(handle);
        assert_eq!(manager.global_available(), 100_000);
        assert_eq!(manager.low_priority_dag_available(), 50_000);
    }

    #[test]
    fn test_low_priority_dag_without_keyspace_quota_fails_when_global_budget_is_exhausted() {
        let mut config = test_cpu_throttle_config();
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 1.0;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 1_000);

        let manager = CpuThrottleManager::new(config, 1_000, keyspace_quotas);
        let deadline = Instant::now() + Duration::from_secs(5);

        let _normal_handle = manager.try_allocate(1, 1_000, deadline).unwrap();
        assert_eq!(manager.global_available(), 0);
        assert_eq!(manager.low_priority_dag_available(), 1_000);

        let result = manager.try_allocate_low_priority_dag(42, 500, deadline);
        assert!(matches!(result, Err(ThrottleError::GlobalCpuExhausted)));
        assert_eq!(manager.low_priority_dag_available(), 1_000);
    }

    #[test]
    fn test_low_priority_dag_can_burst_when_bucket_has_capacity() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = true;
        config.burst_threshold = 0.9;
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.5;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 1_000);

        let manager = CpuThrottleManager::new(config, 100_000, keyspace_quotas);
        manager.update_usage(0.5, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let first_handle = manager
            .try_allocate_low_priority_dag(1, 900, deadline)
            .unwrap();
        assert_eq!(manager.low_priority_dag_available(), 50_000 - 900);

        let (burst_handle, is_burst) = manager
            .try_allocate_low_priority_dag_with_burst(1, 200, deadline)
            .unwrap();
        assert!(is_burst);
        assert_eq!(manager.low_priority_dag_available(), 50_000 - 900 - 200);

        drop(burst_handle);
        assert_eq!(manager.low_priority_dag_available(), 50_000 - 900);

        drop(first_handle);
        assert_eq!(manager.low_priority_dag_available(), 50_000);
    }

    #[test]
    fn test_low_priority_dag_rejected_when_bucket_exhausted_even_if_burst_allowed() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = true;
        config.burst_threshold = 0.9;
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.01;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 1_000);

        let manager = CpuThrottleManager::new(config, 100_000, keyspace_quotas);
        manager.update_usage(0.5, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let _first_handle = manager
            .try_allocate_low_priority_dag(1, 1_000, deadline)
            .unwrap();
        assert_eq!(manager.low_priority_dag_available(), 0);

        let result = manager.try_allocate_low_priority_dag_with_burst(1, 100, deadline);
        assert!(matches!(
            result,
            Err(ThrottleError::LowPriorityDagGlobalCpuExhausted)
        ));
        assert_eq!(manager.low_priority_dag_available(), 0);
    }

    #[test]
    fn test_tenant_cap_auto_creates_bucket_and_can_be_disabled() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);
        let expected_quota = (SysQuota::cpu_cores_quota().max(1.0) * 0.4 * 100.0 * 1000.0) as u64;

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        assert_eq!(manager.keyspace_available(42), Some(expected_quota - 1_000));
        assert_eq!(
            manager.keyspace_quotas.get(&42).map(|quota| *quota.value()),
            Some(expected_quota / 100)
        );
        drop(handle);

        manager.update_tenant_max_cpu_ratio(0.0);
        assert_eq!(manager.keyspace_available(42), None);
        let handle = manager.try_allocate(43, 1_000, deadline).unwrap();
        assert!(handle.skip_throttle);
    }

    #[test]
    fn test_tenant_cap_recreates_missing_bucket_before_global_fallback() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);
        let expected_quota = (SysQuota::cpu_cores_quota().max(1.0) * 0.4 * 100.0 * 1000.0) as u64;

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);

        // Simulate a cleanup race after tenant cap is observed active: the
        // quota map still has the keyspace, but the bucket disappeared before
        // the request reaches the keyspace allocation path.
        manager.keyspace_buckets.remove(&42);
        assert!(manager.keyspace_quotas.contains_key(&42));
        assert_eq!(manager.keyspace_available(42), None);

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        assert!(handle.keyspace_bucket.is_some());
        assert_eq!(handle.keyspace_allocated_us.load(Ordering::Acquire), 1_000);
        assert_eq!(handle.global_allocated_us.load(Ordering::Acquire), 1_000);
        assert_eq!(manager.keyspace_available(42), Some(expected_quota - 1_000));
    }

    #[test]
    fn test_tenant_cap_disable_is_not_resurrected_by_stale_recreate() {
        // Regression for the hot-disable race: an in-flight request that
        // snapshotted a non-zero quota must not be able to resurrect a keyspace
        // bucket after a concurrent disable has cleared the cap. After disable,
        // `ensure_tenant_cap_bucket` re-loads quota == 0 under the lifecycle
        // lock and refuses, so no stale bucket/quota is left behind.
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        // Cap active: keyspace 42 gets an auto-created bucket.
        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);
        assert!(manager.keyspace_buckets.contains_key(&42));

        // Disable the cap (simulates the hot-update disable path winning).
        manager.update_tenant_max_cpu_ratio(0.0);
        assert!(!manager.tenant_cap_enabled());
        assert!(!manager.keyspace_buckets.contains_key(&42));
        assert!(!manager.tenant_cap_keyspaces.contains(&42));

        // An in-flight request that had observed the cap active now tries to
        // (re)create the bucket. With the cap disabled it must refuse rather
        // than reinsert a bucket on the stale quota.
        let recreated = manager.ensure_tenant_cap_bucket(42, Instant::now());
        assert!(!recreated, "disabled cap must not resurrect a bucket");
        assert!(!manager.keyspace_buckets.contains_key(&42));
        assert!(!manager.tenant_cap_keyspaces.contains(&42));
        assert_eq!(manager.keyspace_available(42), None);

        // A fresh request after disable is no longer throttled.
        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(handle.skip_throttle);
    }

    #[test]
    fn test_tenant_cap_concurrent_disable_and_alloc_leaves_no_stale_bucket() {
        // Stress the disable/alloc interleaving: repeatedly allocate (which may
        // auto-create buckets) on one thread while disabling on another. After
        // both settle with the cap disabled, no keyspace bucket/quota may
        // survive.
        use std::sync::Barrier;

        for _ in 0..50 {
            let mut config = test_cpu_throttle_config();
            config.enforce_global_limit = true;
            config.tenant_max_cpu_ratio = 0.4;
            config.enable_pd_service_limit = false;
            config.enable_burst = false;
            config.enable_fair_allocation = false;
            config.enable_dynamic_adjustment = false;
            config.low_priority_dag_global_enabled = false;

            let manager = Arc::new(CpuThrottleManager::new(
                config,
                100_000,
                std::collections::HashMap::new(),
            ));
            let deadline = Instant::now() + Duration::from_secs(5);
            let barrier = Arc::new(Barrier::new(2));

            let alloc_mgr = manager.clone();
            let alloc_barrier = barrier.clone();
            let alloc = std::thread::spawn(move || {
                alloc_barrier.wait();
                for ks in 0..16u32 {
                    let _ = alloc_mgr.try_allocate(ks, 1_000, deadline);
                }
            });

            let disable_mgr = manager.clone();
            let disable_barrier = barrier.clone();
            let disable = std::thread::spawn(move || {
                disable_barrier.wait();
                disable_mgr.update_tenant_max_cpu_ratio(0.0);
            });

            alloc.join().unwrap();
            disable.join().unwrap();

            // Disable was the final lifecycle transition and permanently set
            // quota = 0, so once both threads have joined no allocation can
            // have resurrected a bucket: any bucket created before disable's
            // critical section is in disable's snapshot and removed, and any
            // recreation after it re-loads quota == 0 and refuses. Assert the
            // clean steady state directly (no extra cleanup, which would mask a
            // resurrection bug).
            assert!(!manager.tenant_cap_enabled());
            assert_eq!(
                manager.tenant_cap_keyspaces.len(),
                0,
                "no tenant cap keyspace may survive disable"
            );
            for ks in 0..16u32 {
                assert_eq!(
                    manager.keyspace_available(ks),
                    None,
                    "no stale bucket may survive disable for keyspace {ks}"
                );
            }
        }
    }

    #[test]
    fn test_tenant_cap_hot_update_refreshes_cached_quota() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);
        let cpu_cores = SysQuota::cpu_cores_quota().max(1.0);
        let quota_40 = (cpu_cores * 0.4 * 100.0 * 1000.0) as u64;
        let quota_20 = (cpu_cores * 0.2 * 100.0 * 1000.0) as u64;

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        assert_eq!(manager.keyspace_available(42), Some(quota_40 - 1_000));
        drop(handle);

        manager.update_tenant_max_cpu_ratio(0.2);
        assert_eq!(manager.keyspace_available(42), Some(quota_20));

        let handle = manager.try_allocate(43, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        assert_eq!(manager.keyspace_available(43), Some(quota_20 - 1_000));
    }

    #[test]
    fn test_tenant_cap_rewrites_existing_non_uniform_quota_on_hot_update() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.0;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(42, 1_000);
        let manager = CpuThrottleManager::new(config, 100_000, keyspace_quotas);
        let deadline = Instant::now() + Duration::from_secs(5);
        let cpu_cores = SysQuota::cpu_cores_quota().max(1.0);
        let quota_40 = (cpu_cores * 0.4 * 100.0 * 1000.0) as u64;

        assert_eq!(manager.keyspace_buckets.get(&42).unwrap().capacity(), 1_000);
        manager.update_tenant_max_cpu_ratio(0.4);
        assert_eq!(
            manager.keyspace_buckets.get(&42).unwrap().capacity(),
            quota_40
        );
        assert!(manager.tenant_cap_keyspaces.contains(&42));

        manager.update_keyspace_quota(42, 2_000);
        assert_eq!(
            manager.keyspace_buckets.get(&42).unwrap().capacity(),
            quota_40
        );

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        assert_eq!(handle.keyspace_allocated_us.load(Ordering::Acquire), 1_000);
        assert_eq!(
            manager.keyspace_buckets.get(&42).unwrap().capacity(),
            quota_40
        );
    }

    #[test]
    fn test_config_quota_update_does_not_remove_bucket_in_tenant_cap_mode() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.0;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(42, 1_000);
        let manager = CpuThrottleManager::new(config, 100_000, keyspace_quotas);
        let cpu_cores = SysQuota::cpu_cores_quota().max(1.0);
        let quota_40 = (cpu_cores * 0.4 * 100.0 * 1000.0) as u64;

        manager.update_tenant_max_cpu_ratio(0.4);
        manager.update_config_keyspace_quotas(std::collections::HashMap::new());

        assert_eq!(manager.keyspace_available(42), Some(quota_40));
        assert!(manager.tenant_cap_keyspaces.contains(&42));
        assert!(!manager.config_keyspaces.contains(&42));
    }

    #[test]
    fn test_pd_service_limit_removal_ignored_in_tenant_cap_mode() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);
        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);

        assert!(!manager.remove_keyspace_for_pd_service_limit(42));
        assert!(manager.keyspace_available(42).is_some());
        assert!(manager.tenant_cap_keyspaces.contains(&42));
    }

    #[test]
    fn test_tenant_cap_estimate_is_capped_before_bucket_creation() {
        let mut config = test_cpu_throttle_config();
        config.refill_interval_ms = 10;
        config.burst_window_ms = 100;
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.01;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;
        config.estimated_cpu_per_request_us = u64::MAX;
        config.enable_runtime_token_management = false;

        let manager =
            CpuThrottleManager::new(config, 100_000_000, std::collections::HashMap::new());
        let refill_rate = (SysQuota::cpu_cores_quota().max(1.0) * 0.01 * 10.0 * 1000.0) as u64;
        let expected_capacity = refill_rate_to_capacity(refill_rate, 10, 100);

        assert_eq!(manager.keyspace_available(42), None);
        assert_eq!(
            manager.get_estimated_cpu_per_request_us(42),
            expected_capacity
        );

        let deadline = Instant::now() + Duration::from_secs(5);
        let handle = manager
            .try_allocate(42, manager.get_estimated_cpu_per_request_us(42), deadline)
            .unwrap();
        assert!(!handle.skip_throttle);
        assert_eq!(
            handle.keyspace_allocated_us.load(Ordering::Acquire),
            expected_capacity
        );
        assert_eq!(
            manager.keyspace_buckets.get(&42).unwrap().capacity(),
            expected_capacity
        );
    }

    #[test]
    fn test_tenant_cap_evicts_idle_auto_created_bucket() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;
        config.tenant_cap_idle_ttl_ms = 10 * 60 * 1000;
        let idle_ttl_ms = config.tenant_cap_idle_ttl_ms;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);
        assert!(manager.tenant_cap_keyspaces.contains(&42));
        assert!(manager.keyspace_available(42).is_some());

        let now = Instant::now();
        manager.tenant_cap_keyspace_last_access.insert(
            42,
            now.checked_sub(Duration::from_millis(idle_ttl_ms + 1))
                .unwrap(),
        );
        *manager.last_tenant_cap_cleanup.write().unwrap() = now
            .checked_sub(Duration::from_secs(
                TENANT_CAP_KEYSPACE_CLEANUP_INTERVAL_SECS + 1,
            ))
            .unwrap();

        let handle = manager.try_allocate(43, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        assert!(!manager.tenant_cap_keyspaces.contains(&42));
        assert_eq!(manager.keyspace_available(42), None);
        assert!(manager.tenant_cap_keyspaces.contains(&43));
        assert!(manager.keyspace_available(43).is_some());
    }

    #[test]
    fn test_tenant_cap_idle_ttl_hot_update_controls_eviction() {
        let mut config = test_cpu_throttle_config();
        config.enforce_global_limit = true;
        config.tenant_max_cpu_ratio = 0.4;
        config.enable_pd_service_limit = false;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;
        config.tenant_cap_idle_ttl_ms = 60_000;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let handle = manager.try_allocate(42, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);

        let now = Instant::now();
        manager
            .tenant_cap_keyspace_last_access
            .insert(42, now.checked_sub(Duration::from_millis(30_000)).unwrap());
        *manager.last_tenant_cap_cleanup.write().unwrap() = now
            .checked_sub(Duration::from_secs(
                TENANT_CAP_KEYSPACE_CLEANUP_INTERVAL_SECS + 1,
            ))
            .unwrap();

        manager.update_tenant_cap_idle_ttl(60_000);
        let handle = manager.try_allocate(43, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);
        assert!(manager.tenant_cap_keyspaces.contains(&42));
        assert!(manager.keyspace_available(42).is_some());

        manager.update_tenant_cap_idle_ttl(1_000);
        let handle = manager.try_allocate(44, 1_000, deadline).unwrap();
        assert!(!handle.skip_throttle);
        drop(handle);
        assert!(!manager.tenant_cap_keyspaces.contains(&42));
        assert_eq!(manager.keyspace_available(42), None);
    }

    #[test]
    fn test_low_priority_dag_estimate_is_capped_by_low_priority_bucket() {
        let mut config = test_cpu_throttle_config();
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 0.005;
        config.estimated_cpu_per_request_us = 1_000;
        config.enable_runtime_token_management = false;

        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());
        assert_eq!(manager.low_priority_dag_capacity(), 500);
        assert_eq!(manager.get_estimated_cpu_per_request_us(1), 1_000);
        assert_eq!(
            manager.get_estimated_cpu_per_low_priority_dag_request_us(1),
            500
        );

        let deadline = Instant::now() + Duration::from_secs(5);
        let handle = manager
            .try_allocate_low_priority_dag(
                1,
                manager.get_estimated_cpu_per_low_priority_dag_request_us(1),
                deadline,
            )
            .unwrap();
        assert_eq!(manager.low_priority_dag_available(), 0);

        drop(handle);
        assert_eq!(manager.low_priority_dag_available(), 500);
    }

    #[test]
    fn test_low_priority_dag_hot_update_preserves_inflight_tokens() {
        let mut config = test_cpu_throttle_config();
        config.low_priority_dag_global_enabled = true;
        config.low_priority_dag_global_ratio = 1.0;
        config.enable_runtime_token_management = false;

        let manager = CpuThrottleManager::new(config, 1_000, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let handle = manager
            .try_allocate_low_priority_dag(1, 600, deadline)
            .unwrap();
        assert_eq!(manager.low_priority_dag_available(), 400);
        assert_eq!(
            manager
                .low_priority_dag_inflight_allocated_us
                .load(Ordering::Acquire),
            600
        );

        manager.update_low_priority_dag_settings(false, 0.0);
        assert_eq!(manager.low_priority_dag_capacity(), 0);
        assert_eq!(manager.low_priority_dag_available(), 0);

        manager.update_low_priority_dag_settings(true, 1.0);
        assert_eq!(manager.low_priority_dag_capacity(), 1_000);
        assert_eq!(manager.low_priority_dag_available(), 400);
        assert!(matches!(
            manager.try_allocate_low_priority_dag(1, 500, deadline),
            Err(ThrottleError::LowPriorityDagGlobalCpuExhausted)
        ));

        drop(handle);
        assert_eq!(
            manager
                .low_priority_dag_inflight_allocated_us
                .load(Ordering::Acquire),
            0
        );
        assert_eq!(manager.low_priority_dag_available(), 1_000);
    }

    #[test]
    fn test_fair_allocation_below_threshold() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = true;
        config.fair_allocation_threshold = 0.7;
        config.enforce_global_limit = true; // Use legacy mode for this test

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 50000); // 50ms per interval = 500 millicores
        keyspace_quotas.insert(2, 30000); // 30ms per interval = 300 millicores

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // CPU usage below threshold, fair allocation not triggered
        manager.update_usage(0.6, std::collections::HashMap::new());

        // Both allocations should succeed with normal quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle1 = manager.try_allocate(1, 40000, deadline).unwrap();
        let _handle2 = manager.try_allocate(2, 25000, deadline).unwrap();

        // Check available tokens
        assert_eq!(manager.global_available(), 35000);
        assert_eq!(manager.keyspace_available(1), Some(10000));
        assert_eq!(manager.keyspace_available(2), Some(5000));
    }

    #[test]
    fn test_fair_allocation_above_threshold() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = true;
        config.fair_allocation_threshold = 0.7;
        config.enable_burst = true;
        config.enforce_global_limit = true; // Use legacy mode for this test

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 50000); // 50ms per interval = 500 millicores
        keyspace_quotas.insert(2, 30000); // 30ms per interval = 300 millicores

        let manager = CpuThrottleManager::new(config, 200000, keyspace_quotas);

        // CPU usage above threshold (75%)
        manager.update_usage(0.75, std::collections::HashMap::new());

        // Exhaust keyspace quota to force burst allocation
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle = manager.try_allocate(1, 50000, deadline).unwrap();

        // Per-interval burst budget: 180ms
        // Fair share for keyspace 1: 180ms * (500/(500+300)) = 112.5ms
        // Burst allocation exceeding fair share should fail
        assert!(matches!(
            manager.try_allocate_with_burst(1, 120000, deadline),
            Err(ThrottleError::KeyspaceFairShareExceeded { .. })
        ));
    }

    #[test]
    fn test_fair_allocation_uses_refill_interval_budget() {
        let mut config = test_cpu_throttle_config();
        config.refill_interval_ms = 10;
        config.burst_window_ms = 100;
        config.enable_fair_allocation = true;
        config.enable_burst = true;
        config.burst_threshold = 0.9;

        let manager = CpuThrottleManager::new(config, 10_000, [(1, 1_000), (2, 1_000)].into());
        let interval_budget_us = manager.global_bucket.interval_allocation_limit();

        assert_eq!(manager.global_capacity(), 100_000);
        assert_eq!(interval_budget_us, 9_000);
        assert_eq!(
            manager.calculate_fair_share(1, interval_budget_us),
            Some(4_500)
        );
    }

    #[test]
    fn test_unlimited_global_refill_disables_interval_fair_cap() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = true;
        config.enable_fair_allocation = true;
        let manager = CpuThrottleManager::new(config, u64::MAX, [(1, 100)].into());

        let interval_budget_us = manager.global_bucket.interval_allocation_limit();
        assert_eq!(manager.global_capacity(), u64::MAX);
        assert_eq!(interval_budget_us, u64::MAX);
        assert_eq!(manager.calculate_fair_share(1, interval_budget_us), None);
    }

    #[test]
    fn test_fair_share_rejection_rolls_back_interval_reservation() {
        let mut config = test_cpu_throttle_config();
        // Keep the accounting window stable even if the test is descheduled.
        // The 10x burst/refill ratio matches the production 10ms/100ms pair.
        config.refill_interval_ms = 60_000;
        config.burst_window_ms = 600_000;
        config.enable_fair_allocation = true;
        config.fair_allocation_threshold = 0.7;
        config.enable_burst = true;
        config.burst_threshold = 0.9;
        config.enforce_global_limit = true;

        let manager = CpuThrottleManager::new(config, 600_000, [(1, 60_000), (2, 60_000)].into());
        manager.update_usage(0.75, std::collections::HashMap::new());
        let deadline = Instant::now() + Duration::from_secs(5);

        let _initial1 = manager.try_allocate(1, 600_000, deadline).unwrap();
        let _initial2 = manager.try_allocate(2, 600_000, deadline).unwrap();
        let _burst1 = manager
            .try_allocate_with_burst(1, 270_000, deadline)
            .unwrap();
        assert_eq!(manager.global_bucket.allocated_in_interval(), 270_000);

        assert!(matches!(
            manager.try_allocate_with_burst(1, 1, deadline),
            Err(ThrottleError::KeyspaceFairShareExceeded { .. })
        ));
        assert_eq!(manager.global_bucket.allocated_in_interval(), 270_000);

        // A failed request from keyspace 1 must not consume keyspace 2's half
        // of the interval budget.
        let _burst2 = manager
            .try_allocate_with_burst(2, 270_000, deadline)
            .unwrap();
        assert_eq!(manager.global_bucket.allocated_in_interval(), 540_000);
    }

    #[test]
    fn test_burst_mode_below_threshold() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = false;
        config.enable_burst = true;
        config.burst_threshold = 0.9;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 10000); // 10ms per interval

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // CPU usage at 85% (below burst threshold)
        manager.update_usage(0.85, std::collections::HashMap::new());

        // Exhaust keyspace quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle = manager.try_allocate(1, 9000, deadline).unwrap();
        assert_eq!(manager.keyspace_available(1), Some(1000));

        // Try to allocate more - should burst since global CPU < 90%
        let result = manager.try_allocate_with_burst(1, 2000, deadline);
        assert!(result.is_ok());
        let (_handle, is_burst) = result.unwrap();
        assert!(is_burst); // Should be in burst mode
    }

    #[test]
    fn test_burst_mode_above_threshold() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = false;
        config.enable_burst = true;
        config.burst_threshold = 0.9;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 10000); // 10ms per interval

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // CPU usage at 92% (above burst threshold)
        manager.update_usage(0.92, std::collections::HashMap::new());

        // Exhaust keyspace quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle = manager.try_allocate(1, 9000, deadline).unwrap();

        // Try to allocate more - should fail since burst is not allowed
        let result = manager.try_allocate_with_burst(1, 2000, deadline);
        assert!(matches!(
            result,
            Err(ThrottleError::KeyspaceCpuExhaustedNoBurst { .. })
        ));
    }

    #[test]
    fn test_burst_global_budget_cap() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = false;
        config.enable_burst = true;
        config.burst_threshold = 0.5;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 10000); // 10ms per interval

        let manager = CpuThrottleManager::new(config, 1000, keyspace_quotas);

        // CPU usage below burst threshold
        manager.update_usage(0.1, std::collections::HashMap::new());

        // Exhaust keyspace quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle = manager.try_allocate(1, 10000, deadline).unwrap();

        // Burst up to the per-interval cap (500us)
        manager.try_allocate_with_burst(1, 400, deadline).unwrap();

        // Next burst exceeds the cap, should be rejected
        let result = manager.try_allocate_with_burst(1, 200, deadline);
        assert!(matches!(
            result,
            Err(ThrottleError::KeyspaceCpuExhaustedNoBurst { .. })
        ));
    }

    #[test]
    fn test_burst_allocate_more_respects_threshold() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = true;
        config.burst_threshold = 0.9;

        let global_bucket = Arc::new(CpuTokenBucket::new(
            10000,
            10000,
            Duration::from_millis(100),
            0,
        ));
        let cpu_usage_monitor = Arc::new(RwLock::new(CpuUsageMonitor {
            global_usage_ratio: 0.95,
            ..CpuUsageMonitor::default()
        }));

        let handle = CpuTokenHandle {
            keyspace_id: 1,
            allocated_us: AtomicU64::new(0),
            keyspace_allocated_us: AtomicU64::new(0),
            global_allocated_us: AtomicU64::new(0),
            actual_used_us: AtomicU64::new(0),
            global_bucket,
            keyspace_bucket: None,
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: AtomicBool::new(true),
            request_deadline: Instant::now() + Duration::from_secs(5),
            skip_throttle: false,
            keyspace_burst_allowed: true,
            cpu_usage_monitor: Some(cpu_usage_monitor),
            config: Some(Arc::new(config)),
            adaptive_estimator: None,
            burst_usage_token: None,
            burst_usage_tracker: None,
            last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
            last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
        };

        assert!(!handle.allocate_more(1000));
    }

    #[test]
    fn test_burst_runtime_allocation_is_capped_by_interval_budget() {
        let mut config = test_cpu_throttle_config();
        config.enable_burst = true;
        config.burst_threshold = 0.5;

        let global_bucket = Arc::new(CpuTokenBucket::new(
            100_000,
            10_000,
            Duration::from_millis(10),
            5_000,
        ));
        let handle = CpuTokenHandle {
            keyspace_id: 1,
            allocated_us: AtomicU64::new(0),
            keyspace_allocated_us: AtomicU64::new(0),
            global_allocated_us: AtomicU64::new(0),
            actual_used_us: AtomicU64::new(0),
            global_bucket,
            keyspace_bucket: None,
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: AtomicBool::new(true),
            request_deadline: Instant::now() + Duration::from_secs(5),
            skip_throttle: false,
            keyspace_burst_allowed: true,
            cpu_usage_monitor: None,
            config: Some(Arc::new(config)),
            adaptive_estimator: None,
            burst_usage_token: None,
            burst_usage_tracker: None,
            last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
            last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
        };

        assert_eq!(handle.max_single_runtime_allocation_us(), 5_000);
    }

    #[test]
    fn test_burst_allocate_more_concurrent_updates_are_atomic() {
        use std::sync::Barrier;

        let thread_count = 16;
        let allocations_per_thread = 1000;
        let expected_total = (thread_count * allocations_per_thread) as u64;

        for _ in 0..4 {
            let mut config = test_cpu_throttle_config();
            config.enable_burst = true;
            config.burst_threshold = 0.9;

            let capacity = expected_total + 10_000;
            let global_bucket = Arc::new(CpuTokenBucket::new(
                capacity,
                capacity,
                Duration::from_millis(100),
                0,
            ));
            let cpu_usage_monitor = Arc::new(RwLock::new(CpuUsageMonitor {
                global_usage_ratio: 0.1,
                ..CpuUsageMonitor::default()
            }));

            let handle = Arc::new(CpuTokenHandle {
                keyspace_id: 1,
                allocated_us: AtomicU64::new(0),
                keyspace_allocated_us: AtomicU64::new(0),
                global_allocated_us: AtomicU64::new(0),
                actual_used_us: AtomicU64::new(0),
                global_bucket: global_bucket.clone(),
                keyspace_bucket: None,
                low_priority_bucket: None,
                keyspace_inflight_state: None,
                low_priority_allocated_us: AtomicU64::new(0),
                low_priority_inflight_allocated_us: None,
                is_burst: AtomicBool::new(true),
                request_deadline: Instant::now() + Duration::from_secs(5),
                skip_throttle: false,
                keyspace_burst_allowed: true,
                cpu_usage_monitor: Some(cpu_usage_monitor),
                config: Some(Arc::new(config)),
                adaptive_estimator: None,
                burst_usage_token: None,
                burst_usage_tracker: None,
                last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
                last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
            });

            let barrier = Arc::new(Barrier::new(thread_count + 1));
            let workers: Vec<_> = (0..thread_count)
                .map(|_| {
                    let barrier = barrier.clone();
                    let handle = handle.clone();
                    thread::spawn(move || {
                        barrier.wait();
                        for _ in 0..allocations_per_thread {
                            assert!(handle.allocate_more(1));
                        }
                    })
                })
                .collect();

            barrier.wait();

            for worker in workers {
                worker.join().unwrap();
            }

            assert_eq!(handle.allocated(), expected_total);
            assert_eq!(
                handle.global_allocated_us.load(Ordering::Acquire),
                expected_total
            );
            assert_eq!(global_bucket.available(), capacity - expected_total);

            drop(handle);
            assert_eq!(global_bucket.available(), capacity);
        }
    }

    #[test]
    fn test_normal_allocate_more_concurrent_updates_are_atomic() {
        use std::sync::Barrier;

        let thread_count = 16;
        let allocations_per_thread = 1000;
        let expected_total = (thread_count * allocations_per_thread) as u64;

        for _ in 0..4 {
            let mut config = test_cpu_throttle_config();
            config.enable_burst = false;
            config.enforce_global_limit = true;

            let capacity = expected_total + 10_000;
            let global_bucket = Arc::new(CpuTokenBucket::new(
                capacity,
                capacity,
                Duration::from_millis(100),
                0,
            ));
            let keyspace_bucket = Arc::new(CpuTokenBucket::new(
                capacity,
                capacity,
                Duration::from_millis(100),
                0,
            ));
            let keyspace_inflight_state = Arc::new(KeyspaceInflightState::new(1));

            let handle = Arc::new(CpuTokenHandle {
                keyspace_id: 1,
                allocated_us: AtomicU64::new(0),
                keyspace_allocated_us: AtomicU64::new(0),
                global_allocated_us: AtomicU64::new(0),
                actual_used_us: AtomicU64::new(0),
                global_bucket: global_bucket.clone(),
                keyspace_bucket: Some(keyspace_bucket.clone()),
                low_priority_bucket: None,
                keyspace_inflight_state: Some(keyspace_inflight_state.clone()),
                low_priority_allocated_us: AtomicU64::new(0),
                low_priority_inflight_allocated_us: None,
                is_burst: AtomicBool::new(false),
                request_deadline: Instant::now() + Duration::from_secs(5),
                skip_throttle: false,
                keyspace_burst_allowed: false,
                cpu_usage_monitor: None,
                config: Some(Arc::new(config)),
                adaptive_estimator: None,
                burst_usage_token: None,
                burst_usage_tracker: None,
                last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
                last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
            });

            let barrier = Arc::new(Barrier::new(thread_count + 1));
            let workers: Vec<_> = (0..thread_count)
                .map(|_| {
                    let barrier = barrier.clone();
                    let handle = handle.clone();
                    thread::spawn(move || {
                        barrier.wait();
                        for _ in 0..allocations_per_thread {
                            assert!(handle.allocate_more(1));
                        }
                    })
                })
                .collect();

            barrier.wait();

            for worker in workers {
                worker.join().unwrap();
            }

            assert_eq!(handle.allocated(), expected_total);
            assert_eq!(
                handle.keyspace_allocated_us.load(Ordering::Acquire),
                expected_total
            );
            assert_eq!(
                handle.global_allocated_us.load(Ordering::Acquire),
                expected_total
            );
            assert_eq!(global_bucket.available(), capacity - expected_total);
            assert_eq!(keyspace_bucket.available(), capacity - expected_total);

            let inflight = keyspace_inflight_state.allocated_us.load(Ordering::Acquire);
            assert_eq!(inflight, expected_total);

            drop(handle);

            assert_eq!(global_bucket.available(), capacity);
            assert_eq!(keyspace_bucket.available(), capacity);

            let inflight_after_drop = keyspace_inflight_state.allocated_us.load(Ordering::Acquire);
            assert_eq!(inflight_after_drop, 0);
        }
    }

    #[test]
    fn test_burst_disabled_for_keyspace() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = false;
        config.enable_burst = true;
        config.burst_threshold = 0.9;
        config.keyspace_burst_enabled = "1:false".to_string();

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 10000); // 10ms per interval

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // CPU usage below burst threshold, but keyspace is explicitly disabled.
        manager.update_usage(0.85, std::collections::HashMap::new());

        // Exhaust keyspace quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle = manager.try_allocate(1, 9000, deadline).unwrap();

        // Try to allocate more - should fail since burst is disabled for keyspace
        let result = manager.try_allocate_with_burst(1, 2000, deadline);
        assert!(matches!(
            result,
            Err(ThrottleError::KeyspaceCpuExhaustedNoBurst { .. })
        ));
    }

    #[test]
    fn test_burst_mode_with_no_global_tokens() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = false;
        config.enable_burst = true;
        config.burst_threshold = 0.9;

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 10000); // 10ms per interval

        let manager = CpuThrottleManager::new(config, 100000, keyspace_quotas);

        // CPU usage at 85% (below burst threshold)
        manager.update_usage(0.85, std::collections::HashMap::new());

        // Exhaust global quota
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _global_handle = manager.try_allocate(1, 90000, deadline).unwrap();
        assert_eq!(manager.global_available(), 10000);

        // Exhaust keyspace quota
        let _keyspace_handle = manager.try_allocate(1, 9000, deadline).unwrap();

        // Try to allocate more - should fail because global bucket is nearly empty
        // (No tokens left for burst even though burst is allowed by watermark)
        let result = manager.try_allocate_with_burst(1, 5000, deadline);
        // This might fail due to global exhaustion
        let _ = result;
    }

    #[test]
    fn test_fair_allocation_with_burst() {
        let mut config = test_cpu_throttle_config();
        config.enable_fair_allocation = true;
        config.fair_allocation_threshold = 0.7;
        config.enable_burst = true;
        config.burst_threshold = 0.9;
        config.enforce_global_limit = true; // Use legacy mode for this test

        let mut keyspace_quotas = std::collections::HashMap::new();
        keyspace_quotas.insert(1, 50000); // 500 millicores
        keyspace_quotas.insert(2, 30000); // 300 millicores

        let manager = CpuThrottleManager::new(config, 200000, keyspace_quotas);

        // CPU usage at 75% (triggers fair allocation)
        manager.update_usage(0.75, std::collections::HashMap::new());

        // Exhaust keyspace quotas to force burst allocation
        let deadline = Instant::now() + std::time::Duration::from_secs(5);
        let _handle1 = manager.try_allocate(1, 50000, deadline).unwrap();
        let _handle2 = manager.try_allocate(2, 30000, deadline).unwrap();

        // Per-interval burst budget: 180ms
        // Fair shares are 112.5ms and 67.5ms respectively.
        assert!(matches!(
            manager.try_allocate_with_burst(1, 113000, deadline),
            Err(ThrottleError::KeyspaceFairShareExceeded { .. })
        ));
        assert!(matches!(
            manager.try_allocate_with_burst(2, 68000, deadline),
            Err(ThrottleError::KeyspaceFairShareExceeded { .. })
        ));
    }

    // -----------------------------------------------------------------------
    // Concurrent tests for BurstUsageTracker::global_window (Mutex-protected)
    // -----------------------------------------------------------------------

    /// Verify that concurrent `try_reserve_global` calls never allow the
    /// accumulated `used_us` to exceed `cap_us`, and that `used_us` exactly
    /// matches the number of successful reservations (no lost or double
    /// increments).
    ///
    /// There are no window transitions here — the focus is on the correctness
    /// of the in-window CAS logic under contention.
    #[test]
    fn test_burst_global_concurrent_cap_invariant() {
        use std::sync::atomic::AtomicU64;

        // Large enough window so it never advances during the test.
        let tracker = Arc::new(BurstUsageTracker::new(100_000));
        let cap_us: u64 = 1_000;
        let amount: u64 = 100;
        // 20 threads × 100 µs = 2 000 µs > cap, so roughly half will be
        // rejected.  The invariant is that successes × amount ≤ cap_us.
        let n_threads = 20;

        let success_count = Arc::new(AtomicU64::new(0));
        let window_id = tracker.current_window_id();

        let handles: Vec<_> = (0..n_threads)
            .map(|_| {
                let tracker = tracker.clone();
                let success_count = success_count.clone();
                thread::spawn(move || {
                    if tracker
                        .try_reserve_global(window_id, amount, cap_us)
                        .is_some()
                    {
                        success_count.fetch_add(1, Ordering::Relaxed);
                    }
                })
            })
            .collect();

        for h in handles {
            h.join().unwrap();
        }

        let successes = success_count.load(Ordering::Relaxed);
        // At most cap_us / amount threads can succeed.
        assert!(
            successes <= cap_us / amount,
            "over-reserved: {successes} threads succeeded, max allowed {}",
            cap_us / amount,
        );
        // used_us must equal exactly what was reserved — no lost updates.
        let used = tracker.global_window.lock().unwrap().used_us;
        assert_eq!(
            used,
            successes * amount,
            "used_us {used} does not match successful reservations {}",
            successes * amount,
        );
    }

    #[test]
    fn test_burst_usage_stale_window_does_not_reset_newer_usage() {
        let tracker = BurstUsageTracker::new(60_000);
        let old_window = tracker.current_window_id();
        let new_window = old_window + 1;

        assert_eq!(
            tracker.try_reserve_global(new_window, 100, 1_000),
            Some(new_window)
        );
        assert_eq!(
            tracker.try_reserve_global(old_window, 100, 1_000),
            Some(new_window)
        );
        assert_eq!(tracker.global_window.lock().unwrap().used_us, 200);
        tracker.release_global(old_window, 100);
        assert_eq!(tracker.global_window.lock().unwrap().used_us, 200);

        tracker.add_usage_in_window(1, new_window, 100);
        tracker.add_usage_in_window(1, old_window, 50);
        assert_eq!(tracker.usage_in_window(1, old_window), 0);
        assert_eq!(tracker.usage_in_window(1, new_window), 100);
    }

    /// Reproduce the original non-atomic window-transition race.
    ///
    /// Strategy:
    ///  1. Fill `used_us` close to `cap_us` in window N.
    ///  2. Sleep until window N+1 begins.
    ///  3. Release all threads simultaneously via a `Barrier`.
    ///  4. Each thread tries to reserve a small amount in the new window.
    ///
    /// With the old split-atomic code a thread could observe the new
    /// `window_id` (after the CAS) but the stale `used_us` from window N
    /// (before the `store(0)`), causing spurious rejections even though the
    /// new window is effectively empty.
    ///
    /// With the Mutex fix both fields are updated atomically, so every thread
    /// that lands in the new window must see `used_us = 0` and therefore
    /// succeed (total reserved 800 µs ≪ cap 9 999 µs).
    #[test]
    fn test_burst_global_concurrent_window_transition() {
        use std::sync::{Barrier, atomic::AtomicU64};

        let interval_ms: u64 = 10;
        let tracker = Arc::new(BurstUsageTracker::new(interval_ms));
        // cap is tight: old used_us (9 900) + one reservation (100) = 10 000
        // which exceeds cap (9 999), so stale reads produce false rejections.
        let cap_us: u64 = 9_999;
        let amount: u64 = 100;
        let n_threads: usize = 8;

        // Step 1: fill window N up to 9 900 µs (99 × 100).
        let old_window = tracker.current_window_id();
        for _ in 0..99 {
            assert!(
                tracker
                    .try_reserve_global(old_window, amount, cap_us)
                    .is_some(),
                "pre-fill should always succeed"
            );
        }
        assert_eq!(
            tracker.global_window.lock().unwrap().used_us,
            9_900,
            "pre-condition: old window used_us must be 9 900"
        );

        // Step 2: wait for window N+1 and capture it once.
        let new_window = {
            let deadline = Instant::now() + Duration::from_millis(interval_ms * 5);
            loop {
                let window = tracker.current_window_id();
                if window > old_window {
                    break window;
                }
                if Instant::now() >= deadline {
                    panic!("window did not advance within expected time");
                }
                thread::sleep(Duration::from_millis(1));
            }
        };

        // Step 3: all threads start simultaneously.
        let barrier = Arc::new(Barrier::new(n_threads + 1));
        let success_count = Arc::new(AtomicU64::new(0));

        let handles: Vec<_> = (0..n_threads)
            .map(|_| {
                let tracker = tracker.clone();
                let barrier = barrier.clone();
                let success_count = success_count.clone();
                thread::spawn(move || {
                    barrier.wait();
                    if tracker
                        .try_reserve_global(new_window, amount, cap_us)
                        .is_some()
                    {
                        success_count.fetch_add(1, Ordering::Relaxed);
                    }
                })
            })
            .collect();

        barrier.wait();

        for h in handles {
            h.join().unwrap();
        }

        // Step 4: assert.
        // Total attempted: 8 × 100 = 800 µs, cap = 9 999 µs → all must succeed.
        // With the stale-read bug, threads seeing used_us = 9 900 would compute
        // 9 900 + 100 = 10 000 > 9 999 and return false — a spurious rejection.
        let successes = success_count.load(Ordering::Relaxed);
        assert_eq!(
            successes, n_threads as u64,
            "all {n_threads} threads should succeed in a fresh window, \
             but only {successes} did (likely stale used_us from old window)"
        );
    }

    #[test]
    fn test_update_dynamic_adjustment() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 100_000, std::collections::HashMap::new());

        manager.update_dynamic_adjustment(true, 0.77, 0.33);

        assert!(manager.dynamic_adjustment_enabled.load(Ordering::Acquire));
        assert_eq!(
            f64::from_bits(manager.high_watermark_bits.load(Ordering::Acquire)),
            0.77
        );
        assert_eq!(
            f64::from_bits(manager.low_watermark_bits.load(Ordering::Acquire)),
            0.33
        );
    }

    #[test]
    fn test_update_burst_settings() {
        let config = test_cpu_throttle_config();
        let manager = CpuThrottleManager::new(config, 1000, std::collections::HashMap::new());

        let mut keyspace_flags = std::collections::HashMap::new();
        keyspace_flags.insert(1, true);
        keyspace_flags.insert(2, false);

        manager.update_burst_settings(true, 0.5, keyspace_flags);

        assert!(manager.burst_enabled.load(Ordering::Acquire));
        assert_eq!(
            f64::from_bits(manager.burst_threshold_bits.load(Ordering::Acquire)),
            0.5
        );
        assert_eq!(
            manager
                .global_bucket
                .interval_allocation_limit
                .load(Ordering::Acquire),
            burst_cap_us(manager.global_bucket.base_refill_rate(), 0.5)
        );
        assert_eq!(
            manager.keyspace_burst_enabled.get(&1).map(|v| *v.value()),
            Some(true)
        );
        assert_eq!(
            manager.keyspace_burst_enabled.get(&2).map(|v| *v.value()),
            Some(false)
        );

        let mut updated_flags = std::collections::HashMap::new();
        updated_flags.insert(2, true);
        manager.update_burst_settings(false, 0.9, updated_flags);

        assert!(!manager.burst_enabled.load(Ordering::Acquire));
        assert_eq!(
            f64::from_bits(manager.burst_threshold_bits.load(Ordering::Acquire)),
            0.9
        );
        assert_eq!(
            manager
                .global_bucket
                .interval_allocation_limit
                .load(Ordering::Acquire),
            burst_cap_us(manager.global_bucket.base_refill_rate(), 0.9)
        );
        assert!(manager.keyspace_burst_enabled.get(&1).is_none());
        assert_eq!(
            manager.keyspace_burst_enabled.get(&2).map(|v| *v.value()),
            Some(true)
        );
    }

    #[test]
    fn test_update_config_keyspace_quotas() {
        let config = test_cpu_throttle_config();
        let mut initial_quotas = std::collections::HashMap::new();
        initial_quotas.insert(1, 10_000);
        initial_quotas.insert(2, 20_000);
        let manager = CpuThrottleManager::new(config, 100_000, initial_quotas);

        let mut new_quotas = std::collections::HashMap::new();
        new_quotas.insert(2, 30_000);
        new_quotas.insert(3, 15_000);

        manager.update_config_keyspace_quotas(new_quotas);

        assert!(manager.keyspace_buckets.get(&1).is_none());
        assert!(!manager.config_keyspaces.contains(&1));

        assert_eq!(manager.keyspace_buckets.get(&2).unwrap().capacity(), 30_000);
        assert!(manager.config_keyspaces.contains(&2));

        assert_eq!(manager.keyspace_buckets.get(&3).unwrap().capacity(), 15_000);
        assert!(manager.config_keyspaces.contains(&3));

        let refill_interval_ms = manager.refill_interval_ms();
        assert_eq!(
            manager.keyspace_quotas.get(&2).map(|v| *v.value()),
            Some(30_000 / refill_interval_ms.max(1))
        );
        assert_eq!(
            manager.keyspace_quotas.get(&3).map(|v| *v.value()),
            Some(15_000 / refill_interval_ms.max(1))
        );
    }
}