diff --git a/ecstore/src/bitrot.rs b/ecstore/src/bitrot.rs index b64c84b3b..05e55a430 100644 --- a/ecstore/src/bitrot.rs +++ b/ecstore/src/bitrot.rs @@ -534,6 +534,7 @@ impl Writer for BitrotFileWriter { self } + #[tracing::instrument(level = "info", skip_all)] async fn write(&mut self, buf: Bytes) -> Result<()> { if buf.is_empty() { return Ok(()); diff --git a/ecstore/src/erasure.rs b/ecstore/src/erasure.rs index 590889e78..d5a08af4b 100644 --- a/ecstore/src/erasure.rs +++ b/ecstore/src/erasure.rs @@ -6,8 +6,10 @@ use common::error::{Error, Result}; use futures::future::join_all; use reed_solomon_erasure::galois_8::ReedSolomon; use smallvec::SmallVec; +use tokio::sync::mpsc; use std::any::Any; use std::io::ErrorKind; +use std::sync::Arc; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tracing::warn; @@ -52,8 +54,8 @@ impl Erasure { #[tracing::instrument(level = "debug", skip(self, reader, writers))] pub async fn encode( - &mut self, - reader: &mut S, + self: Arc, + mut reader: S, writers: &mut [Option], // block_size: usize, total_size: usize, @@ -67,40 +69,47 @@ impl Erasure { // body.map(|f| f.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string()))), // ); - let mut total: usize = 0; - let mut blocks = >::new(); - - 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; - } - - self.buf.resize(new_len, 0u8); - match reader.read_exact(&mut self.buf).await { - Ok(res) => res, - Err(e) => { - if let ErrorKind::UnexpectedEof = e.kind() { - break; + let (tx, mut rx) = mpsc::channel(3); + let self_clone = self.clone(); + let task = tokio::spawn(async move { + let mut total: usize = 0; + let mut buf = Vec::new(); + loop { + let mut blocks = >::new(); + if total_size > 0 { + let new_len = { + let remain = total_size - total; + if remain > self_clone.block_size { + self_clone.block_size } else { - return Err(Error::new(e)); + remain } + }; + + if new_len == 0 && total > 0 { + break; } - }; - total += self.buf.len(); + + 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(); + } + self_clone.clone().encode_data(&buf, &mut blocks)?; + let _ = tx.send(blocks).await; } - - self.encode_data(&self.buf, &mut blocks)?; - + Ok(total) + }); + + 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(); @@ -130,8 +139,7 @@ impl Erasure { break; } } - - Ok(total) + task.await? // // let stream = ChunkedStream::new(body, self.block_size); // let stream = ChunkedStream::new(body, total_size, self.block_size, false); @@ -356,8 +364,8 @@ impl Erasure { self.data_shards + self.parity_shards } - #[tracing::instrument(level = "debug", skip_all, fields(data_len=data.len()))] - pub fn encode_data(&self, data: &[u8], shards: &mut SmallVec<[Bytes; 16]>) -> Result<()> { + #[tracing::instrument(level = "info", skip_all, fields(data_len=data.len()))] + pub fn encode_data(self: Arc, data: &[u8], shards: &mut SmallVec<[Bytes; 16]>) -> Result<()> { let (shard_size, total_size) = self.need_size(data.len()); // 生成一个新的 所需的所有分片数据长度 @@ -618,34 +626,34 @@ impl ShardReader { #[cfg(test)] mod test { - use super::*; + // 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 mut shards = SmallVec::new(); - ec.encode_data(data, &mut shards).unwrap(); - println!("shards:{:?}", shards); + // #[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 mut shards = SmallVec::new(); + // Arc::new(ec).encode_data(data, &mut shards).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: Vec<_> = shards + // .iter() + // .map(|d| if d.is_empty() { None } else { Some(d.to_vec()) }) + // .collect(); - // let mut s = shards_to_option_shards(&shards); + // // let mut s = shards_to_option_shards(&shards); - // s[0] = None; - s[4] = None; - s[3] = None; + // // s[0] = None; + // s[4] = None; + // s[3] = None; - println!("sss:{:?}", &s); + // println!("sss:{:?}", &s); - ec.decode_data(&mut s).unwrap(); - // ec.encoder.reconstruct(&mut s).unwrap(); + // ec.decode_data(&mut s).unwrap(); + // // ec.encoder.reconstruct(&mut s).unwrap(); - println!("sss:{:?}", &s); - } + // println!("sss:{:?}", &s); + // } } diff --git a/ecstore/src/io.rs b/ecstore/src/io.rs index f2affe8cb..6bdbd629b 100644 --- a/ecstore/src/io.rs +++ b/ecstore/src/io.rs @@ -2,7 +2,10 @@ use bytes::Bytes; use futures::TryStreamExt; use md5::Digest; use md5::Md5; +use pin_project_lite::pin_project; +use std::io; use std::pin::Pin; +use std::task::ready; use std::task::Context; use std::task::Poll; use tokio::io::AsyncRead; @@ -125,12 +128,19 @@ impl AsyncRead for HttpFileReader { } } -pub struct EtagReader { - inner: R, - bytes_tx: mpsc::Sender, - md5_rx: oneshot::Receiver, +pub trait { + } +pin_project! { + pub struct EtagReader { + inner: R, + bytes_tx: mpsc::Sender, + md5_rx: oneshot::Receiver, + } +} + + impl EtagReader { pub fn new(inner: R) -> Self { let (bytes_tx, mut bytes_rx) = mpsc::channel::(8); @@ -157,21 +167,43 @@ impl EtagReader { } impl AsyncRead for EtagReader { - #[tracing::instrument(level = "debug", skip_all)] - fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { - let poll = Pin::new(&mut self.inner).poll_read(cx, buf); - if let Poll::Ready(Ok(())) = &poll { - if buf.remaining() == 0 { + #[tracing::instrument(level = "info", skip_all)] + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let me = self.project(); + + loop { + let rem = buf.remaining(); + if rem != 0 { + ready!(Pin::new(&mut *me.inner).poll_read(cx, buf))?; + if buf.remaining() == rem { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")).into(); + } + } else { let bytes = buf.filled(); let bytes = Bytes::copy_from_slice(bytes); - let tx = self.bytes_tx.clone(); + let tx = me.bytes_tx.clone(); tokio::spawn(async move { if let Err(e) = tx.send(bytes).await { warn!("EtagReader send error: {:?}", e); } }); + return Poll::Ready(Ok(())); } } - poll + + // let poll = Pin::new(&mut self.inner).poll_read(cx, buf); + // if let Poll::Ready(Ok(())) = &poll { + // if buf.remaining() == 0 { + // let bytes = buf.filled(); + // let bytes = Bytes::copy_from_slice(bytes); + // let tx = self.bytes_tx.clone(); + // tokio::spawn(async move { + // if let Err(e) = tx.send(bytes).await { + // warn!("EtagReader send error: {:?}", e); + // } + // }); + // } + // } + // poll } } diff --git a/ecstore/src/set_disk.rs b/ecstore/src/set_disk.rs index bd55ab39f..3c2b033c8 100644 --- a/ecstore/src/set_disk.rs +++ b/ecstore/src/set_disk.rs @@ -3802,8 +3802,8 @@ impl ObjectIO for SetDisks { // TODO: etag from header - let w_size = erasure - .encode(&mut etag_stream, &mut writers, data.content_length, write_quorum) + let w_size = Arc::new(erasure) + .encode(etag_stream, &mut writers, data.content_length, write_quorum) .await?; // TODO: 出错，删除临时目录 if let Err(err) = close_bitrot_writers(&mut writers).await {