// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

mod context;
mod future_pool;
mod metrics;
mod tokio_pool;

use std::{
    array,
    cell::RefCell,
    collections::HashMap,
    mem::size_of,
    sync::{
        Arc,
        atomic::{AtomicU64, AtomicUsize, Ordering},
    },
    time::Instant as StdInstant,
};

use cpu_time::ThreadTime;
use crossbeam::queue::SegQueue;
use fail::fail_point;
pub use future_pool::{Full, FuturePool};
use lazy_static::lazy_static;
use parking_lot::{Mutex, RwLock};
use prometheus::{Counter, Histogram, IntGauge, local::LocalHistogram};
use rand::prelude::*;
pub use tokio_pool::{ScalableTokioHandle, ScalableTokioRuntime};
use yatp::{
    ThreadPool,
    pool::{CloneRunnerBuilder, Local, Runner},
    queue::{
        self, CustomBuilder, CustomConfig, Extras, PopResult, QueueType, TaskCell as YatpTaskCell,
        TaskQueue as YatpTaskQueue, multilevel,
    },
    task::future::{Runner as FutureRunner, TaskCell},
};

pub use crate::yatp_pool::context::YatpDispatchContext;
use crate::{
    fairness::{FairnessRecorder, GroupInnerQueue, MultiGroupFairQueue},
    thread_group::GroupProperties,
    time::{Duration, Instant, LazyInstant},
    yatp_pool::context::YatpDispatchPrepareContext,
};

pub(crate) const TICK_INTERVAL: Duration = Duration::from_secs(1);
const FAIR_CPU_CHECKPOINT_MIN_INTERVAL: Duration = Duration::from_millis(1);
const TASK_METADATA_KEYSPACE_ID_LEN: usize = size_of::<u32>();
const TASK_METADATA_FLAG_FAIR_CPU_GATED: u8 = 0b0000_0001;
const TASK_METADATA_FLAG_FAIR_CPU_HAS_DEADLINE: u8 = 0b0000_0010;
const TASK_METADATA_WITH_FAIR_CPU_ADMISSION_LEN: usize =
    size_of::<u32>() + size_of::<u8>() + size_of::<u8>() + size_of::<u64>();
const DEFAULT_TASK_KEYSPACE_ID: u32 = 0;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FairCpuAdmissionMetadata {
    pub task_type: u8,
    pub deadline: Option<StdInstant>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TaskMetadata {
    pub keyspace_id: u32,
    pub fair_cpu_admission: Option<FairCpuAdmissionMetadata>,
}

lazy_static! {
    static ref PROCESS_EPOCH: StdInstant = StdInstant::now();
}

fn encode_deadline(deadline: StdInstant) -> u64 {
    deadline
        .saturating_duration_since(*PROCESS_EPOCH)
        .as_nanos() as u64
}

fn decode_deadline(deadline_ns: u64) -> StdInstant {
    *PROCESS_EPOCH + Duration::from_nanos(deadline_ns)
}

pub fn set_task_meta(extras: &mut Extras, meta: &TaskMetadata) {
    match meta.fair_cpu_admission {
        Some(fair_cpu_admission) => {
            let mut metadata = Vec::with_capacity(TASK_METADATA_WITH_FAIR_CPU_ADMISSION_LEN);
            metadata.extend_from_slice(&meta.keyspace_id.to_be_bytes());
            let mut flags = TASK_METADATA_FLAG_FAIR_CPU_GATED;
            let deadline_ns = fair_cpu_admission.deadline.map_or(0, |deadline| {
                flags |= TASK_METADATA_FLAG_FAIR_CPU_HAS_DEADLINE;
                encode_deadline(deadline)
            });
            metadata.push(flags);
            metadata.push(fair_cpu_admission.task_type);
            metadata.extend_from_slice(&deadline_ns.to_be_bytes());
            extras.set_metadata(metadata);
        }
        None => extras.set_metadata(meta.keyspace_id.to_be_bytes().to_vec()),
    }
}

fn task_metadata_from_metadata(metadata: &[u8]) -> TaskMetadata {
    let keyspace_id = if metadata.len() == TASK_METADATA_KEYSPACE_ID_LEN
        || metadata.len() == TASK_METADATA_WITH_FAIR_CPU_ADMISSION_LEN
    {
        u32::from_be_bytes([metadata[0], metadata[1], metadata[2], metadata[3]])
    } else {
        DEFAULT_TASK_KEYSPACE_ID
    };
    let flags = if metadata.len() == TASK_METADATA_WITH_FAIR_CPU_ADMISSION_LEN {
        metadata[4]
    } else {
        0
    };
    let is_fair_cpu_gated = flags & TASK_METADATA_FLAG_FAIR_CPU_GATED != 0;
    let has_deadline = flags & TASK_METADATA_FLAG_FAIR_CPU_HAS_DEADLINE != 0;
    let fair_cpu_admission =
        if metadata.len() == TASK_METADATA_WITH_FAIR_CPU_ADMISSION_LEN && is_fair_cpu_gated {
            Some(FairCpuAdmissionMetadata {
                task_type: metadata[5],
                deadline: has_deadline.then(|| {
                    decode_deadline(u64::from_be_bytes([
                        metadata[6],
                        metadata[7],
                        metadata[8],
                        metadata[9],
                        metadata[10],
                        metadata[11],
                        metadata[12],
                        metadata[13],
                    ]))
                }),
            })
        } else {
            None
        };
    TaskMetadata {
        keyspace_id,
        fair_cpu_admission,
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FairCpuBudgetAction {
    Continue,
    YieldToScheduler,
}

pub trait FairCpuPermit: Send + 'static {
    /// Consume actual CPU used since the previous dispatch checkpoint.
    ///
    /// This operation must not allocate memory or wait. Returning
    /// [`FairCpuBudgetAction::YieldToScheduler`] asks the DAG executor to yield
    /// after the current `next_batch()` completes. Any CPU used after the last
    /// executor checkpoint is consumed when the YATP dispatch guard drops.
    fn consume(&mut self, cpu_delta_us: u64) -> FairCpuBudgetAction;
}

pub enum FairAdmitOutcome {
    Bypassed,
    Admitted(Box<dyn FairCpuPermit>),
}

thread_local! {
    static FAIR_CPU_PERMIT: RefCell<Option<Box<dyn FairCpuPermit>>> = const { RefCell::new(None) };
    static FAIR_CPU_DISPATCH: RefCell<Option<FairCpuDispatch>> = const { RefCell::new(None) };
}

fn install_fair_cpu_permit(permit: Box<dyn FairCpuPermit>) {
    FAIR_CPU_PERMIT.with(|slot| {
        let old = slot.borrow_mut().replace(permit);
        debug_assert!(
            old.is_none(),
            "stale fair CPU permit should be swept before install"
        );
        drop(old);
    });
}

fn take_fair_cpu_permit() -> Option<Box<dyn FairCpuPermit>> {
    FAIR_CPU_PERMIT.with(|slot| slot.borrow_mut().take())
}

fn has_fair_cpu_permit() -> bool {
    FAIR_CPU_PERMIT.with(|slot| slot.borrow().is_some())
}

/// Whether the current YATP task is running with fair CPU dispatch accounting.
#[inline]
pub fn is_fair_cpu_dispatch_active() -> bool {
    FAIR_CPU_DISPATCH.with(|slot| slot.borrow().is_some())
}

struct FairCpuDispatch {
    cpu_timer: ThreadTime,
    consumed_cpu_us: u64,
    last_checkpoint_at: StdInstant,
}

impl FairCpuDispatch {
    fn new() -> Self {
        Self {
            cpu_timer: ThreadTime::now(),
            consumed_cpu_us: 0,
            last_checkpoint_at: StdInstant::now(),
        }
    }

    /// Returns CPU consumed since the previous completed checkpoint.
    fn take_cpu_delta_us(&mut self) -> u64 {
        let total_cpu_us = self.cpu_timer.elapsed().as_micros() as u64;
        let cpu_delta_us = total_cpu_us.saturating_sub(self.consumed_cpu_us);
        self.consumed_cpu_us = total_cpu_us;
        cpu_delta_us
    }

    /// Returns accumulated CPU only after the minimum checkpoint interval.
    ///
    /// The wall-clock comparison avoids the more expensive thread CPU clock,
    /// token-bucket, and metrics path when safe points are very frequent.
    fn take_cpu_delta_us_if_due(&mut self, now: StdInstant) -> Option<u64> {
        if now.saturating_duration_since(self.last_checkpoint_at) < FAIR_CPU_CHECKPOINT_MIN_INTERVAL
        {
            return None;
        }
        self.last_checkpoint_at = now;
        Some(self.take_cpu_delta_us())
    }
}

fn consume_fair_cpu_delta(cpu_delta_us: u64) -> FairCpuBudgetAction {
    if cpu_delta_us == 0 {
        return FairCpuBudgetAction::Continue;
    }
    FAIR_CPU_PERMIT.with(|slot| {
        slot.borrow_mut()
            .as_deref_mut()
            .map(|permit| permit.consume(cpu_delta_us))
            .unwrap_or(FairCpuBudgetAction::Continue)
    })
}

/// Consume CPU used by the current YATP dispatch when its checkpoint interval
/// has elapsed.
///
/// The function is a no-op outside a fair-gated YATP dispatch. Calls made less
/// than one millisecond after the previous completed checkpoint return before
/// reading the thread CPU clock or touching the shared token bucket. Callers
/// should yield to YATP when it returns
/// [`FairCpuBudgetAction::YieldToScheduler`].
pub fn fair_cpu_checkpoint(now: &mut LazyInstant) -> FairCpuBudgetAction {
    let cpu_delta_us = FAIR_CPU_DISPATCH.with(|slot| {
        let mut slot = slot.borrow_mut();
        let Some(dispatch) = slot.as_mut() else {
            return 0;
        };
        dispatch.take_cpu_delta_us_if_due(now.now()).unwrap_or(0)
    });
    consume_fair_cpu_delta(cpu_delta_us)
}

/// Yields the current YATP task when its fair CPU dispatch budget is
/// exhausted.
///
/// This function is a no-op when called outside a fair-gated YATP dispatch.
///
/// Returns whether the task yielded. Callers must not reuse `now` after this
/// function returns `true`, because the await may resume in a new poll.
pub async fn yield_if_fair_cpu_exhausted(now: &mut LazyInstant) -> bool {
    if !matches!(
        fair_cpu_checkpoint(now),
        FairCpuBudgetAction::YieldToScheduler
    ) {
        return false;
    }
    yatp::task::future::yield_to_scheduler().await;
    true
}

/// Unconditionally consumes CPU remaining at the end of a dispatch.
///
/// Tail settlement bypasses the time limit so dispatches shorter than one
/// millisecond cannot lose their final CPU usage.
fn settle_fair_cpu_tail() {
    let cpu_delta_us = FAIR_CPU_DISPATCH.with(|slot| {
        slot.borrow_mut()
            .as_mut()
            .map(FairCpuDispatch::take_cpu_delta_us)
            .unwrap_or(0)
    });
    // The dispatch is ending, so a scheduling action cannot be used here.
    let _ = consume_fair_cpu_delta(cpu_delta_us);
}

struct FairCpuGuard;

impl FairCpuGuard {
    fn new() -> Self {
        FAIR_CPU_DISPATCH.with(|slot| {
            let old = slot.borrow_mut().replace(FairCpuDispatch::new());
            debug_assert!(old.is_none(), "fair CPU dispatch must not be nested");
        });
        Self
    }
}

impl Drop for FairCpuGuard {
    fn drop(&mut self) {
        settle_fair_cpu_tail();
        FAIR_CPU_DISPATCH.with(|slot| {
            slot.borrow_mut().take();
        });
        drop(take_fair_cpu_permit());
    }
}

fn maybe_fair_cpu_guard() -> Option<FairCpuGuard> {
    has_fair_cpu_permit().then(FairCpuGuard::new)
}

fn clear_stale_fair_cpu_permit() {
    debug_assert!(FAIR_CPU_DISPATCH.with(|slot| slot.borrow().is_none()));
    drop(take_fair_cpu_permit());
}

fn fair_cpu_metadata(task_cell: &mut TaskCell) -> Option<(u8, Option<StdInstant>)> {
    let metadata = task_metadata(task_cell);
    metadata
        .fair_cpu_admission
        .map(|admission| (admission.task_type, admission.deadline))
}

fn task_keyspace_id_from_metadata(metadata: &[u8]) -> u32 {
    task_metadata_from_metadata(metadata).keyspace_id
}

fn task_keyspace_id(task_cell: &mut TaskCell) -> u32 {
    task_keyspace_id_from_metadata(task_cell.mut_extras().metadata())
}

fn task_metadata(task_cell: &mut TaskCell) -> TaskMetadata {
    task_metadata_from_metadata(task_cell.mut_extras().metadata())
}

fn tick_interval() -> Duration {
    fail_point!("mock_tick_interval", |_| { Duration::from_millis(1) });
    TICK_INTERVAL
}

#[cfg(test)]
thread_local! {
    static MOCK_NOW_COARSE_SECS: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
}

#[cfg(test)]
fn set_now_coarse_secs_for_test(now_secs: Option<u64>) {
    MOCK_NOW_COARSE_SECS.set(now_secs);
}

pub trait PoolTicker: Send + Clone + 'static {
    fn on_tick(&mut self);
}

pub enum FairQueueAdmissionResult {
    Ready,
    Bypassed,
    Pending { retry_at: std::time::Instant },
}

pub trait FairQueueAdmissionController: Send + Sync + 'static {
    fn precheck(&self) -> FairQueueAdmissionResult;

    fn admit(
        &self,
        keyspace_id: u32,
        task_type: u8,
        deadline: Option<StdInstant>,
    ) -> FairAdmitOutcome;
}

