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

use std::{
    borrow::Cow,
    cell::RefCell,
    fmt,
    future::Future,
    pin::Pin,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
        mpsc::SyncSender,
    },
    task::{Context, Poll},
    time::{Duration, Instant as StdInstant},
};

use cpu_time::ThreadTime;
use dashmap::DashMap;
use file_system::{IoType, set_io_type};
use futures::{channel::oneshot, future::TryFutureExt};
use kvproto::kvrpcpb::CommandPri;
use online_config::{ConfigChange, ConfigManager, ConfigValue, Result as CfgResult};
use pin_project::pin_project;
use prometheus::{IntCounter, IntGauge};
use resource_control::{CpuThrottleManager, ReadLimiter};
use thiserror::Error;
use tikv_util::{
    fairness::FairnessRecorder,
    sys::{SysQuota, cpu_time::ProcessStat, thread::ThreadBuildWrapper},
    task_reschedule::PollRescheduleContextGuard,
    time::Instant,
    worker::{Runnable, RunnableWithTimer, Scheduler, Worker},
    yatp_pool::{
        self, FairAdmitOutcome, FairCpuBudgetAction, FairCpuPermit, FairQueueAdmissionController,
        FairQueueAdmissionResult, FuturePool, KeyspaceInnerMetrics, KeyspaceInnerMetricsHandle,
        KeyspaceMetrics, PoolTicker, ScalableTokioHandle, ScalableTokioRuntime,
        YatpDispatchContext, YatpKeyspaceFairQueue, YatpPoolBuilder,
    },
};
use tracker::{PoolTaskMetrics, TrackedFuture, with_tls_tracker};
use yatp::{
    metrics::MULTILEVEL_LEVEL_ELAPSED, pool::Remote, queue::Extras, task::future::TaskCell,
};

use self::metrics::*;
use crate::{
    config::{UNIFIED_READPOOL_MIN_CONCURRENCY, UnifiedReadPoolConfig},
    storage::kv::{Engine, FlowStatsReporter, destroy_tls_engine, set_tls_engine},
};

// the duration to check auto-scale unified-thread-pool's thread
const READ_POOL_THREAD_CHECK_DURATION: Duration = Duration::from_secs(10);
// consider scale out read pool size if the average thread cpu usage is higher
// than this threshold.
const READ_POOL_THREAD_HIGH_THRESHOLD: f64 = 0.8;
// consider scale in read pool size if the average thread cpu usage is lower
// than this threshold.
const READ_POOL_THREAD_LOW_THRESHOLD: f64 = 0.7;
// avg running tasks per-thread that indicates read-pool is busy
const RUNNING_TASKS_PER_THREAD_THRESHOLD: i64 = 3;
const TOKIO_REPORTER_TICK_CHECK_INTERVAL: Duration = Duration::from_millis(100);
const READ_POOL_REJECT_REASON_GLOBAL_FULL: &str = "global_full";
const READ_POOL_REJECT_REASON_KEYSPACE_FULL: &str = "keyspace_full";
const KEYSPACE_TASK_COUNTER_CLEANUP_THRESHOLD: usize = 8192;

fn record_unified_read_pool_reject(keyspace_id: Option<u32>, reason: &str) {
    let _ = record_unified_read_pool_reject_and_get_counter(keyspace_id, reason);
}

fn record_unified_read_pool_reject_and_get_counter(
    keyspace_id: Option<u32>,
    reason: &str,
) -> IntCounter {
    let counter = unified_read_pool_reject_counter(keyspace_id, reason);
    counter.inc();
    counter
}

fn unified_read_pool_reject_keyspace_label(keyspace_id: Option<u32>) -> Cow<'static, str> {
    match keyspace_id {
        Some(keyspace_id) => Cow::Owned(keyspace_id.to_string()),
        None => Cow::Borrowed("-"),
    }
}

fn unified_read_pool_reject_counter(keyspace_id: Option<u32>, reason: &str) -> IntCounter {
    let pool_name = get_unified_read_pool_name();
    let keyspace_label = unified_read_pool_reject_keyspace_label(keyspace_id);
    UNIFIED_READ_POOL_REJECT_TASKS.with_label_values(&[
        pool_name.as_str(),
        reason,
        keyspace_label.as_ref(),
    ])
}

fn set_yatp_task_metadata_if_fair(
    extras: &mut Extras,
    keyspace_fairness_enabled: bool,
    keyspace_id: u32,
    fair_cpu_admission: Option<yatp_pool::FairCpuAdmissionMetadata>,
) {
    if keyspace_fairness_enabled {
        yatp_pool::set_task_meta(
            extras,
            &yatp_pool::TaskMetadata {
                keyspace_id,
                fair_cpu_admission,
            },
        );
    }
}

#[derive(Clone)]
pub struct KeyspaceTaskLimiter {
    max_tasks_per_keyspace: Arc<AtomicUsize>,
    task_counts: Arc<DashMap<u32, Arc<AtomicUsize>>>,
    zero_transitions_since_cleanup: Arc<AtomicUsize>,
}

impl KeyspaceTaskLimiter {
    fn new(max_tasks_per_keyspace: usize) -> Self {
        Self {
            max_tasks_per_keyspace: Arc::new(AtomicUsize::new(max_tasks_per_keyspace)),
            task_counts: Arc::new(DashMap::new()),
            zero_transitions_since_cleanup: Arc::new(AtomicUsize::new(0)),
        }
    }

    fn set_max_tasks_per_keyspace(&self, max_tasks_per_keyspace: usize) {
        self.max_tasks_per_keyspace
            .store(max_tasks_per_keyspace, Ordering::Relaxed);
    }

    fn try_acquire(
        &self,
        keyspace_id: u32,
        bypass_limit: bool,
    ) -> Result<KeyspaceTaskPermit, KeyspaceTaskLimitError> {
        self.cleanup_zero_counters_if_needed();

        // The DashMap entry guard serializes check+increment with removal on the
        // same shard. Once incremented, a live permit keeps this entry's count
        // above zero, so Drop will decrement the same Arc it acquired.
        if let Some(task_count) = self.task_counts.get(&keyspace_id) {
            self.try_increment(task_count.value(), bypass_limit)?;
            return Ok(KeyspaceTaskPermit {
                keyspace_id,
                task_count: Arc::clone(task_count.value()),
                limiter: self.clone(),
            });
        }

        let task_count = self
            .task_counts
            .entry(keyspace_id)
            .or_insert_with(|| Arc::new(AtomicUsize::new(0)));
        self.try_increment(task_count.value(), bypass_limit)?;

        Ok(KeyspaceTaskPermit {
            keyspace_id,
            task_count: Arc::clone(task_count.value()),
            limiter: self.clone(),
        })
    }

    fn try_increment(
        &self,
        task_count: &AtomicUsize,
        bypass_limit: bool,
    ) -> Result<(), KeyspaceTaskLimitError> {
        loop {
            let current = task_count.load(Ordering::Relaxed);
            if !bypass_limit && current >= self.max_tasks_per_keyspace.load(Ordering::Relaxed) {
                return Err(KeyspaceTaskLimitError::Full);
            }
            if task_count
                .compare_exchange_weak(current, current + 1, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                return Ok(());
            }
        }
    }

    fn cleanup_zero_counters_if_needed(&self) {
        if self.zero_transitions_since_cleanup.load(Ordering::Relaxed)
            < KEYSPACE_TASK_COUNTER_CLEANUP_THRESHOLD
        {
            return;
        }

        let transitions = self
            .zero_transitions_since_cleanup
            .swap(0, Ordering::Relaxed);
        if transitions >= KEYSPACE_TASK_COUNTER_CLEANUP_THRESHOLD {
            self.cleanup_zero_counters();
        } else if transitions > 0 {
            self.zero_transitions_since_cleanup
                .fetch_add(transitions, Ordering::Relaxed);
        }
    }

    fn cleanup_zero_counters(&self) {
        let keyspace_ids: Vec<u32> = self.task_counts.iter().map(|entry| *entry.key()).collect();

        for keyspace_id in keyspace_ids {
            self.task_counts.remove_if(&keyspace_id, |_, task_count| {
                task_count.load(Ordering::Relaxed) == 0
            });
        }
    }

    #[cfg(test)]
    fn task_count(&self, keyspace_id: u32) -> usize {
        self.task_counts
            .get(&keyspace_id)
            .map(|task_count| task_count.load(Ordering::Relaxed))
            .unwrap_or(0)
    }

    #[cfg(test)]
    fn has_task_counter(&self, keyspace_id: u32) -> bool {
        self.task_counts.contains_key(&keyspace_id)
    }

    #[cfg(test)]
    fn max_tasks_per_keyspace(&self) -> usize {
        self.max_tasks_per_keyspace.load(Ordering::Relaxed)
    }
}

#[derive(Debug)]
enum KeyspaceTaskLimitError {
    Full,
}

struct KeyspaceTaskPermit {
    keyspace_id: u32,
    task_count: Arc<AtomicUsize>,
    limiter: KeyspaceTaskLimiter,
}

impl Drop for KeyspaceTaskPermit {
    fn drop(&mut self) {
        match self
            .task_count
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
                current.checked_sub(1)
            }) {
            Ok(1) => {
                self.limiter
                    .zero_transitions_since_cleanup
                    .fetch_add(1, Ordering::Relaxed);
            }
            Ok(_) => {}
            Err(current) => {
                warn!(
                    "read pool keyspace task counter underflow";
                    "keyspace_id" => self.keyspace_id,
                    "current" => current,
                );
            }
        }
    }
}

pub enum ReadPool {
    FuturePools {
        read_pool_high: FuturePool,
        read_pool_normal: FuturePool,
        read_pool_low: FuturePool,
    },
    Yatp {
        pool: yatp::ThreadPool<TaskCell>,
        keyspace_fair_queue: Option<Arc<YatpKeyspaceFairQueue>>,
        fair_global_cpu_admission_controller: Option<Arc<FairGlobalCpuAdmissionController>>,
        fairness_use_wall_time: Arc<AtomicBool>,
        fairness_slice_precharge_us: Arc<AtomicU64>,
        fairness_max_time_slice_us: Arc<AtomicU64>,
        keyspace_metrics: Option<Arc<KeyspaceMetrics>>,
        running_tasks: IntGauge,
        running_threads: IntGauge,
        max_tasks: Arc<AtomicUsize>,
        keyspace_task_limiter: Option<KeyspaceTaskLimiter>,
        pool_size: Arc<AtomicUsize>,
    },
    Tokio {
        runtime: ScalableTokioRuntime,
        running_tasks: IntGauge,
        max_tasks: Arc<AtomicUsize>,
        pool_size: Arc<AtomicUsize>,
        task_monitor: tokio_metrics::TaskMonitor,
        pool_name: String,
    },
}

/// Consumer for read-pool task poll CPU time.
pub trait ReadPoolTaskPollCpuRecorder: fmt::Debug + Send + Sync {
    /// Records thread CPU time consumed by one poll call.
    fn observe_poll_cpu_time(&self, cpu_time: Duration);

    /// Called when read-pool task observation finishes.
    fn finish(&self) {}
}

/// Read request types that can participate in fair global CPU admission.
///
/// The numeric values are stored in YATP task metadata and must remain stable.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum FairCpuTaskType {
    DagSelect = 1,
    DagIndex = 2,
    KvBatchGet = 3,
    KvGet = 4,
    KvBatchGetCommand = 5,
    KvScan = 6,
    KvScanLock = 7,
}

/// Optional scheduling behavior attached to one read-pool task.
#[derive(Default)]
pub struct ReadPoolTaskOptions {
    poll_cpu_recorder: Option<Arc<dyn ReadPoolTaskPollCpuRecorder>>,
    fair_cpu_admission: Option<yatp_pool::FairCpuAdmissionMetadata>,
    poll_tracker_observer: Option<Box<dyn YatpPollEventTracker + Send>>,
}

impl ReadPoolTaskOptions {
    /// Observes CPU time consumed by each poll.
    pub fn with_poll_cpu_recorder(
        mut self,
        poll_cpu_recorder: Arc<dyn ReadPoolTaskPollCpuRecorder>,
    ) -> Self {
        self.poll_cpu_recorder = Some(poll_cpu_recorder);
        self
    }

    /// Admits the task through the fair global CPU gate.
    pub fn with_fair_cpu_admission(
        mut self,
        task_type: FairCpuTaskType,
        request_deadline: Option<StdInstant>,
    ) -> Self {
        self.fair_cpu_admission = Some(yatp_pool::FairCpuAdmissionMetadata {
            task_type: task_type as u8,
            deadline: request_deadline,
        });
        self
    }

    /// Observes YATP dispatch and poll events.
    pub fn with_poll_tracker_observer(
        mut self,
        poll_tracker_observer: Box<dyn YatpPollEventTracker + Send>,
    ) -> Self {
        self.poll_tracker_observer = Some(poll_tracker_observer);
        self
    }
}

struct FairGlobalCpuPermit {
    manager: Arc<CpuThrottleManager>,
    read_limiter: ReadLimiter,
}

impl FairCpuPermit for FairGlobalCpuPermit {
    fn consume(&mut self, cpu_delta_us: u64) -> FairCpuBudgetAction {
        let has_debt = self.manager.consume_fair_global_cpu(cpu_delta_us);
        // Keep actual usage and debt observable in dry-run mode without
        // enforcing the simulated debt on the running task.
        if has_debt && !self.read_limiter.is_dry_run() {
            FairCpuBudgetAction::YieldToScheduler
        } else {
            FairCpuBudgetAction::Continue
        }
    }
}

pub struct FairGlobalCpuAdmissionController {
    manager: Arc<CpuThrottleManager>,
    read_limiter: ReadLimiter,
}

impl FairGlobalCpuAdmissionController {
    /// Creates a fair-queue admission controller backed by the supplied global
    /// CPU token manager.
    pub fn new(manager: Arc<CpuThrottleManager>, read_limiter: ReadLimiter) -> Self {
        manager.register_fair_global_cpu_gate();
        Self {
            manager,
            read_limiter,
        }
    }

    fn update_max_read_cpu_ratio(&self, max_read_cpu_ratio: f64) {
        self.manager
            .update_fair_max_read_cpu_ratio(max_read_cpu_ratio);
    }
}

impl FairQueueAdmissionController for FairGlobalCpuAdmissionController {
    fn precheck(&self) -> FairQueueAdmissionResult {
        // Refill and publish the bucket state even in dry-run mode, but never
        // block queue admission there.
        let Some(retry_at) = self.manager.fair_global_cpu_dispatch_retry_at() else {
            return FairQueueAdmissionResult::Ready;
        };
        if self.read_limiter.is_dry_run() {
            return FairQueueAdmissionResult::Ready;
        }
        // TODO(#5480): add fair-gate observability before enabling this path in
        // production. The fair path should expose admit/pending counters and a
        // per-keyspace token-wait histogram instead of logging expected global
        // bucket waits.
        FairQueueAdmissionResult::Pending { retry_at }
    }

    fn admit(
        &self,
        _keyspace_id: u32,
        _task_type: u8,
        deadline: Option<StdInstant>,
    ) -> FairAdmitOutcome {
        if deadline.is_some_and(|deadline| StdInstant::now() >= deadline) {
            return FairAdmitOutcome::Bypassed;
        }
        FairAdmitOutcome::Admitted(Box::new(FairGlobalCpuPermit {
            manager: self.manager.clone(),
            read_limiter: self.read_limiter.clone(),
        }))
    }
}

struct KeyspaceFairnessRecorder {
    keyspace_id: u32,
    recorder: Arc<dyn FairnessRecorder<u32>>,
    use_wall_time: bool,
    slice_precharge_us: u64,
    /// Wall-time slice each poll of this task may run before it cooperatively
    /// reschedules. Like the other fields, it is snapshotted when the task is
    /// spawned, so an online change applies to later tasks.
    max_time_slice: Duration,
}

impl KeyspaceFairnessRecorder {
    #[inline]
    pub fn precharge(&self) {
        if self.slice_precharge_us > 0 {
            self.recorder
                .increase_vt(self.keyspace_id, self.slice_precharge_us);
        }
    }

    #[inline]
    pub fn observe_poll_profile(&self, profile: &TaskPollProfile) {
        let used_us = if self.use_wall_time {
            profile.wall_time.as_micros()
        } else {
            profile.cpu_time.as_micros()
        } as u64;

        if used_us >= self.slice_precharge_us {
            let extra_used_us = used_us - self.slice_precharge_us;
            if extra_used_us > 0 {
                self.recorder.increase_vt(self.keyspace_id, extra_used_us);
            }
        } else {
            self.recorder
                .decrease_vt(self.keyspace_id, self.slice_precharge_us - used_us);
        }
    }
}

/// Observer for the dispatch and execution intervals of a YATP task.
pub trait YatpPollEventTracker {
    /// Called immediately before one invocation of the task future's `poll`.
    fn on_poll_start(&mut self, profile: &TaskDispatchProfile);

    /// Called after one invocation of the task future's `poll`.
    fn on_poll_end(&mut self, task_finished: bool, profile: &TaskPollProfile);
}

/// Scheduling information for the YATP dispatch containing one task poll.
#[derive(Clone, Copy, Debug)]
pub struct TaskDispatchProfile {
    /// When the task was scheduled into the YATP queue for this dispatch.
    pub schedule_time: Option<StdInstant>,
    /// When a YATP worker started handling this dispatch.
    pub dispatch_start_time: StdInstant,
    /// CPU time consumed by the current worker thread when this poll started.
    pub poll_start_thread_time: ThreadTime,
    /// Monotonic wall-clock time when this poll started.
    pub poll_start_time: StdInstant,
    /// Time from entering the queue until the worker started this dispatch.
    pub queue_wait: Option<Duration>,
    /// Fair-queue slices dispatched ahead of this task slice.
    pub fair_queue_waited_task_slices: Option<u64>,
}

impl TaskDispatchProfile {
    fn new(
        ctx: Option<&YatpDispatchContext>,
        poll_start_thread_time: ThreadTime,
        poll_start_time: StdInstant,
    ) -> Self {
        match ctx {
            Some(ctx) => Self {
                schedule_time: ctx.schedule_time(),
                dispatch_start_time: ctx.dispatch_start_time(),
                poll_start_thread_time,
                poll_start_time,
                queue_wait: ctx.ready_to_dispatch_duration(),
                fair_queue_waited_task_slices: ctx.fair_queue_waited_task_slices(),
            },
            None => Self {
                schedule_time: None,
                dispatch_start_time: poll_start_time,
                poll_start_thread_time,
                poll_start_time,
                queue_wait: None,
                fair_queue_waited_task_slices: None,
            },
        }
    }

    /// Returns the identifier for the YATP dispatch containing this poll.
    ///
    /// This currently uses the dispatch start time. Keeping the representation
    /// behind this method lets callers identify immediate repolls without
    /// depending directly on that implementation detail.
    #[inline]
    pub fn dispatch_id(&self) -> StdInstant {
        self.dispatch_start_time
    }
}

#[derive(Clone)]
struct TrackerMetricsObserverState {
    #[cfg(feature = "testexport")]
    spawn_task_id: u64,
    #[cfg(feature = "testexport")]
    tracker_sync_sequence: Arc<AtomicU64>,
    metrics: PoolTaskMetrics,
    task_start_time: StdInstant,
    last_dispatch_id: Option<StdInstant>,
    last_poll_end_time: Option<StdInstant>,
    has_queue_wait: bool,
    has_wake_wait: bool,
}

impl Default for TrackerMetricsObserverState {
    fn default() -> Self {
        Self {
            #[cfg(feature = "testexport")]
            spawn_task_id: 0,
            #[cfg(feature = "testexport")]
            tracker_sync_sequence: Arc::new(AtomicU64::new(0)),
            metrics: PoolTaskMetrics::default(),
            task_start_time: StdInstant::now(),
            last_dispatch_id: None,
            last_poll_end_time: None,
            has_queue_wait: false,
            has_wake_wait: false,
        }
    }
}

impl TrackerMetricsObserverState {
    #[inline]
    fn record_duration(
        total_nanos: &mut u64,
        max_nanos: &mut u64,
        min_nanos: &mut u64,
        has_sample: bool,
        duration: Duration,
    ) {
        let nanos = duration.as_nanos().min(u64::MAX as u128) as u64;
        Self::record_value(total_nanos, max_nanos, min_nanos, has_sample, nanos);
    }

    #[inline]
    fn record_value(total: &mut u64, max: &mut u64, min: &mut u64, has_sample: bool, value: u64) {
        *total = total.saturating_add(value);
        *max = (*max).max(value);
        if !has_sample || value < *min {
            *min = value;
        }
    }

    #[inline]
    fn sync_to_current_tracker(&self) {
        with_tls_tracker(|tracker| {
            #[cfg(feature = "testexport")]
            {
                let tracker_sync_sequence =
                    self.tracker_sync_sequence.fetch_add(1, Ordering::Relaxed) + 1;
                test_util::observe_pool_task_metrics_sync(
                    self.spawn_task_id,
                    tracker_sync_sequence,
                    &tracker.req_info,
                    &self.metrics,
                );
            }
            tracker.metrics.read_pool.clone_from(&self.metrics);
        });
    }

