perf(storage): optimize internode RPC transfer path (#2262)

Co-authored-by: momoda693 <momoda693@gmail.com>

crates/ecstore/src/erasure_coding/bitrot.rs

@@ -15,6 +15,7 @@
 use bytes::Bytes;
 use pin_project_lite::pin_project;
 use rustfs_utils::HashAlgorithm;
+use std::io::IoSlice;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use tracing::error;
 use uuid::Uuid;
@@ -155,23 +156,49 @@ where
                 error!("bitrot writer write hash error: hash is empty");
                 return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "hash is empty"));
             }
-            self.inner.write_all(hash.as_ref()).await?;
+            write_all_vectored(&mut self.inner, hash.as_ref(), buf).await?;
+        } else {
+            self.inner.write_all(buf).await?;
         }
 
-        self.inner.write_all(buf).await?;
-
-        self.inner.flush().await?;
-
         let n = buf.len();
 
         Ok(n)
     }
 
     pub async fn shutdown(&mut self) -> std::io::Result<()> {
+        self.inner.flush().await?;
         self.inner.shutdown().await
     }
 }
 
+async fn write_all_vectored<W>(writer: &mut W, hash: &[u8], data: &[u8]) -> std::io::Result<()>
+where
+    W: AsyncWrite + Unpin,
+{
+    let mut hash_offset = 0;
+    let mut data_offset = 0;
+
+    while hash_offset < hash.len() || data_offset < data.len() {
+        let slices = [IoSlice::new(&hash[hash_offset..]), IoSlice::new(&data[data_offset..])];
+        let written = writer.write_vectored(&slices).await?;
+        if written == 0 {
+            return Err(std::io::Error::new(std::io::ErrorKind::WriteZero, "failed to write hash and data"));
+        }
+
+        let hash_remaining = hash.len() - hash_offset;
+        if written < hash_remaining {
+            hash_offset += written;
+            continue;
+        }
+
+        hash_offset = hash.len();
+        data_offset += written - hash_remaining;
+    }
+
+    Ok(())
+}
+
 pub fn bitrot_shard_file_size(size: usize, shard_size: usize, algo: HashAlgorithm) -> usize {
     if algo != HashAlgorithm::HighwayHash256S && algo != HashAlgorithm::HighwayHash256SLegacy {
         return size;
@@ -292,6 +319,33 @@ impl AsyncWrite for CustomWriter {
             }
         }
     }
+
+    fn poll_write_vectored(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> std::task::Poll<std::io::Result<usize>> {
+        match self.get_mut() {
+            Self::InlineBuffer(data) => {
+                let total = bufs.iter().map(|buf| buf.len()).sum::<usize>();
+                for buf in bufs {
+                    data.extend_from_slice(buf);
+                }
+                std::task::Poll::Ready(Ok(total))
+            }
+            Self::Other(writer) => {
+                let pinned_writer = std::pin::Pin::new(writer.as_mut());
+                pinned_writer.poll_write_vectored(cx, bufs)
+            }
+        }
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        match self {
+            Self::InlineBuffer(_) => true,
+            Self::Other(writer) => writer.is_write_vectored(),
+        }
+    }
 }
 
 /// Wrapper around BitrotWriter that uses our custom writer
