feat(services/hdfs): Atomic write for hdfs

core/src/services/hdfs/backend.rs

@@ -25,6 +25,7 @@ use async_trait::async_trait;
 use futures::AsyncWriteExt;
 use log::debug;
 use serde::Deserialize;
+use uuid::Uuid;
 
 use super::lister::HdfsLister;
 use super::writer::HdfsWriter;
@@ -48,6 +49,8 @@ pub struct HdfsConfig {
     pub user: Option<String>,
     /// enable the append capacity
     pub enable_append: bool,
+    /// atomic_write_dir of this backend
+    pub atomic_write_dir: Option<String>,
 }
 
 impl Debug for HdfsConfig {
@@ -133,6 +136,21 @@ impl HdfsBuilder {
         self.config.enable_append = enable_append;
         self
     }
+
+    /// Set temp dir for atomic write.
+    ///
+    /// # Notes
+    ///
+    /// - When append is enabled, we will not use atomic write
+    /// to avoid data loss and performance issue.
+    pub fn atomic_write_dir(&mut self, dir: &str) -> &mut Self {
+        self.config.atomic_write_dir = if dir.is_empty() {
+            None
+        } else {
+            Some(String::from(dir))
+        };
+        self
+    }
 }
 
 impl Builder for HdfsBuilder {
@@ -181,19 +199,40 @@ impl Builder for HdfsBuilder {
             }
         }
 
+        let atomic_write_dir = self.config.atomic_write_dir.take();
+
+        // If atomic write dir is not exist, we must create it.
+        if let Some(d) = &atomic_write_dir {
+            if let Err(e) = client.metadata(&d) {
+                if e.kind() == io::ErrorKind::NotFound {
+                    client.create_dir(&d).map_err(new_std_io_error)?
+                }
+            }
+        }
+
         debug!("backend build finished: {:?}", &self);
         Ok(HdfsBackend {
             root,
+            atomic_write_dir,
             client: Arc::new(client),
             enable_append: self.config.enable_append,
         })
     }
 }
 
+#[inline]
+fn tmp_file_of(path: &str) -> String {
+    let name = get_basename(path);
+    let uuid = Uuid::new_v4().to_string();
+
+    format!("{name}.{uuid}")
+}
+
 /// Backend for hdfs services.
 #[derive(Debug, Clone)]
 pub struct HdfsBackend {
     root: String,
+    atomic_write_dir: Option<String>,
     client: Arc<hdrs::Client>,
     enable_append: bool,
 }
