From 4479e1f65a8f8d3d7928960194af4b973b5b2636 Mon Sep 17 00:00:00 2001 From: TennyZhuang <9161438+TennyZhuang@users.noreply.github.com> Date: Tue, 28 Apr 2026 02:58:05 +0800 Subject: [PATCH 1/2] fix(layers/throttle): limit read bandwidth --- core/Cargo.lock | 1 + core/layers/throttle/Cargo.toml | 1 + core/layers/throttle/src/lib.rs | 101 ++++++++++++++++++++++++-------- 3 files changed, 80 insertions(+), 23 deletions(-) diff --git a/core/Cargo.lock b/core/Cargo.lock index 4353fed27917..bd008efe0e65 100644 --- a/core/Cargo.lock +++ b/core/Cargo.lock @@ -6490,6 +6490,7 @@ version = "0.56.0" dependencies = [ "governor", "opendal-core", + "tokio", ] [[package]] diff --git a/core/layers/throttle/Cargo.toml b/core/layers/throttle/Cargo.toml index 9fecc533d1b7..0192b111cd19 100644 --- a/core/layers/throttle/Cargo.toml +++ b/core/layers/throttle/Cargo.toml @@ -36,3 +36,4 @@ opendal-core = { path = "../../core", version = "0.56.0", default-features = fal [dev-dependencies] opendal-core = { path = "../../core", version = "0.56.0" } +tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/core/layers/throttle/src/lib.rs b/core/layers/throttle/src/lib.rs index 3e8d7b644d64..46a263027e67 100644 --- a/core/layers/throttle/src/lib.rs +++ b/core/layers/throttle/src/lib.rs @@ -166,36 +166,39 @@ impl ThrottleWrapper { } } +async fn wait_for_quota(limiter: &SharedRateLimiter, len: usize) -> Result<()> { + if len == 0 { + return Ok(()); + } + + if len > u32::MAX as usize { + return Err(Error::new( + ErrorKind::RateLimited, + "request size exceeds throttle quota capacity", + )); + } + + let buf_length = NonZeroU32::new(len as u32).expect("len is non-zero so NonZeroU32 must exist"); + + limiter.until_n_ready(buf_length).await.map_err(|_| { + Error::new( + ErrorKind::RateLimited, + "burst size is smaller than the request size", + ) + }) +} + impl oio::Read for ThrottleWrapper { async fn read(&mut self) -> Result { - self.inner.read().await + let bs = self.inner.read().await?; + wait_for_quota(&self.limiter, bs.len()).await?; + Ok(bs) } } impl oio::Write for ThrottleWrapper { async fn write(&mut self, bs: Buffer) -> Result<()> { - let len = bs.len(); - if len == 0 { - return self.inner.write(bs).await; - } - - if len > u32::MAX as usize { - return Err(Error::new( - ErrorKind::RateLimited, - "request size exceeds throttle quota capacity", - )); - } - - let buf_length = - NonZeroU32::new(len as u32).expect("len is non-zero so NonZeroU32 must exist"); - - self.limiter.until_n_ready(buf_length).await.map_err(|_| { - Error::new( - ErrorKind::RateLimited, - "burst size is smaller than the request size", - ) - })?; - + wait_for_quota(&self.limiter, bs.len()).await?; self.inner.write(bs).await } @@ -207,3 +210,55 @@ impl oio::Write for ThrottleWrapper { self.inner.close().await } } + +#[cfg(test)] +mod tests { + use super::*; + use opendal_core::raw::oio::Read; + + struct SingleRead { + data: Option, + } + + impl oio::Read for SingleRead { + async fn read(&mut self) -> Result { + Ok(self.data.take().unwrap_or_default()) + } + } + + fn new_limiter(bandwidth: u32, burst: u32) -> SharedRateLimiter { + Arc::new(RateLimiter::direct( + Quota::per_second(NonZeroU32::new(bandwidth).unwrap()) + .allow_burst(NonZeroU32::new(burst).unwrap()), + )) + } + + #[tokio::test] + async fn read_allows_chunks_within_burst() { + let mut reader = ThrottleWrapper::new( + SingleRead { + data: Some(Buffer::from(vec![0; 2])), + }, + new_limiter(2, 2), + ); + + let bs = reader.read().await.expect("read should pass throttle"); + assert_eq!(bs.len(), 2); + } + + #[tokio::test] + async fn read_rejects_chunks_larger_than_burst() { + let mut reader = ThrottleWrapper::new( + SingleRead { + data: Some(Buffer::from(vec![0; 2])), + }, + new_limiter(1, 1), + ); + + let err = reader + .read() + .await + .expect_err("chunk larger than burst should fail"); + assert_eq!(err.kind(), ErrorKind::RateLimited); + } +} From 103cdcaa1e716c71a8a817653f02e09426477462 Mon Sep 17 00:00:00 2001 From: TennyZhuang <9161438+TennyZhuang@users.noreply.github.com> Date: Tue, 28 Apr 2026 12:26:20 +0800 Subject: [PATCH 2/2] fix(layers/throttle): precheck bounded read quota --- core/layers/throttle/src/lib.rs | 130 +++++++++++++++++++++----------- 1 file changed, 85 insertions(+), 45 deletions(-) diff --git a/core/layers/throttle/src/lib.rs b/core/layers/throttle/src/lib.rs index 46a263027e67..facb4cbc66b5 100644 --- a/core/layers/throttle/src/lib.rs +++ b/core/layers/throttle/src/lib.rs @@ -115,7 +115,7 @@ pub struct ThrottleAccessor { impl LayeredAccess for ThrottleAccessor { type Inner = A; - type Reader = ThrottleWrapper; + type Reader = A::Reader; type Writer = ThrottleWrapper; type Lister = A::Lister; type Deleter = A::Deleter; @@ -125,12 +125,11 @@ impl LayeredAccess for ThrottleAccessor { } async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> { - let limiter = self.rate_limiter.clone(); + if let Some(size) = args.range().size() { + wait_for_quota(&self.rate_limiter, size).await?; + } - self.inner - .read(path, args) - .await - .map(|(rp, r)| (rp, ThrottleWrapper::new(r, limiter))) + self.inner.read(path, args).await } async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> { @@ -166,12 +165,12 @@ impl ThrottleWrapper { } } -async fn wait_for_quota(limiter: &SharedRateLimiter, len: usize) -> Result<()> { +async fn wait_for_quota(limiter: &SharedRateLimiter, len: u64) -> Result<()> { if len == 0 { return Ok(()); } - if len > u32::MAX as usize { + if len > u32::MAX as u64 { return Err(Error::new( ErrorKind::RateLimited, "request size exceeds throttle quota capacity", @@ -188,17 +187,9 @@ async fn wait_for_quota(limiter: &SharedRateLimiter, len: usize) -> Result<()> { }) } -impl oio::Read for ThrottleWrapper { - async fn read(&mut self) -> Result { - let bs = self.inner.read().await?; - wait_for_quota(&self.limiter, bs.len()).await?; - Ok(bs) - } -} - impl oio::Write for ThrottleWrapper { async fn write(&mut self, bs: Buffer) -> Result<()> { - wait_for_quota(&self.limiter, bs.len()).await?; + wait_for_quota(&self.limiter, bs.len() as u64).await?; self.inner.write(bs).await } @@ -214,7 +205,8 @@ impl oio::Write for ThrottleWrapper { #[cfg(test)] mod tests { use super::*; - use opendal_core::raw::oio::Read; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; struct SingleRead { data: Option, @@ -226,39 +218,87 @@ mod tests { } } - fn new_limiter(bandwidth: u32, burst: u32) -> SharedRateLimiter { - Arc::new(RateLimiter::direct( - Quota::per_second(NonZeroU32::new(bandwidth).unwrap()) - .allow_burst(NonZeroU32::new(burst).unwrap()), - )) + #[derive(Debug)] + struct MockService { + read_count: Arc, + } + + impl Access for MockService { + type Reader = oio::Reader; + type Writer = oio::Writer; + type Lister = oio::Lister; + type Deleter = oio::Deleter; + + fn info(&self) -> Arc { + let info = AccessorInfo::default(); + info.set_native_capability(Capability { + read: true, + ..Default::default() + }); + + info.into() + } + + async fn read(&self, _: &str, _: OpRead) -> Result<(RpRead, Self::Reader)> { + self.read_count.fetch_add(1, Ordering::SeqCst); + Ok((RpRead::new(), Box::new(SingleRead { data: None }))) + } + } + + fn new_throttle_accessor( + bandwidth: u32, + burst: u32, + ) -> (ThrottleAccessor, Arc) { + let read_count = Arc::new(AtomicUsize::new(0)); + let accessor = ThrottleLayer::new(bandwidth, burst).layer(MockService { + read_count: read_count.clone(), + }); + + (accessor, read_count) } #[tokio::test] - async fn read_allows_chunks_within_burst() { - let mut reader = ThrottleWrapper::new( - SingleRead { - data: Some(Buffer::from(vec![0; 2])), - }, - new_limiter(2, 2), - ); - - let bs = reader.read().await.expect("read should pass throttle"); - assert_eq!(bs.len(), 2); + async fn bounded_read_allows_request_within_burst() { + let (accessor, read_count) = new_throttle_accessor(2, 2); + + LayeredAccess::read( + &accessor, + "path", + OpRead::new().with_range(BytesRange::new(0, Some(2))), + ) + .await + .expect("bounded read within burst should pass throttle"); + + assert_eq!(read_count.load(Ordering::SeqCst), 1); } #[tokio::test] - async fn read_rejects_chunks_larger_than_burst() { - let mut reader = ThrottleWrapper::new( - SingleRead { - data: Some(Buffer::from(vec![0; 2])), - }, - new_limiter(1, 1), - ); - - let err = reader - .read() - .await - .expect_err("chunk larger than burst should fail"); + async fn bounded_read_rejects_before_sending_request() { + let (accessor, read_count) = new_throttle_accessor(1, 1); + + let err = match LayeredAccess::read( + &accessor, + "path", + OpRead::new().with_range(BytesRange::new(0, Some(2))), + ) + .await + { + Ok(_) => panic!("bounded read larger than burst should fail"), + Err(err) => err, + }; + assert_eq!(err.kind(), ErrorKind::RateLimited); + assert_eq!(read_count.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn unbounded_read_is_forwarded_without_precheck() { + let (accessor, read_count) = new_throttle_accessor(1, 1); + + LayeredAccess::read(&accessor, "path", OpRead::new()) + .await + .expect("unbounded read size is unknown before request"); + + assert_eq!(read_count.load(Ordering::SeqCst), 1); } }