rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-07-07 06:00:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP: page_cache_priming_writer: plumb through RequestContext for previous commit, yet more churn -,-

Browse Source
This commit is contained in:
Christian Schwarz
2024-03-14 16:06:04 +00:00
parent f1f0452722
commit eb1ccd7988
5 changed files with 127 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
pageserver/src/tenant/ephemeral_file.rs
View File

@@ -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<u64, io::Error> {
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

4
pageserver/src/virtual_file.rs
View File

@@ -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<usize> {
async fn write_all_borrowed(&mut self, _buf: &[u8], _ctx: &RequestContext) -> std::io::Result<usize> {
// TODO: ensure this through the type system
panic!("this should not happen");
}

48
pageserver/src/virtual_file/owned_buffers_io/util/page_cache_priming_writer.rs
View File

@@ -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<const N: usize, W: OwnedAsyncWriter, C: Cache<N, W>> {
}
pub trait Cache<const PAGE_SZ: usize, W> {
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<const N: usize, W, C> Writer<N, W, C>
@@ -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<usize> {
async fn write_all_borrowed(&mut self, _: &[u8], _: &RequestContext) -> std::io::Result<usize> {
// 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<usize> {
async fn write_all_borrowed(&mut self, _: &[u8], _: &RequestContext) -> std::io::Result<usize> {
// TODO: use type system to ensure this doesn't happen at runtime
panic!("don't put the types together this way")
}

15
pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs
View File

@@ -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<W> {
@@ -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<usize> {
let nwritten = self.dst.write_all_borrowed(buf).await?;
async fn write_all_borrowed(
&mut self,
buf: &[u8],
ctx: &RequestContext,
) -> std::io::Result<usize> {
let nwritten = self.dst.write_all_borrowed(buf, ctx).await?;
self.bytes_amount += u64::try_from(nwritten).unwrap();
Ok(nwritten)
}

97
pageserver/src/virtual_file/owned_buffers_io/write.rs
View File

@@ -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<usize>;
async fn write_all_borrowed(
&mut self,
buf: &[u8],
ctx: &RequestContext,
) -> std::io::Result<usize>;
}
/// 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<W> {
self.flush().await?;
pub async fn flush_and_into_inner(mut self, ctx: &RequestContext) -> std::io::Result<W> {
self.flush(ctx).await?;
let Self { buf, writer } = self;
assert!(buf.is_some());
Ok(writer)
}
pub async fn write_buffered<B: IoBuf>(&mut self, chunk: Slice<B>) -> std::io::Result<(usize, B)>
pub async fn write_buffered<B: IoBuf>(
&mut self,
chunk: Slice<B>,
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<usize> {
pub async fn write_all_borrowed(
&mut self,
mut chunk: &[u8],
ctx: &RequestContext,
) -> std::io::Result<usize> {
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<const BUFFER_SIZE: usize, W: OwnedAsyncWriter> 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<usize> {
BufferedWriter::write_all_borrowed(self, buf).await
async fn write_all_borrowed(
&mut self,
buf: &[u8],
ctx: &RequestContext,
) -> std::io::Result<usize> {
BufferedWriter::write_all_borrowed(self, buf, ctx).await
}
}
@@ -180,6 +199,7 @@ impl OwnedAsyncWriter for Vec<u8> {
>(
&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<u8> {
Ok((buf.len(), Slice::into_inner(buf)))
}
async fn write_all_borrowed(&mut self, buf: &[u8]) -> std::io::Result<usize> {
async fn write_all_borrowed(
&mut self,
buf: &[u8],
_ctx: &RequestContext,
) -> std::io::Result<usize> {
self.extend_from_slice(&buf[..]);
Ok(buf.len())
}
@@ -199,6 +223,8 @@ impl OwnedAsyncWriter for Vec<u8> {
#[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<usize> {
async fn write_all_borrowed(
&mut self,
buf: &[u8],
_: &RequestContext,
) -> std::io::Result<usize> {
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,
{