@@ -361,7 +415,74 @@ mod tests {
     use super::BitrotReader;
     use super::BitrotWriter;
     use rustfs_utils::HashAlgorithm;
-    use std::io::Cursor;
+    use std::io::{Cursor, IoSlice};
+    use std::sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    };
+    use std::task::{Context, Poll};
+    use tokio::io::AsyncWrite;
+
+    #[derive(Default)]
+    struct VectoredCountingWriter {
+        vectored_writes: Arc<AtomicUsize>,
+        writes: Vec<u8>,
+    }
+
+    impl AsyncWrite for VectoredCountingWriter {
+        fn poll_write(self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &[u8]) -> Poll<std::io::Result<usize>> {
+            Poll::Ready(Err(std::io::Error::other("poll_write should not be used")))
+        }
+
+        fn poll_flush(self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_shutdown(self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_write_vectored(
+            mut self: std::pin::Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+            bufs: &[IoSlice<'_>],
+        ) -> Poll<std::io::Result<usize>> {
+            self.vectored_writes.fetch_add(1, Ordering::SeqCst);
+            let total = bufs.iter().map(|buf| buf.len()).sum::<usize>();
+            for buf in bufs {
+                self.writes.extend_from_slice(buf);
+            }
+            Poll::Ready(Ok(total))
+        }
+
+        fn is_write_vectored(&self) -> bool {
+            true
+        }
+    }
+
+    #[derive(Default)]
+    struct CountingWriter {
+        flushes: Arc<AtomicUsize>,
+        shutdowns: Arc<AtomicUsize>,
+        writes: Vec<u8>,
+    }
+
+    impl AsyncWrite for CountingWriter {
+        fn poll_write(mut self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
+            self.writes.extend_from_slice(buf);
+            Poll::Ready(Ok(buf.len()))
+        }
+
+        fn poll_flush(self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            self.flushes.fetch_add(1, Ordering::SeqCst);
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_shutdown(self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            self.shutdowns.fetch_add(1, Ordering::SeqCst);
+            Poll::Ready(Ok(()))
+        }
+    }
 
     #[tokio::test]
     async fn test_bitrot_read_write_ok() {
@@ -471,4 +592,41 @@ mod tests {
         assert_eq!(n, data_size);
         assert_eq!(data, &out[..]);
     }
+
+    #[tokio::test]
+    async fn test_bitrot_writer_flushes_once_on_shutdown() {
+        let flushes = Arc::new(AtomicUsize::new(0));
+        let shutdowns = Arc::new(AtomicUsize::new(0));
+        let writer = CountingWriter {
+            flushes: flushes.clone(),
+            shutdowns: shutdowns.clone(),
+            writes: Vec::new(),
+        };
+        let mut bitrot_writer = BitrotWriter::new(writer, 8, HashAlgorithm::None);
+
+        bitrot_writer.write(b"12345678").await.unwrap();
+        bitrot_writer.write(b"abc").await.unwrap();
+
+        assert_eq!(flushes.load(Ordering::SeqCst), 0);
+        assert_eq!(shutdowns.load(Ordering::SeqCst), 0);
+
+        bitrot_writer.shutdown().await.unwrap();
+
+        assert_eq!(flushes.load(Ordering::SeqCst), 1);
+        assert_eq!(shutdowns.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn test_bitrot_writer_uses_vectored_write_for_hash_and_data() {
+        let vectored_writes = Arc::new(AtomicUsize::new(0));
+        let writer = VectoredCountingWriter {
+            vectored_writes: vectored_writes.clone(),
+            writes: Vec::new(),
+        };
+        let mut bitrot_writer = BitrotWriter::new(writer, 8, HashAlgorithm::HighwayHash256);
+
+        bitrot_writer.write(b"payload").await.unwrap();
+
+        assert!(vectored_writes.load(Ordering::SeqCst) > 0);
+    }
 }

crates/ecstore/src/erasure_coding/encode.rs

@@ -112,11 +112,66 @@ impl<'a> MultiWriter<'a> {
         )))
     }
 
-    pub async fn _shutdown(&mut self) -> std::io::Result<()> {
-        for writer in self.writers.iter_mut().flatten() {
-            writer.shutdown().await?;
+    async fn shutdown_writer(writer_opt: &mut Option<BitrotWriterWrapper>, err: &mut Option<Error>) {
+        match writer_opt {
+            Some(writer) => match writer.shutdown().await {
+                Ok(()) => {
+                    *err = None;
+                }
+                Err(e) => {
+                    *err = Some(Error::from(e));
+                    *writer_opt = None;
+                }
+            },
+            None => {
+                *err = Some(Error::DiskNotFound);
+            }
         }
-        Ok(())
+    }
+
+    pub async fn shutdown(&mut self) -> std::io::Result<()> {
+        {
+            let mut futures = FuturesUnordered::new();
+            for (writer_opt, err) in self.writers.iter_mut().zip(self.errs.iter_mut()) {
+                if err.is_some() {
+                    continue;
+                }
+                futures.push(Self::shutdown_writer(writer_opt, err));
+            }
+            while let Some(()) = futures.next().await {}
+        }
+
+        let nil_count = self.errs.iter().filter(|&e| e.is_none()).count();
+        if nil_count >= self.write_quorum {
+            return Ok(());
+        }
+
+        if let Some(write_err) = reduce_write_quorum_errs(&self.errs, OBJECT_OP_IGNORED_ERRS, self.write_quorum) {
+            error!(
+                "reduce_write_quorum_errs during shutdown: {:?}, offline-disks={}/{}, errs={:?}",
+                write_err,
+                count_errs(&self.errs, &Error::DiskNotFound),
+                self.writers.len(),
+                self.errs
+            );
+            return Err(std::io::Error::other(format!(
+                "Failed to shutdown writers: {} (offline-disks={}/{})",
+                write_err,
+                count_errs(&self.errs, &Error::DiskNotFound),
+                self.writers.len()
+            )));
+        }
+
+        Err(std::io::Error::other(format!(
+            "Failed to shutdown writers: (offline-disks={}/{}): {}",
+            count_errs(&self.errs, &Error::DiskNotFound),
+            self.writers.len(),
+            self.errs
+                .iter()
+                .map(|e| e.as_ref().map_or("<nil>".to_string(), |e| e.to_string()))
+                .collect::<Vec<_>>()
+                .join(", ")
+        )))
     }
 }
 
@@ -176,7 +231,69 @@ impl Erasure {
         }
 
         let (reader, total) = task.await??;
-        // writers.shutdown().await?;
+        writers.shutdown().await?;
         Ok((reader, total))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::erasure_coding::{BitrotWriterWrapper, CustomWriter};
+    use rustfs_utils::HashAlgorithm;
+    use std::pin::Pin;
+    use std::sync::{Arc, Mutex};
+    use std::task::{Context, Poll};
+    use tokio::io::AsyncWrite;
+
+    #[derive(Clone, Default)]
+    struct DeferredCommitWriter {
+        buffered: Vec<u8>,
+        committed: Arc<Mutex<Vec<u8>>>,
+    }
+
+    impl DeferredCommitWriter {
+        fn new(committed: Arc<Mutex<Vec<u8>>>) -> Self {
+            Self {
+                buffered: Vec::new(),
+                committed,
+            }
+        }
+    }
+
+    impl AsyncWrite for DeferredCommitWriter {
+        fn poll_write(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
+            self.buffered.extend_from_slice(buf);
+            Poll::Ready(Ok(buf.len()))
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_shutdown(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            let buffered = std::mem::take(&mut self.buffered);
+            let mut committed = self.committed.lock().unwrap();
+            committed.extend_from_slice(&buffered);
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    #[tokio::test]
+    async fn encode_shutdowns_writers_after_small_shards() {
+        let committed = Arc::new(Mutex::new(Vec::new()));
+        let writer = DeferredCommitWriter::new(committed.clone());
+        let mut writers = vec![Some(BitrotWriterWrapper::new(
+            CustomWriter::new_tokio_writer(writer),
+            16,
+            HashAlgorithm::HighwayHash256S,
+        ))];
+
+        let erasure = Arc::new(Erasure::new(1, 0, 16));
+        let reader = tokio::io::BufReader::new(std::io::Cursor::new(b"small payload".to_vec()));
+        let (_reader, written) = erasure.encode(reader, &mut writers, 1).await.unwrap();
+
+        assert_eq!(written, b"small payload".len());
+        assert!(!committed.lock().unwrap().is_empty());
+    }
+}