// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

#![feature(test)]

extern crate test;

#[cfg(target_os = "macos")]
use std::convert::TryFrom;
use std::{
    env,
    fmt::Write,
    future::Future,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
        mpsc,
    },
    time::{Duration, Instant},
};

use cpu_time::ThreadTime;
use futures::{executor::block_on, future::join_all};
use kvproto::kvrpcpb::CommandPri;
use prometheus::{CounterVec, HistogramVec, IntGauge};
use rand::{Rng, SeedableRng, rngs::StdRng};
use test::Bencher;
use tikv::read_pool::{ReadPool, ReadPoolHandle};
use tikv_util::{
    task_reschedule::DEFAULT_MAX_TIME_SLICE_US,
    yatp_pool::{DefaultTicker, KeyspaceInnerMetrics, KeyspaceMetrics, YatpPoolBuilder},
};
use tracker::{
    GLOBAL_TRACKERS, RequestInfo, Tracker, TrackerToken, get_tls_tracker_token,
    set_tls_tracker_token,
};

const ENV_FAIRNESS: &str = "TIKV_UTIL_YATP_BENCH_FAIRNESS";
const ENV_WORKERS: &str = "TIKV_UTIL_YATP_BENCH_WORKERS";
const ENV_SLICES: &str = "TIKV_UTIL_YATP_BENCH_SLICES";
const ENV_SLICE_CPU_US: &str = "TIKV_UTIL_YATP_BENCH_SLICE_CPU_US";
const ENV_ENABLE_YATP_KEYSPACE_METRICS: &str = "TIKV_UTIL_YATP_BENCH_ENABLE_YATP_KEYSPACE_METRICS";
const ENV_CLIENTS: &str = "TIKV_UTIL_YATP_BENCH_CLIENTS";
const ENV_KEYSPACES: &str = "TIKV_UTIL_YATP_BENCH_KEYSPACES";
const ENV_CHANNEL_CAPACITY: &str = "TIKV_UTIL_YATP_BENCH_CHANNEL_CAPACITY";

#[derive(Clone, Copy)]
struct BenchConfig {
    fairness: bool,
    workers: usize,
    slices_per_task: usize,
    slice_cpu_time: Duration,
    enable_yatp_keyspace_metrics: bool,
    clients: usize,
    keyspaces: usize,
    channel_capacity: usize,
}

impl BenchConfig {
    fn from_env() -> Self {
        let slices_per_task = read_env_usize(ENV_SLICES, 10).max(1);
        let clients = read_env_usize(ENV_CLIENTS, 16).max(1);
        Self {
            fairness: read_env_bool(ENV_FAIRNESS, true),
            workers: read_env_usize(ENV_WORKERS, 4).max(1),
            slices_per_task,
            slice_cpu_time: Duration::from_micros(read_env_u64(ENV_SLICE_CPU_US, 100).max(1)),
            enable_yatp_keyspace_metrics: read_env_bool(ENV_ENABLE_YATP_KEYSPACE_METRICS, true),
            clients,
            keyspaces: read_env_usize(ENV_KEYSPACES, 5).max(1),
            channel_capacity: read_env_usize(ENV_CHANNEL_CAPACITY, clients).max(1),
        }
    }
}

struct BenchPool {
    _pool: ReadPool,
    handle: ReadPoolHandle,
    metrics_counters: Option<MetricCounters>,
    thread_cpu_sampler: ThreadCpuSampler,
}

