diff --git a/pageserver/src/tenant/blob_io.rs b/pageserver/src/tenant/blob_io.rs index 8cf3c548c9..5bfea0d1c1 100644 --- a/pageserver/src/tenant/blob_io.rs +++ b/pageserver/src/tenant/blob_io.rs @@ -90,10 +90,18 @@ impl Header { } } +#[derive(Debug, thiserror::Error)] +pub enum BlobWriterError { + #[error("flush task cancelled")] + Cancelled, + #[error(transparent)] + Other(anyhow::Error), +} + #[derive(Debug, thiserror::Error)] pub enum WriteBlobError { #[error(transparent)] - Flush(FlushTaskError), + Flush(BlobWriterError), #[error("blob too large ({len} bytes)")] BlobTooLarge { len: usize }, #[error(transparent)] @@ -238,14 +246,16 @@ where cancel: CancellationToken, ctx: &RequestContext, flush_task_span: tracing::Span, - ) -> anyhow::Result { + ) -> Result { + let gate_token = gate.enter().map_err(|e| BlobWriterError::Other(e.into()))?; + Ok(Self { io_buf: Some(BytesMut::new()), writer: BufferedWriter::new( file, start_offset, || IoBufferMut::with_capacity(Self::CAPACITY), - gate.enter()?, + gate_token, cancel, ctx, flush_task_span, @@ -265,13 +275,16 @@ where &mut self, src_buf: FullSlice, ctx: &RequestContext, - ) -> (FullSlice, Result<(), FlushTaskError>) { + ) -> (FullSlice, Result<(), BlobWriterError>) { let res = self .writer // TODO: why are we taking a FullSlice if we're going to pass a borrow downstack? // Can remove all the complexity around owned buffers upstack .write_buffered_borrowed(&src_buf, ctx) .await + .map_err(|e| match e { + FlushTaskError::Cancelled => BlobWriterError::Cancelled, + }) .map(|len| { self.offset += len as u64; }); @@ -418,8 +431,10 @@ where self, mode: BufferedWriterShutdownMode, ctx: &RequestContext, - ) -> Result { - let (_, file) = self.writer.shutdown(mode, ctx).await?; + ) -> Result { + let (_, file) = self.writer.shutdown(mode, ctx).await.map_err(|e| match e { + FlushTaskError::Cancelled => BlobWriterError::Cancelled, + })?; Ok(file) } } @@ -467,8 +482,11 @@ pub(crate) mod tests { .await?, gate.enter()?, ); - let mut wtr = - BlobWriter::new(file, 0, &gate, cancel.clone(), ctx, info_span!("test")).unwrap(); + let mut wtr = BlobWriter::new(file, 0, &gate, cancel.clone(), ctx, info_span!("test")) + .map_err(|e| match e { + BlobWriterError::Cancelled => anyhow::anyhow!("flush task cancelled"), + BlobWriterError::Other(err) => err, + })?; for blob in blobs.iter() { let (_, res) = if compression { let res = wtr @@ -490,7 +508,11 @@ pub(crate) mod tests { BufferedWriterShutdownMode::ZeroPadToNextMultiple(PAGE_SZ), ctx, ) - .await?; + .await + .map_err(|e| match e { + BlobWriterError::Cancelled => anyhow::anyhow!("flush task cancelled"), + BlobWriterError::Other(err) => err, + })?; file.disarm_into_inner() }; Ok((temp_dir, pathbuf, offsets)) diff --git a/pageserver/src/tenant/storage_layer/delta_layer.rs b/pageserver/src/tenant/storage_layer/delta_layer.rs index 607b0d513c..5594a596b1 100644 --- a/pageserver/src/tenant/storage_layer/delta_layer.rs +++ b/pageserver/src/tenant/storage_layer/delta_layer.rs @@ -35,6 +35,8 @@ use std::str::FromStr; use std::sync::Arc; use std::sync::atomic::AtomicU64; +use crate::tenant::blob_io::BlobWriterError; + use anyhow::{Context, Result, bail, ensure}; use camino::{Utf8Path, Utf8PathBuf}; use futures::StreamExt; @@ -448,7 +450,11 @@ impl DeltaLayerWriterInner { cancel, ctx, info_span!(parent: None, "delta_layer_writer_flush_task", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), timeline_id=%timeline_id, path = %path), - )?; + ) + .map_err(|e| match e { + BlobWriterError::Cancelled => anyhow::anyhow!("flush task cancelled"), + BlobWriterError::Other(err) => err, + })?; // Initialize the b-tree index builder let block_buf = BlockBuf::new(); diff --git a/pageserver/src/tenant/storage_layer/image_layer.rs b/pageserver/src/tenant/storage_layer/image_layer.rs index 2f7c5715bb..8050152b33 100644 --- a/pageserver/src/tenant/storage_layer/image_layer.rs +++ b/pageserver/src/tenant/storage_layer/image_layer.rs @@ -33,6 +33,8 @@ use std::str::FromStr; use std::sync::Arc; use std::sync::atomic::AtomicU64; +use crate::tenant::blob_io::BlobWriterError; + use anyhow::{Context, Result, bail, ensure}; use bytes::Bytes; use camino::{Utf8Path, Utf8PathBuf}; @@ -814,7 +816,11 @@ impl ImageLayerWriterInner { cancel, ctx, info_span!(parent: None, "image_layer_writer_flush_task", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), timeline_id=%timeline_id, path = %path), - )?; + ) + .map_err(|e| match e { + BlobWriterError::Cancelled => anyhow::anyhow!("flush task cancelled"), + BlobWriterError::Other(err) => err, + })?; // Initialize the b-tree index builder let block_buf = BlockBuf::new();