try(localfile): Eliminate data copy when flushing

Author: zuston

src/config.rs

@@ -135,7 +135,7 @@ fn as_default_healthy_check_min_disks() -> i32 {
     1
 }
 fn as_default_disk_write_buf_capacity() -> String {
-    "1M".to_string()
+    "1G".to_string()
 }
 fn as_default_disk_read_buf_capacity() -> String {
     "1M".to_string()

src/store/local/delegator.rs

@@ -11,6 +11,7 @@ use crate::readable_size::ReadableSize;
 use crate::runtime::manager::RuntimeManager;
 use crate::store::local::sync_io::SyncLocalIO;
 use crate::store::local::{FileStat, LocalDiskStorage, LocalIO};
+use crate::store::BytesWrapper;
 use anyhow::Result;
 use async_trait::async_trait;
 use await_tree::InstrumentAwait;
@@ -227,7 +228,7 @@ impl LocalIO for LocalDiskDelegator {
         self.inner.io_handler.create_dir(dir).await
     }
 
-    async fn append(&self, path: &str, data: Bytes) -> Result<()> {
+    async fn append(&self, path: &str, data: BytesWrapper) -> Result<()> {
         // todo: add the concurrency limitation. do we need? may be not.
 
         let timer = LOCALFILE_DISK_APPEND_OPERATION_DURATION

src/store/local/mod.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::store::BytesWrapper;
 use anyhow::Result;
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -30,7 +31,7 @@ pub struct FileStat {
 #[async_trait]
 pub trait LocalIO: Clone {
     async fn create_dir(&self, dir: &str) -> Result<()>;
-    async fn append(&self, path: &str, data: Bytes) -> Result<()>;
+    async fn append(&self, path: &str, data: BytesWrapper) -> Result<()>;
     async fn read(&self, path: &str, offset: i64, length: Option<i64>) -> Result<Bytes>;
     async fn delete(&self, path: &str) -> Result<()>;
     async fn write(&self, path: &str, data: Bytes) -> Result<()>;

src/store/local/sync_io.rs

@@ -1,13 +1,14 @@
 use crate::runtime::RuntimeRef;
 use crate::store::local::{FileStat, LocalIO};
+use crate::store::BytesWrapper;
 use anyhow::anyhow;
 use async_trait::async_trait;
 use bytes::Bytes;
-use std::fs;
 use std::fs::{File, OpenOptions};
 use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
 use std::path::Path;
 use std::sync::Arc;
+use std::{fs, io};
 
 #[derive(Clone)]
 pub struct SyncLocalIO {
@@ -60,28 +61,31 @@ impl LocalIO for SyncLocalIO {
         Ok(())
     }
 
-    async fn append(&self, path: &str, data: Bytes) -> anyhow::Result<()> {
+    async fn append(&self, path: &str, data: BytesWrapper) -> anyhow::Result<()> {
         let path = self.with_root(path);
-        let buf = self.inner.buf_writer_capacity.clone();
+        let buffer_capacity = self.inner.buf_writer_capacity.clone();
 
         let r = self
             .inner
             .write_runtime_ref
             .spawn_blocking(move || {
                 let path = Path::new(&path);
                 let mut file = OpenOptions::new().append(true).create(true).open(path)?;
+                let mut buf_writer = match buffer_capacity {
+                    Some(capacity) => BufWriter::with_capacity(capacity, file),
+                    _ => BufWriter::new(file),
+                };
 
-                match buf {
-                    Some(capacity) => {
-                        let mut buf_writer = BufWriter::with_capacity(capacity, file);
-                        buf_writer.write_all(&data)?;
-                        buf_writer.flush()
-                    }
-                    _ => {
-                        file.write_all(&data)?;
-                        file.flush()
+                match data {
+                    BytesWrapper::Direct(bytes) => buf_writer.write_all(&bytes)?,
+                    BytesWrapper::Composed(composed) => {
+                        for data in composed.iter() {
+                            buf_writer.write_all(data)?;
+                        }
                     }
                 }
+                buf_writer.flush()?;
+                Ok::<(), io::Error>(())
             })
             .await
             .map_err(|e| anyhow!(e))??;
@@ -207,9 +211,12 @@ mod test {
         );
 
         // append
-        base_runtime_ref.block_on(io_handler.append(data_file_name, Bytes::from(vec![0; 1000])))?;
-        base_runtime_ref.block_on(io_handler.append(data_file_name, Bytes::from(vec![0; 1000])))?;
-        base_runtime_ref.block_on(io_handler.append(data_file_name, Bytes::from(vec![0; 1000])))?;
+        base_runtime_ref
+            .block_on(io_handler.append(data_file_name, Bytes::from(vec![0; 1000]).into()))?;
+        base_runtime_ref
+            .block_on(io_handler.append(data_file_name, Bytes::from(vec![0; 1000]).into()))?;
+        base_runtime_ref
+            .block_on(io_handler.append(data_file_name, Bytes::from(vec![0; 1000]).into()))?;
 
         // stat
         let stat = base_runtime_ref.block_on(io_handler.file_stat(data_file_name))?;

src/store/localfile.rs

@@ -244,57 +244,17 @@ impl LocalFileStore {
             }
         }
 
-        let mut index_bytes_holder = BytesMut::new();
-        let mut data_bytes_holder = BytesMut::new();
-
-        let mut total_size = 0u64;
-        for block in blocks {
-            let block_id = block.block_id;
-            let length = block.length;
-            let uncompress_len = block.uncompress_length;
-            let task_attempt_id = block.task_attempt_id;
-            let crc = block.crc;
-
-            total_size += length as u64;
-
-            index_bytes_holder.put_i64(next_offset);
-            index_bytes_holder.put_i32(length);
-            index_bytes_holder.put_i32(uncompress_len);
-            index_bytes_holder.put_i64(crc);
-            index_bytes_holder.put_i64(block_id);
-            index_bytes_holder.put_i64(task_attempt_id);
-
-            let data = &block.data;
-
-            // let actual = get_crc(data);
-            // if crc != actual {
-            //     error!("Found not correct crc value while writing. expected: {}, actual: {} for uid: {:?}", crc, actual, &uid);
-            // }
-
-            data_bytes_holder.put(data.clone());
-            next_offset += length as i64;
-        }
-
-        let data = data_bytes_holder.freeze();
-        if data.len() != total_size as usize {
-            error!(
-                "Found not correct data len. expected: {}. actual: {} for uid: {:?}",
-                total_size,
-                data.len(),
-                &uid
-            );
-        }
-
+        let shuffle_file_format = self.generate_shuffle_file_format(blocks, next_offset)?;
         local_disk
-            .append(&data_file_path, data)
+            .append(&data_file_path, shuffle_file_format.data)
             .instrument_await(format!("data flushing. path: {}", &data_file_path))
             .await?;
         local_disk
-            .append(&index_file_path, index_bytes_holder.freeze())
+            .append(&index_file_path, shuffle_file_format.index)
             .instrument_await(format!("index flushing. path: {}", &index_file_path))
             .await?;
 
-        TOTAL_LOCALFILE_USED.inc_by(total_size);
+        TOTAL_LOCALFILE_USED.inc_by(shuffle_file_format.len as u64);
 
         locked_obj
             .deref()

src/store/mod.rs

@@ -26,8 +26,8 @@ pub mod memory;
 mod spill;
 
 use crate::app::{
-    PartitionedUId, PurgeDataContext, ReadingIndexViewContext, ReadingViewContext,
-    RegisterAppContext, ReleaseTicketContext, RequireBufferContext, WritingViewContext,
+    PurgeDataContext, ReadingIndexViewContext, ReadingViewContext, RegisterAppContext,
+    ReleaseTicketContext, RequireBufferContext, WritingViewContext,
 };
 use crate::config::{Config, StorageType};
 use crate::error::WorkerError;
@@ -38,12 +38,12 @@ use std::fmt::{Display, Formatter};
 use crate::util::now_timestamp_as_sec;
 use anyhow::Result;
 use async_trait::async_trait;
-use bytes::Bytes;
+use bytes::{BufMut, Bytes, BytesMut};
 
 use crate::composed_bytes::ComposedBytes;
 use crate::runtime::manager::RuntimeManager;
-use crate::store::mem::buffer::BatchMemoryBlock;
 use crate::store::spill::SpillWritingViewContext;
+use crate::store::BytesWrapper::{Composed, Direct};
 use std::sync::Arc;
 
 #[derive(Debug)]
@@ -248,6 +248,53 @@ pub trait Store {
     async fn name(&self) -> StorageType;
 
     async fn spill_insert(&self, ctx: SpillWritingViewContext) -> Result<(), WorkerError>;
+
+    fn generate_shuffle_file_format(
+        &self,
+        blocks: Vec<&Block>,
+        offset: i64,
+    ) -> Result<ShuffleFileFormat> {
+        let mut offset = offset;
+
+        let mut index_bytes_holder = BytesMut::new();
+        let mut data_chain = Vec::with_capacity(blocks.len());
+
+        let mut total_size = 0;
+        for block in blocks {
+            let block_id = block.block_id;
+            let length = block.length;
+            let uncompress_len = block.uncompress_length;
+            let task_attempt_id = block.task_attempt_id;
+            let crc = block.crc;
+
+            total_size += length as usize;
+
+            index_bytes_holder.put_i64(offset);
+            index_bytes_holder.put_i32(length);
+            index_bytes_holder.put_i32(uncompress_len);
+            index_bytes_holder.put_i64(crc);
+            index_bytes_holder.put_i64(block_id);
+            index_bytes_holder.put_i64(task_attempt_id);
+
+            let data = &block.data;
+            data_chain.push(data.clone());
+            offset += length as i64;
+        }
+
+        Ok(ShuffleFileFormat {
+            data: Composed(ComposedBytes::from(data_chain, total_size)),
+            index: Direct(index_bytes_holder.into()),
+            len: total_size,
+            offset,
+        })
+    }
+}
+
+pub struct ShuffleFileFormat {
+    data: BytesWrapper,
+    index: BytesWrapper,
+    len: usize,
+    offset: i64,
 }
 
 pub trait Persistent {}