    fn record_poll_start(&mut self, profile: &TaskDispatchProfile) {
        let dispatch_id = profile.dispatch_id();
        if self.last_dispatch_id == Some(dispatch_id) {
            return;
        }

        if let Some(queue_wait) = profile.queue_wait {
            Self::record_duration(
                &mut self.metrics.total_queue_wait_nanos,
                &mut self.metrics.max_queue_wait_nanos,
                &mut self.metrics.min_queue_wait_nanos,
                self.has_queue_wait,
                queue_wait,
            );
            self.has_queue_wait = true;

            if let (Some(last_poll_end_time), Some(schedule_time)) =
                (self.last_poll_end_time, profile.schedule_time)
            {
                Self::record_duration(
                    &mut self.metrics.total_wake_wait_nanos,
                    &mut self.metrics.max_wake_wait_nanos,
                    &mut self.metrics.min_wake_wait_nanos,
                    self.has_wake_wait,
                    schedule_time.saturating_duration_since(last_poll_end_time),
                );
                self.has_wake_wait = true;
            }
        }

        if let Some(waited_task_slices) = profile.fair_queue_waited_task_slices {
            Self::record_value(
                &mut self.metrics.total_fair_queue_waited_task_slices,
                &mut self.metrics.max_fair_queue_waited_task_slices,
                &mut self.metrics.min_fair_queue_waited_task_slices,
                self.metrics.fair_queue_enabled,
                waited_task_slices,
            );
            self.metrics.fair_queue_enabled = true;
        }

        self.metrics.dispatch_count = self.metrics.dispatch_count.saturating_add(1);
        self.last_dispatch_id = Some(dispatch_id);
    }

    fn record_poll_end(&mut self, task_finished: bool, profile: &TaskPollProfile) {
        let poll_end_time = StdInstant::now();
        let has_poll = self.metrics.poll_count > 0;
        Self::record_duration(
            &mut self.metrics.poll_cpu_nanos,
            &mut self.metrics.max_poll_cpu_nanos,
            &mut self.metrics.min_poll_cpu_nanos,
            has_poll,
            profile.cpu_time,
        );
        Self::record_duration(
            &mut self.metrics.poll_wall_nanos,
            &mut self.metrics.max_poll_wall_nanos,
            &mut self.metrics.min_poll_wall_nanos,
            has_poll,
            profile.wall_time,
        );
        self.metrics.poll_count = self.metrics.poll_count.saturating_add(1);
        self.last_poll_end_time = Some(poll_end_time);

        if task_finished {
            self.metrics.total_wall_nanos = poll_end_time
                .saturating_duration_since(self.task_start_time)
                .as_nanos()
                .min(u64::MAX as u128) as u64;
        }
    }
}

#[cfg(feature = "testexport")]
pub mod test_util {
    use std::sync::Arc;

    use parking_lot::RwLock;
    use tracker::{PoolTaskMetrics, RequestInfo};

    type PoolTaskMetricsSyncObserver =
        Arc<dyn Fn(u64, u64, &RequestInfo, &PoolTaskMetrics) + Send + Sync>;
    type SpawnIdMock = Arc<dyn Fn(&[u8], u64) -> Option<u64> + Send + Sync>;

    lazy_static::lazy_static! {
        static ref POOL_TASK_METRICS_SYNC_OBSERVER: RwLock<Option<PoolTaskMetricsSyncObserver>> =
            RwLock::new(None);
        static ref SPAWN_TASK_ID_MOCK: RwLock<Option<SpawnIdMock>> = RwLock::new(None);
    }

    /// Keeps a process-wide read-pool metrics sync observer installed until
    /// the guard is dropped.
    pub struct PoolTaskMetricsSyncObserverGuard;

    impl Drop for PoolTaskMetricsSyncObserverGuard {
        fn drop(&mut self) {
            *POOL_TASK_METRICS_SYNC_OBSERVER.write() = None;
        }
    }

    /// Installs an observer at the boundary where one read-pool task copies
    /// its metrics into the tracker selected by the current TLS token.
    ///
    /// `spawn_task_id` belongs to the read-pool future and is captured before
    /// dispatch. `tracker_sync_sequence` identifies a publication to a tracker
    /// from a shared task, such as one child in `batch_get_command`.
    /// `request_info` belongs to the destination tracker, so tests can detect
    /// both selecting the wrong tracker and syncing the right task metrics into
    /// the wrong tracker.
    pub fn set_pool_task_metrics_sync_observer<F>(observer: F) -> PoolTaskMetricsSyncObserverGuard
    where
        F: Fn(u64, u64, &RequestInfo, &PoolTaskMetrics) + Send + Sync + 'static,
    {
        let mut current = POOL_TASK_METRICS_SYNC_OBSERVER.write();
        assert!(
            current.is_none(),
            "read-pool task metrics sync observer already set"
        );
        *current = Some(Arc::new(observer));
        PoolTaskMetricsSyncObserverGuard
    }

    pub(super) fn observe_pool_task_metrics_sync(
        spawn_task_id: u64,
        tracker_sync_sequence: u64,
        request_info: &RequestInfo,
        metrics: &PoolTaskMetrics,
    ) {
        let observer = POOL_TASK_METRICS_SYNC_OBSERVER.read().clone();
        if let Some(observer) = observer {
            observer(spawn_task_id, tracker_sync_sequence, request_info, metrics);
        }
    }

    /// Keeps a process-wide read-pool spawn ID mock installed until the guard
    /// is dropped.
    pub struct SpawnIdMockGuard;

    impl Drop for SpawnIdMockGuard {
        fn drop(&mut self) {
            *SPAWN_TASK_ID_MOCK.write() = None;
        }
    }

    /// Installs a mock that can provide the ID passed to a read-pool spawn.
    ///
    /// The callback receives a request's resource-group tag and region ID.
    /// Returning `None` leaves that request's original spawn ID unchanged.
    pub fn set_spawn_task_id_mock<F>(mock_spawn_task_id: F) -> SpawnIdMockGuard
    where
        F: Fn(&[u8], u64) -> Option<u64> + Send + Sync + 'static,
    {
        let mut current = SPAWN_TASK_ID_MOCK.write();
        assert!(current.is_none(), "read-pool spawn ID mock already set");
        *current = Some(Arc::new(mock_spawn_task_id));
        SpawnIdMockGuard
    }

    pub(crate) fn mock_spawn_task_id(resource_group_tag: &[u8], region_id: u64) -> Option<u64> {
        let mock_spawn_task_id = SPAWN_TASK_ID_MOCK.read().clone();
        mock_spawn_task_id
            .and_then(|mock_spawn_task_id| mock_spawn_task_id(resource_group_tag, region_id))
    }
}

#[derive(Clone)]
struct ActiveTrackerMetricsObserverState {
    state: TrackerMetricsObserverState,
    poll_cpu_timer: ThreadTime,
    poll_wall_timer: StdInstant,
}

impl ActiveTrackerMetricsObserverState {
    fn new(
        state: TrackerMetricsObserverState,
        poll_cpu_timer: ThreadTime,
        poll_wall_timer: StdInstant,
    ) -> Self {
        Self {
            state,
            poll_cpu_timer,
            poll_wall_timer,
        }
    }

    fn record_task_finished(&mut self) {
        self.state.record_poll_end(
            true,
            &TaskPollProfile {
                cpu_time: self.poll_cpu_timer.elapsed(),
                wall_time: self.poll_wall_timer.elapsed(),
            },
        );
    }
}

thread_local! {
    static TLS_TRACKER_METRICS_OBSERVER_STATE: RefCell<Option<ActiveTrackerMetricsObserverState>> =
        const { RefCell::new(None) };
}

struct TrackerMetricsObserver {
    enable_req_metrics: bool,
    extra_observer: Option<Box<dyn YatpPollEventTracker + Send>>,
    state: Option<TrackerMetricsObserverState>,
}

impl Default for TrackerMetricsObserver {
    fn default() -> Self {
        Self {
            enable_req_metrics: false,
            extra_observer: None,
            state: Some(TrackerMetricsObserverState::default()),
        }
    }
}

impl TrackerMetricsObserver {
    #[inline]
    fn enabled(&self) -> bool {
        self.enable_req_metrics || self.extra_observer.is_some()
    }
}

impl YatpPollEventTracker for TrackerMetricsObserver {
    fn on_poll_start(&mut self, profile: &TaskDispatchProfile) {
        if let Some(observer) = self.extra_observer.as_mut() {
            observer.on_poll_start(profile);
        }

        if !self.enable_req_metrics {
            return;
        }

        if let Some(mut state) = self.state.take() {
            state.record_poll_start(profile);
            TLS_TRACKER_METRICS_OBSERVER_STATE.with(|current| {
                let old = current
                    .borrow_mut()
                    .replace(ActiveTrackerMetricsObserverState::new(
                        state,
                        profile.poll_start_thread_time,
                        profile.poll_start_time,
                    ));
                debug_assert!(
                    old.is_none(),
                    "read-pool tracker metrics poll must not be nested"
                );
            });
        }
    }

    fn on_poll_end(&mut self, task_finished: bool, profile: &TaskPollProfile) {
        if let Some(observer) = self.extra_observer.as_mut() {
            observer.on_poll_end(task_finished, profile);
        }

        if !self.enable_req_metrics {
            return;
        }

        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|current| {
            if let Some(mut active) = current.borrow_mut().take() {
                active.state.record_poll_end(task_finished, profile);
                if task_finished {
                    // This is a fallback for tasks that do not publish their result through a
                    // path that calls `sync_read_pool_metrics_to_current_tracker(true)`.
                    active.state.sync_to_current_tracker();
                } else {
                    self.state = Some(active.state);
                }
            }
        });
    }
}

/// Synchronizes request-level read-pool metrics to the tracker selected by the
/// current TLS tracker token.
///
/// This function is intended to be called from inside the observed read-pool
/// future immediately before making a result externally visible. The snapshot
/// includes the current poll up to this call. When `done` is false, the
/// observer continues collecting and later snapshots remain cumulative. When
/// `done` is true, the observer state is consumed and later tracker-metrics
/// callbacks become no-ops.
pub fn sync_read_pool_metrics_to_current_tracker(done: bool) {
    TLS_TRACKER_METRICS_OBSERVER_STATE.with(|current| {
        let mut active = if done {
            current.borrow_mut().take()
        } else {
            current.borrow().as_ref().cloned()
        };
        if let Some(active) = active.as_mut() {
            active.record_task_finished();
            active.state.sync_to_current_tracker();
        }
    });
}

/// Execution time measured around one invocation of a task future's `poll`.
pub struct TaskPollProfile {
    /// Thread CPU time consumed by this poll.
    pub cpu_time: Duration,
    /// Wall time spent inside this poll.
    pub wall_time: Duration,
}

#[pin_project(PinnedDrop)]
pub struct ObserveReadPoolTaskFuture<F> {
    #[pin]
    delegate: F,
    keyspace_id: u32,
    fairness_recorder: Option<KeyspaceFairnessRecorder>,
    poll_cpu_recorder: Option<Arc<dyn ReadPoolTaskPollCpuRecorder>>,
    keyspace_inner_metrics: Option<Arc<KeyspaceInnerMetricsHandle>>,
    req_tracker_observer: TrackerMetricsObserver,
}

impl<F> ObserveReadPoolTaskFuture<F> {
    #[inline]
    fn new(delegate: F, keyspace_id: u32) -> Self {
        Self {
            delegate,
            keyspace_id,
            fairness_recorder: None,
            poll_cpu_recorder: None,
            keyspace_inner_metrics: None,
            req_tracker_observer: TrackerMetricsObserver::default(),
        }
    }

    #[inline]
    fn set_fairness_recorder(
        &mut self,
        fairness_recorder: Option<Arc<dyn FairnessRecorder<u32>>>,
        use_wall_time: bool,
        slice_precharge_us: u64,
        max_time_slice_us: u64,
    ) {
        self.fairness_recorder = fairness_recorder.map(|recorder| KeyspaceFairnessRecorder {
            keyspace_id: self.keyspace_id,
            recorder,
            use_wall_time,
            slice_precharge_us,
            max_time_slice: Duration::from_micros(max_time_slice_us),
        });
    }

    fn set_poll_cpu_recorder(
        &mut self,
        poll_cpu_recorder: Option<Arc<dyn ReadPoolTaskPollCpuRecorder>>,
    ) {
        self.poll_cpu_recorder = poll_cpu_recorder
    }

    #[inline]
    fn set_keyspace_inner_metrics(&mut self, metrics: Option<Arc<KeyspaceInnerMetricsHandle>>) {
        self.keyspace_inner_metrics = metrics;
    }

    #[inline]
    fn set_enable_req_tracker_metrics(&mut self, enabled: bool) {
        self.req_tracker_observer.enable_req_metrics = enabled;
    }

    #[inline]
    #[cfg(feature = "testexport")]
    fn set_req_tracker_spawn_task_id(&mut self, spawn_task_id: u64) {
        self.req_tracker_observer
            .state
            .as_mut()
            .expect("read-pool tracker metrics state must exist before polling")
            .spawn_task_id = spawn_task_id;
    }

    #[inline]
    fn set_extra_tracker_observer(
        &mut self,
        observer: Option<Box<dyn YatpPollEventTracker + Send>>,
    ) {
        self.req_tracker_observer.extra_observer = observer;
    }

    #[inline]
    fn poll_delegate<D: Future>(
        delegate: Pin<&mut D>,
        cx: &mut Context<'_>,
        poll_started_at: StdInstant,
        max_time_slice: Option<Duration>,
    ) -> Poll<D::Output> {
        let reschedule_context_guard = max_time_slice
            .map(|max_time_slice| PollRescheduleContextGuard::new(poll_started_at, max_time_slice));
        let result = delegate.poll(cx);
        drop(reschedule_context_guard);
        result
    }

    #[inline]
    fn poll_with_profile<D: Future>(
        delegate: Pin<&mut D>,
        cx: &mut Context<'_>,
        req_tracker_observer: &mut TrackerMetricsObserver,
        max_time_slice: Option<Duration>,
        on_start: impl FnOnce(&mut TrackerMetricsObserver, ThreadTime, StdInstant),
        on_end: impl FnOnce(&mut TrackerMetricsObserver, bool, TaskPollProfile),
    ) -> Poll<D::Output> {
        let cpu_timer = ThreadTime::now();
        let wall_timer = StdInstant::now();
        on_start(req_tracker_observer, cpu_timer, wall_timer);
        let result = Self::poll_delegate(delegate, cx, wall_timer, max_time_slice);
        on_end(
            req_tracker_observer,
            result.is_ready(),
            TaskPollProfile {
                cpu_time: cpu_timer.elapsed(),
                wall_time: wall_timer.elapsed(),
            },
        );
        result
    }
}

impl<F: Future> Future for ObserveReadPoolTaskFuture<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        if this.fairness_recorder.is_none()
            && this.poll_cpu_recorder.is_none()
            && this.keyspace_inner_metrics.is_none()
            && !this.req_tracker_observer.enabled()
        {
            // No observer means keyspace fairness is disabled. Preserve the
            // legacy poll path without installing the new reschedule context.
            return this.delegate.poll(cx);
        }

        // Charge this poll to the keyspace's virtual time. The resulting charge
        // affects which keyspace the fair queue dispatches next.
        let fairness_recorder = this.fairness_recorder.as_ref();
        // Keep the compatibility policy at the read-pool boundary: the new
        // wall-time reschedule behavior is enabled only with keyspace
        // fairness. Read-path helpers only inspect whether the context exists.
        let max_time_slice = fairness_recorder.map(|recorder| recorder.max_time_slice);

        // Accumulate actual poll CPU time for CPU-token reconciliation and
        // keyspace resource-usage publishing when the task finishes.
        let poll_cpu_recorder = this.poll_cpu_recorder.as_ref();

        // Update per-keyspace task-slice CPU distribution, wall-time counter,
        // and the activity timestamp used by KeyspaceMetrics for monitoring.
        let keyspace_inner_metrics = this.keyspace_inner_metrics.as_ref();

        // Measure the delegate poll once, then feed the same execution interval
        // to request diagnostics, for example, slow query.
        Self::poll_with_profile(
            this.delegate.as_mut(),
            cx,
            this.req_tracker_observer,
            max_time_slice,
            |req_tracker_observer, cpu_timer, wall_timer| {
                fairness_recorder.inspect(|r| r.precharge());
                // Capture the request's queue, wake, and poll breakdown in TLS-backed
                // RequestMetrics for per-request diagnostics.
                if req_tracker_observer.enabled() {
                    let profile = YatpDispatchContext::with_tls(|ctx| {
                        TaskDispatchProfile::new(ctx, cpu_timer, wall_timer)
                    });
                    req_tracker_observer.on_poll_start(&profile);
                }
            },
            move |req_tracker_observer, finished, profile| {
                fairness_recorder.inspect(|r| r.observe_poll_profile(&profile));
                poll_cpu_recorder.inspect(|r| r.observe_poll_cpu_time(profile.cpu_time));
                keyspace_inner_metrics.inspect(|h| {
                    h.observe_task_slice_cpu_duration(profile.cpu_time);
                    h.observe_task_wall_time(profile.wall_time);
                });
                if req_tracker_observer.enabled() {
                    req_tracker_observer.on_poll_end(finished, &profile);
                }
            },
        )
    }
}

#[pin_project::pinned_drop]
impl<F> PinnedDrop for ObserveReadPoolTaskFuture<F> {
    fn drop(self: Pin<&mut Self>) {
        let this = self.project();
        if let Some(recorder) = &this.poll_cpu_recorder {
            recorder.finish();
        }
    }
}

impl ReadPool {
    pub fn handle(&self) -> ReadPoolHandle {
        match self {
            ReadPool::FuturePools {
                read_pool_high,
                read_pool_normal,
                read_pool_low,
            } => ReadPoolHandle::FuturePools {
                read_pool_high: read_pool_high.clone(),
                read_pool_normal: read_pool_normal.clone(),
                read_pool_low: read_pool_low.clone(),
            },
            ReadPool::Yatp {
                pool,
                keyspace_fair_queue,
                fair_global_cpu_admission_controller,
                fairness_use_wall_time,
                fairness_slice_precharge_us,
                fairness_max_time_slice_us,
                keyspace_metrics,
                running_tasks,
                running_threads,
                max_tasks,
                keyspace_task_limiter,
                pool_size,
            } => ReadPoolHandle::Yatp {
                remote: pool.remote().clone(),
                keyspace_fair_queue: keyspace_fair_queue.clone(),
                fair_global_cpu_admission_enabled: fair_global_cpu_admission_controller.is_some(),
                fair_global_cpu_admission_controller: fair_global_cpu_admission_controller.clone(),
                fairness_use_wall_time: fairness_use_wall_time.clone(),
                fairness_slice_precharge_us: fairness_slice_precharge_us.clone(),
                fairness_max_time_slice_us: fairness_max_time_slice_us.clone(),
                keyspace_metrics: keyspace_metrics.clone(),
                running_tasks: running_tasks.clone(),
                running_threads: running_threads.clone(),
                max_tasks: max_tasks.clone(),
                keyspace_task_limiter: keyspace_task_limiter.clone(),
                pool_size: pool_size.clone(),
            },
            ReadPool::Tokio {
                runtime,
                running_tasks,
                max_tasks,
                pool_size,
                task_monitor,
                pool_name,
            } => ReadPoolHandle::Tokio {
                // here, the runtime must not be stop or dropped, so it safe
                // to unwrap.
                handle: runtime.handle().unwrap(),
                running_tasks: running_tasks.clone(),
                max_tasks: max_tasks.clone(),
                pool_size: pool_size.clone(),
                task_monitor: task_monitor.clone(),
                pool_name: pool_name.clone(),
            },
        }
    }

    pub fn shutdown(self) {
        if let ReadPool::Tokio { mut runtime, .. } = self {
            runtime.shutdown();
        }
    }
}

#[derive(Clone)]
pub enum ReadPoolHandle {
    FuturePools {
        read_pool_high: FuturePool,
        read_pool_normal: FuturePool,
        read_pool_low: FuturePool,
    },
    Yatp {
        remote: Remote<TaskCell>,
        keyspace_fair_queue: Option<Arc<YatpKeyspaceFairQueue>>,
        fair_global_cpu_admission_enabled: bool,
        fair_global_cpu_admission_controller: Option<Arc<FairGlobalCpuAdmissionController>>,
        fairness_use_wall_time: Arc<AtomicBool>,
        fairness_slice_precharge_us: Arc<AtomicU64>,
        fairness_max_time_slice_us: Arc<AtomicU64>,
        keyspace_metrics: Option<Arc<KeyspaceMetrics>>,
        running_tasks: IntGauge,
        running_threads: IntGauge,
        max_tasks: Arc<AtomicUsize>,
        keyspace_task_limiter: Option<KeyspaceTaskLimiter>,
        pool_size: Arc<AtomicUsize>,
    },
    Tokio {
        handle: ScalableTokioHandle,
        running_tasks: IntGauge,
        max_tasks: Arc<AtomicUsize>,
        pool_size: Arc<AtomicUsize>,
        task_monitor: tokio_metrics::TaskMonitor,
        pool_name: String,
    },
}

impl ReadPoolHandle {
    pub fn spawn<F>(
        &self,
        f: F,
        priority: CommandPri,
        keyspace_id: u32,
        task_id: u64,
    ) -> Result<(), ReadPoolError>
    where
        F: Future<Output = ()> + Send + 'static,
    {
        self.spawn_with_options(
            f,
            priority,
            keyspace_id,
            task_id,
            ReadPoolTaskOptions::default(),
        )
    }

