WIP: page_cache_priming_writer (is it really worth it?)

pageserver/src/tenant/ephemeral_file.rs

@@ -3,8 +3,9 @@
 
 use crate::config::PageServerConf;
 use crate::context::RequestContext;
-use crate::page_cache::{self, PAGE_SZ};
+use crate::page_cache::{self, PageCache, 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;
@@ -31,7 +32,13 @@ pub struct EphemeralFile {
     file: owned_buffers_io::util::size_tracking_writer::Writer<
         owned_buffers_io::write::BufferedWriter<
             { Self::TAIL_SZ },
-            owned_buffers_io::util::size_tracking_writer::Writer<VirtualFile>,
+            owned_buffers_io::util::size_tracking_writer::Writer<
+                owned_buffers_io::util::page_cache_priming_writer::Writer<
+                    PAGE_SZ,
+                    owned_buffers_io::util::page_cache_priming_writer::VirtualFileAdaptor,
+                    &'static crate::page_cache::PageCache,
+                >,
+            >,
         >,
     >,
 }
@@ -62,12 +69,19 @@ impl EphemeralFile {
                 .create(true),
         )
         .await?;
+
+        let page_cache_file_id = page_cache::next_file_id();
+        let file = owned_buffers_io::util::page_cache_priming_writer::VirtualFileAdaptor::new(
+            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::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);
 
         Ok(EphemeralFile {
-            page_cache_file_id: page_cache::next_file_id(),
+            page_cache_file_id,
             _tenant_shard_id: tenant_shard_id,
             _timeline_id: timeline_id,
             file,
@@ -110,7 +124,13 @@ impl EphemeralFile {
                         format!(
                             "ephemeral file: read immutable page #{}: {}: {:#}",
                             blknum,
-                            self.file.as_inner().as_inner().as_inner().path,
+                            self.file
+                                .as_inner()
+                                .as_inner()
+                                .as_inner()
+                                .as_inner()
+                                .as_inner()
+                                .path,
                             e,
                         ),
                     )
@@ -124,6 +144,8 @@ impl EphemeralFile {
                         .as_inner()
                         .as_inner()
                         .as_inner()
+                        .as_inner()
+                        .as_inner()
                         .read_exact_at_page(write_guard, blknum as u64 * PAGE_SZ as u64)
                         .await?;
                     let read_guard = write_guard.mark_valid();
@@ -182,8 +204,6 @@ impl EphemeralFile {
         // Write the payload
         self.file.write_all_borrowed(&srcbuf).await?;
 
-        // TODO: bring back pre-warming of page cache, using another sandwich layer
-
         Ok(pos)
     }
 }
@@ -203,7 +223,7 @@ 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().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
@@ -212,7 +232,13 @@ impl Drop for EphemeralFile {
                 // not found files might also be related to https://github.com/neondatabase/neon/issues/2442
                 error!(
                     "could not remove ephemeral file '{}': {}",
-                    self.file.as_inner().as_inner().as_inner().path,
+                    self.file
+                        .as_inner()
+                        .as_inner()
+                        .as_inner()
+                        .as_inner()
+                        .as_inner()
+                        .path,
                     e
                 );
             }

pageserver/src/virtual_file.rs

@@ -48,6 +48,7 @@ pub(crate) mod owned_buffers_io {
 
     pub(crate) mod write;
     pub(crate) mod util {
+        pub(crate) mod page_cache_priming_writer;
         pub(crate) mod size_tracking_writer;
     }
 }
@@ -1084,7 +1085,11 @@ impl Drop for VirtualFile {
 
 impl OwnedAsyncWriter for VirtualFile {
     #[inline(always)]
-    async fn write_all<B: BoundedBuf<Buf = Buf>, Buf: IoBuf + Send>(
+    async fn write_all<
+        B: BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
         &mut self,
         buf: B,
     ) -> std::io::Result<(usize, B::Buf)> {

pageserver/src/virtual_file/owned_buffers_io/util/page_cache_priming_writer.rs

@@ -0,0 +1,153 @@
+use std::ops::{Deref, Range};
+
+use tokio_epoll_uring::BoundedBuf;
+use tracing::error;
+
+use crate::{
+    page_cache::{self, PAGE_SZ},
+    virtual_file::{owned_buffers_io::write::OwnedAsyncWriter, VirtualFile},
+};
+
+pub struct Writer<const N: usize, W: OwnedAsyncWriter, C: Cache<N, W>> {
+    under: W,
+    written: u64,
+    cache: C,
+}
+
+pub trait Cache<const PAGE_SZ: usize, W> {
+    async fn fill_cache(&self, under: &W, page_no: u64, contents: &[u8; PAGE_SZ]);
+}
+
+impl<const N: usize, W, C> Writer<N, W, C>
+where
+    C: Cache<N, W>,
+    W: OwnedAsyncWriter,
+{
+    pub fn new(under: W, cache: C) -> Self {
+        Self {
+            under,
+            written: 0,
+            cache,
+        }
+    }
+    pub fn as_inner(&self) -> &W {
+        &self.under
+    }
+    fn invariants(&self) {
+        assert_eq!(
+            self.written % (N as u64),
+            0,
+            "writes must happen in multiples of N"
+        );
+    }
+}
+
+impl<const N: usize, W, C> OwnedAsyncWriter for Writer<N, W, C>
+where
+    W: OwnedAsyncWriter,
+    C: Cache<N, W>,
+{
+    async fn write_all<
+        B: tokio_epoll_uring::BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: tokio_epoll_uring::IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
+        &mut self,
+        buf: B,
+    ) -> 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 = if let Ok((nwritten, buf)) = res {
+            assert_eq!(nwritten % N, 0);
+            let buf = tokio_epoll_uring::Slice::from_buf_bounds(buf, saved_bounds);
+            if let Some(before) = debug_assert_contents_eq {
+                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!();
+                self.cache
+                    .fill_cache(&self.under, self.written / (N as u64), page)
+                    .await;
+                self.written += (N as u64);
+            }
+            self.written += nwritten as u64;
+            Ok((nwritten, tokio_epoll_uring::Slice::into_inner(buf)))
+        } else {
+            res
+        };
+        self.invariants();
+        res
+    }
+
+    async fn write_all_borrowed(&mut self, buf: &[u8]) -> 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")
+    }
+}
+
+pub struct VirtualFileAdaptor {
+    file: VirtualFile,
+    file_id: page_cache::FileId,
+}
+
+impl VirtualFileAdaptor {
+    pub fn new(file: VirtualFile, file_id: page_cache::FileId) -> Self {
+        Self { file, file_id }
+    }
+    pub fn as_inner(&self) -> &VirtualFile {
+        &self.file
+    }
+}
+
+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]) {
+        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"),
+            )
+            .await
+        {
+            Ok(crate::page_cache::ReadBufResult::Found(guard)) => {
+                debug_assert_eq!(guard.deref(), contents);
+            }
+            Ok(crate::page_cache::ReadBufResult::NotFound(guard)) => {
+                guard.copy_from_slice(contents);
+                guard.mark_valid();
+            }
+            Err(e) => {
+                error!("failed to get immutable buf to pre-warm page cache: {e:?}");
+            }
+        }
+    }
+}
+
+impl OwnedAsyncWriter for VirtualFileAdaptor {
+    async fn write_all<
+        B: tokio_epoll_uring::BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: tokio_epoll_uring::IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
+        &mut self,
+        buf: B,
+    ) -> std::io::Result<(usize, B::Buf)> {
+        OwnedAsyncWriter::write_all(&mut self.file, buf).await
+    }
+
+    async fn write_all_borrowed(&mut self, buf: &[u8]) -> 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")
+    }
+}