impl BenchPool {
    fn new(config: BenchConfig) -> Self {
        let metrics_counters = config
            .enable_yatp_keyspace_metrics
            .then(|| MetricCounters::new(config.keyspaces));
        let keyspace_metrics = metrics_counters
            .as_ref()
            .map(MetricCounters::keyspace_metrics);
        let thread_cpu_sampler = ThreadCpuSampler::new();
        let register_thread_cpu = thread_cpu_sampler.clone();
        let mut builder = YatpPoolBuilder::new(DefaultTicker);
        builder
            .name_prefix("bench-yatp-fair-pool")
            .thread_count(config.workers, config.workers, config.workers)
            .keyspace_metrics(keyspace_metrics.clone())
            .after_start(move || register_thread_cpu.register_current_thread());

        let (pool, keyspace_fair_queue) = if config.fairness {
            let (pool, keyspace_fair_queue) = builder.build_keyspace_fair_pool(None);
            (pool, Some(keyspace_fair_queue))
        } else {
            (builder.build_multi_level_pool(), None)
        };
        let read_pool = ReadPool::Yatp {
            pool,
            keyspace_fair_queue,
            fair_global_cpu_admission_controller: None,
            fairness_use_wall_time: Arc::new(AtomicBool::new(true)),
            fairness_slice_precharge_us: Arc::new(AtomicU64::new(0)),
            fairness_max_time_slice_us: Arc::new(AtomicU64::new(DEFAULT_MAX_TIME_SLICE_US)),
            keyspace_metrics,
            running_tasks: IntGauge::new(
                "bench_yatp_fair_pool_running_tasks",
                "bench yatp fair pool running tasks",
            )
            .unwrap(),
            running_threads: IntGauge::new(
                "bench_yatp_fair_pool_running_threads",
                "bench yatp fair pool running threads",
            )
            .unwrap(),
            max_tasks: Arc::new(AtomicUsize::new(usize::MAX)),
            keyspace_task_limiter: None,
            pool_size: Arc::new(AtomicUsize::new(config.workers)),
        };
        let handle = read_pool.handle();
        thread_cpu_sampler.wait_for_threads(config.workers);
        Self {
            _pool: read_pool,
            handle,
            metrics_counters,
            thread_cpu_sampler,
        }
    }
}

#[derive(Clone)]
struct ThreadCpuSampler {
    inner: Arc<ThreadCpuSamplerInner>,
}

struct ThreadCpuSamplerInner {
    thread_handles: Mutex<Vec<usize>>,
}

impl ThreadCpuSampler {
    fn new() -> Self {
        Self {
            inner: Arc::new(ThreadCpuSamplerInner {
                thread_handles: Mutex::new(Vec::new()),
            }),
        }
    }

    fn register_current_thread(&self) {
        if let Some(handle) = current_native_thread_handle() {
            self.inner.thread_handles.lock().unwrap().push(handle);
        }
    }

    fn wait_for_threads(&self, expected_threads: usize) {
        if !native_thread_cpu_supported() {
            return;
        }
        let deadline = Instant::now() + Duration::from_secs(5);
        while self.inner.thread_handles.lock().unwrap().len() < expected_threads {
            assert!(
                Instant::now() < deadline,
                "timed out waiting for read pool workers to start"
            );
            std::thread::sleep(Duration::from_millis(1));
        }
    }

    fn snapshot(&self) -> Option<ThreadCpuSnapshot> {
        let thread_handles = self.inner.thread_handles.lock().unwrap().clone();
        if thread_handles.is_empty() {
            return None;
        }
        let mut total_cpu_nanos = 0;
        for thread_handle in thread_handles {
            total_cpu_nanos += read_native_thread_cpu_nanos(thread_handle)?;
        }
        Some(ThreadCpuSnapshot { total_cpu_nanos })
    }
}

struct ThreadCpuSnapshot {
    total_cpu_nanos: u64,
}

#[derive(Default)]
struct WorkloadStats {
    submitted_tasks: AtomicU64,
    total_tasks: AtomicU64,
    total_slices: AtomicU64,
    total_cpu_nanos: AtomicU64,
    submitted_keyspace_tasks: Vec<AtomicU64>,
    keyspace_tasks: Vec<AtomicU64>,
    keyspace_slices: Vec<AtomicU64>,
    keyspace_cpu_nanos: Vec<AtomicU64>,
}

impl WorkloadStats {
    fn new(keyspaces: usize) -> Self {
        Self {
            submitted_tasks: AtomicU64::new(0),
            total_tasks: AtomicU64::new(0),
            total_slices: AtomicU64::new(0),
            total_cpu_nanos: AtomicU64::new(0),
            submitted_keyspace_tasks: new_atomic_vec(keyspaces),
            keyspace_tasks: new_atomic_vec(keyspaces),
            keyspace_slices: new_atomic_vec(keyspaces),
            keyspace_cpu_nanos: new_atomic_vec(keyspaces),
        }
    }

