// Copyright 2024 PingCAP, Inc.

use std::{
    marker::PhantomData,
    sync::{
        Arc, Mutex,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use prometheus::{IntCounter, IntGauge};
use tikv_util::{
    memory::{MemoryLimiter, MemoryLimiterGuard},
    sys::memory_usage_would_reach_throttling_level,
    time::{Instant, Limiter},
};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};

use crate::{
    KvEngineConfig, Options, ShardTag,
    metrics::{
        ENGINE_DFS_LOAD_MEMORY_WAIT_DURATION, ENGINE_DFS_LOAD_MEMORY_WAITING_REQUIRED_LOADS,
        ENGINE_LIMITER_THROTTLE_ACTION_COUNTER, ENGINE_LIMITER_THROTTLE_COUNTER,
    },
};

/// Default ratio applied to `hard_limit` to define the bursting threshold
/// (i.e. `hard_limit * (1.0 + DEFAULT_BURST_RATIO)`).
const DEFAULT_BURST_RATIO: f64 = 0.1;

// Cache reservations up to T/16 use current required reservations in the
// process-memory projection instead of reserving the complete future required
// quota T. This bounded unreserved allowance is a safety tradeoff, not a
// derived optimal value.
const DFS_CACHE_MEMORY_UNRESERVED_ALLOWANCE_DIVISOR: u64 = 16;

/// All members are in bytes.
#[derive(Clone, Default, Debug)]
pub struct LimiterOptions {
    pub enable: bool,
    pub soft_limit: u64,
    pub hard_limit: u64,
    pub max_speed_limit: u64,
    pub min_speed_limit: u64,
}

pub trait LimiterTypeTrait {
    const TAG: &'static str;
}

#[derive(Clone)]
pub struct StoreMemTable {}

impl LimiterTypeTrait for StoreMemTable {
    const TAG: &'static str = "store-memtable";
}

pub type StoreLimiter = WriteRateLimiter<StoreMemTable>;

#[derive(Clone)]
pub struct RegionMemTable {}

impl LimiterTypeTrait for RegionMemTable {
    const TAG: &'static str = "region-memtable";
}

#[derive(Clone)]
pub struct RegionL0Table {}

impl LimiterTypeTrait for RegionL0Table {
    const TAG: &'static str = "region-l0table";
}

struct StoreMetrics {
    speed_metric: IntGauge,
    last_record_time: Mutex<Instant>,
}

// `Ty` makes limiters be different types and avoids misuse.
#[derive(Clone)]
pub struct WriteRateLimiter<Ty: LimiterTypeTrait> {
    options: LimiterOptions,
    limiter: Arc<Limiter>,
    store_metrics: Option<Arc<StoreMetrics>>,
    throttle_metric: IntCounter,
    _phantom: PhantomData<Ty>,
}

impl WriteRateLimiter<StoreMemTable> {
    pub fn new_with_metric(options: LimiterOptions, speed_metric: IntGauge) -> Self {
        Self::new_impl(options, Some(speed_metric))
    }
}

/// Implement interfaces for `storage::txn::flow_controller::FlowController`.
impl<Ty: LimiterTypeTrait> WriteRateLimiter<Ty> {
    pub fn should_drop(&self, _region_id: u64) -> bool {
        // TODO: early drop ?
        false
    }

    pub fn discard_ratio(&self, _region_id: u64) -> f64 {
        0.0
    }

    pub fn enable(&self, _enable: bool) {
        // TODO: support online enable/disable
    }

    pub fn enabled(&self) -> bool {
        self.options.enable
    }

    pub fn consume(&self, _region_id: u64, bytes: usize) -> Duration {
        let dur = self.limiter.consume_duration(bytes);
        if !dur.is_zero() {
            self.throttle_metric.inc();
        }
        dur
    }

    pub fn unconsume(&self, _region_id: u64, bytes: usize) {
        self.limiter.unconsume(bytes);
    }

    pub fn is_unlimited(&self, _region_id: u64) -> bool {
        self.limiter.speed_limit() == f64::INFINITY
    }

    pub fn total_bytes_consumed(&self, _region_id: u64) -> usize {
        self.limiter.total_bytes_consumed()
    }

    /// Returns the hard resource limit in bytes.
    ///
    /// Returns the configured hard limit when enabled, or `0` when disabled.
    pub fn resource_max_limit(&self) -> u64 {
        if self.enabled() {
            self.options.hard_limit
        } else {
            0
        }
    }

    pub fn speed_limit(&self) -> f64 {
        self.limiter.speed_limit()
    }

    pub fn set_speed_limit(&self, _region_id: u64, speed_limit: f64) {
        self.limiter.set_speed_limit(speed_limit);
    }
}

/// Calculate the speed limit for the bursting mechanism.
fn calculate_bursting_speed_limit(
    min_speed_limit: u64,
    max_speed_limit: u64,
    excess: u64,         // excess beyond hard_limit
    bursting_range: u64, // range beyond hard_limit
) -> f64 {
    // Calculate the throttle value at hard_limit (from the exponential curve)
    let throttle_at_hard_limit = min_speed_limit as f64
        + (max_speed_limit as f64 - min_speed_limit as f64) * (-2.0f64).exp(); // exp(-2.0) when normalized = 1.0

    // Calculate the excess beyond hard_limit
    let bursting_ratio = excess as f64 / (bursting_range as f64).max(1.0);

    // Use a smooth sigmoid-like curve for the bursting mechanism
    // This provides a gradual deceleration from throttle_at_hard_limit to
    // min_speed_limit Using a cubic easing function for smoothness:
    // 3t^2 - 2t^3
    let t = bursting_ratio.min(1.0);
    let smooth_factor = 3.0 * t * t - 2.0 * t * t * t;

    // Interpolate smoothly from throttle_at_hard_limit to min_speed_limit
    let throttle =
        throttle_at_hard_limit - (throttle_at_hard_limit - min_speed_limit as f64) * smooth_factor;
    throttle.max(min_speed_limit as f64)
}

impl<Ty: LimiterTypeTrait> WriteRateLimiter<Ty> {
    pub fn new(options: LimiterOptions) -> Self {
        Self::new_impl(options, None)
    }

    /// New from another one.
    pub fn new_from(another: &WriteRateLimiter<Ty>) -> Self {
        let limiter = Self::new_impl(another.options.clone(), None);
        limiter.set_speed_limit(0, another.speed_limit());
        limiter
    }

    fn new_impl(options: LimiterOptions, speed_metric: Option<IntGauge>) -> Self {
        let limiter = Arc::new(
            <Limiter>::builder(f64::INFINITY)
                .refill(Duration::from_millis(1))
                .build(),
        );
        let store_metrics = speed_metric.map(|speed_metric| {
            Arc::new(StoreMetrics {
                speed_metric,
                last_record_time: Mutex::new(Instant::now_coarse()),
            })
        });
        let throttle_metric = ENGINE_LIMITER_THROTTLE_COUNTER.with_label_values(&[Ty::TAG]);
        Self {
            options,
            limiter,
            store_metrics,
            throttle_metric,
            _phantom: PhantomData,
        }
    }

    pub fn dummy() -> Self {
        let options = LimiterOptions {
            enable: false,
            ..Default::default()
        };
        Self::new_impl(options, None)
    }

    fn update_speed_limit(&self, throttle: f64) {
        self.limiter.set_speed_limit(throttle);
        self.update_statistics();
    }

    fn update_statistics(&self) {
        if let Some(metrics) = self.store_metrics.as_ref() {
            let mut last_record_time = metrics.last_record_time.lock().unwrap();
            let dur = last_record_time.saturating_elapsed_secs();
            if dur < f64::EPSILON {
                return;
            }

            let total = self.limiter.total_bytes_consumed();
            self.limiter.reset_statistics();
            *last_record_time = Instant::now_coarse();
            drop(last_record_time);

            let rate = total as f64 / dur;
            debug!("{}: update_statistics", Ty::TAG;
                "rate" => rate,
                "total" => total,
                "dur" => ?dur,
            );
            metrics.speed_metric.set(rate as i64);
        }
    }

    pub fn update_usage(
        &self,
        tag: &ShardTag,
        usage: u64,
        over_throttling_level: impl Fn() -> Option<(u64, u64)>,
    ) {
        if !self.options.enable {
            return;
        }

        let throttle = Self::calculate_throttle(
            usage,
            self.options.soft_limit,
            self.options.hard_limit,
            self.options.max_speed_limit,
            self.options.min_speed_limit,
            over_throttling_level,
        );
        if !throttle.is_infinite() {
            debug!("{} {}: update_usage", tag, Ty::TAG;
                "usage" => usage,
                "throttle" => throttle,
            );
        }
        self.update_speed_limit(throttle);
    }

    /// Calculate throttle according to current usage using a smooth exponential
    /// curve with bursting mechanism.
    ///
    /// All the parameters are in bytes or bytes/s.
    ///
    /// `usage` is the current usage of component (e.g. memtables) to be
    /// throttled.
    ///
    /// If `soft_limit` is `0`, throttling is disabled and the limiter remains
    /// unlimited.
    ///
    /// Otherwise the throttling algorithm has six zones:
    /// 1. When global memory is slightly above the throttling level and local
    ///    usage is still below `hard_limit`: use a short bursting curve
    /// 2. When global memory is well above the throttling level, or local usage
    ///    already reaches `hard_limit`: use `min_speed_limit`
    /// 3. Below soft_limit: no throttling (INFINITY)
    /// 4. Between soft_limit and hard_limit: exponential throttling curve
    /// 5. Between hard_limit and hard_limit * 1.1: bursting mechanism for
    ///    smoother throttling
    /// 6. Above hard_limit * 1.1: minimum speed limit (hard cutoff)
    pub(crate) fn calculate_throttle(
        usage: u64,
        soft_limit: u64,
        hard_limit: u64,
        max_speed_limit: u64,
        min_speed_limit: u64,
        over_throttling_level: impl Fn() -> Option<(u64, u64)>,
    ) -> f64 {
        debug_assert!(hard_limit >= soft_limit);
        debug_assert!(max_speed_limit >= min_speed_limit);

        // Early exit: if soft_limit is 0, no throttling
        if soft_limit == 0 {
            return f64::INFINITY;
        }

        // If there is global memory pressure, calculate the overshoot speed limit.
        // Otherwise, calculate the normal speed limit.
        let overshoot_speed_limit = Self::calculate_overshoot_speed_limit(
            usage,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            over_throttling_level,
        );
        if overshoot_speed_limit.is_finite() {
            return overshoot_speed_limit;
        }
        Self::calculate_normal_speed_limit(
            usage,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
        )
    }

    /// Calculate the speed limit for the overshoot mechanism.
    fn calculate_overshoot_speed_limit(
        usage: u64,
        hard_limit: u64,
        max_speed_limit: u64,
        min_speed_limit: u64,
        over_throttling_level: impl Fn() -> Option<(u64, u64)>,
    ) -> f64 {
        if let Some((overshoot, throttling_level)) = over_throttling_level() {
            let overshoot_bursting_range =
                (throttling_level as f64 * (DEFAULT_BURST_RATIO / 2.0_f64)) as u64;
            if usage < hard_limit && overshoot <= overshoot_bursting_range {
                // Zone 1: Even though it reaches the throttling level, the overshoot is below
                // bursting threshold and the usage is below hard_limit, we still need to use
                // the bursting mechanism to smooth the throttling.
                ENGINE_LIMITER_THROTTLE_ACTION_COUNTER
                    .with_label_values(&[Ty::TAG, "overshoot.bursting"])
                    .inc();
                calculate_bursting_speed_limit(
                    min_speed_limit,
                    max_speed_limit,
                    overshoot,
                    overshoot_bursting_range,
                )
            } else {
                // Zone 2: Forcibly throttling in emergency situation, such as high memory
                // pressure, use minimum speed limit
                ENGINE_LIMITER_THROTTLE_ACTION_COUNTER
                    .with_label_values(&[Ty::TAG, "overshoot.hard_cutoff"])
                    .inc();
                min_speed_limit as f64
            }
        } else {
            f64::INFINITY
        }
    }

    /// Calculate the speed limit for the normal mechanism.
    fn calculate_normal_speed_limit(
        usage: u64,
        soft_limit: u64,
        hard_limit: u64,
        max_speed_limit: u64,
        min_speed_limit: u64,
    ) -> f64 {
        let burst_extra = (hard_limit as f64 * DEFAULT_BURST_RATIO) as u64;
        let bursting_threshold = hard_limit.saturating_add(burst_extra);
        if usage < soft_limit {
            // Zone 3: No throttling below soft_limit
            f64::INFINITY
        } else if usage < hard_limit {
            // Zone 4: Exponential throttling curve between soft_limit and hard_limit
            ENGINE_LIMITER_THROTTLE_ACTION_COUNTER
                .with_label_values(&[Ty::TAG, "smoothing"])
                .inc();

            let range = (hard_limit - soft_limit) as f64;
            let position = (usage - soft_limit) as f64;
            let normalized = position / range.max(1.0);

            // Exponential curve: e^(-k*x) where k controls the curve steepness
            // Using k=2.0 provides a good balance between smoothness and responsiveness
            let exp_factor = (-normalized * 2.0).exp();

            // Interpolate from max_speed_limit to min_speed_limit using exponential curve
            let throttle = min_speed_limit as f64
                + (max_speed_limit as f64 - min_speed_limit as f64) * exp_factor;
            throttle.max(min_speed_limit as f64)
        } else if usage < bursting_threshold {
            // Zone 5: Bursting mechanism between hard_limit and hard_limit * 1.1
            // This provides a smoother transition from the exponential curve result
            // at hard_limit down to min_speed_limit at bursting_threshold
            ENGINE_LIMITER_THROTTLE_ACTION_COUNTER
                .with_label_values(&[Ty::TAG, "bursting"])
                .inc();

            calculate_bursting_speed_limit(
                min_speed_limit,
                max_speed_limit,
                usage.saturating_sub(hard_limit),
                bursting_threshold.saturating_sub(hard_limit),
            )
        } else {
            // Zone 6: Above bursting threshold, use minimum speed limit
            ENGINE_LIMITER_THROTTLE_ACTION_COUNTER
                .with_label_values(&[Ty::TAG, "hard_cutoff"])
                .inc();
            min_speed_limit as f64
        }
    }
}

#[derive(Clone)]
pub struct RegionLimiter {
    memtable: WriteRateLimiter<RegionMemTable>,
    l0table: WriteRateLimiter<RegionL0Table>,
}

impl RegionLimiter {
    pub fn enabled(&self) -> bool {
        self.memtable.options.enable || self.l0table.options.enable
    }

    pub fn consume(&self, region_id: u64, bytes: usize) -> Duration {
        std::cmp::max(
            self.memtable.consume(region_id, bytes),
            self.l0table.consume(region_id, bytes),
        )
    }

    pub fn unconsume(&self, region_id: u64, bytes: usize) {
        self.memtable.unconsume(region_id, bytes);
        self.l0table.unconsume(region_id, bytes);
    }

    pub fn is_unlimited(&self, region_id: u64) -> bool {
        self.memtable.is_unlimited(region_id) && self.l0table.is_unlimited(region_id)
    }

    pub fn new(memtable_options: LimiterOptions, l0table_options: LimiterOptions) -> Self {
        Self {
            memtable: WriteRateLimiter::<RegionMemTable>::new(memtable_options),
            l0table: WriteRateLimiter::<RegionL0Table>::new(l0table_options),
        }
    }

    pub fn new_from(another: &RegionLimiter) -> Self {
        Self {
            memtable: WriteRateLimiter::new_from(&another.memtable),
            l0table: WriteRateLimiter::new_from(&another.l0table),
        }
    }

    pub fn update_usage(
        &self,
        tag: &ShardTag,
        memtable_usage: u64,
        l0table_usage: u64,
        over_throttling_level: impl Fn() -> Option<(u64, u64)>,
    ) {
        self.memtable
            .update_usage(tag, memtable_usage, over_throttling_level);
        self.l0table.update_usage(tag, l0table_usage, || None);
    }
}

/// Bounds memory used by required DFS loads and best-effort cache reads.
#[derive(Clone)]
pub(crate) struct DfsLoadLimiter {
    semaphore: Arc<Semaphore>,
    required_memory_limiter: MemoryLimiter,
    cache_memory_limiter: MemoryLimiter,
    // Mirrors the RAII metric for deterministic waiter lifecycle tests. Cache
    // admission must remain independent from this state.
    waiting_required_loads: Arc<AtomicUsize>,
}

pub(crate) type DfsLoadLimiterPermit = OwnedSemaphorePermit;

/// Reason a DFS cache-read memory reservation was rejected.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DfsCacheMemoryRejection {
    /// One file is larger than the complete cache burst quota.
    Oversized,
    /// The cache burst quota is currently exhausted.
    QuotaExhausted,
    /// Process memory would reach the existing global throttling level.
    MemoryPressure,
}

