perf: avoid transmitting parity shards when the object is good (#322)

crates/ecstore/Cargo.toml

@@ -50,7 +50,7 @@ serde.workspace = true
 time.workspace = true
 bytesize.workspace = true
 serde_json.workspace = true
-quick-xml.workspace = true
+quick-xml = { workspace = true, features = ["serialize", "async-tokio"] }
 s3s.workspace = true
 http.workspace = true
 url.workspace = true

crates/ecstore/src/erasure_coding/decode.rs

@@ -16,7 +16,7 @@ use super::BitrotReader;
 use super::Erasure;
 use crate::disk::error::Error;
 use crate::disk::error_reduce::reduce_errs;
-use futures::future::join_all;
+use futures::stream::{FuturesUnordered, StreamExt};
 use pin_project_lite::pin_project;
 use std::io;
 use std::io::ErrorKind;
@@ -69,6 +69,7 @@ where
         // if self.readers.len() != self.total_shards {
         //     return Err(io::Error::new(ErrorKind::InvalidInput, "Invalid number of readers"));
         // }
+        let num_readers = self.readers.len();
 
         let shard_size = if self.offset + self.shard_size > self.shard_file_size {
             self.shard_file_size - self.offset
@@ -77,14 +78,16 @@ where
         };
 
         if shard_size == 0 {
-            return (vec![None; self.readers.len()], vec![None; self.readers.len()]);
+            return (vec![None; num_readers], vec![None; num_readers]);
         }
 
-        // 使用并发读取所有分片
+        let mut shards: Vec<Option<Vec<u8>>> = vec![None; num_readers];
-        let mut read_futs = Vec::with_capacity(self.readers.len());
+        let mut errs = vec![None; num_readers];
 
-        for (i, opt_reader) in self.readers.iter_mut().enumerate() {
+        let mut futures = Vec::with_capacity(self.total_shards);
-            let future = if let Some(reader) = opt_reader.as_mut() {
+        let reader_iter: std::slice::IterMut<'_, Option<BitrotReader<R>>> = self.readers.iter_mut();
+        for (i, reader) in reader_iter.enumerate() {
+            let future = if let Some(reader) = reader {
                 Box::pin(async move {
                     let mut buf = vec![0u8; shard_size];
                     match reader.read(&mut buf).await {
@@ -100,30 +103,41 @@ where
                 Box::pin(async move { (i, Err(Error::FileNotFound)) })
                     as std::pin::Pin<Box<dyn std::future::Future<Output = (usize, Result<Vec<u8>, Error>)> + Send>>
             };
-            read_futs.push(future);
+
+            futures.push(future);
         }
 
-        let results = join_all(read_futs).await;
+        if futures.len() >= self.data_shards {
+            let mut fut_iter = futures.into_iter();
+            let mut sets = FuturesUnordered::new();
+            for _ in 0..self.data_shards {
+                if let Some(future) = fut_iter.next() {
+                    sets.push(future);
+                }
+            }
 
-        let mut shards: Vec<Option<Vec<u8>>> = vec![None; self.readers.len()];
+            let mut success = 0;
-        let mut errs = vec![None; self.readers.len()];
+            while let Some((i, result)) = sets.next().await {
+                match result {
+                    Ok(v) => {
+                        shards[i] = Some(v);
+                        success += 1;
+                    }
+                    Err(e) => {
+                        errs[i] = Some(e);
 
-        for (i, shard) in results.into_iter() {
+                        if let Some(future) = fut_iter.next() {
-            match shard {
+                            sets.push(future);
-                Ok(data) => {
+                        }
-                    if !data.is_empty() {
-                        shards[i] = Some(data);
                     }
                 }
-                Err(e) => {
+
-                    // error!("Error reading shard {}: {}", i, e);
+                if success >= self.data_shards {
-                    errs[i] = Some(e);
+                    break;
                 }
             }
         }
 
-        self.offset += shard_size;
-
         (shards, errs)
     }
 
@@ -294,3 +308,151 @@ impl Erasure {
         (written, ret_err)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use rustfs_utils::HashAlgorithm;
+
+    use crate::{disk::error::DiskError, erasure_coding::BitrotWriter};
+
+    use super::*;
+    use std::io::Cursor;
+
+    #[tokio::test]
+    async fn test_parallel_reader_normal() {
+        const BLOCK_SIZE: usize = 64;
+        const NUM_SHARDS: usize = 2;
+        const DATA_SHARDS: usize = 8;
+        const PARITY_SHARDS: usize = 4;
+        const SHARD_SIZE: usize = BLOCK_SIZE / DATA_SHARDS;
+
+        let reader_offset = 0;
+        let mut readers = vec![];
+        for i in 0..(DATA_SHARDS + PARITY_SHARDS) {
+            readers.push(Some(
+                create_reader(SHARD_SIZE, NUM_SHARDS, (i % 256) as u8, &HashAlgorithm::HighwayHash256, false).await,
+            ));
+        }
+
+        let erausre = Erasure::new(DATA_SHARDS, PARITY_SHARDS, BLOCK_SIZE);
+        let mut parallel_reader = ParallelReader::new(readers, erausre, reader_offset, NUM_SHARDS * BLOCK_SIZE);
+
+        for _ in 0..NUM_SHARDS {
+            let (bufs, errs) = parallel_reader.read().await;
+
+            bufs.into_iter().enumerate().for_each(|(index, buf)| {
+                if index < DATA_SHARDS {
+                    assert!(buf.is_some());
+                    let buf = buf.unwrap();
+                    assert_eq!(SHARD_SIZE, buf.len());
+                    assert_eq!(index as u8, buf[0]);
+                } else {
+                    assert!(buf.is_none());
+                }
+            });
+
+            assert!(errs.iter().filter(|err| err.is_some()).count() == 0);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_parallel_reader_with_offline_disks() {
+        const OFFLINE_DISKS: usize = 2;
+        const NUM_SHARDS: usize = 2;
+        const BLOCK_SIZE: usize = 64;
+        const DATA_SHARDS: usize = 8;
+        const PARITY_SHARDS: usize = 4;
+        const SHARD_SIZE: usize = BLOCK_SIZE / DATA_SHARDS;
+
+        let reader_offset = 0;
+        let mut readers = vec![];
+        for i in 0..(DATA_SHARDS + PARITY_SHARDS) {
+            if i < OFFLINE_DISKS {
+                // Two disks are offline
+                readers.push(None);
+            } else {
+                readers.push(Some(
+                    create_reader(SHARD_SIZE, NUM_SHARDS, (i % 256) as u8, &HashAlgorithm::HighwayHash256, false).await,
+                ));
+            }
+        }
+
+        let erausre = Erasure::new(DATA_SHARDS, PARITY_SHARDS, BLOCK_SIZE);
+        let mut parallel_reader = ParallelReader::new(readers, erausre, reader_offset, NUM_SHARDS * BLOCK_SIZE);
+
+        for _ in 0..NUM_SHARDS {
+            let (bufs, errs) = parallel_reader.read().await;
+
+            assert_eq!(DATA_SHARDS, bufs.iter().filter(|buf| buf.is_some()).count());
+            assert_eq!(OFFLINE_DISKS, errs.iter().filter(|err| err.is_some()).count());
+        }
+    }
+
+    #[tokio::test]
+    async fn test_parallel_reader_with_bitrots() {
+        const BITROT_DISKS: usize = 2;
+        const NUM_SHARDS: usize = 2;
+        const BLOCK_SIZE: usize = 64;
+        const DATA_SHARDS: usize = 8;
+        const PARITY_SHARDS: usize = 4;
+        const SHARD_SIZE: usize = BLOCK_SIZE / DATA_SHARDS;
+
+        let reader_offset = 0;
+        let mut readers = vec![];
+        for i in 0..(DATA_SHARDS + PARITY_SHARDS) {
+            readers.push(Some(
+                create_reader(SHARD_SIZE, NUM_SHARDS, (i % 256) as u8, &HashAlgorithm::HighwayHash256, i < BITROT_DISKS).await,
+            ));
+        }
+
+        let erausre = Erasure::new(DATA_SHARDS, PARITY_SHARDS, BLOCK_SIZE);
+        let mut parallel_reader = ParallelReader::new(readers, erausre, reader_offset, NUM_SHARDS * BLOCK_SIZE);
+
+        for _ in 0..NUM_SHARDS {
+            let (bufs, errs) = parallel_reader.read().await;
+
+            assert_eq!(DATA_SHARDS, bufs.iter().filter(|buf| buf.is_some()).count());
+            assert_eq!(
+                BITROT_DISKS,
+                errs.iter()
+                    .filter(|err| {
+                        match err {
+                            Some(DiskError::Io(err)) => {
+                                err.kind() == std::io::ErrorKind::InvalidData && err.to_string().contains("bitrot")
+                            }
+                            _ => false,
+                        }
+                    })
+                    .count()
+            );
+        }
+    }
+
+    async fn create_reader(
+        shard_size: usize,
+        num_shards: usize,
+        value: u8,
+        hash_algo: &HashAlgorithm,
+        bitrot: bool,
+    ) -> BitrotReader<Cursor<Vec<u8>>> {
+        let len = (hash_algo.size() + shard_size) * num_shards;
+        let buf = Cursor::new(vec![0u8; len]);
+
+        let mut writer = BitrotWriter::new(buf, shard_size, hash_algo.clone());
+        for _ in 0..num_shards {
+            writer.write(vec![value; shard_size].as_slice()).await.unwrap();
+        }
+
+        let mut buf = writer.into_inner().into_inner();
+
+        if bitrot {
+            for i in 0..num_shards {
+                // Rot one bit for each shard
+                buf[i * (hash_algo.size() + shard_size)] ^= 1;
+            }
+        }
+
+        let reader_cursor = Cursor::new(buf);
+        BitrotReader::new(reader_cursor, shard_size, hash_algo.clone())
+    }
+}