feat: add compress support

crates/rio/src/compress_reader.rs

@@ -1,13 +1,18 @@
-use crate::compress::{CompressionAlgorithm, compress_block, decompress_block};
+use crate::compress_index::{Index, TryGetIndex};
 use crate::{EtagResolvable, HashReaderDetector};
 use crate::{HashReaderMut, Reader};
 use pin_project_lite::pin_project;
+use rustfs_utils::compress::{CompressionAlgorithm, compress_block, decompress_block};
 use rustfs_utils::{put_uvarint, put_uvarint_len, uvarint};
 use std::io::{self};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::io::{AsyncRead, ReadBuf};
 
+const COMPRESS_TYPE_COMPRESSED: u8 = 0x00;
+const COMPRESS_TYPE_UNCOMPRESSED: u8 = 0x01;
+const COMPRESS_TYPE_END: u8 = 0xFF;
+
 pin_project! {
     #[derive(Debug)]
     /// A reader wrapper that compresses data on the fly using DEFLATE algorithm.
@@ -19,6 +24,11 @@ pin_project! {
         done: bool,
         block_size: usize,
         compression_algorithm: CompressionAlgorithm,
+        index: Index,
+        written: usize,
+        uncomp_written: usize,
+        temp_buffer: Vec<u8>,
+        temp_pos: usize,
     }
 }
 
@@ -34,6 +44,11 @@ where
             done: false,
             compression_algorithm,
             block_size: 1 << 20, // Default 1MB
+            index: Index::new(),
+            written: 0,
+            uncomp_written: 0,
+            temp_buffer: Vec::with_capacity(1 << 20), // 预分配1MB容量
+            temp_pos: 0,
         }
     }
 
@@ -46,10 +61,24 @@ where
             done: false,
             compression_algorithm,
             block_size,
+            index: Index::new(),
+            written: 0,
+            uncomp_written: 0,
+            temp_buffer: Vec::with_capacity(block_size),
+            temp_pos: 0,
         }
     }
 }
 
+impl<R> TryGetIndex for CompressReader<R>
+where
+    R: Reader,
+{
+    fn try_get_index(&self) -> Option<&Index> {
+        Some(&self.index)
+    }
+}
+
 impl<R> AsyncRead for CompressReader<R>
 where
     R: AsyncRead + Unpin + Send + Sync,
@@ -72,69 +101,99 @@ where
             return Poll::Ready(Ok(()));
         }
 