impl DfsCacheMemoryRejection {
    /// Stable low-cardinality label for DFS admission metrics.
    pub fn metric_label(self) -> &'static str {
        match self {
            Self::Oversized => "oversized",
            Self::QuotaExhausted => "quota_exhausted",
            Self::MemoryPressure => "memory_pressure",
        }
    }
}

struct RequiredDfsLoadWaitGuard {
    waiting_required_loads: Arc<AtomicUsize>,
}

impl RequiredDfsLoadWaitGuard {
    fn new(waiting_required_loads: Arc<AtomicUsize>) -> Self {
        waiting_required_loads.fetch_add(1, Ordering::AcqRel);
        ENGINE_DFS_LOAD_MEMORY_WAITING_REQUIRED_LOADS.inc();
        Self {
            waiting_required_loads,
        }
    }
}

impl Drop for RequiredDfsLoadWaitGuard {
    fn drop(&mut self) {
        self.waiting_required_loads.fetch_sub(1, Ordering::AcqRel);
        ENGINE_DFS_LOAD_MEMORY_WAITING_REQUIRED_LOADS.dec();
    }
}

impl DfsLoadLimiter {
    pub(crate) fn new(cfg: &KvEngineConfig, opts: &Options) -> DfsLoadLimiter {
        let num_cores = tikv_util::sys::SysQuota::cpu_cores_quota();
        let global_concurrency = (num_cores.max(1.0) as usize) * cfg.dfs_load_concurrency_per_core;
        let semaphore = Arc::new(Semaphore::new(global_concurrency));
        let required_memory_limiter = MemoryLimiter::new(
            opts.dfs_load_memory_limit,
            Some(crate::metrics::ENGINE_DFS_LOAD_MEMORY_USAGE.clone()),
        );
        // Reuse the required-load quota as the cache burst ceiling. This keeps
        // burst memory bounded without another configuration value; the
        // process-wide throttling level below can reduce the usable amount.
        let cache_memory_limiter = MemoryLimiter::new(opts.dfs_load_memory_limit, None);
        Self {
            semaphore,
            required_memory_limiter,
            cache_memory_limiter,
            waiting_required_loads: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub(crate) async fn acquire_permit(&self) -> DfsLoadLimiterPermit {
        let global_semaphore = self.semaphore.clone();
        // We never close the semaphore, so it is safe to unwrap.
        global_semaphore.acquire_owned().await.unwrap()
    }

    pub(crate) fn available_permits(&self) -> usize {
        self.semaphore.available_permits()
    }

    /// Tries to reserve memory for a best-effort cache read.
    ///
    /// Unlike [`Self::acquire_memory_blocking`], this never waits and never
    /// bypasses accounting for an oversized file. Cache reads use
    /// process-memory headroom under a burst ceiling equal to the
    /// required-load quota. Small reservations use current required memory;
    /// larger reservations reserve the complete future required quota in
    /// the headroom check.
    pub(crate) fn try_acquire_cache_memory(
        &self,
        file_size: u64,
    ) -> Result<MemoryLimiterGuard, DfsCacheMemoryRejection> {
        if file_size > self.cache_memory_limiter.capacity() {
            return Err(DfsCacheMemoryRejection::Oversized);
        }
        if file_size == 0 {
            return Ok(self
                .cache_memory_limiter
                .acquire(0)
                .expect("acquiring 0 bytes should always succeed"));
        }

        let guard = match self.cache_memory_limiter.acquire(file_size) {
            Ok(guard) => guard,
            Err(_) => return Err(DfsCacheMemoryRejection::QuotaExhausted),
        };

        let cache_bytes = self.cache_memory_limiter.used();
        let unreserved_allowance =
            self.required_memory_limiter.capacity() / DFS_CACHE_MEMORY_UNRESERVED_ALLOWANCE_DIVISOR;

        // Reserve before checking RSS so concurrent cache admissions cannot
        // all pass using the same stale snapshot. RSS is refreshed
        // periodically, so include cache reservations it may not reflect yet.
        // The bounded unreserved allowance uses current required reservations;
        // burst cache admission reserves the complete required quota so later
        // required growth cannot consume the same apparent headroom. Existing
        // reservations already reflected in RSS can be counted twice; that is
        // intentionally conservative for optional burst work.
        let required_bytes = if cache_bytes > unreserved_allowance {
            self.required_memory_limiter.capacity()
        } else {
            self.required_memory_limiter.used()
        };
        let projected_dfs_bytes = required_bytes.saturating_add(cache_bytes);
        if memory_usage_would_reach_throttling_level(projected_dfs_bytes) {
            return Err(DfsCacheMemoryRejection::MemoryPressure);
        }
        Ok(guard)
    }

    /// Acquires memory for loading one DFS file.
    ///
    /// For large file loading we cannot proceed without memory budget; this
    /// waits until enough budget is available. If one file itself is larger
    /// than total capacity, it skips memory accounting for that file to avoid
    /// deadlock.
    pub(crate) async fn acquire_memory_blocking(&self, file_size: u64) -> MemoryLimiterGuard {
        let capacity = self.required_memory_limiter.capacity();
        if file_size > capacity {
            warn!(
                "file size exceeds dfs load memory capacity, skipping memory limit for this file";
                "file_size" => file_size,
                "capacity" => capacity,
                "file_size_mb" => file_size / (1024 * 1024),
                "capacity_mb" => capacity / (1024 * 1024),
            );
            return self
                .required_memory_limiter
                .acquire(0)
                .expect("acquiring 0 bytes should always succeed");
        }

        let start = Instant::now();
        if let Ok(guard) = self.required_memory_limiter.acquire(file_size) {
            return guard;
        }
        let _waiter = RequiredDfsLoadWaitGuard::new(self.waiting_required_loads.clone());
        loop {
            match self.required_memory_limiter.acquire(file_size) {
                Ok(guard) => {
                    ENGINE_DFS_LOAD_MEMORY_WAIT_DURATION.observe(start.saturating_elapsed_secs());
                    return guard;
                }
                Err(_) => {
                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    static GLOBAL_MEMORY_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());

    fn no_global_memory_pressure() -> Option<(u64, u64)> {
        None
    }

    #[test]
    fn test_calculate_throttle() {
        let soft_limit = 256 << 20; // 256 MB
        let hard_limit = 768 << 20; // 768 MB
        let max_speed_limit = 50 << 20; // 50 MB/s
        let min_speed_limit = 1 << 20; // 1 MB/s
        let bursting_threshold = (hard_limit as f64 * 1.1) as u64; // 110% of hard_limit
        // Test 1: Below soft_limit - should be INFINITY
        let result_below = StoreLimiter::calculate_throttle(
            0,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert_eq!(
            result_below,
            f64::INFINITY,
            "Below soft_limit should be INFINITY"
        );

        let result_just_below_soft = StoreLimiter::calculate_throttle(
            soft_limit - 1,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert_eq!(
            result_just_below_soft,
            f64::INFINITY,
            "Just below soft_limit should be INFINITY"
        );

        // Test 2: At soft_limit - should start throttling (not INFINITY)
        let result_at_soft = StoreLimiter::calculate_throttle(
            soft_limit,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert!(
            result_at_soft.is_finite()
                && result_at_soft <= max_speed_limit as f64
                && result_at_soft >= min_speed_limit as f64,
            "At soft_limit should start throttling between min and max: {}",
            result_at_soft
        );

        // Test 3: Between soft_limit and hard_limit - should use exponential curve
        let result_mid = StoreLimiter::calculate_throttle(
            (soft_limit + hard_limit) / 2,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert!(
            result_mid.is_finite()
                && result_mid < result_at_soft
                && result_mid >= min_speed_limit as f64,
            "Between soft_limit and hard_limit should throttle: {}",
            result_mid
        );

        // Test 4: At hard_limit - should be in bursting zone
        let result_at_hard = StoreLimiter::calculate_throttle(
            hard_limit,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert!(
            result_at_hard >= min_speed_limit as f64,
            "At hard_limit should be at least min_speed_limit: {}",
            result_at_hard
        );

        // Test 5: In bursting zone (between hard_limit and hard_limit * 1.1)
        let result_in_bursting = StoreLimiter::calculate_throttle(
            hard_limit + (bursting_threshold - hard_limit) / 2,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert!(
            result_in_bursting >= min_speed_limit as f64 && result_in_bursting <= result_at_hard,
            "In bursting zone should throttle smoothly: {}",
            result_in_bursting
        );

        // Test 6: At bursting threshold (hard_limit * 1.1) - should be min_speed_limit
        let result_at_bursting_threshold = StoreLimiter::calculate_throttle(
            bursting_threshold,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert_eq!(
            result_at_bursting_threshold, min_speed_limit as f64,
            "At bursting threshold should be min_speed_limit: {}",
            result_at_bursting_threshold
        );

        // Test 7: Above bursting threshold - should be min_speed_limit (hard cutoff)
        let result_above_bursting = StoreLimiter::calculate_throttle(
            bursting_threshold + (100 << 20), // 100 MB above bursting threshold
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert_eq!(
            result_above_bursting, min_speed_limit as f64,
            "Above bursting threshold should be min_speed_limit: {}",
            result_above_bursting
        );

        // Test 8: Monotonicity - throttle should decrease as usage increases
        let usages = vec![
            soft_limit - 1,
            soft_limit,
            (soft_limit + hard_limit) / 2,
            hard_limit,
            hard_limit + (bursting_threshold - hard_limit) / 2,
            bursting_threshold,
            bursting_threshold + (100 << 20),
        ];
        let mut prev_throttle = f64::INFINITY;
        for usage in usages {
            let throttle = StoreLimiter::calculate_throttle(
                usage,
                soft_limit,
                hard_limit,
                max_speed_limit,
                min_speed_limit,
                no_global_memory_pressure,
            );
            assert!(
                throttle <= prev_throttle || (prev_throttle.is_infinite() && throttle.is_finite()),
                "Throttle should be monotonic decreasing: usage={}, prev={}, current={}",
                usage,
                prev_throttle,
                throttle
            );
            prev_throttle = throttle;
        }

        // Test 9: Corner case - soft_limit == hard_limit
        let result_equal = StoreLimiter::calculate_throttle(
            512 << 20,
            512 << 20,
            512 << 20,
            50 << 20,
            1 << 20,
            no_global_memory_pressure,
        );
        assert!(
            result_equal >= min_speed_limit as f64 && result_equal <= max_speed_limit as f64,
            "When soft_limit == hard_limit, should be between min and max: {}",
            result_equal
        );

        // Test 10: Corner case - soft_limit == 0 (should return INFINITY)
        let result_zero_soft = StoreLimiter::calculate_throttle(
            100 << 20,
            0,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            no_global_memory_pressure,
        );
        assert_eq!(
            result_zero_soft,
            f64::INFINITY,
            "When soft_limit is 0, should return INFINITY"
        );
    }

    #[test]
    fn test_calculate_throttle_with_global_memory_pressure() {
        let soft_limit = 256 << 20; // 256 MB
        let hard_limit = 768 << 20; // 768 MB
        let max_speed_limit = 50 << 20; // 50 MB/s
        let min_speed_limit = 1 << 20; // 1 MB/s
        let throttling_level = 1024 << 20; // 1 GB
        let overshoot_bursting_range =
            (throttling_level as f64 * (DEFAULT_BURST_RATIO / 2.0_f64)) as u64;

        let result_with_small_overshoot = StoreLimiter::calculate_throttle(
            soft_limit - 1,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            || Some((1, throttling_level)),
        );
        assert!(
            result_with_small_overshoot.is_finite()
                && result_with_small_overshoot > min_speed_limit as f64,
            "Small global overshoot should stay on the smoothing curve: {}",
            result_with_small_overshoot
        );

        let result_with_mid_overshoot = StoreLimiter::calculate_throttle(
            soft_limit - 1,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            || Some((overshoot_bursting_range / 2, throttling_level)),
        );
        assert!(
            result_with_mid_overshoot.is_finite()
                && result_with_mid_overshoot >= min_speed_limit as f64
                && result_with_mid_overshoot < result_with_small_overshoot,
            "Larger overshoot inside the smoothing window should further reduce throttle: {} vs {}",
            result_with_mid_overshoot,
            result_with_small_overshoot
        );

        let result_at_smoothing_boundary = StoreLimiter::calculate_throttle(
            soft_limit - 1,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            || Some((overshoot_bursting_range, throttling_level)),
        );
        assert_eq!(
            result_at_smoothing_boundary, min_speed_limit as f64,
            "At the overshoot smoothing boundary the throttle should reach min_speed_limit"
        );

        let result_above_smoothing_boundary = StoreLimiter::calculate_throttle(
            soft_limit - 1,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            || Some((overshoot_bursting_range + 1, throttling_level)),
        );
        assert_eq!(
            result_above_smoothing_boundary, min_speed_limit as f64,
            "Above the overshoot smoothing boundary the throttle should clamp directly"
        );

        let result_at_local_hard_limit = StoreLimiter::calculate_throttle(
            hard_limit,
            soft_limit,
            hard_limit,
            max_speed_limit,
            min_speed_limit,
            || Some((1, throttling_level)),
        );
        assert_eq!(
            result_at_local_hard_limit, min_speed_limit as f64,
            "Local usage at hard_limit should bypass overshoot smoothing"
        );
    }

    #[tokio::test]
    async fn test_acquire_memory_blocking() {
        use crate::Options;

        let _memory_lock = GLOBAL_MEMORY_TEST_LOCK.lock().await;

        // Test case 1: file_size within capacity - should acquire memory successfully
        {
            let config = KvEngineConfig {
                dfs_load_concurrency_per_core: 64,
                dfs_load_memory_ratio: 0.25,
                ..Default::default()
            };
            let mut opts = Options::default();
            opts.dfs_load_memory_limit = 1024 * 1024; // 1 MB capacity

            let limiter = DfsLoadLimiter::new(&config, &opts);

            // Try to acquire 512 KB (within capacity)
            let file_size = 512 * 1024;
            let guard = limiter.acquire_memory_blocking(file_size).await;

            // Should succeed and track the memory
            assert_eq!(limiter.required_memory_limiter.used(), file_size);

            // Drop the guard to release memory
            drop(guard);
            assert_eq!(limiter.required_memory_limiter.used(), 0);
        }

        // Test case 2: file_size exceeds capacity - should skip memory limiting
        {
            let config = KvEngineConfig {
                dfs_load_concurrency_per_core: 64,
                dfs_load_memory_ratio: 0.25,
                ..Default::default()
            };
            let mut opts = Options::default();
            opts.dfs_load_memory_limit = 1024 * 1024; // 1 MB capacity

            let limiter = DfsLoadLimiter::new(&config, &opts);

            // Try to acquire 2 MB (exceeds capacity)
            let file_size = 2 * 1024 * 1024;
            let guard = limiter.acquire_memory_blocking(file_size).await;

            // Should succeed but not track the memory (returns 0-size guard)
            // Memory usage should be 0 because we skipped memory limiting
            assert_eq!(limiter.required_memory_limiter.used(), 0);

            // Drop the guard - should still be 0
            drop(guard);
            assert_eq!(limiter.required_memory_limiter.used(), 0);
        }

        // Test case 3: Multiple concurrent acquisitions within capacity
        {
            let config = KvEngineConfig {
                dfs_load_concurrency_per_core: 64,
                dfs_load_memory_ratio: 0.25,
                ..Default::default()
            };
            let mut opts = Options::default();
            opts.dfs_load_memory_limit = 1024 * 1024; // 1 MB capacity

            let limiter = DfsLoadLimiter::new(&config, &opts);

            // Acquire 400 KB three times in sequence (total 1.2 MB would exceed)
            let file_size = 400 * 1024;

            let guard1 = limiter.acquire_memory_blocking(file_size).await;
            assert_eq!(limiter.required_memory_limiter.used(), file_size);

            let guard2 = limiter.acquire_memory_blocking(file_size).await;
            assert_eq!(limiter.required_memory_limiter.used(), file_size * 2);

            // Release first guard
            drop(guard1);
            assert_eq!(limiter.required_memory_limiter.used(), file_size);

            // Now we can acquire another one
            let guard3 = limiter.acquire_memory_blocking(file_size).await;
            assert_eq!(limiter.required_memory_limiter.used(), file_size * 2);

            drop(guard2);
            drop(guard3);
            assert_eq!(limiter.required_memory_limiter.used(), 0);
        }

        // Test case 4: Verify wait duration metric is recorded when blocking occurs
        {
            let config = KvEngineConfig {
                dfs_load_concurrency_per_core: 64,
                dfs_load_memory_ratio: 0.25,
                ..Default::default()
            };
            let mut opts = Options::default();
            opts.dfs_load_memory_limit = 500 * 1024; // 500 KB capacity

            let limiter = DfsLoadLimiter::new(&config, &opts);

            // Get the initial metric count
            let initial_count =
                crate::metrics::ENGINE_DFS_LOAD_MEMORY_WAIT_DURATION.get_sample_count();

            // Acquire 400 KB - should succeed immediately (no wait)
            let guard1 = limiter.acquire_memory_blocking(400 * 1024).await;
            assert_eq!(limiter.required_memory_limiter.used(), 400 * 1024);

            // Spawn a task that will try to acquire 300 KB (total would be 700 KB > 500 KB)
            // This should block and wait
            let limiter_clone = limiter.clone();
            let handle =
                tokio::spawn(
                    async move { limiter_clone.acquire_memory_blocking(300 * 1024).await },
                );

            // Give the task time to start and begin waiting
            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;

            // Release the first guard to allow the second task to proceed
            drop(guard1);

            // Wait for the second task to complete
            let guard2 = handle.await.unwrap();
            assert_eq!(limiter.required_memory_limiter.used(), 300 * 1024);

            // Verify the wait duration metric was recorded
            let final_count =
                crate::metrics::ENGINE_DFS_LOAD_MEMORY_WAIT_DURATION.get_sample_count();
            assert!(
                final_count > initial_count,
                "Wait duration metric should be recorded when blocking occurs. Initial: {}, Final: {}",
                initial_count,
                final_count
            );

            drop(guard2);
            assert_eq!(limiter.required_memory_limiter.used(), 0);
        }
    }

    #[tokio::test]
    async fn test_cache_memory_admission_is_independent_from_required_waiting() {
        use tikv_util::sys::{
            get_global_memory_usage, memory_usage_throttling_level,
            register_memory_usage_throttling_level, set_memory_usage_for_test,
        };

        use crate::Options;

        let _memory_lock = GLOBAL_MEMORY_TEST_LOCK.lock().await;
        let previous_usage = get_global_memory_usage();
        let previous_throttling_level = memory_usage_throttling_level();
        set_memory_usage_for_test(0);
        register_memory_usage_throttling_level(3 * 1024 * 1024);

        let config = KvEngineConfig {
            dfs_load_concurrency_per_core: 64,
            dfs_load_memory_ratio: 0.25,
            ..Default::default()
        };
        let mut opts = Options::default();
        opts.dfs_load_memory_limit = 1024 * 1024;
        let limiter = DfsLoadLimiter::new(&config, &opts);

        let blocking_guard = limiter.acquire_memory_blocking(1024 * 1024).await;
        let waiting_limiter = limiter.clone();
        let waiting_task =
            tokio::spawn(async move { waiting_limiter.acquire_memory_blocking(400 * 1024).await });
        tokio::time::timeout(Duration::from_secs(5), async {
            while limiter.waiting_required_loads.load(Ordering::Acquire) == 0 {
                tokio::task::yield_now().await;
            }
        })
        .await
        .unwrap();
        assert_eq!(limiter.waiting_required_loads.load(Ordering::Acquire), 1);

        // Required and cache memory have independent quotas. A required waiter
        // does not reject cache work that passes the full-T headroom check.
        assert_eq!(limiter.required_memory_limiter.used(), 1024 * 1024);
        let full_cache_guard = limiter.try_acquire_cache_memory(1024 * 1024).unwrap();
        assert_eq!(limiter.cache_memory_limiter.used(), 1024 * 1024);
        assert!(matches!(
            limiter.try_acquire_cache_memory(1),
            Err(DfsCacheMemoryRejection::QuotaExhausted)
        ));
        drop(full_cache_guard);
        set_memory_usage_for_test(2 * 1024 * 1024);
        assert!(matches!(
            limiter.try_acquire_cache_memory(1),
            Err(DfsCacheMemoryRejection::MemoryPressure)
        ));
        set_memory_usage_for_test(0);

        // Canceling a required waiter must release both its internal count and
        // the externally visible gauge through the same RAII guard.
        waiting_task.abort();
        match waiting_task.await {
            Err(join_err) => assert!(join_err.is_cancelled()),
            Ok(_) => panic!("required DFS load waiter was not cancelled"),
        }
        assert_eq!(limiter.waiting_required_loads.load(Ordering::Acquire), 0);
        drop(blocking_guard);
        assert_eq!(limiter.cache_memory_limiter.used(), 0);

        assert!(matches!(
            limiter.try_acquire_cache_memory(1024 * 1024 + 1),
            Err(DfsCacheMemoryRejection::Oversized)
        ));

        set_memory_usage_for_test(previous_usage);
        register_memory_usage_throttling_level(previous_throttling_level);
    }

    #[tokio::test]
    async fn test_cache_memory_admission_checks_throttling_level() {
        use tikv_util::sys::{
            get_global_memory_usage, memory_usage_throttling_level,
            register_memory_usage_throttling_level, set_memory_usage_for_test,
        };

        use crate::Options;

        let _memory_lock = GLOBAL_MEMORY_TEST_LOCK.lock().await;
        let previous_usage = get_global_memory_usage();
        let previous_throttling_level = memory_usage_throttling_level();
        set_memory_usage_for_test(600 * 1024);
        register_memory_usage_throttling_level(1024 * 1024);

        let config = KvEngineConfig {
            dfs_load_concurrency_per_core: 64,
            dfs_load_memory_ratio: 0.25,
            ..Default::default()
        };
        let mut opts = Options::default();
        opts.dfs_load_memory_limit = 256 * 1024;
        let limiter = DfsLoadLimiter::new(&config, &opts);

        // Burst admission reserves the complete 256 KiB required quota even
        // though required usage is currently zero.
        assert!(matches!(
            limiter.try_acquire_cache_memory(168 * 1024),
            Err(DfsCacheMemoryRejection::MemoryPressure)
        ));
        let guard = limiter.try_acquire_cache_memory(167 * 1024).unwrap();
        // The stale RSS snapshot does not include the first reservation. The
        // next admission must still account for it.
        assert!(matches!(
            limiter.try_acquire_cache_memory(1024),
            Err(DfsCacheMemoryRejection::MemoryPressure)
        ));
        drop(guard);

        // The unreserved allowance never overrides the global memory
        // throttling level.
        set_memory_usage_for_test(1008 * 1024);
        assert!(matches!(
            limiter.try_acquire_cache_memory(16 * 1024),
            Err(DfsCacheMemoryRejection::MemoryPressure)
        ));

        set_memory_usage_for_test(previous_usage);
        register_memory_usage_throttling_level(previous_throttling_level);
    }
}