    /// Spawns a read-pool task with explicit scheduling options.
    pub fn spawn_with_options<F>(
        &self,
        f: F,
        priority: CommandPri,
        keyspace_id: u32,
        task_id: u64,
        options: ReadPoolTaskOptions,
    ) -> Result<(), ReadPoolError>
    where
        F: Future<Output = ()> + Send + 'static,
    {
        let ReadPoolTaskOptions {
            poll_cpu_recorder,
            fair_cpu_admission,
            poll_tracker_observer,
        } = options;
        let mut f = ObserveReadPoolTaskFuture::new(f, keyspace_id);
        #[cfg(feature = "testexport")]
        f.set_req_tracker_spawn_task_id(task_id);
        f.set_poll_cpu_recorder(poll_cpu_recorder);
        f.set_extra_tracker_observer(poll_tracker_observer);
        match self {
            ReadPoolHandle::FuturePools {
                read_pool_high,
                read_pool_normal,
                read_pool_low,
                ..
            } => {
                let pool = match priority {
                    CommandPri::High => read_pool_high,
                    CommandPri::Normal => read_pool_normal,
                    CommandPri::Low => read_pool_low,
                };

                pool.spawn(f)?;
            }
            ReadPoolHandle::Yatp {
                remote,
                keyspace_fair_queue,
                fair_global_cpu_admission_enabled,
                fairness_use_wall_time,
                fairness_slice_precharge_us,
                fairness_max_time_slice_us,
                keyspace_metrics,
                running_tasks,
                max_tasks,
                keyspace_task_limiter,
                ..
            } => {
                let running_tasks = running_tasks.clone();
                let keyspace_task_permit = match keyspace_task_limiter {
                    Some(limiter) => Some(
                        // Keyspace fairness replaces the old pool-global max
                        // tasks admission check with this per-keyspace cap.
                        // High-priority tasks bypass the per-keyspace cap check,
                        // but still hold a permit so their occupancy backpressures
                        // later normal-priority tasks from the same keyspace.
                        limiter
                            .try_acquire(keyspace_id, priority == CommandPri::High)
                            .map_err(|KeyspaceTaskLimitError::Full| {
                                record_unified_read_pool_reject(
                                    Some(keyspace_id),
                                    READ_POOL_REJECT_REASON_KEYSPACE_FULL,
                                );
                                ReadPoolError::UnifiedReadPoolFull
                            })?,
                    ),
                    None => {
                        // Note that the running task number limit is not strict.
                        // If several tasks are spawned at the same time while the running task
                        // number is close to the limit, they may all pass this check and the
                        // number of running tasks may exceed the limit.
                        if priority != CommandPri::High
                            && running_tasks.get() as usize >= max_tasks.load(Ordering::Relaxed)
                        {
                            record_unified_read_pool_reject(
                                None,
                                READ_POOL_REJECT_REASON_GLOBAL_FULL,
                            );
                            return Err(ReadPoolError::UnifiedReadPoolFull);
                        }
                        None
                    }
                };

                if let Some(keyspace_metrics) = keyspace_metrics.as_ref() {
                    f.set_keyspace_inner_metrics(Some(
                        keyspace_metrics.get_keyspace_metrics(keyspace_id),
                    ));
                }
                running_tasks.inc();
                // With keyspace fairness enabled, this fixed level only affects
                // ordering inside the keyspace selected by the outer fair queue.
                // `High` does not bypass virtual-time ordering across keyspaces.
                let fixed_level = match priority {
                    CommandPri::High => Some(0),
                    CommandPri::Normal => None,
                    CommandPri::Low => Some(2),
                };
                let mut extras = Extras::new_multilevel(task_id, fixed_level);
                let fair_cpu_admission = if let Some(queue) = keyspace_fair_queue {
                    // only to enable request tracker metrics only if the yatp fairness enabled to
                    // reduce the regression for old path
                    f.set_enable_req_tracker_metrics(true);
                    f.set_fairness_recorder(
                        Some(queue.get_fairness_recorder()),
                        fairness_use_wall_time.load(Ordering::Relaxed),
                        fairness_slice_precharge_us.load(Ordering::Relaxed),
                        fairness_max_time_slice_us.load(Ordering::Relaxed),
                    );
                    if *fair_global_cpu_admission_enabled {
                        fair_cpu_admission
                    } else {
                        None
                    }
                } else {
                    None
                };
                set_yatp_task_metadata_if_fair(
                    &mut extras,
                    keyspace_fair_queue.is_some(),
                    keyspace_id,
                    fair_cpu_admission,
                );
                let task_cell = TaskCell::new(
                    TrackedFuture::new(async move {
                        let _keyspace_task_permit = keyspace_task_permit;
                        tikv_util::init_task_local(f).await;
                        running_tasks.dec();
                    }),
                    extras,
                );
                remote.spawn(task_cell);
            }
            ReadPoolHandle::Tokio {
                handle,
                running_tasks,
                max_tasks,
                task_monitor,
                ..
            } => {
                let running_tasks = running_tasks.clone();
                // Note that the running task number limit is not strict.
                // If several tasks are spawned at the same time while the running task number
                // is close to the limit, they may all pass this check and the number of running
                // tasks may exceed the limit.
                if priority != CommandPri::High
                    && running_tasks.get() as usize >= max_tasks.load(Ordering::Relaxed)
                {
                    record_unified_read_pool_reject(None, READ_POOL_REJECT_REASON_GLOBAL_FULL);
                    return Err(ReadPoolError::UnifiedReadPoolFull);
                }
                running_tasks.inc();
                let tracked = TrackedFuture::new(async move {
                    tikv_util::init_task_local(f).await;
                    try_tick_reporter_ticker_if_due();
                    running_tasks.dec();
                });
                // Instrument with tokio-metrics to track schedule wait time
                let instrumented = task_monitor.instrument(tracked);
                handle.spawn(instrumented);
            }
        }
        Ok(())
    }

    pub fn spawn_handle<F, T>(
        &self,
        f: F,
        priority: CommandPri,
        keyspace_id: u32,
        task_id: u64,
    ) -> impl Future<Output = Result<T, ReadPoolError>>
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        self.spawn_handle_with_options(
            f,
            priority,
            keyspace_id,
            task_id,
            ReadPoolTaskOptions::default(),
        )
    }

    /// Spawns a read-pool task with explicit scheduling options and returns a
    /// future for its result.
    pub fn spawn_handle_with_options<F, T>(
        &self,
        f: F,
        priority: CommandPri,
        keyspace_id: u32,
        task_id: u64,
        options: ReadPoolTaskOptions,
    ) -> impl Future<Output = Result<T, ReadPoolError>>
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        let (tx, rx) = oneshot::channel::<T>();
        let res = self.spawn_with_options(
            async move {
                let res = f.await;
                sync_read_pool_metrics_to_current_tracker(true);
                let _ = tx.send(res);
            },
            priority,
            keyspace_id,
            task_id,
            options,
        );
        async move {
            res?;
            rx.map_err(ReadPoolError::from).await
        }
    }

    pub fn get_normal_pool_size(&self) -> usize {
        match self {
            ReadPoolHandle::FuturePools {
                read_pool_normal, ..
            } => read_pool_normal.get_pool_size(),
            ReadPoolHandle::Yatp { pool_size, .. } => pool_size.load(Ordering::Relaxed),
            ReadPoolHandle::Tokio { pool_size, .. } => pool_size.load(Ordering::Relaxed),
        }
    }

    pub fn get_queue_size_per_worker(&self) -> usize {
        match self {
            ReadPoolHandle::FuturePools {
                read_pool_normal, ..
            } => read_pool_normal.get_running_task_count() / read_pool_normal.get_pool_size(),
            ReadPoolHandle::Yatp {
                running_tasks,
                pool_size,
                ..
            } => running_tasks.get() as usize / pool_size.load(Ordering::Relaxed),
            ReadPoolHandle::Tokio {
                running_tasks,
                pool_size,
                ..
            } => running_tasks.get() as usize / pool_size.load(Ordering::Relaxed),
        }
    }

    /// Whether this handle enforces global CPU tokens in the keyspace fair
    /// read-pool wrapper.
    pub fn uses_fair_global_cpu_gate(&self) -> bool {
        matches!(
            self,
            ReadPoolHandle::Yatp {
                keyspace_fair_queue: Some(_),
                fair_global_cpu_admission_enabled: true,
                ..
            }
        )
    }

    /// Whether this handle schedules work through the keyspace fair queue.
    pub fn uses_keyspace_fair_queue(&self) -> bool {
        matches!(
            self,
            ReadPoolHandle::Yatp {
                keyspace_fair_queue: Some(_),
                ..
            }
        )
    }

    /// Whether coprocessor requests should retain the pre-fairness global
    /// concurrency limit. The keyspace fair queue uses its keyspace-aware task
    /// limiter instead, so the process-global semaphore must not sit in front
    /// of fair scheduling.
    pub(crate) fn uses_coprocessor_concurrency_limiter(&self) -> bool {
        matches!(
            self,
            ReadPoolHandle::Yatp {
                keyspace_fair_queue: None,
                ..
            }
        )
    }

    pub fn scale_pool_size(&mut self, max_thread_count: usize) {
        match self {
            ReadPoolHandle::FuturePools { .. } => {
                unreachable!()
            }
            ReadPoolHandle::Yatp {
                remote,
                running_threads,
                max_tasks,
                pool_size,
                ..
            } => {
                remote.scale_workers(max_thread_count);
                let max_task_count = max_tasks
                    .load(Ordering::Relaxed)
                    .saturating_div(pool_size.load(Ordering::Relaxed))
                    .saturating_mul(max_thread_count);
                max_tasks.store(max_task_count, Ordering::Relaxed);
                running_threads.set(max_thread_count as i64);
                pool_size.store(max_thread_count, Ordering::Relaxed);
            }
            ReadPoolHandle::Tokio {
                handle,
                max_tasks,
                pool_size,
                ..
            } => {
                handle.scale_pool_size(max_thread_count);
                let max_task_count = max_tasks
                    .load(Ordering::Relaxed)
                    .saturating_div(pool_size.load(Ordering::Relaxed))
                    .saturating_mul(max_thread_count);
                max_tasks.store(max_task_count, Ordering::Relaxed);
                pool_size.store(max_thread_count, Ordering::Relaxed);
            }
        }
    }

    fn set_fairness_use_wall_time(&self, use_wall_time: bool) {
        if let ReadPoolHandle::Yatp {
            fairness_use_wall_time,
            ..
        } = self
        {
            fairness_use_wall_time.store(use_wall_time, Ordering::Relaxed);
        }
    }

    fn set_fairness_slice_precharge_us(&self, slice_precharge_us: u64) {
        if let ReadPoolHandle::Yatp {
            fairness_slice_precharge_us,
            ..
        } = self
        {
            fairness_slice_precharge_us.store(slice_precharge_us, Ordering::Relaxed);
        }
    }

    fn set_fairness_max_read_cpu_ratio(&self, max_read_cpu_ratio: f64) {
        if let ReadPoolHandle::Yatp {
            fair_global_cpu_admission_controller: Some(controller),
            ..
        } = self
        {
            controller.update_max_read_cpu_ratio(max_read_cpu_ratio);
        }
    }

    fn set_fairness_max_time_slice_us(&self, max_time_slice_us: u64) {
        if let ReadPoolHandle::Yatp {
            fairness_max_time_slice_us,
            ..
        } = self
        {
            fairness_max_time_slice_us.store(max_time_slice_us, Ordering::Relaxed);
        }
    }

    fn maintain_keyspaces_if_needed(&self) {
        if let ReadPoolHandle::Yatp {
            keyspace_fair_queue,
            keyspace_metrics,
            ..
        } = self
        {
            if let Some(metrics) = keyspace_metrics {
                metrics.clear_inactive_keyspace_metrics_if_needed();
            }
            if let Some(keyspace_fair_queue) = keyspace_fair_queue {
                keyspace_fair_queue.maintain_queue_if_needed();
            }
        }
    }

    fn set_max_tasks_per_keyspace(&self, max_tasks_per_keyspace: usize) {
        if let ReadPoolHandle::Yatp {
            keyspace_task_limiter: Some(limiter),
            ..
        } = self
        {
            limiter.set_max_tasks_per_keyspace(max_tasks_per_keyspace);
        }
    }
}

#[derive(Clone)]
pub struct ReporterTicker<R: FlowStatsReporter> {
    reporter: R,
}

impl<R: FlowStatsReporter> PoolTicker for ReporterTicker<R> {
    fn on_tick(&mut self) {
        self.flush_metrics_on_tick();
    }
}

impl<R: FlowStatsReporter> ReporterTicker<R> {
    fn flush_metrics_on_tick(&mut self) {
        crate::storage::metrics::tls_flush(&self.reporter);
        crate::coprocessor::metrics::tls_flush(&self.reporter);
    }
}

#[cfg(test)]
fn get_unified_read_pool_name() -> String {
    use std::sync::atomic::AtomicU64;

    static COUNTER: AtomicU64 = AtomicU64::new(0);
    format!(
        "unified-read-pool-test-{}",
        COUNTER.fetch_add(1, Ordering::Relaxed)
    )
}

#[cfg(not(test))]
fn get_unified_read_pool_name() -> String {
    "unified-read-pool".to_string()
}

fn create_keyspace_metrics_from_tikv_prom(pool_name: &str) -> Arc<KeyspaceMetrics> {
    let create_pool_name = pool_name.to_owned();
    let remove_pool_name = pool_name.to_owned();
    KeyspaceMetrics::new(
        metrics::YATP_POOL_KEYSPACE_FAIR_GROUP_COUNT_GAUGE_VEC.with_label_values(&[pool_name]),
        move |keyspace_id| {
            let keyspace_id = keyspace_id.to_string();
            let labels = [create_pool_name.as_str(), keyspace_id.as_str()];
            KeyspaceInnerMetrics::new(
                metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_LENGTH_HISTOGRAM_VEC
                    .with_label_values(&labels),
                metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_WAITS_HISTOGRAM_VEC
                    .with_label_values(&labels),
                metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_WAIT_DURATION_HISTOGRAM_VEC
                    .with_label_values(&labels),
                metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_SCHEDULE_INTERVAL_HISTOGRAM_VEC
                    .with_label_values(&labels),
                metrics::YATP_POOL_KEYSPACE_TASK_SLICE_CPU_DURATION_HISTOGRAM_VEC
                    .with_label_values(&labels),
                metrics::YATP_POOL_KEYSPACE_TASK_WALL_TIME_COUNTER_VEC.with_label_values(&labels),
            )
        },
        move |keyspace_id| {
            let keyspace_id = keyspace_id.to_string();
            let labels = [remove_pool_name.as_str(), keyspace_id.as_str()];
            let _ = metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_LENGTH_HISTOGRAM_VEC
                .remove_label_values(&labels);
            let _ = metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_WAITS_HISTOGRAM_VEC
                .remove_label_values(&labels);
            let _ = metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_WAIT_DURATION_HISTOGRAM_VEC
                .remove_label_values(&labels);
            let _ =
                metrics::YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_SCHEDULE_INTERVAL_HISTOGRAM_VEC
                    .remove_label_values(&labels);
            let _ = metrics::YATP_POOL_KEYSPACE_TASK_SLICE_CPU_DURATION_HISTOGRAM_VEC
                .remove_label_values(&labels);
            let _ =
                metrics::YATP_POOL_KEYSPACE_TASK_WALL_TIME_COUNTER_VEC.remove_label_values(&labels);
        },
    )
}

pub fn build_yatp_read_pool<E: Engine, R: FlowStatsReporter>(
    config: &UnifiedReadPoolConfig,
    reporter: R,
    engine: E,
    read_limiter: Option<ReadLimiter>,
) -> ReadPool {
    let unified_read_pool_name = get_unified_read_pool_name();
    let mut builder = YatpPoolBuilder::new(ReporterTicker { reporter });
    let raftkv = Arc::new(Mutex::new(engine));
    let keyspace_metrics = config
        .enable_yatp_keyspace_metrics
        .then(|| create_keyspace_metrics_from_tikv_prom(&unified_read_pool_name));
    builder
        .name_prefix(&unified_read_pool_name)
        .stack_size(config.stack_size.0 as usize)
        .thread_count(
            config.min_thread_count,
            config.max_thread_count,
            std::cmp::max(
                UNIFIED_READPOOL_MIN_CONCURRENCY,
                SysQuota::cpu_cores_quota() as usize,
            ),
        )
        .keyspace_metrics(keyspace_metrics.clone())
        .after_start(move || {
            let engine = raftkv.lock().unwrap().clone();
            set_tls_engine(engine);
            set_io_type(IoType::ForegroundRead);
        })
        .before_stop(|| unsafe {
            destroy_tls_engine::<E>();
        });

    // TODO(#5480): this controller is captured when the read pool is built. If
    // CPU throttle is enabled dynamically after boot, the fair global CPU gate
    // requires a read-pool restart/rebuild before it becomes active.
    let fair_global_cpu_admission_controller = if config.keyspace_fairness.enabled {
        read_limiter.and_then(|limiter| {
            let manager = limiter.get_cpu_throttle_manager()?;
            let controller = Arc::new(FairGlobalCpuAdmissionController::new(manager, limiter));
            controller.update_max_read_cpu_ratio(config.keyspace_fairness.max_read_cpu_ratio);
            Some(controller)
        })
    } else {
        None
    };
    let (pool, keyspace_fair_queue) = if config.keyspace_fairness.enabled {
        let controller = fair_global_cpu_admission_controller
            .as_ref()
            .map(|controller| controller.clone() as Arc<dyn FairQueueAdmissionController>);
        let (pool, queue) = builder.build_keyspace_fair_pool(controller);
        (pool, Some(queue))
    } else {
        (builder.build_multi_level_pool(), None)
    };
    let max_tasks = config
        .max_tasks_per_worker
        .saturating_mul(config.max_thread_count);
    ReadPool::Yatp {
        pool,
        keyspace_fair_queue,
        fair_global_cpu_admission_controller,
        fairness_use_wall_time: Arc::new(AtomicBool::new(config.keyspace_fairness.use_wall_time)),
        fairness_slice_precharge_us: Arc::new(AtomicU64::new(
            config.keyspace_fairness.slice_precharge_us,
        )),
        fairness_max_time_slice_us: Arc::new(AtomicU64::new(
            config.keyspace_fairness.max_time_slice_us,
        )),
        keyspace_metrics,
        running_tasks: UNIFIED_READ_POOL_RUNNING_TASKS
            .with_label_values(&[&unified_read_pool_name]),
        running_threads: UNIFIED_READ_POOL_RUNNING_THREADS
            .with_label_values(&[&unified_read_pool_name]),
        max_tasks: Arc::new(AtomicUsize::new(max_tasks)),
        keyspace_task_limiter: if config.keyspace_fairness.enabled {
            Some(KeyspaceTaskLimiter::new(
                config.keyspace_fairness.max_tasks_per_keyspace,
            ))
        } else {
            None
        },
        pool_size: Arc::new(AtomicUsize::new(config.max_thread_count)),
    }
}

pub fn build_tokio_pool<E: Engine, R: FlowStatsReporter>(
    config: &UnifiedReadPoolConfig,
    reporter: R,
    engine: E,
) -> ReadPool {
    let unified_read_pool_name = get_unified_read_pool_name();
    let thread_name_prefix = unified_read_pool_name.clone();
    let ticker = yatp_pool::TickerWrapper::new(ReporterTicker { reporter });
    let raftkv = Arc::new(Mutex::new(engine));
    let props = tikv_util::thread_group::current_properties();
    // The config thread count can exceed cpu quota when the cpu quota is low.
    let total_thread_count = std::cmp::max(
        SysQuota::cpu_cores_quota().ceil() as usize,
        UNIFIED_READPOOL_MIN_CONCURRENCY,
    );
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .thread_name_fn(move || {
            static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
            let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
            format!("{}-{}", thread_name_prefix, id)
        })
        .worker_threads(total_thread_count)
        .after_start_wrapper(move || {
            let engine = raftkv.lock().unwrap().clone();
            set_tls_engine(engine);
            set_reporter_ticker(ticker.clone());
            tikv_util::thread_group::set_properties(props.clone());
            tikv_util::sys::thread::add_thread_name_to_map();
        })
        .before_stop_wrapper(|| unsafe {
            destroy_tls_engine::<E>();
            destroy_reporter_ticker();
        })
        .on_thread_park(move || {
            try_tick_reporter_ticker();
        })
        .enable_all()
        .build()
        .unwrap();

    // Create TaskMonitor for schedule wait time metrics
    // TODO: Consider using separate TaskMonitors for different task types
    // (e.g., storage read, coprocessor, write cmds) for finer-grained debugging.
    // See tokio-metrics best practices: distinct monitors per task kind.
    let task_monitor = tokio_metrics::TaskMonitor::new();

    // Spawn background task to collect and export metrics
    let metrics_task_monitor = task_monitor.clone();
    let metrics_pool_name = unified_read_pool_name.clone();
    runtime.spawn(async move {
        let mut intervals = metrics_task_monitor.intervals();
        loop {
            tokio::time::sleep(Duration::from_secs(1)).await;
            if let Some(metrics) = intervals.next() {
                TOKIO_POOL_MEAN_FIRST_POLL_DELAY
                    .with_label_values(&[&metrics_pool_name])
                    .set(metrics.mean_first_poll_delay().as_secs_f64());
                TOKIO_POOL_FIRST_POLL_COUNT
                    .with_label_values(&[&metrics_pool_name])
                    .inc_by(metrics.first_poll_count);
                TOKIO_POOL_MEAN_IDLE_DURATION
                    .with_label_values(&[&metrics_pool_name])
                    .set(metrics.mean_idle_duration().as_secs_f64());
                TOKIO_POOL_MEAN_SCHEDULED_DURATION
                    .with_label_values(&[&metrics_pool_name])
                    .set(metrics.mean_scheduled_duration().as_secs_f64());
                TOKIO_POOL_MEAN_POLL_DURATION
                    .with_label_values(&[&metrics_pool_name])
                    .set(metrics.mean_poll_duration().as_secs_f64());
            }
        }
    });

    let scalable_runtime =
        ScalableTokioRuntime::new(runtime, total_thread_count, config.max_thread_count);
    let max_tasks = config
        .max_tasks_per_worker
        .saturating_mul(config.max_thread_count);
    ReadPool::Tokio {
        runtime: scalable_runtime,
        running_tasks: UNIFIED_READ_POOL_RUNNING_TASKS
            .with_label_values(&[&unified_read_pool_name]),
        max_tasks: Arc::new(AtomicUsize::new(max_tasks)),
        pool_size: Arc::new(AtomicUsize::new(config.max_thread_count)),
        task_monitor,
        pool_name: unified_read_pool_name,
    }
}

