From 4f9ea76a38949c9dee3b0f6d2111ee387e1a9bd6 Mon Sep 17 00:00:00 2001 From: Ahmed Farghal Date: Mon, 8 Dec 2025 10:41:28 +0000 Subject: [PATCH] Less memory per clock and enable rehydrating buckets Also introduces a new RawTokenBucket for full control from user code. --- Cargo.lock | 2 +- Cargo.toml | 2 +- README.md | 30 ++-- benches/throughput.rs | 47 ++--- benches/time_ops.rs | 12 +- examples/basic.rs | 10 +- examples/fast_clock.rs | 6 +- examples/streams.rs | 6 +- examples/weighted_stream.rs | 6 +- src/bucket.rs | 209 +++++++--------------- src/clock.rs | 145 +++++++-------- src/error.rs | 8 +- src/futures/mod.rs | 12 +- src/futures/stream.rs | 35 ++-- src/lib.rs | 10 +- src/limit.rs | 4 +- src/raw_bucket.rs | 330 +++++++++++++++++++++++++++++++++++ src/storage.rs | 2 +- src/storage/atomic.rs | 4 +- src/storage/local.rs | 2 +- src/storage/padded_atomic.rs | 8 +- src/tokens.rs | 4 +- 22 files changed, 555 insertions(+), 339 deletions(-) create mode 100644 src/raw_bucket.rs diff --git a/Cargo.lock b/Cargo.lock index d1a8b1d..b787c10 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -275,7 +275,7 @@ dependencies = [ [[package]] name = "gardal" -version = "0.0.1-alpha.7" +version = "0.0.1-alpha.8" dependencies = [ "criterion", "futures", diff --git a/Cargo.toml b/Cargo.toml index ec0a215..03065ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "gardal" -version = "0.0.1-alpha.7" +version = "0.0.1-alpha.8" edition = "2024" license = "Apache-2.0 OR MIT" authors = ["Ahmed Farghal "] diff --git a/README.md b/README.md index c21cea7..9f308a6 100644 --- a/README.md +++ b/README.md @@ -25,29 +25,29 @@ Add this to your `Cargo.toml`: ```toml [dependencies] -gardal = "0.0.1-alpha.4" +gardal = "0.0.1-alpha.8" # For async support -gardal = { version = "0.0.1-alpha.4", features = ["async"] } +gardal = { version = "0.0.1-alpha.8", features = ["async"] } # For high-performance timing -gardal = { version = "0.0.1-alpha.4", features = ["quanta"] } +gardal = { version = "0.0.1-alpha.8", features = ["quanta"] } # For high-resolution async timers -gardal = { version = "0.0.1-alpha.4", features = ["async", "tokio-hrtime"] } +gardal = { version = "0.0.1-alpha.8", features = ["async", "tokio-hrtime"] } ``` ### Basic Usage ```rust -use gardal::{Limit, TokenBucket}; +use gardal::{Limit, LocalTokenBucket, StdClock}; use nonzero_ext::nonzero; // Create a token bucket: 10 tokens per second, burst of 20 -let bucket = TokenBucket::new(Limit::per_second_and_burst( +let bucket = LocalTokenBucket::new(Limit::per_second_and_burst( nonzero!(10u32), nonzero!(20u32), -)); +), StdClock); // Consume 5 tokens match bucket.consume(nonzero!(5u32)) { @@ -61,15 +61,15 @@ match bucket.consume(nonzero!(5u32)) { ```rust use futures::{StreamExt, stream}; use gardal::futures::StreamExt as GardalStreamExt; -use gardal::{Limit, TokenBucket, AtomicSharedStorage, QuantaClock}; +use gardal::{Limit, SharedTokenBucket, QuantaClock}; use nonzero_ext::nonzero; #[tokio::main] async fn main() { let limit = Limit::per_second(nonzero!(5u32)); - let bucket = TokenBucket::::from_parts( + let bucket = SharedTokenBucket::new( limit, - QuantaClock::default() + QuantaClock, ); let mut stream = stream::iter(1..=100) @@ -115,13 +115,13 @@ Choose the appropriate storage strategy for your use case: - **`LocalStorage`**: Thread-local storage for single-threaded applications ```rust -use gardal::{TokenBucket, AtomicSharedStorage, Limit}; +use gardal::{TokenBucket, StdClock, AtomicSharedStorage, Limit}; use nonzero_ext::nonzero; // Explicitly specify storage type -let bucket = TokenBucket::::from_parts( +let bucket = TokenBucket::::new( Limit::per_second(nonzero!(10u32)), - gardal::StdClock::default() + StdClock ); ``` @@ -141,13 +141,13 @@ For applications requiring precise timing in async contexts, enable the `tokio-h ```rust use futures::{StreamExt, stream}; use gardal::futures::StreamExt as GardalStreamExt; -use gardal::{Limit, TokenBucket}; +use gardal::{Limit, SharedTokenBucket, TokioClock}; use nonzero_ext::nonzero; #[tokio::main] async fn main() { let limit = Limit::per_second(nonzero!(1000u32)); // High-frequency rate limiting - let bucket = TokenBucket::new(limit); + let bucket = SharedTokenBucket::new(limit, TokioClock); let mut stream = stream::iter(1..=10000) .throttle(bucket) diff --git a/benches/throughput.rs b/benches/throughput.rs index 45eb9cb..4e65848 100644 --- a/benches/throughput.rs +++ b/benches/throughput.rs @@ -9,17 +9,14 @@ use gardal::{ use nonzero_ext::nonzero; fn bench_consume(c: &mut Criterion) { - let clock = quanta::Clock::new(); let limit = Limit::per_second(nonzero!(10_000u32)); - let _quanta_thread = quanta::Upkeep::new_with_clock(Duration::from_micros(10), clock.clone()) + let _quanta_thread = quanta::Upkeep::new(Duration::from_micros(10)) .start() .unwrap(); - let clock = FastClock::new(clock); - let quanta_tb = - TokenBucket::::from_parts(limit, QuantaClock::default()); - let std_tb = TokenBucket::::from_parts(limit, StdClock::default()); - let fast_tb = TokenBucket::::from_parts(limit, clock.clone()); - let fast_tb_padded = TokenBucket::::from_parts(limit, clock.clone()); + let quanta_tb = TokenBucket::::with_datum(limit, QuantaClock); + let std_tb = TokenBucket::::with_datum(limit, StdClock); + let fast_tb = TokenBucket::::with_datum(limit, FastClock); + let fast_tb_padded = TokenBucket::::with_datum(limit, FastClock); std::thread::sleep(Duration::from_secs(1)); let mut group = c.benchmark_group("tokenbucket"); group @@ -27,7 +24,7 @@ fn bench_consume(c: &mut Criterion) { .sample_size(100) .bench_function("consume-mock-clock-local-storage", |b| { let clock = ManualClock::default(); - let tb = TokenBucket::::from_parts(limit, &clock); + let tb = TokenBucket::::with_datum(limit, &clock); clock.set(10.0); b.iter(|| { let _x = std::hint::black_box(tb.try_consume_one()); @@ -35,7 +32,7 @@ fn bench_consume(c: &mut Criterion) { }) .bench_function("consume-mock-clock-atomic-storage", |b| { let clock = ManualClock::default(); - let tb = TokenBucket::::from_parts(limit, &clock); + let tb = TokenBucket::::with_datum(limit, &clock); clock.set(10.0); b.iter(|| { tb.consume_one(); @@ -43,7 +40,7 @@ fn bench_consume(c: &mut Criterion) { }) .bench_function("consume-mock-clock-padded-atomic-storage", |b| { let clock = ManualClock::default(); - let tb = TokenBucket::::from_parts(limit, &clock); + let tb = TokenBucket::::with_datum(limit, &clock); clock.set(10.0); b.iter(|| { tb.consume_one(); @@ -75,19 +72,16 @@ fn bench_consume(c: &mut Criterion) { const THREADS: u32 = 24; fn multi_threaded(c: &mut Criterion) { - let clock = quanta::Clock::new(); - let _quanta_thread = quanta::Upkeep::new_with_clock(Duration::from_micros(100), clock.clone()) + let _quanta_thread = quanta::Upkeep::new(Duration::from_micros(100)) .start() .unwrap(); - let clock = FastClock::new(clock); let limit = Limit::per_second(nonzero!(10_000u32)); let mut group = c.benchmark_group("multi_threaded"); group .throughput(Throughput::Elements(1)) .bench_function("padded", |b| { - let tb = Arc::new(TokenBucket::::from_parts( - limit, - clock.clone(), + let tb = Arc::new(TokenBucket::::with_datum( + limit, FastClock, )); b.iter_custom(|iters| { let mut children = vec![]; @@ -108,9 +102,8 @@ fn multi_threaded(c: &mut Criterion) { }) }) .bench_function("atomic", |b| { - let tb = Arc::new(TokenBucket::::from_parts( - limit, - clock.clone(), + let tb = Arc::new(TokenBucket::::with_datum( + limit, FastClock, )); b.iter_custom(|iters| { let mut children = vec![]; @@ -133,20 +126,17 @@ fn multi_threaded(c: &mut Criterion) { } fn multi_threaded2(c: &mut Criterion) { - let clock = quanta::Clock::new(); - let _quanta_thread = quanta::Upkeep::new_with_clock(Duration::from_micros(10), clock.clone()) + let _quanta_thread = quanta::Upkeep::new(Duration::from_micros(10)) .start() .unwrap(); - let clock = FastClock::new(clock); let limit = Limit::per_second(nonzero!(50u32)); let mut group = c.benchmark_group("multi_threaded2"); group .throughput(Throughput::Elements(1)) .bench_function("padded", |b| { b.iter_custom(|iters| { - let tb = Arc::new(TokenBucket::::from_parts( - limit, - clock.clone(), + let tb = Arc::new(TokenBucket::::with_datum( + limit, FastClock, )); let mut children = vec![]; let start = std::time::Instant::now(); @@ -166,9 +156,8 @@ fn multi_threaded2(c: &mut Criterion) { }) .bench_function("atomic", |b| { b.iter_custom(|iters| { - let tb = Arc::new(TokenBucket::::from_parts( - limit, - clock.clone(), + let tb = Arc::new(TokenBucket::::with_datum( + limit, FastClock, )); let mut children = vec![]; let start = std::time::Instant::now(); diff --git a/benches/time_ops.rs b/benches/time_ops.rs index 4073eed..cd3d155 100644 --- a/benches/time_ops.rs +++ b/benches/time_ops.rs @@ -4,25 +4,21 @@ use criterion::{Criterion, criterion_group, criterion_main}; use gardal::{Clock, FastClock, QuantaClock, StdClock}; fn time_single_threaded(c: &mut Criterion) { - let c_clock = quanta::Clock::new(); // 1KHz - let _quanta_thread = quanta::Upkeep::new_with_clock(Duration::from_micros(10), c_clock.clone()) + let _quanta_thread = quanta::Upkeep::new(Duration::from_micros(10)) .start() .unwrap(); let mut group = c.benchmark_group("gardal"); group .sample_size(100) .bench_function("std-time-getting-instant", |b| { - let clock = StdClock::default(); - b.iter(|| clock.now()); + b.iter(|| std::hint::black_box(StdClock.now())); }) .bench_function("quanta-time-getting-instant", |b| { - let clock = QuantaClock::default(); - b.iter(|| clock.now()); + b.iter(|| std::hint::black_box(QuantaClock.now())); }) .bench_function("quanta-fast-getting-instant", |b| { - let clock = FastClock::new(c_clock.clone()); - b.iter(|| clock.now()); + b.iter(|| std::hint::black_box(FastClock.now())); }); group.finish(); } diff --git a/examples/basic.rs b/examples/basic.rs index 7537a9b..17065e5 100644 --- a/examples/basic.rs +++ b/examples/basic.rs @@ -1,13 +1,13 @@ use std::time::Duration; -use gardal::{Limit, TokenBucket}; +use gardal::{AtomicTokenBucket, Limit, StdClock}; use nonzero_ext::nonzero; fn main() { - let tb = TokenBucket::new(Limit::per_second_and_burst( - nonzero!(10u32), - nonzero!(20u32), - )); + let tb = AtomicTokenBucket::new( + Limit::per_second_and_burst(nonzero!(10u32), nonzero!(20u32)), + StdClock, + ); // after two seconds bucket should be full std::thread::sleep(Duration::from_secs(2)); assert_eq!(5, tb.consume(nonzero!(5u32)).unwrap().as_u64()); diff --git a/examples/fast_clock.rs b/examples/fast_clock.rs index d03cb01..8d61793 100644 --- a/examples/fast_clock.rs +++ b/examples/fast_clock.rs @@ -4,14 +4,12 @@ use gardal::{FastClock, Limit, TokenBucket}; use nonzero_ext::nonzero; fn main() { - let clock = quanta::Clock::new(); // Updates at 1Khz - let _quanta_thread = quanta::Upkeep::new_with_clock(Duration::from_millis(1), clock.clone()) + let _quanta_thread = quanta::Upkeep::new(Duration::from_millis(1)) .start() .unwrap(); - let clock = FastClock::new(clock); let limit = Limit::per_second_and_burst(nonzero!(10u32), nonzero!(20u32)); - let tb = TokenBucket::with_clock(limit, clock); + let tb = TokenBucket::with_clock(limit, FastClock); // after two seconds bucket should be full println!("sleeping for 2 seconds..."); std::thread::sleep(Duration::from_secs(2)); diff --git a/examples/streams.rs b/examples/streams.rs index b650424..289d7b4 100644 --- a/examples/streams.rs +++ b/examples/streams.rs @@ -4,16 +4,16 @@ use std::time::Duration; use futures::StreamExt; use futures::stream; +use gardal::SharedTokenBucket; use gardal::futures::StreamExt as GardalStreamExt; -use gardal::{Limit, PaddedAtomicSharedStorage, TokenBucket, TokioClock}; +use gardal::{Limit, TokioClock}; use nonzero_ext::nonzero; use tokio::task::JoinSet; #[tokio::main(flavor = "multi_thread")] async fn main() { let limit = Limit::per_second_and_burst(nonzero!(1000000u32), nonzero!(100u32)); - let bucket = - TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = SharedTokenBucket::new(limit, TokioClock); let mut print_start = tokio::time::Instant::now(); let program_start = print_start; diff --git a/examples/weighted_stream.rs b/examples/weighted_stream.rs index 8139953..4b5f429 100644 --- a/examples/weighted_stream.rs +++ b/examples/weighted_stream.rs @@ -1,6 +1,6 @@ use futures::stream; use gardal::futures::{StreamExt as GardalStreamExt, WeightedStream}; -use gardal::{Limit, LocalStorage, TokenBucket, TokioClock}; +use gardal::{Limit, LocalTokenBucket, TokioClock}; use nonzero_ext::nonzero; use std::num::NonZeroU32; use tokio_stream::StreamExt; @@ -38,7 +38,7 @@ async fn main() { // Create a throttling limit: 5 tokens per second with a burst of 25 let limit = Limit::per_second_and_burst(nonzero!(5u32), nonzero!(25u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = LocalTokenBucket::new(limit, TokioClock); // Create a weighted stream where each task consumes tokens based on its size let weighted_stream = WeightedStream::new(stream, bucket, |task: &Task| { @@ -76,7 +76,7 @@ async fn main() { let stream2 = stream::iter(tasks2); let limit2 = Limit::per_second_and_burst(nonzero!(3u32), nonzero!(25u32)); - let bucket2 = TokenBucket::::from_parts(limit2, TokioClock::default()); + let bucket2 = LocalTokenBucket::new(limit2, TokioClock); // Use the extension trait to create a weighted stream let weighted_stream2 = stream2.throttle_weighted(bucket2, |text: &&str| { diff --git a/src/bucket.rs b/src/bucket.rs index f790230..5c7b4ee 100644 --- a/src/bucket.rs +++ b/src/bucket.rs @@ -1,14 +1,11 @@ use std::num::NonZeroU32; use std::time::Duration; -use likely_stable::unlikely; - use crate::clock::Nanos; use crate::error::{ExceededBurstCapacity, RateLimited}; -use crate::storage::atomic::AtomicStorage; +use crate::storage::TimeStorage; use crate::storage::padded_atomic::PaddedAtomicStorage; -use crate::storage::{TimeStorage, TokenAcquisition, TokenBucketStorage}; -use crate::{Clock, Limit, StdClock, Tokens}; +use crate::{Clock, Limit, RawTokenBucket, StdClock, Tokens}; pub const UNLIMITED_BUCKET: Option = const { None }; @@ -21,19 +18,19 @@ pub const UNLIMITED_BUCKET: Option = const { None }; /// # Type Parameters /// /// - `S`: Storage strategy (default: [`PaddedAtomicStorage`] for concurrent access) -/// - `C`: Clock implementation (default: [`StdClock`] for standard timing) +/// - `C`: Clock implementation (use [`StdClock`] for standard library timing) /// /// # Examples /// /// ```rust -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second_and_burst( /// NonZeroU32::new(10).unwrap(), /// NonZeroU32::new(20).unwrap() /// ); -/// let bucket = TokenBucket::new(limit); +/// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// // Try to consume 5 tokens /// if let Some(tokens) = bucket.consume(NonZeroU32::new(5).unwrap()) { @@ -42,34 +39,11 @@ pub const UNLIMITED_BUCKET: Option = const { None }; /// ``` #[derive(Clone)] pub struct TokenBucket { - bucket: TokenBucketStorage, + bucket: RawTokenBucket, clock: C, limit: Limit, } -impl TokenBucket { - /// Creates a new token bucket with the specified rate limit using atomic storage and standard clock. - /// - /// This is the most common constructor for basic use cases requiring thread-safe access. - /// - /// # Arguments - /// - /// * `limit` - The rate and burst configuration for the bucket - /// - /// # Examples - /// - /// ```rust - /// use gardal::{TokenBucket, Limit}; - /// use std::num::NonZeroU32; - /// - /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); - /// let bucket = TokenBucket::new(limit); - /// ``` - pub fn new(limit: Limit) -> Self { - TokenBucket::::from_parts(limit, StdClock::default()) - } -} - impl TokenBucket { /// Creates a new token bucket with a custom clock implementation. /// @@ -92,9 +66,8 @@ impl TokenBucket { /// let bucket = TokenBucket::with_clock(limit, clock); /// ``` pub fn with_clock(limit: Limit, clock: C) -> Self { - let storage = PaddedAtomicStorage::new(clock.now()); Self { - bucket: TokenBucketStorage::::new(storage), + bucket: RawTokenBucket::::new(&clock), clock, limit, } @@ -102,6 +75,23 @@ impl TokenBucket { } impl TokenBucket { + /// Creates a token bucket + /// + /// The time datum defaults to 0.0 which means that the bucket will be initially full + /// in case you are using one of the standard clocks provided by this crate. + /// + /// # Arguments + /// + /// * `limit` - The rate and burst configuration for the bucket + /// * `clock` - The clock implementation to use for timing. + pub fn new(limit: Limit, clock: C) -> Self { + Self { + bucket: RawTokenBucket::with_zero_time(0.0), + clock, + limit, + } + } + /// Creates a token bucket from custom storage and clock implementations. /// /// This provides maximum flexibility for advanced use cases requiring specific @@ -110,12 +100,12 @@ impl TokenBucket { /// # Arguments /// /// * `limit` - The rate and burst configuration for the bucket - /// * `clock` - The clock implementation to use for timing - pub fn from_parts(limit: Limit, clock: C) -> Self { - // let storage = S::new(clock.now()); - let storage = S::new(0.0); + /// * `clock` - The clock implementation to use for timing. It will use the clock's + /// datum point as the zero time. This means that the bucket will be considered + /// empty initially. + pub fn with_datum(limit: Limit, clock: C) -> Self { Self { - bucket: TokenBucketStorage::new(storage), + bucket: RawTokenBucket::new(&clock), clock, limit, } @@ -130,12 +120,9 @@ impl TokenBucket { /// * `limit` - The new rate and burst configuration pub fn reset(&mut self, limit: Limit) { let now = self.clock.now(); - let available = self - .bucket - .balance(self.limit.rate, self.limit.burst, now) - .max(0.0); + let available = self.bucket.balance_at(now, &self.limit); self.limit = limit; - self.set_capacity(available, now); + self.bucket.set_capacity(available, now, limit.rate); } /// Consumes the bucket and returns a new one with updated rate limits. @@ -151,24 +138,16 @@ impl TokenBucket { /// A new token bucket with the updated configuration pub fn update_limit(self, limit: Limit) -> Self { let now = self.clock.now(); - let available = self - .bucket - .balance(self.limit.rate, self.limit.burst, now) - .max(0.0); + let available = self.bucket.balance_at(now, &self.limit); let mut new = Self { bucket: self.bucket, clock: self.clock, limit, }; - new.set_capacity(available, now); + new.bucket.set_capacity(available, now, new.limit.rate); new } - /// Set the number of tokens currently available in the bucket. - fn set_capacity(&mut self, tokens: f64, now: f64) { - self.bucket.reset(now - tokens / self.limit.rate); - } - /// Attempts to consume exactly the specified number of tokens. /// /// This is the fastest consumption method. Returns `Some(tokens)` if successful, @@ -188,11 +167,11 @@ impl TokenBucket { /// # Examples /// /// ```rust - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::new(limit); + /// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// if let Some(tokens) = bucket.consume(NonZeroU32::new(5).unwrap()) { /// println!("Consumed {} tokens", tokens.as_u64()); @@ -201,16 +180,7 @@ impl TokenBucket { /// } /// ``` pub fn consume(&self, to_consume: impl Into) -> Option { - let now = self.clock.now(); - let to_consume: NonZeroU32 = to_consume.into(); - let to_consume: f64 = to_consume.get() as f64; - - let consumed = self - .bucket - .consume(self.limit.rate, self.limit.burst, now, |avail| { - if avail < to_consume { 0.0 } else { to_consume } - }); - Tokens::new_checked(consumed) + self.bucket.consume(to_consume, &self.clock, &self.limit) } /// Attempts to consume exactly one token. @@ -222,7 +192,8 @@ impl TokenBucket { /// * `Some(Tokens)` - Successfully consumed one token /// * `None` - No tokens available pub fn consume_one(&self) -> Option { - self.consume(NonZeroU32::new(1u32).unwrap()) + self.bucket + .consume(NonZeroU32::new(1u32).unwrap(), &self.clock, &self.limit) } /// Attempts to consume one token with wait time information. @@ -234,7 +205,8 @@ impl TokenBucket { /// * `Ok(Tokens)` - Successfully consumed one token /// * `Err(RateLimited)` - Rate limited with suggested wait time pub fn try_consume_one(&self) -> Result { - self.try_consume(NonZeroU32::new(1u32).unwrap()) + self.bucket + .try_consume(NonZeroU32::new(1u32).unwrap(), &self.clock, &self.limit) } /// Attempts to consume tokens with detailed rate limiting information. @@ -254,11 +226,11 @@ impl TokenBucket { /// # Examples /// /// ```rust - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::new(limit); + /// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// match bucket.try_consume(NonZeroU32::new(5).unwrap()) { /// Ok(tokens) => println!("Consumed {} tokens", tokens.as_u64()), @@ -269,24 +241,8 @@ impl TokenBucket { /// } /// ``` pub fn try_consume(&self, to_consume: impl Into) -> Result { - let to_consume: NonZeroU32 = to_consume.into(); - let to_consume: f64 = to_consume.get() as f64; - let now = self.clock.now(); - let consumed = self - .bucket - .consume2(self.limit.rate, self.limit.burst, now, |avail| { - if avail < to_consume { 0.0 } else { to_consume } - }); - match consumed { - TokenAcquisition::Acquired(consumed) => Ok(Tokens::new_unchecked(consumed)), - TokenAcquisition::ZeroedAt(zero_time) => { - let est_time = zero_time - now + to_consume / self.limit.rate; - debug_assert!(est_time >= 0.0); - Err(RateLimited { - earliest_retry_time: Nanos::from_secs_f64_unchecked(est_time), - }) - } - } + self.bucket + .try_consume(to_consume, &self.clock, &self.limit) } /// Consumes up to the requested number of tokens, returning whatever is available. @@ -306,11 +262,11 @@ impl TokenBucket { /// # Examples /// /// ```rust - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::new(limit); + /// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// // Request 100 tokens, but only get what's available /// if let Some(tokens) = bucket.saturating_consume(NonZeroU32::new(100).unwrap()) { @@ -318,10 +274,8 @@ impl TokenBucket { /// } /// ``` pub fn saturating_consume(&self, to_consume: impl Into) -> Option { - let now = self.clock.now(); - let to_consume: NonZeroU32 = to_consume.into(); - let to_consume: f64 = to_consume.get() as f64; - Tokens::new_checked(self.saturating_consume_inner(to_consume, now)) + self.bucket + .saturating_consume(to_consume, &self.clock, &self.limit) } /// Returns unused tokens to the bucket or manually adds tokens. @@ -336,19 +290,17 @@ impl TokenBucket { /// # Examples /// /// ```rust - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::new(limit); + /// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// // Return 5 tokens to the bucket /// bucket.add_tokens(5.0); /// ``` pub fn add_tokens(&self, tokens: impl Into) { - let tokens = tokens.into(); - debug_assert!(tokens > 0.0); - self.bucket.return_tokens(tokens, self.limit.rate); + self.bucket.add_tokens(tokens, &self.limit) } /// Consumes tokens by borrowing from future capacity. @@ -369,14 +321,14 @@ impl TokenBucket { /// # Examples /// /// ```rust - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second_and_burst( /// NonZeroU32::new(10).unwrap(), /// NonZeroU32::new(20).unwrap() /// ); - /// let bucket = TokenBucket::new(limit); + /// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// match bucket.consume_with_borrow(NonZeroU32::new(15).unwrap()) { /// Ok(Some(wait_time)) => { @@ -390,24 +342,8 @@ impl TokenBucket { &self, to_consume: impl Into, ) -> Result, ExceededBurstCapacity> { - let now = self.clock.now(); - let to_consume: NonZeroU32 = to_consume.into(); - let mut to_consume: f64 = to_consume.get() as f64; - if unlikely(self.limit.burst < to_consume) { - return Err(ExceededBurstCapacity); - } - while to_consume > 0.0 { - let consumed = self.saturating_consume_inner(to_consume, now); - if consumed > 0.0 { - to_consume -= consumed; - } else { - self.bucket.return_tokens(-to_consume, self.limit.rate); - let debt_paid = self.bucket.time_when_bucket(self.limit.rate, 0.0); - let nap_time = (debt_paid - now).max(0.0); - return Ok(Nanos::new_checked(nap_time)); - } - } - Ok(None) + self.bucket + .consume_with_borrow(to_consume, &self.clock, &self.limit) } /// Consumes tokens with borrowing, limited to burst capacity. @@ -428,26 +364,8 @@ impl TokenBucket { &self, to_consume: impl Into, ) -> (Option, Duration) { - let now = self.clock.now(); - let to_consume: NonZeroU32 = to_consume.into(); - let mut to_consume: f64 = to_consume.get() as f64; - to_consume = to_consume.min(self.limit.burst); - let actual_to_be_consumed = to_consume; - while to_consume > 0.0 { - let consumed = self.saturating_consume_inner(to_consume, now); - if consumed > 0.0 { - to_consume -= consumed; - } else { - self.bucket.return_tokens(-to_consume, self.limit.rate); - let debt_paid = self.bucket.time_when_bucket(self.limit.rate, 0.0); - let nap_time = (debt_paid - now).max(0.0); - return ( - Tokens::new_checked(actual_to_be_consumed), - Duration::from_secs_f64(nap_time), - ); - } - } - (None, Duration::ZERO) + self.bucket + .saturating_consume_with_borrow(to_consume, &self.clock, &self.limit) } /// Returns the number of tokens currently available for consumption. @@ -459,7 +377,7 @@ impl TokenBucket { /// /// Number of tokens available for immediate consumption pub fn available(&self) -> f64 { - self.balance().max(0.0) + self.bucket.available(&self.clock, &self.limit) } /// Returns the current token balance, which may be negative if in debt. @@ -471,8 +389,7 @@ impl TokenBucket { /// /// Current token balance (negative indicates debt) pub fn balance(&self) -> f64 { - self.bucket - .balance(self.limit.rate, self.limit.burst, self.clock.now()) + self.bucket.balance(&self.clock, &self.limit) } /// Returns a reference to the current rate limit configuration. @@ -484,11 +401,9 @@ impl TokenBucket { &self.limit } - fn saturating_consume_inner(&self, to_consume: f64, now: f64) -> f64 { - self.bucket - .consume(self.limit.rate, self.limit.burst, now, |avail| { - avail.max(0.0).min(to_consume) - }) + /// Returns the internal timepoint of the bucket. + pub fn get_zero_time(&self) -> f64 { + self.bucket.get_zero_time() } } diff --git a/src/clock.rs b/src/clock.rs index b82145e..d927dd6 100644 --- a/src/clock.rs +++ b/src/clock.rs @@ -1,7 +1,28 @@ use std::fmt::{Debug, Display}; use std::num::NonZeroU64; -use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::sync::{Arc, LazyLock, Mutex}; +use std::time::{Duration, SystemTime}; + +struct Datum { + dur_since_epoch: f64, + std_instant: std::time::Instant, + #[cfg(feature = "tokio")] + tokio_instant: tokio::time::Instant, + #[cfg(feature = "quanta")] + quanta_instant: quanta::Instant, +} + +static CLOCK_ORIGIN: LazyLock = LazyLock::new(|| Datum { + dur_since_epoch: SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_secs_f64(), + std_instant: std::time::Instant::now(), + #[cfg(feature = "tokio")] + tokio_instant: tokio::time::Instant::now(), + #[cfg(feature = "quanta")] + quanta_instant: quanta::Instant::now(), +}); /// Represents a non-zero time duration in nanoseconds. #[derive(Clone, Copy, PartialEq, Eq)] @@ -50,6 +71,12 @@ impl Display for Nanos { /// Implementations must provide monotonic time that never goes backwards. /// The time is measured in seconds as floating-point values. pub trait Clock { + /// The origin time point of the clock. + /// + /// The origin is the point where now() is guaranteed to be after. + fn datum(&self) -> f64 { + 0.0 + } /// Returns the current time in seconds since an arbitrary epoch. /// /// The returned value must be monotonic (never decrease) and should @@ -77,24 +104,19 @@ pub trait Clock { /// let clock = StdClock::default(); /// let bucket = TokenBucket::with_clock(limit, clock); /// ``` -#[derive(Clone)] -pub struct StdClock { - origin: std::time::Instant, -} +#[derive(Clone, Default)] +pub struct StdClock; -impl Default for StdClock { - fn default() -> Self { - Self { - origin: std::time::Instant::now(), - } +impl Clock for StdClock { + fn datum(&self) -> f64 { + CLOCK_ORIGIN.dur_since_epoch } -} -impl Clock for StdClock { fn now(&self) -> f64 { std::time::Instant::now() - .duration_since(self.origin) + .duration_since(CLOCK_ORIGIN.std_instant) .as_secs_f64() + + CLOCK_ORIGIN.dur_since_epoch } } @@ -117,35 +139,16 @@ impl Clock for StdClock { /// # } /// ``` #[cfg(feature = "quanta")] -#[derive(Clone)] -pub struct QuantaClock { - origin: quanta::Instant, -} - -#[cfg(feature = "quanta")] -impl Default for QuantaClock { - fn default() -> Self { - Self::new(quanta::Clock::new()) - } -} - -#[cfg(feature = "quanta")] -impl QuantaClock { - /// Creates a new `QuantaClock` from a `quanta::Clock` instance. - /// - /// # Arguments - /// - /// * `clock` - The quanta clock instance to use - pub fn new(clock: quanta::Clock) -> Self { - let origin = clock.now(); - Self { origin } - } -} +#[derive(Clone, Default)] +pub struct QuantaClock; #[cfg(feature = "quanta")] impl Clock for QuantaClock { fn now(&self) -> f64 { - self.origin.elapsed().as_secs_f64() + quanta::Instant::now() + .duration_since(CLOCK_ORIGIN.quanta_instant) + .as_secs_f64() + + CLOCK_ORIGIN.dur_since_epoch } } @@ -162,29 +165,24 @@ impl Clock for QuantaClock { /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); -/// let clock = TokioClock::default(); -/// let bucket = TokenBucket::with_clock(limit, clock); +/// let bucket = TokenBucket::with_clock(limit, TokioClock); /// # } /// ``` #[cfg(feature = "tokio")] -#[derive(Clone)] -pub struct TokioClock { - origin: tokio::time::Instant, -} +#[derive(Clone, Default)] +pub struct TokioClock; #[cfg(feature = "tokio")] -impl Default for TokioClock { - fn default() -> Self { - Self { - origin: tokio::time::Instant::now(), - } +impl Clock for TokioClock { + fn datum(&self) -> f64 { + CLOCK_ORIGIN.dur_since_epoch } -} -#[cfg(feature = "tokio")] -impl Clock for TokioClock { fn now(&self) -> f64 { - self.origin.elapsed().as_secs_f64() + tokio::time::Instant::now() + .duration_since(CLOCK_ORIGIN.tokio_instant) + .as_secs_f64() + + CLOCK_ORIGIN.dur_since_epoch } } @@ -209,43 +207,24 @@ impl Clock for TokioClock { /// /// let limit = Limit::per_second(NonZeroU32::new(1000).unwrap()); /// let clock = FastClock::default(); -/// let bucket = TokenBucket::with_clock(limit, clock); +/// let bucket = TokenBucket::with_clock(limit, FastClock); /// # } /// ``` #[cfg(feature = "quanta")] -#[derive(Clone)] -pub struct FastClock { - clock: quanta::Clock, - origin: quanta::Instant, -} - -#[cfg(feature = "quanta")] -impl Default for FastClock { - fn default() -> Self { - Self::new(quanta::Clock::new()) - } -} +#[derive(Clone, Default)] +pub struct FastClock; #[cfg(feature = "quanta")] -impl FastClock { - /// Creates a new `FastClock` from a `quanta::Clock` instance. - /// - /// **Important**: Ensure the clock's upkeep thread is running, otherwise - /// the token bucket will not observe clock changes and timing will be incorrect. - /// - /// # Arguments - /// - /// * `clock` - The quanta clock instance to use - pub fn new(clock: quanta::Clock) -> Self { - let origin = clock.recent(); - Self { clock, origin } +impl Clock for FastClock { + fn datum(&self) -> f64 { + CLOCK_ORIGIN.dur_since_epoch } -} -#[cfg(feature = "quanta")] -impl Clock for FastClock { fn now(&self) -> f64 { - (self.clock.recent() - self.origin).as_secs_f64() + quanta::Instant::recent() + .duration_since(CLOCK_ORIGIN.quanta_instant) + .as_secs_f64() + + CLOCK_ORIGIN.dur_since_epoch } } diff --git a/src/error.rs b/src/error.rs index f56ae89..2549206 100644 --- a/src/error.rs +++ b/src/error.rs @@ -10,11 +10,11 @@ use crate::clock::Nanos; /// # Examples /// /// ```rust -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(1).unwrap()); -/// let bucket = TokenBucket::new(limit); +/// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// match bucket.try_consume_one() { /// Ok(tokens) => println!("Got tokens: {}", tokens.as_u64()), @@ -36,14 +36,14 @@ pub struct RateLimited { /// # Examples /// /// ```rust -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second_and_burst( /// NonZeroU32::new(10).unwrap(), /// NonZeroU32::new(20).unwrap() /// ); -/// let bucket = TokenBucket::new(limit); +/// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// // This will fail because 25 > burst capacity of 20 /// match bucket.consume_with_borrow(NonZeroU32::new(25).unwrap()) { diff --git a/src/futures/mod.rs b/src/futures/mod.rs index a363e61..7efb68b 100644 --- a/src/futures/mod.rs +++ b/src/futures/mod.rs @@ -13,14 +13,14 @@ //! ```rust //! # #[cfg(feature = "async")] //! # { -//! use gardal::{TokenBucket, Limit}; +//! use gardal::{AtomicTokenBucket, Limit, StdClock}; //! use gardal::futures::StreamExt; //! use futures::stream; //! use std::num::NonZeroU32; //! //! # async fn example() { //! let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); -//! let bucket = TokenBucket::new(limit); +//! let bucket = AtomicTokenBucket::new(limit, StdClock); //! //! let stream = stream::iter(0..100) //! .throttle(Some(bucket)); @@ -49,14 +49,14 @@ use crate::{Clock, TokenBucket}; /// ```rust /// # #[cfg(feature = "async")] /// # { -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{AtomicTokenBucket, Limit, StdClock}; /// use gardal::futures::StreamExt; /// use futures::stream; /// use std::num::NonZeroU32; /// /// # async fn example() { /// let limit = Limit::per_second(NonZeroU32::new(5).unwrap()); -/// let bucket = TokenBucket::new(limit); +/// let bucket = AtomicTokenBucket::new(limit, StdClock); /// /// let throttled = stream::iter(1..=10) /// .throttle(Some(bucket)); @@ -100,14 +100,14 @@ where /// ```rust /// # #[cfg(feature = "async")] /// # { - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{AtomicTokenBucket, Limit, StdClock}; /// use gardal::futures::StreamExt; /// use futures::stream; /// use std::num::NonZeroU32; /// /// # async fn example() { /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::new(limit); + /// let bucket = AtomicTokenBucket::new(limit, StdClock); /// /// let throttled = stream::iter(vec!["small", "large", "medium"]) /// .throttle_weighted(bucket, |item: &&str| { diff --git a/src/futures/stream.rs b/src/futures/stream.rs index 7be0363..341cc36 100644 --- a/src/futures/stream.rs +++ b/src/futures/stream.rs @@ -20,13 +20,13 @@ pin_project! { /// # Examples /// /// ```rust - /// use gardal::{TokenBucket, Limit}; + /// use gardal::{AtomicTokenBucket, Limit, StdClock}; /// use gardal::futures::ThrottledStream; /// use futures::stream; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::new(limit); + /// let bucket = AtomicTokenBucket::new(limit, StdClock); /// let stream = stream::iter(0..100); /// /// let throttled = ThrottledStream::new(stream, bucket); @@ -161,12 +161,12 @@ where /// /// ```rust /// use gardal::futures::WeightedStream; - /// use gardal::{LocalStorage, Limit, TokioClock, TokenBucket}; + /// use gardal::{LocalStorage, Limit, TokioClock, SharedTokenBucket}; /// use futures::stream; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second_and_burst(NonZeroU32::new(10).unwrap(), NonZeroU32::new(10).unwrap()); - /// let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + /// let bucket = SharedTokenBucket::new(limit, TokioClock); /// /// let stream = stream::iter(vec!["small", "large", "medium"]); /// let weighted_stream = WeightedStream::new(stream, bucket, |item: &&str| { @@ -327,7 +327,7 @@ mod tests { let start = tokio::time::Instant::now(); let stream = stream::iter(vec![1, 2, 3, 4, 5]); let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(1u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut throttled_stream = std::pin::pin!(ThrottledStream::new(stream, bucket)); @@ -345,7 +345,7 @@ mod tests { async fn test_throttled_stream_burst() { let stream = stream::iter(vec![1, 2, 3, 4, 5]); let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(3u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut throttled_stream = std::pin::pin!(ThrottledStream::new(stream, bucket)); @@ -364,7 +364,7 @@ mod tests { async fn test_throttled_stream_all_ready() { let stream = stream::iter(vec![1, 2, 3, 4, 5]); let limit = Limit::per_second(nonzero!(100000u32)).with_burst(nonzero!(1u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut throttled_stream = std::pin::pin!(ThrottledStream::new(stream, bucket)); @@ -385,7 +385,7 @@ mod tests { .throttle(Duration::from_secs(2)) .chain(stream::iter(vec![6, 7, 8, 9])); let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(3u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut throttled_stream = std::pin::pin!(ThrottledStream::new(stream, bucket)); @@ -405,7 +405,7 @@ mod tests { let start = tokio::time::Instant::now(); let stream = stream::iter(vec![1, 2, 3, 4, 5]); let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(1u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |_| nonzero!(1u32))); @@ -425,7 +425,7 @@ mod tests { let start = tokio::time::Instant::now(); let stream = stream::iter(vec![1, 2, 3, 4, 5]); let limit = Limit::per_second_and_burst(nonzero!(2u32), nonzero!(2u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |&item| { if item % 2 == 0 { @@ -451,7 +451,7 @@ mod tests { async fn test_weighted_stream_with_burst() { let stream = stream::iter(vec![1, 2, 3, 4, 5]); let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(5u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |&item| { NonZeroU32::new(item as u32).unwrap_or(nonzero!(1u32)) @@ -474,7 +474,7 @@ mod tests { async fn test_weighted_stream_empty() { let stream = stream::iter(Vec::::new()); let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(1u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |_| nonzero!(1u32))); @@ -491,7 +491,7 @@ mod tests { async fn test_weighted_stream_single_item() { let stream = stream::iter(vec![42]); let limit = Limit::per_second_and_burst(nonzero!(10u32), nonzero!(10u32)); - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |_| nonzero!(3u32))); @@ -513,7 +513,7 @@ mod tests { // This test verifies that expensive items are properly delayed before being returned let stream = stream::iter(vec![10]); // Single expensive item let limit = Limit::per_second_and_burst(nonzero!(1u32), nonzero!(10u32)); // Enough burst for the item - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |&item| { NonZeroU32::new(item as u32).unwrap_or(nonzero!(1u32)) // Item value determines token cost @@ -532,7 +532,10 @@ mod tests { // The expensive item (10 tokens) should be delayed significantly // With 1 token/sec rate and 10 token burst, it should take ~9 seconds - assert!(delivery_time >= Duration::from_secs(9)); + assert!( + delivery_time >= Duration::from_secs(9), + ":{delivery_time:?}" + ); } #[tokio::test(start_paused = true)] @@ -540,7 +543,7 @@ mod tests { // Test that cheap items come quickly and expensive items are delayed appropriately let stream = stream::iter(vec![1, 5, 1]); // cheap, expensive, cheap let limit = Limit::per_second_and_burst(nonzero!(2u32), nonzero!(10u32)); // Enough burst capacity - let bucket = TokenBucket::::from_parts(limit, TokioClock::default()); + let bucket = TokenBucket::::with_datum(limit, TokioClock::default()); let mut weighted_stream = std::pin::pin!(WeightedStream::new(stream, bucket, |&item| { NonZeroU32::new(item as u32).unwrap_or(nonzero!(1u32)) diff --git a/src/lib.rs b/src/lib.rs index 1ce4d22..2536ade 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,7 +12,7 @@ //! ```rust //! use std::num::NonZeroU32; //! -//! use gardal::{TokenBucket, Limit}; +//! use gardal::{AtomicTokenBucket, Limit, StdClock}; //! //! // Create a rate limit: 10 tokens per second, burst of 20 //! let limit = Limit::per_second_and_burst( @@ -21,7 +21,7 @@ //! ); //! //! // Create token bucket -//! let bucket = TokenBucket::new(limit); +//! let bucket = AtomicTokenBucket::new(limit, StdClock); //! //! // Try to consume tokens //! if let Some(tokens) = bucket.consume(NonZeroU32::new(5).unwrap()) { @@ -35,6 +35,7 @@ mod error; #[cfg(feature = "async")] pub mod futures; mod limit; +mod raw_bucket; mod storage; mod tokens; @@ -48,6 +49,7 @@ pub use error::*; #[cfg(feature = "async")] pub use futures::StreamExt; pub use limit::Limit; +pub use raw_bucket::RawTokenBucket; pub use tokens::Tokens; pub use storage::{ @@ -55,6 +57,10 @@ pub use storage::{ padded_atomic::PaddedAtomicSharedStorage, padded_atomic::PaddedAtomicStorage, }; +pub type LocalTokenBucket = TokenBucket; +pub type AtomicTokenBucket = TokenBucket; +pub type SharedTokenBucket = TokenBucket; + pub(crate) mod private { pub trait Sealed {} } diff --git a/src/limit.rs b/src/limit.rs index f1352e7..10386ec 100644 --- a/src/limit.rs +++ b/src/limit.rs @@ -115,7 +115,7 @@ impl Limit { pub const fn per_minute(rate: NonZeroU32) -> Self { Self { rate: rate.get() as f64 / SECONDS_PER_MINUTE, - burst: rate.get() as f64 / SECONDS_PER_MINUTE, + burst: (rate.get() as f64).max(1.0), } } @@ -140,7 +140,7 @@ impl Limit { pub const fn per_hour(rate: NonZeroU32) -> Self { Self { rate: rate.get() as f64 / SECONDS_PER_HOUR, - burst: rate.get() as f64 / SECONDS_PER_HOUR, + burst: (rate.get() as f64).max(1.0), } } diff --git a/src/raw_bucket.rs b/src/raw_bucket.rs new file mode 100644 index 0000000..484d5b9 --- /dev/null +++ b/src/raw_bucket.rs @@ -0,0 +1,330 @@ +use std::marker::PhantomData; +use std::num::NonZeroU32; +use std::time::Duration; + +use likely_stable::unlikely; + +use crate::clock::Nanos; +use crate::error::{ExceededBurstCapacity, RateLimited}; +use crate::storage::{TimeStorage, TokenAcquisition, TokenBucketStorage}; +use crate::{Clock, Limit, Tokens}; + +#[derive(Clone)] +pub struct RawTokenBucket { + bucket: TokenBucketStorage, + _clock: PhantomData, +} + +impl RawTokenBucket { + /// Creates a token bucket from custom storage and clock implementations. + /// + /// This provides maximum flexibility for advanced use cases requiring specific + /// storage strategies or clock implementations. + /// + /// # Arguments + /// + /// * `clock` - The clock implementation to use for setting the origin datum time + pub fn new(clock: &C) -> Self { + let storage = S::new(clock.datum()); + Self { + bucket: TokenBucketStorage::new(storage), + _clock: PhantomData, + } + } + + /// Returns the internal timepoint of the bucket. + /// + /// In combination with [`set_zero_time`](Self::set_zero_time), this can be used to + /// create a new bucket with a different origin time. For instance, if you'd like + /// to reconstruct a bucket with the exact same state as another. + pub fn get_zero_time(&self) -> f64 { + self.bucket.inner.load() + } + + pub fn with_zero_time(time_point: f64) -> Self { + let storage = S::new(time_point); + Self { + bucket: TokenBucketStorage::new(storage), + _clock: PhantomData, + } + } + + /// Updates the rate limit configuration while preserving available tokens. + /// + /// The current token balance is maintained proportionally when changing limits. + /// + /// # Arguments + /// + /// * `clock` - The clock used originally to create the bucket + /// * `limit` - The new rate and burst configuration + pub fn reset(&mut self, clock: &C, limit: &Limit) { + let now = clock.now(); + let available = self.bucket.balance(limit.rate, limit.burst, now).max(0.0); + self.set_capacity(available, now, limit.rate); + } + + /// Set the number of tokens currently available in the bucket. + pub(crate) fn set_capacity(&mut self, tokens: f64, now: f64, rate: f64) { + self.bucket.reset(now - tokens / rate); + } + + /// Attempts to consume exactly the specified number of tokens. + /// + /// This is the fastest consumption method. Returns `Some(tokens)` if successful, + /// or `None` if insufficient tokens are available. + /// + /// For wait time estimates when tokens are unavailable, use [`try_consume`](Self::try_consume). + /// + /// # Arguments + /// + /// * `to_consume` - Number of tokens to consume + /// * `clock` - Clock implementation to use for timing + /// * `limit` - Rate and burst configuration for the bucket + /// + /// # Returns + /// + /// * `Some(Tokens)` - Successfully consumed tokens + /// * `None` - Insufficient tokens available + pub fn consume( + &self, + to_consume: impl Into, + clock: &C, + limit: &Limit, + ) -> Option { + let now = clock.now(); + let to_consume: NonZeroU32 = to_consume.into(); + let to_consume: f64 = to_consume.get() as f64; + + let consumed = self.bucket.consume(limit.rate, limit.burst, now, |avail| { + if avail < to_consume { 0.0 } else { to_consume } + }); + Tokens::new_checked(consumed) + } + + /// Attempts to consume exactly one token. + /// + /// Convenience method equivalent to `consume(1)`. + /// + /// # Returns + /// + /// * `Some(Tokens)` - Successfully consumed one token + /// * `None` - No tokens available + pub fn consume_one(&self, clock: &C, limit: &Limit) -> Option { + self.consume(NonZeroU32::new(1u32).unwrap(), clock, limit) + } + + /// Attempts to consume one token with wait time information. + /// + /// Convenience method equivalent to `try_consume(1)`. + /// + /// # Returns + /// + /// * `Ok(Tokens)` - Successfully consumed one token + /// * `Err(RateLimited)` - Rate limited with suggested wait time + pub fn try_consume_one(&self, clock: &C, limit: &Limit) -> Result { + self.try_consume(NonZeroU32::new(1u32).unwrap(), clock, limit) + } + + /// Attempts to consume tokens with detailed rate limiting information. + /// + /// Unlike [`consume`](Self::consume), this method provides an estimate of how long + /// to wait before retrying when tokens are unavailable. + /// + /// # Arguments + /// + /// * `to_consume` - Number of tokens to consume + /// * `clock` - Clock implementation to use for timing + /// * `limit` - Rate and burst configuration for the bucket + /// + /// # Returns + /// + /// * `Ok(Tokens)` - Successfully consumed tokens + /// * `Err(RateLimited)` - Rate limited with suggested retry time + pub fn try_consume( + &self, + to_consume: impl Into, + clock: &C, + limit: &Limit, + ) -> Result { + let to_consume: NonZeroU32 = to_consume.into(); + let to_consume: f64 = to_consume.get() as f64; + let now = clock.now(); + let consumed = self.bucket.consume2(limit.rate, limit.burst, now, |avail| { + if avail < to_consume { 0.0 } else { to_consume } + }); + match consumed { + TokenAcquisition::Acquired(consumed) => Ok(Tokens::new_unchecked(consumed)), + TokenAcquisition::ZeroedAt(zero_time) => { + let est_time = zero_time - now + to_consume / limit.rate; + debug_assert!(est_time >= 0.0); + Err(RateLimited { + earliest_retry_time: Nanos::from_secs_f64_unchecked(est_time), + }) + } + } + } + + /// Consumes up to the requested number of tokens, returning whatever is available. + /// + /// This method will consume as many tokens as possible up to the requested amount, + /// without waiting. Returns `None` if no tokens are available. + /// + /// # Arguments + /// + /// * `to_consume` - Maximum number of tokens to consume + /// * `clock` - Clock implementation to use for timing + /// * `limit` - Rate and burst configuration for the bucket + /// + /// + /// # Returns + /// + /// * `Some(Tokens)` - Number of tokens actually consumed (may be less than requested) + /// * `None` - No tokens available + pub fn saturating_consume( + &self, + to_consume: impl Into, + clock: &C, + limit: &Limit, + ) -> Option { + let now = clock.now(); + let to_consume: NonZeroU32 = to_consume.into(); + let to_consume: f64 = to_consume.get() as f64; + Tokens::new_checked(self.saturating_consume_inner(to_consume, now, limit)) + } + + /// Returns unused tokens to the bucket or manually adds tokens. + /// + /// This operation respects the bucket's burst capacity and will not cause overflow. + /// Useful for returning tokens from cancelled operations. + /// + /// # Arguments + /// + /// * `tokens` - Number of tokens to add back to the bucket + /// * `limit` - Rate and burst configuration for the bucket + pub fn add_tokens(&self, tokens: impl Into, limit: &Limit) { + let tokens = tokens.into(); + debug_assert!(tokens > 0.0); + self.bucket.return_tokens(tokens, limit.rate); + } + + /// Consumes tokens by borrowing from future capacity. + /// + /// This allows consuming more tokens than currently available by going into "debt". + /// The bucket will need time to replenish before more tokens can be consumed. + /// + /// # Arguments + /// + /// * `to_consume` - Number of tokens to consume + /// * `clock` - Clock implementation to use for timing + /// * `limit` - Rate and burst configuration for the bucket + /// + /// # Returns + /// + /// * `Ok(None)` - Tokens consumed immediately without borrowing + /// * `Ok(Some(duration))` - Tokens consumed with borrowing; wait time until debt is paid + /// * `Err(ExceededBurstCapacity)` - Cannot borrow more than burst capacity + pub fn consume_with_borrow( + &self, + to_consume: impl Into, + clock: &C, + limit: &Limit, + ) -> Result, ExceededBurstCapacity> { + let now = clock.now(); + let to_consume: NonZeroU32 = to_consume.into(); + let mut to_consume: f64 = to_consume.get() as f64; + if unlikely(limit.burst < to_consume) { + return Err(ExceededBurstCapacity); + } + while to_consume > 0.0 { + let consumed = self.saturating_consume_inner(to_consume, now, limit); + if consumed > 0.0 { + to_consume -= consumed; + } else { + self.bucket.return_tokens(-to_consume, limit.rate); + let debt_paid = self.bucket.time_when_bucket(limit.rate, 0.0); + let nap_time = (debt_paid - now).max(0.0); + return Ok(Nanos::new_checked(nap_time)); + } + } + Ok(None) + } + + /// Consumes tokens with borrowing, limited to burst capacity. + /// + /// Similar to [`consume_with_borrow`](Self::consume_with_borrow) but automatically + /// limits the request to the burst capacity instead of returning an error. + /// + /// # Arguments + /// + /// * `to_consume` - Number of tokens to consume (capped at burst capacity) + /// * `clock` - Clock implementation to use for timing + /// * `limit` - Rate and burst configuration for the bucket + /// + /// # Returns + /// + /// A tuple of: + /// * `Option` - Number of tokens consumed (None if no borrowing occurred) + /// * `Duration` - Wait time until the debt is paid (zero if no borrowing) + pub fn saturating_consume_with_borrow( + &self, + to_consume: impl Into, + clock: &C, + limit: &Limit, + ) -> (Option, Duration) { + let now = clock.now(); + let to_consume: NonZeroU32 = to_consume.into(); + let mut to_consume: f64 = to_consume.get() as f64; + to_consume = to_consume.min(limit.burst); + let actual_to_be_consumed = to_consume; + while to_consume > 0.0 { + let consumed = self.saturating_consume_inner(to_consume, now, limit); + if consumed > 0.0 { + to_consume -= consumed; + } else { + self.bucket.return_tokens(-to_consume, limit.rate); + let debt_paid = self.bucket.time_when_bucket(limit.rate, 0.0); + let nap_time = (debt_paid - now).max(0.0); + return ( + Tokens::new_checked(actual_to_be_consumed), + Duration::from_secs_f64(nap_time), + ); + } + } + (None, Duration::ZERO) + } + + /// Returns the number of tokens currently available for consumption. + /// + /// This value is always non-negative. If the bucket is in debt from borrowing, + /// this returns zero. + /// + /// # Returns + /// + /// Number of tokens available for immediate consumption + pub fn available(&self, clock: &C, limit: &Limit) -> f64 { + self.balance(clock, limit).max(0.0) + } + + /// Returns the current token balance, which may be negative if in debt. + /// + /// Unlike [`available`](Self::available), this can return negative values + /// when tokens have been borrowed from future capacity. + /// + /// # Returns + /// + /// Current token balance (negative indicates debt) + pub fn balance(&self, clock: &C, limit: &Limit) -> f64 { + self.bucket.balance(limit.rate, limit.burst, clock.now()) + } + + pub fn balance_at(&self, time: f64, limit: &Limit) -> f64 { + self.bucket.balance(limit.rate, limit.burst, time) + } + + #[inline] + fn saturating_consume_inner(&self, to_consume: f64, now: f64, limit: &Limit) -> f64 { + self.bucket.consume(limit.rate, limit.burst, now, |avail| { + avail.max(0.0).min(to_consume) + }) + } +} diff --git a/src/storage.rs b/src/storage.rs index aec17e7..09c2f30 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -31,7 +31,7 @@ pub(crate) enum TokenAcquisition { /// Heavily inspired by folly's TokenBucket algorithm. #[derive(Debug, Clone)] pub(crate) struct TokenBucketStorage { - inner: S, + pub(crate) inner: S, } impl TokenBucketStorage { diff --git a/src/storage/atomic.rs b/src/storage/atomic.rs index 72ab925..669b198 100644 --- a/src/storage/atomic.rs +++ b/src/storage/atomic.rs @@ -60,7 +60,7 @@ impl Debug for AtomicF64 { /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); -/// let bucket = TokenBucket::::from_parts(limit, StdClock::default()); +/// let bucket = TokenBucket::::new(limit, StdClock); /// ``` #[derive(Debug)] pub struct AtomicStorage(AtomicF64); @@ -101,7 +101,7 @@ impl TimeStorage for AtomicStorage { /// /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); /// let clock = Arc::new(ManualClock::new(0.0)); -/// let bucket1 = TokenBucket::::from_parts(limit, Arc::clone(&clock)); +/// let bucket1 = TokenBucket::::new(limit, Arc::clone(&clock)); /// let bucket2 = bucket1.clone(); // Shares the same token state /// ``` #[derive(Debug, Clone)] diff --git a/src/storage/local.rs b/src/storage/local.rs index 65d0557..18e7fbc 100644 --- a/src/storage/local.rs +++ b/src/storage/local.rs @@ -20,7 +20,7 @@ use super::TimeStorage; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); -/// let bucket = TokenBucket::::from_parts(limit, StdClock::default()); +/// let bucket = TokenBucket::::new(limit, StdClock); /// ``` #[derive(Debug)] pub struct LocalStorage(Cell); diff --git a/src/storage/padded_atomic.rs b/src/storage/padded_atomic.rs index 3876869..b054ad1 100644 --- a/src/storage/padded_atomic.rs +++ b/src/storage/padded_atomic.rs @@ -20,12 +20,12 @@ use super::cache_padded::CachePadded; /// # Examples /// /// ```rust -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{SharedTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// // PaddedAtomicStorage is the default storage type /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); -/// let bucket = TokenBucket::new(limit); +/// let bucket = SharedTokenBucket::new(limit, StdClock); /// ``` pub struct PaddedAtomicStorage(CachePadded); @@ -65,12 +65,12 @@ impl TimeStorage for PaddedAtomicStorage { /// # Examples /// /// ```rust -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{SharedTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// // PaddedAtomicStorage is the default storage type /// let limit = Limit::per_second(NonZeroU32::new(100).unwrap()); -/// let bucket = TokenBucket::new(limit); +/// let bucket = SharedTokenBucket::new(limit, StdClock); /// ``` #[derive(Clone)] pub struct PaddedAtomicSharedStorage(Arc>); diff --git a/src/tokens.rs b/src/tokens.rs index 3558b4f..ddee2ac 100644 --- a/src/tokens.rs +++ b/src/tokens.rs @@ -8,11 +8,11 @@ use std::num::NonZero; /// # Examples /// /// ```rust -/// use gardal::{TokenBucket, Limit}; +/// use gardal::{LocalTokenBucket, Limit, StdClock}; /// use std::num::NonZeroU32; /// /// let limit = Limit::per_second(NonZeroU32::new(10).unwrap()); -/// let bucket = TokenBucket::new(limit); +/// let bucket = LocalTokenBucket::new(limit, StdClock); /// /// if let Some(tokens) = bucket.consume_one() { /// println!("Consumed {} tokens", tokens.as_u64());