    fn record_submit(&self, keyspace_id: usize) {
        self.submitted_tasks.fetch_add(1, Ordering::Relaxed);
        self.submitted_keyspace_tasks[keyspace_id].fetch_add(1, Ordering::Relaxed);
    }

    fn record_slice(&self, keyspace_id: usize, cpu_time: Duration) {
        let cpu_nanos = cpu_time.as_nanos() as u64;
        self.total_slices.fetch_add(1, Ordering::Relaxed);
        self.total_cpu_nanos.fetch_add(cpu_nanos, Ordering::Relaxed);
        self.keyspace_slices[keyspace_id].fetch_add(1, Ordering::Relaxed);
        self.keyspace_cpu_nanos[keyspace_id].fetch_add(cpu_nanos, Ordering::Relaxed);
    }

    fn record_task(&self, keyspace_id: usize) {
        self.total_tasks.fetch_add(1, Ordering::Relaxed);
        self.keyspace_tasks[keyspace_id].fetch_add(1, Ordering::Relaxed);
    }
}

struct WorkloadReport {
    /// Wall-clock time from the first measured benchmark submission until all
    /// measured tasks complete.
    elapsed: Duration,
    /// Number of task submissions accepted by the bounded benchmark channel.
    /// This includes the initial pipeline-fill submissions used to keep the
    /// read pool saturated before standard benchmark iterations start.
    submitted_tasks: u64,
    /// Number of read-pool tasks that completed after draining the channel.
    total_tasks: u64,
    /// Number of cooperative CPU slices completed by all tasks.
    total_slices: u64,
    /// Thread CPU time consumed by all slices, measured inside the benchmark
    /// task body.
    total_cpu_nanos: u64,
    /// Total CPU time consumed by read-pool worker threads during the measured
    /// workload window. This includes task execution, polling, scheduling,
    /// queue operations, metrics, and other worker-side overhead.
    readpool_cpu_nanos: Option<u64>,
    /// Submitted task count per keyspace. Index is the keyspace id used by the
    /// benchmark.
    submitted_keyspace_tasks: Vec<u64>,
    /// Completed task count per keyspace.
    keyspace_tasks: Vec<u64>,
    /// Completed slice count per keyspace.
    keyspace_slices: Vec<u64>,
    /// Thread CPU time consumed per keyspace, measured inside the benchmark
    /// task body.
    keyspace_cpu_nanos: Vec<u64>,
    /// Per-keyspace CPU time observed from YATP keyspace metrics, when metrics
    /// are enabled.
    metrics_cpu_nanos: Option<Vec<u64>>,
    /// Per-keyspace wall time observed from YATP keyspace metrics, when metrics
    /// are enabled.
    metrics_wall_nanos: Option<Vec<u64>>,
}

impl WorkloadReport {
    fn from_stats(
        elapsed: Duration,
        stats: &WorkloadStats,
        metrics_before: Option<MetricSnapshot>,
        metrics_after: Option<MetricSnapshot>,
        readpool_cpu_before: Option<ThreadCpuSnapshot>,
        readpool_cpu_after: Option<ThreadCpuSnapshot>,
    ) -> Self {
        let (metrics_cpu_nanos, metrics_wall_nanos) = match (metrics_before, metrics_after) {
            (Some(before), Some(after)) => (
                Some(delta_f64_nanos(&before.cpu_seconds, &after.cpu_seconds)),
                Some(delta_f64_nanos(&before.wall_seconds, &after.wall_seconds)),
            ),
            _ => (None, None),
        };
        let readpool_cpu_nanos = match (readpool_cpu_before, readpool_cpu_after) {
            (Some(before), Some(after)) => {
                Some(after.total_cpu_nanos.saturating_sub(before.total_cpu_nanos))
            }
            _ => None,
        };
        Self {
            elapsed,
            submitted_tasks: stats.submitted_tasks.load(Ordering::Relaxed),
            total_tasks: stats.total_tasks.load(Ordering::Relaxed),
            total_slices: stats.total_slices.load(Ordering::Relaxed),
            total_cpu_nanos: stats.total_cpu_nanos.load(Ordering::Relaxed),
            readpool_cpu_nanos,
            submitted_keyspace_tasks: load_atomic_vec(&stats.submitted_keyspace_tasks),
            keyspace_tasks: load_atomic_vec(&stats.keyspace_tasks),
            keyspace_slices: load_atomic_vec(&stats.keyspace_slices),
            keyspace_cpu_nanos: load_atomic_vec(&stats.keyspace_cpu_nanos),
            metrics_cpu_nanos,
            metrics_wall_nanos,
        }
    }