@@ -263,15 +302,35 @@ impl Accessor for HdfsBackend {
     }
 
     async fn write(&self, path: &str, op: OpWrite) -> Result<(RpWrite, Self::Writer)> {
-        let p = build_rooted_abs_path(&self.root, path);
+        let (target_path, tmp_path) = if let Some(atomic_write_dir) = &self.atomic_write_dir {
+            let target_path = build_rooted_abs_path(&self.root, path);
+            let tmp_path = build_rooted_abs_path(atomic_write_dir, &tmp_file_of(path));
+
+            // If the target file exists, we should append to the end of it directly.
+            if op.append() && self.client.metadata(&target_path).is_ok() {
+                debug!("target_path set : {:?} and tmp_path is none", target_path);
+                (target_path, None)
+            } else {
+                debug!(
+                    "target_path set : {:?} and tmp_path is {:?}",
+                    target_path, tmp_path
+                );
+                (target_path, Some(tmp_path))
+            }
+        } else {
+            let p = build_rooted_abs_path(&self.root, path);
+            debug!("target_path set : {:?} and tmp_path is none", p);
+            (p, None)
+        };
 
-        if let Err(err) = self.client.metadata(&p) {
+        if let Err(err) = self.client.metadata(&target_path) {
+            debug!("Error in metadata : {:?}", err);
             // Early return if other error happened.
             if err.kind() != io::ErrorKind::NotFound {
                 return Err(new_std_io_error(err));
             }
 
-            let parent = get_parent(&p);
+            let parent = get_parent(&target_path);
 
             self.client.create_dir(parent).map_err(new_std_io_error)?;
 
@@ -280,7 +339,7 @@ impl Accessor for HdfsBackend {
                 .open_file()
                 .create(true)
                 .write(true)
-                .async_open(&p)
+                .async_open(&target_path)
                 .await
                 .map_err(new_std_io_error)?;
             f.close().await.map_err(new_std_io_error)?;
@@ -295,11 +354,14 @@ impl Accessor for HdfsBackend {
         }
 
         let f = open_options
-            .async_open(&p)
+            .async_open(tmp_path.as_ref().unwrap_or(&target_path))
             .await
             .map_err(new_std_io_error)?;
 
-        Ok((RpWrite::new(), HdfsWriter::new(f)))
+        Ok((
+            RpWrite::new(),
+            HdfsWriter::new(target_path, tmp_path, f, Arc::clone(&self.client)),
+        ))
     }
 
     async fn rename(&self, from: &str, to: &str, _args: OpRename) -> Result<RpRename> {
@@ -438,23 +500,37 @@ impl Accessor for HdfsBackend {
     }
 
     fn blocking_write(&self, path: &str, op: OpWrite) -> Result<(RpWrite, Self::BlockingWriter)> {
-        let p = build_rooted_abs_path(&self.root, path);
+        let (target_path, tmp_path) = if let Some(atomic_write_dir) = &self.atomic_write_dir {
+            let target_path = build_rooted_abs_path(&self.root, path);
+            let tmp_path = build_rooted_abs_path(atomic_write_dir, &tmp_file_of(path));
 
-        if let Err(err) = self.client.metadata(&p) {
+            // If the target file exists, we should append to the end of it directly.
+            if op.append() && self.client.metadata(&target_path).is_ok() {
+                (target_path, None)
+            } else {
+                (target_path, Some(tmp_path))
+            }
+        } else {
+            let p = build_rooted_abs_path(&self.root, path);
+
+            (p, None)
+        };
+
+        if let Err(err) = self.client.metadata(&target_path) {
             // Early return if other error happened.
             if err.kind() != io::ErrorKind::NotFound {
                 return Err(new_std_io_error(err));
             }
 
-            let parent = get_parent(&p);
+            let parent = get_parent(&target_path);
 
             self.client.create_dir(parent).map_err(new_std_io_error)?;
 
             self.client
                 .open_file()
                 .create(true)
                 .write(true)
-                .open(&p)
+                .open(tmp_path.as_ref().unwrap_or(&target_path))
                 .map_err(new_std_io_error)?;
         }
 
@@ -466,9 +542,14 @@ impl Accessor for HdfsBackend {
             open_options.write(true);
         }
 
-        let f = open_options.open(&p).map_err(new_std_io_error)?;
+        let f = open_options
+            .open(tmp_path.as_ref().unwrap_or(&target_path))
+            .map_err(new_std_io_error)?;
 
-        Ok((RpWrite::new(), HdfsWriter::new(f)))
+        Ok((
+            RpWrite::new(),
+            HdfsWriter::new(target_path, tmp_path, f, Arc::clone(&self.client)),
+        ))
     }
 
     fn blocking_rename(&self, from: &str, to: &str, _args: OpRename) -> Result<RpRename> {

core/src/services/hdfs/writer.rs

@@ -15,56 +15,115 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use futures::future::BoxFuture;
 use std::io::Write;
 use std::pin::Pin;
-use std::task::Context;
+use std::sync::Arc;
 use std::task::Poll;
+use std::task::{ready, Context};
 
 use async_trait::async_trait;
-use futures::AsyncWrite;
+use futures::{AsyncWrite, FutureExt};
 
 use crate::raw::*;
 use crate::*;
 
 pub struct HdfsWriter<F> {
-    f: F,
+    target_path: String,
+    tmp_path: Option<String>,
+    f: Option<F>,
+    client: Arc<hdrs::Client>,
+    fut: Option<BoxFuture<'static, Result<()>>>,
 }
 
+/// # Safety
+///
+/// We will only take `&mut Self` reference for HdfsWriter.
+unsafe impl<F> Sync for HdfsWriter<F> {}
+
 impl<F> HdfsWriter<F> {
-    pub fn new(f: F) -> Self {
-        Self { f }
+    pub fn new(
+        target_path: String,
+        tmp_path: Option<String>,
+        f: F,
+        client: Arc<hdrs::Client>,
+    ) -> Self {
+        Self {
+            target_path,
+            tmp_path,
+            f: Some(f),
+            client,
+            fut: None,
+        }
     }
 }
 
 #[async_trait]
 impl oio::Write for HdfsWriter<hdrs::AsyncFile> {
     fn poll_write(&mut self, cx: &mut Context<'_>, bs: &dyn oio::WriteBuf) -> Poll<Result<usize>> {
-        Pin::new(&mut self.f)
+        let f = self.f.as_mut().expect("HdfsWriter must be initialized");
+
+        Pin::new(f)
             .poll_write(cx, bs.chunk())
             .map_err(new_std_io_error)
     }
 
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        loop {
+            if let Some(mut fut) = self.fut.as_mut() {
+                let res = ready!(fut.poll_unpin(cx));
+                self.fut = None;
+                return Poll::Ready(res);
+            }
+
+            let mut f = self.f.take().expect("HdfsWriter must be initialized");
+            // Clone client to allow move into the future.
+            let tmp_path = self.tmp_path.clone();
+            let client = self.client.clone();
+            let target_path = self.target_path.clone();
+            // Clone the necessary parts of the context
+            let waker = cx.waker().clone();
+
+            self.fut = Some(Box::pin(async move {
+                // Now use the cloned waker in the async block
+                let mut pinned = std::pin::pin!(f);
+                let _ = pinned.as_mut().poll_close(&mut Context::from_waker(&waker));
+
+                if let Some(tmp_path) = tmp_path {
+                    client
+                        .rename_file(&tmp_path, &target_path)
+                        .map_err(new_std_io_error)?;
+                }
+
+                Ok(())
+            }));
+        }
+    }
+
     fn poll_abort(&mut self, _: &mut Context<'_>) -> Poll<Result<()>> {
         Poll::Ready(Err(Error::new(
             ErrorKind::Unsupported,
             "HdfsWriter doesn't support abort",
         )))
     }
-
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        Pin::new(&mut self.f)
-            .poll_close(cx)
-            .map_err(new_std_io_error)
-    }
 }
 
 impl oio::BlockingWrite for HdfsWriter<hdrs::File> {
     fn write(&mut self, bs: &dyn oio::WriteBuf) -> Result<usize> {
-        self.f.write(bs.chunk()).map_err(new_std_io_error)
+        let f = self.f.as_mut().expect("HdfsWriter must be initialized");
+        f.write(bs.chunk()).map_err(new_std_io_error)
     }
 
     fn close(&mut self) -> Result<()> {
-        self.f.flush().map_err(new_std_io_error)?;
+        let f = self.f.as_mut().expect("HdfsWriter must be initialized");
+        f.flush().map_err(new_std_io_error)?;
+
+        if let Some(tmp_path) = &self.tmp_path {
+            let client = Arc::as_ref(&self.client);
+            client
+                .rename_file(tmp_path, &self.target_path)
+                .map_err(new_std_io_error)?;
+        }
 
         Ok(())
     }