diff --git a/core/Cargo.lock b/core/Cargo.lock index c5e953b134f2..d0b40f474219 100644 --- a/core/Cargo.lock +++ b/core/Cargo.lock @@ -6637,6 +6637,7 @@ version = "0.57.0" dependencies = [ "governor", "opendal-core", + "tokio", ] [[package]] diff --git a/core/layers/throttle/Cargo.toml b/core/layers/throttle/Cargo.toml index 869b64c52497..3efce9c849e0 100644 --- a/core/layers/throttle/Cargo.toml +++ b/core/layers/throttle/Cargo.toml @@ -36,3 +36,4 @@ opendal-core = { path = "../../core", version = "0.57.0", default-features = fal [dev-dependencies] opendal-core = { path = "../../core", version = "0.57.0" } +tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/core/layers/throttle/src/lib.rs b/core/layers/throttle/src/lib.rs index df0ffcd0ac82..9f116377493e 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; @@ -126,12 +126,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)> { @@ -177,44 +176,140 @@ impl ThrottleWrapper { } } -impl oio::Read for ThrottleWrapper { - async fn read(&mut self) -> Result { - self.inner.read().await +async fn wait_for_quota(limiter: &SharedRateLimiter, len: u64) -> Result<()> { + if len == 0 { + return Ok(()); + } + + if len > u32::MAX as u64 { + 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::Write for ThrottleWrapper { async fn write(&mut self, bs: Buffer) -> Result<()> { - let len = bs.len(); - if len == 0 { - return self.inner.write(bs).await; + wait_for_quota(&self.limiter, bs.len() as u64).await?; + self.inner.write(bs).await + } + + async fn close(&mut self) -> Result { + self.inner.close().await + } + + async fn abort(&mut self) -> Result<()> { + self.inner.abort().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + struct SingleRead { + data: Option, + } + + impl oio::Read for SingleRead { + async fn read(&mut self) -> Result { + Ok(self.data.take().unwrap_or_default()) } + } + + #[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() + }); - if len > u32::MAX as usize { - return Err(Error::new( - ErrorKind::RateLimited, - "request size exceeds throttle quota capacity", - )); + info.into() } - let buf_length = - NonZeroU32::new(len as u32).expect("len is non-zero so NonZeroU32 must exist"); + 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 }))) + } + } - self.limiter.until_n_ready(buf_length).await.map_err(|_| { - Error::new( - ErrorKind::RateLimited, - "burst size is smaller than the request size", - ) - })?; + 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(), + }); - self.inner.write(bs).await + (accessor, read_count) } - async fn close(&mut self) -> Result { - self.inner.close().await + #[tokio::test] + 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); } - async fn abort(&mut self) -> Result<()> { - self.inner.abort().await + #[tokio::test] + 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); } }