    fn total_qps(&self) -> f64 {
        self.total_tasks as f64 / self.elapsed.as_secs_f64()
    }

    fn slice_qps(&self) -> f64 {
        self.total_slices as f64 / self.elapsed.as_secs_f64()
    }

    fn pool_wall_time_per_slice_us(&self, config: BenchConfig) -> f64 {
        let slice_qps = self.slice_qps();
        if slice_qps == 0.0 {
            0.0
        } else {
            config.workers as f64 * 1_000_000.0 / slice_qps
        }
    }

    fn cpu_time_per_slice_us(&self) -> f64 {
        nanos_per_slice_to_us(self.total_cpu_nanos, self.total_slices)
    }

    fn readpool_cpu_time_per_slice_us(&self) -> Option<f64> {
        self.readpool_cpu_nanos
            .map(|cpu_nanos| nanos_per_slice_to_us(cpu_nanos, self.total_slices))
    }

    fn readpool_cpu_utilization(&self, config: BenchConfig) -> Option<f64> {
        let worker_cpu_capacity_secs = self.elapsed.as_secs_f64() * config.workers as f64;
        self.readpool_cpu_nanos.map(|cpu_nanos| {
            if worker_cpu_capacity_secs == 0.0 {
                0.0
            } else {
                nanos_to_secs(cpu_nanos) / worker_cpu_capacity_secs
            }
        })
    }

    fn metrics_cpu_time_per_slice_us(&self) -> Option<f64> {
        self.metrics_cpu_nanos
            .as_ref()
            .map(|metrics| nanos_per_slice_to_us(metrics.iter().sum(), self.total_slices))
    }

    fn metrics_cpu_utilization(&self, config: BenchConfig) -> Option<f64> {
        let worker_cpu_capacity_secs = self.elapsed.as_secs_f64() * config.workers as f64;
        self.metrics_cpu_nanos.as_ref().map(|metrics| {
            if worker_cpu_capacity_secs == 0.0 {
                0.0
            } else {
                nanos_to_secs(metrics.iter().sum()) / worker_cpu_capacity_secs
            }
        })
    }

    fn metrics_wall_time_per_slice_us(&self) -> Option<f64> {
        self.metrics_wall_nanos
            .as_ref()
            .map(|metrics| nanos_per_slice_to_us(metrics.iter().sum(), self.total_slices))
    }