#[derive(Clone)]
pub struct TickerWrapper<T: PoolTicker> {
    ticker: T,
    last_tick_time: Instant,
}

impl<T: PoolTicker> TickerWrapper<T> {
    pub fn new(ticker: T) -> Self {
        Self {
            ticker,
            last_tick_time: Instant::now_coarse(),
        }
    }

    // Returns whether tick has been triggered.
    pub fn try_tick(&mut self) -> bool {
        let now = Instant::now_coarse();
        if now.saturating_duration_since(self.last_tick_time) < tick_interval() {
            return false;
        }
        self.last_tick_time = now;
        self.ticker.on_tick();
        true
    }

    pub fn on_tick(&mut self) {
        self.ticker.on_tick();
    }
}

#[derive(Clone, Copy)]
pub struct DefaultTicker;

impl PoolTicker for DefaultTicker {
    fn on_tick(&mut self) {}
}

#[derive(Debug, Clone, Copy)]
pub struct Config {
    pub workers: usize,
    pub max_tasks_per_worker: usize,
    pub stack_size: usize,
}

impl Config {
    pub fn default_for_test() -> Self {
        Self {
            workers: 2,
            max_tasks_per_worker: usize::MAX,
            stack_size: 2_000_000,
        }
    }
}

#[derive(Clone)]
pub struct YatpPoolRunner<T: PoolTicker> {
    inner: FutureRunner,
    ticker: TickerWrapper<T>,
    props: Option<GroupProperties>,
    after_start: Option<Arc<dyn Fn() + Send + Sync>>,
    before_stop: Option<Arc<dyn Fn() + Send + Sync>>,
    before_pause: Option<Arc<dyn Fn() + Send + Sync>>,

    // Statistics about the schedule wait duration.
    schedule_wait_duration: LocalHistogram,
}

impl<T: PoolTicker> Runner for YatpPoolRunner<T> {
    type TaskCell = TaskCell;

    fn start(&mut self, local: &mut Local<Self::TaskCell>) {
        crate::sys::thread::add_thread_name_to_map();
        if let Some(props) = self.props.take() {
            crate::thread_group::set_properties(Some(props));
        }
        self.inner.start(local);
        if let Some(f) = self.after_start.take() {
            f();
        }
        unsafe {
            tikv_alloc::add_thread_memory_accessor();
        }
    }

    fn handle(&mut self, local: &mut Local<Self::TaskCell>, mut task_cell: Self::TaskCell) -> bool {
        let _ctx_guard = YatpDispatchContext::install(&mut task_cell);
        let extras = task_cell.mut_extras();
        if let Some(schedule_time) = extras.schedule_time() {
            self.schedule_wait_duration
                .observe(schedule_time.elapsed().as_secs_f64());
        }
        let _fair_cpu_guard = maybe_fair_cpu_guard();
        let finished = self.inner.handle(local, task_cell);
        if self.ticker.try_tick() {
            self.schedule_wait_duration.flush();
        }
        finished
    }

    fn pause(&mut self, local: &mut Local<Self::TaskCell>) -> bool {
        if let Some(f) = self.before_pause.as_ref() {
            f();
        }
        self.inner.pause(local)
    }

    fn resume(&mut self, local: &mut Local<Self::TaskCell>) {
        self.inner.resume(local)
    }

    fn end(&mut self, local: &mut Local<Self::TaskCell>) {
        if let Some(f) = self.before_stop.as_ref() {
            f();
        }
        self.ticker.on_tick();
        self.inner.end(local);
        tikv_alloc::remove_thread_memory_accessor();
        crate::sys::thread::remove_thread_name_from_map()
    }
}

impl<T: PoolTicker> YatpPoolRunner<T> {
    pub fn new(
        inner: FutureRunner,
        ticker: TickerWrapper<T>,
        after_start: Option<Arc<dyn Fn() + Send + Sync>>,
        before_stop: Option<Arc<dyn Fn() + Send + Sync>>,
        before_pause: Option<Arc<dyn Fn() + Send + Sync>>,
        schedule_wait_duration: Histogram,
    ) -> Self {
        YatpPoolRunner {
            inner,
            ticker,
            props: crate::thread_group::current_properties(),
            after_start,
            before_stop,
            before_pause,
            schedule_wait_duration: schedule_wait_duration.local(),
        }
    }
}

pub struct YatpPoolBuilder<T: PoolTicker> {
    name_prefix: Option<String>,
    ticker: TickerWrapper<T>,
    after_start: Option<Arc<dyn Fn() + Send + Sync>>,
    before_stop: Option<Arc<dyn Fn() + Send + Sync>>,
    before_pause: Option<Arc<dyn Fn() + Send + Sync>>,
    keyspace_metrics: Option<Arc<KeyspaceMetrics>>,
    min_thread_count: usize,
    core_thread_count: usize,
    max_thread_count: usize,
    stack_size: usize,
    max_tasks: usize,
}

impl<T: PoolTicker> YatpPoolBuilder<T> {
    pub fn new(ticker: T) -> Self {
        Self {
            ticker: TickerWrapper::new(ticker),
            name_prefix: None,
            after_start: None,
            before_stop: None,
            before_pause: None,
            keyspace_metrics: None,
            min_thread_count: 1,
            core_thread_count: 1,
            max_thread_count: 1,
            stack_size: 0,
            max_tasks: usize::MAX,
        }
    }

    pub fn config(&mut self, config: Config) -> &mut Self {
        // TODO: maybe we should use (1, num_cpu) for min and max thread count.
        self.thread_count(config.workers, config.workers, config.workers)
            .stack_size(config.stack_size)
            .max_tasks(config.workers.saturating_mul(config.max_tasks_per_worker))
    }

    pub fn stack_size(&mut self, val: usize) -> &mut Self {
        self.stack_size = val;
        self
    }

    pub fn name_prefix(&mut self, val: impl Into<String>) -> &mut Self {
        let name = val.into();
        self.name_prefix = Some(name);
        self
    }

    pub fn thread_count(
        &mut self,
        min_thread_count: usize,
        core_thread_count: usize,
        max_thread_count: usize,
    ) -> &mut Self {
        self.min_thread_count = min_thread_count;
        self.core_thread_count = core_thread_count;
        self.max_thread_count = max_thread_count;
        self
    }

    pub fn max_tasks(&mut self, tasks: usize) -> &mut Self {
        self.max_tasks = tasks;
        self
    }

    /// Sets metrics used by the keyspace-fair pool.
    pub fn keyspace_metrics(&mut self, metrics: Option<Arc<KeyspaceMetrics>>) -> &mut Self {
        self.keyspace_metrics = metrics;
        self
    }

    pub fn before_stop<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.before_stop = Some(Arc::new(f));
        self
    }

    pub fn after_start<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.after_start = Some(Arc::new(f));
        self
    }

    pub fn before_pause<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.before_pause = Some(Arc::new(f));
        self
    }

    pub fn build_future_pool(&mut self) -> FuturePool {
        let pool = self.build_single_level_pool();
        let name = self.name_prefix.as_deref().unwrap_or("yatp_pool");
        FuturePool::from_pool(pool, name, self.core_thread_count, self.max_tasks)
    }

    pub fn build_single_level_pool(&mut self) -> ThreadPool<TaskCell> {
        let (builder, runner) = self.create_builder();
        builder.build_with_queue_and_runner(
            yatp::queue::QueueType::SingleLevel,
            yatp::pool::CloneRunnerBuilder(runner),
        )
    }

    pub fn build_multi_level_pool(&mut self) -> ThreadPool<TaskCell> {
        let (builder, read_pool_runner) = self.create_builder();
        let name = self.name_prefix.as_deref().unwrap_or("yatp_pool");
        let multilevel_builder =
            multilevel::Builder::new(multilevel::Config::default().name(Some(name)));
        let runner_builder =
            multilevel_builder.runner_builder(CloneRunnerBuilder(read_pool_runner));
        builder
            .build_with_queue_and_runner(QueueType::Multilevel(multilevel_builder), runner_builder)
    }

    pub fn build_keyspace_fair_pool(
        &mut self,
        admission_controller: Option<Arc<dyn FairQueueAdmissionController>>,
    ) -> (ThreadPool<TaskCell>, Arc<YatpKeyspaceFairQueue>) {
        let (builder, read_pool_runner) = self.create_builder();
        let name = self.name_prefix.as_deref().unwrap_or("yatp_pool");
        let queue = Arc::new(YatpKeyspaceFairQueue::new(
            self.keyspace_metrics.clone(),
            admission_controller,
        ));
        let keyspace_fair_queue = queue.clone();
        let custom_builder = CustomBuilder::new(CustomConfig::default().name(Some(name)), queue);
        let runner_builder = custom_builder.runner_builder(CloneRunnerBuilder(read_pool_runner));
        (
            builder.build_with_queue_and_runner(QueueType::Custom(custom_builder), runner_builder),
            keyspace_fair_queue,
        )
    }

    fn create_builder(&mut self) -> (yatp::Builder, YatpPoolRunner<T>) {
        let name = self.name_prefix.as_deref().unwrap_or("yatp_pool");
        let mut builder = yatp::Builder::new(thd_name!(name));
        builder
            .stack_size(self.stack_size)
            .min_thread_count(self.min_thread_count)
            .core_thread_count(self.core_thread_count)
            .max_thread_count(self.max_thread_count);

        let after_start = self.after_start.take();
        let before_stop = self.before_stop.take();
        let before_pause = self.before_pause.take();
        let schedule_wait_duration =
            metrics::YATP_POOL_SCHEDULE_WAIT_DURATION_VEC.with_label_values(&[name]);
        let read_pool_runner = YatpPoolRunner::new(
            Default::default(),
            self.ticker.clone(),
            after_start,
            before_stop,
            before_pause,
            schedule_wait_duration,
        );
        (builder, read_pool_runner)
    }
}

/// Task cell stored inside [`YatpKeyspaceFairQueue`].
///
/// The wrapper keeps queue-local bookkeeping with the task while preserving the
/// original YATP future [`TaskCell`] for execution.
struct KeyspaceFairTaskCell {
    task_cell: TaskCell,
    fair_queue_popped_task_count_at_push: u64,
    // current_pop_waits only presents when task cell is pop
    current_pop_waits: Option<u64>,
}

impl KeyspaceFairTaskCell {
    fn new(task_cell: TaskCell, fair_queue_popped_task_count_at_push: u64) -> Self {
        Self {
            task_cell,
            fair_queue_popped_task_count_at_push,
            current_pop_waits: None,
        }
    }

