From 81309d17fbe6d6db28665694de831bbcc58a48e1 Mon Sep 17 00:00:00 2001 From: Zhang Jingqiang Date: Thu, 31 Jul 2025 15:19:56 +0800 Subject: [PATCH] g3bench: allow more flexible rate limit config --- Cargo.lock | 13 ++-- g3bench/CHANGELOG | 3 + g3bench/Cargo.toml | 2 +- g3bench/src/opts.rs | 11 +--- g3bench/src/target/mod.rs | 2 +- lib/g3-clap/Cargo.toml | 2 + lib/g3-clap/src/lib.rs | 3 + lib/g3-clap/src/limit/mod.rs | 7 +++ lib/g3-clap/src/limit/rate.rs | 59 +++++++++++++++++++ lib/g3-types/src/limit/rate_limit_quota.rs | 28 ++++----- .../g3bench/target_keyless_openssl.sh | 10 ++++ 11 files changed, 108 insertions(+), 32 deletions(-) create mode 100644 lib/g3-clap/src/limit/mod.rs create mode 100644 lib/g3-clap/src/limit/rate.rs diff --git a/Cargo.lock b/Cargo.lock index b9346d9a..c7d0aef9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -588,18 +588,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.41" +version = "4.5.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be92d32e80243a54711e5d7ce823c35c41c9d929dc4ab58e1276f625841aadf9" +checksum = "ed87a9d530bb41a67537289bafcac159cb3ee28460e0a4571123d2a778a6a882" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.5.41" +version = "4.5.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "707eab41e9622f9139419d573eca0900137718000c517d47da73045f54331c3d" +checksum = "64f4f3f3c77c94aff3c7e9aac9a2ca1974a5adf392a8bb751e827d6d127ab966" dependencies = [ "anstream", "anstyle", @@ -1026,6 +1026,7 @@ version = "0.2.0" dependencies = [ "anyhow", "clap", + "governor", "http", "humanize-rs", ] @@ -3365,9 +3366,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.29" +version = "0.23.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2491382039b29b9b11ff08b76ff6c97cf287671dbb74f0be44bda389fffe9bd1" +checksum = "c0ebcbd2f03de0fc1122ad9bb24b127a5a6cd51d72604a3f3c50ac459762b6cc" dependencies = [ "aws-lc-rs", "brotli", diff --git a/g3bench/CHANGELOG b/g3bench/CHANGELOG index a846883e..dcd6c04c 100644 --- a/g3bench/CHANGELOG +++ b/g3bench/CHANGELOG @@ -1,4 +1,7 @@ +v0.9.7: + - Feature: allow to set more flexible rate limit value + v0.9.6: - BUG FIX: fix wake in cloudflare keyless multiplex task - Feature: add support for thrift tcp target diff --git a/g3bench/Cargo.toml b/g3bench/Cargo.toml index ba700a8a..6db4da06 100644 --- a/g3bench/Cargo.toml +++ b/g3bench/Cargo.toml @@ -41,7 +41,7 @@ hickory-proto.workspace = true g3-runtime.workspace = true g3-std-ext.workspace = true g3-types = { workspace = true, features = ["openssl", "rustls"] } -g3-clap = { workspace = true, features = ["http"] } +g3-clap = { workspace = true, features = ["http", "limit"] } g3-socket.workspace = true g3-http.workspace = true g3-socks.workspace = true diff --git a/g3bench/src/opts.rs b/g3bench/src/opts.rs index da3ad619..b2e12dca 100644 --- a/g3bench/src/opts.rs +++ b/g3bench/src/opts.rs @@ -14,12 +14,12 @@ use std::time::Duration; use ahash::AHashMap; use anyhow::{Context, anyhow}; use clap::{Arg, ArgAction, ArgMatches, Command, ValueHint, value_parser}; +use governor::Quota; use g3_runtime::blended::BlendedRuntimeConfig; use g3_runtime::unaided::UnaidedRuntimeConfig; use g3_statsd_client::{StatsdBackend, StatsdClient, StatsdClientConfig}; use g3_types::collection::{SelectivePickPolicy, SelectiveVec, SelectiveVecBuilder, WeightedValue}; -use g3_types::limit::RateLimitQuotaConfig; use g3_types::metrics::NodeName; use g3_types::net::{TcpSockSpeedLimitConfig, UdpSockSpeedLimitConfig, UpstreamAddr}; @@ -57,7 +57,7 @@ pub struct ProcArgs { pub(super) latency: Option, pub(super) requests: Option, pub(super) time_limit: Option, - pub(super) rate_limit: Option, + pub(super) rate_limit: Option, pub(super) log_error_count: usize, pub(super) ignore_fatal_error: bool, pub(super) task_unconstrained: bool, @@ -453,12 +453,7 @@ pub fn parse_global_args(args: &ArgMatches) -> anyhow::Result { } proc_args.time_limit = g3_clap::humanize::get_duration(args, GLOBAL_ARG_TIME_LIMIT)?; - - if let Some(v) = args.get_one::(GLOBAL_ARG_RATE_LIMIT) { - let rate_limit = - RateLimitQuotaConfig::from_str(v).context("invalid request rate limit value")?; - proc_args.rate_limit = Some(rate_limit); - } + proc_args.rate_limit = g3_clap::limit::get_rate_limit(args, GLOBAL_ARG_RATE_LIMIT)?; if args.get_flag(GLOBAL_ARG_UNAIDED) { proc_args.use_unaided_worker = true; diff --git a/g3bench/src/target/mod.rs b/g3bench/src/target/mod.rs index 44f8641b..bc807daf 100644 --- a/g3bench/src/target/mod.rs +++ b/g3bench/src/target/mod.rs @@ -185,7 +185,7 @@ where let rate_limit = proc_args .rate_limit .as_ref() - .map(|c| Arc::new(RateLimiter::direct(c.get_inner()))); + .map(|q| Arc::new(RateLimiter::direct(*q))); for i in 0..proc_args.concurrency.get() { let sem = Arc::clone(&sync_sem); let barrier = Arc::clone(&sync_barrier); diff --git a/lib/g3-clap/Cargo.toml b/lib/g3-clap/Cargo.toml index d3ca37a5..69742519 100644 --- a/lib/g3-clap/Cargo.toml +++ b/lib/g3-clap/Cargo.toml @@ -11,7 +11,9 @@ anyhow.workspace = true clap.workspace = true humanize-rs.workspace = true http = { workspace = true, optional = true } +governor = { workspace = true, optional = true } [features] default = [] http = ["dep:http"] +limit = ["dep:governor"] diff --git a/lib/g3-clap/src/lib.rs b/lib/g3-clap/src/lib.rs index 4aec1900..740013a7 100644 --- a/lib/g3-clap/src/lib.rs +++ b/lib/g3-clap/src/lib.rs @@ -5,5 +5,8 @@ pub mod humanize; +#[cfg(feature = "limit")] +pub mod limit; + #[cfg(feature = "http")] pub mod http; diff --git a/lib/g3-clap/src/limit/mod.rs b/lib/g3-clap/src/limit/mod.rs new file mode 100644 index 00000000..1cdbaa45 --- /dev/null +++ b/lib/g3-clap/src/limit/mod.rs @@ -0,0 +1,7 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2025 ByteDance and/or its affiliates. + */ + +mod rate; +pub use rate::get_rate_limit; diff --git a/lib/g3-clap/src/limit/rate.rs b/lib/g3-clap/src/limit/rate.rs new file mode 100644 index 00000000..475ecdfa --- /dev/null +++ b/lib/g3-clap/src/limit/rate.rs @@ -0,0 +1,59 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2025 ByteDance and/or its affiliates. + */ + +use std::num::NonZeroU32; +use std::str::FromStr; +use std::time::Duration; + +use anyhow::anyhow; +use clap::ArgMatches; +use governor::Quota; + +pub fn get_rate_limit(args: &ArgMatches, id: &str) -> anyhow::Result> { + if let Some(v) = args.get_one::(id) { + match v.split_once('/') { + Some((v1, v2)) => { + let burst = NonZeroU32::from_str(v1.trim()) + .map_err(|e| anyhow!("invalid burst value: {e}"))?; + let interval_s = v2.trim(); + if let Ok(seconds) = u64::from_str(interval_s) { + let quota = if burst.get() > 1000000 { + let replenish_nanos = seconds * 1000000000 / (burst.get() as u64); + Quota::with_period(Duration::from_nanos(replenish_nanos)) + } else { + let replenish_micros = seconds * 1000000 / (burst.get() as u64); + Quota::with_period(Duration::from_micros(replenish_micros)) + }; + return Ok(quota); + } + if let Ok(interval) = humanize_rs::duration::parse(interval_s) { + let quota = if burst.get() > 1000000 { + let nanos = interval.as_nanos() as u64; + let replenish_nanos = nanos / (burst.get() as u64); + Quota::with_period(Duration::from_nanos(replenish_nanos)) + } else { + let micros = interval.as_micros() as u64; + let replenish_micros = micros / (burst.get() as u64); + Quota::with_period(Duration::from_micros(replenish_micros)) + }; + return Ok(quota); + } + match interval_s { + "s" => Ok(Some(Quota::per_second(burst))), + "m" => Ok(Some(Quota::per_minute(burst))), + "h" => Ok(Some(Quota::per_hour(burst))), + _ => Err(anyhow!("invalid interval value {v}")), + } + } + None => { + let burst = + NonZeroU32::from_str(v).map_err(|e| anyhow!("invalid burst value: {e}"))?; + Ok(Some(Quota::per_second(burst))) + } + } + } else { + Ok(None) + } +} diff --git a/lib/g3-types/src/limit/rate_limit_quota.rs b/lib/g3-types/src/limit/rate_limit_quota.rs index cd52856e..047ff332 100644 --- a/lib/g3-types/src/limit/rate_limit_quota.rs +++ b/lib/g3-types/src/limit/rate_limit_quota.rs @@ -35,26 +35,22 @@ impl FromStr for RateLimitQuotaConfig { type Err = anyhow::Error; fn from_str(s: &str) -> Result { - let ss: Vec<&str> = s.splitn(2, '/').collect(); - - match ss.len() { - 1 => { - let u = NonZeroU32::from_str(ss[0]) - .map_err(|e| anyhow!("invalid non-zero u32 string: {e}"))?; - Ok(RateLimitQuotaConfig(Quota::per_second(u))) - } - 2 => { - let u = NonZeroU32::from_str(ss[0].trim()) + match s.split_once('/') { + Some((v1, v2)) => { + let u = NonZeroU32::from_str(v1.trim()) .map_err(|_| anyhow!("invalid non-zero u32 string as the first part"))?; - let dur = ss[1].trim(); - match dur.to_ascii_lowercase().chars().next() { - Some('s') => Ok(RateLimitQuotaConfig(Quota::per_second(u))), - Some('m') => Ok(RateLimitQuotaConfig(Quota::per_minute(u))), - Some('h') => Ok(RateLimitQuotaConfig(Quota::per_hour(u))), + match v2 { + "s" => Ok(RateLimitQuotaConfig(Quota::per_second(u))), + "m" => Ok(RateLimitQuotaConfig(Quota::per_minute(u))), + "h" => Ok(RateLimitQuotaConfig(Quota::per_hour(u))), _ => Err(anyhow!("invalid unit in second part")), } } - _ => Err(anyhow!("invalid string value")), + None => { + let u = NonZeroU32::from_str(s) + .map_err(|e| anyhow!("invalid non-zero u32 string: {e}"))?; + Ok(RateLimitQuotaConfig(Quota::per_second(u))) + } } } } diff --git a/scripts/coverage/g3bench/target_keyless_openssl.sh b/scripts/coverage/g3bench/target_keyless_openssl.sh index af7875a4..80ba6466 100644 --- a/scripts/coverage/g3bench/target_keyless_openssl.sh +++ b/scripts/coverage/g3bench/target_keyless_openssl.sh @@ -15,3 +15,13 @@ g3bench keyless openssl --cert "${TEST_RSA_CERT_FILE}" --rsa-padding PKCS1 --rsa # EC g3bench keyless openssl --key "${TEST_EC_KEY_FILE}" --sign --digest-type sha256 --verify "4d4dfb668f8c6ddd0227c03907515c58779914098a1bf8c169faafdea4d1b91d" + +# global args + +TARGET_PARAMS="keyless openssl --key ${TEST_EC_KEY_FILE} --sign --digest-type sha256 --verify 4d4dfb668f8c6ddd0227c03907515c58779914098a1bf8c169faafdea4d1b91d" + +g3bench -c 2 -n 4 ${TARGET_PARAMS} +g3bench -c 2 -l 5 -r 10/100ms -t 4 ${TARGET_PARAMS} +g3bench -c 2 -l 5 -r 100/s -t 4 ${TARGET_PARAMS} +g3bench -c 2 -l 5 -r 100 -t 4 ${TARGET_PARAMS} +g3bench -c 1 -t 4 --unaided --emit-metrics ${TARGET_PARAMS}