    fn format(&self, config: BenchConfig) -> String {
        let elapsed_secs = self.elapsed.as_secs_f64();
        let total_cpu_secs = nanos_to_secs(self.total_cpu_nanos);
        let slice_qps = self.slice_qps();
        let mut output = String::new();
        let _ = writeln!(
            output,
            "yatp_fair_pool config fairness={} workers={} slices_per_task={} slice_cpu_us={} enable_yatp_keyspace_metrics={} clients={} keyspaces={} channel_capacity={}",
            config.fairness,
            config.workers,
            config.slices_per_task,
            config.slice_cpu_time.as_micros(),
            config.enable_yatp_keyspace_metrics,
            config.clients,
            config.keyspaces,
            config.channel_capacity,
        );
        let _ = write!(output, "summary slice_qps={:.3}", slice_qps,);
        if let (Some(readpool_cpu_util), Some(readpool_cpu_per_slice)) = (
            self.readpool_cpu_utilization(config),
            self.readpool_cpu_time_per_slice_us(),
        ) {
            let _ = write!(
                output,
                " readpool_cpu_util={:.1}% readpool_cpu_per_slice={:.3}us",
                readpool_cpu_util * 100.0,
                readpool_cpu_per_slice,
            );
        } else {
            let _ = write!(
                output,
                " readpool_cpu_util=unsupported readpool_cpu_per_slice=unsupported",
            );
        }
        let _ = write!(
            output,
            " pool_wall_per_slice={:.3}us spawn_cpu_per_slice={:.3}us",
            self.pool_wall_time_per_slice_us(config),
            self.cpu_time_per_slice_us(),
        );
        if let (Some(metrics_cpu), Some(metrics_wall), Some(metrics_cpu_util)) = (
            self.metrics_cpu_time_per_slice_us(),
            self.metrics_wall_time_per_slice_us(),
            self.metrics_cpu_utilization(config),
        ) {
            let _ = write!(
                output,
                " metrics_cpu_per_slice={:.3}us metrics_wall_per_slice={:.3}us metrics_cpu_util={:.1}%",
                metrics_cpu,
                metrics_wall,
                metrics_cpu_util * 100.0,
            );
        } else {
            let _ = write!(
                output,
                " metrics_cpu_per_slice=disabled metrics_wall_per_slice=disabled metrics_cpu_util=disabled",
            );
        }
        output.push('\n');
        let readpool_cpu = self
            .readpool_cpu_nanos
            .map(|cpu_nanos| {
                format!(
                    "{:.6}s readpool_cpu_rate={:.3} cpu-s/s",
                    nanos_to_secs(cpu_nanos),
                    nanos_to_secs(cpu_nanos) / elapsed_secs,
                )
            })
            .unwrap_or_else(|| "unsupported readpool_cpu_rate=unsupported".to_owned());
        let _ = writeln!(
            output,
            "total submitted={} completed={} slices={} elapsed={:.6}s cpu={:.6}s cpu_rate={:.3} cpu-s/s readpool_cpu={} task_qps={:.3} slice_qps={:.3}",
            self.submitted_tasks,
            self.total_tasks,
            self.total_slices,
            elapsed_secs,
            total_cpu_secs,
            total_cpu_secs / elapsed_secs,
            readpool_cpu,
            self.total_qps(),
            slice_qps,
        );
        for keyspace_id in 0..config.keyspaces {
            let cpu_secs = nanos_to_secs(self.keyspace_cpu_nanos[keyspace_id]);
            let _ = write!(
                output,
                "keyspace={} submitted={} completed={} slices={} cpu={:.6}s cpu_rate={:.3} cpu-s/s task_qps={:.3} slice_qps={:.3}",
                keyspace_id,
                self.submitted_keyspace_tasks[keyspace_id],
                self.keyspace_tasks[keyspace_id],
                self.keyspace_slices[keyspace_id],
                cpu_secs,
                cpu_secs / elapsed_secs,
                self.keyspace_tasks[keyspace_id] as f64 / elapsed_secs,
                self.keyspace_slices[keyspace_id] as f64 / elapsed_secs,
            );
            if let (Some(cpu), Some(wall)) = (&self.metrics_cpu_nanos, &self.metrics_wall_nanos) {
                let _ = write!(
                    output,
                    " metrics_cpu={:.6}s metrics_wall={:.6}s",
                    nanos_to_secs(cpu[keyspace_id]),
                    nanos_to_secs(wall[keyspace_id]),
                );
            }
            output.push('\n');
        }
        output
    }
}

struct MetricCounters {
    queue_length: HistogramVec,
    queue_task_slice_waits: HistogramVec,
    queue_task_slice_wait_duration: HistogramVec,
    task_slice_schedule_interval: HistogramVec,
    task_slice_cpu_duration: HistogramVec,
    task_wall_time: CounterVec,
    group_count: IntGauge,
    keyspaces: usize,
}