    fn into_task_cell(self) -> TaskCell {
        self.task_cell
    }

    #[inline]
    fn set_current_pop_waits(&mut self, current_pop_waits: u64) {
        self.current_pop_waits = Some(current_pop_waits);
    }
}

impl YatpTaskCell for KeyspaceFairTaskCell {
    fn mut_extras(&mut self) -> &mut Extras {
        self.task_cell.mut_extras()
    }
}

/// Per-keyspace inner queue used by [`YatpKeyspaceFairQueue`].
///
/// The outer [`MultiGroupFairQueue`] decides which keyspace should run next by
/// virtual time. Once a keyspace is selected, this inner queue chooses a task
/// inside that keyspace with YATP-style multilevel feedback scheduling.
///
/// Tasks are assigned to levels by [`multilevel::TaskLevelManager`], using the
/// same default thresholds as YATP's built-in multilevel queue: level 1 after
/// 5ms accumulated running time, level 2 after 100ms, and a 10s elapsed-map
/// cleanup interval. Pop selection follows the built-in queue's initial level
/// preference: level 0 is selected with 80% probability, while the remaining
/// levels are selected with a 4:1 chance ratio. If the selected level is empty,
/// pop scans the following levels and returns the first available task.
///
/// This is intentionally a lightweight group-local queue, not a full copy of
/// YATP's built-in multilevel queue. It has no worker-local queues, and it does
/// not run the built-in dynamic `maybe_adjust_chance` logic that adjusts
/// `level0_chance` and the lowest-priority global steal batch size. The outer
/// fair queue already provides keyspace fairness; this type preserves the
/// usual multilevel priority ordering within one selected keyspace.
///
/// Concurrency follows the [`GroupInnerQueue`] contract: after `push` returns,
/// the task is counted by `len` until it is popped. The count is used by the
/// outer fair queue as a readiness hint for this keyspace, so queue operations
/// should stay short and non-blocking.
struct MultiLevelInnerQueue {
    /// Metric handles for the keyspace represented by this inner queue.
    metrics: Option<Arc<KeyspaceInnerMetricsHandle>>,
    /// Number of task slices popped from the keyspace fair queue for
    /// scheduling.
    fair_queue_popped_task_count: Arc<AtomicU64>,
    /// Fixed time origin captured when this inner queue is created. It is the
    /// encoding origin for `last_task_slice_pop_time_nanos`.
    task_slice_time_origin: StdInstant,
    /// Latest task-slice pop time as nanoseconds after the origin, plus one.
    /// Zero means this keyspace has not popped a slice yet.
    last_task_slice_pop_time_nanos: AtomicU64,
    /// Tracks cumulative task running time and updates each task's level.
    level_manager: multilevel::TaskLevelManager,
    /// One concurrent FIFO queue per YATP multilevel priority level.
    levels: [SegQueue<KeyspaceFairTaskCell>; multilevel::LEVEL_NUM],
    /// Count of tasks covered by this queue; it may include an in-progress push
    /// before the task is visible in levels.
    len: AtomicUsize,
}

impl MultiLevelInnerQueue {
    fn new_with_config(
        metrics: Option<Arc<KeyspaceInnerMetricsHandle>>,
        fair_queue_popped_task_count: Arc<AtomicU64>,
        level_time_threshold: [Duration; multilevel::LEVEL_NUM - 1],
        cleanup_interval: Option<Duration>,
    ) -> Self {
        Self {
            metrics,
            fair_queue_popped_task_count,
            task_slice_time_origin: StdInstant::now(),
            last_task_slice_pop_time_nanos: AtomicU64::new(0),
            level_manager: multilevel::TaskLevelManager::new(
                level_time_threshold,
                cleanup_interval,
            ),
            levels: array::from_fn(|_| SegQueue::new()),
            len: AtomicUsize::new(0),
        }
    }
}

impl GroupInnerQueue<KeyspaceFairTaskCell> for MultiLevelInnerQueue {
    fn push(&self, mut task_cell: KeyspaceFairTaskCell) {
        self.level_manager.adjust_task_level(&mut task_cell);
        let level = task_cell.mut_extras().current_level() as usize;
        // Count before publishing into the level queue. If the task were
        // published first, a racing pop could take it and decrement len before
        // this push increments it.
        let queue_len = self.len.fetch_add(1, Ordering::AcqRel);
        if let Some(metrics) = self.metrics.as_ref() {
            metrics.observe_fair_queue_length(queue_len);
        }
        self.levels[level].push(task_cell);
    }

    fn pop(&self) -> Option<KeyspaceFairTaskCell> {
        let expected_level = self.next_level();
        for level in expected_level..expected_level + multilevel::LEVEL_NUM {
            if let Some(task_cell) = self.levels[level % multilevel::LEVEL_NUM].pop() {
                self.len.fetch_sub(1, Ordering::AcqRel);
                let mut task_cell = task_cell;
                let popped_before_current = self
                    .fair_queue_popped_task_count
                    .fetch_add(1, Ordering::Relaxed);
                task_cell.set_current_pop_waits(
                    popped_before_current
                        .saturating_sub(task_cell.fair_queue_popped_task_count_at_push),
                );
                if let Some(metrics) = self.metrics.as_ref() {
                    let enqueue_time = task_cell.mut_extras().schedule_time();
                    let schedule_interval =
                        self.task_slice_schedule_interval(enqueue_time, StdInstant::now());

                    metrics.observe_fair_queue_popped_task_slice(&mut task_cell);
                    if let Some(schedule_interval) = schedule_interval {
                        metrics.observe_fair_queue_task_slice_schedule_interval(schedule_interval);
                    }
                }
                return Some(task_cell);
            }
        }
        None
    }

    fn len(&self) -> usize {
        self.len.load(Ordering::Acquire)
    }
}

impl MultiLevelInnerQueue {
    const CHANCE_RATIO: u32 = 4;
    const LEVEL0_CHANCE: f64 = 0.8;

    fn next_level(&self) -> usize {
        let mut rng = thread_rng();
        if rng.gen::<f64>() < Self::LEVEL0_CHANCE {
            0
        } else {
            (1..multilevel::LEVEL_NUM - 1)
                .find(|_| rng.gen_ratio(Self::CHANCE_RATIO, Self::CHANCE_RATIO + 1))
                .unwrap_or(multilevel::LEVEL_NUM - 1)
        }
    }

    #[inline]
    fn task_slice_schedule_interval(
        &self,
        enqueue_time: Option<StdInstant>,
        pop_time: StdInstant,
    ) -> Option<Duration> {
        let encode_time = |time: StdInstant| {
            time.saturating_duration_since(self.task_slice_time_origin)
                .as_nanos()
                .min((u64::MAX - 1) as u128) as u64
                + 1
        };
        let pop_time_nanos = encode_time(pop_time);
        let previous_pop_time_nanos = self
            .last_task_slice_pop_time_nanos
            .fetch_max(pop_time_nanos, Ordering::Relaxed);
        let enqueue_time_nanos = enqueue_time.map(encode_time);
        let interval_start_nanos = enqueue_time_nanos.unwrap_or(0).max(previous_pop_time_nanos);
        if interval_start_nanos == 0 {
            None
        } else {
            Some(Duration::from_nanos(
                pop_time_nanos.saturating_sub(interval_start_nanos),
            ))
        }
    }
}

pub struct KeyspaceMetrics {
    /// Gauge tracking the current number of keyspace groups in the pool.
    fair_queue_keyspace_count_gauge: IntGauge,
    // keyspace metrics handles
    handles: RwLock<HashMap<u32, Arc<KeyspaceInnerMetricsHandle>>>,
    // per-slot locks serialize metric create/unregister for mapped keyspaces.
    slots: [Mutex<()>; KeyspaceMetrics::KEYSPACE_METRICS_SLOT_COUNT],
    // create KeyspaceInnerMetrics function
    create_keyspace_metrics: Box<dyn Fn(u32) -> KeyspaceInnerMetrics + Send + Sync>,
    // unregister KeyspaceInnerMetrics function
    unregister_keyspace_metrics: Box<dyn Fn(u32) + Send + Sync>,
    /// Last second when inactive inner metrics cleanup was attempted.
    last_inactive_metrics_cleanup_secs: AtomicU64,
}

impl KeyspaceMetrics {
    const KEYSPACE_METRICS_SLOT_COUNT: usize = 4096;
    const MAX_KEYSPACE_METRICS_INACTIVE_SECS: u64 = 600;
    const KEYSPACE_METRICS_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);

    pub fn new(
        fair_queue_keyspace_count_gauge: IntGauge,
        create_keyspace_metrics: impl Fn(u32) -> KeyspaceInnerMetrics + Send + Sync + 'static,
        unregister_keyspace_metrics: impl Fn(u32) + Send + Sync + 'static,
    ) -> Arc<Self> {
        fair_queue_keyspace_count_gauge.set(0);
        Arc::new(Self {
            fair_queue_keyspace_count_gauge,
            handles: RwLock::new(HashMap::new()),
            slots: array::from_fn(|_| Mutex::new(())),
            create_keyspace_metrics: Box::new(create_keyspace_metrics),
            unregister_keyspace_metrics: Box::new(unregister_keyspace_metrics),
            last_inactive_metrics_cleanup_secs: AtomicU64::new(now_coarse_secs()),
        })
    }

    #[inline]
    pub fn set_fair_queue_keyspace_count(&self, cnt: usize) {
        self.fair_queue_keyspace_count_gauge.set(cnt as i64);
    }

    #[inline]
    pub fn get_keyspace_metrics(
        self: &Arc<Self>,
        keyspace_id: u32,
    ) -> Arc<KeyspaceInnerMetricsHandle> {
        if let Some(metrics) = self.handles.read().get(&keyspace_id) {
            metrics.renew_active();
            return metrics.clone();
        }

        let _slot_guard = self.lock_metrics_slot(keyspace_id);
        if let Some(metrics) = self.handles.read().get(&keyspace_id) {
            metrics.renew_active();
            return metrics.clone();
        }

        let metrics = Arc::new(KeyspaceInnerMetricsHandle::new(self.clone(), keyspace_id));
        metrics.renew_active();
        let mut guard = self.handles.write();
        guard.insert(keyspace_id, metrics.clone());
        metrics
    }

    pub fn clear_inactive_keyspace_metrics_if_needed(&self) -> usize {
        let now_secs = now_coarse_secs();
        let last_cleanup_secs = self
            .last_inactive_metrics_cleanup_secs
            .load(Ordering::Relaxed);
        if now_secs.saturating_sub(last_cleanup_secs)
            < Self::KEYSPACE_METRICS_CLEANUP_INTERVAL.as_secs()
        {
            return 0;
        }
        if self
            .last_inactive_metrics_cleanup_secs
            .compare_exchange(
                last_cleanup_secs,
                now_secs,
                Ordering::Relaxed,
                Ordering::Relaxed,
            )
            .is_err()
        {
            return 0;
        }
        self.clear_inactive_keyspace_metrics(now_secs, Self::MAX_KEYSPACE_METRICS_INACTIVE_SECS)
    }

    fn clear_inactive_keyspace_metrics(&self, now_secs: u64, inactive_secs: u64) -> usize {
        let inactive_keyspaces = {
            let metrics = self.handles.read();
            metrics
                .iter()
                .filter_map(|(keyspace_id, metrics)| {
                    (Arc::strong_count(metrics) == 1
                        && metrics.secs_since_last_active(now_secs) > inactive_secs)
                        .then_some(*keyspace_id)
                })
                .collect::<Vec<_>>()
        };
        if inactive_keyspaces.is_empty() {
            return 0;
        }

        let mut removed_keyspace_count = 0;
        for keyspace_id in inactive_keyspaces {
            let _slot_guard = self.lock_metrics_slot(keyspace_id);
            let removed = {
                let mut metrics = self.handles.write();
                let should_remove = match metrics.get(&keyspace_id) {
                    Some(metrics) => {
                        Arc::strong_count(metrics) == 1
                            && metrics.secs_since_last_active(now_secs) > inactive_secs
                    }
                    None => false,
                };
                if should_remove {
                    metrics.remove(&keyspace_id);
                    true
                } else {
                    false
                }
            };

            if removed {
                (self.unregister_keyspace_metrics)(keyspace_id);
                removed_keyspace_count += 1;
            }
        }
        removed_keyspace_count
    }

    #[inline]
    fn lock_metrics_slot(&self, keyspace_id: u32) -> parking_lot::MutexGuard<'_, ()> {
        self.slots[Self::keyspace_metrics_slot(keyspace_id)].lock()
    }

    #[inline]
    fn keyspace_metrics_slot(keyspace_id: u32) -> usize {
        keyspace_id as usize % Self::KEYSPACE_METRICS_SLOT_COUNT
    }
}

