feat(perf): add large PUT tuning and encode optimization (#3816)

crates/ecstore/src/erasure_coding/encode.rs

@@ -25,16 +25,20 @@ use futures::stream::FuturesUnordered;
 use std::sync::Arc;
 use std::vec;
 use tokio::io::AsyncRead;
+use tokio::runtime::RuntimeFlavor;
 use tokio::sync::mpsc;
 use tracing::error;
 
 const ENV_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BYTES: &str = "RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BYTES";
+const ENV_RUSTFS_ERASURE_ENCODE_BATCH_BLOCKS: &str = "RUSTFS_ERASURE_ENCODE_BATCH_BLOCKS";
 const DEFAULT_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BYTES: usize = 32 * 1024 * 1024;
 const DEFAULT_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BLOCKS: usize = 32;
+const DEFAULT_RUSTFS_ERASURE_ENCODE_BATCH_BLOCKS: usize = 4;
 
 /// Cached value of `RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BYTES` env var.
 /// Read once at first use via `OnceLock` to avoid per-encode syscall.
 static CACHED_MAX_INFLIGHT_BYTES: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
+static CACHED_BATCH_BLOCKS: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
 
 fn encode_channel_capacity(expanded_block_bytes: usize, max_inflight_bytes: usize) -> usize {
     if expanded_block_bytes == 0 {
@@ -46,6 +50,13 @@ fn encode_channel_capacity(expanded_block_bytes: usize, max_inflight_bytes: usiz
         .clamp(1, DEFAULT_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BLOCKS)
 }
 
+fn encode_batch_block_count() -> usize {
+    *CACHED_BATCH_BLOCKS.get_or_init(|| {
+        rustfs_utils::get_env_usize(ENV_RUSTFS_ERASURE_ENCODE_BATCH_BLOCKS, DEFAULT_RUSTFS_ERASURE_ENCODE_BATCH_BLOCKS)
+            .clamp(1, DEFAULT_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BLOCKS)
+    })
+}
+
 fn queued_block_bytes(block: &[Bytes]) -> usize {
     block.iter().map(Bytes::len).sum()
 }
@@ -56,6 +67,16 @@ async fn drain_queued_inflight_bytes(rx: &mut mpsc::Receiver<Vec<Bytes>>) {
     }
 }
 
+fn queued_batch_bytes(batch: &[Vec<Bytes>]) -> usize {
+    batch.iter().map(|block| queued_block_bytes(block)).sum()
+}
+
+async fn drain_queued_batched_inflight_bytes(rx: &mut mpsc::Receiver<Vec<Vec<Bytes>>>) {
+    while let Some(batch) = rx.recv().await {
+        rustfs_io_metrics::remove_ec_encode_inflight_bytes(queued_batch_bytes(&batch));
+    }
+}
+
 fn dominant_error_summary_label(summary: &WriteQuorumFailureSummary) -> &'static str {
     summary.dominant_error_label
 }
@@ -215,6 +236,25 @@ impl<'a> MultiWriter<'a> {
 }
 
 impl Erasure {
+    async fn encode_block(self: Arc<Self>, encode_buf: Vec<u8>, len: usize) -> std::io::Result<(Vec<Bytes>, Vec<u8>)> {
+        let encode_once = move || {
+            let res = self.encode_data(&encode_buf[..len]);
+            (res, encode_buf)
+        };
+
+        let (res, returned_buf) = match tokio::runtime::Handle::current().runtime_flavor() {
+            RuntimeFlavor::MultiThread => tokio::task::block_in_place(encode_once),
+            RuntimeFlavor::CurrentThread => tokio::task::spawn_blocking(encode_once)
+                .await
+                .map_err(|err| std::io::Error::other(format!("EC encode task failed: {err}")))?,
+            _ => tokio::task::spawn_blocking(encode_once)
+                .await
+                .map_err(|err| std::io::Error::other(format!("EC encode task failed: {err}")))?,
+        };
+
+        Ok((res?, returned_buf))
+    }
+
     async fn encode_small_direct<R>(
         self: Arc<Self>,
         mut reader: R,
@@ -294,16 +334,9 @@ impl Erasure {
                     Ok(Some(n)) => {
                         debug_assert!(n > 0, "non-zero block_size prevents zero-length reads");
                         total += n;
-                        let erasure = self.clone();
                         let encode_buf = std::mem::take(&mut buf);
-                        let (res, returned_buf) = tokio::task::spawn_blocking(move || {
-                            let res = erasure.encode_data(&encode_buf[..n]);
-                            (res, encode_buf)
-                        })
-                        .await
-                        .map_err(|err| std::io::Error::other(format!("EC encode task failed: {err}")))?;
+                        let (res, returned_buf) = self.clone().encode_block(encode_buf, n).await?;
                         buf = returned_buf;
-                        let res = res?;
                         let queued_bytes = queued_block_bytes(&res);
                         rustfs_io_metrics::add_ec_encode_inflight_bytes(queued_bytes);
                         if let Err(err) = tx.send(res).await {
@@ -363,6 +396,121 @@ impl Erasure {
         Ok((reader, total))
     }
 
+    pub async fn encode_batched<R>(
+        self: Arc<Self>,
+        mut reader: R,
+        writers: &mut [Option<BitrotWriterWrapper>],
+        quorum: usize,
+    ) -> std::io::Result<(R, usize)>
+    where
+        R: AsyncRead + Send + Sync + Unpin + 'static,
+    {
+        if self.block_size == 0 {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "erasure block_size must be non-zero",
+            ));
+        }
+
+        let expanded_block_bytes = self.shard_size().saturating_mul(self.total_shard_count());
+        let max_inflight_bytes = *CACHED_MAX_INFLIGHT_BYTES.get_or_init(|| {
+            rustfs_utils::get_env_usize(
+                ENV_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BYTES,
+                DEFAULT_RUSTFS_ERASURE_ENCODE_MAX_INFLIGHT_BYTES,
+            )
+        });
+        let inflight_blocks = encode_channel_capacity(expanded_block_bytes, max_inflight_bytes);
+        let batch_blocks = encode_batch_block_count().min(inflight_blocks);
+        let channel_capacity = inflight_blocks.div_ceil(batch_blocks).max(1);
+        let (tx, mut rx) = mpsc::channel::<Vec<Vec<Bytes>>>(channel_capacity);
+
+        let task = tokio::spawn(async move {
+            let block_size = self.block_size;
+            let mut total = 0;
+            let mut buf = vec![0u8; block_size];
+            let mut pending_batch = Vec::with_capacity(batch_blocks);
+            let mut pending_batch_bytes = 0usize;
+            loop {
+                match rustfs_utils::read_full_or_eof(&mut reader, &mut buf).await {
+                    Ok(Some(n)) => {
+                        debug_assert!(n > 0, "non-zero block_size prevents zero-length reads");
+                        total += n;
+                        let encode_buf = std::mem::take(&mut buf);
+                        let (res, returned_buf) = self.clone().encode_block(encode_buf, n).await?;
+                        buf = returned_buf;
+                        let queued_bytes = queued_block_bytes(&res);
+                        pending_batch_bytes = pending_batch_bytes.saturating_add(queued_bytes);
+                        pending_batch.push(res);
+
+                        if pending_batch.len() >= batch_blocks {
+                            rustfs_io_metrics::add_ec_encode_inflight_bytes(pending_batch_bytes);
+                            if let Err(err) = tx.send(pending_batch).await {
+                                rustfs_io_metrics::remove_ec_encode_inflight_bytes(pending_batch_bytes);
+                                return Err(std::io::Error::other(format!("Failed to send encoded data : {err}")));
+                            }
+                            pending_batch = Vec::with_capacity(batch_blocks);
+                            pending_batch_bytes = 0;
+                        }
+                    }
+                    Ok(None) => {
+                        break;
+                    }
+                    Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
+                        if let Some(inner) = e.get_ref()
+                            && rustfs_rio::is_checksum_mismatch(inner)
+                        {
+                            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()));
+                        }
+                        return Err(e);
+                    }
+                    Err(e) => {
+                        return Err(e);
+                    }
+                }
+            }
+
+            if !pending_batch.is_empty() {
+                rustfs_io_metrics::add_ec_encode_inflight_bytes(pending_batch_bytes);
+                if let Err(err) = tx.send(pending_batch).await {
+                    rustfs_io_metrics::remove_ec_encode_inflight_bytes(pending_batch_bytes);
+                    return Err(std::io::Error::other(format!("Failed to send encoded data : {err}")));
+                }
+            }
+
+            Ok((reader, total))
+        });
+
+        let mut writers = MultiWriter::new(writers, quorum);
+        let mut write_err = None;
+
+        while let Some(batch) = rx.recv().await {
+            rustfs_io_metrics::remove_ec_encode_inflight_bytes(queued_batch_bytes(&batch));
+            for block in batch {
+                if let Err(err) = writers.write(block).await {
+                    write_err = Some(err);
+                    break;
+                }
+            }
+            if write_err.is_some() {
+                break;
+            }
+        }
+
+        if let Some(err) = write_err {
+            task.abort();
+            let _ = task.await;
+            drain_queued_batched_inflight_bytes(&mut rx).await;
+            if let Err(shutdown_err) = writers.shutdown().await {
+                error!("failed to shutdown erasure writers after write error: {:?}", shutdown_err);
+            }
+            return Err(err);
+        }
+
+        let (reader, total) = task.await??;
+        writers.shutdown().await?;
+        Ok((reader, total))
+    }
+
     /// Fast path for small inline objects: skip tokio::spawn + mpsc channel.
     /// Reads all data, encodes directly, writes shards sequentially.
     pub async fn encode_inline_small<R>(
@@ -497,6 +645,23 @@ mod tests {
         assert!(err.to_string().contains("block_size"));
     }
 
+    #[tokio::test(flavor = "current_thread")]
+    async fn encode_works_on_current_thread_runtime() {
+        let committed = Arc::new(Mutex::new(Vec::new()));
+        let writer = DeferredCommitWriter::new(committed);
+        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(Cursor::new(b"current-thread payload".to_vec()));
+        let (_reader, written) = erasure.encode(reader, &mut writers, 1).await.unwrap();
+
+        assert_eq!(written, b"current-thread payload".len());
+    }
+
     /// encode_inline_small: empty reader returns (reader, 0) without writing to any shard.
     #[tokio::test]
     async fn encode_inline_small_empty_stream_returns_zero() {