From 78c8e08341140270cfc8f07604a4af48385ae7d1 Mon Sep 17 00:00:00 2001 From: Johanna Larsson Date: Sat, 26 Sep 2026 11:36:39 +0100 Subject: [PATCH] fix: abort multipart upload on errors adds a test that proves it works too! --- crates/tranquil-storage/src/lib.rs | 86 ++++++++++++++--------------- crates/tranquil-storage/tests/s3.rs | 43 ++++++++++++++- 2 files changed, 84 insertions(+), 45 deletions(-) diff --git a/crates/tranquil-storage/src/lib.rs b/crates/tranquil-storage/src/lib.rs index b0d36da..1940381 100644 --- a/crates/tranquil-storage/src/lib.rs +++ b/crates/tranquil-storage/src/lib.rs @@ -329,7 +329,7 @@ mod s3 { .await; }; - let result: Result = { + let result: Result = 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> { diff --git a/crates/tranquil-storage/tests/s3.rs b/crates/tranquil-storage/tests/s3.rs index 0485fb2..0c31afd 100644 --- a/crates/tranquil-storage/tests/s3.rs +++ b/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(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, + ) -> S3Result> { + self.0.create_multipart_upload(req).await + } + + async fn abort_multipart_upload( + &self, + req: S3Request, + ) -> S3Result> { + 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); +}