pub struct KeyspaceInnerMetrics {
    /// The queue-length histogram.
    fair_queue_keyspace_queue_len: Histogram,
    /// The task-slice wait count histogram.
    fair_queue_keyspace_task_slice_waits: Histogram,
    /// The task-slice wait duration histogram.
    fair_queue_keyspace_task_slice_wait_duration: Histogram,
    /// The scheduling interval since the previous pop or current enqueue.
    fair_queue_keyspace_task_slice_schedule_interval: Histogram,
    /// The CPU time of one task slice.
    keyspace_task_slice_cpu_duration: Histogram,
    /// The task wall-time counter.
    keyspace_task_wall_time_counter: Counter,
}

impl KeyspaceInnerMetrics {
    pub fn new(
        fair_queue_keyspace_queue_len: Histogram,
        fair_queue_keyspace_task_slice_waits: Histogram,
        fair_queue_keyspace_task_slice_wait_duration: Histogram,
        fair_queue_keyspace_task_slice_schedule_interval: Histogram,
        keyspace_task_slice_cpu_duration: Histogram,
        keyspace_task_wall_time_counter: Counter,
    ) -> Self {
        Self {
            fair_queue_keyspace_queue_len,
            fair_queue_keyspace_task_slice_waits,
            fair_queue_keyspace_task_slice_wait_duration,
            fair_queue_keyspace_task_slice_schedule_interval,
            keyspace_task_slice_cpu_duration,
            keyspace_task_wall_time_counter,
        }
    }
}

/// Metrics handle for one concrete keyspace.
pub struct KeyspaceInnerMetricsHandle {
    last_active_seconds: AtomicU64,
    metrics: KeyspaceInnerMetrics,
}

impl KeyspaceInnerMetricsHandle {
    fn new(parent: Arc<KeyspaceMetrics>, keyspace_id: u32) -> Self {
        let metrics = (parent.create_keyspace_metrics)(keyspace_id);
        Self {
            last_active_seconds: AtomicU64::new(now_coarse_secs()),
            metrics,
        }
    }

    #[inline]
    pub fn observe_fair_queue_length(&self, queue_len: usize) {
        self.renew_active();
        self.metrics
            .fair_queue_keyspace_queue_len
            .observe(queue_len as f64);
    }

    #[inline]
    pub fn observe_fair_queue_task_slice_waits(&self, task_slice_waits: u64) {
        self.renew_active();
        self.metrics
            .fair_queue_keyspace_task_slice_waits
            .observe(task_slice_waits as f64);
    }

    #[inline]
    pub fn observe_fair_queue_task_slice_wait_duration(&self, task_cell: &mut impl YatpTaskCell) {
        self.renew_active();
        if let Some(schedule_time) = task_cell.mut_extras().schedule_time() {
            self.metrics
                .fair_queue_keyspace_task_slice_wait_duration
                .observe(schedule_time.elapsed().as_secs_f64());
        }
    }

    #[inline]
    fn observe_fair_queue_popped_task_slice(&self, task_cell: &mut KeyspaceFairTaskCell) {
        self.observe_fair_queue_task_slice_wait_duration(task_cell);
        task_cell
            .current_pop_waits
            .inspect(|&waits| self.observe_fair_queue_task_slice_waits(waits));
    }

    #[inline]
    pub fn observe_fair_queue_task_slice_schedule_interval(&self, interval: Duration) {
        self.renew_active();
        self.metrics
            .fair_queue_keyspace_task_slice_schedule_interval
            .observe(interval.as_secs_f64());
    }

    #[inline]
    pub fn observe_task_slice_cpu_duration(&self, cpu_duration: Duration) {
        self.renew_active();
        self.metrics
            .keyspace_task_slice_cpu_duration
            .observe(cpu_duration.as_secs_f64());
    }

    #[inline]
    pub fn observe_task_wall_time(&self, wall_time: Duration) {
        self.renew_active();
        self.metrics
            .keyspace_task_wall_time_counter
            .inc_by(wall_time.as_secs_f64());
    }

    #[inline]
    fn renew_active(&self) {
        self.last_active_seconds
            .store(now_coarse_secs(), Ordering::Relaxed);
    }

    #[inline]
    fn secs_since_last_active(&self, now_secs: u64) -> u64 {
        now_secs.saturating_sub(self.last_active_seconds.load(Ordering::Relaxed))
    }
}

fn now_coarse_secs() -> u64 {
    #[cfg(test)]
    if let Some(now_secs) = MOCK_NOW_COARSE_SECS.get() {
        return now_secs;
    }
    Instant::now_coarse().second().max(0) as u64
}

/// YATP custom queue used by the unified read pool when keyspace fairness is
/// enabled.
///
/// The read-pool submit path stores a `u32` keyspace id in each task's YATP
/// metadata with [`set_task_meta`]. `push` decodes that id and routes
/// the task into a per-keyspace [`MultiLevelInnerQueue`]. Tasks without valid
/// keyspace metadata are assigned to keyspace `0`.
///
/// `pop` delegates fair keyspace selection to the wrapped
/// [`MultiGroupFairQueue`] and then pops one task from the selected keyspace's
/// inner multilevel queue. This type is only the YATP [`YatpTaskQueue`]
/// adapter; the generic fairness and bounded consistency rules are documented
/// on [`MultiGroupFairQueue`].
///
/// Virtual time is accounted in microseconds. This queue does not charge
/// virtual time by itself. Instead,
/// [`YatpKeyspaceFairQueue::get_fairness_recorder`] returns a
/// [`FairnessRecorder`] to the read pool, and `ObserveReadPoolTaskFuture`
/// charges configured precharge plus measured CPU time or wall time in
/// microseconds.
/// `KEYSPACE_FAIR_MAX_VT_LAG` is therefore a microsecond lag window.
///
/// The queue owns the lifecycle of per-keyspace queues. Periodic maintenance
/// refreshes group-count metrics and removes empty keyspaces that have been
/// inactive for `MAX_KEYSPACE_INACTIVE_SECS`, with inactive cleanup
/// additionally throttled by `KEYSPACE_FAIR_INACTIVE_CLEANUP_INTERVAL`.
pub struct YatpKeyspaceFairQueue {
    /// Per-keyspace fair queue. Virtual time is measured in microseconds.
    inner: Arc<MultiGroupFairQueue<u32, KeyspaceFairTaskCell, MultiLevelInnerQueue>>,
    /// Metrics used by keyspace-level read-pool observability.
    keyspace_metrics: Option<Arc<KeyspaceMetrics>>,
    /// Number of task slices popped from the keyspace fair queue for
    /// scheduling.
    fair_queue_popped_task_count: Arc<AtomicU64>,
    /// Last second when periodic queue maintenance was attempted.
    last_maintenance_secs: AtomicU64,
    admission_controller: Option<Arc<dyn FairQueueAdmissionController>>,
    /// Last second when inactive-keyspace cleanup was attempted.
    last_inactive_cleanup_secs: AtomicU64,
}

impl YatpKeyspaceFairQueue {
    const MAX_KEYSPACE_INACTIVE_SECS: u64 = 120;
    const KEYSPACE_FAIR_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(10);
    const KEYSPACE_FAIR_INACTIVE_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
    const KEYSPACE_FAIR_MAX_VT_LAG: u64 = 1_000_000;
    const KEYSPACE_GROUP_COUNT_WARN_THRESHOLD: usize = 1000;

    fn new(
        keyspace_metrics: Option<Arc<KeyspaceMetrics>>,
        admission_controller: Option<Arc<dyn FairQueueAdmissionController>>,
    ) -> Self {
        let inner_keyspace_metrics = keyspace_metrics.clone();
        let fair_queue_popped_task_count = Arc::new(AtomicU64::new(0));
        let inner_fair_queue_popped_task_count = fair_queue_popped_task_count.clone();
        let now_secs = now_coarse_secs();
        Self {
            inner: Arc::new(MultiGroupFairQueue::new(
                Self::KEYSPACE_FAIR_MAX_VT_LAG,
                move |keyspace_id| {
                    let cfg = multilevel::Config::default();
                    let metrics = inner_keyspace_metrics
                        .as_ref()
                        .map(|metrics| metrics.get_keyspace_metrics(keyspace_id));
                    MultiLevelInnerQueue::new_with_config(
                        metrics,
                        inner_fair_queue_popped_task_count.clone(),
                        cfg.get_level_time_threshold(),
                        cfg.get_cleanup_interval(),
                    )
                },
            )),
            keyspace_metrics,
            fair_queue_popped_task_count,
            last_maintenance_secs: AtomicU64::new(now_secs),
            last_inactive_cleanup_secs: AtomicU64::new(now_secs),
            admission_controller,
        }
    }

    pub fn get_fairness_recorder(&self) -> Arc<dyn FairnessRecorder<u32>> {
        self.inner.clone()
    }

    /// Runs periodic queue maintenance when the maintenance interval has
    /// elapsed.
    pub fn maintain_queue_if_needed(&self) {
        let now_secs = now_coarse_secs();
        let last_maintenance_secs = self.last_maintenance_secs.load(Ordering::Relaxed);
        if now_secs.saturating_sub(last_maintenance_secs)
            < Self::KEYSPACE_FAIR_MAINTENANCE_INTERVAL.as_secs()
        {
            return;
        }
        if self
            .last_maintenance_secs
            .compare_exchange(
                last_maintenance_secs,
                now_secs,
                Ordering::Relaxed,
                Ordering::Relaxed,
            )
            .is_err()
        {
            return;
        }

        self.remove_inactive_groups_if_needed(now_secs);
        if let Some(keyspace_metrics) = self.keyspace_metrics.as_ref() {
            keyspace_metrics.set_fair_queue_keyspace_count(self.inner.group_count());
        }
    }

    fn remove_inactive_groups_if_needed(&self, now_secs: u64) {
        let last_inactive_cleanup_secs = self.last_inactive_cleanup_secs.load(Ordering::Relaxed);
        if now_secs.saturating_sub(last_inactive_cleanup_secs)
            < Self::KEYSPACE_FAIR_INACTIVE_CLEANUP_INTERVAL.as_secs()
        {
            return;
        }
        if self
            .last_inactive_cleanup_secs
            .compare_exchange(
                last_inactive_cleanup_secs,
                now_secs,
                Ordering::Relaxed,
                Ordering::Relaxed,
            )
            .is_err()
        {
            return;
        }
        let group_count_before_cleanup = self.inner.group_count();
        self.inner
            .remove_inactive_groups(Self::MAX_KEYSPACE_INACTIVE_SECS);

        if group_count_before_cleanup > Self::KEYSPACE_GROUP_COUNT_WARN_THRESHOLD {
            let group_count_after_cleanup = self.inner.group_count();
            warn!(
                "too many keyspace groups in keyspace fair read pool";
                "group_count" => group_count_before_cleanup,
                "removed_group_count" => group_count_before_cleanup
                    .saturating_sub(group_count_after_cleanup),
                "remaining_group_count" => group_count_after_cleanup,
                "threshold" => Self::KEYSPACE_GROUP_COUNT_WARN_THRESHOLD,
            );
        }
    }

    fn admit_popped_task(&self, task_cell: &mut TaskCell) {
        clear_stale_fair_cpu_permit();
        let Some(controller) = self.admission_controller.as_ref() else {
            return;
        };
        let metadata = task_metadata(task_cell);
        let Some((task_type, deadline)) = fair_cpu_metadata(task_cell) else {
            return;
        };
        match controller.admit(metadata.keyspace_id, task_type, deadline) {
            FairAdmitOutcome::Admitted(permit) => install_fair_cpu_permit(permit),
            FairAdmitOutcome::Bypassed => {}
        }
    }
}

