feat: fs add concurrent write (#4817)

core/Cargo.toml

@@ -134,7 +134,7 @@ services-dbfs = []
 services-dropbox = []
 services-etcd = ["dep:etcd-client", "dep:bb8"]
 services-foundationdb = ["dep:foundationdb"]
-services-fs = ["tokio/fs"]
+services-fs = ["tokio/fs", "internal-tokio-rt"]
 services-ftp = ["dep:suppaftp", "dep:bb8", "dep:async-tls"]
 services-gcs = [
   "dep:reqsign",

core/src/raw/oio/write/position_write.rs

@@ -49,6 +49,12 @@ pub trait PositionWrite: Send + Sync + Unpin + 'static {
         offset: u64,
         buf: Buffer,
     ) -> impl Future<Output = Result<()>> + MaybeSend;
+
+    /// close is used to close the underlying file.
+    fn close(&self) -> impl Future<Output = Result<()>> + MaybeSend;
+
+    /// abort is used to abort the underlying abort.
+    fn abort(&self) -> impl Future<Output = Result<()>> + MaybeSend;
 }
 
 struct WriteInput<W: PositionWrite> {
@@ -151,10 +157,14 @@ impl<W: PositionWrite> oio::Write for PositionWriter<W> {
             self.w.write_all_at(offset, buffer).await?;
             self.cache = None;
         }
+        self.w.close().await?;
         Ok(())
     }
 
     async fn abort(&mut self) -> Result<()> {
+        self.tasks.clear();
+        self.cache = None;
+        self.w.abort().await?;
         Ok(())
     }
 }
@@ -216,6 +226,14 @@ mod tests {
 
             Ok(())
         }
+
+        async fn close(&self) -> Result<()> {
+            Ok(())
+        }
+
+        async fn abort(&self) -> Result<()> {
+            Ok(())
+        }
     }
 
     #[tokio::test]

core/src/services/fs/backend.rs

@@ -32,6 +32,7 @@ use super::core::*;
 use super::lister::FsLister;
 use super::reader::FsReader;
 use super::writer::FsWriter;
+use super::writer::FsWriters;
 
 /// config for file system
 #[derive(Default, Deserialize, Debug)]
@@ -174,7 +175,7 @@ pub struct FsBackend {
 
 impl Access for FsBackend {
     type Reader = FsReader<tokio::fs::File>;
-    type Writer = FsWriter<tokio::fs::File>;
+    type Writer = FsWriters;
     type Lister = Option<FsLister<tokio::fs::ReadDir>>;
     type BlockingReader = FsReader<std::fs::File>;
     type BlockingWriter = FsWriter<std::fs::File>;
@@ -318,7 +319,19 @@ impl Access for FsBackend {
             .await
             .map_err(new_std_io_error)?;
 
-        Ok((RpWrite::new(), FsWriter::new(target_path, tmp_path, f)))
+        let w = FsWriter::new(target_path, tmp_path, f);
+
+        let w = if op.append() {
+            FsWriters::One(w)
+        } else {
+            FsWriters::Two(oio::PositionWriter::new(
+                w,
+                op.executor().cloned(),
+                op.concurrent(),
+            ))
+        };
+
+        Ok((RpWrite::default(), w))
     }
 
     async fn delete(&self, path: &str, _: OpDelete) -> Result<RpDelete> {

core/src/services/fs/writer.rs

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fs::File;
 use std::io::Write;
+
 use std::path::PathBuf;
 
 use bytes::Buf;
@@ -24,6 +26,9 @@ use tokio::io::AsyncWriteExt;
 use crate::raw::*;
 use crate::*;
 
+pub type FsWriters =
+    TwoWays<FsWriter<tokio::fs::File>, oio::PositionWriter<FsWriter<tokio::fs::File>>>;
+
 pub struct FsWriter<F> {
     target_path: PathBuf,
     tmp_path: Option<PathBuf>,
@@ -101,3 +106,79 @@ impl oio::BlockingWrite for FsWriter<std::fs::File> {
         Ok(())
     }
 }
+
+impl oio::PositionWrite for FsWriter<tokio::fs::File> {
+    async fn write_all_at(&self, offset: u64, buf: Buffer) -> Result<()> {
+        let f = self.f.as_ref().expect("FsWriter must be initialized");
+
+        let f = f
+            .try_clone()
+            .await
+            .map_err(new_std_io_error)?
+            .into_std()
+            .await;
+
+        tokio::task::spawn_blocking(move || {
+            let mut buf = buf;
+            let mut offset = offset;
+            while !buf.is_empty() {
+                match write_at(&f, buf.chunk(), offset) {
+                    Ok(n) => {
+                        buf.advance(n);
+                        offset += n as u64
+                    }
+                    Err(e) => return Err(e),
+                }
+            }
+            Ok(())
+        })
+        .await
+        .map_err(new_task_join_error)?
+    }
+
+    async fn close(&self) -> Result<()> {
+        let f = self.f.as_ref().expect("FsWriter must be initialized");
+
+        let mut f = f
+            .try_clone()
+            .await
+            .map_err(new_std_io_error)?
+            .into_std()
+            .await;
+
+        f.flush().map_err(new_std_io_error)?;
+        f.sync_all().map_err(new_std_io_error)?;
+
+        if let Some(tmp_path) = &self.tmp_path {
+            tokio::fs::rename(tmp_path, &self.target_path)
+                .await
+                .map_err(new_std_io_error)?;
+        }
+        Ok(())
+    }
+
+    async fn abort(&self) -> Result<()> {
+        if let Some(tmp_path) = &self.tmp_path {
+            tokio::fs::remove_file(tmp_path)
+                .await
+                .map_err(new_std_io_error)
+        } else {
+            Err(Error::new(
+                ErrorKind::Unsupported,
+                "Fs doesn't support abort if atomic_write_dir is not set",
+            ))
+        }
+    }
+}
+
+#[cfg(windows)]
+fn write_at(f: &File, buf: &[u8], offset: u64) -> Result<usize> {
+    use std::os::windows::fs::FileExt;
+    f.seek_write(buf, offset).map_err(new_std_io_error)
+}
+
+#[cfg(unix)]
+fn write_at(f: &File, buf: &[u8], offset: u64) -> Result<usize> {
+    use std::os::unix::fs::FileExt;
+    f.write_at(buf, offset).map_err(new_std_io_error)
+}