use crate::bitrot::{BitrotReader, BitrotWriter}; use crate::error::clone_err; use crate::io::Etag; use crate::quorum::{object_op_ignored_errs, reduce_write_quorum_errs}; use bytes::{Bytes, BytesMut}; use common::error::{Error, Result}; use futures::future::join_all; use reed_solomon_erasure::galois_8::ReedSolomon; use smallvec::SmallVec; use std::any::Any; use std::io::ErrorKind; use std::sync::Arc; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::mpsc; use tracing::warn; use tracing::{error, info}; // use tracing::debug; use uuid::Uuid; // use reader::reader::ChunkedStream; // use crate::chunk_stream::ChunkedStream; use crate::disk::error::DiskError; #[derive(Default)] pub struct Erasure { data_shards: usize, parity_shards: usize, encoder: Option, pub block_size: usize, _id: Uuid, _buf: Vec, } impl Erasure { pub fn new(data_shards: usize, parity_shards: usize, block_size: usize) -> Self { // debug!( // "Erasure new data_shards {},parity_shards {} block_size {} ", // data_shards, parity_shards, block_size // ); let mut encoder = None; if parity_shards > 0 { encoder = Some(ReedSolomon::new(data_shards, parity_shards).unwrap()); } Erasure { data_shards, parity_shards, block_size, encoder, _id: Uuid::new_v4(), _buf: vec![0u8; block_size], } } #[tracing::instrument(level = "info", skip(self, reader, writers))] pub async fn encode( self: Arc, mut reader: S, writers: &mut [Option], // block_size: usize, total_size: usize, write_quorum: usize, ) -> Result<(usize, String)> where S: AsyncRead + Etag + Unpin + Send + 'static, { let (tx, mut rx) = mpsc::channel(5); let task = tokio::spawn(async move { let mut buf = vec![0u8; self.block_size]; let mut total: usize = 0; loop { if total_size > 0 { let new_len = { let remain = total_size - total; if remain > self.block_size { self.block_size } else { remain } }; if new_len == 0 && total > 0 { break; } buf.resize(new_len, 0u8); match reader.read_exact(&mut buf).await { Ok(res) => res, Err(e) => { if let ErrorKind::UnexpectedEof = e.kind() { break; } else { return Err(Error::new(e)); } } }; total += buf.len(); } let blocks = Arc::new(Box::pin(self.clone().encode_data(&buf)?)); let _ = tx.send(blocks).await; if total_size == 0 { break; } } let etag = reader.etag().await; Ok((total, etag)) }); while let Some(blocks) = rx.recv().await { let write_futures = writers.iter_mut().enumerate().map(|(i, w_op)| { let i_inner = i; let blocks_inner = blocks.clone(); async move { if let Some(w) = w_op { (w.write(blocks_inner[i_inner].clone()).await).err() } else { Some(Error::new(DiskError::DiskNotFound)) } } }); let errs = join_all(write_futures).await; let none_count = errs.iter().filter(|&x| x.is_none()).count(); if none_count >= write_quorum { if total_size == 0 { break; } continue; } if let Some(err) = reduce_write_quorum_errs(&errs, object_op_ignored_errs().as_ref(), write_quorum) { warn!("Erasure encode errs {:?}", &errs); return Err(err); } } task.await? } pub async fn decode( &self, writer: &mut W, readers: Vec>, offset: usize, length: usize, total_length: usize, ) -> (usize, Option) where W: AsyncWriteExt + Send + Unpin + 'static, { if length == 0 { return (0, None); } let mut reader = ShardReader::new(readers, self, offset, total_length); // debug!("ShardReader {:?}", &reader); let start_block = offset / self.block_size; let end_block = (offset + length) / self.block_size; // debug!("decode block from {} to {}", start_block, end_block); let mut bytes_writed = 0; for block_idx in start_block..=end_block { let (block_offset, block_length) = if start_block == end_block { (offset % self.block_size, length) } else if block_idx == start_block { let block_offset = offset % self.block_size; (block_offset, self.block_size - block_offset) } else if block_idx == end_block { (0, (offset + length) % self.block_size) } else { (0, self.block_size) }; if block_length == 0 { // debug!("block_length == 0 break"); break; } // debug!("decode {} block_offset {},block_length {} ", block_idx, block_offset, block_length); let mut bufs = match reader.read().await { Ok(bufs) => bufs, Err(err) => return (bytes_writed, Some(err)), }; if self.parity_shards > 0 { if let Err(err) = self.decode_data(&mut bufs) { return (bytes_writed, Some(err)); } } let writed_n = match self .write_data_blocks(writer, bufs, self.data_shards, block_offset, block_length) .await { Ok(n) => n, Err(err) => { error!("write_data_blocks err {:?}", &err); return (bytes_writed, Some(err)); } }; bytes_writed += writed_n; // debug!("decode {} writed_n {}, total_writed: {} ", block_idx, writed_n, bytes_writed); } if bytes_writed != length { // debug!("bytes_writed != length: {} != {} ", bytes_writed, length); return (bytes_writed, Some(Error::msg("erasure decode less data"))); } (bytes_writed, None) } async fn write_data_blocks( &self, writer: &mut W, bufs: Vec>>, data_blocks: usize, offset: usize, length: usize, ) -> Result where W: AsyncWrite + Send + Unpin + 'static, { if bufs.len() < data_blocks { return Err(Error::msg("read bufs not match data_blocks")); } let data_len: usize = bufs .iter() .take(data_blocks) .filter(|v| v.is_some()) .map(|v| v.as_ref().unwrap().len()) .sum(); if data_len < length { return Err(Error::msg(format!("write_data_blocks data_len < length {} < {}", data_len, length))); } let mut offset = offset; // debug!("write_data_blocks offset {}, length {}", offset, length); let mut write = length; let mut total_writed = 0; for opt_buf in bufs.iter().take(data_blocks) { let buf = opt_buf.as_ref().unwrap(); if offset >= buf.len() { offset -= buf.len(); continue; } let buf = &buf[offset..]; offset = 0; // debug!("write_data_blocks write buf len {}", buf.len()); if write < buf.len() { let buf = &buf[..write]; // debug!("write_data_blocks write buf less len {}", buf.len()); writer.write_all(buf).await?; // debug!("write_data_blocks write done len {}", buf.len()); total_writed += buf.len(); break; } writer.write_all(buf).await?; let n = buf.len(); // debug!("write_data_blocks write done len {}", n); write -= n; total_writed += n; } Ok(total_writed) } pub fn total_shard_count(&self) -> usize { self.data_shards + self.parity_shards } #[tracing::instrument(level = "info", skip_all, fields(data_len=data.len()))] pub fn encode_data(self: Arc, data: &[u8]) -> Result> { let (shard_size, total_size) = self.need_size(data.len()); // 生成一个新的 所需的所有分片数据长度 let mut data_buffer = BytesMut::with_capacity(total_size); // 复制源数据 data_buffer.extend_from_slice(data); data_buffer.resize(total_size, 0u8); { // ec encode, 结果会写进 data_buffer let data_slices: SmallVec<[&mut [u8]; 16]> = data_buffer.chunks_exact_mut(shard_size).collect(); // partiy 数量大于0 才ec if self.parity_shards > 0 { self.encoder.as_ref().unwrap().encode(data_slices)?; } } // 零拷贝分片,所有 shard 引用 data_buffer let mut data_buffer = data_buffer.freeze(); let mut shards = Vec::with_capacity(self.total_shard_count()); for _ in 0..self.total_shard_count() { let shard = data_buffer.split_to(shard_size); shards.push(shard); } Ok(shards) } pub fn decode_data(&self, shards: &mut [Option>]) -> Result<()> { if self.parity_shards > 0 { self.encoder.as_ref().unwrap().reconstruct(shards)?; } Ok(()) } // 每个分片长度,所需要的总长度 fn need_size(&self, data_size: usize) -> (usize, usize) { let shard_size = self.shard_size(data_size); (shard_size, shard_size * (self.total_shard_count())) } // 算出每个分片大小 pub fn shard_size(&self, data_size: usize) -> usize { data_size.div_ceil(self.data_shards) } // returns final erasure size from original size. pub fn shard_file_size(&self, total_size: usize) -> usize { if total_size == 0 { return 0; } let num_shards = total_size / self.block_size; let last_block_size = total_size % self.block_size; let last_shard_size = last_block_size.div_ceil(self.data_shards); num_shards * self.shard_size(self.block_size) + last_shard_size // // 因为写入的时候ec需要补全,所以最后一个长度应该也是一样的 // if last_block_size != 0 { // num_shards += 1 // } // num_shards * self.shard_size(self.block_size) } // where erasure reading begins. pub fn shard_file_offset(&self, start_offset: usize, length: usize, total_length: usize) -> usize { let shard_size = self.shard_size(self.block_size); let shard_file_size = self.shard_file_size(total_length); let end_shard = (start_offset + length) / self.block_size; let mut till_offset = end_shard * shard_size + shard_size; if till_offset > shard_file_size { till_offset = shard_file_size; } till_offset } pub async fn heal( &self, writers: &mut [Option], readers: Vec>, total_length: usize, _prefer: &[bool], ) -> Result<()> { info!( "Erasure heal, writers len: {}, readers len: {}, total_length: {}", writers.len(), readers.len(), total_length ); if writers.len() != self.parity_shards + self.data_shards { return Err(Error::from_string("invalid argument")); } let mut reader = ShardReader::new(readers, self, 0, total_length); let start_block = 0; let mut end_block = total_length / self.block_size; if total_length % self.block_size != 0 { end_block += 1; } let mut errs = Vec::new(); for _ in start_block..end_block { let mut bufs = reader.read().await?; if self.parity_shards > 0 { self.encoder.as_ref().unwrap().reconstruct(&mut bufs)?; } let shards = bufs.into_iter().flatten().map(Bytes::from).collect::>(); if shards.len() != self.parity_shards + self.data_shards { return Err(Error::from_string("can not reconstruct data")); } for (i, w) in writers.iter_mut().enumerate() { if w.is_none() { continue; } match w.as_mut().unwrap().write(shards[i].clone()).await { Ok(_) => {} Err(e) => { info!("write failed, err: {:?}", e); errs.push(e); } } } } if !errs.is_empty() { return Err(clone_err(&errs[0])); } Ok(()) } } #[async_trait::async_trait] pub trait Writer { fn as_any(&self) -> &dyn Any; async fn write(&mut self, buf: Bytes) -> Result<()>; async fn close(&mut self) -> Result<()> { Ok(()) } } #[async_trait::async_trait] pub trait ReadAt { async fn read_at(&mut self, offset: usize, length: usize) -> Result<(Vec, usize)>; } pub struct ShardReader { readers: Vec>, // 磁盘 data_block_count: usize, // 总的分片数量 parity_block_count: usize, shard_size: usize, // 每个分片的块大小 一次读取一块 shard_file_size: usize, // 分片文件总长度 offset: usize, // 在分片中的offset } impl ShardReader { pub fn new(readers: Vec>, ec: &Erasure, offset: usize, total_length: usize) -> Self { Self { readers, data_block_count: ec.data_shards, parity_block_count: ec.parity_shards, shard_size: ec.shard_size(ec.block_size), shard_file_size: ec.shard_file_size(total_length), offset: (offset / ec.block_size) * ec.shard_size(ec.block_size), } } pub async fn read(&mut self) -> Result>>> { // let mut disks = self.readers; let reader_length = self.readers.len(); // 需要读取的块长度 let mut read_length = self.shard_size; if self.offset + read_length > self.shard_file_size { read_length = self.shard_file_size - self.offset } if read_length == 0 { return Ok(vec![None; reader_length]); } // debug!("shard reader read offset {}, shard_size {}", self.offset, read_length); let mut futures = Vec::with_capacity(reader_length); let mut errors = Vec::with_capacity(reader_length); let mut ress = Vec::with_capacity(reader_length); for disk in self.readers.iter_mut() { // if disk.is_none() { // ress.push(None); // errors.push(Some(Error::new(DiskError::DiskNotFound))); // continue; // } // let disk: &mut BitrotReader = disk.as_mut().unwrap(); let offset = self.offset; futures.push(async move { if let Some(disk) = disk { disk.read_at(offset, read_length).await } else { Err(Error::new(DiskError::DiskNotFound)) } }); } let results = join_all(futures).await; for result in results { match result { Ok((res, _)) => { ress.push(Some(res)); errors.push(None); } Err(e) => { ress.push(None); errors.push(Some(e)); } } } if !self.can_decode(&ress) { warn!("ec decode read ress {:?}", &ress); warn!("ec decode read errors {:?}", &errors); return Err(Error::msg("shard reader read faild")); } self.offset += self.shard_size; Ok(ress) } fn can_decode(&self, bufs: &[Option>]) -> bool { let c = bufs.iter().filter(|v| v.is_some()).count(); if self.parity_block_count > 0 { c >= self.data_block_count } else { c == self.data_block_count } } } // fn shards_to_option_shards(shards: &[Vec]) -> Vec>> { // let mut result = Vec::with_capacity(shards.len()); // for v in shards.iter() { // let inner: Vec = v.clone(); // result.push(Some(inner)); // } // result // } #[cfg(test)] mod test { use super::*; #[test] fn test_erasure() { let data_shards = 3; let parity_shards = 2; let data: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]; let ec = Erasure::new(data_shards, parity_shards, 1); let shards = Arc::new(ec).encode_data(data).unwrap(); println!("shards:{:?}", shards); let mut s: Vec<_> = shards .iter() .map(|d| if d.is_empty() { None } else { Some(d.to_vec()) }) .collect(); // let mut s = shards_to_option_shards(&shards); // s[0] = None; s[4] = None; s[3] = None; println!("sss:{:?}", &s); let ec = Erasure::new(data_shards, parity_shards, 1); ec.decode_data(&mut s).unwrap(); // ec.encoder.reconstruct(&mut s).unwrap(); println!("sss:{:?}", &s); } }