refactor(storage): support zstd compression in block (#2960)

* refactor(storage): support zstd compression in block

Signed-off-by: Shmiwy <[email protected]>

* refactor(storage): support zstd compression in block

Signed-off-by: Shmiwy <[email protected]>

* refactor(storage): support zstd compression in block

Signed-off-by: Shmiwy <[email protected]>

Author: Shmiwy

Cargo.lock

@@ -70,6 +70,7 @@ tracing = { version = "0.1" }
 twox-hash = "1"
 value-encoding = { path = "../utils/value-encoding" }
 workspace-hack = { version = "0.1", path = "../workspace-hack" }
+zstd = "0.11.2"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 procinfo = { git = "https://github.com/tikv/procinfo-rs", rev = "6599eb9dca74229b2c1fcc44118bef7eff127128" }

src/storage/Cargo.toml

@@ -17,8 +17,8 @@ use std::io::{Read, Write};
 use std::ops::Range;
 
 use bytes::{Buf, BufMut, Bytes, BytesMut};
-use lz4::Decoder;
 use risingwave_hummock_sdk::VersionedComparator;
+use {lz4, zstd};
 
 use super::utils::{
     bytes_diff, var_u32_len, xxhash64_verify, BufExt, BufMutExt, CompressionAlgorithm,
@@ -45,17 +45,25 @@ impl Block {
 
         // Decompress.
         let compression = CompressionAlgorithm::decode(&mut &buf[buf.len() - 9..buf.len() - 8])?;
+        let compressed_data = &buf[..buf.len() - 9];
         let buf = match compression {
             CompressionAlgorithm::None => buf.slice(..buf.len() - 9),
             CompressionAlgorithm::Lz4 => {
-                let mut decoder = Decoder::new(buf.reader())
-                    .map_err(HummockError::decode_error)
-                    .unwrap();
+                let mut decoder = lz4::Decoder::new(compressed_data.reader())
+                    .map_err(HummockError::decode_error)?;
                 let mut decoded = Vec::with_capacity(DEFAULT_BLOCK_SIZE);
                 decoder
                     .read_to_end(&mut decoded)
-                    .map_err(HummockError::decode_error)
-                    .unwrap();
+                    .map_err(HummockError::decode_error)?;
+                Bytes::from(decoded)
+            }
+            CompressionAlgorithm::Zstd => {
+                let mut decoder = zstd::Decoder::new(compressed_data.reader())
+                    .map_err(HummockError::decode_error)?;
+                let mut decoded = Vec::with_capacity(DEFAULT_BLOCK_SIZE);
+                decoder
+                    .read_to_end(&mut decoded)
+                    .map_err(HummockError::decode_error)?;
                 Bytes::from(decoded)
             }
         };
@@ -298,6 +306,21 @@ impl BlockBuilder {
                 result.map_err(HummockError::encode_error).unwrap();
                 writer.into_inner()
             }
+            CompressionAlgorithm::Zstd => {
+                let mut encoder =
+                    zstd::Encoder::new(BytesMut::with_capacity(self.buf.len()).writer(), 4)
+                        .map_err(HummockError::encode_error)
+                        .unwrap();
+                encoder
+                    .write(&self.buf[..])
+                    .map_err(HummockError::encode_error)
+                    .unwrap();
+                let writer = encoder
+                    .finish()
+                    .map_err(HummockError::encode_error)
+                    .unwrap();
+                writer.into_inner()
+            }
         };
         self.compression_algorithm.encode(&mut buf);
         let checksum = xxhash64_checksum(&buf);
@@ -354,8 +377,13 @@ mod tests {
 
     #[test]
     fn test_compressed_block_enc_dec() {
+        inner_test_compressed(CompressionAlgorithm::Lz4);
+        inner_test_compressed(CompressionAlgorithm::Zstd);
+    }
+
+    fn inner_test_compressed(algo: CompressionAlgorithm) {
         let options = BlockBuilderOptions {
-            compression_algorithm: CompressionAlgorithm::Lz4,
+            compression_algorithm: algo,
             ..Default::default()
         };
         let mut builder = BlockBuilder::new(options);

src/storage/src/hummock/sstable/block.rs

@@ -154,13 +154,15 @@ pub fn get_length_prefixed_slice(buf: &mut &[u8]) -> Vec<u8> {
 pub enum CompressionAlgorithm {
     None,
     Lz4,
+    Zstd,
 }
 
 impl CompressionAlgorithm {
     pub fn encode(&self, buf: &mut impl BufMut) {
         let v = match self {
             Self::None => 0,
             Self::Lz4 => 1,
+            Self::Zstd => 2,
         };
         buf.put_u8(v);
     }
@@ -169,6 +171,7 @@ impl CompressionAlgorithm {
         match buf.get_u8() {
             0 => Ok(Self::None),
             1 => Ok(Self::Lz4),
+            2 => Ok(Self::Zstd),
             _ => Err(HummockError::decode_error(
                 "not valid compression algorithm",
             )),
@@ -181,6 +184,7 @@ impl From<CompressionAlgorithm> for u8 {
         match ca {
             CompressionAlgorithm::None => 0,
             CompressionAlgorithm::Lz4 => 1,
+            CompressionAlgorithm::Zstd => 2,
         }
     }
 }
@@ -190,6 +194,7 @@ impl From<CompressionAlgorithm> for u64 {
         match ca {
             CompressionAlgorithm::None => 0,
             CompressionAlgorithm::Lz4 => 1,
+            CompressionAlgorithm::Zstd => 2,
         }
     }
 }
@@ -201,6 +206,7 @@ impl TryFrom<u8> for CompressionAlgorithm {
         match v {
             0 => Ok(Self::None),
             1 => Ok(Self::Lz4),
+            2 => Ok(Self::Zstd),
             _ => Err(HummockError::decode_error(
                 "not valid compression algorithm",
             )),