fix: abort multipart upload on errors

adds a test that proves it works too!

crates/tranquil-storage/src/lib.rs

@@ -329,7 +329,7 @@ mod s3 {
                     .await;
             };
 
-            let result: Result<UploadState, StorageError> = {
+            let result: Result<StreamUploadResult, StorageError> = async {
                 let mut state = initial_state;
 
                 while let Some(chunk_result) = stream.next().await {
@@ -357,58 +357,56 @@ mod s3 {
                                 state.part_number += 1;
                             }
                         }
-                        Err(e) => {
-                            abort_upload().await;
-                            return Err(StorageError::Io(e));
-                        }
+                        Err(e) => return Err(StorageError::Io(e)),
                     }
                 }
 
-                Ok(state)
-            };
+                if !state.buffer.is_empty() {
+                    let part = upload_part(
+                        &self.client,
+                        &self.bucket,
+                        key,
+                        &upload_id,
+                        state.part_number,
+                        std::mem::take(&mut state.buffer),
+                    )
+                    .await?;
+                    state.completed_parts.push(part);
+                }
 
-            let mut state = result?;
+                if state.completed_parts.is_empty() {
+                    return Err(StorageError::Other("Empty upload".to_string()));
+                }
 
-            if !state.buffer.is_empty() {
-                let part = upload_part(
-                    &self.client,
-                    &self.bucket,
-                    key,
-                    &upload_id,
-                    state.part_number,
-                    std::mem::take(&mut state.buffer),
-                )
-                .await?;
-                state.completed_parts.push(part);
+                let completed_upload = CompletedMultipartUpload::builder()
+                    .set_parts(Some(state.completed_parts))
+                    .build();
+
+                let path = self.resolve_path(key);
+                self.client
+                    .complete_multipart_upload()
+                    .bucket(&self.bucket)
+                    .key(&path)
+                    .upload_id(&upload_id)
+                    .multipart_upload(completed_upload)
+                    .send()
+                    .await
+                    .map_err(|e| {
+                        StorageError::Backend(format!("Failed to complete multipart upload: {}", e))
+                    })?;
+
+                let hash: [u8; 32] = state.hasher.finalize().into();
+                Ok(StreamUploadResult {
+                    sha256_hash: hash,
+                    size: state.total_size,
+                })
             }
+            .await;
 
-            if state.completed_parts.is_empty() {
+            if result.is_err() {
                 abort_upload().await;
-                return Err(StorageError::Other("Empty upload".to_string()));
             }
-
-            let completed_upload = CompletedMultipartUpload::builder()
-                .set_parts(Some(state.completed_parts))
-                .build();
-
-            let path = self.resolve_path(key);
-            self.client
-                .complete_multipart_upload()
-                .bucket(&self.bucket)
-                .key(&path)
-                .upload_id(&upload_id)
-                .multipart_upload(completed_upload)
-                .send()
-                .await
-                .map_err(|e| {
-                    StorageError::Backend(format!("Failed to complete multipart upload: {}", e))
-                })?;
-
-            let hash: [u8; 32] = state.hasher.finalize().into();
-            Ok(StreamUploadResult {
-                sha256_hash: hash,
-                size: state.total_size,
-            })
+            result
         }
 
         async fn copy(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> {

crates/tranquil-storage/tests/s3.rs

@@ -3,7 +3,12 @@ use bytes::Bytes;
 use hyper_util::rt::{TokioExecutor, TokioIo};
 use hyper_util::server::conn::auto::Builder as ConnBuilder;
 use s3s::auth::SimpleAuth;
+use s3s::dto::{
+    AbortMultipartUploadInput, AbortMultipartUploadOutput, CreateMultipartUploadInput,
+    CreateMultipartUploadOutput,
+};
 use s3s::service::S3ServiceBuilder;
+use s3s::{S3, S3Request, S3Response, S3Result};
 use s3s_fs::FileSystem;
 use sha2::{Digest, Sha256};
 use tempfile::TempDir;
@@ -14,10 +19,14 @@ const BUCKET: &str = "bucket";
 const PREFIX: &str = "prefix";
 
 async fn start_s3() -> (TempDir, S3BlobStorage) {
+    start_s3_with(|fs| fs).await
+}
+
+async fn start_s3_with<T: S3>(wrap: impl FnOnce(FileSystem) -> T) -> (TempDir, S3BlobStorage) {
     let root = tempfile::tempdir().unwrap();
     std::fs::create_dir(root.path().join(BUCKET)).unwrap();
 
-    let mut builder = S3ServiceBuilder::new(FileSystem::new(root.path()).unwrap());
+    let mut builder = S3ServiceBuilder::new(wrap(FileSystem::new(root.path()).unwrap()));
     builder.set_auth(SimpleAuth::from_single("test", "test"));
     let service = builder.build();
 
@@ -127,3 +136,35 @@ async fn put_stream_multipart() {
     assert_eq!(result.sha256_hash[..], Sha256::digest(&expected)[..]);
     assert_eq!(storage.get_bytes("key").await.unwrap(), expected);
 }
+
+struct FailingUploadPart(FileSystem);
+
+#[async_trait::async_trait]
+impl S3 for FailingUploadPart {
+    async fn create_multipart_upload(
+        &self,
+        req: S3Request<CreateMultipartUploadInput>,
+    ) -> S3Result<S3Response<CreateMultipartUploadOutput>> {
+        self.0.create_multipart_upload(req).await
+    }
+
+    async fn abort_multipart_upload(
+        &self,
+        req: S3Request<AbortMultipartUploadInput>,
+    ) -> S3Result<S3Response<AbortMultipartUploadOutput>> {
+        self.0.abort_multipart_upload(req).await
+    }
+}
+
+#[tokio::test]
+async fn put_stream_part_failure_aborts_upload() {
+    let (root, storage) = start_s3_with(FailingUploadPart).await;
+    let chunks = [Ok(Bytes::from("hello"))];
+
+    let result = storage
+        .put_stream("key", Box::pin(futures::stream::iter(chunks)))
+        .await;
+
+    assert!(result.is_err());
+    assert_eq!(std::fs::read_dir(root.path()).unwrap().count(), 1);
+}