diff --git a/libs/utils/src/sync/duplex/mpsc.rs b/libs/utils/src/sync/duplex/mpsc.rs index 56b4e6d2b3..7c0183e615 100644 --- a/libs/utils/src/sync/duplex/mpsc.rs +++ b/libs/utils/src/sync/duplex/mpsc.rs @@ -26,6 +26,10 @@ impl Duplex { self.tx.send(x).await } + pub fn try_send(&self, x: S) -> Result<(), mpsc::error::TrySendError> { + self.tx.try_send(x) + } + /// Receives the next value for this receiver. /// /// This method returns `None` if the channel has been closed and there are diff --git a/pageserver/benches/bench_ingest.rs b/pageserver/benches/bench_ingest.rs index d66cc01f7d..3108b5351f 100644 --- a/pageserver/benches/bench_ingest.rs +++ b/pageserver/benches/bench_ingest.rs @@ -13,6 +13,7 @@ use pageserver::{page_cache, virtual_file}; use pageserver_api::key::Key; use pageserver_api::shard::TenantShardId; use pageserver_api::value::Value; +use tokio_util::sync::CancellationToken; use utils::bin_ser::BeSer; use utils::id::{TenantId, TimelineId}; use wal_decoder::serialized_batch::SerializedValueBatch; @@ -61,8 +62,18 @@ async fn ingest( RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error).with_scope_debug_tools(); let gate = utils::sync::gate::Gate::default(); + let cancel = CancellationToken::new(); - let layer = InMemoryLayer::create(conf, timeline_id, tenant_shard_id, lsn, &gate, &ctx).await?; + let layer = InMemoryLayer::create( + conf, + timeline_id, + tenant_shard_id, + lsn, + &gate, + &cancel, + &ctx, + ) + .await?; let data = Value::Image(Bytes::from(vec![0u8; put_size])); let data_ser_size = data.serialized_size().unwrap() as usize; @@ -115,7 +126,7 @@ async fn ingest( max_concurrency: NonZeroUsize::new(1).unwrap(), }); let (_desc, path) = layer - .write_to_disk(&ctx, None, l0_flush_state.inner(), &gate) + .write_to_disk(&ctx, None, l0_flush_state.inner(), &gate, cancel.clone()) .await? .unwrap(); tokio::fs::remove_file(path).await?; diff --git a/pageserver/src/tenant/blob_io.rs b/pageserver/src/tenant/blob_io.rs index 7031e94d04..0e0c5a99e0 100644 --- a/pageserver/src/tenant/blob_io.rs +++ b/pageserver/src/tenant/blob_io.rs @@ -15,7 +15,6 @@ //! len >= 128: 1CCCXXXX XXXXXXXX XXXXXXXX XXXXXXXX //! use std::cmp::min; -use std::io::{Error, ErrorKind}; use std::sync::Arc; use async_compression::Level; @@ -23,6 +22,7 @@ use bytes::{BufMut, BytesMut}; use pageserver_api::models::ImageCompressionAlgorithm; use tokio::io::AsyncWriteExt; use tokio_epoll_uring::IoBuf; +use tokio_util::sync::CancellationToken; use tracing::warn; use crate::context::RequestContext; @@ -31,7 +31,7 @@ use crate::tenant::block_io::BlockCursor; use crate::virtual_file::IoBufferMut; use crate::virtual_file::VirtualFile; use crate::virtual_file::owned_buffers_io::io_buf_ext::{FullSlice, IoBufExt}; -use crate::virtual_file::owned_buffers_io::write::BufferedWriter; +use crate::virtual_file::owned_buffers_io::write::{BufferedWriter, FlushTaskError}; #[derive(Copy, Clone, Debug)] pub struct CompressionInfo { @@ -39,6 +39,14 @@ pub struct CompressionInfo { pub compressed_size: Option, } +#[derive(Debug, thiserror::Error)] +pub enum WriteBlobError { + #[error(transparent)] + Flush(FlushTaskError), + #[error("blob too large ({len} bytes)")] + BlobTooLarge { len: usize }, +} + impl BlockCursor<'_> { /// Read a blob into a new buffer. pub async fn read_blob( @@ -174,6 +182,7 @@ impl BlobWriter { file: Arc, start_offset: u64, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ctx: &RequestContext, flush_task_span: tracing::Span, ) -> anyhow::Result { @@ -184,6 +193,7 @@ impl BlobWriter { start_offset, || IoBufferMut::with_capacity(Self::CAPACITY), gate.enter()?, + cancel, ctx, flush_task_span, ), @@ -202,7 +212,7 @@ impl BlobWriter { &mut self, src_buf: FullSlice, ctx: &RequestContext, - ) -> (FullSlice, Result<(), Error>) { + ) -> (FullSlice, Result<(), FlushTaskError>) { let res = self .writer // TODO: why are we taking a FullSlice if we're going to pass a borrow downstack? @@ -222,7 +232,7 @@ impl BlobWriter { &mut self, srcbuf: FullSlice, ctx: &RequestContext, - ) -> (FullSlice, Result) { + ) -> (FullSlice, Result) { let (buf, res) = self .write_blob_maybe_compressed(srcbuf, ctx, ImageCompressionAlgorithm::Disabled) .await; @@ -236,7 +246,10 @@ impl BlobWriter { srcbuf: FullSlice, ctx: &RequestContext, algorithm: ImageCompressionAlgorithm, - ) -> (FullSlice, Result<(u64, CompressionInfo), Error>) { + ) -> ( + FullSlice, + Result<(u64, CompressionInfo), WriteBlobError>, + ) { let offset = self.offset; let mut compression_info = CompressionInfo { written_compressed: false, @@ -252,17 +265,16 @@ impl BlobWriter { if len < 128 { // Short blob. Write a 1-byte length header io_buf.put_u8(len as u8); - (self.write_all(io_buf.slice_len(), ctx).await, srcbuf) + let (slice, res) = self.write_all(io_buf.slice_len(), ctx).await; + let res = res.map_err(WriteBlobError::Flush); + ((slice, res), srcbuf) } else { // Write a 4-byte length header if len > MAX_SUPPORTED_BLOB_LEN { return ( ( io_buf.slice_len(), - Err(Error::new( - ErrorKind::Other, - format!("blob too large ({len} bytes)"), - )), + Err(WriteBlobError::BlobTooLarge { len }), ), srcbuf, ); @@ -296,7 +308,9 @@ impl BlobWriter { assert_eq!(len_buf[0] & 0xf0, 0); len_buf[0] |= high_bit_mask; io_buf.extend_from_slice(&len_buf[..]); - (self.write_all(io_buf.slice_len(), ctx).await, srcbuf) + let (slice, res) = self.write_all(io_buf.slice_len(), ctx).await; + let res = res.map_err(WriteBlobError::Flush); + ((slice, res), srcbuf) } } .await; @@ -311,6 +325,7 @@ impl BlobWriter { } else { self.write_all(srcbuf, ctx).await }; + let res = res.map_err(WriteBlobError::Flush); (srcbuf, res.map(|_| (offset, compression_info))) } @@ -324,7 +339,7 @@ impl BlobWriter { pub async fn into_inner( self, handle_tail: impl FnMut(IoBufferMut) -> Option, - ) -> Result { + ) -> Result { let (_, file) = self.writer.shutdown(handle_tail).await?; Ok(file) } @@ -342,7 +357,7 @@ pub(crate) mod tests { use crate::task_mgr::TaskKind; use crate::tenant::block_io::BlockReaderRef; - async fn round_trip_test(blobs: &[Vec]) -> Result<(), Error> { + async fn round_trip_test(blobs: &[Vec]) -> anyhow::Result<()> { round_trip_test_compressed(blobs, false).await } @@ -350,16 +365,18 @@ pub(crate) mod tests { blobs: &[Vec], compression: bool, ctx: &RequestContext, - ) -> Result<(Utf8TempDir, Utf8PathBuf, Vec), Error> { + ) -> anyhow::Result<(Utf8TempDir, Utf8PathBuf, Vec)> { let temp_dir = camino_tempfile::tempdir()?; let pathbuf = temp_dir.path().join("file"); let gate = utils::sync::gate::Gate::default(); + let cancel = CancellationToken::new(); // Write part (in block to drop the file) let mut offsets = Vec::new(); { let file = Arc::new(VirtualFile::create_v2(pathbuf.as_path(), ctx).await?); - let mut wtr = BlobWriter::new(file, 0, &gate, ctx, info_span!("test")).unwrap(); + let mut wtr = + BlobWriter::new(file, 0, &gate, cancel.clone(), ctx, info_span!("test")).unwrap(); for blob in blobs.iter() { let (_, res) = if compression { let res = wtr @@ -396,7 +413,10 @@ pub(crate) mod tests { Ok((temp_dir, pathbuf, offsets)) } - async fn round_trip_test_compressed(blobs: &[Vec], compression: bool) -> Result<(), Error> { + async fn round_trip_test_compressed( + blobs: &[Vec], + compression: bool, + ) -> anyhow::Result<()> { let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test(); let (_temp_dir, pathbuf, offsets) = @@ -422,14 +442,14 @@ pub(crate) mod tests { } #[tokio::test] - async fn test_one() -> Result<(), Error> { + async fn test_one() -> anyhow::Result<()> { let blobs = &[vec![12, 21, 22]]; round_trip_test(blobs).await?; Ok(()) } #[tokio::test] - async fn test_hello_simple() -> Result<(), Error> { + async fn test_hello_simple() -> anyhow::Result<()> { let blobs = &[ vec![0, 1, 2, 3], b"Hello, World!".to_vec(), @@ -442,7 +462,7 @@ pub(crate) mod tests { } #[tokio::test] - async fn test_really_big_array() -> Result<(), Error> { + async fn test_really_big_array() -> anyhow::Result<()> { let blobs = &[ b"test".to_vec(), random_array(10 * PAGE_SZ), @@ -457,7 +477,7 @@ pub(crate) mod tests { } #[tokio::test] - async fn test_arrays_inc() -> Result<(), Error> { + async fn test_arrays_inc() -> anyhow::Result<()> { let blobs = (0..PAGE_SZ / 8) .map(|v| random_array(v * 16)) .collect::>(); @@ -466,7 +486,7 @@ pub(crate) mod tests { } #[tokio::test] - async fn test_arrays_random_size() -> Result<(), Error> { + async fn test_arrays_random_size() -> anyhow::Result<()> { let mut rng = rand::rngs::StdRng::seed_from_u64(42); let blobs = (0..1024) .map(|_| { @@ -483,7 +503,7 @@ pub(crate) mod tests { } #[tokio::test] - async fn test_arrays_page_boundary() -> Result<(), Error> { + async fn test_arrays_page_boundary() -> anyhow::Result<()> { let blobs = &[ random_array(PAGE_SZ - 4), random_array(PAGE_SZ - 4), diff --git a/pageserver/src/tenant/ephemeral_file.rs b/pageserver/src/tenant/ephemeral_file.rs index 7ed9c3cc11..19b5165d0f 100644 --- a/pageserver/src/tenant/ephemeral_file.rs +++ b/pageserver/src/tenant/ephemeral_file.rs @@ -9,6 +9,7 @@ use camino::Utf8PathBuf; use num_traits::Num; use pageserver_api::shard::TenantShardId; use tokio_epoll_uring::{BoundedBuf, Slice}; +use tokio_util::sync::CancellationToken; use tracing::{error, info_span}; use utils::id::TimelineId; @@ -19,7 +20,7 @@ use crate::page_cache; use crate::tenant::storage_layer::inmemory_layer::vectored_dio_read::File; use crate::virtual_file::owned_buffers_io::io_buf_aligned::IoBufAlignedMut; use crate::virtual_file::owned_buffers_io::slice::SliceMutExt; -use crate::virtual_file::owned_buffers_io::write::Buffer; +use crate::virtual_file::owned_buffers_io::write::{Buffer, FlushTaskError}; use crate::virtual_file::{self, IoBufferMut, VirtualFile, owned_buffers_io}; pub struct EphemeralFile { @@ -40,6 +41,7 @@ impl EphemeralFile { tenant_shard_id: TenantShardId, timeline_id: TimelineId, gate: &utils::sync::gate::Gate, + cancel: &CancellationToken, ctx: &RequestContext, ) -> anyhow::Result { static NEXT_FILENAME: AtomicU64 = AtomicU64::new(1); @@ -76,6 +78,7 @@ impl EphemeralFile { 0, || IoBufferMut::with_capacity(TAIL_SZ), gate.enter()?, + cancel.child_token(), ctx, info_span!(parent: None, "ephemeral_file_buffered_writer", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), timeline_id=%timeline_id, path = %filename), ), @@ -102,6 +105,14 @@ impl Drop for EphemeralFile { } } +#[derive(Debug, thiserror::Error)] +pub(crate) enum EphemeralFileWriteError { + #[error("{0}")] + TooLong(String), + #[error("cancelled")] + Cancelled, +} + impl EphemeralFile { pub(crate) fn len(&self) -> u64 { self.bytes_written @@ -134,7 +145,7 @@ impl EphemeralFile { &mut self, srcbuf: &[u8], ctx: &RequestContext, - ) -> std::io::Result { + ) -> Result { let (pos, control) = self.write_raw_controlled(srcbuf, ctx).await?; if let Some(control) = control { control.release().await; @@ -146,24 +157,24 @@ impl EphemeralFile { &mut self, srcbuf: &[u8], ctx: &RequestContext, - ) -> std::io::Result<(u64, Option)> { + ) -> Result<(u64, Option), EphemeralFileWriteError> { let pos = self.bytes_written; let new_bytes_written = pos.checked_add(srcbuf.len().into_u64()).ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::Other, - format!( - "write would grow EphemeralFile beyond u64::MAX: len={pos} writen={srcbuf_len}", - srcbuf_len = srcbuf.len(), - ), - ) + EphemeralFileWriteError::TooLong(format!( + "write would grow EphemeralFile beyond u64::MAX: len={pos} writen={srcbuf_len}", + srcbuf_len = srcbuf.len(), + )) })?; // Write the payload let (nwritten, control) = self .buffered_writer .write_buffered_borrowed_controlled(srcbuf, ctx) - .await?; + .await + .map_err(|e| match e { + FlushTaskError::Cancelled => EphemeralFileWriteError::Cancelled, + })?; assert_eq!( nwritten, srcbuf.len(), @@ -185,8 +196,14 @@ impl super::storage_layer::inmemory_layer::vectored_dio_read::File for Ephemeral ) -> std::io::Result<(tokio_epoll_uring::Slice, usize)> { let submitted_offset = self.buffered_writer.bytes_submitted(); - let mutable = self.buffered_writer.inspect_mutable(); - let mutable = &mutable[0..mutable.pending()]; + let mutable = match self.buffered_writer.inspect_mutable() { + Some(mutable) => &mutable[0..mutable.pending()], + None => { + // Timeline::cancel and hence buffered writer flush was cancelled. + // Remain read-available while timeline is shutting down. + &[] + } + }; let maybe_flushed = self.buffered_writer.inspect_maybe_flushed(); @@ -217,7 +234,6 @@ impl super::storage_layer::inmemory_layer::vectored_dio_read::File for Ephemeral let (written_range, maybe_flushed_range) = { if maybe_flushed.is_some() { // [ written ][ maybe_flushed ][ mutable ] - // <- TAIL_SZ -><- TAIL_SZ -> // ^ // `submitted_offset` // <++++++ on disk +++++++????????????????> @@ -233,7 +249,6 @@ impl super::storage_layer::inmemory_layer::vectored_dio_read::File for Ephemeral ) } else { // [ written ][ mutable ] - // <- TAIL_SZ -> // ^ // `submitted_offset` // <++++++ on disk +++++++++++++++++++++++> @@ -367,8 +382,9 @@ mod tests { harness("ephemeral_file_holds_gate_open").unwrap(); let gate = utils::sync::gate::Gate::default(); + let cancel = CancellationToken::new(); - let file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &ctx) + let file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &cancel, &ctx) .await .unwrap(); @@ -398,12 +414,13 @@ mod tests { let (conf, tenant_id, timeline_id, ctx) = harness("test_ephemeral_file_basics").unwrap(); let gate = utils::sync::gate::Gate::default(); + let cancel = CancellationToken::new(); - let mut file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &ctx) + let mut file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &cancel, &ctx) .await .unwrap(); - let mutable = file.buffered_writer.inspect_mutable(); + let mutable = file.buffered_writer.mutable(); let cap = mutable.capacity(); let align = mutable.align(); @@ -446,7 +463,7 @@ mod tests { let maybe_flushed_buffer_contents = file.buffered_writer.inspect_maybe_flushed().unwrap(); assert_eq!(&maybe_flushed_buffer_contents[..], &content[cap..cap * 2]); - let mutable_buffer_contents = file.buffered_writer.inspect_mutable(); + let mutable_buffer_contents = file.buffered_writer.mutable(); assert_eq!(mutable_buffer_contents, &content[cap * 2..write_nbytes]); } @@ -455,13 +472,13 @@ mod tests { let (conf, tenant_id, timeline_id, ctx) = harness("test_flushes_do_happen").unwrap(); let gate = utils::sync::gate::Gate::default(); - - let mut file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &ctx) + let cancel = CancellationToken::new(); + let mut file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &cancel, &ctx) .await .unwrap(); // mutable buffer and maybe_flushed buffer each has `cap` bytes. - let cap = file.buffered_writer.inspect_mutable().capacity(); + let cap = file.buffered_writer.mutable().capacity(); let content: Vec = rand::thread_rng() .sample_iter(rand::distributions::Standard) @@ -471,10 +488,8 @@ mod tests { file.write_raw(&content, &ctx).await.unwrap(); // assert the state is as this test expects it to be - assert_eq!( - &file.load_to_io_buf(&ctx).await.unwrap(), - &content[0..cap * 2 + cap / 2] - ); + let load_io_buf_res = file.load_to_io_buf(&ctx).await.unwrap(); + assert_eq!(&load_io_buf_res[..], &content[0..cap * 2 + cap / 2]); let md = file.buffered_writer.as_inner().path().metadata().unwrap(); assert_eq!( md.len(), @@ -486,7 +501,7 @@ mod tests { &content[cap..cap * 2] ); assert_eq!( - &file.buffered_writer.inspect_mutable()[0..cap / 2], + &file.buffered_writer.mutable()[0..cap / 2], &content[cap * 2..cap * 2 + cap / 2] ); } @@ -502,12 +517,13 @@ mod tests { harness("test_read_split_across_file_and_buffer").unwrap(); let gate = utils::sync::gate::Gate::default(); + let cancel = CancellationToken::new(); - let mut file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &ctx) + let mut file = EphemeralFile::create(conf, tenant_id, timeline_id, &gate, &cancel, &ctx) .await .unwrap(); - let mutable = file.buffered_writer.inspect_mutable(); + let mutable = file.buffered_writer.mutable(); let cap = mutable.capacity(); let align = mutable.align(); let content: Vec = rand::thread_rng() diff --git a/pageserver/src/tenant/layer_map.rs b/pageserver/src/tenant/layer_map.rs index 2b04e53f10..96cee922ff 100644 --- a/pageserver/src/tenant/layer_map.rs +++ b/pageserver/src/tenant/layer_map.rs @@ -1112,6 +1112,7 @@ mod tests { }; use pageserver_api::key::DBDIR_KEY; use pageserver_api::keyspace::{KeySpace, KeySpaceRandomAccum}; + use tokio_util::sync::CancellationToken; use utils::id::{TenantId, TimelineId}; use utils::shard::TenantShardId; @@ -1193,6 +1194,7 @@ mod tests { async fn ranged_search() { let harness = TenantHarness::create("ranged_search").await.unwrap(); let (tenant, ctx) = harness.load().await; + let cancel = CancellationToken::new(); let timeline_id = TimelineId::generate(); // Create the timeline such that the in-memory layers can be written // to the timeline directory. @@ -1209,6 +1211,7 @@ mod tests { harness.tenant_shard_id, lsn_range.start, &gate, + &cancel, &ctx, ) .await diff --git a/pageserver/src/tenant/remote_timeline_client/download.rs b/pageserver/src/tenant/remote_timeline_client/download.rs index c59081c40e..32fce8cae3 100644 --- a/pageserver/src/tenant/remote_timeline_client/download.rs +++ b/pageserver/src/tenant/remote_timeline_client/download.rs @@ -207,6 +207,7 @@ async fn download_object( } #[cfg(target_os = "linux")] crate::virtual_file::io_engine::IoEngine::TokioEpollUring => { + use crate::virtual_file::owned_buffers_io::write::FlushTaskError; use std::sync::Arc; use crate::virtual_file::{IoBufferMut, owned_buffers_io}; @@ -231,6 +232,7 @@ async fn download_object( 0, || IoBufferMut::with_capacity(super::BUFFER_SIZE), gate.enter().map_err(|_| DownloadError::Cancelled)?, + cancel.child_token(), ctx, tracing::info_span!(parent: None, "download_object_buffered_writer", %dst_path), ); @@ -243,12 +245,14 @@ async fn download_object( { let chunk = match res { Ok(chunk) => chunk, - Err(e) => return Err(DownloadError::Other(anyhow::Error::new(e).context("download next chunk"))), + Err(e) => return Err(DownloadError::from(e)), }; - buffered.write_buffered_borrowed(&chunk, ctx) + buffered + .write_buffered_borrowed(&chunk, ctx) .await - .context("write_buffered_borrowed") - .map_err(DownloadError::Other)?; + .map_err(|e| match e { + FlushTaskError::Cancelled => DownloadError::Cancelled, + })?; } let mut pad_amount = None; let (bytes_amount, destination_file) = buffered @@ -269,8 +273,9 @@ async fn download_object( Some(buf) }) .await - .context("buffered writer shutdown") - .map_err(DownloadError::Other)?; + .map_err(|e| match e { + FlushTaskError::Cancelled => DownloadError::Cancelled, + })?; let pad_amount = pad_amount.expect("shutdown always invokes the closure").into_u64(); let set_len_arg = bytes_amount - pad_amount; diff --git a/pageserver/src/tenant/storage_layer/batch_split_writer.rs b/pageserver/src/tenant/storage_layer/batch_split_writer.rs index 84db79a83e..3c295f6aeb 100644 --- a/pageserver/src/tenant/storage_layer/batch_split_writer.rs +++ b/pageserver/src/tenant/storage_layer/batch_split_writer.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use bytes::Bytes; use pageserver_api::key::{KEY_SIZE, Key}; use pageserver_api::value::Value; +use tokio_util::sync::CancellationToken; use utils::id::TimelineId; use utils::lsn::Lsn; use utils::shard::TenantShardId; @@ -188,6 +189,7 @@ pub struct SplitImageLayerWriter { tenant_shard_id: TenantShardId, batches: BatchLayerWriter, start_key: Key, + cancel: CancellationToken, } impl SplitImageLayerWriter { @@ -200,6 +202,7 @@ impl SplitImageLayerWriter { lsn: Lsn, target_layer_size: u64, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ctx: &RequestContext, ) -> anyhow::Result { Ok(Self { @@ -211,6 +214,7 @@ impl SplitImageLayerWriter { &(start_key..Key::MAX), lsn, gate, + cancel.clone(), ctx, ) .await?, @@ -220,6 +224,7 @@ impl SplitImageLayerWriter { batches: BatchLayerWriter::new(conf).await?, lsn, start_key, + cancel, }) } @@ -244,6 +249,7 @@ impl SplitImageLayerWriter { &(key..Key::MAX), self.lsn, gate, + self.cancel.clone(), ctx, ) .await?; @@ -305,6 +311,7 @@ pub struct SplitDeltaLayerWriter { lsn_range: Range, last_key_written: Key, batches: BatchLayerWriter, + cancel: CancellationToken, } impl SplitDeltaLayerWriter { @@ -314,6 +321,7 @@ impl SplitDeltaLayerWriter { tenant_shard_id: TenantShardId, lsn_range: Range, target_layer_size: u64, + cancel: CancellationToken, ) -> anyhow::Result { Ok(Self { target_layer_size, @@ -324,6 +332,7 @@ impl SplitDeltaLayerWriter { lsn_range, last_key_written: Key::MIN, batches: BatchLayerWriter::new(conf).await?, + cancel, }) } @@ -351,6 +360,7 @@ impl SplitDeltaLayerWriter { key, self.lsn_range.clone(), gate, + self.cancel.clone(), ctx, ) .await?, @@ -370,6 +380,7 @@ impl SplitDeltaLayerWriter { key, self.lsn_range.clone(), gate, + self.cancel.clone(), ctx, ) .await?; @@ -478,6 +489,7 @@ mod tests { Lsn(0x18), 4 * 1024 * 1024, &tline.gate, + tline.cancel.clone(), &ctx, ) .await @@ -489,6 +501,7 @@ mod tests { tenant.tenant_shard_id, Lsn(0x18)..Lsn(0x20), 4 * 1024 * 1024, + tline.cancel.clone(), ) .await .unwrap(); @@ -562,6 +575,7 @@ mod tests { Lsn(0x18), 4 * 1024 * 1024, &tline.gate, + tline.cancel.clone(), &ctx, ) .await @@ -572,6 +586,7 @@ mod tests { tenant.tenant_shard_id, Lsn(0x18)..Lsn(0x20), 4 * 1024 * 1024, + tline.cancel.clone(), ) .await .unwrap(); @@ -666,6 +681,7 @@ mod tests { Lsn(0x18), 4 * 1024, &tline.gate, + tline.cancel.clone(), &ctx, ) .await @@ -677,6 +693,7 @@ mod tests { tenant.tenant_shard_id, Lsn(0x18)..Lsn(0x20), 4 * 1024, + tline.cancel.clone(), ) .await .unwrap(); @@ -765,6 +782,7 @@ mod tests { tenant.tenant_shard_id, Lsn(0x10)..Lsn(N as u64 * 16 + 0x10), 4 * 1024 * 1024, + tline.cancel.clone(), ) .await .unwrap(); diff --git a/pageserver/src/tenant/storage_layer/delta_layer.rs b/pageserver/src/tenant/storage_layer/delta_layer.rs index ddb140b997..f2a34ceec9 100644 --- a/pageserver/src/tenant/storage_layer/delta_layer.rs +++ b/pageserver/src/tenant/storage_layer/delta_layer.rs @@ -47,6 +47,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicU64; use tokio::sync::OnceCell; use tokio_epoll_uring::IoBuf; +use tokio_util::sync::CancellationToken; use tracing::*; use utils::bin_ser::BeSer; use utils::bin_ser::SerializeError; @@ -412,6 +413,7 @@ impl DeltaLayerWriterInner { /// /// Start building a new delta layer. /// + #[allow(clippy::too_many_arguments)] async fn new( conf: &'static PageServerConf, timeline_id: TimelineId, @@ -419,6 +421,7 @@ impl DeltaLayerWriterInner { key_start: Key, lsn_range: Range, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ctx: &RequestContext, ) -> anyhow::Result { // Create the file initially with a temporary filename. We don't know @@ -437,6 +440,7 @@ impl DeltaLayerWriterInner { file, PAGE_SZ as u64, gate, + 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), )?; @@ -664,6 +668,7 @@ impl DeltaLayerWriter { /// /// Start building a new delta layer. /// + #[allow(clippy::too_many_arguments)] pub async fn new( conf: &'static PageServerConf, timeline_id: TimelineId, @@ -671,6 +676,7 @@ impl DeltaLayerWriter { key_start: Key, lsn_range: Range, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ctx: &RequestContext, ) -> anyhow::Result { Ok(Self { @@ -682,6 +688,7 @@ impl DeltaLayerWriter { key_start, lsn_range, gate, + cancel, ctx, ) .await?, @@ -768,11 +775,7 @@ impl Drop for DeltaLayerWriter { .. } = inner; - let vfile = match blob_writer - .into_inner(|_| None) - .await - .maybe_fatal_err("failed to access inner virtual file") - { + let vfile = match blob_writer.into_inner(|_| None).await { Ok(vfile) => vfile, Err(e) => { error!(err=%e, "failed to remove delta layer writer file"); @@ -1946,6 +1949,7 @@ pub(crate) mod test { entries_meta.key_range.start, entries_meta.lsn_range.clone(), &timeline.gate, + timeline.cancel.clone(), &ctx, ) .await?; @@ -2141,6 +2145,7 @@ pub(crate) mod test { Key::MIN, Lsn(0x11)..truncate_at, &branch.gate, + branch.cancel.clone(), ctx, ) .await @@ -2276,6 +2281,7 @@ pub(crate) mod test { *key_start, (*lsn_min)..lsn_end, &tline.gate, + tline.cancel.clone(), ctx, ) .await?; diff --git a/pageserver/src/tenant/storage_layer/image_layer.rs b/pageserver/src/tenant/storage_layer/image_layer.rs index 6674838169..1626742771 100644 --- a/pageserver/src/tenant/storage_layer/image_layer.rs +++ b/pageserver/src/tenant/storage_layer/image_layer.rs @@ -46,6 +46,7 @@ use pageserver_api::value::Value; use serde::{Deserialize, Serialize}; use tokio::sync::OnceCell; use tokio_stream::StreamExt; +use tokio_util::sync::CancellationToken; use tracing::*; use utils::bin_ser::BeSer; use utils::bin_ser::SerializeError; @@ -761,6 +762,7 @@ impl ImageLayerWriterInner { /// /// Start building a new image layer. /// + #[allow(clippy::too_many_arguments)] async fn new( conf: &'static PageServerConf, timeline_id: TimelineId, @@ -768,6 +770,7 @@ impl ImageLayerWriterInner { key_range: &Range, lsn: Lsn, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ctx: &RequestContext, ) -> anyhow::Result { // Create the file initially with a temporary filename. @@ -800,6 +803,7 @@ impl ImageLayerWriterInner { file, PAGE_SZ as u64, gate, + 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), )?; @@ -1025,6 +1029,7 @@ impl ImageLayerWriter { /// /// Start building a new image layer. /// + #[allow(clippy::too_many_arguments)] pub async fn new( conf: &'static PageServerConf, timeline_id: TimelineId, @@ -1032,6 +1037,7 @@ impl ImageLayerWriter { key_range: &Range, lsn: Lsn, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ctx: &RequestContext, ) -> anyhow::Result { Ok(Self { @@ -1043,6 +1049,7 @@ impl ImageLayerWriter { key_range, lsn, gate, + cancel, ctx, ) .await?, @@ -1107,11 +1114,7 @@ impl Drop for ImageLayerWriter { .. } = inner; - let vfile = match blob_writer - .into_inner(|_| None) - .await - .maybe_fatal_err("failed to access inner virtual file") - { + let vfile = match blob_writer.into_inner(|_| None).await { Ok(vfile) => vfile, Err(e) => { error!(err=%e, "failed to remove image layer writer file"); @@ -1279,6 +1282,7 @@ mod test { &range, lsn, &timeline.gate, + timeline.cancel.clone(), &ctx, ) .await @@ -1345,6 +1349,7 @@ mod test { &range, lsn, &timeline.gate, + timeline.cancel.clone(), &ctx, ) .await @@ -1424,6 +1429,7 @@ mod test { &key_range, lsn, &tline.gate, + tline.cancel.clone(), ctx, ) .await?; diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs index f3a45a4976..5a07389dc6 100644 --- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs +++ b/pageserver/src/tenant/storage_layer/inmemory_layer.rs @@ -19,6 +19,7 @@ use pageserver_api::keyspace::KeySpace; use pageserver_api::models::InMemoryLayerInfo; use pageserver_api::shard::TenantShardId; use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; use tracing::*; use utils::id::TimelineId; use utils::lsn::Lsn; @@ -552,13 +553,15 @@ impl InMemoryLayer { tenant_shard_id: TenantShardId, start_lsn: Lsn, gate: &utils::sync::gate::Gate, + cancel: &CancellationToken, ctx: &RequestContext, ) -> Result { trace!( "initializing new empty InMemoryLayer for writing on timeline {timeline_id} at {start_lsn}" ); - let file = EphemeralFile::create(conf, tenant_shard_id, timeline_id, gate, ctx).await?; + let file = + EphemeralFile::create(conf, tenant_shard_id, timeline_id, gate, cancel, ctx).await?; let key = InMemoryLayerFileId(file.page_cache_file_id()); Ok(InMemoryLayer { @@ -717,6 +720,7 @@ impl InMemoryLayer { key_range: Option>, l0_flush_global_state: &l0_flush::Inner, gate: &utils::sync::gate::Gate, + cancel: CancellationToken, ) -> Result> { // Grab the lock in read-mode. We hold it over the I/O, but because this // layer is not writeable anymore, no one should be trying to acquire the @@ -758,6 +762,7 @@ impl InMemoryLayer { Key::MIN, self.start_lsn..end_lsn, gate, + cancel, ctx, ) .await?; diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 8ea4549304..51b78dcaba 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -4186,6 +4186,7 @@ impl Timeline { self.timeline_id, self.tenant_shard_id, &self.gate, + &self.cancel, ctx, ) .await?; @@ -4711,6 +4712,7 @@ impl Timeline { key_range, self_clone.l0_flush_global_state.inner(), &self_clone.gate, + self_clone.cancel.clone(), ) .await? else { @@ -5250,6 +5252,7 @@ impl Timeline { &img_range, lsn, &self.gate, + self.cancel.clone(), ctx, ) .await?; @@ -6599,6 +6602,7 @@ impl Timeline { &(min_key..end_key), lsn, &self.gate, + self.cancel.clone(), ctx, ) .await?; @@ -6661,6 +6665,7 @@ impl Timeline { deltas.key_range.start, deltas.lsn_range, &self.gate, + self.cancel.clone(), ctx, ) .await?; @@ -6750,6 +6755,8 @@ impl Timeline { self.tenant_shard_id, in_memory.lsn_range.start, &self.gate, + // TODO: if we ever use this function in production code, we need to pass the real cancellation token + &CancellationToken::new(), ctx, ) .await diff --git a/pageserver/src/tenant/timeline/compaction.rs b/pageserver/src/tenant/timeline/compaction.rs index 5189ae418c..4ef2a1c38d 100644 --- a/pageserver/src/tenant/timeline/compaction.rs +++ b/pageserver/src/tenant/timeline/compaction.rs @@ -1218,6 +1218,7 @@ impl Timeline { &layer.layer_desc().key_range, layer.layer_desc().image_layer_lsn(), &self.gate, + self.cancel.clone(), ctx, ) .await @@ -1819,6 +1820,7 @@ impl Timeline { lsn_range.clone() }, &self.gate, + self.cancel.clone(), ctx, ) .await @@ -3014,6 +3016,7 @@ impl Timeline { lowest_retain_lsn, self.get_compaction_target_size(), &self.gate, + self.cancel.clone(), ctx, ) .await @@ -3030,6 +3033,7 @@ impl Timeline { self.tenant_shard_id, lowest_retain_lsn..end_lsn, self.get_compaction_target_size(), + self.cancel.clone(), ) .await .context("failed to create delta layer writer") @@ -3129,6 +3133,7 @@ impl Timeline { desc.key_range.start, desc.lsn_range.clone(), &self.gate, + self.cancel.clone(), ctx, ) .await @@ -3147,6 +3152,7 @@ impl Timeline { job_desc.compaction_key_range.end, desc.lsn_range.clone(), &self.gate, + self.cancel.clone(), ctx, ) .await @@ -3716,6 +3722,7 @@ impl CompactionJobExecutor for TimelineAdaptor { key_range.start, lsn_range.clone(), &self.timeline.gate, + self.timeline.cancel.clone(), ctx, ) .await?; @@ -3792,6 +3799,7 @@ impl TimelineAdaptor { key_range, lsn, &self.timeline.gate, + self.timeline.cancel.clone(), ctx, ) .await?; diff --git a/pageserver/src/tenant/timeline/detach_ancestor.rs b/pageserver/src/tenant/timeline/detach_ancestor.rs index 0014ab9113..66e6d1e848 100644 --- a/pageserver/src/tenant/timeline/detach_ancestor.rs +++ b/pageserver/src/tenant/timeline/detach_ancestor.rs @@ -675,6 +675,7 @@ async fn copy_lsn_prefix( layer.layer_desc().key_range.start, layer.layer_desc().lsn_range.start..end_lsn, &target_timeline.gate, + target_timeline.cancel.clone(), ctx, ) .await diff --git a/pageserver/src/tenant/timeline/import_pgdata/flow.rs b/pageserver/src/tenant/timeline/import_pgdata/flow.rs index c8f151b56b..c6d2944769 100644 --- a/pageserver/src/tenant/timeline/import_pgdata/flow.rs +++ b/pageserver/src/tenant/timeline/import_pgdata/flow.rs @@ -739,6 +739,7 @@ impl ChunkProcessingJob { &self.range, self.pgdata_lsn, &self.timeline.gate, + self.timeline.cancel.clone(), ctx, ) .await?; diff --git a/pageserver/src/tenant/timeline/layer_manager.rs b/pageserver/src/tenant/timeline/layer_manager.rs index 1b489028dc..ed92ea28ce 100644 --- a/pageserver/src/tenant/timeline/layer_manager.rs +++ b/pageserver/src/tenant/timeline/layer_manager.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use anyhow::{Context, bail, ensure}; use itertools::Itertools; use pageserver_api::shard::TenantShardId; +use tokio_util::sync::CancellationToken; use tracing::trace; use utils::id::TimelineId; use utils::lsn::{AtomicLsn, Lsn}; @@ -193,6 +194,7 @@ impl OpenLayerManager { /// Open a new writable layer to append data if there is no open layer, otherwise return the /// current open layer, called within `get_layer_for_write`. + #[allow(clippy::too_many_arguments)] pub(crate) async fn get_layer_for_write( &mut self, lsn: Lsn, @@ -200,6 +202,7 @@ impl OpenLayerManager { timeline_id: TimelineId, tenant_shard_id: TenantShardId, gate: &utils::sync::gate::Gate, + cancel: &CancellationToken, ctx: &RequestContext, ) -> anyhow::Result> { ensure!(lsn.is_aligned()); @@ -227,9 +230,16 @@ impl OpenLayerManager { timeline_id, start_lsn, lsn ); - let new_layer = - InMemoryLayer::create(conf, timeline_id, tenant_shard_id, start_lsn, gate, ctx) - .await?; + let new_layer = InMemoryLayer::create( + conf, + timeline_id, + tenant_shard_id, + start_lsn, + gate, + cancel, + ctx, + ) + .await?; let layer = Arc::new(new_layer); self.layer_map.open_layer = Some(layer.clone()); diff --git a/pageserver/src/tenant/vectored_blob_io.rs b/pageserver/src/tenant/vectored_blob_io.rs index ef42940481..8b30e70e41 100644 --- a/pageserver/src/tenant/vectored_blob_io.rs +++ b/pageserver/src/tenant/vectored_blob_io.rs @@ -677,7 +677,6 @@ impl StreamingVectoredReadPlanner { #[cfg(test)] mod tests { - use anyhow::Error; use super::super::blob_io::tests::{random_array, write_maybe_compressed}; use super::*; @@ -960,7 +959,10 @@ mod tests { } } - async fn round_trip_test_compressed(blobs: &[Vec], compression: bool) -> Result<(), Error> { + async fn round_trip_test_compressed( + blobs: &[Vec], + compression: bool, + ) -> anyhow::Result<()> { let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error).with_scope_unit_test(); let (_temp_dir, pathbuf, offsets) = @@ -1003,7 +1005,7 @@ mod tests { } #[tokio::test] - async fn test_really_big_array() -> Result<(), Error> { + async fn test_really_big_array() -> anyhow::Result<()> { let blobs = &[ b"test".to_vec(), random_array(10 * PAGE_SZ), @@ -1018,7 +1020,7 @@ mod tests { } #[tokio::test] - async fn test_arrays_inc() -> Result<(), Error> { + async fn test_arrays_inc() -> anyhow::Result<()> { let blobs = (0..PAGE_SZ / 8) .map(|v| random_array(v * 16)) .collect::>(); diff --git a/pageserver/src/virtual_file/owned_buffers_io/write.rs b/pageserver/src/virtual_file/owned_buffers_io/write.rs index af408c97c3..077bd1bc06 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/write.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/write.rs @@ -3,10 +3,11 @@ use std::sync::Arc; use super::io_buf_aligned::{IoBufAligned, IoBufAlignedMut}; use bytes::BufMut; -use flush::FlushHandle; -use tokio_epoll_uring::IoBuf; - pub(crate) use flush::FlushControl; +use flush::FlushHandle; +pub(crate) use flush::FlushTaskError; +use tokio_epoll_uring::IoBuf; +use tokio_util::sync::CancellationToken; use super::io_buf_ext::{FullSlice, IoBufExt}; use crate::context::RequestContext; @@ -42,11 +43,19 @@ pub trait OwnedAsyncWriter { // since we would avoid copying majority of the data into the internal buffer. pub struct BufferedWriter { writer: Arc, - /// invariant: always remains Some(buf) except - /// - while IO is ongoing => goes back to Some() once the IO completed successfully - /// - after an IO error => stays `None` forever - /// - /// In these exceptional cases, it's `None`. + /// Clone of the buffer that was last submitted to the flush loop. + /// `None` if no flush request has been submitted, Some forever after. + pub(super) maybe_flushed: Option>, + /// New writes are accumulated here. + /// `None` only during submission while we wait for flush loop to accept + /// the full dirty buffer in exchange for a clean buffer. + /// If that exchange fails with an [`FlushTaskError`], the write path + /// bails and leaves this as `None`. + /// Subsequent writes will panic if attempted. + /// The read path continues to work without error because [`Self::maybe_flushed`] + /// and [`Self::bytes_submitted`] are advanced before the flush loop exchange starts, + /// so, they will never try to read from [`Self::mutable`] anyway, because it's past + /// the [`Self::maybe_flushed`] point. mutable: Option, /// A handle to the background flush task for writting data to disk. flush_handle: FlushHandle, @@ -68,16 +77,19 @@ where start_offset: u64, buf_new: impl Fn() -> B, gate_guard: utils::sync::gate::GateGuard, + cancel: CancellationToken, ctx: &RequestContext, flush_task_span: tracing::Span, ) -> Self { Self { writer: writer.clone(), mutable: Some(buf_new()), + maybe_flushed: None, flush_handle: FlushHandle::spawn_new( writer, buf_new(), gate_guard, + cancel, ctx.attached_child(), flush_task_span, ), @@ -95,23 +107,24 @@ where } /// Panics if used after any of the write paths returned an error - pub fn inspect_mutable(&self) -> &B { - self.mutable() + pub fn inspect_mutable(&self) -> Option<&B> { + self.mutable.as_ref() } /// Gets a reference to the maybe flushed read-only buffer. /// Returns `None` if the writer has not submitted any flush request. pub fn inspect_maybe_flushed(&self) -> Option<&FullSlice> { - self.flush_handle.maybe_flushed.as_ref() + self.maybe_flushed.as_ref() } #[cfg_attr(target_os = "macos", allow(dead_code))] pub async fn shutdown( self, mut handle_tail: impl FnMut(B) -> Option, - ) -> std::io::Result<(u64, W)> { + ) -> Result<(u64, W), FlushTaskError> { let Self { mutable: tail, + maybe_flushed: _, writer, mut flush_handle, bytes_submitted: submit_offset, @@ -123,18 +136,17 @@ where let mut bytes_amount = submit_offset; if let Some(buf) = handle_tail(buf) { bytes_amount += buf.pending() as u64; + // TODO: infinite retries + maybe_fatal_err like we do in the flush loop; can we just send this + // as work into the flush loop, as part of flush_handle.shutdown() let (_, res) = writer.write_all_at(buf.flush(), submit_offset, &ctx).await; - let _: () = res?; + let _: () = res.unwrap(); // DO NOT MERGE, THIS CAN FAIL E.G. on ENOSPC } Ok((bytes_amount, writer)) } - /// Gets a reference to the mutable in-memory buffer. - #[inline(always)] - fn mutable(&self) -> &B { - self.mutable - .as_ref() - .expect("must not use after we returned an error") + #[cfg(test)] + pub(crate) fn mutable(&self) -> &B { + self.mutable.as_ref().expect("must not use after an error") } #[cfg_attr(target_os = "macos", allow(dead_code))] @@ -142,7 +154,7 @@ where &mut self, chunk: &[u8], ctx: &RequestContext, - ) -> std::io::Result { + ) -> Result { let (len, control) = self.write_buffered_borrowed_controlled(chunk, ctx).await?; if let Some(control) = control { control.release().await; @@ -155,7 +167,7 @@ where &mut self, mut chunk: &[u8], ctx: &RequestContext, - ) -> std::io::Result<(usize, Option)> { + ) -> Result<(usize, Option), FlushTaskError> { let chunk_len = chunk.len(); let mut control: Option = None; while !chunk.is_empty() { @@ -176,17 +188,47 @@ where Ok((chunk_len, control)) } + /// This function can only error if the flush task got cancelled. + /// In that case, we leave [`Self::mutable`] intentionally as `None`. + /// + /// The read path continues to function correctly; it can read up to the + /// point where it could read before, i.e., including what was in [`Self::mutable`] + /// before the call to this function, because that's now stored in [`Self::maybe_flushed`]. + /// + /// The write path becomes unavailable and will panic if used. + /// The only correct solution to retry writes is to discard the entire [`BufferedWriter`], + /// which upper layers of pageserver write path currently do not support. + /// It is in fact quite hard to reason about what exactly happens in today's code. + /// Best case we accumulate junk in the EphemeralFile, worst case is data corruption. #[must_use = "caller must explcitly check the flush control"] - async fn flush(&mut self, _ctx: &RequestContext) -> std::io::Result> { + async fn flush( + &mut self, + _ctx: &RequestContext, + ) -> Result, FlushTaskError> { let buf = self.mutable.take().expect("must not use after an error"); let buf_len = buf.pending(); if buf_len == 0 { self.mutable = Some(buf); return Ok(None); } - let (recycled, flush_control) = self.flush_handle.flush(buf, self.bytes_submitted).await?; + // Prepare the buffer for read while flushing. + let slice = buf.flush(); + // NB: this assignment also drops thereference to the old buffer, allowing us to re-own & make it mutable below. + self.maybe_flushed = Some(slice.cheap_clone()); + let offset = self.bytes_submitted; self.bytes_submitted += u64::try_from(buf_len).unwrap(); + + // If we return/panic here or later, we'll leave mutable = None, breaking further + // writers, but the read path should still work. + let (recycled, flush_control) = self.flush_handle.flush(slice, offset).await?; + + // The only other place that could hold a reference to the recycled buffer + // is in `Self::maybe_flushed`, but we have already replace it with the new buffer. + let recycled = Buffer::reuse_after_flush(recycled.into_raw_slice().into_inner()); + + // We got back some recycled buffer, can open up for more writes again. self.mutable = Some(recycled); + Ok(Some(flush_control)) } } @@ -311,11 +353,13 @@ mod tests { let ctx = &ctx; let recorder = Arc::new(RecorderWriter::default()); let gate = utils::sync::gate::Gate::default(); + let cancel = CancellationToken::new(); let mut writer = BufferedWriter::<_, RecorderWriter>::new( recorder, 0, || IoBufferMut::with_capacity(2), gate.enter()?, + cancel, ctx, tracing::Span::none(), ); diff --git a/pageserver/src/virtual_file/owned_buffers_io/write/flush.rs b/pageserver/src/virtual_file/owned_buffers_io/write/flush.rs index cdc8537b13..c417776139 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/write/flush.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/write/flush.rs @@ -1,7 +1,6 @@ use std::ops::ControlFlow; use std::{marker::PhantomData, sync::Arc}; -use once_cell::sync::Lazy; use tokio_util::sync::CancellationToken; use tracing::{Instrument, info, info_span, warn}; use utils::sync::duplex; @@ -15,9 +14,6 @@ use crate::virtual_file::owned_buffers_io::io_buf_ext::FullSlice; /// A handle to the flush task. pub struct FlushHandle { inner: Option>, - /// Immutable buffer for serving tail reads. - /// `None` if no flush request has been submitted. - pub(super) maybe_flushed: Option>, } pub struct FlushHandleInner { @@ -25,8 +21,9 @@ pub struct FlushHandleInner { /// and receives recyled buffer. channel: duplex::mpsc::Duplex, FullSlice>, /// Join handle for the background flush task. - join_handle: tokio::task::JoinHandle>, + join_handle: tokio::task::JoinHandle>, + // TODO: get rit of the type parameter? _phantom: PhantomData, } @@ -119,27 +116,31 @@ where { /// Spawns a new background flush task and obtains a handle. /// - /// Note: The background task so we do not need to explicitly maintain a queue of buffers. + /// Handle and background task are connected through a duplex channel. + /// Dirty buffers are sent to the background task for flushing. + /// Clean buffers are sent back to the handle for reuse. + /// + /// The queue depth is 1, and the passed-in `buf` seeds the queue depth. + /// I.e., the passed-in buf is immediately available to the handle as a recycled buffer. pub fn spawn_new( file: Arc, buf: B, gate_guard: utils::sync::gate::GateGuard, + cancel: CancellationToken, ctx: RequestContext, span: tracing::Span, ) -> Self where B: Buffer + Send + 'static, { - // It is fine to buffer up to only 1 message. We only 1 message in-flight at a time. let (front, back) = duplex::mpsc::channel(1); + back.try_send(buf.flush()) + .expect("we just created it with capacity 1"); let join_handle = tokio::spawn( - async move { - FlushBackgroundTask::new(back, file, gate_guard, ctx) - .run(buf.flush()) - .await - } - .instrument(span), + FlushBackgroundTask::new(back, file, gate_guard, cancel, ctx) + .run() + .instrument(span), ); FlushHandle { @@ -148,7 +149,6 @@ where join_handle, _phantom: PhantomData, }), - maybe_flushed: None, } } @@ -156,15 +156,11 @@ where /// Returns a buffer that completed flushing for re-use, length reset to 0, capacity unchanged. /// If `save_buf_for_read` is true, then we save the buffer in `Self::maybe_flushed`, otherwise /// clear `maybe_flushed`. - pub async fn flush(&mut self, buf: B, offset: u64) -> std::io::Result<(B, FlushControl)> - where - B: Buffer + Send + 'static, - { - let slice = buf.flush(); - - // Saves a buffer for read while flushing. This also removes reference to the old buffer. - self.maybe_flushed = Some(slice.cheap_clone()); - + pub async fn flush( + &mut self, + slice: FullSlice, + offset: u64, + ) -> Result<(FullSlice, FlushControl), FlushTaskError> { let (request, flush_control) = new_flush_op(slice, offset); // Submits the buffer to the background task. @@ -180,13 +176,10 @@ where return self.handle_error().await; }; - // The only other place that could hold a reference to the recycled buffer - // is in `Self::maybe_flushed`, but we have already replace it with the new buffer. - let recycled = Buffer::reuse_after_flush(recycled.into_raw_slice().into_inner()); Ok((recycled, flush_control)) } - async fn handle_error(&mut self) -> std::io::Result { + async fn handle_error(&mut self) -> Result { Err(self .shutdown() .await @@ -195,7 +188,7 @@ where } /// Cleans up the channel, join the flush task. - pub async fn shutdown(&mut self) -> std::io::Result { + pub async fn shutdown(&mut self) -> Result { let handle: FlushHandleInner = self .inner .take() @@ -221,10 +214,17 @@ pub struct FlushBackgroundTask { /// A writter for persisting data to disk. writer: Arc, ctx: RequestContext, + cancel: CancellationToken, /// Prevent timeline from shuting down until the flush background task finishes flushing all remaining buffers to disk. _gate_guard: utils::sync::gate::GateGuard, } +#[derive(Debug, thiserror::Error)] +pub enum FlushTaskError { + #[error("flush task cancelled")] + Cancelled, +} + impl FlushBackgroundTask where Buf: IoBufAligned + Send + Sync, @@ -235,25 +235,20 @@ where channel: duplex::mpsc::Duplex, FlushRequest>, file: Arc, gate_guard: utils::sync::gate::GateGuard, + cancel: CancellationToken, ctx: RequestContext, ) -> Self { FlushBackgroundTask { channel, writer: file, _gate_guard: gate_guard, + cancel, ctx, } } /// Runs the background flush task. - /// The passed in slice is immediately sent back to the flush handle through the duplex channel. - async fn run(mut self, slice: FullSlice) -> std::io::Result { - // Sends the extra buffer back to the handle. - // TODO: can this ever await and or fail? I think not. - self.channel.send(slice).await.map_err(|_| { - std::io::Error::new(std::io::ErrorKind::BrokenPipe, "flush handle closed early") - })?; - + async fn run(mut self) -> Result { // Exit condition: channel is closed and there is no remaining buffer to be flushed while let Some(request) = self.channel.recv().await { #[cfg(test)] @@ -271,15 +266,13 @@ where // (The upper layers of the Pageserver write path are not equipped to retry write errors // becasuse they often deallocate the buffers that were already written). // - // TODO: cancellation sensitiity. - // Without it, if we hit a bug where retrying is never successful, - // then we can't shut down the timeline/tenant/pageserver cleanly because - // layers of the Pageserver write path are holding the gate open for EphemeralFile. - // // TODO: use utils::backoff::retry once async closures are actually usable // let mut slice_storage = Some(request.slice); for attempt in 1.. { + if self.cancel.is_cancelled() { + return Err(FlushTaskError::Cancelled); + } let result = async { if attempt > 1 { info!("retrying flush"); @@ -287,6 +280,11 @@ where let slice = slice_storage.take().expect( "likely previous invocation of this future didn't get polled to completion", ); + // Don't cancel this write by doing tokio::select with self.cancel.cancelled(). + // The underlying tokio-epoll-uring slot / kernel operation is still ongoing and occupies resources. + // If we retry indefinitely, we'll deplete those resources. + // Future: teach tokio-epoll-uring io_uring operation cancellation, but still, + // wait for cancelled ops to complete and discard their error. let (slice, res) = self.writer.write_all_at(slice, request.offset, &self.ctx).await; slice_storage = Some(slice); let res = res.maybe_fatal_err("owned_buffers_io flush"); @@ -294,8 +292,7 @@ where return ControlFlow::Break(()); }; warn!(%err, "error flushing buffered writer buffer to disk, retrying after backoff"); - static NO_CANCELLATION: Lazy = Lazy::new(CancellationToken::new); - utils::backoff::exponential_backoff(attempt, 1.0, 10.0, &NO_CANCELLATION).await; + utils::backoff::exponential_backoff(attempt, 1.0, 10.0, &self.cancel).await; ControlFlow::Continue(()) } .instrument(info_span!("flush_attempt", %attempt))