pageserver/src/virtual_file/owned_buffers_io/util/size_tracking_writer.rs

@@ -1,3 +1,5 @@
+use std::ops::Range;
+
 use crate::virtual_file::owned_buffers_io::write::OwnedAsyncWriter;
 use tokio_epoll_uring::{BoundedBuf, IoBuf};
 
@@ -31,7 +33,11 @@ where
     W: OwnedAsyncWriter,
 {
     #[inline(always)]
-    async fn write_all<B: BoundedBuf<Buf = Buf>, Buf: IoBuf + Send>(
+    async fn write_all<
+        B: BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
         &mut self,
         buf: B,
     ) -> std::io::Result<(usize, B::Buf)> {

pageserver/src/virtual_file/owned_buffers_io/write.rs

@@ -1,9 +1,15 @@
+use std::ops::Range;
+
 use tokio_epoll_uring::{BoundedBuf, IoBuf, Slice};
 
 /// A trait for doing owned-buffer write IO.
 /// Think [`tokio::io::AsyncWrite`] but with owned buffers.
 pub trait OwnedAsyncWriter {
-    async fn write_all<B: BoundedBuf<Buf = Buf>, Buf: IoBuf + Send>(
+    async fn write_all<
+        B: BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
         &mut self,
         buf: B,
     ) -> std::io::Result<(usize, B::Buf)>;
@@ -145,7 +151,11 @@ impl<const BUFFER_SIZE: usize, W: OwnedAsyncWriter> OwnedAsyncWriter
     for BufferedWriter<BUFFER_SIZE, W>
 {
     #[inline(always)]
-    async fn write_all<B: BoundedBuf<Buf = Buf>, Buf: IoBuf + Send>(
+    async fn write_all<
+        B: BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
         &mut self,
         buf: B,
     ) -> std::io::Result<(usize, B::Buf)> {
@@ -163,7 +173,11 @@ impl<const BUFFER_SIZE: usize, W: OwnedAsyncWriter> OwnedAsyncWriter
 }
 
 impl OwnedAsyncWriter for Vec<u8> {
-    async fn write_all<B: BoundedBuf<Buf = Buf>, Buf: IoBuf + Send>(
+    async fn write_all<
+        B: BoundedBuf<Buf = Buf, Bounds = Bounds>,
+        Buf: IoBuf + Send,
+        Bounds: std::ops::RangeBounds<usize>,
+    >(
         &mut self,
         buf: B,
     ) -> std::io::Result<(usize, B::Buf)> {
@@ -191,7 +205,11 @@ mod tests {
         writes: Vec<Vec<u8>>,
     }
     impl OwnedAsyncWriter for RecorderWriter {
-        async fn write_all<B: BoundedBuf<Buf = Buf>, Buf: IoBuf + Send>(
+        async fn write_all<
+            B: BoundedBuf<Buf = Buf, Bounds = Bounds>,
+            Buf: IoBuf + Send,
+            Bounds: std::ops::RangeBounds<usize>,
+        >(
             &mut self,
             buf: B,
         ) -> std::io::Result<(usize, B::Buf)> {