-        // Read from inner, only read block_size bytes each time
-        let mut temp = vec![0u8; *this.block_size];
-        let mut temp_buf = ReadBuf::new(&mut temp);
-        match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
-            Poll::Pending => Poll::Pending,
-            Poll::Ready(Ok(())) => {
-                let n = temp_buf.filled().len();
-                if n == 0 {
-                    // EOF, write end header
-                    let mut header = [0u8; 8];
-                    header[0] = 0xFF;
-                    *this.buffer = header.to_vec();
-                    *this.pos = 0;
-                    *this.done = true;
-                    let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
-                    buf.put_slice(&this.buffer[..to_copy]);
-                    *this.pos += to_copy;
-                    Poll::Ready(Ok(()))
-                } else {
-                    let uncompressed_data = &temp_buf.filled()[..n];
+        // 如果临时缓冲区未满，继续读取数据
+        while this.temp_buffer.len() < *this.block_size {
+            let remaining = *this.block_size - this.temp_buffer.len();
+            let mut temp = vec![0u8; remaining];
+            let mut temp_buf = ReadBuf::new(&mut temp);
 
-                    let crc = crc32fast::hash(uncompressed_data);
-                    let compressed_data = compress_block(uncompressed_data, *this.compression_algorithm);
-
-                    let uncompressed_len = n;
-                    let compressed_len = compressed_data.len();
-                    let int_len = put_uvarint_len(uncompressed_len as u64);
-
-                    let len = compressed_len + int_len + 4; // 4 bytes for CRC32
-
-                    // Header: 8 bytes
-                    // 0: type (0 = compressed, 1 = uncompressed, 0xFF = end)
-                    // 1-3: length (little endian u24)
-                    // 4-7: crc32 (little endian u32)
-                    let mut header = [0u8; 8];
-                    header[0] = 0x00; // 0 = compressed
-                    header[1] = (len & 0xFF) as u8;
-                    header[2] = ((len >> 8) & 0xFF) as u8;
-                    header[3] = ((len >> 16) & 0xFF) as u8;
-                    header[4] = (crc & 0xFF) as u8;
-                    header[5] = ((crc >> 8) & 0xFF) as u8;
-                    header[6] = ((crc >> 16) & 0xFF) as u8;
-                    header[7] = ((crc >> 24) & 0xFF) as u8;
-
-                    // Combine header(4+4) + uncompressed_len + compressed
-                    let mut out = Vec::with_capacity(len + 4);
-                    out.extend_from_slice(&header);
-
-                    let mut uncompressed_len_buf = vec![0u8; int_len];
-                    put_uvarint(&mut uncompressed_len_buf, uncompressed_len as u64);
-                    out.extend_from_slice(&uncompressed_len_buf);
-
-                    out.extend_from_slice(&compressed_data);
-
-                    *this.buffer = out;
-                    *this.pos = 0;
-                    let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
-                    buf.put_slice(&this.buffer[..to_copy]);
-                    *this.pos += to_copy;
-                    Poll::Ready(Ok(()))
+            match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
+                Poll::Pending => {
+                    // 如果临时缓冲区为空，返回 Pending
+                    if this.temp_buffer.is_empty() {
+                        return Poll::Pending;
+                    }
+                    // 否则继续处理已读取的数据
+                    break;
                 }
+                Poll::Ready(Ok(())) => {
+                    let n = temp_buf.filled().len();
+                    if n == 0 {
+                        // EOF
+                        if this.temp_buffer.is_empty() {
+                            // // 如果没有累积的数据，写入结束标记
+                            // let mut header = [0u8; 8];
+                            // header[0] = 0xFF;
+                            // *this.buffer = header.to_vec();
+                            // *this.pos = 0;
+                            // *this.done = true;
+                            // let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
+                            // buf.put_slice(&this.buffer[..to_copy]);
+                            // *this.pos += to_copy;
+                            return Poll::Ready(Ok(()));
+                        }
+                        // 有累积的数据，处理它
+                        break;
+                    }
+                    this.temp_buffer.extend_from_slice(&temp[..n]);
+                }
+                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
             }
-            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+        }
+
+        // 处理累积的数据
+        if !this.temp_buffer.is_empty() {
+            let uncompressed_data = &this.temp_buffer;
+            let crc = crc32fast::hash(uncompressed_data);
+            let compressed_data = compress_block(uncompressed_data, *this.compression_algorithm);
+
+            let uncompressed_len = uncompressed_data.len();
+            let compressed_len = compressed_data.len();
+            let int_len = put_uvarint_len(uncompressed_len as u64);
+
+            let len = compressed_len + int_len;
+            let header_len = 8;
+
+            let mut header = [0u8; 8];
+            header[0] = COMPRESS_TYPE_COMPRESSED;
+            header[1] = (len & 0xFF) as u8;
+            header[2] = ((len >> 8) & 0xFF) as u8;
+            header[3] = ((len >> 16) & 0xFF) as u8;
+            header[4] = (crc & 0xFF) as u8;
+            header[5] = ((crc >> 8) & 0xFF) as u8;
+            header[6] = ((crc >> 16) & 0xFF) as u8;
+            header[7] = ((crc >> 24) & 0xFF) as u8;
+
+            let mut out = Vec::with_capacity(len + header_len);
+            out.extend_from_slice(&header);
+
+            let mut uncompressed_len_buf = vec![0u8; int_len];
+            put_uvarint(&mut uncompressed_len_buf, uncompressed_len as u64);
+            out.extend_from_slice(&uncompressed_len_buf);
+
+            out.extend_from_slice(&compressed_data);
+
+            *this.written += out.len();
+            *this.uncomp_written += uncompressed_len;
+
+            this.index.add(*this.written as i64, *this.uncomp_written as i64)?;
+
+            *this.buffer = out;
+            *this.pos = 0;
+            this.temp_buffer.clear();
+
+            let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
+            buf.put_slice(&this.buffer[..to_copy]);
+            *this.pos += to_copy;
+            if *this.pos == this.buffer.len() {
+                this.buffer.clear();
+                *this.pos = 0;
+            }
+
+            // println!("write block, to_copy: {}, pos: {}, buffer_len: {}", to_copy, this.pos, this.buffer.len());
+            Poll::Ready(Ok(()))
+        } else {
+            Poll::Pending
         }
     }
 }