impl YatpTaskQueue<TaskCell> for YatpKeyspaceFairQueue {
    fn push(&self, mut task_cell: TaskCell) {
        let keyspace_id = task_keyspace_id(&mut task_cell);
        let task_cell = KeyspaceFairTaskCell::new(
            task_cell,
            self.fair_queue_popped_task_count.load(Ordering::Relaxed),
        );
        self.inner.push(keyspace_id, task_cell);
    }

    fn pop(&self) -> PopResult<TaskCell> {
        if self.inner.is_empty() {
            return PopResult::Empty;
        }
        if let Some(controller) = self.admission_controller.as_ref() {
            if let FairQueueAdmissionResult::Pending { retry_at } = controller.precheck() {
                return PopResult::Pending { retry_at };
            }
        }
        match self.inner.pop() {
            Some(task) => {
                YatpDispatchContext::prepare(YatpDispatchPrepareContext {
                    fair_queue_waited_task_slices: task.current_pop_waits,
                });
                let mut task_cell = task.into_task_cell();
                self.admit_popped_task(&mut task_cell);
                let schedule_time = task_cell
                    .mut_extras()
                    .schedule_time()
                    .unwrap_or_else(std::time::Instant::now);
                PopResult::Ready(queue::Pop {
                    task_cell,
                    schedule_time,
                    from_local: false,
                })
            }
            None => PopResult::Empty,
        }
    }

    fn drain(&self) {
        self.inner.drain();
    }

