diff --git a/vortex-io/src/object_store/write.rs b/vortex-io/src/object_store/write.rs index 3cb73cec331..1d7f42d5df9 100644 --- a/vortex-io/src/object_store/write.rs +++ b/vortex-io/src/object_store/write.rs @@ -21,10 +21,16 @@ use crate::VortexWrite; /// Adapter type to write data through a [`ObjectStore`] instance. /// /// After writing, the caller must make sure to call `shutdown`, in order to ensure the data is actually persisted. +/// +/// A write that fails aborts the upload before returning, because S3 and GCS keep the parts that +/// already landed when a [`MultipartUpload`] is dropped without one. pub struct ObjectStoreWrite { upload: Box, buffer: BytesMut, put_result: Option, + /// Set once the upload has been aborted. Aborting an already-aborted upload is + /// implementation-defined behaviour, so no path may do it twice. + aborted: bool, } const CHUNK_SIZE: usize = 16 * 1024 * 1024; @@ -37,12 +43,27 @@ impl ObjectStoreWrite { upload, buffer: BytesMut::with_capacity(CHUNK_SIZE), put_result: None, + aborted: false, }) } pub fn put_result(&self) -> Option<&PutResult> { self.put_result.as_ref() } + + /// Abort the upload and hand back `error`, which is the one the caller needs to see. + /// + /// A failing abort is only logged: it cannot be recovered from here, and returning it instead + /// would hide why the write failed. + async fn abort_with(&mut self, error: io::Error) -> io::Error { + if !self.aborted { + self.aborted = true; + if let Err(abort_error) = self.upload.abort().await { + tracing::warn!("failed to abort multipart upload: {abort_error}"); + } + } + error + } } impl VortexWrite for ObjectStoreWrite { @@ -61,7 +82,9 @@ impl VortexWrite for ObjectStoreWrite { } } - parts.try_collect::>().await?; + if let Err(error) = parts.try_collect::>().await { + return Err(self.abort_with(error.into()).await); + } Ok(buffer) } @@ -76,31 +99,45 @@ impl VortexWrite for ObjectStoreWrite { parts.push(part_fut); } - parts.try_collect::>().await?; + if let Err(error) = parts.try_collect::>().await { + return Err(self.abort_with(error.into()).await); + } Ok(()) } async fn shutdown(&mut self) -> io::Result<()> { + // `flush` aborts the upload itself when it fails. self.flush().await?; if !self.buffer.is_empty() { let payload = std::mem::take(&mut self.buffer).freeze(); - self.upload - .put_part(PutPayload::from_bytes(payload)) - .await?; + let part = self.upload.put_part(PutPayload::from_bytes(payload)); + if let Err(error) = part.await { + return Err(self.abort_with(error.into()).await); + } } - self.put_result = Some(self.upload.complete().await?); - Ok(()) + let completed = self.upload.complete().await; + match completed { + Ok(put_result) => { + self.put_result = Some(put_result); + Ok(()) + } + Err(error) => Err(self.abort_with(error.into()).await), + } } } #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use async_trait::async_trait; use object_store::ObjectStore; + use object_store::UploadPart; use object_store::local::LocalFileSystem; use object_store::memory::InMemory; use object_store::path::Path; @@ -109,6 +146,94 @@ mod tests { use super::*; + /// A [`MultipartUpload`] that refuses every part and every completion, counting how often it + /// is aborted. S3 and GCS keep already-uploaded parts when an upload is dropped without an + /// abort, and neither `InMemory` nor `LocalFileSystem` exposes staged parts, so the count is + /// the only observable for that cleanup. + #[derive(Debug)] + struct RefusingUpload { + aborts: Arc, + } + + #[async_trait] + impl MultipartUpload for RefusingUpload { + fn put_part(&mut self, _data: PutPayload) -> UploadPart { + Box::pin(std::future::ready(Err(object_store::Error::Generic { + store: "refusing", + source: "put_part refused".into(), + }))) + } + + async fn complete(&mut self) -> object_store::Result { + Err(object_store::Error::Generic { + store: "refusing", + source: "complete refused".into(), + }) + } + + async fn abort(&mut self) -> object_store::Result<()> { + self.aborts.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + fn refusing_writer(aborts: &Arc) -> ObjectStoreWrite { + ObjectStoreWrite { + upload: Box::new(RefusingUpload { + aborts: Arc::clone(aborts), + }), + buffer: BytesMut::with_capacity(CHUNK_SIZE), + put_result: None, + aborted: false, + } + } + + /// The trailing part `shutdown` uploads is the one a small file goes out as, so its failure + /// must still abort. + #[tokio::test] + async fn shutdown_aborts_when_the_trailing_part_fails() -> anyhow::Result<()> { + let aborts = Arc::new(AtomicUsize::new(0)); + let mut writer = refusing_writer(&aborts); + + writer.write_all(vec![0u8; 8]).await?; + let error = writer.shutdown().await.expect_err("the part is refused"); + + assert!(error.to_string().contains("put_part refused")); + assert_eq!(aborts.load(Ordering::SeqCst), 1); + assert!(writer.put_result().is_none()); + Ok(()) + } + + /// Nothing has been staged when only `complete` fails, but the upload is still open, so it + /// has to be aborted too. + #[tokio::test] + async fn shutdown_aborts_when_complete_fails() { + let aborts = Arc::new(AtomicUsize::new(0)); + let mut writer = refusing_writer(&aborts); + + let error = writer.shutdown().await.expect_err("completion is refused"); + + assert!(error.to_string().contains("complete refused")); + assert_eq!(aborts.load(Ordering::SeqCst), 1); + assert!(writer.put_result().is_none()); + } + + /// `flush` aborts on its own failure, and a caller that goes on to `shutdown` must not abort a + /// second time — `object_store` leaves that implementation-defined. + #[tokio::test] + async fn a_second_failure_does_not_abort_twice() -> anyhow::Result<()> { + let aborts = Arc::new(AtomicUsize::new(0)); + let mut writer = refusing_writer(&aborts); + + writer.write_all(vec![0u8; CHUNK_SIZE + 1]).await?; + writer.flush().await.expect_err("the part is refused"); + assert_eq!(aborts.load(Ordering::SeqCst), 1); + + writer.shutdown().await.expect_err("the part is refused"); + assert_eq!(aborts.load(Ordering::SeqCst), 1); + Ok(()) + } + // Note: Concurrent writes test removed because &mut self in write_all already ensures // exclusive access. Multiple writers would need to be created with separate buffers, // which is not the intended use case.