From a5102918a1000eb1a9d92077401f25264dc0a24d Mon Sep 17 00:00:00 2001
From: hoslo <ys952572463@gmail.com>
Date: Mon, 4 Mar 2024 15:57:56 +0800
Subject: [PATCH] fix(services/hdfs): fix poll_close when retry (#4309)

---
 core/src/services/hdfs/writer.rs | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/core/src/services/hdfs/writer.rs b/core/src/services/hdfs/writer.rs
index 6c77097d842f..101200a21753 100644
--- a/core/src/services/hdfs/writer.rs
+++ b/core/src/services/hdfs/writer.rs
@@ -36,7 +36,7 @@ pub struct HdfsWriter<F> {
     tmp_path: Option<String>,
     f: Option<F>,
     client: Arc<hdrs::Client>,
-    fut: Option<BoxFuture<'static, Result<()>>>,
+    fut: Option<BoxFuture<'static, (F, Result<()>)>>,
 }
 
 /// # Safety
@@ -76,7 +76,11 @@ impl oio::Write for HdfsWriter<hdrs::AsyncFile> {
             if let Some(fut) = self.fut.as_mut() {
                 let res = ready!(fut.poll_unpin(cx));
                 self.fut = None;
-                return Poll::Ready(res);
+                if let Err(e) = res.1 {
+                    self.f = Some(res.0);
+                    return Poll::Ready(Err(e));
+                }
+                return Poll::Ready(Ok(()));
             }
 
             let mut f = self.f.take().expect("HdfsWriter must be initialized");
@@ -86,15 +90,21 @@ impl oio::Write for HdfsWriter<hdrs::AsyncFile> {
             let client = self.client.clone();
 
             self.fut = Some(Box::pin(async move {
-                f.close().await.map_err(new_std_io_error)?;
+                if let Err(e) = f.close().await.map_err(new_std_io_error) {
+                    // Reserve the original file handle for retry.
+                    return (f, Err(e));
+                }
 
                 if let Some(tmp_path) = tmp_path {
-                    client
+                    if let Err(e) = client
                         .rename_file(&tmp_path, &target_path)
-                        .map_err(new_std_io_error)?;
+                        .map_err(new_std_io_error)
+                    {
+                        return (f, Err(e));
+                    }
                 }
 
-                Ok(())
+                (f, Ok(()))
             }));
         }
     }