impl MetricCounters {
    fn new(keyspaces: usize) -> Self {
        Self {
            queue_length: HistogramVec::new(
                prometheus::HistogramOpts::new(
                    "bench_yatp_fair_pool_keyspace_queue_length",
                    "bench yatp fair pool keyspace queue length",
                ),
                &["name", "keyspace_id"],
            )
            .unwrap(),
            queue_task_slice_waits: HistogramVec::new(
                prometheus::HistogramOpts::new(
                    "bench_yatp_fair_pool_keyspace_task_slice_waits",
                    "bench yatp fair pool keyspace task slice waits",
                ),
                &["name", "keyspace_id"],
            )
            .unwrap(),
            queue_task_slice_wait_duration: HistogramVec::new(
                prometheus::HistogramOpts::new(
                    "bench_yatp_fair_pool_keyspace_task_slice_wait_duration_seconds",
                    "bench yatp fair pool keyspace task slice wait duration",
                ),
                &["name", "keyspace_id"],
            )
            .unwrap(),
            task_slice_schedule_interval: HistogramVec::new(
                prometheus::HistogramOpts::new(
                    "bench_yatp_fair_pool_keyspace_task_slice_schedule_interval_seconds",
                    "bench yatp fair pool task slice schedule interval",
                ),
                &["name", "keyspace_id"],
            )
            .unwrap(),
            task_slice_cpu_duration: HistogramVec::new(
                prometheus::HistogramOpts::new(
                    "bench_yatp_fair_pool_keyspace_task_slice_cpu_duration_seconds",
                    "bench yatp fair pool keyspace task cpu time",
                ),
                &["name", "keyspace_id"],
            )
            .unwrap(),
            task_wall_time: CounterVec::new(
                prometheus::Opts::new(
                    "bench_yatp_fair_pool_keyspace_task_wall_seconds_total",
                    "bench yatp fair pool keyspace task wall time",
                ),
                &["name", "keyspace_id"],
            )
            .unwrap(),
            group_count: IntGauge::new(
                "bench_yatp_fair_pool_keyspace_group_count",
                "bench yatp fair pool keyspace group count",
            )
            .unwrap(),
            keyspaces,
        }
    }

    fn keyspace_metrics(&self) -> Arc<KeyspaceMetrics> {
        let queue_length = self.queue_length.clone();
        let queue_task_slice_waits = self.queue_task_slice_waits.clone();
        let queue_task_slice_wait_duration = self.queue_task_slice_wait_duration.clone();
        let task_slice_schedule_interval = self.task_slice_schedule_interval.clone();
        let task_slice_cpu_duration = self.task_slice_cpu_duration.clone();
        let task_wall_time = self.task_wall_time.clone();
        KeyspaceMetrics::new(
            self.group_count.clone(),
            move |keyspace_id| {
                let keyspace_id = keyspace_id.to_string();
                let labels = ["bench-yatp-fair-pool", keyspace_id.as_str()];
                KeyspaceInnerMetrics::new(
                    queue_length.with_label_values(&labels),
                    queue_task_slice_waits.with_label_values(&labels),
                    queue_task_slice_wait_duration.with_label_values(&labels),
                    task_slice_schedule_interval.with_label_values(&labels),
                    task_slice_cpu_duration.with_label_values(&labels),
                    task_wall_time.with_label_values(&labels),
                )
            },
            |_| {},
        )
    }

    fn snapshot(&self) -> MetricSnapshot {
        let mut cpu_seconds = Vec::with_capacity(self.keyspaces);
        let mut wall_seconds = Vec::with_capacity(self.keyspaces);
        for keyspace_id in 0..self.keyspaces {
            let keyspace_id = keyspace_id.to_string();
            let labels = ["bench-yatp-fair-pool", keyspace_id.as_str()];
            cpu_seconds.push(
                self.task_slice_cpu_duration
                    .with_label_values(&labels)
                    .get_sample_sum(),
            );
            wall_seconds.push(self.task_wall_time.with_label_values(&labels).get());
        }
        MetricSnapshot {
            cpu_seconds,
            wall_seconds,
        }
    }
}

struct MetricSnapshot {
    cpu_seconds: Vec<f64>,
    wall_seconds: Vec<f64>,
}

struct ReadTask {
    keyspace_index: usize,
    task_id: u64,
}

async fn run_clients(
    config: BenchConfig,
    handle: ReadPoolHandle,
    stats: Arc<WorkloadStats>,
    task_receiver: async_channel::Receiver<ReadTask>,
    completion_sender: mpsc::Sender<()>,
) {
    let clients = (0..config.clients).map(|_| {
        run_client(
            config,
            handle.clone(),
            stats.clone(),
            task_receiver.clone(),
            completion_sender.clone(),
        )
    });
    join_all(clients).await;
}

