From 082367408f6e84d44218da4cde7c90b1b7e12119 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Tue, 1 Jul 2025 10:31:23 +0800 Subject: [PATCH 01/10] feat: add support --qps --- Cargo.lock | 4 +- Cargo.toml | 2 + python/resp_benchmark/__init__.py | 2 +- python/resp_benchmark/cli.py | 5 +- python/resp_benchmark/wrapper.py | 58 +- src/async_flag.rs | 14 +- src/auto_connection.rs | 2 +- src/bench.rs | 159 ++- src/histogram.rs | 7 +- src/lib.rs | 103 +- src/qps_limiter.rs | 1674 +++++++++++++++++++++++++++++ src/shared_context.rs | 25 +- 12 files changed, 2025 insertions(+), 30 deletions(-) create mode 100644 src/qps_limiter.rs diff --git a/Cargo.lock b/Cargo.lock index d1ffc82..887ca75 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" @@ -679,6 +679,8 @@ dependencies = [ "ctrlc", "enum_delegate", "nom", + "parking_lot", + "pin-project-lite", "pyo3", "rand", "redis", diff --git a/Cargo.toml b/Cargo.toml index b6444b3..7a85f79 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,3 +27,5 @@ colored = "2.1.0" enum_delegate = "0.2.0" ctrlc = "3.4.4" urlencoding = "2.1.3" +parking_lot = "0.12.1" +pin-project-lite = "0.2.14" diff --git a/python/resp_benchmark/__init__.py b/python/resp_benchmark/__init__.py index d7716ab..bd9ec83 100644 --- a/python/resp_benchmark/__init__.py +++ b/python/resp_benchmark/__init__.py @@ -1 +1 @@ -from .wrapper import Benchmark, Result +from .wrapper import Benchmark, Result, BenchmarkWorkerPool diff --git a/python/resp_benchmark/cli.py b/python/resp_benchmark/cli.py index 2720a22..e6b147a 100644 --- a/python/resp_benchmark/cli.py +++ b/python/resp_benchmark/cli.py @@ -18,6 +18,7 @@ def parse_args(): parser.add_argument("--cores", type=str, default=f"", help="Comma-separated list of CPU cores to use (default all)") parser.add_argument("--cluster", action="store_true", help="Use cluster mode (default false)") parser.add_argument("-n", metavar="requests", type=int, default=0, help="Total number of requests (default 0), 0 for unlimited.") + parser.add_argument("-t", metavar="target", type=int, default=0, help="Target number of requests per second, 0 for unlimited, recommended to be multiple of 100.") parser.add_argument("-s", metavar="seconds", type=int, default=0, help="Total time in seconds (default 0), 0 for unlimited.") parser.add_argument("-P", metavar="pipeline", type=int, default=1, help="Pipeline requests. Default 1 (no pipeline).") # parser.add_argument("--tls", action="store_true", help="Use TLS for connection (default false)") @@ -34,9 +35,9 @@ def main(): args = parse_args() bm = Benchmark(host=args.h, port=args.p, username=args.u, password=args.a, cluster=args.cluster, cores=args.cores, timeout=30) if args.load: - bm.load_data(command=args.command, connections=args.c, pipeline=args.P, count=args.n) + bm.load_data(command=args.command, connections=args.c, pipeline=args.P, count=args.n, target=args.t) else: - bm.bench(command=args.command, connections=args.c, pipeline=args.P, count=args.n, seconds=args.s) + bm.bench(command=args.command, connections=args.c, pipeline=args.P, count=args.n, target=args.t, seconds=args.s) if __name__ == "__main__": diff --git a/python/resp_benchmark/wrapper.py b/python/resp_benchmark/wrapper.py index 3cfd1fa..63d171f 100644 --- a/python/resp_benchmark/wrapper.py +++ b/python/resp_benchmark/wrapper.py @@ -7,6 +7,7 @@ from .cores import parse_cores_string +from ._resp_benchmark_rust_lib import BenchmarkWorkerPool @dataclass class Result: @@ -22,6 +23,7 @@ class Result: qps: float avg_latency_ms: float p99_latency_ms: float + max_latency_ms: float connections: int @@ -69,6 +71,7 @@ def bench( connections: int = 0, pipeline: int = 1, count: int = 0, + target: int = 0, seconds: int = 0, quiet: bool = False, ) -> Result: @@ -80,6 +83,7 @@ def bench( connections (int): The number of parallel connections. pipeline (int): The number of commands to pipeline. count (int): The total number of requests to make. + target(int): The number of requests to send per second. seconds (int): The duration of the test in seconds. quiet: (bool): Whether to suppress output. Returns: @@ -100,6 +104,7 @@ def bench( connections=connections, pipeline=pipeline, count=count, + target=target, seconds=seconds, load=False, quiet=quiet, @@ -108,18 +113,20 @@ def bench( qps=ret.qps, avg_latency_ms=ret.avg_latency_ms, p99_latency_ms=ret.p99_latency_ms, + max_latency_ms=ret.max_latency_ms, connections=ret.connections ) return result - def load_data(self, command: str, count: int, connections: int = 128, pipeline: int = 10, quiet: bool = False): + def load_data(self, command: str, count: int, target: int = 0, connections: int = 128, pipeline: int = 10, quiet: bool = False): """ Load data into the Redis server using the specified command. Args: command (str): The Redis command to use for loading data. count (int): The total number of requests to make. + target(int): The number of requests to send per second. connections (int): The number of parallel connections. pipeline (int): The number of commands to pipeline quiet: (bool): Whether to suppress output. @@ -140,6 +147,7 @@ def load_data(self, command: str, count: int, connections: int = 128, pipeline: connections=connections, pipeline=pipeline, count=count, + target=target, seconds=0, load=True, quiet=quiet, @@ -151,3 +159,51 @@ def flushall(self): """ r = redis.Redis(host=self.host, port=self.port, username=self.username, password=self.password) r.flushall() + + class AsyncBenchmarkContext: + def __init__(self, raw_benchmark_context): + self.raw_benchmark_context = raw_benchmark_context + + def stop(self): + self.raw_benchmark_context.stop() + + def join(self): + self.raw_benchmark_context.join() + + def try_join(self): + return self.raw_benchmark_context.try_join() + + def current_result(self) -> Result: + ret = self.raw_benchmark_context.current_result() + return Result( + qps=ret.qps, + avg_latency_ms=ret.avg_latency_ms, + p99_latency_ms=ret.p99_latency_ms, + max_latency_ms=ret.max_latency_ms, + connections=ret.connections + ) + + def async_bench(self, benchmark_pool: BenchmarkWorkerPool, command: str, connections: int = 0, pipeline: int = 1, count: int = 0, target: int = 0, seconds: int = 0, quiet: bool = False) -> AsyncBenchmarkContext: + from . import _resp_benchmark_rust_lib + raw_ctx = _resp_benchmark_rust_lib.async_benchmark( + worker_pool=benchmark_pool, + + host=self.host, + port=self.port, + username=self.username, + password=self.password, + cluster=self.cluster, + tls=False, # TODO: Implement TLS support + timeout=self.timeout, + cores=self.cores, + + command=command, + connections=connections, + pipeline=pipeline, + count=count, + target=target, + seconds=seconds, + load=False, + quiet=quiet, + ) + return self.AsyncBenchmarkContext(raw_ctx) diff --git a/src/async_flag.rs b/src/async_flag.rs index 139cf1e..9d15869 100644 --- a/src/async_flag.rs +++ b/src/async_flag.rs @@ -1,23 +1,32 @@ -use std::sync::Arc; +use std::sync::{atomic::AtomicBool, Arc}; use tokio::sync::watch; pub struct AsyncFlag { receiver: watch::Receiver, sender: Arc>, + flag: Arc, } impl AsyncFlag { pub fn new() -> Self { let (sender, receiver) = watch::channel(false); - AsyncFlag { receiver, sender: Arc::new(sender) } + AsyncFlag { receiver, sender: Arc::new(sender), flag: Arc::new(AtomicBool::new(false)) } } pub async fn wait_flag(&mut self) { + if self.flag.load(std::sync::atomic::Ordering::SeqCst) { + return; + } self.receiver.changed().await.unwrap(); } pub fn set_flag(&self) { self.sender.send(true).unwrap(); + self.flag.store(true, std::sync::atomic::Ordering::SeqCst); + } + + pub fn flag(&self) -> bool { + self.flag.load(std::sync::atomic::Ordering::SeqCst) } } @@ -26,6 +35,7 @@ impl Clone for AsyncFlag { Self { receiver: self.receiver.clone(), sender: self.sender.clone(), + flag: self.flag.clone(), } } } diff --git a/src/auto_connection.rs b/src/auto_connection.rs index 40645fe..9a2b288 100644 --- a/src/auto_connection.rs +++ b/src/auto_connection.rs @@ -36,7 +36,7 @@ impl ConnLimiter { if active_conn >= target_conn { continue; } - let old_value = self.active_conn.compare_exchange(active_conn, active_conn + 1, std::sync::atomic::Ordering::SeqCst, std::sync::atomic::Ordering::SeqCst).unwrap(); + let old_value = self.active_conn.compare_exchange(active_conn, active_conn + 1, std::sync::atomic::Ordering::SeqCst, std::sync::atomic::Ordering::SeqCst).map_or_else(|e| e, |v| v); if old_value != active_conn { continue; } diff --git a/src/bench.rs b/src/bench.rs index dd73f05..7a7bf8e 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -1,37 +1,41 @@ -use std::io::Write; -use std::sync::Arc; use awaitgroup::WaitGroup; use colored::Colorize; -use tokio::{select, task}; +use std::io::Write; +use std::sync::Arc; +use tokio::{select}; -use crate::BenchmarkResult; +use crate::auto_connection::{AutoConnection, ConnLimiter}; use crate::client::ClientConfig; use crate::command::Command; -use crate::auto_connection::{AutoConnection, ConnLimiter}; use crate::shared_context::SharedContext; +use crate::BenchmarkResult; +use crate::AsyncBenchmarkContext; +use crate::qps_limiter::RateLimiter; #[derive(Clone)] pub struct Case { pub command: Command, pub connections: u64, pub count: u64, + pub target: u64, pub seconds: u64, pub pipeline: u64, } -async fn run_commands_on_single_thread(limiter: Arc, config: ClientConfig, case: Case, context: SharedContext) { - let local = task::LocalSet::new(); - for _ in 0..limiter.total_conn { - let limiter = limiter.clone(); +async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limiter: Arc>, config: ClientConfig, case: Case, context: SharedContext) { + let mut local = vec![]; + for _ in 0..conn_limiter.total_conn { + let conn_limiter = conn_limiter.clone(); + let qps_limiter = qps_limiter.clone(); let config = config.clone(); let case = case.clone(); let mut context = context.clone(); - local.spawn_local(async move { + let join_handle = tokio::task::spawn(async move { let mut client = config.get_client().await; let mut cmd = case.command.clone(); - let limiter = limiter.clone(); + let conn_limiter = conn_limiter.clone(); select! { - _ = limiter.wait_new_conn() =>{} + _ = conn_limiter.wait_new_conn() =>{} _ = context.wait_stop() => { return; } @@ -56,12 +60,21 @@ async fn run_commands_on_single_thread(limiter: Arc, config: Client client.run_commands(p).await; let duration = instant.elapsed().as_micros() as u64; for _ in 0..pipeline_cnt { - context.histogram.record(duration); + context.record(duration); + } + match qps_limiter.as_ref() { + Some(limiter) => { + limiter.acquire(pipeline_cnt as usize).await; + } + None => {} } } }); + local.push(join_handle); + } + for join_handle in local { + let _ = join_handle.await; } - local.await; } fn wait_finish(case: &Case, mut auto_connection: AutoConnection, mut context: SharedContext, mut wg: WaitGroup, quiet: bool) -> BenchmarkResult { @@ -124,6 +137,7 @@ fn wait_finish(case: &Case, mut auto_connection: AutoConnection, mut context: Sh }; result.avg_latency_ms = histogram.avg() as f64 / 1_000.0; result.p99_latency_ms = histogram.percentile(0.99) as f64 / 1_000.0; + result.max_latency_ms = histogram.max() as f64 / 1_000.0; result.connections = conn; }); return result; @@ -134,6 +148,7 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo println!("{}: {}", "command".bold().blue(), case.command.to_string().green().bold()); println!("{}: {}", "connections".bold().blue(), if case.connections == 0 { "auto".to_string() } else { case.connections.to_string() }); println!("{}: {}", "count".bold().blue(), case.count); + println!("{}: {}", "target".bold().blue(), if case.target == 0 { "unlimited".to_string() } else { case.target.to_string() }); println!("{}: {}", "seconds".bold().blue(), case.seconds); println!("{}: {}", "pipeline".bold().blue(), case.pipeline); } @@ -141,6 +156,19 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo // calc connections let auto_connection = AutoConnection::new(case.connections, cores.len() as u64); + // calc target qps + let qps_limiter = Arc::new(if case.target > 0 { + Some( + RateLimiter::builder() + .max(case.target as usize * 5) + .interval(tokio::time::Duration::from_millis(1)) + .refill(case.target as usize / 1000) + .build() + ) + } else { + None + }); + let mut thread_handlers = Vec::new(); let wg = WaitGroup::new(); let core_ids = core_affinity::get_core_ids().unwrap(); @@ -151,12 +179,13 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo let context = context.clone(); let wk = wg.worker(); let core_id = core_ids[cores[inx] as usize]; - let limiter = auto_connection.limiters[inx].clone(); + let conn_limiter = auto_connection.limiters[inx].clone(); + let qps_limiter = qps_limiter.clone(); let thread_handler = std::thread::spawn(move || { core_affinity::set_for_current(core_id); // not work on Apple Silicon let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); rt.block_on(async { - run_commands_on_single_thread(limiter, client_config, case, context).await; + run_commands_on_single_thread(conn_limiter, qps_limiter, client_config, case, context).await; wk.done(); }); }); @@ -178,4 +207,100 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo } return result; -} \ No newline at end of file +} + +async fn async_cron(mut auto_connection: AutoConnection, mut context: SharedContext, mut wg: WaitGroup) { + let histogram = context.histogram.clone(); + // calc overall qps + // for log + let mut interval = tokio::time::interval(std::time::Duration::from_secs(1)); + + let mut log_instance = std::time::Instant::now(); + let mut log_last_cnt = histogram.cnt(); + + if auto_connection.ready { + context.start_timer(); + } + + loop { + select! { + _ = interval.tick() => {} + _ = wg.wait() => {break;} + } + let cnt = histogram.cnt(); + let qps = (cnt - log_last_cnt) as f64 / log_instance.elapsed().as_secs_f64(); + if !auto_connection.ready { + auto_connection.adjust(&histogram); + if auto_connection.ready { + context.start_timer(); + } + } + log_last_cnt = cnt; + log_instance = std::time::Instant::now(); + + let mut result = BenchmarkResult::default(); + let current_histogram_id = context.latest_histogramactive_index.fetch_add(1, std::sync::atomic::Ordering::SeqCst) % 2; + let current_histogram = &context.latest_histogram[current_histogram_id as usize]; + result.qps = qps; + result.connections = auto_connection.active_conn(); + result.avg_latency_ms = current_histogram.avg() as f64 / 1_000.0; + result.p99_latency_ms = current_histogram.percentile(0.99) as f64 / 1_000.0; + result.max_latency_ms = current_histogram.max() as f64 / 1_000.0; + *context.latest_result.lock().unwrap() = result; + } +} + +pub fn do_benchmark_async(pool: Arc, client_config: ClientConfig, cores: Vec, case: Case, load: bool, quiet: bool) -> crate::AsyncBenchmarkContext { + if !quiet { + println!("{}: {}", "command".bold().blue(), case.command.to_string().green().bold()); + println!("{}: {}", "connections".bold().blue(), if case.connections == 0 { "auto".to_string() } else { case.connections.to_string() }); + println!("{}: {}", "count".bold().blue(), case.count); + println!("{}: {}", "target".bold().blue(), if case.target == 0 { "unlimited".to_string() } else { case.target.to_string() }); + println!("{}: {}", "seconds".bold().blue(), case.seconds); + println!("{}: {}", "pipeline".bold().blue(), case.pipeline); + } + + let n_parallel = std::cmp::max(1, std::cmp::min(cores.len(), case.connections as usize)); + + // calc connections + let auto_connection = AutoConnection::new(case.connections, n_parallel as u64); + + // calc target qps + let qps_limiter = Arc::new(if case.target > 0 { + Some( + RateLimiter::builder() + .max(case.target as usize * 5) + .interval(tokio::time::Duration::from_millis(1)) + .refill(case.target as usize / 1000) + .build() + ) + } else { + None + }); + + let wg = WaitGroup::new(); + + let context = SharedContext::new(case.count, case.seconds, load); + for inx in 0..n_parallel { + let client_config = client_config.clone(); + let case = case.clone(); + let context = context.clone(); + let wk = wg.worker(); + let conn_limiter = auto_connection.limiters[inx].clone(); + let qps_limiter = qps_limiter.clone(); + pool.spawn(async move { + run_commands_on_single_thread(conn_limiter, qps_limiter, client_config, case, context).await; + wk.done(); + }); + } + + let cron_ctx = context.clone(); + let join_handle = pool.spawn(async move { + async_cron(auto_connection, cron_ctx, wg).await; + }); + + return AsyncBenchmarkContext { + ctx: context, + join_handle: Option::Some(join_handle), + }; +} diff --git a/src/histogram.rs b/src/histogram.rs index 815556b..1bec81f 100644 --- a/src/histogram.rs +++ b/src/histogram.rs @@ -103,6 +103,10 @@ impl Histogram { 0 } + pub fn max(&self) -> u64 { + return self.percentile(1.0); + } + fn humanize_us(latency_us: u64) -> String { match latency_us { 0 => "<0.01ms".to_string(), @@ -124,8 +128,9 @@ impl Display for Histogram { } let avg = self.avg(); let p99 = self.percentile(0.99); + let max = self.max(); - write!(f, "cnt: {}, avg: {}, p99: {}", cnt, Histogram::humanize_us(avg), Histogram::humanize_us(p99)) + write!(f, "cnt: {}, avg: {}, p99: {}, max: {}", cnt, Histogram::humanize_us(avg), Histogram::humanize_us(p99), Histogram::humanize_us(max)) } } diff --git a/src/lib.rs b/src/lib.rs index c6232b4..5586190 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,32 +4,38 @@ mod command; mod auto_connection; mod shared_context; mod histogram; +mod qps_limiter; mod async_flag; use ctrlc; use pyo3::prelude::*; use pyo3::wrap_pyfunction; use crate::command::Command; +use crate::shared_context::SharedContext; /// A Python module implemented in Rust. #[pymodule] fn _resp_benchmark_rust_lib(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(benchmark, m)?)?; + m.add_function(wrap_pyfunction!(async_benchmark, m)?)?; + m.add_class::()?; Ok(()) } #[pyclass] -#[derive(Default)] +#[derive(Default, Copy, Clone)] struct BenchmarkResult { #[pyo3(get, set)] pub qps: f64, #[pyo3(get, set)] pub avg_latency_ms: f64, #[pyo3(get, set)] pub p99_latency_ms: f64, + #[pyo3(get, set)] pub max_latency_ms: f64, #[pyo3(get, set)] pub connections: u64, } #[pyfunction] fn benchmark( + py: Python<'_>, host: String, port: u16, username: String, @@ -40,6 +46,7 @@ fn benchmark( cores: Vec, command: String, connections: u64, + target: u64, pipeline: u64, count: u64, seconds: u64, @@ -68,8 +75,100 @@ fn benchmark( connections, pipeline, count, + target, seconds, }; - let result = bench::do_benchmark(client_config, cores, case, load, quiet); + let result = py.allow_threads(|| bench::do_benchmark(client_config, cores, case, load, quiet)); + Ok(result) +} + +#[pyclass] +struct BenchmarkWorkerPool { + pool: std::sync::Arc, +} + +#[pymethods] +impl BenchmarkWorkerPool { + #[new] + fn new() -> PyResult { + let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build().map_err(|e| PyErr::new::(e.to_string()))?; + Ok(Self { pool: std::sync::Arc::new(rt) }) + } +} + +#[pyclass] +pub struct AsyncBenchmarkContext { + ctx: SharedContext, + join_handle: Option>, +} + +#[pymethods] +impl AsyncBenchmarkContext { + fn stop(&mut self) { + self.ctx.stop(); + } + + fn join(&mut self) { + match self.join_handle.take() { + None => {} + Some(handle) => { + let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); + let _ = rt.block_on(handle); + } + } + } + + fn try_join(&mut self) -> bool { + self.join_handle.as_ref().map_or(true, |v| v.is_finished()) + } + + fn current_result(&self) -> BenchmarkResult { + let result = self.ctx.latest_result.lock().unwrap().clone(); + return result; + } +} + +#[pyfunction] +fn async_benchmark( + worker_pool: &BenchmarkWorkerPool, + host: String, + port: u16, + username: String, + password: String, + cluster: bool, + tls: bool, + timeout: u64, + cores: Vec, + command: String, + connections: u64, + target: u64, + pipeline: u64, + count: u64, + seconds: u64, + load: bool, + quiet: bool, +) -> PyResult { + if load { + return Err(PyErr::new::("count must be greater than 0".to_string())); + } + + let client_config = client::ClientConfig { + cluster, + address: format!("{}:{}", host, port), + username, + password, + tls, + timeout, + }; + let case = bench::Case { + command: Command::new(command.as_str()), + connections, + pipeline, + count, + target, + seconds, + }; + let pool = worker_pool.pool.clone(); + let result = bench::do_benchmark_async(pool, client_config, cores, case, load, quiet); Ok(result) } \ No newline at end of file diff --git a/src/qps_limiter.rs b/src/qps_limiter.rs new file mode 100644 index 0000000..5d40689 --- /dev/null +++ b/src/qps_limiter.rs @@ -0,0 +1,1674 @@ +// Copyright (c) 2019 John-John Tedro + +// Permission is hereby granted, free of charge, to any +// person obtaining a copy of this software and associated +// documentation files (the "Software"), to deal in the +// Software without restriction, including without +// limitation the rights to use, copy, modify, merge, +// publish, distribute, sublicense, and/or sell copies of +// the Software, and to permit persons to whom the Software +// is furnished to do so, subject to the following +// conditions: + +// The above copyright notice and this permission notice +// shall be included in all copies or substantial portions +// of the Software. + +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWAREo +#![allow(dead_code)] +use std::cell::UnsafeCell; +use std::convert::TryFrom as _; +use std::fmt; +use std::future::Future; +use std::mem::{self, ManuallyDrop}; +use std::ops::{Deref, DerefMut}; +use std::pin::Pin; +use std::ptr; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll, Waker}; +use std::marker; +use std::ops; +use std::sync::Arc; + +use parking_lot::{Mutex, MutexGuard}; +use pin_project_lite::pin_project; +use tokio::time::{self, Duration, Instant}; + + +/// Default factor for how to calculate max refill value. +const DEFAULT_REFILL_MAX_FACTOR: usize = 10; + +/// Interval to bump the shared mutex guard to allow other parts of the system +/// to make process. Processes which loop should use this number to determine +/// how many times it should loop before calling [Guard::bump]. +/// +/// If we do not respect this limit we might inadvertently end up starving other +/// tasks from making progress so that they can unblock. +const BUMP_LIMIT: usize = 16; + +/// The maximum supported balance. +const MAX_BALANCE: usize = isize::MAX as usize; + +/// Marker trait which indicates that a type represents a unique held critical section. +trait IsCritical {} +impl IsCritical for Critical {} +impl IsCritical for Guard<'_> {} + +/// Linked task state. +struct Task { + /// Remaining tokens that need to be satisfied. + remaining: usize, + /// If this node has been released or not. We make this an atomic to permit + /// access to it without synchronization. + complete: AtomicBool, + /// The waker associated with the node. + waker: Option, +} + +impl Task { + /// Construct a new task state with the given permits remaining. + const fn new() -> Self { + Self { + remaining: 0, + complete: AtomicBool::new(false), + waker: None, + } + } + + /// Test if the current node is completed. + fn is_completed(&self) -> bool { + self.remaining == 0 + } + + /// Fill the current node from the given pool of tokens and modify it. + fn fill(&mut self, current: &mut usize) { + let removed = usize::min(self.remaining, *current); + self.remaining -= removed; + *current -= removed; + } +} + +/// A borrowed rate limiter. +struct BorrowedRateLimiter<'a>(&'a RateLimiter); + +impl Deref for BorrowedRateLimiter<'_> { + type Target = RateLimiter; + + #[inline] + fn deref(&self) -> &RateLimiter { + self.0 + } +} + +struct Critical { + /// Waiter list. + waiters: LinkedList, + /// The deadline for when more tokens can be be added. + deadline: Instant, +} + +#[repr(transparent)] +struct Guard<'a> { + critical: MutexGuard<'a, Critical>, +} + +impl Guard<'_> { + #[inline] + fn bump(this: &mut Guard<'_>) { + MutexGuard::bump(&mut this.critical) + } +} + +impl Deref for Guard<'_> { + type Target = Critical; + + #[inline] + fn deref(&self) -> &Critical { + &self.critical + } +} + +impl DerefMut for Guard<'_> { + #[inline] + fn deref_mut(&mut self) -> &mut Critical { + &mut self.critical + } +} + +impl Critical { + #[inline] + fn push_task_front(&mut self, task: &mut Node) { + // SAFETY: We both have mutable access to the node being pushed, and + // mutable access to the critical section through `self`. So we know we + // have exclusive tampering rights to the waiter queue. + unsafe { + self.waiters.push_front(task.into()); + } + } + + #[inline] + fn push_task(&mut self, task: &mut Node) { + // SAFETY: We both have mutable access to the node being pushed, and + // mutable access to the critical section through `self`. So we know we + // have exclusive tampering rights to the waiter queue. + unsafe { + self.waiters.push_back(task.into()); + } + } + + #[inline] + fn remove_task(&mut self, task: &mut Node) { + // SAFETY: We both have mutable access to the node being pushed, and + // mutable access to the critical section through `self`. So we know we + // have exclusive tampering rights to the waiter queue. + unsafe { + self.waiters.remove(task.into()); + } + } + + /// Release the current core. Beyond this point the current task may no + /// longer interact exclusively with the core. + fn release(&mut self, state: &mut State<'_>) { + state.available = true; + + // Find another task that might take over as core. Once it has acquired + // core status it will have to make sure it is no longer linked into the + // wait queue. + unsafe { + if let Some(node) = self.waiters.front() { + if let Some(ref waker) = node.as_ref().waker { + waker.wake_by_ref(); + } + } + } + } +} + +#[derive(Debug)] +struct State<'a> { + /// Original state. + state: usize, + /// If the core is available or not. + available: bool, + /// The balance. + balance: usize, + /// The rate limiter the state is associated with. + lim: &'a RateLimiter, +} + +impl<'a> State<'a> { + fn try_fast_path(mut self, permits: usize) -> bool { + let mut attempts = 0; + + // Fast path where we just try to nab any available permit without + // locking. + // + // We do have to race against anyone else grabbing permits here when + // storing the state back. + while self.balance >= permits { + // Abandon fast path if we've tried too many times. + if attempts == BUMP_LIMIT { + break; + } + + self.balance -= permits; + + if let Err(new_state) = self.try_save() { + self = new_state; + attempts += 1; + continue; + } + + return true; + } + + false + } + + /// Add tokens and release any pending tasks. + #[inline] + fn add_tokens(&mut self, critical: &mut Guard<'_>, tokens: usize, f: F) -> O + where + F: FnOnce(&mut Guard<'_>, &mut State) -> O, + { + if tokens > 0 { + debug_assert!( + tokens <= MAX_BALANCE, + "Additional tokens {} must be less than {}", + tokens, + MAX_BALANCE + ); + + self.balance = (self.balance + tokens).min(self.lim.max); + drain_wait_queue(critical, self); + let output = f(critical, self); + return output; + } + + f(critical, self) + } + + #[inline] + fn decode(state: usize, lim: &'a RateLimiter) -> Self { + State { + state, + available: state & 1 == 1, + balance: state >> 1, + lim, + } + } + + #[inline] + fn encode(&self) -> usize { + (self.balance << 1) | usize::from(self.available) + } + + /// Try to save the state, but only succeed if it hasn't been modified. + #[inline] + fn try_save(self) -> Result<(), Self> { + let this = ManuallyDrop::new(self); + + match this.lim.state.compare_exchange( + this.state, + this.encode(), + Ordering::Release, + Ordering::Relaxed, + ) { + Ok(_) => Ok(()), + Err(state) => Err(State::decode(state, this.lim)), + } + } +} + +impl Drop for State<'_> { + #[inline] + fn drop(&mut self) { + self.lim.state.store(self.encode(), Ordering::Release); + } +} + +/// A token-bucket rate limiter. +pub struct RateLimiter { + /// Tokens to add every `per` duration. + refill: usize, + /// Interval in milliseconds to add tokens. + interval: Duration, + /// Max number of tokens associated with the rate limiter. + max: usize, + /// If the rate limiter is fair or not. + fair: bool, + /// The state of the rate limiter. + state: AtomicUsize, + /// Critical state of the rate limiter. + critical: Mutex, +} + +impl RateLimiter { + /// Construct a new [`Builder`] for a [`RateLimiter`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use tokio::time::Duration; + /// + /// let limiter = RateLimiter::builder() + /// .initial(100) + /// .refill(100) + /// .max(1000) + /// .interval(Duration::from_millis(250)) + /// .fair(false) + /// .build(); + /// ``` + pub fn builder() -> Builder { + Builder::default() + } + + /// Get the refill amount of this rate limiter as set through + /// [`Builder::refill`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .refill(1024) + /// .build(); + /// + /// assert_eq!(limiter.refill(), 1024); + /// ``` + pub fn refill(&self) -> usize { + self.refill + } + + /// Get the refill interval of this rate limiter as set through + /// [`Builder::interval`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use tokio::time::Duration; + /// + /// let limiter = RateLimiter::builder() + /// .interval(Duration::from_millis(1000)) + /// .build(); + /// + /// assert_eq!(limiter.interval(), Duration::from_millis(1000)); + /// ``` + pub fn interval(&self) -> Duration { + self.interval + } + + /// Get the max value of this rate limiter as set through [`Builder::max`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .max(1024) + /// .build(); + /// + /// assert_eq!(limiter.max(), 1024); + /// ``` + pub fn max(&self) -> usize { + self.max + } + + /// Test if the current rate limiter is fair as specified through + /// [`Builder::fair`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .fair(true) + /// .build(); + /// + /// assert_eq!(limiter.is_fair(), true); + /// ``` + pub fn is_fair(&self) -> bool { + self.fair + } + + /// Get the current token balance. + /// + /// This indicates how many tokens can be requested without blocking. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = RateLimiter::builder() + /// .initial(100) + /// .build(); + /// + /// assert_eq!(limiter.balance(), 100); + /// limiter.acquire(10).await; + /// assert_eq!(limiter.balance(), 90); + /// # } + /// ``` + pub fn balance(&self) -> usize { + self.state.load(Ordering::Acquire) >> 1 + } + + /// Acquire a single permit. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = RateLimiter::builder() + /// .initial(10) + /// .build(); + /// + /// limiter.acquire_one().await; + /// # } + /// ``` + pub fn acquire_one(&self) -> Acquire<'_> { + self.acquire(1) + } + + /// Acquire the given number of permits, suspending the current task until + /// they are available. + /// + /// If zero permits are specified, this function never suspends the current + /// task. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = RateLimiter::builder() + /// .initial(10) + /// .build(); + /// + /// limiter.acquire(10).await; + /// # } + /// ``` + pub fn acquire(&self, permits: usize) -> Acquire<'_> { + Acquire { + inner: AcquireFut::new(BorrowedRateLimiter(self), permits), + } + } + + /// Try to acquire the given number of permits, returning `true` if the + /// given number of permits were successfully acquired. + /// + /// If the scheduler is fair, and there are pending tasks waiting to acquire + /// tokens this method will return `false`. + /// + /// If zero permits are specified, this method returns `true`. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use tokio::time; + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = RateLimiter::builder().refill(1).initial(1).build(); + /// + /// assert!(limiter.try_acquire(1)); + /// assert!(!limiter.try_acquire(1)); + /// assert!(limiter.try_acquire(0)); + /// + /// time::sleep(limiter.interval() * 2).await; + /// + /// assert!(limiter.try_acquire(1)); + /// assert!(limiter.try_acquire(1)); + /// assert!(!limiter.try_acquire(1)); + /// # } + /// ``` + pub fn try_acquire(&self, permits: usize) -> bool { + if self.try_fast_path(permits) { + return true; + } + + let mut critical = self.lock(); + + // Reload the state while we are under the critical lock, this + // ensures that the `available` flag is up-to-date since it is only + // ever modified while holding the critical lock. + let mut state = self.take(); + + // The core is *not* available, which also implies that there are tasks + // ahead which are busy. + if !state.available { + return false; + } + + let now = Instant::now(); + + // Here we try to assume core duty temporarily to see if we can + // release a sufficient number of tokens to allow the current task + // to proceed. + if let Some((tokens, deadline)) = self.calculate_drain(critical.deadline, now) { + state.balance = (state.balance + tokens).min(self.max); + critical.deadline = deadline; + } + + if state.balance >= permits { + state.balance -= permits; + return true; + } + + false + } + + /// Acquire a permit using an owned future. + /// + /// If zero permits are specified, this function never suspends the current + /// task. + /// + /// This required the [`RateLimiter`] to be wrapped inside of an + /// [`std::sync::Arc`] but will in contrast permit the acquire operation to + /// be owned by another struct making it more suitable for embedding. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use std::sync::Arc; + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = Arc::new(RateLimiter::builder().initial(10).build()); + /// + /// limiter.acquire_owned(10).await; + /// # } + /// ``` + /// + /// Example when embedded into another future. This wouldn't be possible + /// with [`RateLimiter::acquire`] since it would otherwise hold a reference + /// to the corresponding [`RateLimiter`] instance. + /// + /// ``` + /// use std::future::Future; + /// use std::pin::Pin; + /// use std::sync::Arc; + /// use std::task::{Context, Poll}; + /// + /// use leaky_bucket::{AcquireOwned, RateLimiter}; + /// use pin_project::pin_project; + /// + /// #[pin_project] + /// struct MyFuture { + /// limiter: Arc, + /// #[pin] + /// acquire: Option, + /// } + /// + /// impl Future for MyFuture { + /// type Output = (); + /// + /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + /// let mut this = self.project(); + /// + /// loop { + /// if let Some(acquire) = this.acquire.as_mut().as_pin_mut() { + /// futures::ready!(acquire.poll(cx)); + /// return Poll::Ready(()); + /// } + /// + /// this.acquire.set(Some(this.limiter.clone().acquire_owned(100))); + /// } + /// } + /// } + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = Arc::new(RateLimiter::builder().initial(100).build()); + /// + /// let future = MyFuture { limiter, acquire: None }; + /// future.await; + /// # } + /// ``` + pub fn acquire_owned(self: Arc, permits: usize) -> AcquireOwned { + AcquireOwned { + inner: AcquireFut::new(self, permits), + } + } + + /// Lock the critical section of the rate limiter and return the associated guard. + fn lock(&self) -> Guard<'_> { + Guard { + critical: self.critical.lock(), + } + } + + /// Load the current state. + fn load(&self) -> State<'_> { + State::decode(self.state.load(Ordering::Acquire), self) + } + + /// Take the current state, leaving the core state intact. + fn take(&self) -> State<'_> { + State::decode(self.state.swap(0, Ordering::Acquire), self) + } + + /// Try to use fast path. + fn try_fast_path(&self, permits: usize) -> bool { + if permits == 0 { + return true; + } + + if self.fair { + return false; + } + + self.load().try_fast_path(permits) + } + + /// Calculate refill amount. Returning a tuple of how much to fill and remaining + /// duration to sleep until the next refill time if appropriate. + /// + /// The maximum number of additional tokens this method will ever return is + /// limited to [`MAX_BALANCE`] to ensure that addition with an existing + /// balance will never overflow. + fn calculate_drain(&self, deadline: Instant, now: Instant) -> Option<(usize, Instant)> { + if now < deadline { + return None; + } + + // Time elapsed in milliseconds since the last deadline. + let millis = self.interval.as_millis(); + let since = now.saturating_duration_since(deadline).as_millis(); + + let periods = usize::try_from(since / millis + 1).unwrap_or(usize::MAX); + + let tokens = periods + .checked_mul(self.refill) + .unwrap_or(MAX_BALANCE) + .min(MAX_BALANCE); + + let rem = u64::try_from(since % millis).unwrap_or(u64::MAX); + + // Calculated time remaining until the next deadline. + let next = millis as u64 * periods as u64 - rem; + let deadline = deadline.checked_add(time::Duration::from_millis(next)).unwrap(); + + Some((tokens, deadline)) + } +} + +impl fmt::Debug for RateLimiter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RateLimiter") + .field("refill", &self.refill) + .field("interval", &self.interval) + .field("max", &self.max) + .field("fair", &self.fair) + .finish_non_exhaustive() + } +} + +/// Refill the wait queue with the given number of tokens. +fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) { + let mut bump = 0; + + // SAFETY: we're holding the lock guard to all the waiters so we can be + // sure that we have exclusive access to the wait queue. + unsafe { + while state.balance > 0 { + let mut node = match critical.waiters.pop_back() { + Some(node) => node, + None => break, + }; + + let n = node.as_mut(); + n.fill(&mut state.balance); + + if !n.is_completed() { + critical.waiters.push_back(node); + break; + } + + n.complete.store(true, Ordering::Release); + + if let Some(waker) = n.waker.take() { + waker.wake(); + } + + bump += 1; + + if bump == BUMP_LIMIT { + Guard::bump(critical); + bump = 0; + } + } + } +} + +// SAFETY: All the internals of acquire is thread safe and correctly +// synchronized. The embedded waiter queue doesn't have anything inherently +// unsafe in it. +unsafe impl Send for RateLimiter {} +unsafe impl Sync for RateLimiter {} + +/// A builder for a [`RateLimiter`]. +pub struct Builder { + /// The max number of tokens. + max: Option, + /// The initial count of tokens. + initial: usize, + /// Tokens to add every `per` duration. + refill: usize, + /// Interval to add tokens in milliseconds. + interval: Duration, + /// If the rate limiter is fair or not. + fair: bool, +} + +impl Builder { + /// Configure the max number of tokens to use. + /// + /// If unspecified, this will default to be 10 times the [`refill`] or the + /// [`initial`] value, whichever is largest. + /// + /// The maximum supported balance is limited to [`isize::MAX`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .max(10_000) + /// .build(); + /// ``` + /// + /// [`refill`]: Builder::refill + /// [`initial`]: Builder::initial + pub fn max(&mut self, max: usize) -> &mut Self { + self.max = Some(max); + self + } + + /// Configure the initial number of tokens to configure. The default value + /// is `0`. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .initial(10) + /// .build(); + /// ``` + pub fn initial(&mut self, initial: usize) -> &mut Self { + self.initial = initial; + self + } + + /// Configure the time duration between which we add [`refill`] number to + /// the bucket rate limiter. + /// + /// This is 100ms by default. + /// + /// # Panics + /// + /// This panics if the provided interval does not fit within the millisecond + /// bounds of a [usize] or is zero. + /// + /// ```should_panic + /// use leaky_bucket::RateLimiter; + /// use tokio::time::Duration; + /// + /// let limiter = RateLimiter::builder() + /// .interval(Duration::from_secs(u64::MAX)) + /// .build(); + /// ``` + /// + /// ```should_panic + /// use leaky_bucket::RateLimiter; + /// use tokio::time::Duration; + /// + /// let limiter = RateLimiter::builder() + /// .interval(Duration::from_millis(0)) + /// .build(); + /// ``` + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use tokio::time::Duration; + /// + /// let limiter = RateLimiter::builder() + /// .interval(Duration::from_millis(100)) + /// .build(); + /// ``` + /// + /// [`refill`]: Builder::refill + pub fn interval(&mut self, interval: Duration) -> &mut Self { + assert! { + interval.as_millis() != 0, + "interval must be non-zero", + }; + assert! { + u64::try_from(interval.as_millis()).is_ok(), + "interval must fit within a 64-bit integer" + }; + self.interval = interval; + self + } + + /// The number of tokens to add at each [`interval`] interval. The default + /// value is `1`. + /// + /// # Panics + /// + /// Panics if a refill amount of `0` is specified. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .refill(100) + /// .build(); + /// ``` + /// + /// [`interval`]: Builder::interval + pub fn refill(&mut self, refill: usize) -> &mut Self { + assert!(refill > 0, "refill amount cannot be zero"); + self.refill = refill; + self + } + + /// Configure the rate limiter to be fair. + /// + /// Fairness is enabled by deafult. + /// + /// Fairness ensures that tasks make progress in the order that they acquire + /// even when the rate limiter is under contention. An unfair scheduler + /// might have a higher total throughput. + /// + /// Fair scheduling also affects the behavior of + /// [`RateLimiter::try_acquire`] which will return `false` if there are any + /// pending tasks since they should be given priority. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// + /// let limiter = RateLimiter::builder() + /// .refill(100) + /// .fair(false) + /// .build(); + /// ``` + pub fn fair(&mut self, fair: bool) -> &mut Self { + self.fair = fair; + self + } + + /// Construct a new [`RateLimiter`]. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use tokio::time::Duration; + /// + /// let limiter = RateLimiter::builder() + /// .refill(100) + /// .interval(Duration::from_millis(200)) + /// .max(10_000) + /// .build(); + /// ``` + pub fn build(&self) -> RateLimiter { + let deadline = Instant::now() + self.interval; + + let initial = self.initial.min(MAX_BALANCE); + let refill = self.refill.min(MAX_BALANCE); + + let max = match self.max { + Some(max) => max.min(MAX_BALANCE), + None => refill + .max(initial) + .saturating_mul(DEFAULT_REFILL_MAX_FACTOR) + .min(MAX_BALANCE), + }; + + let initial = initial.min(max); + + RateLimiter { + refill, + interval: self.interval, + max, + fair: self.fair, + state: AtomicUsize::new(initial << 1 | 1), + critical: Mutex::new(Critical { + waiters: LinkedList::new(), + deadline, + }), + } + } +} + +/// Construct a new builder with default options. +/// +/// # Examples +/// +/// ``` +/// use leaky_bucket::Builder; +/// +/// let limiter = Builder::default().build(); +/// ``` +impl Default for Builder { + fn default() -> Self { + Self { + max: None, + initial: 0, + refill: 1, + interval: Duration::from_millis(100), + fair: true, + } + } +} + +/// The state of an acquire operation. +#[derive(Debug, Clone, Copy)] +enum AcquireFutState { + /// Initial unconfigured state. + Initial, + /// The acquire is waiting to be released by the core. + Waiting, + /// The operation is completed. + Complete, + /// The task is currently the core. + Core, +} + +/// Inner state and methods of the acquire. +#[repr(transparent)] +struct AcquireFutInner { + /// Aliased task state. + node: UnsafeCell>, +} + +impl AcquireFutInner { + const fn new() -> AcquireFutInner { + AcquireFutInner { + node: UnsafeCell::new(Node::new(Task::new())), + } + } + + /// Access the completion flag. + pub fn complete(&self) -> &AtomicBool { + // SAFETY: This is always safe to access since it's atomic. + unsafe { &*ptr::addr_of!((*self.node.get()).complete) } + } + + /// Get the underlying task mutably. + /// + /// We prove that the caller does indeed have mutable access to the node by + /// passing in a mutable reference to the critical section. + #[inline] + pub fn get_task<'crit, C>( + self: Pin<&'crit mut Self>, + critical: &'crit mut C, + ) -> (&'crit mut C, &'crit mut Node) + where + C: IsCritical, + { + // SAFETY: Caller has exclusive access to the critical section, since + // it's passed in as a mutable argument. We can also ensure that none of + // the borrows outlive the provided closure. + unsafe { (critical, &mut *self.node.get()) } + } + + /// Update the waiting state for this acquisition task. This might require + /// that we update the associated waker. + fn update(self: Pin<&mut Self>, critical: &mut Guard<'_>, waker: &Waker) { + let (critical, task) = self.get_task(critical); + + if !task.is_linked() { + critical.push_task_front(task); + } + + let new_waker = match task.waker { + None => true, + Some(ref w) => !w.will_wake(waker), + }; + + if new_waker { + task.waker = Some(waker.clone()); + } + } + + /// Ensure that the current core task is correctly linked up if needed. + fn link_core(self: Pin<&mut Self>, critical: &mut Critical, lim: &RateLimiter) { + let (critical, task) = self.get_task(critical); + + match (lim.fair, task.is_linked()) { + (true, false) => { + // Fair scheduling needs to ensure that the core is part of the wait + // queue, and will be woken up in-order with other tasks. + critical.push_task(task); + } + (false, true) => { + // Unfair scheduling will not wake the core in order, so + // don't bother having it linked. + critical.remove_task(task); + } + _ => {} + } + } + + /// Release any remaining tokens which are associated with this particular task. + fn release_remaining( + self: Pin<&mut Self>, + critical: &mut Guard<'_>, + state: &mut State<'_>, + permits: usize, + ) { + let (critical, task) = self.get_task(critical); + + if task.is_linked() { + critical.remove_task(task); + } + + // Hand back permits which we've acquired so far. + let release = permits.saturating_sub(task.remaining); + state.add_tokens(critical, release, |_, _| ()); + } + + /// Drain the given number of tokens through the core. Returns `true` if the + /// core has been completed. + fn drain_core( + self: Pin<&mut Self>, + critical: &mut Guard<'_>, + state: &mut State<'_>, + tokens: usize, + ) -> bool { + let completed = state.add_tokens(critical, tokens, |critical, state| { + let (_, task) = self.get_task(critical); + + // If the limiter is not fair, we need to in addition to draining + // remaining tokens from linked nodes, drain it from ourselves. We + // fill the current holder of the core last (self). To ensure that + // it stays around for as long as possible. + if !state.lim.fair { + task.fill(&mut state.balance); + } + + task.is_completed() + }); + + if completed { + // Everything was drained, including the current core (if + // appropriate). So we can release it now. + critical.release(state); + } + + completed + } + + /// Assume the current core and calculate how long we must sleep for in + /// order to do it. + /// + /// # Safety + /// + /// This might link the current task into the task queue, so the caller must + /// ensure that it is pinned. + fn assume_core( + mut self: Pin<&mut Self>, + critical: &mut Guard<'_>, + state: &mut State<'_>, + now: Instant, + ) -> bool { + self.as_mut().link_core(critical, state.lim); + + let (tokens, deadline) = match state.lim.calculate_drain(critical.deadline, now) { + Some(tokens) => tokens, + None => return true, + }; + + // It is appropriate to update the deadline. + critical.deadline = deadline; + !self.drain_core(critical, state, tokens) + } +} + +pin_project! { + /// The future associated with acquiring permits from a rate limiter using + /// [`RateLimiter::acquire`]. + #[project(!Unpin)] + pub struct Acquire<'a> { + #[pin] + inner: AcquireFut>, + } +} + +impl Acquire<'_> { + /// Test if this acquire task is currently coordinating the rate limiter. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use std::future::Future; + /// use std::sync::Arc; + /// use std::task::Context; + /// + /// struct Waker; + /// # impl std::task::Wake for Waker { fn wake(self: Arc) { } } + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = RateLimiter::builder().build(); + /// + /// let waker = Arc::new(Waker).into(); + /// let mut cx = Context::from_waker(&waker); + /// + /// let a1 = limiter.acquire(1); + /// tokio::pin!(a1); + /// + /// assert!(!a1.is_core()); + /// assert!(a1.as_mut().poll(&mut cx).is_pending()); + /// assert!(a1.is_core()); + /// + /// a1.as_mut().await; + /// + /// // After completion this is no longer a core. + /// assert!(!a1.is_core()); + /// # } + /// ``` + pub fn is_core(&self) -> bool { + self.inner.is_core() + } +} + +impl Future for Acquire<'_> { + type Output = (); + + #[inline] + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.project().inner.poll(cx) + } +} + +pin_project! { + /// The future associated with acquiring permits from a rate limiter using + /// [`RateLimiter::acquire_owned`]. + #[project(!Unpin)] + pub struct AcquireOwned { + #[pin] + inner: AcquireFut>, + } +} + +impl AcquireOwned { + /// Test if this acquire task is currently coordinating the rate limiter. + /// + /// # Examples + /// + /// ``` + /// use leaky_bucket::RateLimiter; + /// use std::future::Future; + /// use std::sync::Arc; + /// use std::task::Context; + /// + /// struct Waker; + /// # impl std::task::Wake for Waker { fn wake(self: Arc) { } } + /// + /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { + /// let limiter = Arc::new(RateLimiter::builder().build()); + /// + /// let waker = Arc::new(Waker).into(); + /// let mut cx = Context::from_waker(&waker); + /// + /// let a1 = limiter.acquire_owned(1); + /// tokio::pin!(a1); + /// + /// assert!(!a1.is_core()); + /// assert!(a1.as_mut().poll(&mut cx).is_pending()); + /// assert!(a1.is_core()); + /// + /// a1.as_mut().await; + /// + /// // After completion this is no longer a core. + /// assert!(!a1.is_core()); + /// # } + /// ``` + pub fn is_core(&self) -> bool { + self.inner.is_core() + } +} + +impl Future for AcquireOwned { + type Output = (); + + #[inline] + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.project().inner.poll(cx) + } +} + +pin_project! { + #[project(!Unpin)] + #[project = AcquireFutProj] + struct AcquireFut + where + T: Deref, + { + lim: T, + permits: usize, + state: AcquireFutState, + #[pin] + sleep: Option, + #[pin] + inner: AcquireFutInner, + } + + impl PinnedDrop for AcquireFut + where + T: Deref, + { + fn drop(this: Pin<&mut Self>) { + let AcquireFutProj { lim, permits, state, inner, .. } = this.project(); + + let is_core = match *state { + AcquireFutState::Waiting => false, + AcquireFutState::Core { .. } => true, + _ => return, + }; + + let mut critical = lim.lock(); + let mut s = lim.take(); + inner.release_remaining(&mut critical, &mut s, *permits); + + if is_core { + critical.release(&mut s); + } + + *state = AcquireFutState::Complete; + } + } +} + +impl AcquireFut +where + T: Deref, +{ + #[inline] + const fn new(lim: T, permits: usize) -> Self { + Self { + lim, + permits, + state: AcquireFutState::Initial, + sleep: None, + inner: AcquireFutInner::new(), + } + } + + fn is_core(&self) -> bool { + matches!(&self.state, AcquireFutState::Core { .. }) + } +} + +// SAFETY: All the internals of acquire is thread safe and correctly +// synchronized. The embedded waiter queue doesn't have anything inherently +// unsafe in it. +unsafe impl Send for AcquireFut where T: Send + Deref {} +unsafe impl Sync for AcquireFut where T: Sync + Deref {} + +impl Future for AcquireFut +where + T: Deref, +{ + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let AcquireFutProj { + lim, + permits, + state, + mut sleep, + inner: mut internal, + .. + } = self.project(); + + // Hold onto the critical lock for core operations, but only acquire it + // when strictly necessary. + let mut critical; + + // Shared state. + // + // Once we are holding onto the critical lock, we take the entire state + // to ensure that any fast-past negotiators do not observe any available + // permits while potential core work is ongoing. + let mut s; + + // Hold onto any call to `Instant::now` which we might perform, so we + // don't have to get the current time multiple times. + let outer_now; + + match *state { + AcquireFutState::Complete => { + return Poll::Ready(()); + } + AcquireFutState::Initial => { + // If the rate limiter is not fair, try to oppurtunistically + // just acquire a permit through the known atomic state. + // + // This is known as the fast path, but requires acquire to raise + // against other tasks when storing the state back. + if lim.try_fast_path(*permits) { + *state = AcquireFutState::Complete; + return Poll::Ready(()); + } + + critical = lim.lock(); + s = lim.take(); + + let now = Instant::now(); + + // If we've hit a deadline, calculate the number of tokens + // to drain and perform it in line here. This is necessary + // because the core isn't aware of how long we sleep between + // each acquire, so we need to perform some of the drain + // work here in order to avoid acruing a debt that needs to + // be filled later in. + // + // If we didn't do this, and the process slept for a long + // time, the next time a core is acquired it would be very + // far removed from the expected deadline and has no idea + // when permits were acquired, so it would over-eagerly + // release a lot of acquires and accumulate permits. + // + // This is tested for in the `test_idle` suite of tests. + let tokens = + if let Some((tokens, deadline)) = lim.calculate_drain(critical.deadline, now) { + // We pre-emptively update the deadline of the core + // since it might bump, and we don't want other + // processes to observe that the deadline has been + // reached. + critical.deadline = deadline; + tokens + } else { + 0 + }; + + let completed = s.add_tokens(&mut critical, tokens, |critical, s| { + let (_, task) = internal.as_mut().get_task(critical); + task.remaining = *permits; + task.fill(&mut s.balance); + task.is_completed() + }); + + if completed { + *state = AcquireFutState::Complete; + return Poll::Ready(()); + } + + // Try to take over as core. If we're unsuccessful we just + // ensure that we're linked into the wait queue. + if !mem::take(&mut s.available) { + internal.as_mut().update(&mut critical, cx.waker()); + *state = AcquireFutState::Waiting; + return Poll::Pending; + } + + // SAFETY: This is done in a pinned section, so we know that + // the linked section stays alive for the duration of this + // future due to pinning guarantees. + internal.as_mut().link_core(&mut critical, lim); + Guard::bump(&mut critical); + *state = AcquireFutState::Core; + outer_now = Some(now); + } + AcquireFutState::Waiting => { + // If we are complete, then return as ready. + // + // This field is atomic, so we can safely read it under shared + // access and do not require a lock. + if internal.complete().load(Ordering::Acquire) { + *state = AcquireFutState::Complete; + return Poll::Ready(()); + } + + // Note: we need to operate under this lock to ensure that + // the core acquired here (or elsewhere) observes that the + // current task has been linked up. + critical = lim.lock(); + s = lim.take(); + + // Try to take over as core. If we're unsuccessful we + // just ensure that we're linked into the wait queue. + if !mem::take(&mut s.available) { + internal.update(&mut critical, cx.waker()); + return Poll::Pending; + } + + let now = Instant::now(); + + // This is done in a pinned section, so we know that the linked + // section stays alive for the duration of this future due to + // pinning guarantees. + if !internal.as_mut().assume_core(&mut critical, &mut s, now) { + // Marks as completed. + *state = AcquireFutState::Complete; + return Poll::Ready(()); + } + + Guard::bump(&mut critical); + *state = AcquireFutState::Core; + outer_now = Some(now); + } + AcquireFutState::Core => { + critical = lim.lock(); + s = lim.take(); + outer_now = None; + } + } + + let mut sleep = match sleep.as_mut().as_pin_mut() { + Some(mut sleep) => { + if sleep.deadline() != critical.deadline { + sleep.as_mut().reset(critical.deadline); + } + + sleep + } + None => { + sleep.set(Some(time::sleep_until(critical.deadline))); + sleep.as_mut().as_pin_mut().unwrap() + } + }; + + if sleep.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + + critical.deadline = outer_now.unwrap_or(sleep.deadline()) + lim.interval; + + if internal.drain_core(&mut critical, &mut s, lim.refill) { + *state = AcquireFutState::Complete; + return Poll::Ready(()); + } + + cx.waker().wake_by_ref(); + Poll::Pending + } +} + +pub struct Node { + /// The next node. + next: Option>>, + /// The previous node. + prev: Option>>, + /// If we are linked or not. + linked: bool, + /// The value inside of the node. + value: T, + /// Avoids noalias heuristics from kicking in on references to a `Node` + /// struct. + _pin: marker::PhantomPinned, +} + +impl Node { + /// Construct a new unlinked node. + const fn new(value: T) -> Self { + Self { + next: None, + prev: None, + linked: false, + value, + _pin: marker::PhantomPinned, + } + } + + #[inline(always)] + fn is_linked(&self) -> bool { + self.linked + } + + /// Set the next node. + #[inline(always)] + unsafe fn set_next(&mut self, node: Option>) { + ptr::addr_of_mut!(self.next).write(node); + } + + /// Take the next node. + #[inline(always)] + unsafe fn take_next(&mut self) -> Option> { + ptr::addr_of_mut!(self.next).replace(None) + } + + /// Set the previous node. + #[inline(always)] + unsafe fn set_prev(&mut self, node: Option>) { + ptr::addr_of_mut!(self.prev).write(node); + } + + /// Take the previous node. + #[inline(always)] + unsafe fn take_prev(&mut self) -> Option> { + ptr::addr_of_mut!(self.prev).replace(None) + } +} + +impl ops::Deref for Node { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.value + } +} + +impl ops::DerefMut for Node +where + T: Unpin, +{ + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.value + } +} + +pub struct LinkedList { + head: Option>>, + tail: Option>>, +} + +impl LinkedList { + /// Construct a new empty list. + const fn new() -> Self { + Self { + head: None, + tail: None, + } + } + + unsafe fn push_front(&mut self, mut node: ptr::NonNull>) { + debug_assert!(node.as_ref().next.is_none()); + debug_assert!(node.as_ref().prev.is_none()); + debug_assert!(!node.as_ref().linked); + + if let Some(mut head) = self.head.take() { + node.as_mut().set_next(Some(head)); + head.as_mut().set_prev(Some(node)); + self.head = Some(node); + } else { + self.head = Some(node); + self.tail = Some(node); + } + + node.as_mut().linked = true; + } + + unsafe fn push_back(&mut self, mut node: ptr::NonNull>) { + + debug_assert!(node.as_ref().next.is_none()); + debug_assert!(node.as_ref().prev.is_none()); + debug_assert!(!node.as_ref().linked); + + if let Some(mut tail) = self.tail.take() { + node.as_mut().set_prev(Some(tail)); + tail.as_mut().set_next(Some(node)); + self.tail = Some(node); + } else { + self.head = Some(node); + self.tail = Some(node); + } + + node.as_mut().linked = true; + } + + #[cfg(test)] + unsafe fn pop_front(&mut self) -> Option>> { + let mut head = self.head?; + debug_assert!(head.as_ref().linked); + + if let Some(mut next) = head.as_mut().take_next() { + next.as_mut().set_prev(None); + self.head = Some(next); + } else { + debug_assert_eq!(self.tail, Some(head)); + self.head = None; + self.tail = None; + } + + debug_assert!(head.as_ref().prev.is_none()); + debug_assert!(head.as_ref().next.is_none()); + head.as_mut().linked = false; + Some(head) + } + + /// Pop the back element from the list. + unsafe fn pop_back(&mut self) -> Option>> { + let mut tail = self.tail?; + debug_assert!(tail.as_ref().linked); + + if let Some(mut prev) = tail.as_mut().take_prev() { + prev.as_mut().set_next(None); + self.tail = Some(prev); + } else { + debug_assert_eq!(self.head, Some(tail)); + self.head = None; + self.tail = None; + } + + debug_assert!(tail.as_ref().prev.is_none()); + debug_assert!(tail.as_ref().next.is_none()); + tail.as_mut().linked = false; + Some(tail) + } + + unsafe fn remove(&mut self, mut node: ptr::NonNull>) { + debug_assert!(node.as_ref().linked); + + let next = node.as_mut().take_next(); + let prev = node.as_mut().take_prev(); + + if let Some(mut next) = next { + next.as_mut().set_prev(prev); + } else { + debug_assert_eq!(self.tail, Some(node)); + self.tail = prev; + } + + if let Some(mut prev) = prev { + prev.as_mut().set_next(next); + } else { + debug_assert_eq!(self.head, Some(node)); + self.head = next; + } + + node.as_mut().linked = false; + } + + unsafe fn front(&mut self) -> Option>> { + self.head + } +} + +impl fmt::Debug for LinkedList { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LinkedList") + .field("head", &self.head) + .field("tail", &self.tail) + .finish() + } +} diff --git a/src/shared_context.rs b/src/shared_context.rs index 95bebc1..bcae4c7 100644 --- a/src/shared_context.rs +++ b/src/shared_context.rs @@ -1,9 +1,10 @@ use crate::async_flag::AsyncFlag; use crate::histogram::Histogram; +use crate::BenchmarkResult; use std::cmp::min; use std::option::Option; -use std::sync::atomic::AtomicU64; -use std::sync::{Arc, RwLock}; +use std::sync::atomic::{AtomicU64, AtomicU8}; +use std::sync::{Arc, Mutex, RwLock}; use std::time::Instant; #[derive(Clone)] @@ -22,6 +23,11 @@ pub struct SharedContext { // histogram pub histogram: Arc, + + // latest histogram + pub latest_histogram: Arc<[Histogram; 2]>, + pub latest_histogramactive_index: Arc, + pub latest_result: Arc> } impl SharedContext { @@ -35,6 +41,10 @@ impl SharedContext { stop_flag: AsyncFlag::new(), histogram: Arc::new(Histogram::new()), + + latest_histogram: Arc::new([Histogram::new(), Histogram::new()]), + latest_histogramactive_index: Arc::new(AtomicU8::new(0)), + latest_result: Arc::new(Mutex::new(BenchmarkResult::default())), } } @@ -67,6 +77,17 @@ impl SharedContext { return 0; } } + + if self.stop_flag.flag() { + return 0; + } + return result; } + + pub fn record(&self, latency_us: u64) { + self.histogram.record(latency_us); + let current_index = self.latest_histogramactive_index.load(std::sync::atomic::Ordering::Relaxed) as usize % 2; + self.latest_histogram[current_index].record(latency_us); + } } From bdf6b6abc04042aad398decdfa2b99ec09f37e24 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Fri, 4 Jul 2025 16:25:42 +0800 Subject: [PATCH 02/10] upgrade version --- Cargo.lock | 27 +++++++++++++-------------- Cargo.toml | 8 ++++---- pyproject.toml | 2 +- python/resp_benchmark/__init__.py | 2 ++ 4 files changed, 20 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 887ca75..53f6af0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -533,11 +533,10 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.22.2" +version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "831e8e819a138c36e212f3af3fd9eeffed6bf1510a805af35b0edee5ffa59433" +checksum = "8970a78afe0628a3e3430376fc5fd76b6b45c4d43360ffd6cdd40bdde72b682a" dependencies = [ - "cfg-if", "indoc", "libc", "memoffset", @@ -551,9 +550,9 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.22.2" +version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e8730e591b14492a8945cdff32f089250b05f5accecf74aeddf9e8272ce1fa8" +checksum = "458eb0c55e7ece017adeba38f2248ff3ac615e53660d7c71a238d7d2a01c7598" dependencies = [ "once_cell", "target-lexicon", @@ -561,9 +560,9 @@ dependencies = [ [[package]] name = "pyo3-ffi" -version = "0.22.2" +version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e97e919d2df92eb88ca80a037969f44e5e70356559654962cbb3316d00300c6" +checksum = "7114fe5457c61b276ab77c5055f206295b812608083644a5c5b2640c3102565c" dependencies = [ "libc", "pyo3-build-config", @@ -571,9 +570,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.22.2" +version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb57983022ad41f9e683a599f2fd13c3664d7063a3ac5714cae4b7bee7d3f206" +checksum = "a8725c0a622b374d6cb051d11a0983786448f7785336139c3c94f5aa6bef7e50" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -583,9 +582,9 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.22.2" +version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec480c0c51ddec81019531705acac51bcdbeae563557c982aa8263bb96880372" +checksum = "4109984c22491085343c05b0dbc54ddc405c3cf7b4374fc533f5c3313a572ccc" dependencies = [ "heck", "proc-macro2", @@ -671,7 +670,7 @@ dependencies = [ [[package]] name = "resp-benchmark" -version = "0.1.7" +version = "0.2.0" dependencies = [ "awaitgroup", "colored", @@ -771,9 +770,9 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.12.16" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a" [[package]] name = "tinyvec" diff --git a/Cargo.toml b/Cargo.toml index 7a85f79..daee569 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "resp-benchmark" -version = "0.1.7" -edition = "2021" +version = "0.2.0" +edition = "2024" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [lib] @@ -9,9 +9,9 @@ name = "_resp_benchmark_rust_lib" crate-type = ["cdylib"] [dependencies] -pyo3 = { version = "0.22.2", features = ["extension-module"] } +pyo3 = { version = "0.25.1", features = ["extension-module"] } tokio = { version = "1", features = ["full"] } -redis = { version = "0.26.1", features = [ +redis = { version = "0", features = [ "tokio-comp", "cluster", "cluster-async", diff --git a/pyproject.toml b/pyproject.toml index bc33115..addc3a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "maturin" [project] name = "resp-benchmark" description = "resp-benchmark is a benchmark tool for testing databases that support the RESP protocol, such as Redis, Valkey, and Tair." -requires-python = ">=3.8" +requires-python = ">=3.9" classifiers = [ "Programming Language :: Rust", "Programming Language :: Python :: Implementation :: CPython", diff --git a/python/resp_benchmark/__init__.py b/python/resp_benchmark/__init__.py index bd9ec83..8a54c01 100644 --- a/python/resp_benchmark/__init__.py +++ b/python/resp_benchmark/__init__.py @@ -1 +1,3 @@ from .wrapper import Benchmark, Result, BenchmarkWorkerPool + +__all__ = ["Benchmark", "Result", "BenchmarkWorkerPool"] From 852e5cfc88a4a415e896f2b8c057bb6d4555fdb2 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Fri, 4 Jul 2025 16:48:22 +0800 Subject: [PATCH 03/10] fix for 2024 --- Cargo.lock | 809 ++++++++++++++++++++++-------------- Cargo.toml | 2 +- src/command/distribution.rs | 4 +- src/command/mod.rs | 4 +- src/command/placeholder.rs | 33 +- src/qps_limiter.rs | 199 ++++----- 6 files changed, 613 insertions(+), 438 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 53f6af0..a27647c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,41 +4,24 @@ version = 4 [[package]] name = "addr2line" -version = "0.22.0" +version = "0.24.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e4503c46a5c0c7844e948c9a4d6acd9f50cccb4de1c48eb9e291ea17470c678" +checksum = "dfbe277e56a376000877090da837660b4427aad530e3028d44e0bffe4f89a1c1" dependencies = [ "gimli", ] [[package]] -name = "adler" -version = "1.0.2" +name = "adler2" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" - -[[package]] -name = "arc-swap" -version = "1.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457" - -[[package]] -name = "async-trait" -version = "0.1.81" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e0c28dcc82d7c8ead5cb13beb15405b57b8546e93215673ff8ca0349a028107" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.72", -] +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "autocfg" -version = "1.3.0" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "awaitgroup" @@ -48,63 +31,51 @@ checksum = "a872ceb3db05a391fbe7cf8eba07a1239b2d946eee66f9e942be9bff06206302" [[package]] name = "backtrace" -version = "0.3.73" +version = "0.3.75" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cc23269a4f8976d0a4d2e7109211a419fe30e8d88d677cd60b6bc79c5732e0a" +checksum = "6806a6321ec58106fea15becdad98371e28d92ccbc7c8f1b3b6dd724fe8f1002" dependencies = [ "addr2line", - "cc", "cfg-if", "libc", "miniz_oxide", "object", "rustc-demangle", + "windows-targets", ] [[package]] name = "bitflags" -version = "2.6.0" +version = "2.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" - -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" +checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" [[package]] name = "bytes" -version = "1.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8318a53db07bb3f8dca91a600466bdb3f2eaadeedfdbcf02e1accbad9271ba50" - -[[package]] -name = "cc" -version = "1.1.8" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "504bdec147f2cc13c8b57ed9401fd8a147cc66b67ad5cb241394244f2c947549" +checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268" [[package]] name = "cfg_aliases" -version = "0.1.1" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "colored" -version = "2.1.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbf2150cce219b664a8a70df7a1f933836724b503f8a413af9365b4dcc4d90b8" +checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -123,9 +94,9 @@ dependencies = [ [[package]] name = "core_affinity" -version = "0.8.1" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "622892f5635ce1fc38c8f16dfc938553ed64af482edb5e150bf4caedbfcb2304" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" dependencies = [ "libc", "num_cpus", @@ -140,12 +111,23 @@ checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff" [[package]] name = "ctrlc" -version = "3.4.4" +version = "3.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "672465ae37dc1bc6380a6547a8883d5dd397b0f1faaad4f265726cc7042a5345" +checksum = "46f93780a459b7d656ef7f071fe699c4d3d2cb201c4b24d085b6ddc505276e73" dependencies = [ "nix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", +] + +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", ] [[package]] @@ -168,7 +150,7 @@ checksum = "2e1f6c3800b304a6be0012039e2a45a322a093539c45ab818d9e6895a39c90fe" dependencies = [ "proc-macro2", "quote", - "rand", + "rand 0.8.5", "syn 1.0.109", ] @@ -182,144 +164,207 @@ dependencies = [ ] [[package]] -name = "futures" -version = "0.3.30" +name = "futures-core" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" -dependencies = [ - "futures-channel", - "futures-core", - "futures-executor", - "futures-io", - "futures-sink", - "futures-task", - "futures-util", -] +checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" [[package]] -name = "futures-channel" -version = "0.3.30" +name = "futures-sink" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" -dependencies = [ - "futures-core", - "futures-sink", -] +checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" [[package]] -name = "futures-core" -version = "0.3.30" +name = "futures-task" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" [[package]] -name = "futures-executor" -version = "0.3.30" +name = "futures-util" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ "futures-core", + "futures-sink", "futures-task", - "futures-util", + "pin-project-lite", + "pin-utils", + "slab", ] [[package]] -name = "futures-io" -version = "0.3.30" +name = "getrandom" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" +checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.11.1+wasi-snapshot-preview1", +] [[package]] -name = "futures-macro" -version = "0.3.30" +name = "getrandom" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.72", + "cfg-if", + "libc", + "r-efi", + "wasi 0.14.2+wasi-0.2.4", ] [[package]] -name = "futures-sink" -version = "0.3.30" +name = "gimli" +version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" +checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" [[package]] -name = "futures-task" -version = "0.3.30" +name = "heck" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] -name = "futures-util" -version = "0.3.30" +name = "hermit-abi" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "icu_collections" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "200072f5d0e3614556f94a9930d5dc3e0662a652823904c3a75dc3b0af7fee47" dependencies = [ - "futures-channel", - "futures-core", - "futures-io", - "futures-macro", - "futures-sink", - "futures-task", - "memchr", - "pin-project-lite", - "pin-utils", - "slab", + "displaydoc", + "potential_utf", + "yoke", + "zerofrom", + "zerovec", ] [[package]] -name = "getrandom" -version = "0.2.15" +name = "icu_locale_core" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "0cde2700ccaed3872079a65fb1a78f6c0a36c91570f28755dda67bc8f7d9f00a" dependencies = [ - "cfg-if", - "libc", - "wasi", + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", ] [[package]] -name = "gimli" -version = "0.29.0" +name = "icu_normalizer" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd" +checksum = "436880e8e18df4d7bbc06d58432329d6458cc84531f7ac5f024e93deadb37979" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] [[package]] -name = "heck" -version = "0.5.0" +name = "icu_normalizer_data" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +checksum = "00210d6893afc98edb752b664b8890f0ef174c8adbb8d0be9710fa66fbbf72d3" [[package]] -name = "hermit-abi" -version = "0.3.9" +name = "icu_properties" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +checksum = "016c619c1eeb94efb86809b015c58f479963de65bdb6253345c1a1276f22e32b" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "potential_utf", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "298459143998310acd25ffe6810ed544932242d3f07083eee1084d83a71bd632" + +[[package]] +name = "icu_provider" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03c80da27b5f4187909049ee2d72f276f0d9f99a42c306bd0131ecfe04d8e5af" +dependencies = [ + "displaydoc", + "icu_locale_core", + "stable_deref_trait", + "tinystr", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] [[package]] name = "idna" -version = "0.5.0" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "634d9b1461af396cad843f47fdba5597a4f9e6ddd4bfb6ff5d85028c25cb12f6" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" dependencies = [ - "unicode-bidi", - "unicode-normalization", + "icu_normalizer", + "icu_properties", ] [[package]] name = "indoc" -version = "2.0.5" +version = "2.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" +checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd" + +[[package]] +name = "io-uring" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b86e202f00093dcba4275d4636b93ef9dd75d025ae560d2521b45ea28ab49013" +dependencies = [ + "bitflags", + "cfg-if", + "libc", +] [[package]] name = "itoa" -version = "1.0.11" +version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "lazy_static" @@ -329,15 +374,27 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" [[package]] name = "libc" -version = "0.2.155" +version = "0.2.174" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776" + +[[package]] +name = "libm" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c" +checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" + +[[package]] +name = "litemap" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" [[package]] name = "lock_api" -version = "0.4.12" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +checksum = "96936507f153605bddfcda068dd804796c84324ed2510809e5b2a624c81da765" dependencies = [ "autocfg", "scopeguard", @@ -345,15 +402,15 @@ dependencies = [ [[package]] name = "log" -version = "0.4.22" +version = "0.4.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" [[package]] name = "memchr" -version = "2.7.4" +version = "2.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" [[package]] name = "memoffset" @@ -372,30 +429,29 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "miniz_oxide" -version = "0.7.4" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8a240ddb74feaf34a79a7add65a741f3167852fba007066dcac1ca548d89c08" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" dependencies = [ - "adler", + "adler2", ] [[package]] name = "mio" -version = "1.0.1" +version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4569e456d394deccd22ce1c1913e6ea0e54519f577285001215d33557431afe4" +checksum = "78bed444cc8a2160f01cbcf811ef18cac863ad68ae8ca62092e8db51d51c761c" dependencies = [ - "hermit-abi", "libc", - "wasi", - "windows-sys 0.52.0", + "wasi 0.11.1+wasi-snapshot-preview1", + "windows-sys 0.59.0", ] [[package]] name = "nix" -version = "0.28.0" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab2156c4fce2f8df6c499cc1c763e4394b7482525bf2a9701c9d79d215f519e4" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" dependencies = [ "bitflags", "cfg-if", @@ -439,13 +495,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] name = "num_cpus" -version = "1.16.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" dependencies = [ "hermit-abi", "libc", @@ -453,24 +510,24 @@ dependencies = [ [[package]] name = "object" -version = "0.36.3" +version = "0.36.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27b64972346851a39438c60b341ebc01bba47464ae329e55cf343eb93964efd9" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" dependencies = [ "memchr", ] [[package]] name = "once_cell" -version = "1.19.0" +version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "parking_lot" -version = "0.12.3" +version = "0.12.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +checksum = "70d58bf43669b5795d1576d0641cfb6fbb2057bf629506267a92807158584a13" dependencies = [ "lock_api", "parking_lot_core", @@ -478,15 +535,15 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.10" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +checksum = "bc838d2a56b5b1a6c25f55575dfc605fabb63bb2365f6c2353ef9159aa69e4a5" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -497,9 +554,9 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "pin-project-lite" -version = "0.2.14" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" [[package]] name = "pin-utils" @@ -509,24 +566,33 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "portable-atomic" -version = "1.7.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da544ee218f0d287a911e9c99a39a8c9bc8fcad3cb8db5959940044ecfc67265" +checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483" + +[[package]] +name = "potential_utf" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7c30837279ca13e7c867e9e40053bc68740f988cb07f7ca6df43cc734b585" +dependencies = [ + "zerovec", +] [[package]] name = "ppv-lite86" -version = "0.2.20" +version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" dependencies = [ "zerocopy", ] [[package]] name = "proc-macro2" -version = "1.0.86" +version = "1.0.95" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778" dependencies = [ "unicode-ident", ] @@ -577,7 +643,7 @@ dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 2.0.72", + "syn 2.0.104", ] [[package]] @@ -590,18 +656,24 @@ dependencies = [ "proc-macro2", "pyo3-build-config", "quote", - "syn 2.0.72", + "syn 2.0.104", ] [[package]] name = "quote" -version = "1.0.36" +version = "1.0.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fa76aaf39101c457836aec0ce2316dbdc3ab723cdda1c6bd4e6ad4208acaca7" +checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "rand" version = "0.8.5" @@ -609,8 +681,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.3", ] [[package]] @@ -620,7 +702,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.3", ] [[package]] @@ -629,28 +721,46 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom", + "getrandom 0.2.16", +] + +[[package]] +name = "rand_core" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +dependencies = [ + "getrandom 0.3.3", +] + +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand 0.8.5", ] [[package]] name = "redis" -version = "0.26.1" +version = "0.32.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e902a69d09078829137b4a5d9d082e0490393537badd7c91a3d69d14639e115f" +checksum = "7f0f6a8c53351d89a3869a703459995a0bcadcfa846002707fbc7e5cca235c4a" dependencies = [ - "arc-swap", - "async-trait", "bytes", + "cfg-if", "combine", "crc16", - "futures", + "futures-sink", "futures-util", "itoa", "log", "num-bigint", "percent-encoding", "pin-project-lite", - "rand", + "rand 0.9.1", "ryu", "sha1_smol", "socket2", @@ -661,9 +771,9 @@ dependencies = [ [[package]] name = "redox_syscall" -version = "0.5.3" +version = "0.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a908a6e00f1fdd0dfd9c0eb08ce85126f6d8bbda50017e74bc4a4b7d4a926a4" +checksum = "0d04b7d0ee6b4a0207a0a7adb104d23ecb0b47d6beae7152d0fa34b692b29fd6" dependencies = [ "bitflags", ] @@ -681,24 +791,24 @@ dependencies = [ "parking_lot", "pin-project-lite", "pyo3", - "rand", + "rand 0.8.5", + "rand_distr", "redis", "tokio", "urlencoding", - "zipf", ] [[package]] name = "rustc-demangle" -version = "0.1.24" +version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" +checksum = "989e6739f80c4ad5b13e0fd7fe89531180375b18520cc8c82080e4dc4035b84f" [[package]] name = "ryu" -version = "1.0.18" +version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" [[package]] name = "scopeguard" @@ -706,6 +816,26 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "serde" +version = "1.0.219" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.219" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "sha1_smol" version = "1.0.1" @@ -714,38 +844,41 @@ checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" [[package]] name = "signal-hook-registry" -version = "1.4.2" +version = "1.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" +checksum = "9203b8055f63a2a00e2f593bb0510367fe707d7ff1e5c872de2f537b339e5410" dependencies = [ "libc", ] [[package]] name = "slab" -version = "0.4.9" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" -dependencies = [ - "autocfg", -] +checksum = "04dc19736151f35336d325007ac991178d504a119863a2fcb3758cdb5e52c50d" [[package]] name = "smallvec" -version = "1.13.2" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "socket2" -version = "0.5.7" +version = "0.5.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce305eb0b4296696835b71df73eb912e0f1ffd2556a501fcede6e0c50349191c" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" dependencies = [ "libc", "windows-sys 0.52.0", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + [[package]] name = "syn" version = "1.0.109" @@ -759,15 +892,26 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.72" +version = "2.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc4b9b9bf2add8093d3f2c0204471e951b2285580335de42f9d2534f3ae7a8af" +checksum = "17b6f705963418cdb9927482fa304bc562ece2fdd4f616084c50b7023b435a40" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "target-lexicon" version = "0.13.2" @@ -775,33 +919,30 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a" [[package]] -name = "tinyvec" -version = "1.8.0" +name = "tinystr" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "445e881f4f6d382d5f27c034e25eb92edd7c784ceab92a0937db7f2e9471b938" +checksum = "5d4f6d1145dcb577acf783d4e601bc1d76a13337bb54e6233add580b07344c8b" dependencies = [ - "tinyvec_macros", + "displaydoc", + "zerovec", ] -[[package]] -name = "tinyvec_macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" - [[package]] name = "tokio" -version = "1.39.2" +version = "1.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "daa4fb1bc778bd6f04cbfc4bb2d06a7396a8f299dc33ea1900cedaa316f467b1" +checksum = "1140bb80481756a8cbe10541f37433b459c5aa1e727b4c020fbfebdc25bf3ec4" dependencies = [ "backtrace", "bytes", + "io-uring", "libc", "mio", "parking_lot", "pin-project-lite", "signal-hook-registry", + "slab", "socket2", "tokio-macros", "windows-sys 0.52.0", @@ -809,20 +950,20 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.72", + "syn 2.0.104", ] [[package]] name = "tokio-util" -version = "0.7.11" +version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9cf6b47b3771c49ac75ad09a6162f53ad4b8088b76ac60e8ec1455b31a189fe1" +checksum = "66a539a9ad6d5d281510d5bd368c973d636c02dbf8a67300bfb6b950696ad7df" dependencies = [ "bytes", "futures-core", @@ -831,38 +972,23 @@ dependencies = [ "tokio", ] -[[package]] -name = "unicode-bidi" -version = "0.3.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08f95100a766bf4f8f28f90d77e0a5461bbdb219042e7679bebe79004fed8d75" - [[package]] name = "unicode-ident" -version = "1.0.12" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" - -[[package]] -name = "unicode-normalization" -version = "0.1.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a56d1686db2308d901306f92a263857ef59ea39678a5458e7cb17f01415101f5" -dependencies = [ - "tinyvec", -] +checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" [[package]] name = "unindent" -version = "0.2.3" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" [[package]] name = "url" -version = "2.5.2" +version = "2.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" +checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" dependencies = [ "form_urlencoded", "idna", @@ -875,11 +1001,26 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" +version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasi" +version = "0.14.2+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" +dependencies = [ + "wit-bindgen-rt", +] [[package]] name = "winapi" @@ -903,37 +1044,22 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets 0.48.5", -] - [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] -name = "windows-targets" -version = "0.48.5" +name = "windows-sys" +version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows_aarch64_gnullvm 0.48.5", - "windows_aarch64_msvc 0.48.5", - "windows_i686_gnu 0.48.5", - "windows_i686_msvc 0.48.5", - "windows_x86_64_gnu 0.48.5", - "windows_x86_64_gnullvm 0.48.5", - "windows_x86_64_msvc 0.48.5", + "windows-targets", ] [[package]] @@ -942,46 +1068,28 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", "windows_i686_gnullvm", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -996,78 +1104,137 @@ checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" [[package]] name = "windows_i686_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" [[package]] -name = "windows_i686_msvc" +name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" [[package]] -name = "windows_x86_64_gnu" -version = "0.48.5" +name = "windows_x86_64_gnullvm" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" [[package]] -name = "windows_x86_64_gnu" +name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.5" +name = "wit-bindgen-rt" +version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" +checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" +dependencies = [ + "bitflags", +] [[package]] -name = "windows_x86_64_gnullvm" -version = "0.52.6" +name = "writeable" +version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb" [[package]] -name = "windows_x86_64_msvc" -version = "0.48.5" +name = "yoke" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" +checksum = "5f41bb01b8226ef4bfd589436a297c53d118f65921786300e427be8d487695cc" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] [[package]] -name = "windows_x86_64_msvc" -version = "0.52.6" +name = "yoke-derive" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +checksum = "38da3c9736e16c5d3c8c597a9aaa5d1fa565d0532ae05e27c24aa62fb32c0ab6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", + "synstructure", +] [[package]] name = "zerocopy" -version = "0.7.35" +version = "0.8.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f" dependencies = [ - "byteorder", "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.7.35" +version = "0.8.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + +[[package]] +name = "zerofrom" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn 2.0.72", + "syn 2.0.104", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36f0bbd478583f79edad978b407914f61b2972f5af6fa089686016be8f9af595" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", ] [[package]] -name = "zipf" -version = "7.0.1" +name = "zerovec" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "390e51da0ed8cc3ade001d15fa5ba6f966b99c858fb466ec6b06d1682f1f94dd" +checksum = "4a05eb080e015ba39cc9e23bbe5e7fb04d5fb040350f99f34e338d5fdd294428" dependencies = [ - "rand", + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b96237efa0c878c64bd89c436f661be4e46b2f3eff1ebb976f7ef2321d2f58f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", ] diff --git a/Cargo.toml b/Cargo.toml index daee569..14cf73e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ redis = { version = "0", features = [ # "tokio-native-tls-comp", ] } rand = { version = "0.8.5", features = [] } -zipf = "7.0.1" +rand_distr = "0.4.3" nom = "7.1.3" core_affinity = "0.8.1" awaitgroup = "0.7.0" diff --git a/src/command/distribution.rs b/src/command/distribution.rs index dae8047..1bf2437 100644 --- a/src/command/distribution.rs +++ b/src/command/distribution.rs @@ -5,7 +5,7 @@ use rand::distributions::Distribution; #[derive(Clone, Debug)] pub enum DistributionEnum { Uniform(rand::distributions::Uniform), - Zipfian(zipf::ZipfDistribution), + Zipfian(rand_distr::Zipf), Sequence(SequenceDistribution), } @@ -13,7 +13,7 @@ impl DistributionEnum { pub fn new(s: &str, range: u64) -> Self { match s { "uniform" => Self::Uniform(rand::distributions::Uniform::new(0, range)), - "zipfian" => Self::Zipfian(zipf::ZipfDistribution::new(range as usize, 1.03).unwrap()), + "zipfian" => Self::Zipfian(rand_distr::Zipf::new(range, 1.03).unwrap()), "sequence" => Self::Sequence(SequenceDistribution::new(range)), _ => panic!("Unknown distribution"), } diff --git a/src/command/mod.rs b/src/command/mod.rs index 2e5b17c..0f121a8 100644 --- a/src/command/mod.rs +++ b/src/command/mod.rs @@ -35,7 +35,7 @@ impl Command { let mut cmd = redis::Cmd::new(); let mut cmd_str = String::new(); for ph in self.argv.iter_mut() { - for arg in ph.gen() { + for arg in ph.generate() { cmd_str.push_str(&arg); } } @@ -50,7 +50,7 @@ impl Command { let mut cmd = redis::Cmd::new(); let mut cmd_str = String::new(); for ph in self.argv.iter_mut() { - for arg in ph.gen() { + for arg in ph.generate() { cmd_str.push_str(&arg); } } diff --git a/src/command/placeholder.rs b/src/command/placeholder.rs index 2c56dad..f0086e8 100644 --- a/src/command/placeholder.rs +++ b/src/command/placeholder.rs @@ -2,7 +2,8 @@ use std::cmp::min; use std::process::exit; use crate::command::distribution::DistributionEnum; use std::str::FromStr; -use rand::{distributions::Alphanumeric, thread_rng, Rng}; +use rand::prelude::*; +use rand::distributions::Alphanumeric; #[derive(Debug, Clone)] pub enum PlaceholderEnum { @@ -65,13 +66,13 @@ impl PlaceholderEnum { }; ph } - pub fn gen(&mut self) -> Vec { + pub fn generate(&mut self) -> Vec { match self { - Self::String(p) => vec![p.gen()], - Self::Key(p) => vec![p.gen()], - Self::Value(p) => vec![p.gen()], - Self::Rand(p) => vec![p.gen()], - Self::Range(p) => p.gen(), + Self::String(p) => vec![p.generate()], + Self::Key(p) => vec![p.generate()], + Self::Value(p) => vec![p.generate()], + Self::Rand(p) => vec![p.generate()], + Self::Range(p) => p.generate(), } } } @@ -85,7 +86,7 @@ impl PlaceholderString { pub fn new(value: String) -> Self { Self { value } } - fn gen(&mut self) -> String { + fn generate(&mut self) -> String { self.value.clone() } } @@ -99,7 +100,7 @@ impl PlaceholderKey { fn new(distribution: DistributionEnum) -> Self { Self { distribution } } - fn gen(&mut self) -> String { + fn generate(&mut self) -> String { format!("key_{:010}", self.distribution.sample(&mut rand::thread_rng())) } } @@ -113,9 +114,9 @@ impl PlaceholderValue { pub fn new(size: u64) -> Self { Self { size: size as usize } } - pub fn gen(&self) -> String { - let rng = thread_rng(); - let chars: String = rng.sample_iter(&Alphanumeric).take(self.size).map(char::from).collect(); + pub fn generate(&self) -> String { + let rng = rand::thread_rng(); + let chars: String = rng.sample_iter(Alphanumeric).take(self.size).map(char::from).collect(); chars } } @@ -129,8 +130,8 @@ impl PlaceholderRand { pub fn new(range: u64) -> Self { Self { distribution: DistributionEnum::new("uniform", range) } } - fn gen(&mut self) -> String { - format!("{}", self.distribution.sample(&mut thread_rng())) + fn generate(&mut self) -> String { + format!("{}", self.distribution.sample(&mut rand::thread_rng())) } } @@ -149,8 +150,8 @@ impl PlaceholderRange { width, } } - fn gen(&mut self) -> Vec { - let left = self.distribution.sample(&mut thread_rng()); + fn generate(&mut self) -> Vec { + let left = self.distribution.sample(&mut rand::thread_rng()); let right = min(left + self.width, self.range - 1); vec![left.to_string(), right.to_string()] } diff --git a/src/qps_limiter.rs b/src/qps_limiter.rs index 5d40689..11c9257 100644 --- a/src/qps_limiter.rs +++ b/src/qps_limiter.rs @@ -146,7 +146,7 @@ impl DerefMut for Guard<'_> { impl Critical { #[inline] fn push_task_front(&mut self, task: &mut Node) { - // SAFETY: We both have mutable access to the node being pushed, and + // SAFETY: We both have a mutable access to the node being pushed, and // mutable access to the critical section through `self`. So we know we // have exclusive tampering rights to the waiter queue. unsafe { @@ -156,7 +156,7 @@ impl Critical { #[inline] fn push_task(&mut self, task: &mut Node) { - // SAFETY: We both have mutable access to the node being pushed, and + // SAFETY: We both have a mutable access to the node being pushed, and // mutable access to the critical section through `self`. So we know we // have exclusive tampering rights to the waiter queue. unsafe { @@ -166,7 +166,7 @@ impl Critical { #[inline] fn remove_task(&mut self, task: &mut Node) { - // SAFETY: We both have mutable access to the node being pushed, and + // SAFETY: We both have a mutable access to the node being pushed, and // mutable access to the critical section through `self`. So we know we // have exclusive tampering rights to the waiter queue. unsafe { @@ -696,7 +696,7 @@ fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) { let n = node.as_mut(); n.fill(&mut state.balance); - if !n.is_completed() { + if !n.is_linked() { critical.waiters.push_back(node); break; } @@ -977,14 +977,16 @@ impl AcquireFutInner { /// Access the completion flag. pub fn complete(&self) -> &AtomicBool { - // SAFETY: This is always safe to access since it's atomic. - unsafe { &*ptr::addr_of!((*self.node.get()).complete) } + // SAFETY: The pointer is created in `AcquireFut`, which is guaranteed + // to be alive as long as `self` is. + unsafe { &*ptr::addr_of!((&*self.node.get()).complete) } } - /// Get the underlying task mutably. + /// Get a mutable reference to the underlying task node. /// - /// We prove that the caller does indeed have mutable access to the node by - /// passing in a mutable reference to the critical section. + /// # Safety + /// + /// This is safe to call as long as `self` is alive. #[inline] pub fn get_task<'crit, C>( self: Pin<&'crit mut Self>, @@ -1490,7 +1492,7 @@ pub struct Node { impl Node { /// Construct a new unlinked node. - const fn new(value: T) -> Self { + pub(crate) const fn new(value: T) -> Self { Self { next: None, prev: None, @@ -1500,33 +1502,29 @@ impl Node { } } - #[inline(always)] - fn is_linked(&self) -> bool { + /// Test if the current node is linked. + pub(crate) fn is_linked(&self) -> bool { self.linked } /// Set the next node. - #[inline(always)] - unsafe fn set_next(&mut self, node: Option>) { - ptr::addr_of_mut!(self.next).write(node); + pub(crate) unsafe fn set_next(&mut self, node: Option>) { + self.next = node; } /// Take the next node. - #[inline(always)] - unsafe fn take_next(&mut self) -> Option> { - ptr::addr_of_mut!(self.next).replace(None) + pub(crate) unsafe fn take_next(&mut self) -> Option> { + self.next.take() } /// Set the previous node. - #[inline(always)] - unsafe fn set_prev(&mut self, node: Option>) { - ptr::addr_of_mut!(self.prev).write(node); + pub(crate) unsafe fn set_prev(&mut self, node: Option>) { + self.prev = node; } /// Take the previous node. - #[inline(always)] - unsafe fn take_prev(&mut self) -> Option> { - ptr::addr_of_mut!(self.prev).replace(None) + pub(crate) unsafe fn take_prev(&mut self) -> Option> { + self.prev.take() } } @@ -1562,101 +1560,110 @@ impl LinkedList { } unsafe fn push_front(&mut self, mut node: ptr::NonNull>) { - debug_assert!(node.as_ref().next.is_none()); - debug_assert!(node.as_ref().prev.is_none()); - debug_assert!(!node.as_ref().linked); - - if let Some(mut head) = self.head.take() { - node.as_mut().set_next(Some(head)); - head.as_mut().set_prev(Some(node)); - self.head = Some(node); - } else { - self.head = Some(node); - self.tail = Some(node); - } + unsafe { + debug_assert!(node.as_ref().next.is_none()); + debug_assert!(node.as_ref().prev.is_none()); + debug_assert!(!node.as_ref().linked); + + if let Some(mut head) = self.head.take() { + node.as_mut().set_next(Some(head)); + head.as_mut().set_prev(Some(node)); + self.head = Some(node); + } else { + self.head = Some(node); + self.tail = Some(node); + } - node.as_mut().linked = true; + node.as_mut().linked = true; + } } unsafe fn push_back(&mut self, mut node: ptr::NonNull>) { + unsafe { + debug_assert!(node.as_ref().next.is_none()); + debug_assert!(node.as_ref().prev.is_none()); + debug_assert!(!node.as_ref().linked); + + if let Some(mut tail) = self.tail.take() { + node.as_mut().set_prev(Some(tail)); + tail.as_mut().set_next(Some(node)); + self.tail = Some(node); + } else { + self.head = Some(node); + self.tail = Some(node); + } - debug_assert!(node.as_ref().next.is_none()); - debug_assert!(node.as_ref().prev.is_none()); - debug_assert!(!node.as_ref().linked); - - if let Some(mut tail) = self.tail.take() { - node.as_mut().set_prev(Some(tail)); - tail.as_mut().set_next(Some(node)); - self.tail = Some(node); - } else { - self.head = Some(node); - self.tail = Some(node); + node.as_mut().linked = true; } - - node.as_mut().linked = true; } #[cfg(test)] unsafe fn pop_front(&mut self) -> Option>> { - let mut head = self.head?; - debug_assert!(head.as_ref().linked); - - if let Some(mut next) = head.as_mut().take_next() { - next.as_mut().set_prev(None); - self.head = Some(next); - } else { - debug_assert_eq!(self.tail, Some(head)); - self.head = None; - self.tail = None; - } + unsafe { + let mut head = self.head?; + debug_assert!(head.as_ref().linked); + + if let Some(mut next) = head.as_mut().take_next() { + next.as_mut().set_prev(None); + self.head = Some(next); + } else { + debug_assert_eq!(self.tail, Some(head)); + self.head = None; + self.tail = None; + } - debug_assert!(head.as_ref().prev.is_none()); - debug_assert!(head.as_ref().next.is_none()); - head.as_mut().linked = false; - Some(head) + debug_assert!(head.as_ref().prev.is_none()); + debug_assert!(head.as_ref().next.is_none()); + head.as_mut().linked = false; + Some(head) + } } /// Pop the back element from the list. unsafe fn pop_back(&mut self) -> Option>> { - let mut tail = self.tail?; - debug_assert!(tail.as_ref().linked); - - if let Some(mut prev) = tail.as_mut().take_prev() { - prev.as_mut().set_next(None); - self.tail = Some(prev); - } else { - debug_assert_eq!(self.head, Some(tail)); - self.head = None; - self.tail = None; - } + unsafe { + let mut tail = self.tail?; + debug_assert!(tail.as_ref().linked); + + if let Some(mut prev) = tail.as_mut().take_prev() { + prev.as_mut().set_next(None); + self.tail = Some(prev); + } else { + debug_assert_eq!(self.head, Some(tail)); + self.head = None; + self.tail = None; + } - debug_assert!(tail.as_ref().prev.is_none()); - debug_assert!(tail.as_ref().next.is_none()); - tail.as_mut().linked = false; - Some(tail) + debug_assert!(tail.as_ref().prev.is_none()); + debug_assert!(tail.as_ref().next.is_none()); + tail.as_mut().linked = false; + Some(tail) + } } unsafe fn remove(&mut self, mut node: ptr::NonNull>) { - debug_assert!(node.as_ref().linked); + unsafe { + debug_assert!(node.as_ref().linked); - let next = node.as_mut().take_next(); - let prev = node.as_mut().take_prev(); + let next = node.as_mut().take_next(); + let prev = node.as_mut().take_prev(); - if let Some(mut next) = next { - next.as_mut().set_prev(prev); - } else { - debug_assert_eq!(self.tail, Some(node)); - self.tail = prev; - } + if let Some(mut next) = next { + next.as_mut().set_prev(prev); + } else { + debug_assert_eq!(self.tail, Some(node)); + self.tail = prev; + } - if let Some(mut prev) = prev { - prev.as_mut().set_next(next); - } else { - debug_assert_eq!(self.head, Some(node)); - self.head = next; - } + if let Some(mut prev) = prev { + prev.as_mut().set_next(next); + } else { + debug_assert_eq!(self.head, Some(node)); + self.head = next; + } - node.as_mut().linked = false; + node.as_mut().linked = false; + } } unsafe fn front(&mut self) -> Option>> { From fdade083a7c1260f655c868a6ca2dcd5b5d3695a Mon Sep 17 00:00:00 2001 From: suxb201 Date: Tue, 8 Jul 2025 14:12:35 +0800 Subject: [PATCH 04/10] fix compile --- .cursor/rules/project.mdc | 24 ++ Cargo.toml | 4 +- python/resp_benchmark/__init__.py | 4 +- python/resp_benchmark/wrapper.py | 49 ---- src/bench.rs | 121 +--------- src/lib.rs | 96 +------- src/qps_limiter.rs | 380 ++++++++---------------------- src/shared_context.rs | 21 +- 8 files changed, 135 insertions(+), 564 deletions(-) create mode 100644 .cursor/rules/project.mdc diff --git a/.cursor/rules/project.mdc b/.cursor/rules/project.mdc new file mode 100644 index 0000000..5580589 --- /dev/null +++ b/.cursor/rules/project.mdc @@ -0,0 +1,24 @@ +--- +description: +globs: +alwaysApply: true +--- +# resp-benchmark + +Redis 接口数据库的性能测试套件。 + +## Overview + +1. 使用 Python 编写 cli 界面和 lib 接口；使用 Rust 编写性能测试逻辑；使用 maturin 生态编译。 + +## Code Style + +1. 不需要加太多异常捕获，让代码保持简洁。 + +## Core Components + +### Python Interface +- [python/resp_benchmark/cli.py](mdc:python/resp_benchmark/cli.py): Command-line interface implementation + +### Rust Core +- [src/lib.rs](mdc:src/lib.rs): Entry point for Rust library, defines Python bindings diff --git a/Cargo.toml b/Cargo.toml index 14cf73e..41cb730 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,5 +27,5 @@ colored = "2.1.0" enum_delegate = "0.2.0" ctrlc = "3.4.4" urlencoding = "2.1.3" -parking_lot = "0.12.1" -pin-project-lite = "0.2.14" +parking_lot = "0.12" +pin-project-lite = "0.2" \ No newline at end of file diff --git a/python/resp_benchmark/__init__.py b/python/resp_benchmark/__init__.py index 8a54c01..ba0d3d1 100644 --- a/python/resp_benchmark/__init__.py +++ b/python/resp_benchmark/__init__.py @@ -1,3 +1,3 @@ -from .wrapper import Benchmark, Result, BenchmarkWorkerPool +from .wrapper import Benchmark, Result -__all__ = ["Benchmark", "Result", "BenchmarkWorkerPool"] +__all__ = ["Benchmark", "Result"] diff --git a/python/resp_benchmark/wrapper.py b/python/resp_benchmark/wrapper.py index 63d171f..c5ccdcb 100644 --- a/python/resp_benchmark/wrapper.py +++ b/python/resp_benchmark/wrapper.py @@ -7,7 +7,6 @@ from .cores import parse_cores_string -from ._resp_benchmark_rust_lib import BenchmarkWorkerPool @dataclass class Result: @@ -159,51 +158,3 @@ def flushall(self): """ r = redis.Redis(host=self.host, port=self.port, username=self.username, password=self.password) r.flushall() - - class AsyncBenchmarkContext: - def __init__(self, raw_benchmark_context): - self.raw_benchmark_context = raw_benchmark_context - - def stop(self): - self.raw_benchmark_context.stop() - - def join(self): - self.raw_benchmark_context.join() - - def try_join(self): - return self.raw_benchmark_context.try_join() - - def current_result(self) -> Result: - ret = self.raw_benchmark_context.current_result() - return Result( - qps=ret.qps, - avg_latency_ms=ret.avg_latency_ms, - p99_latency_ms=ret.p99_latency_ms, - max_latency_ms=ret.max_latency_ms, - connections=ret.connections - ) - - def async_bench(self, benchmark_pool: BenchmarkWorkerPool, command: str, connections: int = 0, pipeline: int = 1, count: int = 0, target: int = 0, seconds: int = 0, quiet: bool = False) -> AsyncBenchmarkContext: - from . import _resp_benchmark_rust_lib - raw_ctx = _resp_benchmark_rust_lib.async_benchmark( - worker_pool=benchmark_pool, - - host=self.host, - port=self.port, - username=self.username, - password=self.password, - cluster=self.cluster, - tls=False, # TODO: Implement TLS support - timeout=self.timeout, - cores=self.cores, - - command=command, - connections=connections, - pipeline=pipeline, - count=count, - target=target, - seconds=seconds, - load=False, - quiet=quiet, - ) - return self.AsyncBenchmarkContext(raw_ctx) diff --git a/src/bench.rs b/src/bench.rs index 7a7bf8e..f1d51da 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -2,14 +2,13 @@ use awaitgroup::WaitGroup; use colored::Colorize; use std::io::Write; use std::sync::Arc; -use tokio::{select}; +use tokio::{select, task}; +use crate::BenchmarkResult; use crate::auto_connection::{AutoConnection, ConnLimiter}; use crate::client::ClientConfig; use crate::command::Command; use crate::shared_context::SharedContext; -use crate::BenchmarkResult; -use crate::AsyncBenchmarkContext; use crate::qps_limiter::RateLimiter; #[derive(Clone)] @@ -23,14 +22,14 @@ pub struct Case { } async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limiter: Arc>, config: ClientConfig, case: Case, context: SharedContext) { - let mut local = vec![]; + let local = task::LocalSet::new(); for _ in 0..conn_limiter.total_conn { let conn_limiter = conn_limiter.clone(); let qps_limiter = qps_limiter.clone(); let config = config.clone(); let case = case.clone(); let mut context = context.clone(); - let join_handle = tokio::task::spawn(async move { + local.spawn_local(async move { let mut client = config.get_client().await; let mut cmd = case.command.clone(); let conn_limiter = conn_limiter.clone(); @@ -60,21 +59,18 @@ async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limit client.run_commands(p).await; let duration = instant.elapsed().as_micros() as u64; for _ in 0..pipeline_cnt { - context.record(duration); + context.histogram.record(duration); } match qps_limiter.as_ref() { Some(limiter) => { - limiter.acquire(pipeline_cnt as usize).await; + limiter.acquire(pipeline_cnt as usize * 1000).await; } None => {} } } }); - local.push(join_handle); - } - for join_handle in local { - let _ = join_handle.await; } + local.await; } fn wait_finish(case: &Case, mut auto_connection: AutoConnection, mut context: SharedContext, mut wg: WaitGroup, quiet: bool) -> BenchmarkResult { @@ -160,10 +156,11 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo let qps_limiter = Arc::new(if case.target > 0 { Some( RateLimiter::builder() - .max(case.target as usize * 5) + .max(case.target as usize * 1000) + .initial(0) .interval(tokio::time::Duration::from_millis(1)) - .refill(case.target as usize / 1000) - .build() + .refill(case.target as usize) + .build(), ) } else { None @@ -208,99 +205,3 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo return result; } - -async fn async_cron(mut auto_connection: AutoConnection, mut context: SharedContext, mut wg: WaitGroup) { - let histogram = context.histogram.clone(); - // calc overall qps - // for log - let mut interval = tokio::time::interval(std::time::Duration::from_secs(1)); - - let mut log_instance = std::time::Instant::now(); - let mut log_last_cnt = histogram.cnt(); - - if auto_connection.ready { - context.start_timer(); - } - - loop { - select! { - _ = interval.tick() => {} - _ = wg.wait() => {break;} - } - let cnt = histogram.cnt(); - let qps = (cnt - log_last_cnt) as f64 / log_instance.elapsed().as_secs_f64(); - if !auto_connection.ready { - auto_connection.adjust(&histogram); - if auto_connection.ready { - context.start_timer(); - } - } - log_last_cnt = cnt; - log_instance = std::time::Instant::now(); - - let mut result = BenchmarkResult::default(); - let current_histogram_id = context.latest_histogramactive_index.fetch_add(1, std::sync::atomic::Ordering::SeqCst) % 2; - let current_histogram = &context.latest_histogram[current_histogram_id as usize]; - result.qps = qps; - result.connections = auto_connection.active_conn(); - result.avg_latency_ms = current_histogram.avg() as f64 / 1_000.0; - result.p99_latency_ms = current_histogram.percentile(0.99) as f64 / 1_000.0; - result.max_latency_ms = current_histogram.max() as f64 / 1_000.0; - *context.latest_result.lock().unwrap() = result; - } -} - -pub fn do_benchmark_async(pool: Arc, client_config: ClientConfig, cores: Vec, case: Case, load: bool, quiet: bool) -> crate::AsyncBenchmarkContext { - if !quiet { - println!("{}: {}", "command".bold().blue(), case.command.to_string().green().bold()); - println!("{}: {}", "connections".bold().blue(), if case.connections == 0 { "auto".to_string() } else { case.connections.to_string() }); - println!("{}: {}", "count".bold().blue(), case.count); - println!("{}: {}", "target".bold().blue(), if case.target == 0 { "unlimited".to_string() } else { case.target.to_string() }); - println!("{}: {}", "seconds".bold().blue(), case.seconds); - println!("{}: {}", "pipeline".bold().blue(), case.pipeline); - } - - let n_parallel = std::cmp::max(1, std::cmp::min(cores.len(), case.connections as usize)); - - // calc connections - let auto_connection = AutoConnection::new(case.connections, n_parallel as u64); - - // calc target qps - let qps_limiter = Arc::new(if case.target > 0 { - Some( - RateLimiter::builder() - .max(case.target as usize * 5) - .interval(tokio::time::Duration::from_millis(1)) - .refill(case.target as usize / 1000) - .build() - ) - } else { - None - }); - - let wg = WaitGroup::new(); - - let context = SharedContext::new(case.count, case.seconds, load); - for inx in 0..n_parallel { - let client_config = client_config.clone(); - let case = case.clone(); - let context = context.clone(); - let wk = wg.worker(); - let conn_limiter = auto_connection.limiters[inx].clone(); - let qps_limiter = qps_limiter.clone(); - pool.spawn(async move { - run_commands_on_single_thread(conn_limiter, qps_limiter, client_config, case, context).await; - wk.done(); - }); - } - - let cron_ctx = context.clone(); - let join_handle = pool.spawn(async move { - async_cron(auto_connection, cron_ctx, wg).await; - }); - - return AsyncBenchmarkContext { - ctx: context, - join_handle: Option::Some(join_handle), - }; -} diff --git a/src/lib.rs b/src/lib.rs index 5586190..04a13a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,21 +4,18 @@ mod command; mod auto_connection; mod shared_context; mod histogram; -mod qps_limiter; mod async_flag; +mod qps_limiter; use ctrlc; use pyo3::prelude::*; use pyo3::wrap_pyfunction; use crate::command::Command; -use crate::shared_context::SharedContext; /// A Python module implemented in Rust. #[pymodule] fn _resp_benchmark_rust_lib(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(benchmark, m)?)?; - m.add_function(wrap_pyfunction!(async_benchmark, m)?)?; - m.add_class::()?; Ok(()) } @@ -80,95 +77,4 @@ fn benchmark( }; let result = py.allow_threads(|| bench::do_benchmark(client_config, cores, case, load, quiet)); Ok(result) -} - -#[pyclass] -struct BenchmarkWorkerPool { - pool: std::sync::Arc, -} - -#[pymethods] -impl BenchmarkWorkerPool { - #[new] - fn new() -> PyResult { - let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build().map_err(|e| PyErr::new::(e.to_string()))?; - Ok(Self { pool: std::sync::Arc::new(rt) }) - } -} - -#[pyclass] -pub struct AsyncBenchmarkContext { - ctx: SharedContext, - join_handle: Option>, -} - -#[pymethods] -impl AsyncBenchmarkContext { - fn stop(&mut self) { - self.ctx.stop(); - } - - fn join(&mut self) { - match self.join_handle.take() { - None => {} - Some(handle) => { - let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); - let _ = rt.block_on(handle); - } - } - } - - fn try_join(&mut self) -> bool { - self.join_handle.as_ref().map_or(true, |v| v.is_finished()) - } - - fn current_result(&self) -> BenchmarkResult { - let result = self.ctx.latest_result.lock().unwrap().clone(); - return result; - } -} - -#[pyfunction] -fn async_benchmark( - worker_pool: &BenchmarkWorkerPool, - host: String, - port: u16, - username: String, - password: String, - cluster: bool, - tls: bool, - timeout: u64, - cores: Vec, - command: String, - connections: u64, - target: u64, - pipeline: u64, - count: u64, - seconds: u64, - load: bool, - quiet: bool, -) -> PyResult { - if load { - return Err(PyErr::new::("count must be greater than 0".to_string())); - } - - let client_config = client::ClientConfig { - cluster, - address: format!("{}:{}", host, port), - username, - password, - tls, - timeout, - }; - let case = bench::Case { - command: Command::new(command.as_str()), - connections, - pipeline, - count, - target, - seconds, - }; - let pool = worker_pool.pool.clone(); - let result = bench::do_benchmark_async(pool, client_config, cores, case, load, quiet); - Ok(result) } \ No newline at end of file diff --git a/src/qps_limiter.rs b/src/qps_limiter.rs index 11c9257..3167317 100644 --- a/src/qps_limiter.rs +++ b/src/qps_limiter.rs @@ -1,5 +1,4 @@ // Copyright (c) 2019 John-John Tedro - // Permission is hereby granted, free of charge, to any // person obtaining a copy of this software and associated // documentation files (the "Software"), to deal in the @@ -9,11 +8,9 @@ // the Software, and to permit persons to whom the Software // is furnished to do so, subject to the following // conditions: - // The above copyright notice and this permission notice // shall be included in all copies or substantial portions // of the Software. - // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF // ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED // TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A @@ -37,15 +34,11 @@ use std::task::{Context, Poll, Waker}; use std::marker; use std::ops; use std::sync::Arc; - use parking_lot::{Mutex, MutexGuard}; use pin_project_lite::pin_project; use tokio::time::{self, Duration, Instant}; - - /// Default factor for how to calculate max refill value. const DEFAULT_REFILL_MAX_FACTOR: usize = 10; - /// Interval to bump the shared mutex guard to allow other parts of the system /// to make process. Processes which loop should use this number to determine /// how many times it should loop before calling [Guard::bump]. @@ -53,15 +46,12 @@ const DEFAULT_REFILL_MAX_FACTOR: usize = 10; /// If we do not respect this limit we might inadvertently end up starving other /// tasks from making progress so that they can unblock. const BUMP_LIMIT: usize = 16; - /// The maximum supported balance. const MAX_BALANCE: usize = isize::MAX as usize; - /// Marker trait which indicates that a type represents a unique held critical section. trait IsCritical {} impl IsCritical for Critical {} impl IsCritical for Guard<'_> {} - /// Linked task state. struct Task { /// Remaining tokens that need to be satisfied. @@ -72,7 +62,6 @@ struct Task { /// The waker associated with the node. waker: Option, } - impl Task { /// Construct a new task state with the given permits remaining. const fn new() -> Self { @@ -82,12 +71,10 @@ impl Task { waker: None, } } - /// Test if the current node is completed. fn is_completed(&self) -> bool { self.remaining == 0 } - /// Fill the current node from the given pool of tokens and modify it. fn fill(&mut self, current: &mut usize) { let removed = usize::min(self.remaining, *current); @@ -95,90 +82,76 @@ impl Task { *current -= removed; } } - /// A borrowed rate limiter. struct BorrowedRateLimiter<'a>(&'a RateLimiter); - impl Deref for BorrowedRateLimiter<'_> { type Target = RateLimiter; - #[inline] fn deref(&self) -> &RateLimiter { self.0 } } - struct Critical { /// Waiter list. waiters: LinkedList, /// The deadline for when more tokens can be be added. deadline: Instant, } - #[repr(transparent)] struct Guard<'a> { critical: MutexGuard<'a, Critical>, } - impl Guard<'_> { #[inline] fn bump(this: &mut Guard<'_>) { MutexGuard::bump(&mut this.critical) } } - impl Deref for Guard<'_> { type Target = Critical; - #[inline] fn deref(&self) -> &Critical { &self.critical } } - impl DerefMut for Guard<'_> { #[inline] fn deref_mut(&mut self) -> &mut Critical { &mut self.critical } } - impl Critical { #[inline] fn push_task_front(&mut self, task: &mut Node) { - // SAFETY: We both have a mutable access to the node being pushed, and + // SAFETY: We both have mutable access to the node being pushed, and // mutable access to the critical section through `self`. So we know we // have exclusive tampering rights to the waiter queue. unsafe { self.waiters.push_front(task.into()); } } - #[inline] fn push_task(&mut self, task: &mut Node) { - // SAFETY: We both have a mutable access to the node being pushed, and + // SAFETY: We both have mutable access to the node being pushed, and // mutable access to the critical section through `self`. So we know we // have exclusive tampering rights to the waiter queue. unsafe { self.waiters.push_back(task.into()); } } - #[inline] fn remove_task(&mut self, task: &mut Node) { - // SAFETY: We both have a mutable access to the node being pushed, and + // SAFETY: We both have mutable access to the node being pushed, and // mutable access to the critical section through `self`. So we know we // have exclusive tampering rights to the waiter queue. unsafe { self.waiters.remove(task.into()); } } - /// Release the current core. Beyond this point the current task may no /// longer interact exclusively with the core. fn release(&mut self, state: &mut State<'_>) { state.available = true; - // Find another task that might take over as core. Once it has acquired // core status it will have to make sure it is no longer linked into the // wait queue. @@ -191,7 +164,6 @@ impl Critical { } } } - #[derive(Debug)] struct State<'a> { /// Original state. @@ -203,11 +175,9 @@ struct State<'a> { /// The rate limiter the state is associated with. lim: &'a RateLimiter, } - impl<'a> State<'a> { fn try_fast_path(mut self, permits: usize) -> bool { let mut attempts = 0; - // Fast path where we just try to nab any available permit without // locking. // @@ -218,21 +188,16 @@ impl<'a> State<'a> { if attempts == BUMP_LIMIT { break; } - self.balance -= permits; - if let Err(new_state) = self.try_save() { self = new_state; attempts += 1; continue; } - return true; } - false } - /// Add tokens and release any pending tasks. #[inline] fn add_tokens(&mut self, critical: &mut Guard<'_>, tokens: usize, f: F) -> O @@ -246,16 +211,13 @@ impl<'a> State<'a> { tokens, MAX_BALANCE ); - self.balance = (self.balance + tokens).min(self.lim.max); drain_wait_queue(critical, self); let output = f(critical, self); return output; } - f(critical, self) } - #[inline] fn decode(state: usize, lim: &'a RateLimiter) -> Self { State { @@ -265,17 +227,14 @@ impl<'a> State<'a> { lim, } } - #[inline] fn encode(&self) -> usize { (self.balance << 1) | usize::from(self.available) } - /// Try to save the state, but only succeed if it hasn't been modified. #[inline] fn try_save(self) -> Result<(), Self> { let this = ManuallyDrop::new(self); - match this.lim.state.compare_exchange( this.state, this.encode(), @@ -287,14 +246,12 @@ impl<'a> State<'a> { } } } - impl Drop for State<'_> { #[inline] fn drop(&mut self) { self.lim.state.store(self.encode(), Ordering::Release); } } - /// A token-bucket rate limiter. pub struct RateLimiter { /// Tokens to add every `per` duration. @@ -310,7 +267,6 @@ pub struct RateLimiter { /// Critical state of the rate limiter. critical: Mutex, } - impl RateLimiter { /// Construct a new [`Builder`] for a [`RateLimiter`]. /// @@ -331,7 +287,6 @@ impl RateLimiter { pub fn builder() -> Builder { Builder::default() } - /// Get the refill amount of this rate limiter as set through /// [`Builder::refill`]. /// @@ -349,7 +304,6 @@ impl RateLimiter { pub fn refill(&self) -> usize { self.refill } - /// Get the refill interval of this rate limiter as set through /// [`Builder::interval`]. /// @@ -368,7 +322,6 @@ impl RateLimiter { pub fn interval(&self) -> Duration { self.interval } - /// Get the max value of this rate limiter as set through [`Builder::max`]. /// /// # Examples @@ -385,7 +338,6 @@ impl RateLimiter { pub fn max(&self) -> usize { self.max } - /// Test if the current rate limiter is fair as specified through /// [`Builder::fair`]. /// @@ -403,7 +355,6 @@ impl RateLimiter { pub fn is_fair(&self) -> bool { self.fair } - /// Get the current token balance. /// /// This indicates how many tokens can be requested without blocking. @@ -426,7 +377,6 @@ impl RateLimiter { pub fn balance(&self) -> usize { self.state.load(Ordering::Acquire) >> 1 } - /// Acquire a single permit. /// /// # Examples @@ -445,7 +395,6 @@ impl RateLimiter { pub fn acquire_one(&self) -> Acquire<'_> { self.acquire(1) } - /// Acquire the given number of permits, suspending the current task until /// they are available. /// @@ -470,7 +419,6 @@ impl RateLimiter { inner: AcquireFut::new(BorrowedRateLimiter(self), permits), } } - /// Try to acquire the given number of permits, returning `true` if the /// given number of permits were successfully acquired. /// @@ -503,22 +451,17 @@ impl RateLimiter { if self.try_fast_path(permits) { return true; } - let mut critical = self.lock(); - // Reload the state while we are under the critical lock, this // ensures that the `available` flag is up-to-date since it is only // ever modified while holding the critical lock. let mut state = self.take(); - // The core is *not* available, which also implies that there are tasks // ahead which are busy. if !state.available { return false; } - let now = Instant::now(); - // Here we try to assume core duty temporarily to see if we can // release a sufficient number of tokens to allow the current task // to proceed. @@ -526,15 +469,12 @@ impl RateLimiter { state.balance = (state.balance + tokens).min(self.max); critical.deadline = deadline; } - if state.balance >= permits { state.balance -= permits; return true; } - false } - /// Acquire a permit using an owned future. /// /// If zero permits are specified, this function never suspends the current @@ -606,37 +546,30 @@ impl RateLimiter { inner: AcquireFut::new(self, permits), } } - /// Lock the critical section of the rate limiter and return the associated guard. fn lock(&self) -> Guard<'_> { Guard { critical: self.critical.lock(), } } - /// Load the current state. fn load(&self) -> State<'_> { State::decode(self.state.load(Ordering::Acquire), self) } - /// Take the current state, leaving the core state intact. fn take(&self) -> State<'_> { State::decode(self.state.swap(0, Ordering::Acquire), self) } - /// Try to use fast path. fn try_fast_path(&self, permits: usize) -> bool { if permits == 0 { return true; } - if self.fair { return false; } - self.load().try_fast_path(permits) } - /// Calculate refill amount. Returning a tuple of how much to fill and remaining /// duration to sleep until the next refill time if appropriate. /// @@ -647,28 +580,22 @@ impl RateLimiter { if now < deadline { return None; } - // Time elapsed in milliseconds since the last deadline. let millis = self.interval.as_millis(); let since = now.saturating_duration_since(deadline).as_millis(); - let periods = usize::try_from(since / millis + 1).unwrap_or(usize::MAX); - let tokens = periods .checked_mul(self.refill) .unwrap_or(MAX_BALANCE) .min(MAX_BALANCE); - let rem = u64::try_from(since % millis).unwrap_or(u64::MAX); // Calculated time remaining until the next deadline. let next = millis as u64 * periods as u64 - rem; let deadline = deadline.checked_add(time::Duration::from_millis(next)).unwrap(); - Some((tokens, deadline)) } } - impl fmt::Debug for RateLimiter { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("RateLimiter") @@ -679,11 +606,9 @@ impl fmt::Debug for RateLimiter { .finish_non_exhaustive() } } - /// Refill the wait queue with the given number of tokens. fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) { let mut bump = 0; - // SAFETY: we're holding the lock guard to all the waiters so we can be // sure that we have exclusive access to the wait queue. unsafe { @@ -692,23 +617,17 @@ fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) { Some(node) => node, None => break, }; - let n = node.as_mut(); n.fill(&mut state.balance); - - if !n.is_linked() { + if !n.is_completed() { critical.waiters.push_back(node); break; } - n.complete.store(true, Ordering::Release); - if let Some(waker) = n.waker.take() { waker.wake(); } - bump += 1; - if bump == BUMP_LIMIT { Guard::bump(critical); bump = 0; @@ -716,13 +635,11 @@ fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) { } } } - // SAFETY: All the internals of acquire is thread safe and correctly // synchronized. The embedded waiter queue doesn't have anything inherently // unsafe in it. unsafe impl Send for RateLimiter {} unsafe impl Sync for RateLimiter {} - /// A builder for a [`RateLimiter`]. pub struct Builder { /// The max number of tokens. @@ -736,7 +653,6 @@ pub struct Builder { /// If the rate limiter is fair or not. fair: bool, } - impl Builder { /// Configure the max number of tokens to use. /// @@ -761,7 +677,6 @@ impl Builder { self.max = Some(max); self } - /// Configure the initial number of tokens to configure. The default value /// is `0`. /// @@ -778,7 +693,6 @@ impl Builder { self.initial = initial; self } - /// Configure the time duration between which we add [`refill`] number to /// the bucket rate limiter. /// @@ -831,7 +745,6 @@ impl Builder { self.interval = interval; self } - /// The number of tokens to add at each [`interval`] interval. The default /// value is `1`. /// @@ -855,7 +768,6 @@ impl Builder { self.refill = refill; self } - /// Configure the rate limiter to be fair. /// /// Fairness is enabled by deafult. @@ -882,7 +794,6 @@ impl Builder { self.fair = fair; self } - /// Construct a new [`RateLimiter`]. /// /// # Examples @@ -899,10 +810,8 @@ impl Builder { /// ``` pub fn build(&self) -> RateLimiter { let deadline = Instant::now() + self.interval; - let initial = self.initial.min(MAX_BALANCE); let refill = self.refill.min(MAX_BALANCE); - let max = match self.max { Some(max) => max.min(MAX_BALANCE), None => refill @@ -910,9 +819,7 @@ impl Builder { .saturating_mul(DEFAULT_REFILL_MAX_FACTOR) .min(MAX_BALANCE), }; - let initial = initial.min(max); - RateLimiter { refill, interval: self.interval, @@ -926,7 +833,6 @@ impl Builder { } } } - /// Construct a new builder with default options. /// /// # Examples @@ -947,7 +853,6 @@ impl Default for Builder { } } } - /// The state of an acquire operation. #[derive(Debug, Clone, Copy)] enum AcquireFutState { @@ -958,35 +863,29 @@ enum AcquireFutState { /// The operation is completed. Complete, /// The task is currently the core. - Core, + Core {}, } - /// Inner state and methods of the acquire. #[repr(transparent)] struct AcquireFutInner { /// Aliased task state. node: UnsafeCell>, } - impl AcquireFutInner { const fn new() -> AcquireFutInner { AcquireFutInner { node: UnsafeCell::new(Node::new(Task::new())), } } - /// Access the completion flag. pub fn complete(&self) -> &AtomicBool { - // SAFETY: The pointer is created in `AcquireFut`, which is guaranteed - // to be alive as long as `self` is. - unsafe { &*ptr::addr_of!((&*self.node.get()).complete) } + // SAFETY: This is always safe to access since it's atomic. + unsafe { &*ptr::addr_of!((*self.node.get()).complete) } } - - /// Get a mutable reference to the underlying task node. + /// Get the underlying task mutably. /// - /// # Safety - /// - /// This is safe to call as long as `self` is alive. + /// We prove that the caller does indeed have mutable access to the node by + /// passing in a mutable reference to the critical section. #[inline] pub fn get_task<'crit, C>( self: Pin<&'crit mut Self>, @@ -1000,30 +899,24 @@ impl AcquireFutInner { // the borrows outlive the provided closure. unsafe { (critical, &mut *self.node.get()) } } - /// Update the waiting state for this acquisition task. This might require /// that we update the associated waker. fn update(self: Pin<&mut Self>, critical: &mut Guard<'_>, waker: &Waker) { let (critical, task) = self.get_task(critical); - if !task.is_linked() { critical.push_task_front(task); } - let new_waker = match task.waker { None => true, Some(ref w) => !w.will_wake(waker), }; - if new_waker { task.waker = Some(waker.clone()); } } - /// Ensure that the current core task is correctly linked up if needed. fn link_core(self: Pin<&mut Self>, critical: &mut Critical, lim: &RateLimiter) { let (critical, task) = self.get_task(critical); - match (lim.fair, task.is_linked()) { (true, false) => { // Fair scheduling needs to ensure that the core is part of the wait @@ -1038,7 +931,6 @@ impl AcquireFutInner { _ => {} } } - /// Release any remaining tokens which are associated with this particular task. fn release_remaining( self: Pin<&mut Self>, @@ -1047,16 +939,13 @@ impl AcquireFutInner { permits: usize, ) { let (critical, task) = self.get_task(critical); - if task.is_linked() { critical.remove_task(task); } - // Hand back permits which we've acquired so far. let release = permits.saturating_sub(task.remaining); state.add_tokens(critical, release, |_, _| ()); } - /// Drain the given number of tokens through the core. Returns `true` if the /// core has been completed. fn drain_core( @@ -1067,7 +956,6 @@ impl AcquireFutInner { ) -> bool { let completed = state.add_tokens(critical, tokens, |critical, state| { let (_, task) = self.get_task(critical); - // If the limiter is not fair, we need to in addition to draining // remaining tokens from linked nodes, drain it from ourselves. We // fill the current holder of the core last (self). To ensure that @@ -1075,19 +963,15 @@ impl AcquireFutInner { if !state.lim.fair { task.fill(&mut state.balance); } - task.is_completed() }); - if completed { // Everything was drained, including the current core (if // appropriate). So we can release it now. critical.release(state); } - completed } - /// Assume the current core and calculate how long we must sleep for in /// order to do it. /// @@ -1102,18 +986,15 @@ impl AcquireFutInner { now: Instant, ) -> bool { self.as_mut().link_core(critical, state.lim); - let (tokens, deadline) = match state.lim.calculate_drain(critical.deadline, now) { Some(tokens) => tokens, None => return true, }; - // It is appropriate to update the deadline. critical.deadline = deadline; !self.drain_core(critical, state, tokens) } } - pin_project! { /// The future associated with acquiring permits from a rate limiter using /// [`RateLimiter::acquire`]. @@ -1123,7 +1004,6 @@ pin_project! { inner: AcquireFut>, } } - impl Acquire<'_> { /// Test if this acquire task is currently coordinating the rate limiter. /// @@ -1161,16 +1041,13 @@ impl Acquire<'_> { self.inner.is_core() } } - impl Future for Acquire<'_> { type Output = (); - #[inline] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { self.project().inner.poll(cx) } } - pin_project! { /// The future associated with acquiring permits from a rate limiter using /// [`RateLimiter::acquire_owned`]. @@ -1180,7 +1057,6 @@ pin_project! { inner: AcquireFut>, } } - impl AcquireOwned { /// Test if this acquire task is currently coordinating the rate limiter. /// @@ -1218,16 +1094,13 @@ impl AcquireOwned { self.inner.is_core() } } - impl Future for AcquireOwned { type Output = (); - #[inline] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { self.project().inner.poll(cx) } } - pin_project! { #[project(!Unpin)] #[project = AcquireFutProj] @@ -1243,33 +1116,27 @@ pin_project! { #[pin] inner: AcquireFutInner, } - impl PinnedDrop for AcquireFut where T: Deref, { fn drop(this: Pin<&mut Self>) { let AcquireFutProj { lim, permits, state, inner, .. } = this.project(); - let is_core = match *state { AcquireFutState::Waiting => false, AcquireFutState::Core { .. } => true, _ => return, }; - let mut critical = lim.lock(); let mut s = lim.take(); inner.release_remaining(&mut critical, &mut s, *permits); - if is_core { critical.release(&mut s); } - *state = AcquireFutState::Complete; } } } - impl AcquireFut where T: Deref, @@ -1284,24 +1151,20 @@ where inner: AcquireFutInner::new(), } } - fn is_core(&self) -> bool { matches!(&self.state, AcquireFutState::Core { .. }) } } - // SAFETY: All the internals of acquire is thread safe and correctly // synchronized. The embedded waiter queue doesn't have anything inherently // unsafe in it. unsafe impl Send for AcquireFut where T: Send + Deref {} unsafe impl Sync for AcquireFut where T: Sync + Deref {} - impl Future for AcquireFut where T: Deref, { type Output = (); - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let AcquireFutProj { lim, @@ -1311,22 +1174,18 @@ where inner: mut internal, .. } = self.project(); - // Hold onto the critical lock for core operations, but only acquire it // when strictly necessary. let mut critical; - // Shared state. // // Once we are holding onto the critical lock, we take the entire state // to ensure that any fast-past negotiators do not observe any available // permits while potential core work is ongoing. let mut s; - // Hold onto any call to `Instant::now` which we might perform, so we // don't have to get the current time multiple times. let outer_now; - match *state { AcquireFutState::Complete => { return Poll::Ready(()); @@ -1341,12 +1200,9 @@ where *state = AcquireFutState::Complete; return Poll::Ready(()); } - critical = lim.lock(); s = lim.take(); - let now = Instant::now(); - // If we've hit a deadline, calculate the number of tokens // to drain and perform it in line here. This is necessary // because the core isn't aware of how long we sleep between @@ -1372,19 +1228,16 @@ where } else { 0 }; - let completed = s.add_tokens(&mut critical, tokens, |critical, s| { let (_, task) = internal.as_mut().get_task(critical); task.remaining = *permits; task.fill(&mut s.balance); task.is_completed() }); - if completed { *state = AcquireFutState::Complete; return Poll::Ready(()); } - // Try to take over as core. If we're unsuccessful we just // ensure that we're linked into the wait queue. if !mem::take(&mut s.available) { @@ -1392,13 +1245,12 @@ where *state = AcquireFutState::Waiting; return Poll::Pending; } - // SAFETY: This is done in a pinned section, so we know that // the linked section stays alive for the duration of this // future due to pinning guarantees. internal.as_mut().link_core(&mut critical, lim); Guard::bump(&mut critical); - *state = AcquireFutState::Core; + *state = AcquireFutState::Core {}; outer_now = Some(now); } AcquireFutState::Waiting => { @@ -1410,22 +1262,18 @@ where *state = AcquireFutState::Complete; return Poll::Ready(()); } - // Note: we need to operate under this lock to ensure that // the core acquired here (or elsewhere) observes that the // current task has been linked up. critical = lim.lock(); s = lim.take(); - // Try to take over as core. If we're unsuccessful we // just ensure that we're linked into the wait queue. if !mem::take(&mut s.available) { internal.update(&mut critical, cx.waker()); return Poll::Pending; } - let now = Instant::now(); - // This is done in a pinned section, so we know that the linked // section stays alive for the duration of this future due to // pinning guarantees. @@ -1434,24 +1282,21 @@ where *state = AcquireFutState::Complete; return Poll::Ready(()); } - Guard::bump(&mut critical); - *state = AcquireFutState::Core; + *state = AcquireFutState::Core {}; outer_now = Some(now); } - AcquireFutState::Core => { + AcquireFutState::Core {} => { critical = lim.lock(); s = lim.take(); outer_now = None; } } - let mut sleep = match sleep.as_mut().as_pin_mut() { Some(mut sleep) => { if sleep.deadline() != critical.deadline { sleep.as_mut().reset(critical.deadline); } - sleep } None => { @@ -1459,23 +1304,18 @@ where sleep.as_mut().as_pin_mut().unwrap() } }; - if sleep.as_mut().poll(cx).is_pending() { return Poll::Pending; } - critical.deadline = outer_now.unwrap_or(sleep.deadline()) + lim.interval; - if internal.drain_core(&mut critical, &mut s, lim.refill) { *state = AcquireFutState::Complete; return Poll::Ready(()); } - cx.waker().wake_by_ref(); Poll::Pending } } - pub struct Node { /// The next node. next: Option>>, @@ -1489,10 +1329,9 @@ pub struct Node { /// struct. _pin: marker::PhantomPinned, } - impl Node { /// Construct a new unlinked node. - pub(crate) const fn new(value: T) -> Self { + const fn new(value: T) -> Self { Self { next: None, prev: None, @@ -1501,41 +1340,37 @@ impl Node { _pin: marker::PhantomPinned, } } - - /// Test if the current node is linked. - pub(crate) fn is_linked(&self) -> bool { + #[inline(always)] + fn is_linked(&self) -> bool { self.linked } - /// Set the next node. - pub(crate) unsafe fn set_next(&mut self, node: Option>) { - self.next = node; + #[inline(always)] + unsafe fn set_next(&mut self, node: Option>) { + ptr::addr_of_mut!(self.next).write(node); } - /// Take the next node. - pub(crate) unsafe fn take_next(&mut self) -> Option> { - self.next.take() + #[inline(always)] + unsafe fn take_next(&mut self) -> Option> { + ptr::addr_of_mut!(self.next).replace(None) } - /// Set the previous node. - pub(crate) unsafe fn set_prev(&mut self, node: Option>) { - self.prev = node; + #[inline(always)] + unsafe fn set_prev(&mut self, node: Option>) { + ptr::addr_of_mut!(self.prev).write(node); } - /// Take the previous node. - pub(crate) unsafe fn take_prev(&mut self) -> Option> { - self.prev.take() + #[inline(always)] + unsafe fn take_prev(&mut self) -> Option> { + ptr::addr_of_mut!(self.prev).replace(None) } } - impl ops::Deref for Node { type Target = T; - fn deref(&self) -> &Self::Target { &self.value } } - impl ops::DerefMut for Node where T: Unpin, @@ -1544,12 +1379,10 @@ where &mut self.value } } - pub struct LinkedList { head: Option>>, tail: Option>>, } - impl LinkedList { /// Construct a new empty list. const fn new() -> Self { @@ -1558,119 +1391,90 @@ impl LinkedList { tail: None, } } - unsafe fn push_front(&mut self, mut node: ptr::NonNull>) { - unsafe { - debug_assert!(node.as_ref().next.is_none()); - debug_assert!(node.as_ref().prev.is_none()); - debug_assert!(!node.as_ref().linked); - - if let Some(mut head) = self.head.take() { - node.as_mut().set_next(Some(head)); - head.as_mut().set_prev(Some(node)); - self.head = Some(node); - } else { - self.head = Some(node); - self.tail = Some(node); - } - - node.as_mut().linked = true; + debug_assert!(node.as_ref().next.is_none()); + debug_assert!(node.as_ref().prev.is_none()); + debug_assert!(!node.as_ref().linked); + if let Some(mut head) = self.head.take() { + node.as_mut().set_next(Some(head)); + head.as_mut().set_prev(Some(node)); + self.head = Some(node); + } else { + self.head = Some(node); + self.tail = Some(node); } + node.as_mut().linked = true; } - unsafe fn push_back(&mut self, mut node: ptr::NonNull>) { - unsafe { - debug_assert!(node.as_ref().next.is_none()); - debug_assert!(node.as_ref().prev.is_none()); - debug_assert!(!node.as_ref().linked); - - if let Some(mut tail) = self.tail.take() { - node.as_mut().set_prev(Some(tail)); - tail.as_mut().set_next(Some(node)); - self.tail = Some(node); - } else { - self.head = Some(node); - self.tail = Some(node); - } - - node.as_mut().linked = true; + debug_assert!(node.as_ref().next.is_none()); + debug_assert!(node.as_ref().prev.is_none()); + debug_assert!(!node.as_ref().linked); + if let Some(mut tail) = self.tail.take() { + node.as_mut().set_prev(Some(tail)); + tail.as_mut().set_next(Some(node)); + self.tail = Some(node); + } else { + self.head = Some(node); + self.tail = Some(node); } + node.as_mut().linked = true; } - #[cfg(test)] unsafe fn pop_front(&mut self) -> Option>> { - unsafe { - let mut head = self.head?; - debug_assert!(head.as_ref().linked); - - if let Some(mut next) = head.as_mut().take_next() { - next.as_mut().set_prev(None); - self.head = Some(next); - } else { - debug_assert_eq!(self.tail, Some(head)); - self.head = None; - self.tail = None; - } - - debug_assert!(head.as_ref().prev.is_none()); - debug_assert!(head.as_ref().next.is_none()); - head.as_mut().linked = false; - Some(head) + let mut head = self.head?; + debug_assert!(head.as_ref().linked); + if let Some(mut next) = head.as_mut().take_next() { + next.as_mut().set_prev(None); + self.head = Some(next); + } else { + debug_assert_eq!(self.tail, Some(head)); + self.head = None; + self.tail = None; } + debug_assert!(head.as_ref().prev.is_none()); + debug_assert!(head.as_ref().next.is_none()); + head.as_mut().linked = false; + Some(head) } - /// Pop the back element from the list. unsafe fn pop_back(&mut self) -> Option>> { - unsafe { - let mut tail = self.tail?; - debug_assert!(tail.as_ref().linked); - - if let Some(mut prev) = tail.as_mut().take_prev() { - prev.as_mut().set_next(None); - self.tail = Some(prev); - } else { - debug_assert_eq!(self.head, Some(tail)); - self.head = None; - self.tail = None; - } - - debug_assert!(tail.as_ref().prev.is_none()); - debug_assert!(tail.as_ref().next.is_none()); - tail.as_mut().linked = false; - Some(tail) + let mut tail = self.tail?; + debug_assert!(tail.as_ref().linked); + if let Some(mut prev) = tail.as_mut().take_prev() { + prev.as_mut().set_next(None); + self.tail = Some(prev); + } else { + debug_assert_eq!(self.head, Some(tail)); + self.head = None; + self.tail = None; } + debug_assert!(tail.as_ref().prev.is_none()); + debug_assert!(tail.as_ref().next.is_none()); + tail.as_mut().linked = false; + Some(tail) } - unsafe fn remove(&mut self, mut node: ptr::NonNull>) { - unsafe { - debug_assert!(node.as_ref().linked); - - let next = node.as_mut().take_next(); - let prev = node.as_mut().take_prev(); - - if let Some(mut next) = next { - next.as_mut().set_prev(prev); - } else { - debug_assert_eq!(self.tail, Some(node)); - self.tail = prev; - } - - if let Some(mut prev) = prev { - prev.as_mut().set_next(next); - } else { - debug_assert_eq!(self.head, Some(node)); - self.head = next; - } - - node.as_mut().linked = false; + debug_assert!(node.as_ref().linked); + let next = node.as_mut().take_next(); + let prev = node.as_mut().take_prev(); + if let Some(mut next) = next { + next.as_mut().set_prev(prev); + } else { + debug_assert_eq!(self.tail, Some(node)); + self.tail = prev; + } + if let Some(mut prev) = prev { + prev.as_mut().set_next(next); + } else { + debug_assert_eq!(self.head, Some(node)); + self.head = next; } + node.as_mut().linked = false; } - unsafe fn front(&mut self) -> Option>> { self.head } } - impl fmt::Debug for LinkedList { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("LinkedList") @@ -1678,4 +1482,4 @@ impl fmt::Debug for LinkedList { .field("tail", &self.tail) .finish() } -} +} \ No newline at end of file diff --git a/src/shared_context.rs b/src/shared_context.rs index bcae4c7..75e9230 100644 --- a/src/shared_context.rs +++ b/src/shared_context.rs @@ -1,10 +1,9 @@ use crate::async_flag::AsyncFlag; use crate::histogram::Histogram; -use crate::BenchmarkResult; use std::cmp::min; use std::option::Option; -use std::sync::atomic::{AtomicU64, AtomicU8}; -use std::sync::{Arc, Mutex, RwLock}; +use std::sync::atomic::AtomicU64; +use std::sync::{Arc, RwLock}; use std::time::Instant; #[derive(Clone)] @@ -23,11 +22,6 @@ pub struct SharedContext { // histogram pub histogram: Arc, - - // latest histogram - pub latest_histogram: Arc<[Histogram; 2]>, - pub latest_histogramactive_index: Arc, - pub latest_result: Arc> } impl SharedContext { @@ -41,10 +35,6 @@ impl SharedContext { stop_flag: AsyncFlag::new(), histogram: Arc::new(Histogram::new()), - - latest_histogram: Arc::new([Histogram::new(), Histogram::new()]), - latest_histogramactive_index: Arc::new(AtomicU8::new(0)), - latest_result: Arc::new(Mutex::new(BenchmarkResult::default())), } } @@ -84,10 +74,5 @@ impl SharedContext { return result; } - - pub fn record(&self, latency_us: u64) { - self.histogram.record(latency_us); - let current_index = self.latest_histogramactive_index.load(std::sync::atomic::Ordering::Relaxed) as usize % 2; - self.latest_histogram[current_index].record(latency_us); - } + } From b458758005a3520e5fcd1cd18e87b28c7b8364b1 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Tue, 8 Jul 2025 17:35:08 +0800 Subject: [PATCH 05/10] replace --- Cargo.lock | 225 ++++++- Cargo.toml | 3 +- src/bench.rs | 16 +- src/lib.rs | 1 - src/qps_limiter.rs | 1485 -------------------------------------------- 5 files changed, 229 insertions(+), 1501 deletions(-) delete mode 100644 src/qps_limiter.rs diff --git a/Cargo.lock b/Cargo.lock index a27647c..b54adc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,12 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "autocfg" version = "1.5.0" @@ -50,6 +56,12 @@ version = "2.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" +[[package]] +name = "bumpalo" +version = "3.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" + [[package]] name = "bytes" version = "1.10.1" @@ -109,6 +121,12 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "338089f42c427b86394a5ee60ff321da23a5c89c9d89514c829687b26359fcff" +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "ctrlc" version = "3.4.7" @@ -119,6 +137,20 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -154,6 +186,18 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "form_urlencoded" version = "1.2.1" @@ -181,6 +225,12 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + [[package]] name = "futures-util" version = "0.3.31" @@ -213,9 +263,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi", "wasi 0.14.2+wasi-0.2.4", + "wasm-bindgen", ] [[package]] @@ -224,6 +276,46 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "governor" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbe789d04bf14543f03c4b60cd494148aa79438c8440ae7d81a7778147745c3" +dependencies = [ + "cfg-if", + "dashmap", + "futures-sink", + "futures-timer", + "futures-util", + "getrandom 0.3.3", + "hashbrown 0.15.4", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand 0.9.1", + "smallvec", + "spinning_top", + "web-time", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + [[package]] name = "heck" version = "0.5.0" @@ -366,6 +458,16 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" +[[package]] +name = "js-sys" +version = "0.3.77" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -469,6 +571,12 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "num-bigint" version = "0.4.6" @@ -659,6 +767,21 @@ dependencies = [ "syn 2.0.104", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi 0.11.1+wasi-snapshot-preview1", + "web-sys", + "winapi", +] + [[package]] name = "quote" version = "1.0.40" @@ -743,6 +866,15 @@ dependencies = [ "rand 0.8.5", ] +[[package]] +name = "raw-cpuid" +version = "11.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6df7ab838ed27997ba19a4664507e6f82b41fe6e20be42929332156e5e85146" +dependencies = [ + "bitflags", +] + [[package]] name = "redis" version = "0.32.3" @@ -787,9 +919,8 @@ dependencies = [ "core_affinity", "ctrlc", "enum_delegate", + "governor", "nom", - "parking_lot", - "pin-project-lite", "pyo3", "rand 0.8.5", "rand_distr", @@ -873,6 +1004,15 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + [[package]] name = "stable_deref_trait" version = "1.2.0" @@ -930,9 +1070,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.46.0" +version = "1.46.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1140bb80481756a8cbe10541f37433b459c5aa1e727b4c020fbfebdc25bf3ec4" +checksum = "0cc3a2344dafbe23a245241fe8b09735b521110d30fcefbbd5feb1797ca35d17" dependencies = [ "backtrace", "bytes", @@ -1022,6 +1162,83 @@ dependencies = [ "wit-bindgen-rt", ] +[[package]] +name = "wasm-bindgen" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" +dependencies = [ + "cfg-if", + "once_cell", + "wasm-bindgen-macro", +] + +[[package]] +name = "wasm-bindgen-backend" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" +dependencies = [ + "bumpalo", + "log", + "proc-macro2", + "quote", + "syn 2.0.104", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", + "wasm-bindgen-backend", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.77" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33b6dd2ef9186f1f2072e409e99cd22a975331a6b3591b12c764e0e55c60d5d2" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/Cargo.toml b/Cargo.toml index 41cb730..b4bb1b5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,5 +27,4 @@ colored = "2.1.0" enum_delegate = "0.2.0" ctrlc = "3.4.4" urlencoding = "2.1.3" -parking_lot = "0.12" -pin-project-lite = "0.2" \ No newline at end of file +governor = "0.10.0" diff --git a/src/bench.rs b/src/bench.rs index f1d51da..77c74c6 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -1,6 +1,7 @@ use awaitgroup::WaitGroup; use colored::Colorize; use std::io::Write; +use std::num::NonZeroU32; use std::sync::Arc; use tokio::{select, task}; @@ -9,7 +10,7 @@ use crate::auto_connection::{AutoConnection, ConnLimiter}; use crate::client::ClientConfig; use crate::command::Command; use crate::shared_context::SharedContext; -use crate::qps_limiter::RateLimiter; +use governor::{Quota, DefaultDirectRateLimiter}; #[derive(Clone)] pub struct Case { @@ -21,7 +22,7 @@ pub struct Case { pub pipeline: u64, } -async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limiter: Arc>, config: ClientConfig, case: Case, context: SharedContext) { +async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limiter: Arc>, config: ClientConfig, case: Case, context: SharedContext) { let local = task::LocalSet::new(); for _ in 0..conn_limiter.total_conn { let conn_limiter = conn_limiter.clone(); @@ -63,7 +64,9 @@ async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limit } match qps_limiter.as_ref() { Some(limiter) => { - limiter.acquire(pipeline_cnt as usize * 1000).await; + for _ in 0..pipeline_cnt { + limiter.until_ready().await; + } } None => {} } @@ -155,12 +158,7 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo // calc target qps let qps_limiter = Arc::new(if case.target > 0 { Some( - RateLimiter::builder() - .max(case.target as usize * 1000) - .initial(0) - .interval(tokio::time::Duration::from_millis(1)) - .refill(case.target as usize) - .build(), + DefaultDirectRateLimiter::direct(Quota::per_second(NonZeroU32::new(case.target as u32).unwrap())), ) } else { None diff --git a/src/lib.rs b/src/lib.rs index 04a13a3..7fd1a92 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,7 +5,6 @@ mod auto_connection; mod shared_context; mod histogram; mod async_flag; -mod qps_limiter; use ctrlc; use pyo3::prelude::*; diff --git a/src/qps_limiter.rs b/src/qps_limiter.rs deleted file mode 100644 index 3167317..0000000 --- a/src/qps_limiter.rs +++ /dev/null @@ -1,1485 +0,0 @@ -// Copyright (c) 2019 John-John Tedro -// Permission is hereby granted, free of charge, to any -// person obtaining a copy of this software and associated -// documentation files (the "Software"), to deal in the -// Software without restriction, including without -// limitation the rights to use, copy, modify, merge, -// publish, distribute, sublicense, and/or sell copies of -// the Software, and to permit persons to whom the Software -// is furnished to do so, subject to the following -// conditions: -// The above copyright notice and this permission notice -// shall be included in all copies or substantial portions -// of the Software. -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF -// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED -// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A -// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT -// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY -// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION -// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR -// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -// DEALINGS IN THE SOFTWAREo -#![allow(dead_code)] -use std::cell::UnsafeCell; -use std::convert::TryFrom as _; -use std::fmt; -use std::future::Future; -use std::mem::{self, ManuallyDrop}; -use std::ops::{Deref, DerefMut}; -use std::pin::Pin; -use std::ptr; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Poll, Waker}; -use std::marker; -use std::ops; -use std::sync::Arc; -use parking_lot::{Mutex, MutexGuard}; -use pin_project_lite::pin_project; -use tokio::time::{self, Duration, Instant}; -/// Default factor for how to calculate max refill value. -const DEFAULT_REFILL_MAX_FACTOR: usize = 10; -/// Interval to bump the shared mutex guard to allow other parts of the system -/// to make process. Processes which loop should use this number to determine -/// how many times it should loop before calling [Guard::bump]. -/// -/// If we do not respect this limit we might inadvertently end up starving other -/// tasks from making progress so that they can unblock. -const BUMP_LIMIT: usize = 16; -/// The maximum supported balance. -const MAX_BALANCE: usize = isize::MAX as usize; -/// Marker trait which indicates that a type represents a unique held critical section. -trait IsCritical {} -impl IsCritical for Critical {} -impl IsCritical for Guard<'_> {} -/// Linked task state. -struct Task { - /// Remaining tokens that need to be satisfied. - remaining: usize, - /// If this node has been released or not. We make this an atomic to permit - /// access to it without synchronization. - complete: AtomicBool, - /// The waker associated with the node. - waker: Option, -} -impl Task { - /// Construct a new task state with the given permits remaining. - const fn new() -> Self { - Self { - remaining: 0, - complete: AtomicBool::new(false), - waker: None, - } - } - /// Test if the current node is completed. - fn is_completed(&self) -> bool { - self.remaining == 0 - } - /// Fill the current node from the given pool of tokens and modify it. - fn fill(&mut self, current: &mut usize) { - let removed = usize::min(self.remaining, *current); - self.remaining -= removed; - *current -= removed; - } -} -/// A borrowed rate limiter. -struct BorrowedRateLimiter<'a>(&'a RateLimiter); -impl Deref for BorrowedRateLimiter<'_> { - type Target = RateLimiter; - #[inline] - fn deref(&self) -> &RateLimiter { - self.0 - } -} -struct Critical { - /// Waiter list. - waiters: LinkedList, - /// The deadline for when more tokens can be be added. - deadline: Instant, -} -#[repr(transparent)] -struct Guard<'a> { - critical: MutexGuard<'a, Critical>, -} -impl Guard<'_> { - #[inline] - fn bump(this: &mut Guard<'_>) { - MutexGuard::bump(&mut this.critical) - } -} -impl Deref for Guard<'_> { - type Target = Critical; - #[inline] - fn deref(&self) -> &Critical { - &self.critical - } -} -impl DerefMut for Guard<'_> { - #[inline] - fn deref_mut(&mut self) -> &mut Critical { - &mut self.critical - } -} -impl Critical { - #[inline] - fn push_task_front(&mut self, task: &mut Node) { - // SAFETY: We both have mutable access to the node being pushed, and - // mutable access to the critical section through `self`. So we know we - // have exclusive tampering rights to the waiter queue. - unsafe { - self.waiters.push_front(task.into()); - } - } - #[inline] - fn push_task(&mut self, task: &mut Node) { - // SAFETY: We both have mutable access to the node being pushed, and - // mutable access to the critical section through `self`. So we know we - // have exclusive tampering rights to the waiter queue. - unsafe { - self.waiters.push_back(task.into()); - } - } - #[inline] - fn remove_task(&mut self, task: &mut Node) { - // SAFETY: We both have mutable access to the node being pushed, and - // mutable access to the critical section through `self`. So we know we - // have exclusive tampering rights to the waiter queue. - unsafe { - self.waiters.remove(task.into()); - } - } - /// Release the current core. Beyond this point the current task may no - /// longer interact exclusively with the core. - fn release(&mut self, state: &mut State<'_>) { - state.available = true; - // Find another task that might take over as core. Once it has acquired - // core status it will have to make sure it is no longer linked into the - // wait queue. - unsafe { - if let Some(node) = self.waiters.front() { - if let Some(ref waker) = node.as_ref().waker { - waker.wake_by_ref(); - } - } - } - } -} -#[derive(Debug)] -struct State<'a> { - /// Original state. - state: usize, - /// If the core is available or not. - available: bool, - /// The balance. - balance: usize, - /// The rate limiter the state is associated with. - lim: &'a RateLimiter, -} -impl<'a> State<'a> { - fn try_fast_path(mut self, permits: usize) -> bool { - let mut attempts = 0; - // Fast path where we just try to nab any available permit without - // locking. - // - // We do have to race against anyone else grabbing permits here when - // storing the state back. - while self.balance >= permits { - // Abandon fast path if we've tried too many times. - if attempts == BUMP_LIMIT { - break; - } - self.balance -= permits; - if let Err(new_state) = self.try_save() { - self = new_state; - attempts += 1; - continue; - } - return true; - } - false - } - /// Add tokens and release any pending tasks. - #[inline] - fn add_tokens(&mut self, critical: &mut Guard<'_>, tokens: usize, f: F) -> O - where - F: FnOnce(&mut Guard<'_>, &mut State) -> O, - { - if tokens > 0 { - debug_assert!( - tokens <= MAX_BALANCE, - "Additional tokens {} must be less than {}", - tokens, - MAX_BALANCE - ); - self.balance = (self.balance + tokens).min(self.lim.max); - drain_wait_queue(critical, self); - let output = f(critical, self); - return output; - } - f(critical, self) - } - #[inline] - fn decode(state: usize, lim: &'a RateLimiter) -> Self { - State { - state, - available: state & 1 == 1, - balance: state >> 1, - lim, - } - } - #[inline] - fn encode(&self) -> usize { - (self.balance << 1) | usize::from(self.available) - } - /// Try to save the state, but only succeed if it hasn't been modified. - #[inline] - fn try_save(self) -> Result<(), Self> { - let this = ManuallyDrop::new(self); - match this.lim.state.compare_exchange( - this.state, - this.encode(), - Ordering::Release, - Ordering::Relaxed, - ) { - Ok(_) => Ok(()), - Err(state) => Err(State::decode(state, this.lim)), - } - } -} -impl Drop for State<'_> { - #[inline] - fn drop(&mut self) { - self.lim.state.store(self.encode(), Ordering::Release); - } -} -/// A token-bucket rate limiter. -pub struct RateLimiter { - /// Tokens to add every `per` duration. - refill: usize, - /// Interval in milliseconds to add tokens. - interval: Duration, - /// Max number of tokens associated with the rate limiter. - max: usize, - /// If the rate limiter is fair or not. - fair: bool, - /// The state of the rate limiter. - state: AtomicUsize, - /// Critical state of the rate limiter. - critical: Mutex, -} -impl RateLimiter { - /// Construct a new [`Builder`] for a [`RateLimiter`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use tokio::time::Duration; - /// - /// let limiter = RateLimiter::builder() - /// .initial(100) - /// .refill(100) - /// .max(1000) - /// .interval(Duration::from_millis(250)) - /// .fair(false) - /// .build(); - /// ``` - pub fn builder() -> Builder { - Builder::default() - } - /// Get the refill amount of this rate limiter as set through - /// [`Builder::refill`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .refill(1024) - /// .build(); - /// - /// assert_eq!(limiter.refill(), 1024); - /// ``` - pub fn refill(&self) -> usize { - self.refill - } - /// Get the refill interval of this rate limiter as set through - /// [`Builder::interval`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use tokio::time::Duration; - /// - /// let limiter = RateLimiter::builder() - /// .interval(Duration::from_millis(1000)) - /// .build(); - /// - /// assert_eq!(limiter.interval(), Duration::from_millis(1000)); - /// ``` - pub fn interval(&self) -> Duration { - self.interval - } - /// Get the max value of this rate limiter as set through [`Builder::max`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .max(1024) - /// .build(); - /// - /// assert_eq!(limiter.max(), 1024); - /// ``` - pub fn max(&self) -> usize { - self.max - } - /// Test if the current rate limiter is fair as specified through - /// [`Builder::fair`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .fair(true) - /// .build(); - /// - /// assert_eq!(limiter.is_fair(), true); - /// ``` - pub fn is_fair(&self) -> bool { - self.fair - } - /// Get the current token balance. - /// - /// This indicates how many tokens can be requested without blocking. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = RateLimiter::builder() - /// .initial(100) - /// .build(); - /// - /// assert_eq!(limiter.balance(), 100); - /// limiter.acquire(10).await; - /// assert_eq!(limiter.balance(), 90); - /// # } - /// ``` - pub fn balance(&self) -> usize { - self.state.load(Ordering::Acquire) >> 1 - } - /// Acquire a single permit. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = RateLimiter::builder() - /// .initial(10) - /// .build(); - /// - /// limiter.acquire_one().await; - /// # } - /// ``` - pub fn acquire_one(&self) -> Acquire<'_> { - self.acquire(1) - } - /// Acquire the given number of permits, suspending the current task until - /// they are available. - /// - /// If zero permits are specified, this function never suspends the current - /// task. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = RateLimiter::builder() - /// .initial(10) - /// .build(); - /// - /// limiter.acquire(10).await; - /// # } - /// ``` - pub fn acquire(&self, permits: usize) -> Acquire<'_> { - Acquire { - inner: AcquireFut::new(BorrowedRateLimiter(self), permits), - } - } - /// Try to acquire the given number of permits, returning `true` if the - /// given number of permits were successfully acquired. - /// - /// If the scheduler is fair, and there are pending tasks waiting to acquire - /// tokens this method will return `false`. - /// - /// If zero permits are specified, this method returns `true`. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use tokio::time; - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = RateLimiter::builder().refill(1).initial(1).build(); - /// - /// assert!(limiter.try_acquire(1)); - /// assert!(!limiter.try_acquire(1)); - /// assert!(limiter.try_acquire(0)); - /// - /// time::sleep(limiter.interval() * 2).await; - /// - /// assert!(limiter.try_acquire(1)); - /// assert!(limiter.try_acquire(1)); - /// assert!(!limiter.try_acquire(1)); - /// # } - /// ``` - pub fn try_acquire(&self, permits: usize) -> bool { - if self.try_fast_path(permits) { - return true; - } - let mut critical = self.lock(); - // Reload the state while we are under the critical lock, this - // ensures that the `available` flag is up-to-date since it is only - // ever modified while holding the critical lock. - let mut state = self.take(); - // The core is *not* available, which also implies that there are tasks - // ahead which are busy. - if !state.available { - return false; - } - let now = Instant::now(); - // Here we try to assume core duty temporarily to see if we can - // release a sufficient number of tokens to allow the current task - // to proceed. - if let Some((tokens, deadline)) = self.calculate_drain(critical.deadline, now) { - state.balance = (state.balance + tokens).min(self.max); - critical.deadline = deadline; - } - if state.balance >= permits { - state.balance -= permits; - return true; - } - false - } - /// Acquire a permit using an owned future. - /// - /// If zero permits are specified, this function never suspends the current - /// task. - /// - /// This required the [`RateLimiter`] to be wrapped inside of an - /// [`std::sync::Arc`] but will in contrast permit the acquire operation to - /// be owned by another struct making it more suitable for embedding. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use std::sync::Arc; - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = Arc::new(RateLimiter::builder().initial(10).build()); - /// - /// limiter.acquire_owned(10).await; - /// # } - /// ``` - /// - /// Example when embedded into another future. This wouldn't be possible - /// with [`RateLimiter::acquire`] since it would otherwise hold a reference - /// to the corresponding [`RateLimiter`] instance. - /// - /// ``` - /// use std::future::Future; - /// use std::pin::Pin; - /// use std::sync::Arc; - /// use std::task::{Context, Poll}; - /// - /// use leaky_bucket::{AcquireOwned, RateLimiter}; - /// use pin_project::pin_project; - /// - /// #[pin_project] - /// struct MyFuture { - /// limiter: Arc, - /// #[pin] - /// acquire: Option, - /// } - /// - /// impl Future for MyFuture { - /// type Output = (); - /// - /// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - /// let mut this = self.project(); - /// - /// loop { - /// if let Some(acquire) = this.acquire.as_mut().as_pin_mut() { - /// futures::ready!(acquire.poll(cx)); - /// return Poll::Ready(()); - /// } - /// - /// this.acquire.set(Some(this.limiter.clone().acquire_owned(100))); - /// } - /// } - /// } - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = Arc::new(RateLimiter::builder().initial(100).build()); - /// - /// let future = MyFuture { limiter, acquire: None }; - /// future.await; - /// # } - /// ``` - pub fn acquire_owned(self: Arc, permits: usize) -> AcquireOwned { - AcquireOwned { - inner: AcquireFut::new(self, permits), - } - } - /// Lock the critical section of the rate limiter and return the associated guard. - fn lock(&self) -> Guard<'_> { - Guard { - critical: self.critical.lock(), - } - } - /// Load the current state. - fn load(&self) -> State<'_> { - State::decode(self.state.load(Ordering::Acquire), self) - } - /// Take the current state, leaving the core state intact. - fn take(&self) -> State<'_> { - State::decode(self.state.swap(0, Ordering::Acquire), self) - } - /// Try to use fast path. - fn try_fast_path(&self, permits: usize) -> bool { - if permits == 0 { - return true; - } - if self.fair { - return false; - } - self.load().try_fast_path(permits) - } - /// Calculate refill amount. Returning a tuple of how much to fill and remaining - /// duration to sleep until the next refill time if appropriate. - /// - /// The maximum number of additional tokens this method will ever return is - /// limited to [`MAX_BALANCE`] to ensure that addition with an existing - /// balance will never overflow. - fn calculate_drain(&self, deadline: Instant, now: Instant) -> Option<(usize, Instant)> { - if now < deadline { - return None; - } - // Time elapsed in milliseconds since the last deadline. - let millis = self.interval.as_millis(); - let since = now.saturating_duration_since(deadline).as_millis(); - let periods = usize::try_from(since / millis + 1).unwrap_or(usize::MAX); - let tokens = periods - .checked_mul(self.refill) - .unwrap_or(MAX_BALANCE) - .min(MAX_BALANCE); - let rem = u64::try_from(since % millis).unwrap_or(u64::MAX); - - // Calculated time remaining until the next deadline. - let next = millis as u64 * periods as u64 - rem; - let deadline = deadline.checked_add(time::Duration::from_millis(next)).unwrap(); - Some((tokens, deadline)) - } -} -impl fmt::Debug for RateLimiter { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("RateLimiter") - .field("refill", &self.refill) - .field("interval", &self.interval) - .field("max", &self.max) - .field("fair", &self.fair) - .finish_non_exhaustive() - } -} -/// Refill the wait queue with the given number of tokens. -fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) { - let mut bump = 0; - // SAFETY: we're holding the lock guard to all the waiters so we can be - // sure that we have exclusive access to the wait queue. - unsafe { - while state.balance > 0 { - let mut node = match critical.waiters.pop_back() { - Some(node) => node, - None => break, - }; - let n = node.as_mut(); - n.fill(&mut state.balance); - if !n.is_completed() { - critical.waiters.push_back(node); - break; - } - n.complete.store(true, Ordering::Release); - if let Some(waker) = n.waker.take() { - waker.wake(); - } - bump += 1; - if bump == BUMP_LIMIT { - Guard::bump(critical); - bump = 0; - } - } - } -} -// SAFETY: All the internals of acquire is thread safe and correctly -// synchronized. The embedded waiter queue doesn't have anything inherently -// unsafe in it. -unsafe impl Send for RateLimiter {} -unsafe impl Sync for RateLimiter {} -/// A builder for a [`RateLimiter`]. -pub struct Builder { - /// The max number of tokens. - max: Option, - /// The initial count of tokens. - initial: usize, - /// Tokens to add every `per` duration. - refill: usize, - /// Interval to add tokens in milliseconds. - interval: Duration, - /// If the rate limiter is fair or not. - fair: bool, -} -impl Builder { - /// Configure the max number of tokens to use. - /// - /// If unspecified, this will default to be 10 times the [`refill`] or the - /// [`initial`] value, whichever is largest. - /// - /// The maximum supported balance is limited to [`isize::MAX`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .max(10_000) - /// .build(); - /// ``` - /// - /// [`refill`]: Builder::refill - /// [`initial`]: Builder::initial - pub fn max(&mut self, max: usize) -> &mut Self { - self.max = Some(max); - self - } - /// Configure the initial number of tokens to configure. The default value - /// is `0`. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .initial(10) - /// .build(); - /// ``` - pub fn initial(&mut self, initial: usize) -> &mut Self { - self.initial = initial; - self - } - /// Configure the time duration between which we add [`refill`] number to - /// the bucket rate limiter. - /// - /// This is 100ms by default. - /// - /// # Panics - /// - /// This panics if the provided interval does not fit within the millisecond - /// bounds of a [usize] or is zero. - /// - /// ```should_panic - /// use leaky_bucket::RateLimiter; - /// use tokio::time::Duration; - /// - /// let limiter = RateLimiter::builder() - /// .interval(Duration::from_secs(u64::MAX)) - /// .build(); - /// ``` - /// - /// ```should_panic - /// use leaky_bucket::RateLimiter; - /// use tokio::time::Duration; - /// - /// let limiter = RateLimiter::builder() - /// .interval(Duration::from_millis(0)) - /// .build(); - /// ``` - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use tokio::time::Duration; - /// - /// let limiter = RateLimiter::builder() - /// .interval(Duration::from_millis(100)) - /// .build(); - /// ``` - /// - /// [`refill`]: Builder::refill - pub fn interval(&mut self, interval: Duration) -> &mut Self { - assert! { - interval.as_millis() != 0, - "interval must be non-zero", - }; - assert! { - u64::try_from(interval.as_millis()).is_ok(), - "interval must fit within a 64-bit integer" - }; - self.interval = interval; - self - } - /// The number of tokens to add at each [`interval`] interval. The default - /// value is `1`. - /// - /// # Panics - /// - /// Panics if a refill amount of `0` is specified. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .refill(100) - /// .build(); - /// ``` - /// - /// [`interval`]: Builder::interval - pub fn refill(&mut self, refill: usize) -> &mut Self { - assert!(refill > 0, "refill amount cannot be zero"); - self.refill = refill; - self - } - /// Configure the rate limiter to be fair. - /// - /// Fairness is enabled by deafult. - /// - /// Fairness ensures that tasks make progress in the order that they acquire - /// even when the rate limiter is under contention. An unfair scheduler - /// might have a higher total throughput. - /// - /// Fair scheduling also affects the behavior of - /// [`RateLimiter::try_acquire`] which will return `false` if there are any - /// pending tasks since they should be given priority. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// - /// let limiter = RateLimiter::builder() - /// .refill(100) - /// .fair(false) - /// .build(); - /// ``` - pub fn fair(&mut self, fair: bool) -> &mut Self { - self.fair = fair; - self - } - /// Construct a new [`RateLimiter`]. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use tokio::time::Duration; - /// - /// let limiter = RateLimiter::builder() - /// .refill(100) - /// .interval(Duration::from_millis(200)) - /// .max(10_000) - /// .build(); - /// ``` - pub fn build(&self) -> RateLimiter { - let deadline = Instant::now() + self.interval; - let initial = self.initial.min(MAX_BALANCE); - let refill = self.refill.min(MAX_BALANCE); - let max = match self.max { - Some(max) => max.min(MAX_BALANCE), - None => refill - .max(initial) - .saturating_mul(DEFAULT_REFILL_MAX_FACTOR) - .min(MAX_BALANCE), - }; - let initial = initial.min(max); - RateLimiter { - refill, - interval: self.interval, - max, - fair: self.fair, - state: AtomicUsize::new(initial << 1 | 1), - critical: Mutex::new(Critical { - waiters: LinkedList::new(), - deadline, - }), - } - } -} -/// Construct a new builder with default options. -/// -/// # Examples -/// -/// ``` -/// use leaky_bucket::Builder; -/// -/// let limiter = Builder::default().build(); -/// ``` -impl Default for Builder { - fn default() -> Self { - Self { - max: None, - initial: 0, - refill: 1, - interval: Duration::from_millis(100), - fair: true, - } - } -} -/// The state of an acquire operation. -#[derive(Debug, Clone, Copy)] -enum AcquireFutState { - /// Initial unconfigured state. - Initial, - /// The acquire is waiting to be released by the core. - Waiting, - /// The operation is completed. - Complete, - /// The task is currently the core. - Core {}, -} -/// Inner state and methods of the acquire. -#[repr(transparent)] -struct AcquireFutInner { - /// Aliased task state. - node: UnsafeCell>, -} -impl AcquireFutInner { - const fn new() -> AcquireFutInner { - AcquireFutInner { - node: UnsafeCell::new(Node::new(Task::new())), - } - } - /// Access the completion flag. - pub fn complete(&self) -> &AtomicBool { - // SAFETY: This is always safe to access since it's atomic. - unsafe { &*ptr::addr_of!((*self.node.get()).complete) } - } - /// Get the underlying task mutably. - /// - /// We prove that the caller does indeed have mutable access to the node by - /// passing in a mutable reference to the critical section. - #[inline] - pub fn get_task<'crit, C>( - self: Pin<&'crit mut Self>, - critical: &'crit mut C, - ) -> (&'crit mut C, &'crit mut Node) - where - C: IsCritical, - { - // SAFETY: Caller has exclusive access to the critical section, since - // it's passed in as a mutable argument. We can also ensure that none of - // the borrows outlive the provided closure. - unsafe { (critical, &mut *self.node.get()) } - } - /// Update the waiting state for this acquisition task. This might require - /// that we update the associated waker. - fn update(self: Pin<&mut Self>, critical: &mut Guard<'_>, waker: &Waker) { - let (critical, task) = self.get_task(critical); - if !task.is_linked() { - critical.push_task_front(task); - } - let new_waker = match task.waker { - None => true, - Some(ref w) => !w.will_wake(waker), - }; - if new_waker { - task.waker = Some(waker.clone()); - } - } - /// Ensure that the current core task is correctly linked up if needed. - fn link_core(self: Pin<&mut Self>, critical: &mut Critical, lim: &RateLimiter) { - let (critical, task) = self.get_task(critical); - match (lim.fair, task.is_linked()) { - (true, false) => { - // Fair scheduling needs to ensure that the core is part of the wait - // queue, and will be woken up in-order with other tasks. - critical.push_task(task); - } - (false, true) => { - // Unfair scheduling will not wake the core in order, so - // don't bother having it linked. - critical.remove_task(task); - } - _ => {} - } - } - /// Release any remaining tokens which are associated with this particular task. - fn release_remaining( - self: Pin<&mut Self>, - critical: &mut Guard<'_>, - state: &mut State<'_>, - permits: usize, - ) { - let (critical, task) = self.get_task(critical); - if task.is_linked() { - critical.remove_task(task); - } - // Hand back permits which we've acquired so far. - let release = permits.saturating_sub(task.remaining); - state.add_tokens(critical, release, |_, _| ()); - } - /// Drain the given number of tokens through the core. Returns `true` if the - /// core has been completed. - fn drain_core( - self: Pin<&mut Self>, - critical: &mut Guard<'_>, - state: &mut State<'_>, - tokens: usize, - ) -> bool { - let completed = state.add_tokens(critical, tokens, |critical, state| { - let (_, task) = self.get_task(critical); - // If the limiter is not fair, we need to in addition to draining - // remaining tokens from linked nodes, drain it from ourselves. We - // fill the current holder of the core last (self). To ensure that - // it stays around for as long as possible. - if !state.lim.fair { - task.fill(&mut state.balance); - } - task.is_completed() - }); - if completed { - // Everything was drained, including the current core (if - // appropriate). So we can release it now. - critical.release(state); - } - completed - } - /// Assume the current core and calculate how long we must sleep for in - /// order to do it. - /// - /// # Safety - /// - /// This might link the current task into the task queue, so the caller must - /// ensure that it is pinned. - fn assume_core( - mut self: Pin<&mut Self>, - critical: &mut Guard<'_>, - state: &mut State<'_>, - now: Instant, - ) -> bool { - self.as_mut().link_core(critical, state.lim); - let (tokens, deadline) = match state.lim.calculate_drain(critical.deadline, now) { - Some(tokens) => tokens, - None => return true, - }; - // It is appropriate to update the deadline. - critical.deadline = deadline; - !self.drain_core(critical, state, tokens) - } -} -pin_project! { - /// The future associated with acquiring permits from a rate limiter using - /// [`RateLimiter::acquire`]. - #[project(!Unpin)] - pub struct Acquire<'a> { - #[pin] - inner: AcquireFut>, - } -} -impl Acquire<'_> { - /// Test if this acquire task is currently coordinating the rate limiter. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use std::future::Future; - /// use std::sync::Arc; - /// use std::task::Context; - /// - /// struct Waker; - /// # impl std::task::Wake for Waker { fn wake(self: Arc) { } } - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = RateLimiter::builder().build(); - /// - /// let waker = Arc::new(Waker).into(); - /// let mut cx = Context::from_waker(&waker); - /// - /// let a1 = limiter.acquire(1); - /// tokio::pin!(a1); - /// - /// assert!(!a1.is_core()); - /// assert!(a1.as_mut().poll(&mut cx).is_pending()); - /// assert!(a1.is_core()); - /// - /// a1.as_mut().await; - /// - /// // After completion this is no longer a core. - /// assert!(!a1.is_core()); - /// # } - /// ``` - pub fn is_core(&self) -> bool { - self.inner.is_core() - } -} -impl Future for Acquire<'_> { - type Output = (); - #[inline] - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - self.project().inner.poll(cx) - } -} -pin_project! { - /// The future associated with acquiring permits from a rate limiter using - /// [`RateLimiter::acquire_owned`]. - #[project(!Unpin)] - pub struct AcquireOwned { - #[pin] - inner: AcquireFut>, - } -} -impl AcquireOwned { - /// Test if this acquire task is currently coordinating the rate limiter. - /// - /// # Examples - /// - /// ``` - /// use leaky_bucket::RateLimiter; - /// use std::future::Future; - /// use std::sync::Arc; - /// use std::task::Context; - /// - /// struct Waker; - /// # impl std::task::Wake for Waker { fn wake(self: Arc) { } } - /// - /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() { - /// let limiter = Arc::new(RateLimiter::builder().build()); - /// - /// let waker = Arc::new(Waker).into(); - /// let mut cx = Context::from_waker(&waker); - /// - /// let a1 = limiter.acquire_owned(1); - /// tokio::pin!(a1); - /// - /// assert!(!a1.is_core()); - /// assert!(a1.as_mut().poll(&mut cx).is_pending()); - /// assert!(a1.is_core()); - /// - /// a1.as_mut().await; - /// - /// // After completion this is no longer a core. - /// assert!(!a1.is_core()); - /// # } - /// ``` - pub fn is_core(&self) -> bool { - self.inner.is_core() - } -} -impl Future for AcquireOwned { - type Output = (); - #[inline] - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - self.project().inner.poll(cx) - } -} -pin_project! { - #[project(!Unpin)] - #[project = AcquireFutProj] - struct AcquireFut - where - T: Deref, - { - lim: T, - permits: usize, - state: AcquireFutState, - #[pin] - sleep: Option, - #[pin] - inner: AcquireFutInner, - } - impl PinnedDrop for AcquireFut - where - T: Deref, - { - fn drop(this: Pin<&mut Self>) { - let AcquireFutProj { lim, permits, state, inner, .. } = this.project(); - let is_core = match *state { - AcquireFutState::Waiting => false, - AcquireFutState::Core { .. } => true, - _ => return, - }; - let mut critical = lim.lock(); - let mut s = lim.take(); - inner.release_remaining(&mut critical, &mut s, *permits); - if is_core { - critical.release(&mut s); - } - *state = AcquireFutState::Complete; - } - } -} -impl AcquireFut -where - T: Deref, -{ - #[inline] - const fn new(lim: T, permits: usize) -> Self { - Self { - lim, - permits, - state: AcquireFutState::Initial, - sleep: None, - inner: AcquireFutInner::new(), - } - } - fn is_core(&self) -> bool { - matches!(&self.state, AcquireFutState::Core { .. }) - } -} -// SAFETY: All the internals of acquire is thread safe and correctly -// synchronized. The embedded waiter queue doesn't have anything inherently -// unsafe in it. -unsafe impl Send for AcquireFut where T: Send + Deref {} -unsafe impl Sync for AcquireFut where T: Sync + Deref {} -impl Future for AcquireFut -where - T: Deref, -{ - type Output = (); - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let AcquireFutProj { - lim, - permits, - state, - mut sleep, - inner: mut internal, - .. - } = self.project(); - // Hold onto the critical lock for core operations, but only acquire it - // when strictly necessary. - let mut critical; - // Shared state. - // - // Once we are holding onto the critical lock, we take the entire state - // to ensure that any fast-past negotiators do not observe any available - // permits while potential core work is ongoing. - let mut s; - // Hold onto any call to `Instant::now` which we might perform, so we - // don't have to get the current time multiple times. - let outer_now; - match *state { - AcquireFutState::Complete => { - return Poll::Ready(()); - } - AcquireFutState::Initial => { - // If the rate limiter is not fair, try to oppurtunistically - // just acquire a permit through the known atomic state. - // - // This is known as the fast path, but requires acquire to raise - // against other tasks when storing the state back. - if lim.try_fast_path(*permits) { - *state = AcquireFutState::Complete; - return Poll::Ready(()); - } - critical = lim.lock(); - s = lim.take(); - let now = Instant::now(); - // If we've hit a deadline, calculate the number of tokens - // to drain and perform it in line here. This is necessary - // because the core isn't aware of how long we sleep between - // each acquire, so we need to perform some of the drain - // work here in order to avoid acruing a debt that needs to - // be filled later in. - // - // If we didn't do this, and the process slept for a long - // time, the next time a core is acquired it would be very - // far removed from the expected deadline and has no idea - // when permits were acquired, so it would over-eagerly - // release a lot of acquires and accumulate permits. - // - // This is tested for in the `test_idle` suite of tests. - let tokens = - if let Some((tokens, deadline)) = lim.calculate_drain(critical.deadline, now) { - // We pre-emptively update the deadline of the core - // since it might bump, and we don't want other - // processes to observe that the deadline has been - // reached. - critical.deadline = deadline; - tokens - } else { - 0 - }; - let completed = s.add_tokens(&mut critical, tokens, |critical, s| { - let (_, task) = internal.as_mut().get_task(critical); - task.remaining = *permits; - task.fill(&mut s.balance); - task.is_completed() - }); - if completed { - *state = AcquireFutState::Complete; - return Poll::Ready(()); - } - // Try to take over as core. If we're unsuccessful we just - // ensure that we're linked into the wait queue. - if !mem::take(&mut s.available) { - internal.as_mut().update(&mut critical, cx.waker()); - *state = AcquireFutState::Waiting; - return Poll::Pending; - } - // SAFETY: This is done in a pinned section, so we know that - // the linked section stays alive for the duration of this - // future due to pinning guarantees. - internal.as_mut().link_core(&mut critical, lim); - Guard::bump(&mut critical); - *state = AcquireFutState::Core {}; - outer_now = Some(now); - } - AcquireFutState::Waiting => { - // If we are complete, then return as ready. - // - // This field is atomic, so we can safely read it under shared - // access and do not require a lock. - if internal.complete().load(Ordering::Acquire) { - *state = AcquireFutState::Complete; - return Poll::Ready(()); - } - // Note: we need to operate under this lock to ensure that - // the core acquired here (or elsewhere) observes that the - // current task has been linked up. - critical = lim.lock(); - s = lim.take(); - // Try to take over as core. If we're unsuccessful we - // just ensure that we're linked into the wait queue. - if !mem::take(&mut s.available) { - internal.update(&mut critical, cx.waker()); - return Poll::Pending; - } - let now = Instant::now(); - // This is done in a pinned section, so we know that the linked - // section stays alive for the duration of this future due to - // pinning guarantees. - if !internal.as_mut().assume_core(&mut critical, &mut s, now) { - // Marks as completed. - *state = AcquireFutState::Complete; - return Poll::Ready(()); - } - Guard::bump(&mut critical); - *state = AcquireFutState::Core {}; - outer_now = Some(now); - } - AcquireFutState::Core {} => { - critical = lim.lock(); - s = lim.take(); - outer_now = None; - } - } - let mut sleep = match sleep.as_mut().as_pin_mut() { - Some(mut sleep) => { - if sleep.deadline() != critical.deadline { - sleep.as_mut().reset(critical.deadline); - } - sleep - } - None => { - sleep.set(Some(time::sleep_until(critical.deadline))); - sleep.as_mut().as_pin_mut().unwrap() - } - }; - if sleep.as_mut().poll(cx).is_pending() { - return Poll::Pending; - } - critical.deadline = outer_now.unwrap_or(sleep.deadline()) + lim.interval; - if internal.drain_core(&mut critical, &mut s, lim.refill) { - *state = AcquireFutState::Complete; - return Poll::Ready(()); - } - cx.waker().wake_by_ref(); - Poll::Pending - } -} -pub struct Node { - /// The next node. - next: Option>>, - /// The previous node. - prev: Option>>, - /// If we are linked or not. - linked: bool, - /// The value inside of the node. - value: T, - /// Avoids noalias heuristics from kicking in on references to a `Node` - /// struct. - _pin: marker::PhantomPinned, -} -impl Node { - /// Construct a new unlinked node. - const fn new(value: T) -> Self { - Self { - next: None, - prev: None, - linked: false, - value, - _pin: marker::PhantomPinned, - } - } - #[inline(always)] - fn is_linked(&self) -> bool { - self.linked - } - /// Set the next node. - #[inline(always)] - unsafe fn set_next(&mut self, node: Option>) { - ptr::addr_of_mut!(self.next).write(node); - } - /// Take the next node. - #[inline(always)] - unsafe fn take_next(&mut self) -> Option> { - ptr::addr_of_mut!(self.next).replace(None) - } - /// Set the previous node. - #[inline(always)] - unsafe fn set_prev(&mut self, node: Option>) { - ptr::addr_of_mut!(self.prev).write(node); - } - /// Take the previous node. - #[inline(always)] - unsafe fn take_prev(&mut self) -> Option> { - ptr::addr_of_mut!(self.prev).replace(None) - } -} -impl ops::Deref for Node { - type Target = T; - fn deref(&self) -> &Self::Target { - &self.value - } -} -impl ops::DerefMut for Node -where - T: Unpin, -{ - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.value - } -} -pub struct LinkedList { - head: Option>>, - tail: Option>>, -} -impl LinkedList { - /// Construct a new empty list. - const fn new() -> Self { - Self { - head: None, - tail: None, - } - } - unsafe fn push_front(&mut self, mut node: ptr::NonNull>) { - debug_assert!(node.as_ref().next.is_none()); - debug_assert!(node.as_ref().prev.is_none()); - debug_assert!(!node.as_ref().linked); - if let Some(mut head) = self.head.take() { - node.as_mut().set_next(Some(head)); - head.as_mut().set_prev(Some(node)); - self.head = Some(node); - } else { - self.head = Some(node); - self.tail = Some(node); - } - node.as_mut().linked = true; - } - unsafe fn push_back(&mut self, mut node: ptr::NonNull>) { - debug_assert!(node.as_ref().next.is_none()); - debug_assert!(node.as_ref().prev.is_none()); - debug_assert!(!node.as_ref().linked); - if let Some(mut tail) = self.tail.take() { - node.as_mut().set_prev(Some(tail)); - tail.as_mut().set_next(Some(node)); - self.tail = Some(node); - } else { - self.head = Some(node); - self.tail = Some(node); - } - node.as_mut().linked = true; - } - #[cfg(test)] - unsafe fn pop_front(&mut self) -> Option>> { - let mut head = self.head?; - debug_assert!(head.as_ref().linked); - if let Some(mut next) = head.as_mut().take_next() { - next.as_mut().set_prev(None); - self.head = Some(next); - } else { - debug_assert_eq!(self.tail, Some(head)); - self.head = None; - self.tail = None; - } - debug_assert!(head.as_ref().prev.is_none()); - debug_assert!(head.as_ref().next.is_none()); - head.as_mut().linked = false; - Some(head) - } - /// Pop the back element from the list. - unsafe fn pop_back(&mut self) -> Option>> { - let mut tail = self.tail?; - debug_assert!(tail.as_ref().linked); - if let Some(mut prev) = tail.as_mut().take_prev() { - prev.as_mut().set_next(None); - self.tail = Some(prev); - } else { - debug_assert_eq!(self.head, Some(tail)); - self.head = None; - self.tail = None; - } - debug_assert!(tail.as_ref().prev.is_none()); - debug_assert!(tail.as_ref().next.is_none()); - tail.as_mut().linked = false; - Some(tail) - } - unsafe fn remove(&mut self, mut node: ptr::NonNull>) { - debug_assert!(node.as_ref().linked); - let next = node.as_mut().take_next(); - let prev = node.as_mut().take_prev(); - if let Some(mut next) = next { - next.as_mut().set_prev(prev); - } else { - debug_assert_eq!(self.tail, Some(node)); - self.tail = prev; - } - if let Some(mut prev) = prev { - prev.as_mut().set_next(next); - } else { - debug_assert_eq!(self.head, Some(node)); - self.head = next; - } - node.as_mut().linked = false; - } - unsafe fn front(&mut self) -> Option>> { - self.head - } -} -impl fmt::Debug for LinkedList { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("LinkedList") - .field("head", &self.head) - .field("tail", &self.tail) - .finish() - } -} \ No newline at end of file From 7c6b3e5774074d7970bece0c38b7c2b5d8e3bb18 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Wed, 9 Jul 2025 10:44:47 +0800 Subject: [PATCH 06/10] add jitter --- src/bench.rs | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/src/bench.rs b/src/bench.rs index 77c74c6..a9f6283 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -10,7 +10,7 @@ use crate::auto_connection::{AutoConnection, ConnLimiter}; use crate::client::ClientConfig; use crate::command::Command; use crate::shared_context::SharedContext; -use governor::{Quota, DefaultDirectRateLimiter}; +use governor::{Quota, DefaultDirectRateLimiter, Jitter}; #[derive(Clone)] pub struct Case { @@ -47,6 +47,15 @@ async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limit break; } + // 先进行速率限制检查 + match qps_limiter.as_ref() { + Some(limiter) => { + limiter.until_ready_with_jitter(Jitter::up_to(std::time::Duration::from_micros(10000))).await; + // limiter.until_ready().await; + } + None => {} + } + // prepare pipeline let mut p = Vec::new(); for _ in 0..pipeline_cnt { @@ -62,14 +71,6 @@ async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limit for _ in 0..pipeline_cnt { context.histogram.record(duration); } - match qps_limiter.as_ref() { - Some(limiter) => { - for _ in 0..pipeline_cnt { - limiter.until_ready().await; - } - } - None => {} - } } }); } @@ -157,9 +158,15 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo // calc target qps let qps_limiter = Arc::new(if case.target > 0 { - Some( - DefaultDirectRateLimiter::direct(Quota::per_second(NonZeroU32::new(case.target as u32).unwrap())), - ) + let limiter = DefaultDirectRateLimiter::direct(Quota::per_second(NonZeroU32::new(case.target as u32).unwrap())); + + // 预先消耗令牌以接近从 0 开始的效果 + // 这会消耗一些初始可用的令牌 + while limiter.check().is_ok() { + // 继续消耗直到没有可用令牌 + } + + Some(limiter) } else { None }); From 4d9887af618c1b3f6d47f2b1ad39650e673c5496 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Wed, 9 Jul 2025 11:24:02 +0800 Subject: [PATCH 07/10] fix all! --- src/bench.rs | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/src/bench.rs b/src/bench.rs index a9f6283..f10b6ed 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -10,7 +10,7 @@ use crate::auto_connection::{AutoConnection, ConnLimiter}; use crate::client::ClientConfig; use crate::command::Command; use crate::shared_context::SharedContext; -use governor::{Quota, DefaultDirectRateLimiter, Jitter}; +use governor::{DefaultDirectRateLimiter, Jitter, Quota}; #[derive(Clone)] pub struct Case { @@ -50,8 +50,8 @@ async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limit // 先进行速率限制检查 match qps_limiter.as_ref() { Some(limiter) => { - limiter.until_ready_with_jitter(Jitter::up_to(std::time::Duration::from_micros(10000))).await; - // limiter.until_ready().await; + let j = Jitter::up_to(std::time::Duration::from_micros(1_000_000_u64 * conn_limiter.total_conn / case.target)); + limiter.until_ready_with_jitter(j).await; } None => {} } @@ -159,13 +159,7 @@ pub fn do_benchmark(client_config: ClientConfig, cores: Vec, case: Case, lo // calc target qps let qps_limiter = Arc::new(if case.target > 0 { let limiter = DefaultDirectRateLimiter::direct(Quota::per_second(NonZeroU32::new(case.target as u32).unwrap())); - - // 预先消耗令牌以接近从 0 开始的效果 - // 这会消耗一些初始可用的令牌 - while limiter.check().is_ok() { - // 继续消耗直到没有可用令牌 - } - + while limiter.check().is_ok() {} Some(limiter) } else { None From d05e36dae60d14a2825f742451de12d3bbb1c3c7 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Wed, 9 Jul 2025 11:29:13 +0800 Subject: [PATCH 08/10] remove max --- python/resp_benchmark/wrapper.py | 2 -- src/bench.rs | 1 - src/histogram.rs | 7 ++----- src/lib.rs | 1 - 4 files changed, 2 insertions(+), 9 deletions(-) diff --git a/python/resp_benchmark/wrapper.py b/python/resp_benchmark/wrapper.py index c5ccdcb..a4b1372 100644 --- a/python/resp_benchmark/wrapper.py +++ b/python/resp_benchmark/wrapper.py @@ -22,7 +22,6 @@ class Result: qps: float avg_latency_ms: float p99_latency_ms: float - max_latency_ms: float connections: int @@ -112,7 +111,6 @@ def bench( qps=ret.qps, avg_latency_ms=ret.avg_latency_ms, p99_latency_ms=ret.p99_latency_ms, - max_latency_ms=ret.max_latency_ms, connections=ret.connections ) diff --git a/src/bench.rs b/src/bench.rs index f10b6ed..a603a9a 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -137,7 +137,6 @@ fn wait_finish(case: &Case, mut auto_connection: AutoConnection, mut context: Sh }; result.avg_latency_ms = histogram.avg() as f64 / 1_000.0; result.p99_latency_ms = histogram.percentile(0.99) as f64 / 1_000.0; - result.max_latency_ms = histogram.max() as f64 / 1_000.0; result.connections = conn; }); return result; diff --git a/src/histogram.rs b/src/histogram.rs index 1bec81f..66b5530 100644 --- a/src/histogram.rs +++ b/src/histogram.rs @@ -103,9 +103,7 @@ impl Histogram { 0 } - pub fn max(&self) -> u64 { - return self.percentile(1.0); - } + fn humanize_us(latency_us: u64) -> String { match latency_us { @@ -128,9 +126,8 @@ impl Display for Histogram { } let avg = self.avg(); let p99 = self.percentile(0.99); - let max = self.max(); - write!(f, "cnt: {}, avg: {}, p99: {}, max: {}", cnt, Histogram::humanize_us(avg), Histogram::humanize_us(p99), Histogram::humanize_us(max)) + write!(f, "cnt: {}, avg: {}, p99: {}", cnt, Histogram::humanize_us(avg), Histogram::humanize_us(p99)) } } diff --git a/src/lib.rs b/src/lib.rs index 7fd1a92..7e95f82 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,6 @@ struct BenchmarkResult { #[pyo3(get, set)] pub qps: f64, #[pyo3(get, set)] pub avg_latency_ms: f64, #[pyo3(get, set)] pub p99_latency_ms: f64, - #[pyo3(get, set)] pub max_latency_ms: f64, #[pyo3(get, set)] pub connections: u64, } From 1e5b9690355cd11abe9d94619ebcca165fbf1538 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Wed, 9 Jul 2025 13:46:36 +0800 Subject: [PATCH 09/10] handle THROTTLED --- .cursor/rules/project.mdc | 6 ++++++ src/bench.rs | 8 +++++--- src/client.rs | 14 ++++++++++---- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/.cursor/rules/project.mdc b/.cursor/rules/project.mdc index 5580589..fa5042d 100644 --- a/.cursor/rules/project.mdc +++ b/.cursor/rules/project.mdc @@ -15,6 +15,12 @@ Redis 接口数据库的性能测试套件。 1. 不需要加太多异常捕获，让代码保持简洁。 +## 任务要求 + +1. 使用 maturin develop 来进行编译安装测试 +2. 使用 resp-benchmark cli 来进行测试，不要新建文件测试 + + ## Core Components ### Python Interface diff --git a/src/bench.rs b/src/bench.rs index a603a9a..5c0f9a3 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -66,10 +66,12 @@ async fn run_commands_on_single_thread(conn_limiter: Arc, qps_limit } } let instant = std::time::Instant::now(); - client.run_commands(p).await; + let should_count = client.run_commands(p).await; let duration = instant.elapsed().as_micros() as u64; - for _ in 0..pipeline_cnt { - context.histogram.record(duration); + if should_count { + for _ in 0..pipeline_cnt { + context.histogram.record(duration); + } } } }); diff --git a/src/client.rs b/src/client.rs index 7d6cc26..478e59a 100644 --- a/src/client.rs +++ b/src/client.rs @@ -84,16 +84,22 @@ impl Client { } - pub async fn run_commands(&mut self, cmds: Vec) { + pub async fn run_commands(&mut self, cmds: Vec) -> bool { let mut pipeline = redis::pipe(); for cmd in cmds { pipeline.add_command(cmd).ignore(); } match pipeline.query_async(&mut self.conn).await { - Ok(()) => {} + Ok(()) => true, Err(e) => { - eprintln!("Failed to execute pipeline: {:?}", e); - std::process::exit(1); + let error_msg = format!("{:?}", e); + if error_msg.contains("THROTTLED: insufficient capacity unit") { + // ignore throttled requests - don't count them in statistics + false + } else { + eprintln!("Failed to execute pipeline: {:?}", e); + std::process::exit(1); + } } } } From 76732d3601b81256004301fe6ae10a8ae1f01433 Mon Sep 17 00:00:00 2001 From: suxb201 Date: Thu, 10 Jul 2025 19:01:03 +0800 Subject: [PATCH 10/10] update README --- README.md | 386 ++++++++++++++++++++++++++++++++++++++++---------- README_en.md | 387 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 703 insertions(+), 70 deletions(-) create mode 100644 README_en.md diff --git a/README.md b/README.md index e8bac8a..6ef8cb8 100644 --- a/README.md +++ b/README.md @@ -5,123 +5,369 @@ [![PyPI - Downloads](https://img.shields.io/pypi/dw/resp-benchmark?color=%231ba784)](https://pypi.org/project/resp-benchmark/) [![License](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/tair-opensource/resp-benchmark/blob/main/LICENSE) -resp-benchmark is a benchmark tool for testing databases that support the RESP protocol, -such as [Redis](https://github.com/redis/redis), [Valkey](https://github.com/valkey-io/valkey), -and [Tair](https://www.alibabacloud.com/en/product/tair). It offers both a command-line interface and a Python library. +[English](README_en.md) | 中文 -## Installation +一个基于 Rust 构建的 RESP (Redis 序列化协议) 数据库基准测试工具，提供 Python 绑定。用于测试 Redis、Valkey、Tair 等 RESP 兼容数据库的性能。 + +## 主要功能 + +- 支持自定义命令模板和占位符 +- 多线程并发测试 +- 支持连接池和管道操作 +- 支持 Redis 集群模式 +- 提供 CLI 工具和 Python 库接口 +- 内置 QPS 限流控制 + +## 安装 + +需要 Python 3.8 或更高版本。 -Requires Python 3.8 or higher. ```bash pip install resp-benchmark ``` -## Usage +## 快速开始 -### Command-Line Tool +### 命令行使用 ```bash -resp-benchmark --help +# 基础基准测试 +resp-benchmark -s 10 "SET {key uniform 100000} {value 64}" + +# 加载数据然后测试 +resp-benchmark --load -n 1000000 "SET {key sequence 100000} {value 64}" +resp-benchmark -s 10 "GET {key uniform 100000}" + +# 自定义连接数和管道 +resp-benchmark -c 128 -P 10 -s 30 "SET {key uniform 1000000} {value 128}" ``` -### Python Library +### Python 库使用 ```python from resp_benchmark import Benchmark +# 初始化基准测试 bm = Benchmark(host="127.0.0.1", port=6379) -bm.flushall() -bm.load_data(command="SET {key sequence 10000000} {value 64}", count=1000000, connections=128) -result = bm.bench("GET {key uniform 10000000}", seconds=3, connections=16) -print(result.qps, result.avg_latency_ms, result.p99_latency_ms) + +# 加载测试数据 +bm.load_data( + command="SET {key sequence 1000000} {value 64}", + count=1000000, + connections=128 +) + +# 运行基准测试 +result = bm.bench( + command="GET {key uniform 1000000}", + seconds=30, + connections=64 +) + +print(f"QPS: {result.qps}") +print(f"平均延迟: {result.avg_latency_ms}ms") +print(f"P99 延迟: {result.p99_latency_ms}ms") +``` + +## 命令语法 + +resp-benchmark 使用强大的占位符系统来生成多样化和真实的测试数据： + +### 键占位符 + +- **`{key uniform N}`**: 从 0 到 N-1 的随机键 + - 示例: `{key uniform 100000}` → `key_0000099999` + +- **`{key sequence N}`**: 从 0 到 N-1 的顺序键（适用于加载数据） + - 示例: `{key sequence 100000}` → `key_0000000000`, `key_0000000001`, ... + +- **`{key zipfian N}`**: Zipfian 分布键（模拟真实世界的访问模式） + - 示例: `{key zipfian 100000}` → 遵循指数为 1.03 的 Zipfian 分布 + +### 值占位符 + +- **`{value N}`**: N 字节的随机字符串 + - 示例: `{value 64}` → `a8x9mK2p...` (64 字节) + +- **`{rand N}`**: 0 到 N-1 的随机数 + - 示例: `{rand 1000}` → `742` + +- **`{range N W}`**: 范围 N 内相差 W 的两个数字 + - 示例: `{range 100 10}` → `45 55` + +### 命令示例 + +```bash +# 字符串操作 +SET {key uniform 1000000} {value 64} +GET {key uniform 1000000} +INCR {key uniform 100000} + +# 列表操作 +LPUSH {key uniform 1000} {value 64} +LINDEX {key uniform 1000} {rand 100} + +# 集合操作 +SADD {key uniform 1000} {value 64} +SISMEMBER {key uniform 1000} {value 64} + +# 有序集合操作 +ZADD {key uniform 1000} {rand 1000} {value 64} +ZRANGEBYSCORE {key uniform 1000} {range 1000 100} + +# 哈希操作 +HSET {key uniform 1000} {key uniform 100} {value 64} +HGET {key uniform 1000} {key uniform 100} ``` -## Custom Commands +## 命令行选项 + +| 选项 | 描述 | 默认值 | +|------|------|--------| +| `-h` | 服务器主机名 | 127.0.0.1 | +| `-p` | 服务器端口 | 6379 | +| `-u` | 认证用户名 | "" | +| `-a` | 认证密码 | "" | +| `-c` | 连接数（0 为自动） | 0 | +| `-n` | 总请求数（0 为无限） | 0 | +| `-s` | 持续时间（秒）（0 为无限） | 0 | +| `-t` | 目标 QPS（0 为无限） | 0 | +| `-P` | 管道深度 | 1 | +| `--cores` | 使用的 CPU 核心（逗号分隔） | 全部 | +| `--cluster` | 启用集群模式 | false | +| `--load` | 仅加载数据，不进行基准测试 | false | + +## 高级特性 + +### 连接自动扩展 -resp-benchmark supports custom test commands using placeholder syntax like `SET {key uniform 10000000} {value 64}` which means the SET command will have a key uniformly distributed in the range -0-10000000 and a value of 64 bytes. +当指定 `-c 0` 时，resp-benchmark 根据系统资源和目标 QPS 自动确定最优连接数。 -Supported placeholders include: +### CPU 核心绑定 -- **`{key uniform N}`**: Generates a random number between `0` and `N-1`. For example, `{key uniform 100}` might generate `key_0000000099`. -- **`{key sequence N}`**: Sequentially generates from `0` to `N-1`, ensuring coverage during data loading. For example, `{key sequence 100}` generates `key_0000000000`, `key_0000000001`, etc. -- **`{key zipfian N}`**: Generates according to a Zipfian distribution (exponent 1.03), simulating real-world key distribution. -- **`{value N}`**: Generates a random string of length `N` bytes. For example, `{value 8}` might generate `92xsqdNg`. -- **`{rand N}`**: Generates a random number between `0` and `N-1`. For example, `{rand 100}` might generate `99`. -- **`{range N W}`**: Generates a pair of random numbers within the range `0` to `N-1`, with a difference of `W`, used for testing `*range*` commands. For example, `{range 100 10}` might generate - `89 99`. +将基准测试线程绑定到特定 CPU 核心以获得一致的性能： -## Best Practices +```bash +# 使用核心 0, 1, 2, 3 +resp-benchmark --cores 0,1,2,3 -s 10 "SET {key uniform 100000} {value 64}" +``` + +### 速率限制 + +控制请求速率进行渐进式负载测试： + +```bash +# 目标 10,000 QPS +resp-benchmark -t 10000 -s 30 "SET {key uniform 100000} {value 64}" +``` -Notes: -1. It is recommended to clear the data each time you test to avoid interference from existing data. -2. In actual tests, it is recommended to manually adjust the number of `connections`, such as setting it to 128, which can be achieved through `-c 128`. +### 管道操作 -### Benchmarking network +使用管道进行批量操作： -```shell -# Test PING command -resp-benchmark -s 10 "PING" -# Test ECHO command -resp-benchmark -s 10 "ECHO {value 64}" +```bash +# 每个连接管道 10 个请求 +resp-benchmark -P 10 -c 128 -s 30 "SET {key uniform 100000} {value 64}" ``` -### Benchmarking string +### 集群模式 -```shell -# Test SET command -resp-benchmark -s 10 "SET {key uniform 10000000} {value 64}" +使用自动槽位分布测试 Redis 集群： -# Test GET command -resp-benchmark --load -c 256 -P 10 -n 1000000 "SET {key sequence 10000000} {value 64}" -resp-benchmark -s 10 "GET {key uniform 10000000}" +```bash +resp-benchmark --cluster -h cluster-endpoint -p 7000 -s 30 "SET {key uniform 100000} {value 64}" ``` -### Benchmarking list +## 性能优化 + +### 最佳实践 -```shell -# Test LPUSH command -resp-benchmark -s 10 "LPUSH {key uniform 1000} {value 64}" +1. **预加载数据**: 使用 `--load` 在基准测试前填充测试数据 +2. **适当的连接数**: 从 `-c 128` 开始，根据结果调整 +3. **键分布**: 读取使用 `uniform`，写入使用 `sequence` +4. **明智使用管道**: 批量操作使用 `-P 10`，延迟测试使用 `-P 1` +5. **清洁状态**: 测试之间清除数据以避免干扰 -# Test LINDEX command -resp-benchmark --load -c 256 -P 10 -n 10000000 "LPUSH {key sequence 1000} {value 64}" -resp-benchmark -s 10 "LINDEX {key uniform 1000} {rand 10000}" +### 示例工作流程 + +```bash +# 1. 清除现有数据 +redis-cli FLUSHALL + +# 2. 加载测试数据 +resp-benchmark --load -c 256 -P 10 -n 1000000 "SET {key sequence 100000} {value 64}" + +# 3. 使用不同模式进行基准测试 +resp-benchmark -c 128 -s 30 "GET {key uniform 100000}" # 随机访问 +resp-benchmark -c 128 -s 30 "GET {key zipfian 100000}" # 真实访问模式 ``` -### Benchmarking set +## 完整示例 + +### 字符串操作 -```shell -# Test SADD command -resp-benchmark -s 10 "SADD {key uniform 1000} {value 64}" +```bash +# 基本 SET/GET +resp-benchmark --load -n 1000000 "SET {key sequence 100000} {value 64}" +resp-benchmark -s 10 "GET {key uniform 100000}" + +# 大值 +resp-benchmark -s 10 "SET {key uniform 10000} {value 1024}" -# Test SISMEMBER command -resp-benchmark --load -c 256 -P 10 -n 10007000 "SADD {key sequence 1000} {key sequence 10007}" -resp-benchmark -s 10 "SISMEMBER {key uniform 1000} {key uniform 10007}" +# 递增操作 +resp-benchmark -s 10 "INCR {key uniform 10000}" ``` -### Benchmarking zset +### 列表操作 -```shell -# Test ZADD command -resp-benchmark -s 10 "ZADD {key uniform 1000} {rand 70000} {key uniform 10007}" +```bash +# 构建列表 +resp-benchmark --load -n 1000000 "LPUSH {key sequence 1000} {value 64}" -# Benchmark ZSCORE & ZRANGEBYSCORE -resp-benchmark --load -P 10 -c 256 -n 10007000 "ZADD {key sequence 1000} {rand 70000} {key sequence 10007}" -resp-benchmark -s 10 "ZSCORE {key uniform 1000} {key uniform 10007}" -resp-benchmark -s 10 "ZRANGEBYSCORE {key uniform 1000} {range 70000 10}" +# 随机访问 +resp-benchmark -s 10 "LINDEX {key uniform 1000} {rand 1000}" + +# 范围操作 +resp-benchmark -s 10 "LRANGE {key uniform 1000} {range 1000 10}" ``` -### Benchmarking Lua Scripts +### 集合操作 -```shell +```bash +# 填充集合 +resp-benchmark --load -n 1000000 "SADD {key sequence 1000} {key sequence 1000}" + +# 测试成员关系 +resp-benchmark -s 10 "SISMEMBER {key uniform 1000} {key uniform 1000}" +``` + +### 有序集合操作 + +```bash +# 添加带分数的成员 +resp-benchmark --load -n 1000000 "ZADD {key sequence 1000} {rand 10000} {key sequence 1000}" + +# 分数查询 +resp-benchmark -s 10 "ZSCORE {key uniform 1000} {key uniform 1000}" + +# 范围查询 +resp-benchmark -s 10 "ZRANGEBYSCORE {key uniform 1000} {range 10000 100}" +``` + +### 哈希操作 + +```bash +# 填充哈希 +resp-benchmark --load -n 1000000 "HSET {key sequence 1000} {key sequence 100} {value 64}" + +# 字段访问 +resp-benchmark -s 10 "HGET {key uniform 1000} {key uniform 100}" +``` + +### Lua 脚本 + +```bash +# 加载脚本 redis-cli SCRIPT LOAD "return redis.call('SET', KEYS[1], ARGV[1])" + +# 基准测试脚本执行 resp-benchmark -s 10 "EVALSHA d8f2fad9f8e86a53d2a6ebd960b33c4972cacc37 1 {key uniform 100000} {value 64}" ``` -## Differences with redis-benchmark +## Python 库 API + +### 初始化 + +```python +from resp_benchmark import Benchmark + +# 基本连接 +bm = Benchmark(host="127.0.0.1", port=6379) + +# 带认证 +bm = Benchmark( + host="redis.example.com", + port=6379, + username="user", + password="pass" +) + +# 集群模式 +bm = Benchmark( + host="cluster-endpoint", + port=7000, + cluster=True +) + +# 自定义配置 +bm = Benchmark( + host="127.0.0.1", + port=6379, + cores="0,1,2,3", # 使用特定核心 + timeout=30 # 连接超时 +) +``` + +### 加载数据 + +```python +# 顺序加载（推荐） +bm.load_data( + command="SET {key sequence 1000000} {value 64}", + count=1000000, + connections=128, + pipeline=10 +) + +# 带速率限制 +bm.load_data( + command="SET {key sequence 1000000} {value 64}", + count=1000000, + connections=128, + target=50000 # 50k QPS +) +``` + +### 基准测试 + +```python +# 基于时间的基准测试 +result = bm.bench( + command="GET {key uniform 1000000}", + seconds=30, + connections=64 +) + +# 基于计数的基准测试 +result = bm.bench( + command="GET {key uniform 1000000}", + count=1000000, + connections=64 +) + +# 带管道 +result = bm.bench( + command="SET {key uniform 1000000} {value 64}", + seconds=30, + connections=64, + pipeline=10 +) +``` + +### 结果分析 + +```python +# 访问基准测试结果 +print(f"QPS: {result.qps:.2f}") +print(f"平均延迟: {result.avg_latency_ms:.2f}ms") +print(f"P99 延迟: {result.p99_latency_ms:.2f}ms") +print(f"使用的连接数: {result.connections}") +``` + +## 贡献 + +欢迎贡献！请随时提交拉取请求或开启议题。 -When testing Redis with resp-benchmark and redis-benchmark, you might get different results due to: +## 许可证 -1. redis-benchmark always uses the same value when testing the set command, which does not trigger DB persistence and replication. In contrast, resp-benchmark uses `{value 64}` to generate different data for each command. -2. redis-benchmark always uses the same primary key when testing list/set/zset/hash commands, while resp-benchmark generates different keys using placeholders like `{key uniform 10000000}`. -3. In cluster mode, redis-benchmark sends requests to each node, but all requests target the same slot on every node. +该项目基于 MIT 许可证 - 详见 [LICENSE](LICENSE) 文件。 \ No newline at end of file diff --git a/README_en.md b/README_en.md new file mode 100644 index 0000000..1b0b666 --- /dev/null +++ b/README_en.md @@ -0,0 +1,387 @@ +# resp-benchmark + +[![Python - Version](https://img.shields.io/badge/python-%3E%3D3.8-brightgreen)](https://www.python.org/doc/versions/) +[![PyPI - Version](https://img.shields.io/pypi/v/resp-benchmark?color=%231772b4)](https://pypi.org/project/resp-benchmark/) +[![PyPI - Downloads](https://img.shields.io/pypi/dw/resp-benchmark?color=%231ba784)](https://pypi.org/project/resp-benchmark/) +[![License](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/tair-opensource/resp-benchmark/blob/main/LICENSE) + +English | [中文](README.md) + +A high-performance benchmark tool for testing databases that support the RESP (Redis Serialization Protocol), built with Rust and Python bindings. Designed to provide accurate and realistic performance measurements for Redis, Valkey, Tair, and other RESP-compatible databases. + +## Features + +- **High Performance**: Built with Rust core for maximum throughput and minimal overhead +- **Flexible Commands**: Support for custom command templates with intelligent placeholders +- **Realistic Testing**: Generates varied data to simulate real-world usage patterns +- **Multi-threaded**: Leverages multiple CPU cores for maximum performance +- **Connection Management**: Intelligent connection pooling and auto-scaling +- **Rate Limiting**: Built-in QPS throttling for controlled testing +- **Pipeline Support**: Configurable pipeline depth for bulk operations +- **Cluster Support**: Native Redis cluster mode support +- **Python Integration**: Use as both CLI tool and Python library +- **Comprehensive Metrics**: Detailed latency histograms and performance statistics + +## Installation + +Requires Python 3.8 or higher. + +```bash +pip install resp-benchmark +``` + +## Quick Start + +### Command Line Usage + +```bash +# Basic benchmark +resp-benchmark -s 10 "SET {key uniform 100000} {value 64}" + +# Load data then benchmark +resp-benchmark --load -n 1000000 "SET {key sequence 100000} {value 64}" +resp-benchmark -s 10 "GET {key uniform 100000}" + +# With custom connections and pipeline +resp-benchmark -c 128 -P 10 -s 30 "SET {key uniform 1000000} {value 128}" +``` + +### Python Library Usage + +```python +from resp_benchmark import Benchmark + +# Initialize benchmark +bm = Benchmark(host="127.0.0.1", port=6379) + +# Load test data +bm.load_data( + command="SET {key sequence 1000000} {value 64}", + count=1000000, + connections=128 +) + +# Run benchmark +result = bm.bench( + command="GET {key uniform 1000000}", + seconds=30, + connections=64 +) + +print(f"QPS: {result.qps}") +print(f"Avg Latency: {result.avg_latency_ms}ms") +print(f"P99 Latency: {result.p99_latency_ms}ms") +``` + +## Command Syntax + +resp-benchmark uses a powerful placeholder system to generate varied and realistic test data: + +### Key Placeholders + +- **`{key uniform N}`**: Random key from 0 to N-1 + - Example: `{key uniform 100000}` → `key_0000099999` + +- **`{key sequence N}`**: Sequential keys from 0 to N-1 (ideal for loading) + - Example: `{key sequence 100000}` → `key_0000000000`, `key_0000000001`, ... + +- **`{key zipfian N}`**: Zipfian distribution keys (simulates real-world access patterns) + - Example: `{key zipfian 100000}` → follows Zipfian distribution with exponent 1.03 + +### Value Placeholders + +- **`{value N}`**: Random string of N bytes + - Example: `{value 64}` → `a8x9mK2p...` (64 bytes) + +- **`{rand N}`**: Random number from 0 to N-1 + - Example: `{rand 1000}` → `742` + +- **`{range N W}`**: Two numbers within range N with difference W + - Example: `{range 100 10}` → `45 55` + +### Example Commands + +```bash +# String operations +SET {key uniform 1000000} {value 64} +GET {key uniform 1000000} +INCR {key uniform 100000} + +# List operations +LPUSH {key uniform 1000} {value 64} +LINDEX {key uniform 1000} {rand 100} + +# Set operations +SADD {key uniform 1000} {value 64} +SISMEMBER {key uniform 1000} {value 64} + +# Sorted Set operations +ZADD {key uniform 1000} {rand 1000} {value 64} +ZRANGEBYSCORE {key uniform 1000} {range 1000 100} + +# Hash operations +HSET {key uniform 1000} {key uniform 100} {value 64} +HGET {key uniform 1000} {key uniform 100} +``` + +## Command Line Options + +| Option | Description | Default | +|--------|-------------|---------| +| `-h` | Server hostname | 127.0.0.1 | +| `-p` | Server port | 6379 | +| `-u` | Username for authentication | "" | +| `-a` | Password for authentication | "" | +| `-c` | Number of connections (0 for auto) | 0 | +| `-n` | Total number of requests (0 for unlimited) | 0 | +| `-s` | Duration in seconds (0 for unlimited) | 0 | +| `-t` | Target QPS (0 for unlimited) | 0 | +| `-P` | Pipeline depth | 1 | +| `--cores` | CPU cores to use (comma-separated) | all | +| `--cluster` | Enable cluster mode | false | +| `--load` | Load data only, no benchmark | false | + +## Advanced Features + +### Connection Auto-scaling + +When `-c 0` is specified, resp-benchmark automatically determines the optimal number of connections based on system resources and target QPS. + +### CPU Core Affinity + +Bind benchmark threads to specific CPU cores for consistent performance: + +```bash +# Use cores 0, 1, 2, 3 +resp-benchmark --cores 0,1,2,3 -s 10 "SET {key uniform 100000} {value 64}" +``` + +### Rate Limiting + +Control the request rate for gradual load testing: + +```bash +# Target 10,000 QPS +resp-benchmark -t 10000 -s 30 "SET {key uniform 100000} {value 64}" +``` + +### Pipeline Operations + +Use pipelining for bulk operations: + +```bash +# Pipeline 10 requests per connection +resp-benchmark -P 10 -c 128 -s 30 "SET {key uniform 100000} {value 64}" +``` + +### Cluster Mode + +Test Redis clusters with automatic slot distribution: + +```bash +resp-benchmark --cluster -h cluster-endpoint -p 7000 -s 30 "SET {key uniform 100000} {value 64}" +``` + +## Performance Optimization + +### Best Practices + +1. **Pre-load Data**: Use `--load` to populate test data before benchmarking +2. **Appropriate Connections**: Start with `-c 128` and adjust based on results +3. **Key Distribution**: Use `uniform` for reads, `sequence` for writes +4. **Pipeline Wisely**: Use `-P 10` for bulk operations, `-P 1` for latency testing +5. **Clean State**: Clear data between tests to avoid interference + +### Example Workflow + +```bash +# 1. Clear existing data +redis-cli FLUSHALL + +# 2. Load test data +resp-benchmark --load -c 256 -P 10 -n 1000000 "SET {key sequence 100000} {value 64}" + +# 3. Benchmark with different patterns +resp-benchmark -c 128 -s 30 "GET {key uniform 100000}" # Random access +resp-benchmark -c 128 -s 30 "GET {key zipfian 100000}" # Realistic access +``` + +## Comprehensive Examples + +### String Operations + +```bash +# Basic SET/GET +resp-benchmark --load -n 1000000 "SET {key sequence 100000} {value 64}" +resp-benchmark -s 10 "GET {key uniform 100000}" + +# Large values +resp-benchmark -s 10 "SET {key uniform 10000} {value 1024}" + +# Increment operations +resp-benchmark -s 10 "INCR {key uniform 10000}" +``` + +### List Operations + +```bash +# Build lists +resp-benchmark --load -n 1000000 "LPUSH {key sequence 1000} {value 64}" + +# Random access +resp-benchmark -s 10 "LINDEX {key uniform 1000} {rand 1000}" + +# Range operations +resp-benchmark -s 10 "LRANGE {key uniform 1000} {range 1000 10}" +``` + +### Set Operations + +```bash +# Populate sets +resp-benchmark --load -n 1000000 "SADD {key sequence 1000} {key sequence 1000}" + +# Test membership +resp-benchmark -s 10 "SISMEMBER {key uniform 1000} {key uniform 1000}" +``` + +### Sorted Set Operations + +```bash +# Add scored members +resp-benchmark --load -n 1000000 "ZADD {key sequence 1000} {rand 10000} {key sequence 1000}" + +# Score queries +resp-benchmark -s 10 "ZSCORE {key uniform 1000} {key uniform 1000}" + +# Range queries +resp-benchmark -s 10 "ZRANGEBYSCORE {key uniform 1000} {range 10000 100}" +``` + +### Hash Operations + +```bash +# Populate hashes +resp-benchmark --load -n 1000000 "HSET {key sequence 1000} {key sequence 100} {value 64}" + +# Field access +resp-benchmark -s 10 "HGET {key uniform 1000} {key uniform 100}" +``` + +### Lua Scripts + +```bash +# Load script +redis-cli SCRIPT LOAD "return redis.call('SET', KEYS[1], ARGV[1])" + +# Benchmark script execution +resp-benchmark -s 10 "EVALSHA d8f2fad9f8e86a53d2a6ebd960b33c4972cacc37 1 {key uniform 100000} {value 64}" +``` + +## Python Library API + +### Initialization + +```python +from resp_benchmark import Benchmark + +# Basic connection +bm = Benchmark(host="127.0.0.1", port=6379) + +# With authentication +bm = Benchmark( + host="redis.example.com", + port=6379, + username="user", + password="pass" +) + +# Cluster mode +bm = Benchmark( + host="cluster-endpoint", + port=7000, + cluster=True +) + +# Custom configuration +bm = Benchmark( + host="127.0.0.1", + port=6379, + cores="0,1,2,3", # Use specific cores + timeout=30 # Connection timeout +) +``` + +### Loading Data + +```python +# Sequential loading (recommended) +bm.load_data( + command="SET {key sequence 1000000} {value 64}", + count=1000000, + connections=128, + pipeline=10 +) + +# With rate limiting +bm.load_data( + command="SET {key sequence 1000000} {value 64}", + count=1000000, + connections=128, + target=50000 # 50k QPS +) +``` + +### Benchmarking + +```python +# Time-based benchmark +result = bm.bench( + command="GET {key uniform 1000000}", + seconds=30, + connections=64 +) + +# Count-based benchmark +result = bm.bench( + command="GET {key uniform 1000000}", + count=1000000, + connections=64 +) + +# With pipeline +result = bm.bench( + command="SET {key uniform 1000000} {value 64}", + seconds=30, + connections=64, + pipeline=10 +) +``` + +### Result Analysis + +```python +# Access benchmark results +print(f"QPS: {result.qps:.2f}") +print(f"Average Latency: {result.avg_latency_ms:.2f}ms") +print(f"P99 Latency: {result.p99_latency_ms:.2f}ms") +print(f"Connections Used: {result.connections}") +``` + +## Differences from redis-benchmark + +resp-benchmark provides more realistic testing compared to redis-benchmark: + +1. **Varied Data**: Generates different values for each request, triggering realistic persistence and replication +2. **Distributed Keys**: Uses varied keys instead of single keys for collections +3. **Cluster Awareness**: Properly distributes requests across all slots in cluster mode +4. **Advanced Placeholders**: Supports realistic data distributions (uniform, zipfian, sequence) +5. **Better Metrics**: Detailed latency histograms and connection statistics + +## Contributing + +Contributions are welcome! Please feel free to submit pull requests or open issues. + +## License + +This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. \ No newline at end of file