    fn has_ready_task(&self) -> bool {
        // This is intentionally only a queue-length hint. Under fair global CPU
        // gating, the head task may still be blocked by the global bucket and
        // `pop` may return Pending. Making this admission-aware would require
        // running the precheck here, which would make this hot readiness hint
        // heavier and couple it to token state. The extra preemption/reschedule
        // churn is accepted for now.
        !self.inner.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{
        Arc, Barrier, Mutex,
        atomic::{AtomicBool, AtomicUsize, Ordering},
        mpsc,
    };

    use futures::compat::Future01CompatExt;
    use prometheus::HistogramVec;

    use super::*;
    use crate::{fairness::set_now_coarse_seconds_for_test, timer::GLOBAL_TIMER_HANDLE};

    struct MockNowGuard;

    impl Drop for MockNowGuard {
        fn drop(&mut self) {
            set_now_coarse_secs_for_test(None);
            set_now_coarse_seconds_for_test(None);
        }
    }

    fn mock_now_coarse_secs(now_secs: u64) -> MockNowGuard {
        set_mock_now_coarse_secs(now_secs);
        MockNowGuard
    }

    fn set_mock_now_coarse_secs(now_secs: u64) {
        set_now_coarse_secs_for_test(Some(now_secs));
        set_now_coarse_seconds_for_test(Some(now_secs));
    }

    fn new_test_task(task_id: u64, keyspace_id: Option<u32>) -> TaskCell {
        new_test_task_with_fixed_level(task_id, keyspace_id, None)
    }

    fn new_test_task_with_fixed_level(
        task_id: u64,
        keyspace_id: Option<u32>,
        fixed_level: Option<u8>,
    ) -> TaskCell {
        let mut extras = Extras::new_multilevel(task_id, fixed_level);
        if let Some(keyspace_id) = keyspace_id {
            set_task_meta(
                &mut extras,
                &TaskMetadata {
                    keyspace_id,
                    fair_cpu_admission: None,
                },
            );
        }
        TaskCell::new(std::future::pending(), extras)
    }

    fn new_test_keyspace_fair_task(task_id: u64, keyspace_id: Option<u32>) -> KeyspaceFairTaskCell {
        KeyspaceFairTaskCell::new(new_test_task(task_id, keyspace_id), 0)
    }

    fn new_test_keyspace_fair_task_with_fixed_level(
        task_id: u64,
        keyspace_id: Option<u32>,
        fixed_level: Option<u8>,
    ) -> KeyspaceFairTaskCell {
        KeyspaceFairTaskCell::new(
            new_test_task_with_fixed_level(task_id, keyspace_id, fixed_level),
            0,
        )
    }

    fn new_test_task_with_fair_cpu_admission(
        task_id: u64,
        keyspace_id: u32,
        task_type: u8,
    ) -> TaskCell {
        let mut extras = Extras::new_multilevel(task_id, None);
        set_task_meta(
            &mut extras,
            &TaskMetadata {
                keyspace_id,
                fair_cpu_admission: Some(FairCpuAdmissionMetadata {
                    task_type,
                    deadline: Some(StdInstant::now() + Duration::from_secs(5)),
                }),
            },
        );
        TaskCell::new(std::future::pending(), extras)
    }

    fn pop_task_id(queue: &YatpKeyspaceFairQueue) -> u64 {
        match queue.pop() {
            PopResult::Ready(mut task) => task.task_cell.mut_extras().task_id(),
            PopResult::Pending { .. } => panic!("keyspace fair queue should not delay tasks"),
            PopResult::Empty => panic!("keyspace fair queue should have a ready task"),
        }
    }

    fn pop_task_id_and_keyspace(queue: &YatpKeyspaceFairQueue) -> (u64, u32) {
        match queue.pop() {
            PopResult::Ready(mut task) => {
                let task_id = task.task_cell.mut_extras().task_id();
                let keyspace_id = task_keyspace_id(&mut task.task_cell);
                (task_id, keyspace_id)
            }
            PopResult::Pending { .. } => panic!("keyspace fair queue should not delay tasks"),
            PopResult::Empty => panic!("keyspace fair queue should have a ready task"),
        }
    }

    fn pop_fair_queue_waited_task_slices(queue: &YatpKeyspaceFairQueue) -> Option<u64> {
        match queue.pop() {
            PopResult::Ready(mut task) => {
                let _ctx_guard = YatpDispatchContext::install(&mut task.task_cell);
                YatpDispatchContext::with_tls(|ctx| ctx.unwrap().fair_queue_waited_task_slices())
            }
            PopResult::Pending { .. } => panic!("keyspace fair queue should not delay tasks"),
            PopResult::Empty => panic!("keyspace fair queue should have a ready task"),
        }
    }

    fn pop_inner_queue_task_id(queue: &MultiLevelInnerQueue) -> u64 {
        queue.pop().unwrap().task_cell.mut_extras().task_id()
    }

    fn no_keyspace_inner_metrics() -> Option<Arc<KeyspaceInnerMetricsHandle>> {
        None
    }

    fn new_test_fair_queue_popped_task_count() -> Arc<AtomicU64> {
        Arc::new(AtomicU64::new(0))
    }

    fn new_test_inner_queue() -> MultiLevelInnerQueue {
        let cfg = multilevel::Config::default();
        MultiLevelInnerQueue::new_with_config(
            no_keyspace_inner_metrics(),
            new_test_fair_queue_popped_task_count(),
            cfg.get_level_time_threshold(),
            cfg.get_cleanup_interval(),
        )
    }

    #[test]
    fn test_task_slice_schedule_interval() {
        let queue = new_test_inner_queue();
        let origin = queue.task_slice_time_origin;

        assert_eq!(
            queue.task_slice_schedule_interval(
                Some(origin + Duration::from_millis(10)),
                origin + Duration::from_millis(20),
            ),
            Some(Duration::from_millis(10))
        );
        assert_eq!(
            queue.task_slice_schedule_interval(
                Some(origin + Duration::from_millis(25)),
                origin + Duration::from_millis(30),
            ),
            Some(Duration::from_millis(5))
        );
        assert_eq!(
            queue.task_slice_schedule_interval(
                Some(origin + Duration::from_millis(15)),
                origin + Duration::from_millis(40),
            ),
            Some(Duration::from_millis(10))
        );
    }

    #[allow(clippy::too_many_arguments)]
    fn new_test_keyspace_metrics(
        pool_name: &'static str,
        group_count: IntGauge,
        queue_length: HistogramVec,
        queue_task_slice_waits: HistogramVec,
        queue_task_slice_wait_duration: HistogramVec,
        task_slice_schedule_interval: HistogramVec,
        task_slice_cpu_duration: HistogramVec,
        task_wall_time: prometheus::CounterVec,
        queue_length_resolve_count: Option<Arc<AtomicUsize>>,
        metrics_remove_count: Option<Arc<AtomicUsize>>,
    ) -> Arc<KeyspaceMetrics> {
        let remove_queue_length = queue_length.clone();
        let remove_queue_task_slice_waits = queue_task_slice_waits.clone();
        let remove_queue_task_slice_wait_duration = queue_task_slice_wait_duration.clone();
        let remove_task_slice_schedule_interval = task_slice_schedule_interval.clone();
        let remove_task_slice_cpu_duration = task_slice_cpu_duration.clone();
        let remove_task_wall_time = task_wall_time.clone();
        KeyspaceMetrics::new(
            group_count,
            move |keyspace_id| {
                if let Some(count) = &queue_length_resolve_count {
                    count.fetch_add(1, Ordering::Relaxed);
                }
                let keyspace_id = keyspace_id.to_string();
                KeyspaceInnerMetrics::new(
                    queue_length.with_label_values(&[pool_name, &keyspace_id]),
                    queue_task_slice_waits.with_label_values(&[pool_name, &keyspace_id]),
                    queue_task_slice_wait_duration.with_label_values(&[pool_name, &keyspace_id]),
                    task_slice_schedule_interval.with_label_values(&[pool_name, &keyspace_id]),
                    task_slice_cpu_duration.with_label_values(&[pool_name, &keyspace_id]),
                    task_wall_time.with_label_values(&[pool_name, &keyspace_id]),
                )
            },
            move |keyspace_id| {
                if let Some(count) = &metrics_remove_count {
                    count.fetch_add(1, Ordering::Relaxed);
                }
                let keyspace_id = keyspace_id.to_string();
                let labels = [pool_name, keyspace_id.as_str()];
                let _ = remove_queue_length.remove_label_values(&labels);
                let _ = remove_queue_task_slice_waits.remove_label_values(&labels);
                let _ = remove_queue_task_slice_wait_duration.remove_label_values(&labels);
                let _ = remove_task_slice_schedule_interval.remove_label_values(&labels);
                let _ = remove_task_slice_cpu_duration.remove_label_values(&labels);
                let _ = remove_task_wall_time.remove_label_values(&labels);
            },
        )
    }

    fn assert_inner_queue_task_level(queue: &MultiLevelInnerQueue, expected_level: usize) {
        for level in 0..multilevel::LEVEL_NUM {
            assert_eq!(
                queue.levels[level].is_empty(),
                level != expected_level,
                "level {} emptiness mismatch, expected task in level {}",
                level,
                expected_level
            );
        }
    }

    struct TestMultiLevelTaskQueue {
        inner: Arc<MultiLevelInnerQueue>,
    }

    #[derive(Default)]
    struct TestAdmissionController {
        pending: AtomicBool,
        prechecks: AtomicUsize,
        admits: Mutex<Vec<(u32, u8)>>,
    }

    impl TestAdmissionController {
        fn new() -> Self {
            Self::default()
        }

        fn set_pending(&self, pending: bool) {
            self.pending.store(pending, Ordering::SeqCst);
        }

        fn precheck_count(&self) -> usize {
            self.prechecks.load(Ordering::SeqCst)
        }

        fn admits(&self) -> Vec<(u32, u8)> {
            self.admits.lock().unwrap().clone()
        }
    }

    impl FairQueueAdmissionController for TestAdmissionController {
        fn precheck(&self) -> FairQueueAdmissionResult {
            self.prechecks.fetch_add(1, Ordering::SeqCst);
            if self.pending.load(Ordering::SeqCst) {
                return FairQueueAdmissionResult::Pending {
                    retry_at: std::time::Instant::now() + Duration::from_secs(1),
                };
            }

            FairQueueAdmissionResult::Ready
        }

        fn admit(
            &self,
            keyspace_id: u32,
            task_type: u8,
            _deadline: Option<StdInstant>,
        ) -> FairAdmitOutcome {
            self.admits.lock().unwrap().push((keyspace_id, task_type));
            FairAdmitOutcome::Bypassed
        }
    }

    struct TestFairCpuPermit {
        consumed_cpu_us: Arc<AtomicU64>,
        action: FairCpuBudgetAction,
    }

    impl FairCpuPermit for TestFairCpuPermit {
        fn consume(&mut self, cpu_delta_us: u64) -> FairCpuBudgetAction {
            self.consumed_cpu_us
                .fetch_add(cpu_delta_us, Ordering::SeqCst);
            self.action
        }
    }

    struct TestChargingAdmissionController {
        first_dispatch_cpu_us: Arc<AtomicU64>,
        admission_count: AtomicUsize,
        first_precheck_barrier: Arc<Barrier>,
        block_first_precheck: AtomicBool,
    }

    impl FairQueueAdmissionController for TestChargingAdmissionController {
        fn precheck(&self) -> FairQueueAdmissionResult {
            if self.block_first_precheck.swap(false, Ordering::SeqCst) {
                self.first_precheck_barrier.wait();
            }
            FairQueueAdmissionResult::Ready
        }

        fn admit(
            &self,
            _keyspace_id: u32,
            _task_type: u8,
            _deadline: Option<StdInstant>,
        ) -> FairAdmitOutcome {
            let consumed_cpu_us = if self.admission_count.fetch_add(1, Ordering::SeqCst) == 0 {
                self.first_dispatch_cpu_us.clone()
            } else {
                Arc::new(AtomicU64::new(0))
            };
            FairAdmitOutcome::Admitted(Box::new(TestFairCpuPermit {
                consumed_cpu_us,
                action: FairCpuBudgetAction::Continue,
            }))
        }
    }

    fn burn_thread_cpu(duration: Duration) {
        let timer = ThreadTime::now();
        while timer.elapsed() < duration {
            std::hint::black_box(());
        }
    }

    #[test]
    fn test_fair_cpu_checkpoint_is_rate_limited() {
        let mut dispatch = FairCpuDispatch::new();
        let dispatch_started_at = dispatch.last_checkpoint_at;

        burn_thread_cpu(Duration::from_micros(100));
        assert_eq!(
            dispatch.take_cpu_delta_us_if_due(
                dispatch_started_at + FAIR_CPU_CHECKPOINT_MIN_INTERVAL - Duration::from_nanos(1)
            ),
            None
        );

        let cpu_delta_us = dispatch
            .take_cpu_delta_us_if_due(dispatch_started_at + FAIR_CPU_CHECKPOINT_MIN_INTERVAL)
            .expect("checkpoint should be due");
        assert!(cpu_delta_us > 0);

        burn_thread_cpu(Duration::from_micros(100));
        assert_eq!(
            dispatch
                .take_cpu_delta_us_if_due(dispatch_started_at + FAIR_CPU_CHECKPOINT_MIN_INTERVAL),
            None
        );
    }

    #[test]
    fn test_fair_cpu_guard_charges_short_dispatch_tail() {
        clear_stale_fair_cpu_permit();
        assert!(!is_fair_cpu_dispatch_active());
        let consumed_cpu_us = Arc::new(AtomicU64::new(0));
        install_fair_cpu_permit(Box::new(TestFairCpuPermit {
            consumed_cpu_us: consumed_cpu_us.clone(),
            action: FairCpuBudgetAction::Continue,
        }));
        let guard = FairCpuGuard::new();
        assert!(is_fair_cpu_dispatch_active());

        burn_thread_cpu(Duration::from_micros(100));
        drop(guard);
        assert!(consumed_cpu_us.load(Ordering::SeqCst) > 0);
        assert!(!has_fair_cpu_permit());
        assert!(!is_fair_cpu_dispatch_active());
        assert!(FAIR_CPU_DISPATCH.with(|slot| slot.borrow().is_none()));
    }

    #[test]
    fn test_wall_reschedule_reconciles_dispatch_cpu() {
        let first_dispatch_cpu_us = Arc::new(AtomicU64::new(0));
        let first_precheck_barrier = Arc::new(Barrier::new(2));
        let controller = Arc::new(TestChargingAdmissionController {
            first_dispatch_cpu_us: first_dispatch_cpu_us.clone(),
            admission_count: AtomicUsize::new(0),
            first_precheck_barrier: first_precheck_barrier.clone(),
            block_first_precheck: AtomicBool::new(true),
        });
        let mut builder = YatpPoolBuilder::new(DefaultTicker);
        builder
            .name_prefix("test-wall-reschedule-tail-charge")
            .thread_count(1, 1, 1);
        let (pool, _queue) = builder.build_keyspace_fair_pool(Some(controller));
        let (tx, rx) = mpsc::channel();

        let resumed_after_reschedule = Arc::new(AtomicBool::new(false));
        let target_resumed = resumed_after_reschedule.clone();
        let inner = async move {
            burn_thread_cpu(Duration::from_millis(1));
            let mut now = crate::time::LazyInstant::default();
            assert!(
                crate::task_reschedule::reschedule_if_time_slice_exhausted(&mut now).await,
                "the synthetic poll start should request wall rescheduling"
            );
            target_resumed.store(true, Ordering::SeqCst);
            tx.send(()).unwrap();
        };
        let mut inner = Box::pin(inner);
        let mut poll_count = 0;
        let future = std::future::poll_fn(move |cx| {
            let poll_started_at = if poll_count == 0 {
                StdInstant::now()
                    .checked_sub(Duration::from_millis(4))
                    .unwrap()
            } else {
                StdInstant::now()
            };
            poll_count += 1;
            let _context = crate::task_reschedule::PollRescheduleContextGuard::new(
                poll_started_at,
                Duration::from_micros(crate::task_reschedule::DEFAULT_MAX_TIME_SLICE_US),
            );
            std::future::Future::poll(inner.as_mut(), cx)
        });
        let mut extras = Extras::new_multilevel(1, None);
        set_task_meta(
            &mut extras,
            &TaskMetadata {
                keyspace_id: 42,
                fair_cpu_admission: Some(FairCpuAdmissionMetadata {
                    task_type: 7,
                    deadline: None,
                }),
            },
        );
        pool.spawn(TaskCell::new(future, extras));

        // Keep another task ready so `reschedule` really returns the target to
        // the scheduler instead of being repolled inside the same dispatch.
        let (competitor_tx, competitor_rx) = mpsc::channel();
        let mut competitor_extras = Extras::new_multilevel(2, None);
        set_task_meta(
            &mut competitor_extras,
            &TaskMetadata {
                keyspace_id: 42,
                fair_cpu_admission: None,
            },
        );
        pool.spawn(TaskCell::new(
            async move {
                competitor_tx
                    .send(resumed_after_reschedule.load(Ordering::SeqCst))
                    .unwrap();
            },
            competitor_extras,
        ));
        first_precheck_barrier.wait();

        assert!(
            !competitor_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
            "the ready competitor must run before the target resumes"
        );
        rx.recv_timeout(Duration::from_secs(5)).unwrap();
        drop(pool);
        assert!(
            first_dispatch_cpu_us.load(Ordering::SeqCst) > 0,
            "the dispatch ended by wall rescheduling must reconcile its CPU"
        );
    }

    #[test]
    fn test_fair_cpu_checkpoint_propagates_yield_action() {
        clear_stale_fair_cpu_permit();
        install_fair_cpu_permit(Box::new(TestFairCpuPermit {
            consumed_cpu_us: Arc::new(AtomicU64::new(0)),
            action: FairCpuBudgetAction::YieldToScheduler,
        }));
        let guard = FairCpuGuard::new();
        FAIR_CPU_DISPATCH.with(|slot| {
            slot.borrow_mut()
                .as_mut()
                .expect("fair CPU dispatch should be active")
                .last_checkpoint_at = StdInstant::now()
                .checked_sub(FAIR_CPU_CHECKPOINT_MIN_INTERVAL)
                .unwrap();
        });

        burn_thread_cpu(Duration::from_micros(100));
        let mut now = LazyInstant::default();
        assert_eq!(
            fair_cpu_checkpoint(&mut now),
            FairCpuBudgetAction::YieldToScheduler
        );
        drop(guard);
    }

    impl YatpTaskQueue<TaskCell> for TestMultiLevelTaskQueue {
        fn push(&self, task_cell: TaskCell) {
            self.inner.push(KeyspaceFairTaskCell::new(task_cell, 0));
        }

        fn pop(&self) -> PopResult<TaskCell> {
            match self.inner.pop() {
                Some(task_cell) => PopResult::Ready(queue::Pop {
                    task_cell: task_cell.into_task_cell(),
                    schedule_time: std::time::Instant::now(),
                    from_local: false,
                }),
                None => PopResult::Empty,
            }
        }

        fn drain(&self) {
            while self.inner.pop().is_some() {}
        }

        fn has_ready_task(&self) -> bool {
            self.inner.len() > 0
        }
    }

    fn record_running_time_with_yatp_runner(
        inner: Arc<MultiLevelInnerQueue>,
        name: &str,
        task_id: u64,
        sleep_duration: Duration,
    ) {
        let queue = Arc::new(TestMultiLevelTaskQueue {
            inner: inner.clone(),
        });
        let custom_builder = CustomBuilder::new(CustomConfig::default().name(Some(name)), queue);
        let runner = YatpPoolRunner::new(
            FutureRunner::default(),
            TickerWrapper::new(DefaultTicker),
            None,
            None,
            None,
            metrics::YATP_POOL_SCHEDULE_WAIT_DURATION_VEC.with_label_values(&[name]),
        );
        let runner_builder = custom_builder.runner_builder(CloneRunnerBuilder(runner));
        let mut builder = yatp::Builder::new(thd_name!(name));
        builder
            .min_thread_count(1)
            .core_thread_count(1)
            .max_thread_count(1);
        let pool =
            builder.build_with_queue_and_runner(QueueType::Custom(custom_builder), runner_builder);
        let (tx, rx) = mpsc::channel();

        pool.spawn(TaskCell::new(
            async move {
                std::thread::sleep(sleep_duration);
                tx.send(()).unwrap();
            },
            Extras::new_multilevel(task_id, None),
        ));
        rx.recv_timeout(Duration::from_secs(5)).unwrap();
        drop(pool);
    }

    fn assert_accumulated_running_time_adjusts_level(
        level_time_threshold: [Duration; multilevel::LEVEL_NUM - 1],
        expected_level: usize,
    ) {
        const TASK_ID: u64 = 1;
        const RUNNING_TIME: Duration = Duration::from_millis(20);
        let name = format!(
            "test_multi_level_inner_queue_adjusts_level_to_{}",
            expected_level
        );
        let inner = Arc::new(MultiLevelInnerQueue::new_with_config(
            no_keyspace_inner_metrics(),
            new_test_fair_queue_popped_task_count(),
            level_time_threshold,
            None,
        ));

        record_running_time_with_yatp_runner(inner.clone(), &name, TASK_ID, RUNNING_TIME);
        inner.push(new_test_keyspace_fair_task(TASK_ID, None));

        assert_inner_queue_task_level(&inner, expected_level);
        assert_eq!(pop_inner_queue_task_id(&inner), TASK_ID);
    }

    #[test]
    fn test_task_keyspace_id_metadata_roundtrip() {
        let mut extras = Extras::new_multilevel(1, None);

        set_task_meta(
            &mut extras,
            &TaskMetadata {
                keyspace_id: 42,
                fair_cpu_admission: None,
            },
        );

        assert_eq!(extras.metadata(), 42u32.to_be_bytes());
        assert_eq!(task_keyspace_id_from_metadata(extras.metadata()), 42);

        let deadline = StdInstant::now() + Duration::from_secs(5);
        set_task_meta(
            &mut extras,
            &TaskMetadata {
                keyspace_id: 44,
                fair_cpu_admission: Some(FairCpuAdmissionMetadata {
                    task_type: 7,
                    deadline: Some(deadline),
                }),
            },
        );

        assert_eq!(task_keyspace_id_from_metadata(extras.metadata()), 44);
        let metadata = task_metadata_from_metadata(extras.metadata());
        assert_eq!(metadata.keyspace_id, 44);
        assert_eq!(
            metadata.fair_cpu_admission,
            Some(FairCpuAdmissionMetadata {
                task_type: 7,
                deadline: Some(deadline),
            })
        );

        set_task_meta(
            &mut extras,
            &TaskMetadata {
                keyspace_id: 45,
                fair_cpu_admission: Some(FairCpuAdmissionMetadata {
                    task_type: 8,
                    deadline: None,
                }),
            },
        );

        let metadata = task_metadata_from_metadata(extras.metadata());
        assert_eq!(metadata.keyspace_id, 45);
        assert_eq!(
            metadata.fair_cpu_admission,
            Some(FairCpuAdmissionMetadata {
                task_type: 8,
                deadline: None,
            })
        );
    }

    #[test]
    fn test_task_keyspace_id_falls_back_to_default_for_invalid_metadata() {
        assert_eq!(
            task_keyspace_id_from_metadata(&[]),
            DEFAULT_TASK_KEYSPACE_ID
        );
        assert_eq!(
            task_keyspace_id_from_metadata(&[1, 2, 3]),
            DEFAULT_TASK_KEYSPACE_ID
        );
        assert_eq!(
            task_keyspace_id_from_metadata(&[1, 0, 0, 0, 42]),
            DEFAULT_TASK_KEYSPACE_ID
        );
        assert_eq!(
            task_keyspace_id_from_metadata(&[0, 0, 0, 42, 0, 0, 0, 0, 0, 0, 0, 99]),
            DEFAULT_TASK_KEYSPACE_ID
        );
    }

    #[test]
    fn test_keyspace_fair_queue_routes_missing_metadata_to_default_keyspace() {
        let queue = YatpKeyspaceFairQueue::new(None, None);
        let recorder = queue.get_fairness_recorder();

        queue.push(new_test_task(1, None));
        queue.push(new_test_task(2, Some(2)));
        recorder.increase_vt(DEFAULT_TASK_KEYSPACE_ID, 10);

        assert_eq!(pop_task_id(&queue), 2);
        assert_eq!(pop_task_id(&queue), 1);
        assert!(queue.pop().is_empty());
    }

    #[test]
    fn test_keyspace_fair_queue_recorder_changes_next_popped_keyspace() {
        let queue = YatpKeyspaceFairQueue::new(None, None);
        let recorder = queue.get_fairness_recorder();

        queue.push(new_test_task(1, Some(1)));
        queue.push(new_test_task(2, Some(2)));
        recorder.increase_vt(1, 10);

        assert_eq!(pop_task_id(&queue), 2);
        assert_eq!(pop_task_id(&queue), 1);
        assert!(queue.pop().is_empty());
    }

    #[test]
    fn test_keyspace_fair_queue_propagates_waited_slices_without_metrics() {
        const KEYSPACE_ID: u32 = 1;
        let queue = YatpKeyspaceFairQueue::new(None, None);

        queue.push(new_test_task(1, Some(KEYSPACE_ID)));
        queue.push(new_test_task(2, Some(KEYSPACE_ID)));
        queue.push(new_test_task(3, Some(KEYSPACE_ID)));

        assert_eq!(pop_fair_queue_waited_task_slices(&queue), Some(0));
        assert_eq!(pop_fair_queue_waited_task_slices(&queue), Some(1));
        assert_eq!(pop_fair_queue_waited_task_slices(&queue), Some(2));
        assert!(queue.pop().is_empty());
    }

    #[test]
    fn test_keyspace_fair_queue_selects_victim_after_concurrent_noisy_tasks_are_charged() {
        const NOISY_KEYSPACE_ID: u32 = 1;
        const VICTIM_KEYSPACE_ID: u32 = 2;
        const WORKER_COUNT: u64 = 8;
        const NOISY_TASK_CHARGE_US: u64 = 10;

        let queue = YatpKeyspaceFairQueue::new(None, None);
        let recorder = queue.get_fairness_recorder();

        for task_id in 0..=WORKER_COUNT {
            queue.push(new_test_task(task_id, Some(NOISY_KEYSPACE_ID)));
        }
        queue.push(new_test_task(100, Some(VICTIM_KEYSPACE_ID)));
        recorder.increase_vt(VICTIM_KEYSPACE_ID, 1);

        // Simulate workers popping before any popped noisy task reports its CPU
        // charge. The noisy keyspace can overshoot within this worker window.
        let mut popped_noisy_tasks = 0;
        for _ in 0..WORKER_COUNT {
            let (_, keyspace_id) = pop_task_id_and_keyspace(&queue);
            if keyspace_id == NOISY_KEYSPACE_ID {
                popped_noisy_tasks += 1;
            }
        }
        assert_eq!(popped_noisy_tasks, WORKER_COUNT);

        for _ in 0..popped_noisy_tasks {
            recorder.increase_vt(NOISY_KEYSPACE_ID, NOISY_TASK_CHARGE_US);
        }

        // Once the popped noisy tasks are charged, the remaining queued noisy
        // task should no longer starve the victim keyspace.
        assert_eq!(pop_task_id_and_keyspace(&queue), (100, VICTIM_KEYSPACE_ID));
    }

    #[test]
    fn test_keyspace_fair_queue_records_keyspace_metrics() {
        const POOL_NAME: &str = "test_keyspace_fair_queue_records_keyspace_metrics";
        const KEYSPACE_1: u32 = 11;
        const KEYSPACE_2: u32 = 22;

        let _mock_now = mock_now_coarse_secs(0);
        let queue_length = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_queue_length",
                "test keyspace fair queue length",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let queue_task_slice_waits = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_queue_task_slice_waits",
                "test keyspace fair queue task slice waits",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let queue_task_slice_wait_duration = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_queue_task_slice_wait_duration_seconds",
                "test keyspace fair queue task slice wait duration",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_slice_schedule_interval = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_queue_task_slice_schedule_interval_seconds",
                "test keyspace task slice schedule interval",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_slice_cpu_duration = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_queue_task_slice_cpu_duration_seconds",
                "test keyspace task cpu time",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_wall_time = prometheus::CounterVec::new(
            prometheus::Opts::new(
                "test_keyspace_fair_queue_task_wall_seconds_total",
                "test keyspace task wall time",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let group_count = IntGauge::new(
            "test_keyspace_fair_group_count",
            "test keyspace group count",
        )
        .unwrap();
        let histogram_resolve_count = Arc::new(AtomicUsize::new(0));
        let queue = YatpKeyspaceFairQueue::new(
            Some(new_test_keyspace_metrics(
                POOL_NAME,
                group_count.clone(),
                queue_length.clone(),
                queue_task_slice_waits.clone(),
                queue_task_slice_wait_duration.clone(),
                task_slice_schedule_interval,
                task_slice_cpu_duration.clone(),
                task_wall_time.clone(),
                Some(histogram_resolve_count.clone()),
                None,
            )),
            None,
        );
        let keyspace_1 = KEYSPACE_1.to_string();
        let keyspace_2 = KEYSPACE_2.to_string();
        let histogram_1 = queue_length.with_label_values(&[POOL_NAME, &keyspace_1]);
        let histogram_2 = queue_length.with_label_values(&[POOL_NAME, &keyspace_2]);
        let task_slice_waits_histogram_1 =
            queue_task_slice_waits.with_label_values(&[POOL_NAME, &keyspace_1]);
        let task_slice_waits_histogram_2 =
            queue_task_slice_waits.with_label_values(&[POOL_NAME, &keyspace_2]);
        let task_slice_wait_duration_histogram_1 =
            queue_task_slice_wait_duration.with_label_values(&[POOL_NAME, &keyspace_1]);
        let task_slice_wait_duration_histogram_2 =
            queue_task_slice_wait_duration.with_label_values(&[POOL_NAME, &keyspace_2]);

        assert_eq!(group_count.get(), 0);

        queue.push(new_test_task(1, Some(KEYSPACE_1)));
        queue.push(new_test_task(2, Some(KEYSPACE_1)));
        queue.push(new_test_task(3, Some(KEYSPACE_2)));

        assert_eq!(group_count.get(), 0);
        assert_eq!(histogram_1.get_sample_count(), 2);
        assert_eq!(histogram_1.get_sample_sum(), 1.0);
        assert_eq!(histogram_2.get_sample_count(), 1);
        assert_eq!(histogram_2.get_sample_sum(), 0.0);
        assert_eq!(histogram_resolve_count.load(Ordering::Relaxed), 2);

        set_mock_now_coarse_secs(
            YatpKeyspaceFairQueue::KEYSPACE_FAIR_MAINTENANCE_INTERVAL.as_secs(),
        );
        queue.push(new_test_task(4, Some(KEYSPACE_2)));
        queue.maintain_queue_if_needed();

        assert_eq!(group_count.get(), 2);
        assert_eq!(histogram_2.get_sample_count(), 2);
        assert_eq!(histogram_2.get_sample_sum(), 1.0);
        assert_eq!(histogram_resolve_count.load(Ordering::Relaxed), 2);

        for _ in 0..4 {
            pop_task_id(&queue);
        }
        assert_eq!(
            task_slice_waits_histogram_1.get_sample_count()
                + task_slice_waits_histogram_2.get_sample_count(),
            4
        );
        assert_eq!(
            task_slice_waits_histogram_1.get_sample_sum()
                + task_slice_waits_histogram_2.get_sample_sum(),
            6.0
        );
        assert_eq!(
            task_slice_wait_duration_histogram_1.get_sample_count()
                + task_slice_wait_duration_histogram_2.get_sample_count(),
            0
        );

        set_mock_now_coarse_secs(
            YatpKeyspaceFairQueue::MAX_KEYSPACE_INACTIVE_SECS
                + YatpKeyspaceFairQueue::KEYSPACE_FAIR_INACTIVE_CLEANUP_INTERVAL.as_secs()
                + 1,
        );
        queue.maintain_queue_if_needed();

        assert_eq!(group_count.get(), 0);

        queue.push(new_test_task(5, Some(KEYSPACE_1)));

        assert_eq!(histogram_1.get_sample_count(), 3);
        assert_eq!(histogram_1.get_sample_sum(), 1.0);
        assert_eq!(histogram_resolve_count.load(Ordering::Relaxed), 2);
    }

    #[test]
    fn test_keyspace_metrics_removes_inactive_cached_metrics() {
        const POOL_NAME: &str = "test_keyspace_metrics_removes_inactive_cached_metrics";
        const KEYSPACE_ID: u32 = 42;

        let _mock_now = mock_now_coarse_secs(0);
        let queue_length = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_metrics_cleanup_queue_length",
                "test keyspace metrics cleanup queue length",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let queue_task_slice_waits = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_metrics_cleanup_task_slice_waits",
                "test keyspace metrics cleanup task slice waits",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let queue_task_slice_wait_duration = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_metrics_cleanup_task_slice_wait_duration_seconds",
                "test keyspace metrics cleanup task slice wait duration",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_slice_schedule_interval = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_metrics_cleanup_task_slice_schedule_interval_seconds",
                "test keyspace task slice schedule interval",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_slice_cpu_duration = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_metrics_cleanup_task_slice_cpu_duration_seconds",
                "test keyspace metrics cleanup task cpu time",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_wall_time = prometheus::CounterVec::new(
            prometheus::Opts::new(
                "test_keyspace_metrics_cleanup_task_wall_seconds_total",
                "test keyspace metrics cleanup task wall time",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let group_count = IntGauge::new(
            "test_keyspace_metrics_cleanup_group_count",
            "test keyspace metrics cleanup group count",
        )
        .unwrap();
        let metrics_resolve_count = Arc::new(AtomicUsize::new(0));
        let metrics_remove_count = Arc::new(AtomicUsize::new(0));
        let keyspace_metrics = new_test_keyspace_metrics(
            POOL_NAME,
            group_count,
            queue_length,
            queue_task_slice_waits,
            queue_task_slice_wait_duration,
            task_slice_schedule_interval,
            task_slice_cpu_duration,
            task_wall_time,
            Some(metrics_resolve_count.clone()),
            Some(metrics_remove_count.clone()),
        );

        let metrics = keyspace_metrics.get_keyspace_metrics(KEYSPACE_ID);
        assert_eq!(metrics_resolve_count.load(Ordering::Relaxed), 1);
        assert_eq!(
            keyspace_metrics.clear_inactive_keyspace_metrics_if_needed(),
            0
        );

        set_mock_now_coarse_secs(
            KeyspaceMetrics::MAX_KEYSPACE_METRICS_INACTIVE_SECS
                + KeyspaceMetrics::KEYSPACE_METRICS_CLEANUP_INTERVAL.as_secs()
                + 1,
        );
        assert_eq!(
            keyspace_metrics.clear_inactive_keyspace_metrics_if_needed(),
            0
        );
        assert_eq!(metrics_remove_count.load(Ordering::Relaxed), 0);

        drop(metrics);
        assert_eq!(
            keyspace_metrics.clear_inactive_keyspace_metrics_if_needed(),
            0
        );

        set_mock_now_coarse_secs(
            KeyspaceMetrics::MAX_KEYSPACE_METRICS_INACTIVE_SECS
                + KeyspaceMetrics::KEYSPACE_METRICS_CLEANUP_INTERVAL.as_secs() * 2
                + 1,
        );
        assert_eq!(
            keyspace_metrics.clear_inactive_keyspace_metrics_if_needed(),
            1
        );
        assert_eq!(metrics_remove_count.load(Ordering::Relaxed), 1);

        let _metrics = keyspace_metrics.get_keyspace_metrics(KEYSPACE_ID);
        assert_eq!(metrics_resolve_count.load(Ordering::Relaxed), 2);
    }

    #[test]
    fn test_keyspace_fair_pool_records_queue_task_slice_wait_duration() {
        const POOL_NAME: &str = "test_keyspace_fair_pool_records_queue_task_slice_wait_duration";
        const KEYSPACE_ID: u32 = 33;
        const TASK_COUNT: u64 = 3;

        let queue_length = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_pool_queue_length",
                "test keyspace fair pool queue length",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let queue_task_slice_waits = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_pool_queue_task_slice_waits",
                "test keyspace fair pool queue task slice waits",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let queue_task_slice_wait_duration = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_pool_queue_task_slice_wait_duration_seconds",
                "test keyspace fair pool queue task slice wait duration",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_slice_schedule_interval = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_pool_task_slice_schedule_interval_seconds",
                "test keyspace fair pool task slice schedule interval",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_slice_cpu_duration = HistogramVec::new(
            prometheus::HistogramOpts::new(
                "test_keyspace_fair_pool_task_slice_cpu_duration_seconds",
                "test keyspace fair pool task cpu time",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let task_wall_time = prometheus::CounterVec::new(
            prometheus::Opts::new(
                "test_keyspace_fair_pool_task_wall_seconds_total",
                "test keyspace fair pool task wall time",
            ),
            &["name", "keyspace_id"],
        )
        .unwrap();
        let group_count = IntGauge::new(
            "test_keyspace_fair_pool_group_count",
            "test keyspace fair pool group count",
        )
        .unwrap();
        let mut builder = YatpPoolBuilder::new(DefaultTicker);
        let (pool, _) = builder
            .name_prefix(POOL_NAME)
            .keyspace_metrics(Some(new_test_keyspace_metrics(
                POOL_NAME,
                group_count,
                queue_length.clone(),
                queue_task_slice_waits.clone(),
                queue_task_slice_wait_duration.clone(),
                task_slice_schedule_interval.clone(),
                task_slice_cpu_duration,
                task_wall_time,
                None,
                None,
            )))
            .build_keyspace_fair_pool(None);
        let (tx, rx) = mpsc::channel();

        for task_id in 0..TASK_COUNT {
            let tx = tx.clone();
            let mut extras = Extras::new_multilevel(task_id, None);
            set_task_meta(
                &mut extras,
                &TaskMetadata {
                    keyspace_id: KEYSPACE_ID,
                    fair_cpu_admission: None,
                },
            );
            pool.spawn(TaskCell::new(
                async move {
                    tx.send(()).unwrap();
                },
                extras,
            ));
        }
        drop(tx);
        for _ in 0..TASK_COUNT {
            rx.recv_timeout(Duration::from_secs(5)).unwrap();
        }
        drop(pool);

        let keyspace_id = KEYSPACE_ID.to_string();
        let task_slice_waits_histogram =
            queue_task_slice_waits.with_label_values(&[POOL_NAME, &keyspace_id]);
        let task_slice_wait_duration_histogram =
            queue_task_slice_wait_duration.with_label_values(&[POOL_NAME, &keyspace_id]);
        assert_eq!(task_slice_waits_histogram.get_sample_count(), TASK_COUNT);
        assert_eq!(
            task_slice_wait_duration_histogram.get_sample_count(),
            TASK_COUNT
        );
        assert_eq!(
            task_slice_schedule_interval
                .with_label_values(&[POOL_NAME, &keyspace_id])
                .get_sample_count(),
            TASK_COUNT
        );
    }

    #[test]
    fn test_keyspace_fair_queue_admission_pending_does_not_pop() {
        let controller = Arc::new(TestAdmissionController::new());
        controller.set_pending(true);
        let queue = YatpKeyspaceFairQueue::new(None, Some(controller.clone()));

        queue.push(new_test_task(1, Some(42)));

        match queue.pop() {
            PopResult::Pending { .. } => {}
            PopResult::Ready(_) => panic!("pending admission should not pop a task"),
            PopResult::Empty => panic!("queue should still have a pending task"),
        }
        assert_eq!(controller.precheck_count(), 1);
        assert!(controller.admits().is_empty());

        controller.set_pending(false);
        assert_eq!(pop_task_id(&queue), 1);
        assert_eq!(controller.precheck_count(), 2);
    }

    #[test]
    fn test_keyspace_fair_queue_empty_does_not_precheck_admission() {
        let controller = Arc::new(TestAdmissionController::new());
        controller.set_pending(true);
        let queue = YatpKeyspaceFairQueue::new(None, Some(controller.clone()));

        match queue.pop() {
            PopResult::Empty => {}
            PopResult::Pending { .. } => panic!("empty queue should not wait for admission"),
            PopResult::Ready(_) => panic!("empty queue should not pop a task"),
        }
        assert_eq!(controller.precheck_count(), 0);

        queue.push(new_test_task(1, Some(42)));
        match queue.pop() {
            PopResult::Pending { .. } => {}
            PopResult::Ready(_) => panic!("pending admission should still gate a queued task"),
            PopResult::Empty => panic!("queue should have a gated task"),
        }
        assert_eq!(controller.precheck_count(), 1);
    }

    #[test]
    fn test_keyspace_fair_queue_admits_fair_cpu_marked_task() {
        let controller = Arc::new(TestAdmissionController::new());
        let queue = YatpKeyspaceFairQueue::new(None, Some(controller.clone()));

        queue.push(new_test_task_with_fair_cpu_admission(1, 42, 7));

        assert_eq!(pop_task_id(&queue), 1);
        assert_eq!(controller.precheck_count(), 1);
        assert_eq!(controller.admits(), vec![(42, 7)]);
    }

    #[test]
    fn test_keyspace_fair_queue_does_not_admit_unmarked_task() {
        let controller = Arc::new(TestAdmissionController::new());
        let queue = YatpKeyspaceFairQueue::new(None, Some(controller.clone()));

        queue.push(new_test_task(1, Some(42)));

        assert_eq!(pop_task_id(&queue), 1);
        assert_eq!(controller.precheck_count(), 1);
        assert!(controller.admits().is_empty());
    }

    #[test]
    fn test_multi_level_inner_queue_push_pop_updates_len() {
        let queue = new_test_inner_queue();

        assert_eq!(queue.len(), 0);

        queue.push(new_test_keyspace_fair_task(1, None));
        assert_eq!(queue.len(), 1);

        assert_eq!(pop_inner_queue_task_id(&queue), 1);
        assert_eq!(queue.len(), 0);
        assert!(queue.pop().is_none());
    }

    #[test]
    fn test_multi_level_inner_queue_pushes_task_to_fixed_level() {
        let queue = new_test_inner_queue();
        let fixed_level = (multilevel::LEVEL_NUM - 1) as u8;

        queue.push(new_test_keyspace_fair_task_with_fixed_level(
            1,
            None,
            Some(fixed_level),
        ));

        assert_inner_queue_task_level(&queue, fixed_level as usize);
        assert_eq!(pop_inner_queue_task_id(&queue), 1);
        assert_eq!(queue.len(), 0);
    }

    #[test]
    fn test_multi_level_inner_queue_adjusts_level_by_accumulated_running_time() {
        assert_accumulated_running_time_adjusts_level(
            [Duration::from_secs(10), Duration::from_secs(20)],
            0,
        );
        assert_accumulated_running_time_adjusts_level(
            [Duration::from_millis(1), Duration::from_secs(10)],
            1,
        );
        assert_accumulated_running_time_adjusts_level(
            [Duration::from_millis(1), Duration::from_millis(2)],
            2,
        );
    }

    #[test]
    fn test_keyspace_fair_queue_removes_inactive_keyspaces_by_cleanup_interval() {
        let _mock_now = mock_now_coarse_secs(0);
        let queue = YatpKeyspaceFairQueue::new(None, None);

        queue.push(new_test_task(1, Some(1)));
        assert_eq!(pop_task_id(&queue), 1);
        assert_eq!(queue.inner.group_count(), 1);

        set_mock_now_coarse_secs(
            YatpKeyspaceFairQueue::KEYSPACE_FAIR_INACTIVE_CLEANUP_INTERVAL.as_secs() - 1,
        );
        queue.maintain_queue_if_needed();
        assert_eq!(queue.inner.group_count(), 1);
        assert_eq!(queue.last_inactive_cleanup_secs.load(Ordering::Relaxed), 0);

        set_mock_now_coarse_secs(YatpKeyspaceFairQueue::MAX_KEYSPACE_INACTIVE_SECS + 1);
        queue.maintain_queue_if_needed();
        assert_eq!(queue.inner.group_count(), 0);
        assert_eq!(
            queue.last_inactive_cleanup_secs.load(Ordering::Relaxed),
            YatpKeyspaceFairQueue::MAX_KEYSPACE_INACTIVE_SECS + 1
        );
    }

    #[test]
    fn test_keyspace_fair_queue_removes_only_after_idle_timeout_exceeded() {
        let _mock_now = mock_now_coarse_secs(0);
        let queue = YatpKeyspaceFairQueue::new(None, None);

        queue.push(new_test_task(1, Some(1)));
        assert_eq!(pop_task_id(&queue), 1);
        assert_eq!(queue.inner.group_count(), 1);

        set_mock_now_coarse_secs(YatpKeyspaceFairQueue::MAX_KEYSPACE_INACTIVE_SECS);
        queue.maintain_queue_if_needed();
        assert_eq!(queue.inner.group_count(), 1);

        set_mock_now_coarse_secs(YatpKeyspaceFairQueue::MAX_KEYSPACE_INACTIVE_SECS + 61);
        queue.maintain_queue_if_needed();
        assert_eq!(queue.inner.group_count(), 0);
        assert_eq!(
            queue.last_inactive_cleanup_secs.load(Ordering::Relaxed),
            YatpKeyspaceFairQueue::MAX_KEYSPACE_INACTIVE_SECS + 61
        );
    }

    #[test]
    fn test_record_schedule_wait_duration() {
        let name = "test_record_schedule_wait_duration";
        let pool = YatpPoolBuilder::new(DefaultTicker)
            .name_prefix(name)
            .build_single_level_pool();
        let (tx, rx) = mpsc::channel();
        for _ in 0..3 {
            let tx = tx.clone();
            pool.spawn(async move {
                GLOBAL_TIMER_HANDLE
                    .delay(std::time::Instant::now() + Duration::from_millis(100))
                    .compat()
                    .await
                    .unwrap();
                tx.send(()).unwrap();
            });
        }
        for _ in 0..3 {
            rx.recv().unwrap();
        }
        // Drop the pool so the local metrics are flushed.
        drop(pool);
        let histogram = metrics::YATP_POOL_SCHEDULE_WAIT_DURATION_VEC.with_label_values(&[name]);
        assert_eq!(histogram.get_sample_count() as u32, 6, "{:?}", histogram);
    }
}