async fn run_client(
    config: BenchConfig,
    handle: ReadPoolHandle,
    stats: Arc<WorkloadStats>,
    task_receiver: async_channel::Receiver<ReadTask>,
    completion_sender: mpsc::Sender<()>,
) {
    let tracker_token = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::default()));
    while let Ok(task) = task_receiver.recv().await {
        let keyspace_id = task.keyspace_index as u32;
        spawn_read_pool_task(
            &handle,
            run_read_task(config, stats.clone(), task.keyspace_index),
            CommandPri::Normal,
            keyspace_id,
            task.task_id,
            tracker_token,
        )
        .await;
        completion_sender
            .send(())
            .expect("benchmark completion receiver should stay open");
    }
    let tracker = GLOBAL_TRACKERS
        .remove(tracker_token)
        .expect("benchmark tracker should remain registered until its client finishes");
    assert_eq!(tracker.metrics.read_pool.poll_count > 0, config.fairness);
}

fn spawn_read_pool_task<F>(
    handle: &ReadPoolHandle,
    f: F,
    priority: CommandPri,
    keyspace_id: u32,
    task_id: u64,
    tracker_token: TrackerToken,
) -> impl Future<Output = ()>
where
    F: Future<Output = ()> + Send + 'static,
{
    let previous_tracker = get_tls_tracker_token();
    set_tls_tracker_token(tracker_token);
    let task = handle.spawn_handle(f, priority, keyspace_id, task_id);
    set_tls_tracker_token(previous_tracker);
    async move {
        task.await
            .expect("benchmark read pool task should complete normally");
    }
}

async fn run_read_task(config: BenchConfig, stats: Arc<WorkloadStats>, keyspace_index: usize) {
    for slice_index in 0..config.slices_per_task {
        let cpu_time = consume_cpu(config.slice_cpu_time);
        stats.record_slice(keyspace_index, cpu_time);
        if slice_index + 1 < config.slices_per_task {
            yatp::task::future::yield_to_scheduler().await;
        }
    }
    stats.record_task(keyspace_index);
}

fn run_workload(config: BenchConfig, pool: &BenchPool, b: &mut Bencher) -> WorkloadReport {
    let stats = Arc::new(WorkloadStats::new(config.keyspaces));
    let (task_sender, task_receiver) = async_channel::bounded(config.channel_capacity);
    let final_receiver = task_receiver.clone();
    let (completion_sender, completion_receiver) = mpsc::channel();
    let client_stats = stats.clone();
    let client_handle = pool.handle.clone();
    let client_completion_sender = completion_sender.clone();
    let client_thread = std::thread::spawn(move || {
        block_on(run_clients(
            config,
            client_handle,
            client_stats,
            task_receiver,
            client_completion_sender,
        ));
    });

    let mut rng = StdRng::seed_from_u64(0x000B_17E5_7A7E);
    let mut submitted_tasks = 0;
    let mut completed_tasks = 0;
    let max_in_flight = config.channel_capacity as u64;
    let mut submit_task = || {
        let keyspace_index = rng.gen_range(0..config.keyspaces);
        let task_id = rng.r#gen::<u64>();
        task_sender
            .send_blocking(ReadTask {
                keyspace_index,
                task_id,
            })
            .expect("benchmark task channel should stay open");
        stats.record_submit(keyspace_index);
        submitted_tasks += 1;
        task_id
    };

    let before = pool.metrics_counters.as_ref().map(MetricCounters::snapshot);
    let readpool_cpu_before = pool.thread_cpu_sampler.snapshot();
    let start = std::time::Instant::now();
    for _ in 0..max_in_flight {
        test::black_box(submit_task());
    }
    b.iter(|| {
        let task_id = submit_task();
        completion_receiver
            .recv()
            .expect("benchmark completion sender should stay open");
        completed_tasks += 1;
        test::black_box(task_id);
    });
    for _ in completed_tasks..submitted_tasks {
        completion_receiver
            .recv()
            .expect("benchmark completion sender should stay open");
    }
    drop(task_sender);
    drop(completion_sender);
    client_thread
        .join()
        .expect("benchmark client thread should finish normally");
    assert!(final_receiver.is_empty());
    assert_eq!(
        stats.submitted_tasks.load(Ordering::Relaxed),
        stats.total_tasks.load(Ordering::Relaxed)
    );
    let elapsed = start.elapsed();
    let readpool_cpu_after = pool.thread_cpu_sampler.snapshot();
    let after = pool.metrics_counters.as_ref().map(MetricCounters::snapshot);
    WorkloadReport::from_stats(
        elapsed,
        &stats,
        before,
        after,
        readpool_cpu_before,
        readpool_cpu_after,
    )
}

