fix(layers/throttle): precheck bounded read quota

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

core/layers/throttle/src/lib.rs

@@ -115,7 +115,7 @@ pub struct ThrottleAccessor<A: Access> {
 
 impl<A: Access> LayeredAccess for ThrottleAccessor<A> {
     type Inner = A;
-    type Reader = ThrottleWrapper<A::Reader>;
+    type Reader = A::Reader;
     type Writer = ThrottleWrapper<A::Writer>;
     type Lister = A::Lister;
     type Deleter = A::Deleter;
@@ -126,12 +126,11 @@ impl<A: Access> LayeredAccess for ThrottleAccessor<A> {
     }
 
     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<R> ThrottleWrapper<R> {
     }
 }
 
-impl<R: oio::Read> oio::Read for ThrottleWrapper<R> {
-    async fn read(&mut self) -> Result<Buffer> {
-        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<R: oio::Write> oio::Write for ThrottleWrapper<R> {
     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<Metadata> {
+        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<Buffer>,
+    }
+
+    impl oio::Read for SingleRead {
+        async fn read(&mut self) -> Result<Buffer> {
+            Ok(self.data.take().unwrap_or_default())
         }
+    }
+
+    #[derive(Debug)]
+    struct MockService {
+        read_count: Arc<AtomicUsize>,
+    }
+
+    impl Access for MockService {
+        type Reader = oio::Reader;
+        type Writer = oio::Writer;
+        type Lister = oio::Lister;
+        type Deleter = oio::Deleter;
+
+        fn info(&self) -> Arc<AccessorInfo> {
+            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<MockService>, Arc<AtomicUsize>) {
+        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<Metadata> {
-        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);
     }
 }