impl From<Vec<FuturePool>> for ReadPool {
    fn from(mut v: Vec<FuturePool>) -> ReadPool {
        assert_eq!(v.len(), 3);
        let read_pool_high = v.remove(2);
        let read_pool_normal = v.remove(1);
        let read_pool_low = v.remove(0);
        ReadPool::FuturePools {
            read_pool_high,
            read_pool_normal,
            read_pool_low,
        }
    }
}

struct ReadPoolCpuTimeTracker {
    yatp_total_time_elapsed: IntCounter,
    // the total time duration of each thread busy with handling tasks. This time also includes
    // the time when the threads are off-cpu, so it might be much higher than the actual cpu time.
    prev_total_task_handling_time_us: u64,
    prev_check_time: Instant,
    prev_cpu_per_sec: f64,
}

impl ReadPoolCpuTimeTracker {
    fn new(pool_name: &str) -> Self {
        let prev_check_time = Instant::now_coarse();
        let yatp_total_time_elapsed = MULTILEVEL_LEVEL_ELAPSED
            .get_metric_with_label_values(&[pool_name, "total"])
            .unwrap();
        let prev_total_task_handling_time_us = yatp_total_time_elapsed.get();
        Self {
            yatp_total_time_elapsed,
            prev_total_task_handling_time_us,
            prev_check_time,
            prev_cpu_per_sec: 0.0,
        }
    }

    fn prev_avg_cpu_used(&mut self) -> f64 {
        let check_time = Instant::now_coarse();
        let duration = check_time.saturating_duration_since(self.prev_check_time);
        // if the check duration is too small, just return the latest cached value.
        if duration < Duration::from_millis(100) {
            return self.prev_cpu_per_sec;
        }
        let total_cpu_time = self.yatp_total_time_elapsed.get();
        let total_cpu_per_sec = (total_cpu_time - self.prev_total_task_handling_time_us) as f64
            / duration.as_micros() as f64;
        self.prev_total_task_handling_time_us = total_cpu_time;
        self.prev_check_time = check_time;
        self.prev_cpu_per_sec = total_cpu_per_sec;
        total_cpu_per_sec
    }
}
struct ReadPoolConfigRunner {
    interval: Duration,
    sender: SyncSender<usize>,
    handle: ReadPoolHandle,
    cpu_time_tracker: ReadPoolCpuTimeTracker,
    process_stats: ProcessStat,
    // configed thread pool size, it's the min thread count to be scale
    core_thread_count: usize,
    // the max thread count can be scaled
    max_thread_count: usize,
    // the current active thread count
    cur_thread_count: usize,
    auto_adjust: bool,
}

impl Runnable for ReadPoolConfigRunner {
    type Task = Task;
    fn run(&mut self, task: Self::Task) {
        match task {
            Task::PoolSize(s) => {
                if s != self.core_thread_count {
                    self.handle.scale_pool_size(s);
                    self.core_thread_count = s;
                    self.cur_thread_count = s;
                    self.notify_pool_size_change(s);
                }
            }
            Task::AutoAdjust(s) => {
                self.auto_adjust = s;
                // when auto adjust is disabled, reset to the config pool size.
                if !s && self.cur_thread_count != self.core_thread_count {
                    self.handle.scale_pool_size(self.core_thread_count);
                    self.cur_thread_count = self.core_thread_count;
                    self.notify_pool_size_change(self.cur_thread_count);
                }
            }
            Task::KeyspaceFairnessUseWallTime(s) => {
                self.handle.set_fairness_use_wall_time(s);
            }
            Task::KeyspaceFairnessSlicePrechargeUs(s) => {
                self.handle.set_fairness_slice_precharge_us(s);
            }
            Task::KeyspaceFairnessMaxTasksPerKeyspace(s) => {
                self.handle.set_max_tasks_per_keyspace(s);
            }
            Task::KeyspaceFairnessMaxReadCpuRatio(s) => {
                self.handle.set_fairness_max_read_cpu_ratio(s);
            }
            Task::KeyspaceFairnessMaxTimeSliceUs(s) => {
                self.handle.set_fairness_max_time_slice_us(s);
            }
        }
    }
}

impl RunnableWithTimer for ReadPoolConfigRunner {
    fn get_interval(&self) -> Duration {
        self.interval
    }

    fn on_timeout(&mut self) {
        self.adjust_pool_size();
        self.handle.maintain_keyspaces_if_needed();
    }
}

impl ReadPoolConfigRunner {
    fn running_tasks(&self) -> i64 {
        match &self.handle {
            ReadPoolHandle::Yatp { running_tasks, .. } => running_tasks.get(),
            _ => unreachable!(),
        }
    }

    // only yatp support auto_adjust thread count.
    fn support_auto_adjust(&self) -> bool {
        matches!(&self.handle, ReadPoolHandle::Yatp { .. })
    }

    fn adjust_pool_size(&mut self) {
        if !(self.auto_adjust && self.support_auto_adjust())
            || (self.cur_thread_count == self.max_thread_count
                && self.core_thread_count == self.max_thread_count)
        {
            return;
        }

        let read_pool_cpu = self.cpu_time_tracker.prev_avg_cpu_used();
        let running_tasks = self.running_tasks();
        let process_cpu = match self.process_stats.cpu_usage() {
            Ok(p) => p,
            Err(e) => {
                warn!("fetch process cpu usage failed"; "err" => ?e);
                return;
            }
        };
        let cpu_quota = SysQuota::cpu_cores_quota();

        // scale out the thread pool size by 1 iff:
        // - current thread count is small than the maximum thread count
        // - process cpu is not overloaded after scaling out one more thread
        // - all read pool threads are busy handling tasks(thread busy time >= 80%)
        // - there are enough tasks waiting in the scheduling queue.
        // scale in the thread pool size by 1 iff:
        // - current thread count is bigger than the configed thread count
        // - the average thread usage percent is under the low water mark(70%)
        // - the running tasks in the scheduling queue is under the threshold
        let new_thread_count = if self.cur_thread_count < self.max_thread_count
            && process_cpu * (self.cur_thread_count as f64 + 1.0) / (self.cur_thread_count as f64)
                < cpu_quota
            && read_pool_cpu > self.cur_thread_count as f64 * READ_POOL_THREAD_HIGH_THRESHOLD
            && running_tasks > self.cur_thread_count as i64 * RUNNING_TASKS_PER_THREAD_THRESHOLD
        {
            self.cur_thread_count + 1
        } else if self.cur_thread_count > self.core_thread_count
            && read_pool_cpu < (self.cur_thread_count - 1) as f64 * READ_POOL_THREAD_LOW_THRESHOLD
            && running_tasks < self.cur_thread_count as i64 * RUNNING_TASKS_PER_THREAD_THRESHOLD
        {
            self.cur_thread_count - 1
        } else {
            self.cur_thread_count
        };

        if new_thread_count != self.cur_thread_count {
            self.handle.scale_pool_size(new_thread_count);
            self.notify_pool_size_change(new_thread_count);
            self.cur_thread_count = new_thread_count;
        }
    }

    fn notify_pool_size_change(&self, new_thread_count: usize) {
        // it's unlikely to send failed.
        if let Err(e) = self.sender.try_send(new_thread_count) {
            warn!("notify read pool thread count change failed"; "err" => ?e);
        }
    }
}

enum Task {
    PoolSize(usize),
    AutoAdjust(bool),
    KeyspaceFairnessUseWallTime(bool),
    KeyspaceFairnessSlicePrechargeUs(u64),
    KeyspaceFairnessMaxTasksPerKeyspace(usize),
    KeyspaceFairnessMaxReadCpuRatio(f64),
    KeyspaceFairnessMaxTimeSliceUs(u64),
}

impl std::fmt::Display for Task {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Task::PoolSize(s) => write!(f, "PoolSize({})", *s),
            Task::AutoAdjust(s) => write!(f, "AutoAdjust({})", *s),
            Task::KeyspaceFairnessUseWallTime(s) => {
                write!(f, "KeyspaceFairnessUseWallTime({})", *s)
            }
            Task::KeyspaceFairnessSlicePrechargeUs(s) => {
                write!(f, "KeyspaceFairnessSlicePrechargeUs({})", *s)
            }
            Task::KeyspaceFairnessMaxTasksPerKeyspace(s) => {
                write!(f, "KeyspaceFairnessMaxTasksPerKeyspace({})", *s)
            }
            Task::KeyspaceFairnessMaxReadCpuRatio(s) => {
                write!(f, "KeyspaceFairnessMaxReadCpuRatio({})", *s)
            }
            Task::KeyspaceFairnessMaxTimeSliceUs(s) => {
                write!(f, "KeyspaceFairnessMaxTimeSliceUs({})", *s)
            }
        }
    }
}

pub struct ReadPoolConfigManager {
    scheduler: Scheduler<Task>,
}

impl ReadPoolConfigManager {
    pub fn new(
        handle: ReadPoolHandle,
        sender: SyncSender<usize>,
        worker: &Worker,
        thread_count: usize,
        auto_adjust: bool,
    ) -> Self {
        let max_thread_count = std::cmp::max(
            UNIFIED_READPOOL_MIN_CONCURRENCY,
            SysQuota::cpu_cores_quota().round() as usize,
        );
        let runner = ReadPoolConfigRunner {
            interval: READ_POOL_THREAD_CHECK_DURATION,
            sender,
            handle,
            cpu_time_tracker: ReadPoolCpuTimeTracker::new(&get_unified_read_pool_name()),
            process_stats: ProcessStat::cur_proc_stat().unwrap(),
            core_thread_count: thread_count,
            cur_thread_count: thread_count,
            max_thread_count,
            auto_adjust,
        };
        let scheduler = worker.start_with_timer("read-pool-config-worker", runner);

        Self { scheduler }
    }
}

impl Drop for ReadPoolConfigManager {
    fn drop(&mut self) {
        self.scheduler.stop();
    }
}

impl ConfigManager for ReadPoolConfigManager {
    fn dispatch(&mut self, change: ConfigChange) -> CfgResult<()> {
        if let Some(ConfigValue::Module(unified)) = change.get("unified") {
            if let Some(ConfigValue::Usize(max_thread_count)) = unified.get("max_thread_count") {
                self.scheduler.schedule(Task::PoolSize(*max_thread_count))?;
            }
            if let Some(ConfigValue::Bool(b)) = unified.get("auto_adjust_pool_size") {
                self.scheduler.schedule(Task::AutoAdjust(*b))?;
            }
            if let Some(ConfigValue::Module(keyspace_fairness)) = unified.get("keyspace_fairness") {
                if let Some(ConfigValue::F64(max_read_cpu_ratio)) =
                    keyspace_fairness.get("max_read_cpu_ratio")
                {
                    self.scheduler
                        .schedule(Task::KeyspaceFairnessMaxReadCpuRatio(*max_read_cpu_ratio))?;
                }
                if let Some(ConfigValue::Bool(use_wall_time)) =
                    keyspace_fairness.get("use_wall_time")
                {
                    self.scheduler
                        .schedule(Task::KeyspaceFairnessUseWallTime(*use_wall_time))?;
                }
                if let Some(ConfigValue::U64(slice_precharge_us)) =
                    keyspace_fairness.get("slice_precharge_us")
                {
                    self.scheduler
                        .schedule(Task::KeyspaceFairnessSlicePrechargeUs(*slice_precharge_us))?;
                }
                if let Some(ConfigValue::Usize(max_tasks_per_keyspace)) =
                    keyspace_fairness.get("max_tasks_per_keyspace")
                {
                    self.scheduler
                        .schedule(Task::KeyspaceFairnessMaxTasksPerKeyspace(
                            *max_tasks_per_keyspace,
                        ))?;
                }
                if let Some(ConfigValue::U64(max_time_slice_us)) =
                    keyspace_fairness.get("max_time_slice_us")
                {
                    self.scheduler
                        .schedule(Task::KeyspaceFairnessMaxTimeSliceUs(*max_time_slice_us))?;
                }
            }
        }
        info!(
            "readpool config changed";
            "change" => ?change,
        );
        Ok(())
    }
}