fn consume_cpu(duration: Duration) -> Duration {
    let start = ThreadTime::now();
    loop {
        std::hint::black_box(());
        if start.elapsed() >= duration {
            return start.elapsed();
        }
    }
}

fn read_env_bool(name: &str, default: bool) -> bool {
    env::var(name)
        .ok()
        .and_then(|value| match value.to_ascii_lowercase().as_str() {
            "1" | "true" | "yes" | "on" => Some(true),
            "0" | "false" | "no" | "off" => Some(false),
            _ => None,
        })
        .unwrap_or(default)
}

fn read_env_usize(name: &str, default: usize) -> usize {
    env::var(name)
        .ok()
        .and_then(|value| value.parse().ok())
        .unwrap_or(default)
}

fn read_env_u64(name: &str, default: u64) -> u64 {
    env::var(name)
        .ok()
        .and_then(|value| value.parse().ok())
        .unwrap_or(default)
}

fn new_atomic_vec(len: usize) -> Vec<AtomicU64> {
    (0..len).map(|_| AtomicU64::new(0)).collect()
}

fn load_atomic_vec(values: &[AtomicU64]) -> Vec<u64> {
    values
        .iter()
        .map(|value| value.load(Ordering::Relaxed))
        .collect()
}

fn delta_f64_nanos(before: &[f64], after: &[f64]) -> Vec<u64> {
    before
        .iter()
        .zip(after)
        .map(|(before, after)| ((after - before) * 1_000_000_000.0) as u64)
        .collect()
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
fn native_thread_cpu_supported() -> bool {
    true
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn native_thread_cpu_supported() -> bool {
    false
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
fn current_native_thread_handle() -> Option<usize> {
    Some(unsafe { libc::pthread_self() as usize })
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn current_native_thread_handle() -> Option<usize> {
    None
}

#[cfg(target_os = "linux")]
fn read_native_thread_cpu_nanos(thread_handle: usize) -> Option<u64> {
    let mut clock_id = 0;
    let result =
        unsafe { libc::pthread_getcpuclockid(thread_handle as libc::pthread_t, &mut clock_id) };
    if result != 0 {
        return None;
    }
    let mut time = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    if unsafe { libc::clock_gettime(clock_id, &mut time) } != 0 {
        return None;
    }
    Some(time.tv_sec as u64 * 1_000_000_000 + time.tv_nsec as u64)
}

#[cfg(target_os = "macos")]
fn read_native_thread_cpu_nanos(thread_handle: usize) -> Option<u64> {
    let mach_thread = unsafe { libc::pthread_mach_thread_np(thread_handle as libc::pthread_t) };
    if mach_thread == 0 {
        return None;
    }
    let mut info: libc::thread_basic_info = unsafe { std::mem::zeroed() };
    let mut count = libc::THREAD_BASIC_INFO_COUNT;
    let result = unsafe {
        libc::thread_info(
            mach_thread,
            libc::THREAD_BASIC_INFO as libc::thread_flavor_t,
            &mut info as *mut libc::thread_basic_info as libc::thread_info_t,
            &mut count,
        )
    };
    if result != libc::KERN_SUCCESS {
        return None;
    }
    Some(time_value_to_nanos(info.user_time)? + time_value_to_nanos(info.system_time)?)
}

#[cfg(target_os = "macos")]
fn time_value_to_nanos(time: libc::time_value_t) -> Option<u64> {
    let seconds = u64::try_from(time.seconds).ok()?;
    let microseconds = u64::try_from(time.microseconds).ok()?;
    Some(seconds * 1_000_000_000 + microseconds * 1_000)
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn read_native_thread_cpu_nanos(_: usize) -> Option<u64> {
    None
}

fn nanos_to_secs(nanos: u64) -> f64 {
    nanos as f64 / 1_000_000_000.0
}

fn nanos_per_slice_to_us(nanos: u64, slices: u64) -> f64 {
    if slices == 0 {
        0.0
    } else {
        nanos as f64 / slices as f64 / 1_000.0
    }
}

#[bench]
fn bench_yatp_fair_read_pool(b: &mut Bencher) {
    let config = BenchConfig::from_env();
    let pool = BenchPool::new(config);
    let report = run_workload(config, &pool, b);
    eprintln!("{}", report.format(config));
}