diff --git a/crates/ecstore/Cargo.toml b/crates/ecstore/Cargo.toml index 0e12a8467..307b06ace 100644 --- a/crates/ecstore/Cargo.toml +++ b/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 diff --git a/crates/ecstore/src/erasure_coding/decode.rs b/crates/ecstore/src/erasure_coding/decode.rs index 4af48cd2c..43b809dff 100644 --- a/crates/ecstore/src/erasure_coding/decode.rs +++ b/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 read_futs = Vec::with_capacity(self.readers.len()); + let mut shards: Vec>> = vec![None; num_readers]; + let mut errs = vec![None; num_readers]; - for (i, opt_reader) in self.readers.iter_mut().enumerate() { - let future = if let Some(reader) = opt_reader.as_mut() { + let mut futures = Vec::with_capacity(self.total_shards); + let reader_iter: std::slice::IterMut<'_, Option>> = 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, 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>> = vec![None; self.readers.len()]; - let mut errs = vec![None; self.readers.len()]; + let mut success = 0; + 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() { - match shard { - Ok(data) => { - if !data.is_empty() { - shards[i] = Some(data); + if let Some(future) = fut_iter.next() { + sets.push(future); + } } } - Err(e) => { - // error!("Error reading shard {}: {}", i, e); - errs[i] = Some(e); + + if success >= self.data_shards { + 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>> { + 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()) + } +}