#[derive(Debug, Error)]
pub enum ReadPoolError {
    #[error("{0}")]
    FuturePoolFull(#[from] yatp_pool::Full),

    #[error("Unified read pool is full")]
    UnifiedReadPoolFull,

    #[error("{0}")]
    Canceled(#[from] oneshot::Canceled),
}

/// Bucket boundaries shared by production and simulator per-keyspace YATP
/// task-slice timing histograms.
pub const YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS: &[f64] = &[
    0.0001, 0.00025, 0.0005, 0.001, 0.002, 0.004, 0.008, 0.012, 0.016, 0.020, 0.050, 0.500, 5.0,
    15.0,
];

mod metrics {
    use lazy_static::lazy_static;
    use prometheus::*;

    use super::YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS;

    lazy_static! {
        pub static ref UNIFIED_READ_POOL_RUNNING_TASKS: IntGaugeVec = register_int_gauge_vec!(
            "tikv_unified_read_pool_running_tasks",
            "The number of running tasks in the unified read pool",
            &["name"]
        )
        .unwrap();
        pub static ref UNIFIED_READ_POOL_RUNNING_THREADS: IntGaugeVec = register_int_gauge_vec!(
            "tikv_unified_read_pool_thread_count",
            "The number of running threads in the unified read pool",
            &["name"]
        )
        .unwrap();
        pub static ref UNIFIED_READ_POOL_REJECT_TASKS: IntCounterVec = register_int_counter_vec!(
            "tikv_unified_read_pool_rejected_tasks_total",
            "The number of tasks rejected by the unified read pool",
            &["name", "reason", "keyspace"]
        )
        .unwrap();
        pub static ref TOKIO_POOL_MEAN_FIRST_POLL_DELAY: GaugeVec = register_gauge_vec!(
            "tikv_tokio_pool_mean_first_poll_delay_seconds",
            "Mean time from task spawn to first poll (schedule wait time) for tokio pools",
            &["name"]
        )
        .unwrap();
        pub static ref TOKIO_POOL_FIRST_POLL_COUNT: IntCounterVec = register_int_counter_vec!(
            "tikv_tokio_pool_first_poll_total",
            "Total number of tasks polled for first time in tokio pools",
            &["name"]
        )
        .unwrap();
        pub static ref TOKIO_POOL_MEAN_IDLE_DURATION: GaugeVec = register_gauge_vec!(
            "tikv_tokio_pool_mean_idle_duration_seconds",
            "Mean time tasks spent idle waiting on external events (I/O) for tokio pools",
            &["name"]
        )
        .unwrap();
        pub static ref TOKIO_POOL_MEAN_SCHEDULED_DURATION: GaugeVec = register_gauge_vec!(
            "tikv_tokio_pool_mean_scheduled_duration_seconds",
            "Mean time tasks spent waiting to be polled after being woken for tokio pools",
            &["name"]
        )
        .unwrap();
        pub static ref TOKIO_POOL_MEAN_POLL_DURATION: GaugeVec = register_gauge_vec!(
            "tikv_tokio_pool_mean_poll_duration_seconds",
            "Mean time tasks spent executing (being polled) for tokio pools",
            &["name"]
        )
        .unwrap();
        pub static ref YATP_POOL_KEYSPACE_FAIR_QUEUE_LENGTH_HISTOGRAM_VEC: HistogramVec =
            register_histogram_vec!(
                "tikv_yatp_pool_keyspace_fair_queue_length",
                "Histogram of per-keyspace queued task slices in the YATP keyspace fair read pool before a task slice is pushed.",
                &["name", "keyspace_id"],
                exponential_buckets(1.0, 2.0, 12).unwrap()
            )
            .unwrap();
        pub static ref YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_WAITS_HISTOGRAM_VEC: HistogramVec =
            register_histogram_vec!(
                "tikv_yatp_pool_keyspace_fair_queue_task_slice_waits",
                "Histogram of successful task-slice pops after a task slice is pushed and before it is popped from the YATP keyspace fair read pool.",
                &["name", "keyspace_id"],
                exponential_buckets(1.0, 2.0, 12).unwrap()
            )
            .unwrap();
        pub static ref YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_WAIT_DURATION_HISTOGRAM_VEC: HistogramVec =
            register_histogram_vec!(
                "tikv_yatp_pool_keyspace_fair_queue_task_slice_wait_duration_seconds",
                "Histogram of per-keyspace task-slice wait duration from push to pop in the YATP keyspace fair read pool.",
                &["name", "keyspace_id"],
                YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS.to_vec()
            )
            .unwrap();
        pub static ref YATP_POOL_KEYSPACE_FAIR_QUEUE_TASK_SLICE_SCHEDULE_INTERVAL_HISTOGRAM_VEC: HistogramVec =
            register_histogram_vec!(
                "tikv_yatp_pool_keyspace_fair_queue_task_slice_schedule_interval_seconds",
                "Histogram of per-keyspace task-slice scheduling intervals since the later of the previous pop and current enqueue; the first slice uses its enqueue time.",
                &["name", "keyspace_id"],
                YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS.to_vec()
            )
            .unwrap();
        pub static ref YATP_POOL_KEYSPACE_FAIR_GROUP_COUNT_GAUGE_VEC: IntGaugeVec =
            register_int_gauge_vec!(
                "tikv_yatp_pool_keyspace_fair_queue_group_count",
                "Current number of keyspace groups in the YATP keyspace fair read pool.",
                &["name"]
            )
            .unwrap();
        pub static ref YATP_POOL_KEYSPACE_TASK_SLICE_CPU_DURATION_HISTOGRAM_VEC: HistogramVec =
            register_histogram_vec!(
                "tikv_yatp_pool_keyspace_task_slice_cpu_duration_seconds",
                "Histogram of CPU time spent in one YATP read-pool task slice for each keyspace.",
                &["name", "keyspace_id"],
                YATP_POOL_KEYSPACE_TASK_SLICE_TIMING_BUCKETS.to_vec()
            )
            .unwrap();
        pub static ref YATP_POOL_KEYSPACE_TASK_WALL_TIME_COUNTER_VEC: CounterVec =
            register_counter_vec!(
                "tikv_yatp_pool_keyspace_task_wall_seconds_total",
                "Total wall time consumed by YATP read-pool task polls for each keyspace.",
                &["name", "keyspace_id"]
            )
            .unwrap();
    }
}

trait ReporterTickerHandle {
    fn try_tick(&mut self);
    fn try_tick_if_due(&mut self);
}

struct TokioReporterTicker<R: FlowStatsReporter> {
    ticker: yatp_pool::TickerWrapper<ReporterTicker<R>>,
    next_check_time: Instant,
}

impl<R: FlowStatsReporter> TokioReporterTicker<R> {
    fn new(ticker: yatp_pool::TickerWrapper<ReporterTicker<R>>) -> Self {
        Self {
            ticker,
            next_check_time: Instant::now_coarse() + TOKIO_REPORTER_TICK_CHECK_INTERVAL,
        }
    }
}

impl<R: FlowStatsReporter> ReporterTickerHandle for TokioReporterTicker<R> {
    fn try_tick(&mut self) {
        self.next_check_time = Instant::now_coarse() + TOKIO_REPORTER_TICK_CHECK_INTERVAL;
        let _ = self.ticker.try_tick();
    }

    fn try_tick_if_due(&mut self) {
        let now = Instant::now_coarse();
        if now < self.next_check_time {
            return;
        }
        self.next_check_time = now + TOKIO_REPORTER_TICK_CHECK_INTERVAL;
        let _ = self.ticker.try_tick();
    }
}

thread_local! {
    static REPORTER_TICKER: RefCell<Option<Box<dyn ReporterTickerHandle>>> = RefCell::new(None);
}

fn set_reporter_ticker<R: FlowStatsReporter>(ticker: yatp_pool::TickerWrapper<ReporterTicker<R>>) {
    REPORTER_TICKER.with(move |t| {
        let mut t = t.borrow_mut();
        if t.is_none() {
            *t = Some(Box::new(TokioReporterTicker::new(ticker)));
        }
    });
}

fn destroy_reporter_ticker() {
    REPORTER_TICKER.with(|t| {
        t.borrow_mut().take();
    });
}

fn try_tick_reporter_ticker() {
    REPORTER_TICKER.with(|t| {
        if let Some(ticker) = t.borrow_mut().as_mut() {
            ticker.try_tick();
        }
    });
}

fn try_tick_reporter_ticker_if_due() {
    REPORTER_TICKER.with(|t| {
        if let Some(ticker) = t.borrow_mut().as_mut() {
            ticker.try_tick_if_due();
        }
    });
}

#[cfg(test)]
mod tests {
    use std::{
        future::Future,
        pin::Pin,
        sync::{
            Arc, Barrier, Mutex,
            atomic::{AtomicU64, AtomicUsize, Ordering},
            mpsc,
        },
        task::{Context, Poll, Wake, Waker},
        thread,
        time::{Duration, Instant},
    };

    use api_version::api_v2::DEFAULT_KEYSPACE_ID;
    use futures::{channel::oneshot, executor::block_on};
    use kvproto::pdpb::QueryKind;
    use raftstore::store::{ReadStats, WriteStats};
    use tikv_util::task_reschedule::DEFAULT_MAX_TIME_SLICE_US;
    use tracker::{
        GLOBAL_TRACKERS, RequestInfo, Tracker, TrackerToken, clear_tls_tracker_token,
        set_tls_tracker_token,
    };

    use super::*;
    use crate::storage::TestEngineBuilder;

    #[derive(Clone)]
    struct DummyReporter;

    impl FlowStatsReporter for DummyReporter {
        fn report_read_stats(&self, _read_stats: ReadStats) {}
        fn report_write_stats(&self, _write_stats: WriteStats) {}
    }

    #[derive(Clone)]
    struct CountingReporter {
        read_reports: Arc<AtomicUsize>,
        write_reports: Arc<AtomicUsize>,
    }

    impl FlowStatsReporter for CountingReporter {
        fn report_read_stats(&self, _read_stats: ReadStats) {
            self.read_reports.fetch_add(1, Ordering::SeqCst);
        }

        fn report_write_stats(&self, _write_stats: WriteStats) {
            self.write_reports.fetch_add(1, Ordering::SeqCst);
        }
    }

    #[derive(Debug)]
    struct TestPollCpuRecorder {
        observed_nanos: AtomicU64,
        finishes: AtomicUsize,
    }

    impl ReadPoolTaskPollCpuRecorder for TestPollCpuRecorder {
        fn observe_poll_cpu_time(&self, cpu_time: Duration) {
            self.observed_nanos
                .fetch_add(cpu_time.as_nanos() as u64, Ordering::SeqCst);
        }

        fn finish(&self) {
            self.finishes.fetch_add(1, Ordering::SeqCst);
        }
    }

    #[derive(Default)]
    struct TestFairnessRecorder {
        charges: Mutex<Vec<(u32, i64)>>,
    }

    impl FairnessRecorder<u32> for TestFairnessRecorder {
        fn increase_vt(&self, keyspace_id: u32, vt: u64) {
            self.charges.lock().unwrap().push((keyspace_id, vt as i64));
        }

        fn decrease_vt(&self, keyspace_id: u32, vt: u64) {
            self.charges
                .lock()
                .unwrap()
                .push((keyspace_id, -(vt as i64)));
        }
    }

    impl TestFairnessRecorder {
        fn total_vt(&self) -> i64 {
            self.charges.lock().unwrap().iter().map(|(_, vt)| *vt).sum()
        }

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

    fn keyspace_fairness_yatp_config() -> UnifiedReadPoolConfig {
        UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 1,
            max_tasks_per_worker: 10,
            use_tokio: false,
            enable_yatp_keyspace_metrics: false,
            keyspace_fairness: crate::config::FairnessConfig {
                enabled: true,
                max_read_cpu_ratio: 0.0,
                use_wall_time: false,
                slice_precharge_us: 0,
                max_tasks_per_keyspace: 10,
                ..Default::default()
            },
            ..Default::default()
        }
    }

    fn fair_global_cpu_read_limiter(
        global_cpu_us: u64,
    ) -> (ReadLimiter, Arc<resource_control::CpuThrottleManager>) {
        fair_global_cpu_read_limiter_with_refill(global_cpu_us, 60_000)
    }

    fn fair_global_cpu_read_limiter_with_refill(
        global_cpu_us: u64,
        refill_interval_ms: u64,
    ) -> (ReadLimiter, Arc<resource_control::CpuThrottleManager>) {
        let mut resource_config = resource_control::Config::default();
        resource_config.enabled = true;
        resource_config.cpu_throttle.enabled = true;
        let read_limiter = ReadLimiter::new(resource_config);

        let mut cpu_config = resource_control::CpuThrottleConfig::default();
        cpu_config.enabled = true;
        cpu_config.refill_interval_ms = refill_interval_ms;
        cpu_config.estimated_cpu_per_request_us = global_cpu_us;
        let manager = Arc::new(resource_control::CpuThrottleManager::new(
            cpu_config,
            global_cpu_us,
            std::collections::HashMap::new(),
        ));
        read_limiter.set_cpu_throttle_manager(manager.clone());
        (read_limiter, manager)
    }

    fn yatp_fairness_recorder(handle: &ReadPoolHandle) -> Arc<dyn FairnessRecorder<u32>> {
        match handle {
            ReadPoolHandle::Yatp {
                keyspace_fair_queue: Some(keyspace_fair_queue),
                ..
            } => keyspace_fair_queue.get_fairness_recorder(),
            _ => unreachable!(),
        }
    }

    fn yatp_keyspace_task_limiter(handle: &ReadPoolHandle) -> KeyspaceTaskLimiter {
        match handle {
            ReadPoolHandle::Yatp {
                keyspace_task_limiter: Some(limiter),
                ..
            } => limiter.clone(),
            _ => unreachable!(),
        }
    }

    #[test]
    fn test_yatp_tasks_have_isolated_current_regions() {
        const FIRST_REGION_ID: u64 = 11;
        const SECOND_REGION_ID: u64 = 22;

        let config = UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 1,
            max_tasks_per_worker: 10,
            use_tokio: false,
            ..Default::default()
        };
        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();

        let (first_started_tx, first_started_rx) = mpsc::channel();
        let (resume_first_tx, resume_first_rx) = oneshot::channel();
        let (observed_tx, observed_rx) = mpsc::channel();

        let first_observed_tx = observed_tx.clone();
        handle
            .spawn(
                async move {
                    assert_eq!(tikv_util::get_current_region(), 0);
                    tikv_util::set_current_region(FIRST_REGION_ID);
                    first_started_tx.send(()).unwrap();
                    resume_first_rx.await.unwrap();
                    first_observed_tx
                        .send(("first", tikv_util::get_current_region()))
                        .unwrap();
                },
                CommandPri::Normal,
                DEFAULT_KEYSPACE_ID,
                1,
            )
            .unwrap();
        first_started_rx
            .recv_timeout(Duration::from_secs(5))
            .unwrap();

        handle
            .spawn(
                async move {
                    observed_tx
                        .send(("second_initial", tikv_util::get_current_region()))
                        .unwrap();
                    tikv_util::set_current_region(SECOND_REGION_ID);
                    observed_tx
                        .send(("second", tikv_util::get_current_region()))
                        .unwrap();
                },
                CommandPri::Normal,
                DEFAULT_KEYSPACE_ID,
                2,
            )
            .unwrap();

        assert_eq!(
            observed_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
            ("second_initial", 0)
        );
        assert_eq!(
            observed_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
            ("second", SECOND_REGION_ID)
        );

        resume_first_tx.send(()).unwrap();
        assert_eq!(
            observed_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
            ("first", FIRST_REGION_ID)
        );
    }

    #[test]
    fn test_keyspace_task_limiter_lazy_cleanup_races_with_acquire_and_drop() {
        const WORKERS: usize = 8;
        const ITERS: usize = 512;

        let limiter = Arc::new(KeyspaceTaskLimiter::new(64));

        // Keep keyspace 1 live across the stress window. Cleanup must not
        // remove this counter while a permit still contributes to it.
        let live_permit = limiter.try_acquire(1, false).unwrap();

        // Create zero-count counters so forced cleanup has entries to remove.
        for keyspace_id in 2..32 {
            drop(limiter.try_acquire(keyspace_id, false).unwrap());
        }

        let start = Arc::new(Barrier::new(WORKERS + 1));
        let mut workers = Vec::with_capacity(WORKERS);

        for worker_id in 0..WORKERS {
            let limiter = limiter.clone();
            let start = start.clone();

            workers.push(thread::spawn(move || {
                start.wait();

                for i in 0..ITERS {
                    if i % 17 == 0 {
                        limiter
                            .zero_transitions_since_cleanup
                            .store(KEYSPACE_TASK_COUNTER_CLEANUP_THRESHOLD, Ordering::Relaxed);
                    }

                    let keyspace_id = 1 + ((worker_id + i) % 3) as u32;
                    let permit = limiter.try_acquire(keyspace_id, false).unwrap();

                    if i % 5 == 0 {
                        let extra_keyspace_id = 1000 + worker_id as u32;
                        drop(limiter.try_acquire(extra_keyspace_id, false).unwrap());
                    }

                    drop(permit);
                }
            }));
        }

        start.wait();

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

        // If cleanup removed a live counter or split accounting across Arcs,
        // this would not converge to the long-lived permit's single count.
        assert_eq!(limiter.task_count(1), 1);
        assert!(limiter.has_task_counter(1));

        drop(live_permit);

        limiter
            .zero_transitions_since_cleanup
            .store(KEYSPACE_TASK_COUNTER_CLEANUP_THRESHOLD, Ordering::Relaxed);
        drop(limiter.try_acquire(9999, false).unwrap());

        assert_eq!(limiter.task_count(1), 0);
        assert!(!limiter.has_task_counter(1));
    }

    fn wait_until(timeout: Duration, cond: impl Fn() -> bool) {
        let start = Instant::now();
        while start.elapsed() < timeout {
            if cond() {
                return;
            }
            thread::sleep(Duration::from_millis(10));
        }
        assert!(cond());
    }

    #[derive(Default)]
    struct WakeOnce {
        has_woken: bool,
    }

    impl Future for WakeOnce {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if self.has_woken {
                Poll::Ready(())
            } else {
                self.has_woken = true;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }

    struct BlockingWake {
        wake_started_tx: mpsc::SyncSender<()>,
        release_rx: Mutex<mpsc::Receiver<()>>,
    }

    impl BlockingWake {
        fn block(&self) {
            let _ = self.wake_started_tx.send(());
            let _ = self.release_rx.lock().unwrap().recv();
        }
    }

    impl Wake for BlockingWake {
        fn wake(self: Arc<Self>) {
            self.block();
        }

        fn wake_by_ref(self: &Arc<Self>) {
            self.block();
        }
    }

    #[derive(Debug)]
    struct TestTaskDispatchProfile {
        schedule_time: Option<StdInstant>,
        dispatch_start_time: StdInstant,
        queue_wait: Option<Duration>,
        fair_queue_waited_task_slices: Option<u64>,
    }

    impl From<&TaskDispatchProfile> for TestTaskDispatchProfile {
        fn from(profile: &TaskDispatchProfile) -> Self {
            Self {
                schedule_time: profile.schedule_time,
                dispatch_start_time: profile.dispatch_start_time,
                queue_wait: profile.queue_wait,
                fair_queue_waited_task_slices: profile.fair_queue_waited_task_slices,
            }
        }
    }

    #[derive(Debug)]
    enum TestYatpPollEvent {
        Start(TestTaskDispatchProfile),
        End { task_finished: bool },
    }

    struct TestYatpPollEventTracker {
        event_tx: mpsc::Sender<TestYatpPollEvent>,
    }

    impl YatpPollEventTracker for TestYatpPollEventTracker {
        fn on_poll_start(&mut self, profile: &TaskDispatchProfile) {
            self.event_tx
                .send(TestYatpPollEvent::Start(profile.into()))
                .unwrap();
        }

        fn on_poll_end(&mut self, task_finished: bool, _profile: &TaskPollProfile) {
            self.event_tx
                .send(TestYatpPollEvent::End { task_finished })
                .unwrap();
        }
    }

    fn spawn_with_test_tracker<F>(handle: &ReadPoolHandle, future: F, task_id: u64) -> TrackerToken
    where
        F: Future<Output = ()> + Send + 'static,
    {
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo {
            task_id,
            ..Default::default()
        }));
        set_tls_tracker_token(tracker);
        let result = handle.spawn(future, CommandPri::Normal, DEFAULT_KEYSPACE_ID, task_id);
        clear_tls_tracker_token();
        if let Err(err) = result {
            GLOBAL_TRACKERS.remove(tracker);
            panic!("failed to spawn test read-pool task: {:?}", err);
        }
        tracker
    }

    fn take_pool_task_metrics(tracker: TrackerToken) -> PoolTaskMetrics {
        GLOBAL_TRACKERS
            .remove(tracker)
            .expect("test tracker should exist")
            .metrics
            .read_pool
    }

    fn clone_pool_task_metrics(tracker: TrackerToken) -> PoolTaskMetrics {
        GLOBAL_TRACKERS
            .with_tracker(tracker, |tracker| tracker.metrics.read_pool.clone())
            .expect("test tracker should exist")
    }