@@ -187,7 +246,7 @@ pin_project! {
 
 impl<R> DecompressReader<R>
 where
-    R: Reader,
+    R: AsyncRead + Unpin + Send + Sync,
 {
     pub fn new(inner: R, compression_algorithm: CompressionAlgorithm) -> Self {
         Self {
@@ -212,6 +271,7 @@ where
 {
     fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
         let mut this = self.project();
+
         // Serve from buffer if any
         if *this.buffer_pos < this.buffer.len() {
             let to_copy = std::cmp::min(buf.remaining(), this.buffer.len() - *this.buffer_pos);
@@ -221,6 +281,7 @@ where
                 this.buffer.clear();
                 *this.buffer_pos = 0;
             }
+
             return Poll::Ready(Ok(()));
         }
 
@@ -252,6 +313,10 @@ where
             }
         }
 
+        if !*this.header_done && *this.header_read == 0 {
+            return Poll::Ready(Ok(()));
+        }
+
         let typ = this.header_buf[0];
         let len = (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16);
         let crc = (this.header_buf[4] as u32)
@@ -263,14 +328,9 @@ where
         *this.header_read = 0;
         *this.header_done = true;
 
-        if typ == 0xFF {
-            *this.finished = true;
-            return Poll::Ready(Ok(()));
-        }
-
         // Save compressed block read progress across polls
         if this.compressed_buf.is_none() {
-            *this.compressed_len = len - 4;
+            *this.compressed_len = len;
             *this.compressed_buf = Some(vec![0u8; *this.compressed_len]);
             *this.compressed_read = 0;
         }
@@ -298,7 +358,7 @@ where
         // After reading all, unpack
         let (uncompress_len, uvarint) = uvarint(&compressed_buf[0..16]);
         let compressed_data = &compressed_buf[uvarint as usize..];
-        let decompressed = if typ == 0x00 {
+        let decompressed = if typ == COMPRESS_TYPE_COMPRESSED {
             match decompress_block(compressed_data, *this.compression_algorithm) {
                 Ok(out) => out,
                 Err(e) => {
@@ -308,9 +368,9 @@ where
                     return Poll::Ready(Err(e));
                 }
             }
-        } else if typ == 0x01 {
+        } else if typ == COMPRESS_TYPE_UNCOMPRESSED {
             compressed_data.to_vec()
-        } else if typ == 0xFF {
+        } else if typ == COMPRESS_TYPE_END {
             // Handle end marker
             this.compressed_buf.take();
             *this.compressed_read = 0;
@@ -348,6 +408,11 @@ where
         buf.put_slice(&this.buffer[..to_copy]);
         *this.buffer_pos += to_copy;
 
+        if *this.buffer_pos == this.buffer.len() {
+            this.buffer.clear();
+            *this.buffer_pos = 0;
+        }
+
         Poll::Ready(Ok(()))
     }
 }
@@ -375,6 +440,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use crate::WarpReader;
+
     use super::*;
     use std::io::Cursor;
     use tokio::io::{AsyncReadExt, BufReader};
@@ -383,7 +450,7 @@ mod tests {
     async fn test_compress_reader_basic() {
         let data = b"hello world, hello world, hello world!";
         let reader = Cursor::new(&data[..]);
-        let mut compress_reader = CompressReader::new(reader, CompressionAlgorithm::Gzip);
+        let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Gzip);
 
         let mut compressed = Vec::new();
         compress_reader.read_to_end(&mut compressed).await.unwrap();
@@ -400,7 +467,7 @@ mod tests {
     async fn test_compress_reader_basic_deflate() {
         let data = b"hello world, hello world, hello world!";
         let reader = BufReader::new(&data[..]);
-        let mut compress_reader = CompressReader::new(reader, CompressionAlgorithm::Deflate);
+        let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Deflate);
 
         let mut compressed = Vec::new();
         compress_reader.read_to_end(&mut compressed).await.unwrap();
@@ -417,7 +484,7 @@ mod tests {
     async fn test_compress_reader_empty() {
         let data = b"";
         let reader = BufReader::new(&data[..]);
-        let mut compress_reader = CompressReader::new(reader, CompressionAlgorithm::Gzip);
+        let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Gzip);
 
         let mut compressed = Vec::new();
         compress_reader.read_to_end(&mut compressed).await.unwrap();
@@ -436,7 +503,7 @@ mod tests {
         let mut data = vec![0u8; 1024 * 1024];
         rand::rng().fill(&mut data[..]);
         let reader = Cursor::new(data.clone());
-        let mut compress_reader = CompressReader::new(reader, CompressionAlgorithm::Gzip);
+        let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::Gzip);
 
         let mut compressed = Vec::new();
         compress_reader.read_to_end(&mut compressed).await.unwrap();
@@ -452,15 +519,15 @@ mod tests {
     async fn test_compress_reader_large_deflate() {
         use rand::Rng;
         // Generate 1MB of random bytes
-        let mut data = vec![0u8; 1024 * 1024];
+        let mut data = vec![0u8; 1024 * 1024 * 3 + 512];
         rand::rng().fill(&mut data[..]);
         let reader = Cursor::new(data.clone());
-        let mut compress_reader = CompressReader::new(reader, CompressionAlgorithm::Deflate);
+        let mut compress_reader = CompressReader::new(WarpReader::new(reader), CompressionAlgorithm::default());
 
         let mut compressed = Vec::new();
         compress_reader.read_to_end(&mut compressed).await.unwrap();
 
-        let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::Deflate);
+        let mut decompress_reader = DecompressReader::new(Cursor::new(compressed.clone()), CompressionAlgorithm::default());
         let mut decompressed = Vec::new();
         decompress_reader.read_to_end(&mut decompressed).await.unwrap();