g3bench: allow more flexible rate limit config

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",

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

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

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<Duration>,
     pub(super) requests: Option<usize>,
     pub(super) time_limit: Option<Duration>,
-    pub(super) rate_limit: Option<RateLimitQuotaConfig>,
+    pub(super) rate_limit: Option<Quota>,
     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<ProcArgs> {
     }
 
     proc_args.time_limit = g3_clap::humanize::get_duration(args, GLOBAL_ARG_TIME_LIMIT)?;
-
-    if let Some(v) = args.get_one::<String>(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;

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);

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"]

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;

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;

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<Option<Quota>> {
+    if let Some(v) = args.get_one::<String>(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)
+    }
+}

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<Self, Self::Err> {
-        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)))
+            }
         }
     }
 }

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}