    fn clone_active_pool_task_metrics() -> PoolTaskMetrics {
        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|state| {
            state
                .borrow()
                .as_ref()
                .expect("tracker metrics observer state should be active")
                .state
                .metrics
                .clone()
        })
    }

    fn assert_pool_task_metrics_eq(actual: &PoolTaskMetrics, expected: &PoolTaskMetrics) {
        assert_eq!(actual.total_wall_nanos, expected.total_wall_nanos);
        assert_eq!(
            actual.total_queue_wait_nanos,
            expected.total_queue_wait_nanos
        );
        assert_eq!(actual.max_queue_wait_nanos, expected.max_queue_wait_nanos);
        assert_eq!(actual.min_queue_wait_nanos, expected.min_queue_wait_nanos);
        assert_eq!(actual.total_wake_wait_nanos, expected.total_wake_wait_nanos);
        assert_eq!(actual.max_wake_wait_nanos, expected.max_wake_wait_nanos);
        assert_eq!(actual.min_wake_wait_nanos, expected.min_wake_wait_nanos);
        assert_eq!(actual.fair_queue_enabled, expected.fair_queue_enabled);
        assert_eq!(
            actual.total_fair_queue_waited_task_slices,
            expected.total_fair_queue_waited_task_slices
        );
        assert_eq!(
            actual.max_fair_queue_waited_task_slices,
            expected.max_fair_queue_waited_task_slices
        );
        assert_eq!(
            actual.min_fair_queue_waited_task_slices,
            expected.min_fair_queue_waited_task_slices
        );
        assert_eq!(actual.dispatch_count, expected.dispatch_count);
        assert_eq!(actual.poll_cpu_nanos, expected.poll_cpu_nanos);
        assert_eq!(actual.max_poll_cpu_nanos, expected.max_poll_cpu_nanos);
        assert_eq!(actual.min_poll_cpu_nanos, expected.min_poll_cpu_nanos);
        assert_eq!(actual.poll_wall_nanos, expected.poll_wall_nanos);
        assert_eq!(actual.max_poll_wall_nanos, expected.max_poll_wall_nanos);
        assert_eq!(actual.min_poll_wall_nanos, expected.min_poll_wall_nanos);
        assert_eq!(actual.poll_count, expected.poll_count);
    }

    fn yatp_running_tasks(handle: &ReadPoolHandle) -> IntGauge {
        match handle {
            ReadPoolHandle::Yatp { running_tasks, .. } => running_tasks.clone(),
            _ => panic!("expected a YATP read pool"),
        }
    }

    #[test]
    fn test_spawn_with_options_forwards_yatp_poll_events() {
        const TASK_ID: u64 = 1;
        const EVENT_COUNT: usize = 6;

        let engine = TestEngineBuilder::new().build().unwrap();
        let mut config = keyspace_fairness_yatp_config();
        config.keyspace_fairness.enabled = false;
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();

        let (event_tx, event_rx) = mpsc::channel();
        let (finished_tx, finished_rx) = mpsc::channel();
        let poll_cpu_recorder = Arc::new(TestPollCpuRecorder {
            observed_nanos: AtomicU64::new(0),
            finishes: AtomicUsize::new(0),
        });
        let future = async move {
            // Wake once without leaving the dispatch, then explicitly
            // yield so the final poll runs in a new dispatch.
            WakeOnce::default().await;
            yatp::task::future::yield_to_scheduler().await;
            let start = ThreadTime::now();
            while start.elapsed() < Duration::from_millis(1) {
                std::hint::black_box(());
            }
            finished_tx.send(()).unwrap();
        };

        // Install a request tracker to verify that custom poll observers do
        // not enable RequestMetrics aggregation or TLS synchronization.
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo {
            task_id: TASK_ID,
            ..Default::default()
        }));
        set_tls_tracker_token(tracker);
        let result = handle.spawn_with_options(
            future,
            CommandPri::Normal,
            DEFAULT_KEYSPACE_ID,
            TASK_ID,
            ReadPoolTaskOptions::default()
                .with_poll_cpu_recorder(poll_cpu_recorder.clone())
                .with_poll_tracker_observer(Box::new(TestYatpPollEventTracker { event_tx })),
        );
        clear_tls_tracker_token();
        if let Err(err) = result {
            GLOBAL_TRACKERS.remove(tracker);
            panic!("failed to spawn test read-pool task: {:?}", err);
        }
        finished_rx.recv_timeout(Duration::from_secs(5)).unwrap();

        let events = (0..EVENT_COUNT)
            .map(|_| event_rx.recv_timeout(Duration::from_secs(5)).unwrap())
            .collect::<Vec<_>>();
        wait_until(Duration::from_secs(5), || {
            poll_cpu_recorder.finishes.load(Ordering::SeqCst) == 1
        });
        let metrics = take_pool_task_metrics(tracker);
        drop(pool);

        // Every poll must report a start/end pair. The first two polls are
        // immediate repolls in one dispatch; the final poll runs after the
        // explicit scheduler yield.
        assert!(matches!(&events[0], TestYatpPollEvent::Start(_)));
        assert!(matches!(
            &events[1],
            TestYatpPollEvent::End {
                task_finished: false
            }
        ));
        assert!(matches!(&events[2], TestYatpPollEvent::Start(_)));
        assert!(matches!(
            &events[3],
            TestYatpPollEvent::End {
                task_finished: false
            }
        ));
        assert!(matches!(&events[4], TestYatpPollEvent::Start(_)));
        assert!(matches!(
            &events[5],
            TestYatpPollEvent::End {
                task_finished: true
            }
        ));

        let TestYatpPollEvent::Start(first_dispatch) = &events[0] else {
            unreachable!();
        };
        let TestYatpPollEvent::Start(first_dispatch_repoll) = &events[2] else {
            unreachable!();
        };
        let TestYatpPollEvent::Start(second_dispatch) = &events[4] else {
            unreachable!();
        };
        assert_eq!(
            first_dispatch.dispatch_start_time,
            first_dispatch_repoll.dispatch_start_time
        );
        assert_ne!(
            first_dispatch.dispatch_start_time,
            second_dispatch.dispatch_start_time
        );
        assert!(first_dispatch.schedule_time.is_some());
        assert!(first_dispatch.queue_wait.is_some());
        assert_eq!(first_dispatch.fair_queue_waited_task_slices, None);

        assert_eq!(metrics.total_wall_nanos, 0);
        assert_eq!(metrics.dispatch_count, 0);
        assert_eq!(metrics.poll_count, 0);
        assert!(poll_cpu_recorder.observed_nanos.load(Ordering::SeqCst) > 0);
        assert_eq!(poll_cpu_recorder.finishes.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_tracker_metrics_observer_aggregates_profiles_and_fallback_syncs_when_task_finishes() {
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
        set_tls_tracker_token(tracker);

        let mut observer = TrackerMetricsObserver {
            enable_req_metrics: true,
            ..Default::default()
        };
        observer.state.as_mut().unwrap().task_start_time =
            StdInstant::now() - Duration::from_millis(1);

        // Record the first poll in the first dispatch. A zero fair-queue wait is
        // still a valid sample and marks fairness as enabled.
        let first_dispatch_start_time = StdInstant::now();
        let first_dispatch = TaskDispatchProfile {
            schedule_time: Some(first_dispatch_start_time - Duration::from_nanos(10)),
            dispatch_start_time: first_dispatch_start_time,
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: first_dispatch_start_time,
            queue_wait: Some(Duration::from_nanos(10)),
            fair_queue_waited_task_slices: Some(0),
        };
        observer.on_poll_start(&first_dispatch);
        observer.on_poll_end(
            false,
            &TaskPollProfile {
                cpu_time: Duration::from_nanos(10),
                wall_time: Duration::from_nanos(20),
            },
        );

        // An immediate repoll uses the same dispatch context. It contributes a
        // poll sample but must not duplicate dispatch, queue, or fair-queue
        // metrics.
        let first_dispatch_repoll = TaskDispatchProfile {
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: StdInstant::now(),
            ..first_dispatch
        };
        observer.on_poll_start(&first_dispatch_repoll);
        observer.on_poll_end(
            false,
            &TaskPollProfile {
                cpu_time: Duration::from_nanos(30),
                wall_time: Duration::from_nanos(40),
            },
        );

        // A later dispatch contributes another queue sample and measures wake
        // wait from the previous poll end until this dispatch was scheduled.
        let second_schedule_time =
            observer.state.as_ref().unwrap().last_poll_end_time.unwrap() + Duration::from_nanos(30);
        let second_dispatch = TaskDispatchProfile {
            schedule_time: Some(second_schedule_time),
            dispatch_start_time: second_schedule_time + Duration::from_nanos(20),
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: second_schedule_time + Duration::from_nanos(20),
            queue_wait: Some(Duration::from_nanos(20)),
            fair_queue_waited_task_slices: Some(5),
        };
        observer.on_poll_start(&second_dispatch);
        observer.on_poll_end(
            false,
            &TaskPollProfile {
                cpu_time: Duration::from_nanos(50),
                wall_time: Duration::from_nanos(60),
            },
        );
        let total_wall_before_finish = observer.state.as_ref().unwrap().metrics.total_wall_nanos;

        // A dispatch without schedule time still records dispatch and
        // fair-queue data, but it cannot contribute queue or wake wait. The
        // Ready poll completes total wall-time accounting.
        let third_dispatch = TaskDispatchProfile {
            schedule_time: None,
            dispatch_start_time: second_dispatch.dispatch_start_time + Duration::from_nanos(1),
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: second_dispatch.dispatch_start_time + Duration::from_nanos(1),
            queue_wait: None,
            fair_queue_waited_task_slices: Some(2),
        };
        observer.on_poll_start(&third_dispatch);
        observer.on_poll_end(
            true,
            &TaskPollProfile {
                cpu_time: Duration::from_nanos(70),
                wall_time: Duration::from_nanos(80),
            },
        );

        assert!(observer.state.is_none());
        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|state| {
            assert!(state.borrow().is_none());
        });

        let metrics = GLOBAL_TRACKERS
            .with_tracker(tracker, |tracker| tracker.metrics.read_pool.clone())
            .expect("test tracker should exist");
        clear_tls_tracker_token();
        GLOBAL_TRACKERS.remove(tracker);

        assert_eq!(total_wall_before_finish, 0);
        assert_eq!(metrics.dispatch_count, 3);
        assert_eq!(metrics.total_queue_wait_nanos, 30);
        assert_eq!(metrics.min_queue_wait_nanos, 10);
        assert_eq!(metrics.max_queue_wait_nanos, 20);
        assert_eq!(metrics.total_wake_wait_nanos, 30);
        assert_eq!(metrics.min_wake_wait_nanos, 30);
        assert_eq!(metrics.max_wake_wait_nanos, 30);
        assert!(metrics.fair_queue_enabled);
        assert_eq!(metrics.total_fair_queue_waited_task_slices, 7);
        assert_eq!(metrics.min_fair_queue_waited_task_slices, 0);
        assert_eq!(metrics.max_fair_queue_waited_task_slices, 5);
        assert_eq!(metrics.poll_count, 4);
        assert_eq!(metrics.poll_cpu_nanos, 160);
        assert_eq!(metrics.min_poll_cpu_nanos, 10);
        assert_eq!(metrics.max_poll_cpu_nanos, 70);
        assert_eq!(metrics.poll_wall_nanos, 200);
        assert_eq!(metrics.min_poll_wall_nanos, 20);
        assert_eq!(metrics.max_poll_wall_nanos, 80);
        assert!(metrics.total_wall_nanos >= 1_000_000);
    }

    #[test]
    fn test_sync_read_pool_metrics_to_current_tracker_snapshots_without_mutating_state_until_done()
    {
        let tracker_a = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
        let tracker_b = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
        let tracker_c = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
        let mut observer = TrackerMetricsObserver {
            enable_req_metrics: true,
            ..Default::default()
        };

        let first_poll_start = StdInstant::now();
        let first_poll = TaskDispatchProfile {
            schedule_time: Some(first_poll_start),
            dispatch_start_time: first_poll_start,
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: first_poll_start,
            queue_wait: Some(Duration::ZERO),
            fair_queue_waited_task_slices: Some(0),
        };
        observer.on_poll_start(&first_poll);
        let original_metrics = clone_active_pool_task_metrics();

        set_tls_tracker_token(tracker_a);
        sync_read_pool_metrics_to_current_tracker(false);
        let tracker_a_metrics = clone_pool_task_metrics(tracker_a);
        let active_metrics_after_a = clone_active_pool_task_metrics();

        let start = ThreadTime::now();
        while start.elapsed() < Duration::from_millis(1) {
            std::hint::black_box(());
        }
        set_tls_tracker_token(tracker_b);
        sync_read_pool_metrics_to_current_tracker(false);
        let tracker_b_metrics = clone_pool_task_metrics(tracker_b);
        let active_metrics_after_b = clone_active_pool_task_metrics();

        observer.on_poll_end(
            false,
            &TaskPollProfile {
                cpu_time: Duration::from_nanos(10),
                wall_time: Duration::from_nanos(20),
            },
        );
        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|state| {
            assert!(state.borrow().is_none());
        });
        assert_eq!(observer.state.as_ref().unwrap().metrics.poll_count, 1);

        let second_poll_start = StdInstant::now();
        let second_poll = TaskDispatchProfile {
            schedule_time: Some(second_poll_start),
            dispatch_start_time: second_poll_start,
            poll_start_thread_time: ThreadTime::now(),
            poll_start_time: second_poll_start,
            queue_wait: Some(Duration::ZERO),
            fair_queue_waited_task_slices: Some(0),
        };
        observer.on_poll_start(&second_poll);
        set_tls_tracker_token(tracker_c);
        sync_read_pool_metrics_to_current_tracker(true);
        let final_metrics = clone_pool_task_metrics(tracker_c);

        assert!(observer.state.is_none());
        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|state| {
            assert!(state.borrow().is_none());
        });

        // `done=true` consumes the observer state. Even while the tracker and
        // its TLS token still exist, later synchronization and poll callbacks
        // must not overwrite the final snapshot.
        sync_read_pool_metrics_to_current_tracker(false);
        sync_read_pool_metrics_to_current_tracker(true);
        observer.on_poll_end(
            true,
            &TaskPollProfile {
                cpu_time: Duration::from_secs(10),
                wall_time: Duration::from_secs(10),
            },
        );
        observer.on_poll_start(&second_poll);
        observer.on_poll_end(
            true,
            &TaskPollProfile {
                cpu_time: Duration::from_secs(10),
                wall_time: Duration::from_secs(10),
            },
        );
        let metrics_after_done = clone_pool_task_metrics(tracker_c);

        clear_tls_tracker_token();
        GLOBAL_TRACKERS.remove(tracker_a);
        GLOBAL_TRACKERS.remove(tracker_b);
        GLOBAL_TRACKERS.remove(tracker_c);

        assert_pool_task_metrics_eq(&active_metrics_after_a, &original_metrics);
        assert_pool_task_metrics_eq(&active_metrics_after_b, &original_metrics);
        assert_eq!(tracker_a_metrics.poll_count, 1);
        assert_eq!(tracker_b_metrics.poll_count, 1);
        assert!(tracker_b_metrics.poll_cpu_nanos >= tracker_a_metrics.poll_cpu_nanos);
        assert!(tracker_b_metrics.poll_wall_nanos >= tracker_a_metrics.poll_wall_nanos);
        assert_eq!(final_metrics.poll_count, 2);
        assert_pool_task_metrics_eq(&metrics_after_done, &final_metrics);
    }

    #[test]
    fn test_tracker_metrics_observer_without_yatp_dispatch_context_records_poll_metrics() {
        YatpDispatchContext::with_tls(|ctx| assert!(ctx.is_none()));

        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
        set_tls_tracker_token(tracker);

        let mut future = ObserveReadPoolTaskFuture::new(
            async {
                let start = ThreadTime::now();
                while start.elapsed() < Duration::from_millis(1) {
                    std::hint::black_box(());
                }
                sync_read_pool_metrics_to_current_tracker(true);
            },
            DEFAULT_KEYSPACE_ID,
        );
        future.set_enable_req_tracker_metrics(true);
        block_on(future);

        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|state| {
            assert!(state.borrow().is_none());
        });
        clear_tls_tracker_token();
        let metrics = take_pool_task_metrics(tracker);

        assert_eq!(metrics.dispatch_count, 1);
        assert_eq!(metrics.poll_count, 1);
        assert_eq!(metrics.total_queue_wait_nanos, 0);
        assert_eq!(metrics.total_wake_wait_nanos, 0);
        assert!(!metrics.fair_queue_enabled);
        assert_eq!(metrics.total_fair_queue_waited_task_slices, 0);
        assert!(metrics.poll_cpu_nanos > 0);
        assert!(metrics.poll_wall_nanos > 0);
        assert!(metrics.total_wall_nanos >= metrics.poll_wall_nanos);
    }

    #[test]
    fn test_sync_read_pool_metrics_to_missing_tracker_does_not_panic() {
        YatpDispatchContext::with_tls(|ctx| assert!(ctx.is_none()));

        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
        set_tls_tracker_token(tracker);

        let mut future = ObserveReadPoolTaskFuture::new(
            async move {
                assert!(GLOBAL_TRACKERS.remove(tracker).is_some());
                sync_read_pool_metrics_to_current_tracker(false);
                sync_read_pool_metrics_to_current_tracker(true);
            },
            DEFAULT_KEYSPACE_ID,
        );
        future.set_enable_req_tracker_metrics(true);
        block_on(future);

        TLS_TRACKER_METRICS_OBSERVER_STATE.with(|state| {
            assert!(state.borrow().is_none());
        });
        clear_tls_tracker_token();
        assert!(GLOBAL_TRACKERS.remove(tracker).is_none());
    }

    #[test]
    fn test_spawn_handle_syncs_tracker_metrics_before_publishing_result() {
        const TASK_ID: u64 = 4;
        const RESULT: u64 = 42;

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            None,
        );
        let handle = pool.handle();
        let running_tasks = yatp_running_tasks(&handle);
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo {
            task_id: TASK_ID,
            ..Default::default()
        }));
        let (run_tx, run_rx) = oneshot::channel();

        set_tls_tracker_token(tracker);
        let result = handle.spawn_handle(
            async move {
                run_rx.await.unwrap();
                RESULT
            },
            CommandPri::Normal,
            DEFAULT_KEYSPACE_ID,
            TASK_ID,
        );
        clear_tls_tracker_token();

        let (wake_started_tx, wake_started_rx) = mpsc::sync_channel(0);
        let (release_wake_tx, release_wake_rx) = mpsc::sync_channel(0);
        let waker = Waker::from(Arc::new(BlockingWake {
            wake_started_tx,
            release_rx: Mutex::new(release_wake_rx),
        }));
        let mut cx = Context::from_waker(&waker);
        let mut result = std::pin::pin!(result);
        assert!(result.as_mut().poll(&mut cx).is_pending());

        run_tx.send(()).unwrap();
        wake_started_rx
            .recv_timeout(Duration::from_secs(5))
            .unwrap();

        // The oneshot result is ready, while the worker is still blocked in
        // the sender's wake call and therefore cannot have reached
        // `on_poll_end` yet.
        let response_at_publication = result.as_mut().poll(&mut cx);
        let metrics_at_publication = clone_pool_task_metrics(tracker);

        release_wake_tx.send(()).unwrap();
        wait_until(Duration::from_secs(5), || running_tasks.get() == 0);
        let metrics_after_poll_end = take_pool_task_metrics(tracker);
        drop(pool);

        assert!(matches!(response_at_publication, Poll::Ready(Ok(RESULT))));
        assert!(metrics_at_publication.total_wall_nanos > 0);
        assert!(metrics_at_publication.dispatch_count > 0);
        assert!(metrics_at_publication.poll_count > 0);
        assert_pool_task_metrics_eq(&metrics_after_poll_end, &metrics_at_publication);
    }

    #[test]
    fn test_request_metrics_count_immediate_repoll_as_one_dispatch() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            None,
        );
        let handle = pool.handle();
        let running_tasks = yatp_running_tasks(&handle);
        let tracker = spawn_with_test_tracker(&handle, WakeOnce::default(), 1);

        wait_until(Duration::from_secs(5), || running_tasks.get() == 0);
        let metrics = take_pool_task_metrics(tracker);

        assert_eq!(metrics.dispatch_count, 1);
        assert_eq!(metrics.poll_count, 2);
        assert!(metrics.fair_queue_enabled);
        assert_eq!(metrics.total_fair_queue_waited_task_slices, 0);
        assert_eq!(metrics.min_fair_queue_waited_task_slices, 0);
        assert_eq!(metrics.max_fair_queue_waited_task_slices, 0);
        assert!(metrics.min_poll_wall_nanos <= metrics.max_poll_wall_nanos);
        assert!(metrics.max_poll_wall_nanos <= metrics.poll_wall_nanos);
        assert!(metrics.poll_wall_nanos <= metrics.total_wall_nanos);
    }

    #[test]
    fn test_request_metrics_count_yield_as_two_dispatches() {
        const TASK_ID: u64 = 2;
        const EVENT_COUNT: usize = 4;

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            None,
        );
        let handle = pool.handle();
        let running_tasks = yatp_running_tasks(&handle);
        let (event_tx, event_rx) = mpsc::channel();
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo {
            task_id: TASK_ID,
            ..Default::default()
        }));
        set_tls_tracker_token(tracker);
        let result = handle.spawn_with_options(
            async {
                yatp::task::future::yield_to_scheduler().await;
            },
            CommandPri::Normal,
            DEFAULT_KEYSPACE_ID,
            TASK_ID,
            ReadPoolTaskOptions::default()
                .with_poll_tracker_observer(Box::new(TestYatpPollEventTracker { event_tx })),
        );
        clear_tls_tracker_token();
        if let Err(err) = result {
            GLOBAL_TRACKERS.remove(tracker);
            panic!("failed to spawn test read-pool task: {:?}", err);
        }

        wait_until(Duration::from_secs(5), || running_tasks.get() == 0);
        let events = (0..EVENT_COUNT)
            .map(|_| event_rx.recv_timeout(Duration::from_secs(5)).unwrap())
            .collect::<Vec<_>>();
        let metrics = take_pool_task_metrics(tracker);

        assert!(matches!(&events[0], TestYatpPollEvent::Start(_)));
        assert!(matches!(
            &events[1],
            TestYatpPollEvent::End {
                task_finished: false
            }
        ));
        assert!(matches!(&events[2], TestYatpPollEvent::Start(_)));
        assert!(matches!(
            &events[3],
            TestYatpPollEvent::End {
                task_finished: true
            }
        ));
        let TestYatpPollEvent::Start(first_dispatch) = &events[0] else {
            unreachable!();
        };
        let TestYatpPollEvent::Start(second_dispatch) = &events[2] else {
            unreachable!();
        };
        assert_ne!(
            first_dispatch.dispatch_start_time,
            second_dispatch.dispatch_start_time
        );
        assert_eq!(first_dispatch.fair_queue_waited_task_slices, Some(0));

        assert_eq!(metrics.dispatch_count, 2);
        assert_eq!(metrics.poll_count, 2);
        assert!(metrics.fair_queue_enabled);
        assert!(metrics.min_queue_wait_nanos <= metrics.max_queue_wait_nanos);
        assert!(metrics.max_queue_wait_nanos <= metrics.total_queue_wait_nanos);
        assert!(metrics.min_wake_wait_nanos <= metrics.max_wake_wait_nanos);
        assert!(metrics.max_wake_wait_nanos <= metrics.total_wake_wait_nanos);
        assert!(metrics.poll_wall_nanos <= metrics.total_wall_nanos);
    }

    #[test]
    fn test_yatp_without_fairness_does_not_record_pool_task_metrics() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let mut config = keyspace_fairness_yatp_config();
        config.keyspace_fairness.enabled = false;
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();
        let running_tasks = yatp_running_tasks(&handle);
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo {
            task_id: 3,
            ..Default::default()
        }));
        set_tls_tracker_token(tracker);
        let result = handle.spawn_handle(
            WakeOnce::default(),
            CommandPri::Normal,
            DEFAULT_KEYSPACE_ID,
            3,
        );
        clear_tls_tracker_token();
        block_on(result).unwrap();

        wait_until(Duration::from_secs(5), || running_tasks.get() == 0);
        let metrics = take_pool_task_metrics(tracker);

        assert_eq!(metrics.total_wall_nanos, 0);
        assert_eq!(metrics.total_queue_wait_nanos, 0);
        assert_eq!(metrics.total_wake_wait_nanos, 0);
        assert!(!metrics.fair_queue_enabled);
        assert_eq!(metrics.total_fair_queue_waited_task_slices, 0);
        assert_eq!(metrics.dispatch_count, 0);
        assert_eq!(metrics.poll_cpu_nanos, 0);
        assert_eq!(metrics.poll_wall_nanos, 0);
        assert_eq!(metrics.poll_count, 0);
    }

    #[test]
    fn test_read_pool_task_poll_cpu_recorder_observes_and_finishes_once() {
        let recorder = Arc::new(TestPollCpuRecorder {
            observed_nanos: AtomicU64::new(0),
            finishes: AtomicUsize::new(0),
        });
        let recorder_trait = recorder.clone() as Arc<dyn ReadPoolTaskPollCpuRecorder>;
        let mut f = ObserveReadPoolTaskFuture::new(
            async {
                let start = ThreadTime::now();
                while start.elapsed() < Duration::from_millis(5) {
                    std::hint::black_box(());
                }
            },
            0,
        );
        f.set_poll_cpu_recorder(Some(recorder_trait));
        block_on(f);

        assert!(recorder.observed_nanos.load(Ordering::SeqCst) > 0);
        assert_eq!(recorder.finishes.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_reschedule_context_resets_at_pending_boundaries() {
        // Isolate yatp's reschedule-action TLS from other tests that may run on
        // this harness thread.
        thread::spawn(|| {
            use tikv_util::{
                task_reschedule::{
                    is_reschedule_context_active, reschedule_if_time_slice_exhausted,
                },
                time::LazyInstant,
            };

            assert!(!is_reschedule_context_active());
            let poll_count = Arc::new(AtomicUsize::new(0));
            let delegate = async {
                assert!(is_reschedule_context_active());
                let mut now = LazyInstant::default();
                reschedule_if_time_slice_exhausted(&mut now).await;

                // The reschedule above returned Pending. Reaching this point
                // requires a new outer poll with a fresh context.
                assert!(is_reschedule_context_active());
                let mut now = LazyInstant::default();
                reschedule_if_time_slice_exhausted(&mut now).await;
                WakeOnce::default().await;

                // An unrelated Pending boundary also gets a fresh context.
                assert!(is_reschedule_context_active());
                let mut now = LazyInstant::default();
                reschedule_if_time_slice_exhausted(&mut now).await;
            };
            let mut delegate = Box::pin(delegate);
            let wrapper_poll_count = poll_count.clone();
            let wrapper = std::future::poll_fn(move |cx| {
                assert!(!is_reschedule_context_active());
                let poll_index = wrapper_poll_count.fetch_add(1, Ordering::SeqCst);
                let poll_started_at = if poll_index == 0 {
                    StdInstant::now()
                        .checked_sub(Duration::from_millis(4))
                        .unwrap()
                } else {
                    // Use a synthetic future start to make the assertion
                    // deterministic even if this test thread is descheduled
                    // for longer than the production time slice.
                    StdInstant::now() + Duration::from_secs(1)
                };
                let result = ObserveReadPoolTaskFuture::<std::future::Ready<()>>::poll_delegate(
                    delegate.as_mut(),
                    cx,
                    poll_started_at,
                    Some(Duration::from_micros(DEFAULT_MAX_TIME_SLICE_US)),
                );
                assert!(!is_reschedule_context_active());
                result
            });

            block_on(wrapper);

            assert_eq!(poll_count.load(Ordering::SeqCst), 3);
            assert!(!is_reschedule_context_active());
        })
        .join()
        .unwrap();
    }

    #[test]
    fn test_fairness_recorder_controls_reschedule_context() {
        use tikv_util::task_reschedule::is_reschedule_context_active;

        let fairness_recorder = Arc::new(TestFairnessRecorder::default());
        let mut fair_future = ObserveReadPoolTaskFuture::new(
            async {
                // Wall-time rescheduling follows the fairness rollout policy,
                // even without an active fair-CPU permit.
                assert!(is_reschedule_context_active());
            },
            DEFAULT_KEYSPACE_ID,
        );
        fair_future.set_fairness_recorder(
            Some(fairness_recorder as Arc<dyn FairnessRecorder<u32>>),
            false,
            0,
            DEFAULT_MAX_TIME_SLICE_US,
        );
        block_on(fair_future);

        let poll_cpu_recorder = Arc::new(TestPollCpuRecorder {
            observed_nanos: AtomicU64::new(0),
            finishes: AtomicUsize::new(0),
        });
        let mut legacy_future = ObserveReadPoolTaskFuture::new(
            async {
                // Other observers do not implicitly enable the new behavior.
                assert!(!is_reschedule_context_active());
            },
            DEFAULT_KEYSPACE_ID,
        );
        legacy_future.set_poll_cpu_recorder(Some(
            poll_cpu_recorder as Arc<dyn ReadPoolTaskPollCpuRecorder>,
        ));
        block_on(legacy_future);

        assert!(!is_reschedule_context_active());
    }

    #[test]
    fn test_configured_max_time_slice_reaches_the_polled_task() {
        use tikv_util::{task_reschedule::reschedule_if_time_slice_exhausted, time::LazyInstant};

        // A zero slice is exhausted as soon as the poll starts, and a slice
        // longer than this test's runtime never is. Asserting both directions
        // shows the configured value, not a fixed constant, gates the check.
        for (max_time_slice_us, expect_exhausted) in [(0, true), (u32::MAX as u64, false)] {
            let fairness_recorder = Arc::new(TestFairnessRecorder::default());
            let mut f = ObserveReadPoolTaskFuture::new(
                async move {
                    let mut now = LazyInstant::default();
                    assert_eq!(
                        reschedule_if_time_slice_exhausted(&mut now).await,
                        expect_exhausted,
                    );
                },
                DEFAULT_KEYSPACE_ID,
            );
            f.set_fairness_recorder(
                Some(fairness_recorder as Arc<dyn FairnessRecorder<u32>>),
                false,
                0,
                max_time_slice_us,
            );
            block_on(f);
        }
    }

    #[test]
    fn test_read_pool_task_fairness_recorder_observes_precharge_and_cpu_time() {
        const KEYSPACE_ID: u32 = 42;
        const SLICE_PRECHARGE_US: u64 = 7;
        let fairness_recorder = Arc::new(TestFairnessRecorder::default());
        let poll_cpu_recorder = Arc::new(TestPollCpuRecorder {
            observed_nanos: AtomicU64::new(0),
            finishes: AtomicUsize::new(0),
        });
        let mut f = ObserveReadPoolTaskFuture::new(
            async {
                let start = ThreadTime::now();
                while start.elapsed() < Duration::from_millis(5) {
                    std::hint::black_box(());
                }
            },
            KEYSPACE_ID,
        );
        f.set_fairness_recorder(
            Some(fairness_recorder.clone() as Arc<dyn FairnessRecorder<u32>>),
            false,
            SLICE_PRECHARGE_US,
            DEFAULT_MAX_TIME_SLICE_US,
        );
        f.set_poll_cpu_recorder(Some(
            poll_cpu_recorder.clone() as Arc<dyn ReadPoolTaskPollCpuRecorder>
        ));

        block_on(f);

        let charges = fairness_recorder.charges();
        assert!(
            charges.contains(&(KEYSPACE_ID, SLICE_PRECHARGE_US as i64)),
            "{:?}",
            charges
        );
        assert!(
            charges.len() >= 2 && fairness_recorder.total_vt() > SLICE_PRECHARGE_US as i64,
            "{:?}",
            charges
        );
        assert!(
            charges
                .iter()
                .all(|(keyspace_id, _)| *keyspace_id == KEYSPACE_ID),
            "{:?}",
            charges
        );
        assert!(poll_cpu_recorder.observed_nanos.load(Ordering::SeqCst) > 0);
        assert_eq!(poll_cpu_recorder.finishes.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_read_pool_task_fairness_recorder_corrects_precharge_downward() {
        const KEYSPACE_ID: u32 = 42;
        const SLICE_PRECHARGE_US: u64 = 1_000_000;

        let fairness_recorder = Arc::new(TestFairnessRecorder::default());
        let mut f = ObserveReadPoolTaskFuture::new(async {}, KEYSPACE_ID);
        f.set_fairness_recorder(
            Some(fairness_recorder.clone() as Arc<dyn FairnessRecorder<u32>>),
            false,
            SLICE_PRECHARGE_US,
            DEFAULT_MAX_TIME_SLICE_US,
        );

        block_on(f);

        let charges = fairness_recorder.charges();
        assert!(
            charges.contains(&(KEYSPACE_ID, SLICE_PRECHARGE_US as i64)),
            "{:?}",
            charges
        );
        assert!(
            charges
                .iter()
                .any(|(keyspace_id, vt)| *keyspace_id == KEYSPACE_ID && *vt < 0),
            "{:?}",
            charges
        );
        assert!(fairness_recorder.total_vt() < SLICE_PRECHARGE_US as i64);
    }

    #[test]
    fn test_read_pool_task_fairness_recorder_can_use_wall_time() {
        const KEYSPACE_ID: u32 = 7;
        const SLEEP_TIME: Duration = Duration::from_millis(20);

        let wall_time_recorder = Arc::new(TestFairnessRecorder::default());
        let mut f = ObserveReadPoolTaskFuture::new(
            async {
                std::thread::sleep(SLEEP_TIME);
            },
            KEYSPACE_ID,
        );
        f.set_fairness_recorder(
            Some(wall_time_recorder.clone() as Arc<dyn FairnessRecorder<u32>>),
            true,
            0,
            DEFAULT_MAX_TIME_SLICE_US,
        );
        block_on(f);

        let cpu_time_recorder = Arc::new(TestFairnessRecorder::default());
        let mut f = ObserveReadPoolTaskFuture::new(
            async {
                std::thread::sleep(SLEEP_TIME);
            },
            KEYSPACE_ID,
        );
        f.set_fairness_recorder(
            Some(cpu_time_recorder.clone() as Arc<dyn FairnessRecorder<u32>>),
            false,
            0,
            DEFAULT_MAX_TIME_SLICE_US,
        );
        block_on(f);

        assert!(
            wall_time_recorder.total_vt() >= 10_000,
            "{}",
            wall_time_recorder.total_vt()
        );
        assert!(
            cpu_time_recorder.total_vt() < 10_000,
            "{}",
            cpu_time_recorder.total_vt()
        );
    }

    #[test]
    fn test_yatp_read_pool_builds_keyspace_fairness_recorder_from_config() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 1,
            max_tasks_per_worker: 10,
            use_tokio: false,
            enable_yatp_keyspace_metrics: false,
            keyspace_fairness: crate::config::FairnessConfig {
                enabled: false,
                use_wall_time: false,
                slice_precharge_us: 0,
                ..Default::default()
            },
            ..Default::default()
        };
        match build_yatp_read_pool(&config, DummyReporter, engine, None) {
            ReadPool::Yatp {
                keyspace_fair_queue,
                fairness_use_wall_time,
                fairness_slice_precharge_us,
                keyspace_task_limiter,
                ..
            } => {
                assert!(keyspace_fair_queue.is_none());
                assert!(!fairness_use_wall_time.load(Ordering::Relaxed));
                assert_eq!(fairness_slice_precharge_us.load(Ordering::Relaxed), 0);
                assert!(keyspace_task_limiter.is_none());
            }
            _ => unreachable!(),
        }

        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            keyspace_fairness: crate::config::FairnessConfig {
                enabled: true,
                use_wall_time: true,
                slice_precharge_us: 123,
                max_time_slice_us: 1_500,
                ..Default::default()
            },
            ..config
        };
        match build_yatp_read_pool(&config, DummyReporter, engine, None) {
            ReadPool::Yatp {
                keyspace_fair_queue,
                fairness_use_wall_time,
                fairness_slice_precharge_us,
                fairness_max_time_slice_us,
                keyspace_task_limiter,
                ..
            } => {
                assert!(keyspace_fair_queue.is_some());
                assert!(fairness_use_wall_time.load(Ordering::Relaxed));
                assert_eq!(fairness_slice_precharge_us.load(Ordering::Relaxed), 123);
                assert_eq!(fairness_max_time_slice_us.load(Ordering::Relaxed), 1_500);
                assert!(keyspace_task_limiter.is_some());
            }
            _ => unreachable!(),
        }
    }

    #[test]
    fn test_yatp_keyspace_metrics_follow_explicit_config() {
        for fairness_enabled in [false, true] {
            for metrics_enabled in [false, true] {
                let mut config = keyspace_fairness_yatp_config();
                config.enable_yatp_keyspace_metrics = metrics_enabled;
                config.keyspace_fairness.enabled = fairness_enabled;

                let engine = TestEngineBuilder::new().build().unwrap();
                match build_yatp_read_pool(&config, DummyReporter, engine, None) {
                    ReadPool::Yatp {
                        keyspace_metrics, ..
                    } => assert_eq!(keyspace_metrics.is_some(), metrics_enabled),
                    _ => unreachable!(),
                }
            }
        }
    }

    #[test]
    fn test_yatp_read_pool_fair_global_cpu_gate_waits_for_global_bucket() {
        const KEYSPACE_ID: u32 = 11;
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, manager) = fair_global_cpu_read_limiter_with_refill(1_000, 100);
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            Some(read_limiter),
        );
        let handle = pool.handle();
        assert!(handle.uses_fair_global_cpu_gate());

        assert!(!manager.consume_fair_global_cpu(1_000));
        assert_eq!(manager.global_available(), 0);

        let (started_tx, started_rx) = mpsc::channel();
        let (event_tx, event_rx) = mpsc::channel();
        let recorder = Arc::new(TestPollCpuRecorder {
            observed_nanos: AtomicU64::new(0),
            finishes: AtomicUsize::new(0),
        });
        let request_deadline = StdInstant::now() + Duration::from_secs(5);
        handle
            .spawn_with_options(
                async move {
                    started_tx.send(()).unwrap();
                },
                CommandPri::Normal,
                KEYSPACE_ID,
                1,
                ReadPoolTaskOptions::default()
                    .with_poll_cpu_recorder(recorder)
                    .with_fair_cpu_admission(FairCpuTaskType::DagSelect, Some(request_deadline))
                    .with_poll_tracker_observer(Box::new(TestYatpPollEventTracker { event_tx })),
            )
            .unwrap();

        thread::sleep(Duration::from_millis(20));
        assert!(started_rx.try_recv().is_err());

        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
        assert!(matches!(
            event_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
            TestYatpPollEvent::Start(_)
        ));
        assert!(matches!(
            event_rx.recv_timeout(Duration::from_secs(5)).unwrap(),
            TestYatpPollEvent::End {
                task_finished: true
            }
        ));
    }

    #[test]
    fn test_yatp_fairness_installs_reschedule_context() {
        use tikv_util::task_reschedule::is_reschedule_context_active;

        const KEYSPACE_ID: u32 = 12;
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, _manager) = fair_global_cpu_read_limiter(1_000_000);
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            Some(read_limiter),
        );
        let handle = pool.handle();

        let (release_tx, release_rx) = oneshot::channel();
        let (admitted_tx, admitted_rx) = mpsc::channel();
        handle
            .spawn_with_options(
                async move {
                    admitted_tx.send(is_reschedule_context_active()).unwrap();
                    release_rx.await.unwrap();
                    admitted_tx.send(is_reschedule_context_active()).unwrap();
                },
                CommandPri::Normal,
                KEYSPACE_ID,
                1,
                ReadPoolTaskOptions::default()
                    .with_fair_cpu_admission(FairCpuTaskType::DagSelect, None),
            )
            .unwrap();
        assert!(admitted_rx.recv_timeout(Duration::from_secs(5)).unwrap());

        let (without_admission_tx, without_admission_rx) = mpsc::channel();
        handle
            .spawn(
                async move {
                    without_admission_tx
                        .send(is_reschedule_context_active())
                        .unwrap();
                },
                CommandPri::Normal,
                KEYSPACE_ID,
                2,
            )
            .unwrap();
        // Wall-time rescheduling follows keyspace fairness independently of
        // fair-CPU admission.
        assert!(
            without_admission_rx
                .recv_timeout(Duration::from_secs(5))
                .unwrap()
        );

        release_tx.send(()).unwrap();
        assert!(admitted_rx.recv_timeout(Duration::from_secs(5)).unwrap());
    }

    #[test]
    fn test_yatp_read_pool_fair_global_cpu_gate_releases_unmarked_tasks() {
        const KEYSPACE_ID: u32 = 14;
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, manager) = fair_global_cpu_read_limiter(1_000);
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            Some(read_limiter),
        );
        let handle = pool.handle();
        assert!(handle.uses_fair_global_cpu_gate());
        assert_eq!(manager.global_available(), 1_000);

        let (started_tx, started_rx) = mpsc::channel();
        handle
            .spawn(
                async move {
                    started_tx.send(()).unwrap();
                },
                CommandPri::Normal,
                KEYSPACE_ID,
                1,
            )
            .unwrap();

        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
        assert_eq!(manager.global_available(), 1_000);
    }

    #[test]
    fn test_yatp_read_pool_fair_global_cpu_gate_charges_actual_dispatch_cpu() {
        const KEYSPACE_ID: u32 = 15;
        const GLOBAL_CPU_US: u64 = 1_000_000;
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, manager) = fair_global_cpu_read_limiter(GLOBAL_CPU_US);
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            Some(read_limiter),
        );
        let handle = pool.handle();
        let recorder = Arc::new(TestPollCpuRecorder {
            observed_nanos: AtomicU64::new(0),
            finishes: AtomicUsize::new(0),
        });
        let (finished_tx, finished_rx) = mpsc::channel();

        handle
            .spawn_with_options(
                async move {
                    let timer = ThreadTime::now();
                    while timer.elapsed() < Duration::from_millis(2) {
                        std::hint::black_box(());
                    }
                    finished_tx.send(()).unwrap();
                },
                CommandPri::Normal,
                KEYSPACE_ID,
                1,
                ReadPoolTaskOptions::default()
                    .with_poll_cpu_recorder(recorder.clone())
                    .with_fair_cpu_admission(
                        FairCpuTaskType::DagSelect,
                        Some(StdInstant::now() + Duration::from_secs(5)),
                    ),
            )
            .unwrap();

        finished_rx.recv_timeout(Duration::from_secs(5)).unwrap();
        wait_until(Duration::from_secs(5), || {
            manager.global_available() < GLOBAL_CPU_US
        });
        assert!(recorder.observed_nanos.load(Ordering::SeqCst) > 0);
        assert_eq!(manager.global_debt(), 0);
    }

    #[test]
    fn test_fair_global_cpu_admission_precheck_waits_when_global_bucket_is_empty() {
        let (read_limiter, manager) = fair_global_cpu_read_limiter(1_000);
        let controller = Arc::new(FairGlobalCpuAdmissionController::new(
            manager.clone(),
            read_limiter,
        ));

        assert!(matches!(
            controller.precheck(),
            FairQueueAdmissionResult::Ready
        ));
        assert!(!manager.consume_fair_global_cpu(1_000));
        assert_eq!(manager.global_available(), 0);

        assert!(matches!(
            controller.precheck(),
            FairQueueAdmissionResult::Pending { .. }
        ));
    }

    #[test]
    fn test_fair_global_cpu_admission_dry_run_tracks_debt_without_throttling() {
        const KEYSPACE_ID: u32 = 13;
        const TASK_TYPE: u8 = 1;
        let (read_limiter, manager) = fair_global_cpu_read_limiter(100);
        let controller = Arc::new(FairGlobalCpuAdmissionController::new(
            manager.clone(),
            read_limiter.clone(),
        ));
        let mut resource_config = resource_control::Config {
            enabled: true,
            read_enabled: true,
            ..Default::default()
        };
        let mut subscriber =
            resource_control::ReadSubscriber::new(read_limiter.clone(), resource_config.clone());
        subscriber.set_cpu_throttle_manager(manager.clone());

        assert!(!manager.consume_fair_global_cpu(100));
        assert!(matches!(
            controller.precheck(),
            FairQueueAdmissionResult::Pending { .. }
        ));

        resource_config.dry_run = true;
        subscriber.update_config(&resource_config);
        assert!(read_limiter.is_dry_run());
        assert!(matches!(
            controller.precheck(),
            FairQueueAdmissionResult::Ready
        ));

        let mut permit = match controller.admit(
            KEYSPACE_ID,
            TASK_TYPE,
            Some(StdInstant::now() + Duration::from_secs(5)),
        ) {
            FairAdmitOutcome::Admitted(permit) => permit,
            FairAdmitOutcome::Bypassed => panic!("fair global cpu admission should admit"),
        };
        assert_eq!(permit.consume(10), FairCpuBudgetAction::Continue);
        assert_eq!(manager.global_debt(), 10);

        resource_config.dry_run = false;
        subscriber.update_config(&resource_config);
        assert!(!read_limiter.is_dry_run());
        assert_eq!(manager.global_available(), 100);
        assert_eq!(manager.global_debt(), 0);
        assert!(matches!(
            controller.precheck(),
            FairQueueAdmissionResult::Ready
        ));
    }

    #[test]
    fn test_fair_global_cpu_admission_charges_actual_dispatch_cpu() {
        const KEYSPACE_ID: u32 = 13;
        const TASK_TYPE: u8 = 1;
        let (read_limiter, manager) = fair_global_cpu_read_limiter(100);
        let controller = Arc::new(FairGlobalCpuAdmissionController::new(
            manager.clone(),
            read_limiter,
        ));

        let mut permit = match controller.admit(
            KEYSPACE_ID,
            TASK_TYPE,
            Some(StdInstant::now() + Duration::from_secs(5)),
        ) {
            FairAdmitOutcome::Admitted(permit) => permit,
            FairAdmitOutcome::Bypassed => panic!("fair global cpu admission should admit"),
        };
        assert_eq!(manager.global_available(), 100);
        assert_eq!(manager.global_debt(), 0);

        assert_eq!(permit.consume(60), FairCpuBudgetAction::Continue);
        assert_eq!(manager.global_available(), 40);
        assert_eq!(permit.consume(50), FairCpuBudgetAction::YieldToScheduler);

        assert_eq!(manager.global_available(), 0);
        assert_eq!(manager.global_debt(), 10);
    }

    #[test]
    fn test_read_pool_task_options_build_fair_cpu_admission_metadata() {
        const KEYSPACE_ID: u32 = 12;
        let (read_limiter, manager) = fair_global_cpu_read_limiter(1_000_000);
        let controller = Arc::new(FairGlobalCpuAdmissionController::new(
            manager.clone(),
            read_limiter,
        ));
        let options = ReadPoolTaskOptions::default()
            .with_fair_cpu_admission(FairCpuTaskType::KvBatchGet, None);
        let admission = options.fair_cpu_admission.unwrap();
        assert_eq!(admission.task_type, FairCpuTaskType::KvBatchGet as u8);
        assert_eq!(admission.deadline, None);

        let mut permit =
            match controller.admit(KEYSPACE_ID, admission.task_type, admission.deadline) {
                FairAdmitOutcome::Admitted(permit) => permit,
                FairAdmitOutcome::Bypassed => panic!("fair global cpu admission should admit"),
            };
        assert_eq!(manager.global_available(), 1_000_000);
        assert_eq!(permit.consume(100), FairCpuBudgetAction::Continue);
        assert_eq!(manager.global_available(), 999_900);
    }

    #[test]
    fn test_fair_cpu_task_type_ids_are_stable() {
        assert_eq!(FairCpuTaskType::DagSelect as u8, 1);
        assert_eq!(FairCpuTaskType::DagIndex as u8, 2);
        assert_eq!(FairCpuTaskType::KvBatchGet as u8, 3);
        assert_eq!(FairCpuTaskType::KvGet as u8, 4);
        assert_eq!(FairCpuTaskType::KvBatchGetCommand as u8, 5);
        assert_eq!(FairCpuTaskType::KvScan as u8, 6);
        assert_eq!(FairCpuTaskType::KvScanLock as u8, 7);
    }

    #[test]
    fn test_fair_global_cpu_admission_bypasses_expired_deadline() {
        let (read_limiter, manager) = fair_global_cpu_read_limiter(1_000);
        let controller = FairGlobalCpuAdmissionController::new(manager, read_limiter);

        assert!(matches!(
            controller.admit(1, FairCpuTaskType::DagSelect as u8, Some(StdInstant::now())),
            FairAdmitOutcome::Bypassed
        ));
    }

    #[test]
    fn test_yatp_read_pool_fairness_without_cpu_throttle_skips_global_gate() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let read_limiter = ReadLimiter::new(resource_control::Config::default());
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            Some(read_limiter),
        );
        let handle = pool.handle();
        assert!(!handle.uses_fair_global_cpu_gate());

        let (started_tx, started_rx) = mpsc::channel();
        handle
            .spawn(
                async move {
                    started_tx.send(()).unwrap();
                },
                CommandPri::Normal,
                DEFAULT_KEYSPACE_ID,
                1,
            )
            .unwrap();
        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
    }

    #[test]
    fn test_yatp_read_pool_without_fairness_does_not_use_global_gate() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, manager) = fair_global_cpu_read_limiter(1_000);
        let config = UnifiedReadPoolConfig {
            use_tokio: false,
            keyspace_fairness: crate::config::FairnessConfig {
                enabled: false,
                ..Default::default()
            },
            ..keyspace_fairness_yatp_config()
        };
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, Some(read_limiter));
        let handle = pool.handle();
        assert!(!handle.uses_fair_global_cpu_gate());
        assert!(!manager.consume_fair_global_cpu(1_000));
        assert_eq!(manager.global_available(), 0);

        let (started_tx, started_rx) = mpsc::channel();
        handle
            .spawn(
                async move {
                    started_tx.send(()).unwrap();
                },
                CommandPri::Normal,
                DEFAULT_KEYSPACE_ID,
                1,
            )
            .unwrap();
        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();
    }

    #[test]
    fn test_yatp_read_pool_applies_fairness_max_read_cpu_ratio_at_startup() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, manager) = fair_global_cpu_read_limiter(1_000);
        let config = UnifiedReadPoolConfig {
            keyspace_fairness: crate::config::FairnessConfig {
                max_read_cpu_ratio: 0.3,
                ..keyspace_fairness_yatp_config().keyspace_fairness
            },
            ..keyspace_fairness_yatp_config()
        };

        let pool = build_yatp_read_pool(&config, DummyReporter, engine, Some(read_limiter));

        assert!(pool.handle().uses_fair_global_cpu_gate());
        assert_eq!(manager.fair_charge_scale(), 2.0);
        assert!((manager.enforced_max_read_cpu_ratio() - 0.3).abs() < f64::EPSILON);
    }

    #[test]
    fn test_yatp_read_pool_keyspace_fairness_config_can_change_online() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            keyspace_fairness: crate::config::FairnessConfig {
                slice_precharge_us: 1,
                max_tasks_per_keyspace: 2,
                ..keyspace_fairness_yatp_config().keyspace_fairness
            },
            ..keyspace_fairness_yatp_config()
        };
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();
        let limiter = yatp_keyspace_task_limiter(&handle);

        handle.set_fairness_use_wall_time(true);
        handle.set_fairness_slice_precharge_us(456);
        handle.set_max_tasks_per_keyspace(3);

        match handle {
            ReadPoolHandle::Yatp {
                fairness_use_wall_time,
                fairness_slice_precharge_us,
                ..
            } => {
                assert!(fairness_use_wall_time.load(Ordering::Relaxed));
                assert_eq!(fairness_slice_precharge_us.load(Ordering::Relaxed), 456);
                assert_eq!(limiter.max_tasks_per_keyspace(), 3);
            }
            _ => unreachable!(),
        }
    }

    #[test]
    fn test_read_pool_config_manager_dispatches_keyspace_fairness_online_config() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let (read_limiter, cpu_manager) = fair_global_cpu_read_limiter(1_000);
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            Some(read_limiter),
        );
        let handle = pool.handle();
        let observed_handle = handle.clone();
        let worker = Worker::new("test-read-pool-config-manager-keyspace-fairness");
        let (pool_size_notifier, _pool_size_receiver) = mpsc::sync_channel(10);
        let mut manager = ReadPoolConfigManager::new(handle, pool_size_notifier, &worker, 1, false);

        let mut keyspace_fairness = ConfigChange::default();
        keyspace_fairness.insert("use_wall_time".to_owned(), ConfigValue::Bool(true));
        keyspace_fairness.insert("slice_precharge_us".to_owned(), ConfigValue::U64(789));
        keyspace_fairness.insert("max_tasks_per_keyspace".to_owned(), ConfigValue::Usize(3));
        keyspace_fairness.insert("max_read_cpu_ratio".to_owned(), ConfigValue::F64(0.3));
        keyspace_fairness.insert("max_time_slice_us".to_owned(), ConfigValue::U64(1_234));
        let mut unified = ConfigChange::default();
        unified.insert(
            "keyspace_fairness".to_owned(),
            ConfigValue::Module(keyspace_fairness),
        );
        let mut change = ConfigChange::default();
        change.insert("unified".to_owned(), ConfigValue::Module(unified));

        manager.dispatch(change).unwrap();

        match observed_handle {
            ReadPoolHandle::Yatp {
                fairness_use_wall_time,
                fairness_slice_precharge_us,
                fairness_max_time_slice_us,
                keyspace_task_limiter: Some(limiter),
                ..
            } => wait_until(Duration::from_secs(5), || {
                fairness_use_wall_time.load(Ordering::Relaxed)
                    && fairness_slice_precharge_us.load(Ordering::Relaxed) == 789
                    && limiter.max_tasks_per_keyspace() == 3
                    && cpu_manager.fair_charge_scale() == 2.0
                    && fairness_max_time_slice_us.load(Ordering::Relaxed) == 1_234
                    && cpu_manager.fair_charge_scale() == 2.0
            }),
            _ => unreachable!(),
        }

        let mut keyspace_fairness = ConfigChange::default();
        keyspace_fairness.insert("max_read_cpu_ratio".to_owned(), ConfigValue::F64(0.0));
        let mut unified = ConfigChange::default();
        unified.insert(
            "keyspace_fairness".to_owned(),
            ConfigValue::Module(keyspace_fairness),
        );
        let mut change = ConfigChange::default();
        change.insert("unified".to_owned(), ConfigValue::Module(unified));
        manager.dispatch(change).unwrap();
        wait_until(Duration::from_secs(5), || {
            cpu_manager.fair_charge_scale() == 1.0
        });
    }

    #[test]
    fn test_uses_keyspace_fair_queue_without_global_cpu_gate() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            keyspace_fairness: crate::config::FairnessConfig {
                enabled: false,
                ..Default::default()
            },
            ..keyspace_fairness_yatp_config()
        };
        let handle = build_yatp_read_pool(&config, DummyReporter, engine, None).handle();
        assert!(!handle.uses_keyspace_fair_queue());
        assert!(handle.uses_coprocessor_concurrency_limiter());

        let engine = TestEngineBuilder::new().build().unwrap();
        let fair_config = keyspace_fairness_yatp_config();
        let fair_handle = build_yatp_read_pool(&fair_config, DummyReporter, engine, None).handle();
        assert!(fair_handle.uses_keyspace_fair_queue());
        assert!(!fair_handle.uses_fair_global_cpu_gate());
        assert!(!fair_handle.uses_coprocessor_concurrency_limiter());
    }

    #[test]
    fn test_yatp_read_pool_limits_admitted_tasks_per_keyspace() {
        const KEYSPACE_ID: u32 = 70_001;
        const OTHER_KEYSPACE_ID: u32 = 70_002;
        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            max_tasks_per_worker: 100_000,
            keyspace_fairness: crate::config::FairnessConfig {
                max_tasks_per_keyspace: 2,
                ..keyspace_fairness_yatp_config().keyspace_fairness
            },
            ..keyspace_fairness_yatp_config()
        };
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();
        let limiter = yatp_keyspace_task_limiter(&handle);

        let gen_task = || {
            let (tx, rx) = oneshot::channel::<()>();
            let task = async move {
                let _ = rx.await;
            };
            (task, tx)
        };

        let (task1, tx1) = gen_task();
        let (task2, _tx2) = gen_task();
        let (task3, _tx3) = gen_task();
        let (task4, _tx4) = gen_task();

        handle
            .spawn(task1, CommandPri::Normal, KEYSPACE_ID, 1)
            .unwrap();
        handle
            .spawn(task2, CommandPri::Normal, KEYSPACE_ID, 2)
            .unwrap();
        assert_eq!(limiter.task_count(KEYSPACE_ID), 2);

        match handle.spawn(task3, CommandPri::Normal, KEYSPACE_ID, 3) {
            Err(ReadPoolError::UnifiedReadPoolFull) => {}
            _ => panic!("should return full error for the saturated keyspace"),
        }

        handle
            .spawn(task4, CommandPri::Normal, OTHER_KEYSPACE_ID, 4)
            .unwrap();
        assert_eq!(limiter.task_count(OTHER_KEYSPACE_ID), 1);

        tx1.send(()).unwrap();
        wait_until(Duration::from_secs(5), || {
            limiter.task_count(KEYSPACE_ID) == 1
        });

        handle
            .spawn(async {}, CommandPri::Normal, KEYSPACE_ID, 5)
            .unwrap();
    }

    #[test]
    fn test_yatp_read_pool_keyspace_fairness_skips_global_max_tasks() {
        const KEYSPACE_ID: u32 = 70_005;
        const OTHER_KEYSPACE_ID: u32 = 70_006;
        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            max_tasks_per_worker: 2,
            keyspace_fairness: crate::config::FairnessConfig {
                max_tasks_per_keyspace: 10,
                ..keyspace_fairness_yatp_config().keyspace_fairness
            },
            ..keyspace_fairness_yatp_config()
        };
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();
        let limiter = yatp_keyspace_task_limiter(&handle);

        let gen_task = || {
            let (tx, rx) = oneshot::channel::<()>();
            let task = async move {
                let _ = rx.await;
            };
            (task, tx)
        };

        let (task1, tx1) = gen_task();
        let (task2, tx2) = gen_task();
        let (task3, tx3) = gen_task();

        handle
            .spawn(task1, CommandPri::Normal, KEYSPACE_ID, 1)
            .unwrap();
        handle
            .spawn(task2, CommandPri::Normal, KEYSPACE_ID, 2)
            .unwrap();
        assert_eq!(limiter.task_count(KEYSPACE_ID), 2);

        handle
            .spawn(task3, CommandPri::Normal, OTHER_KEYSPACE_ID, 3)
            .unwrap();
        assert_eq!(limiter.task_count(OTHER_KEYSPACE_ID), 1);

        let _ = tx1.send(());
        let _ = tx2.send(());
        let _ = tx3.send(());
        wait_until(Duration::from_secs(5), || {
            limiter.task_count(KEYSPACE_ID) == 0 && limiter.task_count(OTHER_KEYSPACE_ID) == 0
        });
    }

    #[test]
    fn test_read_pool_reject_metric_labels_keyspace_only_for_fairness() {
        const KEYSPACE_ID: u32 = 70_004;
        assert_eq!(
            unified_read_pool_reject_keyspace_label(Some(KEYSPACE_ID)),
            KEYSPACE_ID.to_string()
        );
        assert_eq!(unified_read_pool_reject_keyspace_label(None), "-");

        let fair_counter = record_unified_read_pool_reject_and_get_counter(
            Some(KEYSPACE_ID),
            READ_POOL_REJECT_REASON_GLOBAL_FULL,
        );
        let non_fair_counter = record_unified_read_pool_reject_and_get_counter(
            None,
            READ_POOL_REJECT_REASON_GLOBAL_FULL,
        );

        assert!(fair_counter.get() > 0);
        assert!(non_fair_counter.get() > 0);
    }

    #[test]
    fn test_non_fair_yatp_task_metadata_is_empty() {
        let mut extras = Extras::new_multilevel(1, None);
        set_yatp_task_metadata_if_fair(&mut extras, false, 42, None);
        assert!(extras.metadata().is_empty());

        set_yatp_task_metadata_if_fair(&mut extras, true, 42, None);
        assert_eq!(extras.metadata(), 42u32.to_be_bytes());
    }

    #[test]
    fn test_yatp_read_pool_max_tasks_per_keyspace_online_increase() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let config = UnifiedReadPoolConfig {
            max_tasks_per_worker: 100_000,
            keyspace_fairness: crate::config::FairnessConfig {
                max_tasks_per_keyspace: 2,
                ..keyspace_fairness_yatp_config().keyspace_fairness
            },
            ..keyspace_fairness_yatp_config()
        };
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();

        let gen_task = || {
            let (tx, rx) = oneshot::channel::<()>();
            let task = async move {
                let _ = rx.await;
            };
            (task, tx)
        };

        let (task1, _tx1) = gen_task();
        let (task2, _tx2) = gen_task();
        let (task3, _tx3) = gen_task();
        handle.spawn(task1, CommandPri::Normal, 1, 1).unwrap();
        handle.spawn(task2, CommandPri::Normal, 1, 2).unwrap();

        match handle.spawn(async {}, CommandPri::Normal, 1, 3) {
            Err(ReadPoolError::UnifiedReadPoolFull) => {}
            _ => panic!("should return full error before online increase"),
        }

        handle.set_max_tasks_per_keyspace(3);
        handle.spawn(task3, CommandPri::Normal, 1, 4).unwrap();
    }

    #[test]
    #[ignore = "CSE incompatibility; see #6158"]
    fn test_yatp_read_pool_spawn_routes_task_by_keyspace_metadata() {
        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(
            &keyspace_fairness_yatp_config(),
            DummyReporter,
            engine,
            None,
        );
        let handle = pool.handle();
        let recorder = yatp_fairness_recorder(&handle);
        let (warmup_tx, warmup_rx) = mpsc::channel();
        handle
            .spawn(
                async move {
                    warmup_tx.send(()).unwrap();
                },
                CommandPri::Normal,
                1,
                0,
            )
            .unwrap();
        warmup_rx.recv_timeout(Duration::from_secs(5)).unwrap();

        let (started_tx, started_rx) = mpsc::channel();
        let (release_tx, release_rx) = mpsc::channel();
        let (order_tx, order_rx) = mpsc::channel();

        handle
            .spawn(
                async move {
                    started_tx.send(()).unwrap();
                    release_rx.recv().unwrap();
                },
                CommandPri::Normal,
                1,
                1,
            )
            .unwrap();
        started_rx.recv_timeout(Duration::from_secs(5)).unwrap();

        let order_tx_1 = order_tx.clone();
        handle
            .spawn(
                async move {
                    order_tx_1.send(1).unwrap();
                },
                CommandPri::Normal,
                1,
                2,
            )
            .unwrap();
        let order_tx_2 = order_tx.clone();
        handle
            .spawn(
                async move {
                    order_tx_2.send(2).unwrap();
                },
                CommandPri::Normal,
                2,
                3,
            )
            .unwrap();

        // Advance keyspace 1 only after both target tasks have been routed into
        // their groups, so the assertion does not depend on new-group VT setup.
        recorder.increase_vt(1, 2_000_000);
        release_tx.send(()).unwrap();

        // Keyspace 2 has the lower VT and must be selected first. This also
        // verifies that spawn propagated each keyspace ID into task metadata.
        assert_eq!(order_rx.recv_timeout(Duration::from_secs(5)).unwrap(), 2);
        assert_eq!(order_rx.recv_timeout(Duration::from_secs(5)).unwrap(), 1);
    }

    #[test]
    fn test_tokio_ticker_ticks_under_load() {
        const TASK_COUNT: usize = 2000;
        const WORK_MS: u64 = 1;
        const TICK_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
        const COMPLETE_WAIT_TIMEOUT: Duration = Duration::from_secs(20);

        let config = UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 1,
            max_tasks_per_worker: TASK_COUNT + 100,
            ..Default::default()
        };

        let read_reports = Arc::new(AtomicUsize::new(0));
        let reporter = CountingReporter {
            read_reports: read_reports.clone(),
            write_reports: Arc::new(AtomicUsize::new(0)),
        };
        let completed = Arc::new(AtomicUsize::new(0));

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_tokio_pool(&config, reporter, engine);
        let handle = pool.handle();

        for task_id in 0..TASK_COUNT {
            let completed = completed.clone();
            let res = handle.spawn(
                async move {
                    crate::storage::metrics::tls_collect_query_count(1, QueryKind::Get, 1);
                    let start = Instant::now();
                    while start.elapsed() < Duration::from_millis(WORK_MS) {
                        std::hint::spin_loop();
                    }
                    completed.fetch_add(1, Ordering::SeqCst);
                },
                CommandPri::Normal,
                DEFAULT_KEYSPACE_ID,
                task_id as u64,
            );
            res.unwrap();
        }

        let wait_start = Instant::now();
        while read_reports.load(Ordering::SeqCst) == 0 && wait_start.elapsed() < TICK_WAIT_TIMEOUT {
            thread::sleep(Duration::from_millis(10));
        }
        assert!(read_reports.load(Ordering::SeqCst) > 0);
        assert!(completed.load(Ordering::SeqCst) <= TASK_COUNT);

        let complete_start = Instant::now();
        while completed.load(Ordering::SeqCst) < TASK_COUNT
            && complete_start.elapsed() < COMPLETE_WAIT_TIMEOUT
        {
            thread::sleep(Duration::from_millis(20));
        }
        assert_eq!(completed.load(Ordering::SeqCst), TASK_COUNT);
    }

    #[test]
    fn test_tokio_ticker_ticks_by_elapsed_time_under_long_tasks() {
        const TASK_COUNT: usize = 2;
        const BASE_WORK_MS: u64 = 1_250;
        const TICK_WAIT_TIMEOUT: Duration = Duration::from_secs(5);
        const COMPLETE_WAIT_TIMEOUT: Duration = Duration::from_secs(10);

        let config = UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 1,
            max_tasks_per_worker: TASK_COUNT + 100,
            ..Default::default()
        };

        let read_reports = Arc::new(AtomicUsize::new(0));
        let reporter = CountingReporter {
            read_reports: read_reports.clone(),
            write_reports: Arc::new(AtomicUsize::new(0)),
        };
        let completed = Arc::new(AtomicUsize::new(0));

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_tokio_pool(&config, reporter, engine);
        let handle = pool.handle();

        for task_id in 0..TASK_COUNT {
            let completed = completed.clone();
            let res = handle.spawn(
                async move {
                    crate::storage::metrics::tls_collect_query_count(1, QueryKind::Get, 1);
                    let start = Instant::now();
                    let work_duration =
                        Duration::from_millis(BASE_WORK_MS.saturating_mul(task_id as u64 + 1));
                    while start.elapsed() < work_duration {
                        std::hint::spin_loop();
                    }
                    completed.fetch_add(1, Ordering::SeqCst);
                },
                CommandPri::Normal,
                DEFAULT_KEYSPACE_ID,
                task_id as u64,
            );
            res.unwrap();
        }

        let wait_start = Instant::now();
        loop {
            let reports = read_reports.load(Ordering::SeqCst);
            let completed_tasks = completed.load(Ordering::SeqCst);
            if reports > 0 && completed_tasks < TASK_COUNT {
                break;
            }
            assert!(
                wait_start.elapsed() < TICK_WAIT_TIMEOUT,
                "report count: {}, completed tasks: {}",
                reports,
                completed_tasks
            );
            thread::sleep(Duration::from_millis(10));
        }

        let complete_start = Instant::now();
        while completed.load(Ordering::SeqCst) < TASK_COUNT
            && complete_start.elapsed() < COMPLETE_WAIT_TIMEOUT
        {
            thread::sleep(Duration::from_millis(20));
        }
        assert_eq!(completed.load(Ordering::SeqCst), TASK_COUNT);
    }

    #[test]
    fn test_yatp_scale_up() {
        let config = UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 2,
            max_tasks_per_worker: 1,
            ..Default::default()
        };
        // max running tasks number should be 2*1 = 2

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);

        let gen_task = || {
            let (tx, rx) = oneshot::channel::<()>();
            let task = async move {
                let _ = rx.await;
            };
            (task, tx)
        };

        let mut handle = pool.handle();
        let (task1, _tx1) = gen_task();
        let (task2, _tx2) = gen_task();
        let (task3, _tx3) = gen_task();
        let (task4, _tx4) = gen_task();
        let (task5, _tx5) = gen_task();

        handle
            .spawn(task1, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 1)
            .unwrap();
        handle
            .spawn(task2, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 2)
            .unwrap();

        thread::sleep(Duration::from_millis(300));
        match handle.spawn(task3, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 3) {
            Err(ReadPoolError::UnifiedReadPoolFull) => {}
            _ => panic!("should return full error"),
        }

        handle.scale_pool_size(3);
        assert_eq!(handle.get_normal_pool_size(), 3);

        handle
            .spawn(task4, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 4)
            .unwrap();

        thread::sleep(Duration::from_millis(300));
        match handle.spawn(task5, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 5) {
            Err(ReadPoolError::UnifiedReadPoolFull) => {}
            _ => panic!("should return full error"),
        }
    }

    #[test]
    fn test_yatp_scale_down() {
        let config = UnifiedReadPoolConfig {
            min_thread_count: 1,
            max_thread_count: 2,
            max_tasks_per_worker: 1,
            ..Default::default()
        };
        // max running tasks number should be 2*1 = 2

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);

        let gen_task = || {
            let (tx, rx) = oneshot::channel::<()>();
            let task = async move {
                let _ = rx.await;
            };
            (task, tx)
        };

        let mut handle = pool.handle();
        let (task1, tx1) = gen_task();
        let (task2, tx2) = gen_task();
        let (task3, _tx3) = gen_task();
        let (task4, _tx4) = gen_task();
        let (task5, _tx5) = gen_task();

        handle
            .spawn(task1, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 1)
            .unwrap();
        handle
            .spawn(task2, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 2)
            .unwrap();

        thread::sleep(Duration::from_millis(300));
        match handle.spawn(task3, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 3) {
            Err(ReadPoolError::UnifiedReadPoolFull) => {}
            _ => panic!("should return full error"),
        }

        tx1.send(()).unwrap();
        tx2.send(()).unwrap();
        thread::sleep(Duration::from_millis(300));

        handle.scale_pool_size(1);
        assert_eq!(handle.get_normal_pool_size(), 1);

        handle
            .spawn(task4, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 4)
            .unwrap();

        thread::sleep(Duration::from_millis(300));
        match handle.spawn(task5, CommandPri::Normal, DEFAULT_KEYSPACE_ID, 5) {
            Err(ReadPoolError::UnifiedReadPoolFull) => {}
            _ => panic!("should return full error"),
        }
    }

    #[test]
    fn test_auto_adjust_disable_notifies_pool_size_change() {
        use std::sync::mpsc::sync_channel;

        let config = UnifiedReadPoolConfig {
            min_thread_count: 2,
            max_thread_count: 8,
            max_tasks_per_worker: 10,
            ..Default::default()
        };

        let engine = TestEngineBuilder::new().build().unwrap();
        let pool = build_yatp_read_pool(&config, DummyReporter, engine, None);
        let handle = pool.handle();

        let (tx, rx) = sync_channel(10);

        let core_thread_count = config.max_thread_count;
        let max_thread_count = config.max_thread_count;
        let process_stats = match tikv_util::sys::cpu_time::ProcessStat::cur_proc_stat() {
            Ok(process_stats) => process_stats,
            Err(_) => return,
        };

        let mut runner = ReadPoolConfigRunner {
            interval: READ_POOL_THREAD_CHECK_DURATION,
            sender: tx,
            handle,
            cpu_time_tracker: ReadPoolCpuTimeTracker::new("test"),
            process_stats,
            core_thread_count,
            max_thread_count,
            cur_thread_count: core_thread_count,
            auto_adjust: true,
        };

        runner.cur_thread_count = 5;
        runner.run(Task::AutoAdjust(false));

        match rx.recv_timeout(Duration::from_millis(100)) {
            Ok(size) => assert_eq!(size, core_thread_count),
            Err(_) => {
                panic!("No pool size change notification received when auto-adjust was disabled.")
            }
        }

        assert_eq!(runner.cur_thread_count, core_thread_count);
    }
}