WriteMultipart Abort on MultipartUpload::complete Error (#5974)

* update * another one * more update * another update * debug * debug * some updates * debug * debug * cleanup * cleanup * simplify * address some comments * cleanup on failure * restore abort method * docs
apache · Jul 2, 2024 · 3b93a4b · 3b93a4b
1 parent 6351674
commit 3b93a4b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/object_store/src/buffered.rs b/object_store/src/buffered.rs
@@ -438,7 +438,12 @@ impl AsyncWrite for BufWriter {
                 }
                 BufWriterState::Flush(f) => return f.poll_unpin(cx).map_err(std::io::Error::from),
                 BufWriterState::Write(x) => {
-                    let upload = x.take().unwrap();
+                    let upload = x.take().ok_or_else(|| {
+                        std::io::Error::new(
+                            ErrorKind::InvalidInput,
+                            "Cannot shutdown a writer that has already been shut down",
+                        )
+                    })?;
                     self.state = BufWriterState::Flush(
                         async move {
                             upload.finish().await?;

diff --git a/object_store/src/local.rs b/object_store/src/local.rs
@@ -724,15 +724,15 @@ struct LocalUpload {
 #[derive(Debug)]
 struct UploadState {
     dest: PathBuf,
-    file: Mutex<Option<File>>,
+    file: Mutex<File>,
 }
 
 impl LocalUpload {
     pub fn new(src: PathBuf, dest: PathBuf, file: File) -> Self {
         Self {
             state: Arc::new(UploadState {
                 dest,
-                file: Mutex::new(Some(file)),
+                file: Mutex::new(file),
             }),
             src: Some(src),
             offset: 0,
@@ -748,8 +748,7 @@ impl MultipartUpload for LocalUpload {
 
         let s = Arc::clone(&self.state);
         maybe_spawn_blocking(move || {
-            let mut f = s.file.lock();
-            let file = f.as_mut().context(AbortedSnafu)?;
+            let mut file = s.file.lock();
             file.seek(SeekFrom::Start(offset))
                 .context(SeekSnafu { path: &s.dest })?;
 
@@ -767,9 +766,9 @@ impl MultipartUpload for LocalUpload {
         let s = Arc::clone(&self.state);
         maybe_spawn_blocking(move || {
             // Ensure no inflight writes
-            let f = s.file.lock().take().context(AbortedSnafu)?;
+            let file = s.file.lock();
             std::fs::rename(&src, &s.dest).context(UnableToRenameFileSnafu)?;
-            let metadata = f.metadata().map_err(|e| Error::Metadata {
+            let metadata = file.metadata().map_err(|e| Error::Metadata {
                 source: e.into(),
                 path: src.to_string_lossy().to_string(),
             })?;

diff --git a/object_store/src/upload.rs b/object_store/src/upload.rs
@@ -226,7 +226,15 @@ impl WriteMultipart {
         }
 
         self.wait_for_capacity(0).await?;
-        self.upload.complete().await
+
+        match self.upload.complete().await {
+            Err(e) => {
+                self.tasks.shutdown().await;
+                self.upload.abort().await?;
+                Err(e)
+            }
+            Ok(result) => Ok(result),
+        }
     }
 }