From eb1ccd798877ea8e590ff4bb1fd59976c08bce93 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Thu, 14 Mar 2024 16:06:04 +0000 Subject: [PATCH] WIP: page_cache_priming_writer: plumb through RequestContext for previous commit, yet more churn -,- --- pageserver/src/tenant/ephemeral_file.rs | 25 +++-- pageserver/src/virtual_file.rs | 4 +- .../util/page_cache_priming_writer.rs | 48 ++++++--- .../util/size_tracking_writer.rs | 15 +-- .../virtual_file/owned_buffers_io/write.rs | 97 +++++++++++++------ 5 files changed, 127 insertions(+), 62 deletions(-) diff --git a/pageserver/src/tenant/ephemeral_file.rs b/pageserver/src/tenant/ephemeral_file.rs index 8851df939d..179239a621 100644 --- a/pageserver/src/tenant/ephemeral_file.rs +++ b/pageserver/src/tenant/ephemeral_file.rs @@ -3,9 +3,8 @@ use crate::config::PageServerConf; use crate::context::RequestContext; -use crate::page_cache::{self, PageCache, PAGE_SZ}; +use crate::page_cache::{self, PAGE_SZ}; use crate::tenant::block_io::{BlockCursor, BlockLease, BlockReader}; -use crate::virtual_file::owned_buffers_io::util::page_cache_priming_writer::VirtualFileAdaptor; use crate::virtual_file::owned_buffers_io::write::OwnedAsyncWriter; use crate::virtual_file::{self, owned_buffers_io, VirtualFile}; use camino::Utf8PathBuf; @@ -75,7 +74,8 @@ impl EphemeralFile { file, page_cache_file_id, ); - let file = owned_buffers_io::util::page_cache_priming_writer::Writer::new(file, page_cache::get()); + let file = + owned_buffers_io::util::page_cache_priming_writer::Writer::new(file, page_cache::get()); let file = owned_buffers_io::util::size_tracking_writer::Writer::new(file); let file = owned_buffers_io::write::BufferedWriter::new(file); let file = owned_buffers_io::util::size_tracking_writer::Writer::new(file); @@ -185,7 +185,7 @@ impl EphemeralFile { pub(crate) async fn write_blob( &mut self, srcbuf: &[u8], - _ctx: &RequestContext, + ctx: &RequestContext, ) -> Result { let pos = self.file.bytes_written(); @@ -194,15 +194,15 @@ impl EphemeralFile { // short one-byte length header let len_buf = [srcbuf.len() as u8]; - self.file.write_all_borrowed(&len_buf).await?; + self.file.write_all_borrowed(&len_buf, ctx).await?; } else { let mut len_buf = u32::to_be_bytes(srcbuf.len() as u32); len_buf[0] |= 0x80; - self.file.write_all_borrowed(&len_buf).await?; + self.file.write_all_borrowed(&len_buf, ctx).await?; } // Write the payload - self.file.write_all_borrowed(&srcbuf).await?; + self.file.write_all_borrowed(&srcbuf, ctx).await?; Ok(pos) } @@ -223,7 +223,16 @@ impl Drop for EphemeralFile { // We leave them there, [`crate::page_cache::PageCache::find_victim`] will evict them when needed. // unlink the file - let res = std::fs::remove_file(&self.file.as_inner().as_inner().as_inner().as_inner().as_inner().path); + let res = std::fs::remove_file( + &self + .file + .as_inner() + .as_inner() + .as_inner() + .as_inner() + .as_inner() + .path, + ); if let Err(e) = res { if e.kind() != std::io::ErrorKind::NotFound { // just never log the not found errors, we cannot do anything for them; on detach diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs index 9054724596..d0a91b48ac 100644 --- a/pageserver/src/virtual_file.rs +++ b/pageserver/src/virtual_file.rs @@ -10,6 +10,7 @@ //! This is similar to PostgreSQL's virtual file descriptor facility in //! src/backend/storage/file/fd.c //! +use crate::context::RequestContext; use crate::metrics::{StorageIoOperation, STORAGE_IO_SIZE, STORAGE_IO_TIME_METRIC}; use crate::page_cache::PageWriteGuard; @@ -1092,13 +1093,14 @@ impl OwnedAsyncWriter for VirtualFile { >( &mut self, buf: B, + _: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { let (buf, res) = VirtualFile::write_all(self, buf).await; res.map(move |v| (v, buf)) } #[inline(always)] - async fn write_all_borrowed(&mut self, _buf: &[u8]) -> std::io::Result { + async fn write_all_borrowed(&mut self, _buf: &[u8], _ctx: &RequestContext) -> std::io::Result { // TODO: ensure this through the type system panic!("this should not happen"); } diff --git a/pageserver/src/virtual_file/owned_buffers_io/util/page_cache_priming_writer.rs b/pageserver/src/virtual_file/owned_buffers_io/util/page_cache_priming_writer.rs index 8d5f557362..350838face 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/util/page_cache_priming_writer.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/util/page_cache_priming_writer.rs @@ -1,10 +1,11 @@ -use std::ops::{Deref, Range}; +use std::ops::Deref; use tokio_epoll_uring::BoundedBuf; use tracing::error; use crate::{ - page_cache::{self, PAGE_SZ}, + context::RequestContext, + page_cache::{self, PageReadGuard, PAGE_SZ}, virtual_file::{owned_buffers_io::write::OwnedAsyncWriter, VirtualFile}, }; @@ -15,7 +16,13 @@ pub struct Writer> { } pub trait Cache { - async fn fill_cache(&self, under: &W, page_no: u64, contents: &[u8; PAGE_SZ]); + async fn fill_cache( + &self, + under: &W, + page_no: u64, + contents: &[u8; PAGE_SZ], + ctx: &RequestContext, + ); } impl Writer @@ -54,18 +61,18 @@ where >( &mut self, buf: B, + ctx: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { let buf = buf.slice_full(); assert_eq!(buf.bytes_init() % N, 0); self.invariants(); - let pre = self.written; let saved_bounds = buf.bounds(); let debug_assert_contents_eq = if cfg!(debug_assertions) { Some(buf[..].to_vec()) } else { None }; - let res = self.under.write_all(buf).await; + let res = self.under.write_all(buf, ctx).await; let res = if let Ok((nwritten, buf)) = res { assert_eq!(nwritten % N, 0); let buf = tokio_epoll_uring::Slice::from_buf_bounds(buf, saved_bounds); @@ -73,13 +80,14 @@ where debug_assert_eq!(&before[..], &buf[..]); } assert_eq!(nwritten, buf.bytes_init()); - let new_written = self.written + (nwritten as u64); for page_no_in_buf in 0..(nwritten / N) { - let page: &[u8; N] = todo!(); + let page: &[u8; N] = (&buf[(page_no_in_buf * N)..((page_no_in_buf + 1) * N)]) + .try_into() + .unwrap(); self.cache - .fill_cache(&self.under, self.written / (N as u64), page) + .fill_cache(&self.under, self.written / (N as u64), page, ctx) .await; - self.written += (N as u64); + self.written += N as u64; } self.written += nwritten as u64; Ok((nwritten, tokio_epoll_uring::Slice::into_inner(buf))) @@ -90,7 +98,7 @@ where res } - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { + async fn write_all_borrowed(&mut self, _: &[u8], _: &RequestContext) -> std::io::Result { // TODO: use type system to ensure this doesn't happen at runtime panic!("don't put the types together this way") } @@ -111,21 +119,28 @@ impl VirtualFileAdaptor { } impl<'c> Cache<{ PAGE_SZ }, VirtualFileAdaptor> for &'c crate::page_cache::PageCache { - async fn fill_cache(&self, under: &VirtualFileAdaptor, page_no: u64, contents: &[u8; PAGE_SZ]) { + async fn fill_cache( + &self, + under: &VirtualFileAdaptor, + page_no: u64, + contents: &[u8; PAGE_SZ], + ctx: &RequestContext, + ) { match self .read_immutable_buf( under.file_id, u32::try_from(page_no).expect("files larger than u32::MAX * 8192 aren't supported"), - todo!("funnel through context"), + ctx, ) .await { Ok(crate::page_cache::ReadBufResult::Found(guard)) => { debug_assert_eq!(guard.deref(), contents); } - Ok(crate::page_cache::ReadBufResult::NotFound(guard)) => { + Ok(crate::page_cache::ReadBufResult::NotFound(mut guard)) => { guard.copy_from_slice(contents); - guard.mark_valid(); + let guard: PageReadGuard<'_> = guard.mark_valid(); + drop(guard); } Err(e) => { error!("failed to get immutable buf to pre-warm page cache: {e:?}"); @@ -142,11 +157,12 @@ impl OwnedAsyncWriter for VirtualFileAdaptor { >( &mut self, buf: B, + ctx: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { - OwnedAsyncWriter::write_all(&mut self.file, buf).await + OwnedAsyncWriter::write_all(&mut self.file, buf, ctx).await } - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { + async fn write_all_borrowed(&mut self, _: &[u8], _: &RequestContext) -> std::io::Result { // TODO: use type system to ensure this doesn't happen at runtime panic!("don't put the types together this way") } diff --git a/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs b/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs index 55c22e72da..dcde4e3ced 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs @@ -1,6 +1,4 @@ -use std::ops::Range; - -use crate::virtual_file::owned_buffers_io::write::OwnedAsyncWriter; +use crate::{context::RequestContext, virtual_file::owned_buffers_io::write::OwnedAsyncWriter}; use tokio_epoll_uring::{BoundedBuf, IoBuf}; pub struct Writer { @@ -40,15 +38,20 @@ where >( &mut self, buf: B, + ctx: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { - let (nwritten, buf) = self.dst.write_all(buf).await?; + let (nwritten, buf) = self.dst.write_all(buf, ctx).await?; self.bytes_amount += u64::try_from(nwritten).unwrap(); Ok((nwritten, buf)) } #[inline(always)] - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { - let nwritten = self.dst.write_all_borrowed(buf).await?; + async fn write_all_borrowed( + &mut self, + buf: &[u8], + ctx: &RequestContext, + ) -> std::io::Result { + let nwritten = self.dst.write_all_borrowed(buf, ctx).await?; self.bytes_amount += u64::try_from(nwritten).unwrap(); Ok(nwritten) } diff --git a/pageserver/src/virtual_file/owned_buffers_io/write.rs b/pageserver/src/virtual_file/owned_buffers_io/write.rs index 0576f149c9..cfa5921e18 100644 --- a/pageserver/src/virtual_file/owned_buffers_io/write.rs +++ b/pageserver/src/virtual_file/owned_buffers_io/write.rs @@ -1,7 +1,7 @@ -use std::ops::Range; - use tokio_epoll_uring::{BoundedBuf, IoBuf, Slice}; +use crate::context::RequestContext; + /// A trait for doing owned-buffer write IO. /// Think [`tokio::io::AsyncWrite`] but with owned buffers. pub trait OwnedAsyncWriter { @@ -12,8 +12,14 @@ pub trait OwnedAsyncWriter { >( &mut self, buf: B, + ctx: &RequestContext, ) -> std::io::Result<(usize, B::Buf)>; - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result; + async fn write_all_borrowed( + &mut self, + buf: &[u8], + + ctx: &RequestContext, + ) -> std::io::Result; } /// A wrapper aorund an [`OwnedAsyncWriter`] that batches smaller writers @@ -67,21 +73,25 @@ where .as_zero_padded_slice() } - pub async fn flush_and_into_inner(mut self) -> std::io::Result { - self.flush().await?; + pub async fn flush_and_into_inner(mut self, ctx: &RequestContext) -> std::io::Result { + self.flush(ctx).await?; let Self { buf, writer } = self; assert!(buf.is_some()); Ok(writer) } - pub async fn write_buffered(&mut self, chunk: Slice) -> std::io::Result<(usize, B)> + pub async fn write_buffered( + &mut self, + chunk: Slice, + ctx: &RequestContext, + ) -> std::io::Result<(usize, B)> where B: IoBuf + Send, { let chunk_len = chunk.len(); // avoid memcpy for the middle of the chunk if chunk.len() >= BUFFER_SIZE { - self.flush().await?; + self.flush(ctx).await?; // do a big write, bypassing `buf` assert_eq!( self.buf @@ -90,7 +100,7 @@ where .len(), 0 ); - let (nwritten, chunk) = self.writer.write_all(chunk).await?; + let (nwritten, chunk) = self.writer.write_all(chunk, ctx).await?; assert_eq!(nwritten, chunk_len); return Ok((nwritten, chunk)); } @@ -106,7 +116,7 @@ where slice = &slice[n..]; if buf.len() >= BUFFER_SIZE { assert_eq!(buf.len(), BUFFER_SIZE); - self.flush().await?; + self.flush(ctx).await?; } } assert!(slice.is_empty(), "by now we should have drained the chunk"); @@ -115,7 +125,11 @@ where /// Always goes through the internal buffer. /// Guaranteed to never invoke [`OwnedAsyncWrite::write_all_borrowed`] on the underlying. - pub async fn write_all_borrowed(&mut self, mut chunk: &[u8]) -> std::io::Result { + pub async fn write_all_borrowed( + &mut self, + mut chunk: &[u8], + ctx: &RequestContext, + ) -> std::io::Result { let chunk_len = chunk.len(); while !chunk.is_empty() { let buf = self.buf.as_mut().expect("must not use after an error"); @@ -126,20 +140,20 @@ where chunk = &chunk[n..]; if buf.len() >= BUFFER_SIZE { assert_eq!(buf.len(), BUFFER_SIZE); - self.flush().await?; + self.flush(ctx).await?; } } Ok(chunk_len) } - async fn flush(&mut self) -> std::io::Result<()> { + async fn flush(&mut self, ctx: &RequestContext) -> std::io::Result<()> { let buf = self.buf.take().expect("must not use after an error"); if buf.is_empty() { self.buf = Some(buf); return std::io::Result::Ok(()); } let buf_len = buf.len(); - let (nwritten, mut buf) = self.writer.write_all(buf).await?; + let (nwritten, mut buf) = self.writer.write_all(buf, ctx).await?; assert_eq!(nwritten, buf_len); buf.clear(); self.buf = Some(buf); @@ -158,17 +172,22 @@ impl OwnedAsyncWriter >( &mut self, buf: B, + ctx: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { let nbytes = buf.bytes_init(); if nbytes == 0 { return Ok((0, Slice::into_inner(buf.slice_full()))); } let slice = buf.slice(0..nbytes); - BufferedWriter::write_buffered(self, slice).await + BufferedWriter::write_buffered(self, slice, ctx).await } #[inline(always)] - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { - BufferedWriter::write_all_borrowed(self, buf).await + async fn write_all_borrowed( + &mut self, + buf: &[u8], + ctx: &RequestContext, + ) -> std::io::Result { + BufferedWriter::write_all_borrowed(self, buf, ctx).await } } @@ -180,6 +199,7 @@ impl OwnedAsyncWriter for Vec { >( &mut self, buf: B, + _: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { let nbytes = buf.bytes_init(); if nbytes == 0 { @@ -190,7 +210,11 @@ impl OwnedAsyncWriter for Vec { Ok((buf.len(), Slice::into_inner(buf))) } - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { + async fn write_all_borrowed( + &mut self, + buf: &[u8], + _ctx: &RequestContext, + ) -> std::io::Result { self.extend_from_slice(&buf[..]); Ok(buf.len()) } @@ -199,6 +223,8 @@ impl OwnedAsyncWriter for Vec { #[cfg(test)] mod tests { use super::*; + use crate::context::{DownloadBehavior, RequestContext}; + use crate::task_mgr::TaskKind; #[derive(Default)] struct RecorderWriter { @@ -212,6 +238,7 @@ mod tests { >( &mut self, buf: B, + _: &RequestContext, ) -> std::io::Result<(usize, B::Buf)> { let nbytes = buf.bytes_init(); if nbytes == 0 { @@ -223,16 +250,24 @@ mod tests { Ok((buf.len(), Slice::into_inner(buf))) } - async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result { + async fn write_all_borrowed( + &mut self, + buf: &[u8], + _: &RequestContext, + ) -> std::io::Result { self.writes.push(Vec::from(buf)); Ok(buf.len()) } } + fn test_ctx() -> RequestContext { + RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Error) + } + macro_rules! write { ($writer:ident, $data:literal) => {{ $writer - .write_buffered(::bytes::Bytes::from_static($data).slice_full()) + .write_buffered(::bytes::Bytes::from_static($data).slice_full(), &test_ctx()) .await?; }}; } @@ -246,7 +281,7 @@ mod tests { write!(writer, b"c"); write!(writer, b"d"); write!(writer, b"e"); - let recorder = writer.flush_and_into_inner().await?; + let recorder = writer.flush_and_into_inner(&test_ctx()).await?; assert_eq!( recorder.writes, vec![Vec::from(b"ab"), Vec::from(b"cd"), Vec::from(b"e")] @@ -262,7 +297,7 @@ mod tests { write!(writer, b"de"); write!(writer, b""); write!(writer, b"fghijk"); - let recorder = writer.flush_and_into_inner().await?; + let recorder = writer.flush_and_into_inner(&test_ctx()).await?; assert_eq!( recorder.writes, vec![Vec::from(b"abc"), Vec::from(b"de"), Vec::from(b"fghijk")] @@ -278,7 +313,7 @@ mod tests { write!(writer, b"bc"); write!(writer, b"d"); write!(writer, b"e"); - let recorder = writer.flush_and_into_inner().await?; + let recorder = writer.flush_and_into_inner(&test_ctx()).await?; assert_eq!( recorder.writes, vec![Vec::from(b"a"), Vec::from(b"bc"), Vec::from(b"de")] @@ -290,16 +325,16 @@ mod tests { async fn test_write_all_borrowed_always_goes_through_buffer() -> std::io::Result<()> { let recorder = RecorderWriter::default(); let mut writer = BufferedWriter::<2, _>::new(recorder); + let ctx = test_ctx(); + writer.write_all_borrowed(b"abc", &ctx).await?; + writer.write_all_borrowed(b"d", &ctx).await?; + writer.write_all_borrowed(b"e", &ctx).await?; + writer.write_all_borrowed(b"fg", &ctx).await?; + writer.write_all_borrowed(b"hi", &ctx).await?; + writer.write_all_borrowed(b"j", &ctx).await?; + writer.write_all_borrowed(b"klmno", &ctx).await?; - writer.write_all_borrowed(b"abc").await?; - writer.write_all_borrowed(b"d").await?; - writer.write_all_borrowed(b"e").await?; - writer.write_all_borrowed(b"fg").await?; - writer.write_all_borrowed(b"hi").await?; - writer.write_all_borrowed(b"j").await?; - writer.write_all_borrowed(b"klmno").await?; - - let recorder = writer.flush_and_into_inner().await?; + let recorder = writer.flush_and_into_inner(&ctx).await?; assert_eq!( recorder.writes, {