feat(storage): add direct chunk GET fast path (#2351)

Signed-off-by: houseme <housemecn@gmail.com>
Co-authored-by: heihutu <heihutu@gmail.com>
Co-authored-by: cxymds <Cxymds@qq.com>

crates/rio/src/lib.rs

@@ -15,6 +15,10 @@
 // Default encryption block size - aligned with system default read buffer size (1MB)
 pub const DEFAULT_ENCRYPTION_BLOCK_SIZE: usize = 1024 * 1024;
 
+use std::future::Future;
+use std::pin::Pin;
+use tokio::io::AsyncReadExt;
+
 macro_rules! delegate_reader_capabilities_generic {
     ($name:ident<$inner_ty:ident>, $inner:ident) => {
         impl<$inner_ty> crate::EtagResolvable for $name<$inner_ty>
@@ -114,14 +118,39 @@ pub use compress_index::{Index, TryGetIndex};
 
 mod etag;
 
+pub type BoxReadBlockFuture<'a> = Pin<Box<dyn Future<Output = std::io::Result<usize>> + Send + 'a>>;
+
+pub trait BlockReadable {
+    fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a>;
+}
+
+fn read_block_via_async_read<'a, R>(reader: &'a mut R, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a>
+where
+    R: tokio::io::AsyncRead + Unpin + Send + Sync + 'a,
+{
+    Box::pin(async move {
+        let mut total = 0;
+
+        while total < buf.len() {
+            match reader.read(&mut buf[total..]).await {
+                Ok(0) => return Ok(total),
+                Ok(n) => total += n,
+                Err(err) => return Err(err),
+            }
+        }
+
+        Ok(total)
+    })
+}
+
 pub trait ReadStream: tokio::io::AsyncRead + Unpin + Send + Sync {}
 impl<T> ReadStream for T where T: tokio::io::AsyncRead + Unpin + Send + Sync {}
 
 pub trait ReaderCapabilities: EtagResolvable + HashReaderDetector + TryGetIndex {}
 impl<T> ReaderCapabilities for T where T: EtagResolvable + HashReaderDetector + TryGetIndex {}
 
-pub trait Reader: ReadStream + ReaderCapabilities {}
-impl<T> Reader for T where T: ReadStream + ReaderCapabilities {}
+pub trait Reader: ReadStream + ReaderCapabilities + BlockReadable {}
+impl<T> Reader for T where T: ReadStream + ReaderCapabilities + BlockReadable {}
 
 pub type DynReader = Box<dyn Reader>;
 
@@ -154,6 +183,42 @@ pub trait HashReaderDetector {
     }
 }
 
+impl<R> BlockReadable for crate::WarpReader<R>
+where
+    R: tokio::io::AsyncRead + Unpin + Send + Sync,
+{
+    fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> {
+        read_block_via_async_read(self, buf)
+    }
+}
+
+impl<R> BlockReadable for tokio::io::BufReader<R>
+where
+    R: tokio::io::AsyncRead + Unpin + Send + Sync,
+{
+    fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> {
+        read_block_via_async_read(self, buf)
+    }
+}
+
+impl<R> BlockReadable for crate::LimitReader<R>
+where
+    R: Reader,
+{
+    fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> {
+        read_block_via_async_read(self, buf)
+    }
+}
+
+impl<R> BlockReadable for crate::DecryptReader<R>
+where
+    R: Reader,
+{
+    fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> {
+        read_block_via_async_read(self, buf)
+    }
+}
+
 pub fn boxed_reader<R>(reader: R) -> DynReader
 where
     R: Reader + 'static,
@@ -198,3 +263,74 @@ where
         self.as_ref().try_get_index()
     }
 }
+
+impl<T> BlockReadable for Box<T>
+where
+    T: BlockReadable + ?Sized,
+{
+    fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> {
+        self.as_mut().read_block(buf)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::VecDeque;
+    use std::io::{self, ErrorKind};
+    use std::pin::Pin;
+    use std::task::{Context, Poll};
+    use tokio::io::{AsyncRead, ReadBuf};
+
+    enum ReadStep {
+        Data(Vec<u8>),
+        Error(ErrorKind),
+        Eof,
+    }
+
+    struct StepReader {
+        steps: VecDeque<ReadStep>,
+    }
+
+    impl StepReader {
+        fn new(steps: impl IntoIterator<Item = ReadStep>) -> Self {
+            Self {
+                steps: steps.into_iter().collect(),
+            }
+        }
+    }
+
+    impl AsyncRead for StepReader {
+        fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
+            match self.steps.pop_front().unwrap_or(ReadStep::Eof) {
+                ReadStep::Data(data) => {
+                    buf.put_slice(&data);
+                    Poll::Ready(Ok(()))
+                }
+                ReadStep::Error(kind) => Poll::Ready(Err(io::Error::new(kind, "synthetic read failure"))),
+                ReadStep::Eof => Poll::Ready(Ok(())),
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_read_block_via_async_read_preserves_midstream_error_kind() {
+        let reader = StepReader::new([ReadStep::Data(b"ab".to_vec()), ReadStep::Error(ErrorKind::ConnectionReset)]);
+        let mut reader = WarpReader::new(reader);
+        let mut buf = [0_u8; 4];
+
+        let err = reader.read_block(&mut buf).await.unwrap_err();
+
+        assert_eq!(err.kind(), ErrorKind::ConnectionReset);
+    }
+
+    #[tokio::test]
+    async fn test_read_block_via_async_read_returns_zero_on_initial_eof() {
+        let mut reader = WarpReader::new(StepReader::new([ReadStep::Eof]));
+        let mut buf = [0_u8; 4];
+
+        let n = reader.read_block(&mut buf).await.unwrap();
+
+        assert_eq!(n, 0);
+    }
+}