Use a chunked byte buffer to hold page versions in in-memory layer.

Change the way the page versions are stored in an in-memory layer.
Instead of storing the original PageVersion structs, serialize them
into an expandable buffer. That has a couple of advantages: first, the
serialized form is more compact, which greatly reduces the memory
usage. Secondly, we can track the memory used more accurately.

This moves the serialization work from the checkpointer to the WAL
receiver. It also means that if a page in an in-memory layer is
accessed, it needs to be deserialized on the GetPage call. That's adds
a bit of latency, but it's not significant compared to the overhead of
replaying WAL records. The simpler memory accounting is worth it.

Per Konstantin's idea a long time ago.

pageserver/src/layered_repository/blob.rs

@@ -1,4 +1,4 @@
-use std::{fs::File, io::Write};
+use std::{fs::File, io::Read, io::Write};
 
 use anyhow::Result;
 use bookfile::{BookWriter, BoundedReader, ChapterId, ChapterWriter};
@@ -28,14 +28,14 @@ impl<W: Write> BlobWriter<W> {
         Self { writer, offset: 0 }
     }
 
-    pub fn write_blob(&mut self, blob: &[u8]) -> Result<BlobRange> {
-        self.writer.write_all(blob)?;
+    pub fn write_blob_from_reader(&mut self, r: &mut impl Read) -> Result<BlobRange> {
+        let len = std::io::copy(r, &mut self.writer)?;
 
         let range = BlobRange {
             offset: self.offset,
-            size: blob.len(),
+            size: len as usize,
         };
-        self.offset += blob.len() as u64;
+        self.offset += len as u64;
         Ok(range)
     }
 

pageserver/src/layered_repository/delta_layer.rs

@@ -39,6 +39,7 @@
 //!
 use crate::layered_repository::blob::BlobWriter;
 use crate::layered_repository::filename::{DeltaFileName, PathOrConf};
+use crate::layered_repository::page_versions::PageVersions;
 use crate::layered_repository::storage_layer::{
     Layer, PageReconstructData, PageReconstructResult, PageVersion, SegmentTag,
 };
@@ -350,14 +351,14 @@ impl DeltaLayer {
     }
 
     /// Create a new delta file, using the given page versions and relsizes.
-    /// The page versions are passed by an iterator; the iterator must return
-    /// page versions in blknum+lsn order.
+    /// The page versions are passed in a PageVersions struct. If 'cutoff' is
+    /// given, only page versions with LSN < cutoff are included.
     ///
-    /// This is used to write the in-memory layer to disk. The in-memory layer uses the same
-    /// data structure with two btreemaps as we do, so passing the btreemaps is currently
-    /// expedient.
+    /// This is used to write the in-memory layer to disk. The page_versions and
+    /// relsizes are thus passed in the same format as they are in the in-memory
+    /// layer, as that's expedient.
     #[allow(clippy::too_many_arguments)]
-    pub fn create<'a>(
+    pub fn create(
         conf: &'static PageServerConf,
         timelineid: ZTimelineId,
         tenantid: ZTenantId,
@@ -365,7 +366,8 @@ impl DeltaLayer {
         start_lsn: Lsn,
         end_lsn: Lsn,
         dropped: bool,
-        page_versions: impl Iterator<Item = (u32, Lsn, &'a PageVersion)>,
+        page_versions: &PageVersions,
+        cutoff: Option<Lsn>,
         relsizes: VecMap<Lsn, u32>,
     ) -> Result<DeltaLayer> {
         if seg.rel.is_blocky() {
@@ -399,9 +401,10 @@ impl DeltaLayer {
 
         let mut page_version_writer = BlobWriter::new(book, PAGE_VERSIONS_CHAPTER);
 
-        for (blknum, lsn, page_version) in page_versions {
-            let buf = PageVersion::ser(page_version)?;
-            let blob_range = page_version_writer.write_blob(&buf)?;
+        let page_versions_iter = page_versions.ordered_page_version_iter(cutoff);
+        for (blknum, lsn, pos) in page_versions_iter {
+            let blob_range =
+                page_version_writer.write_blob_from_reader(&mut page_versions.reader(pos)?)?;
 
             inner
                 .page_version_metas

pageserver/src/layered_repository/inmemory_layer.rs

@@ -184,10 +184,10 @@ impl Layer for InMemoryLayer {
                 .get_block_lsn_range(blknum, ..=lsn)
                 .iter()
                 .rev();
-            for (entry_lsn, pv) in iter {
-                match pv {
+            for (entry_lsn, pos) in iter {
+                match inner.page_versions.get_page_version(*pos)? {
                     PageVersion::Page(img) => {
-                        reconstruct_data.page_img = Some(img.clone());
+                        reconstruct_data.page_img = Some(img);
                         need_image = false;
                         break;
                     }
@@ -285,8 +285,8 @@ impl Layer for InMemoryLayer {
             println!("segsizes {}: {}", k, v);
         }
 
-        for (blknum, lsn, pv) in inner.page_versions.ordered_page_version_iter(None) {
-            let pv_description = match pv {
+        for (blknum, lsn, pos) in inner.page_versions.ordered_page_version_iter(None) {
+            let pv_description = match inner.page_versions.get_page_version(pos)? {
                 PageVersion::Page(_img) => "page",
                 PageVersion::Wal(_rec) => "wal",
             };
@@ -610,7 +610,8 @@ impl InMemoryLayer {
                 self.start_lsn,
                 end_lsn_exclusive,
                 true,
-                inner.page_versions.ordered_page_version_iter(None),
+                &inner.page_versions,
+                None,
                 inner.segsizes.clone(),
             )?;
             trace!(
@@ -627,13 +628,9 @@ impl InMemoryLayer {
 
         // Since `end_lsn` is inclusive, subtract 1.
         // We want to make an ImageLayer for the last included LSN,
-        // so the DeltaLayer should exlcude that LSN.
+        // so the DeltaLayer should exclude that LSN.
         let end_lsn_inclusive = Lsn(end_lsn_exclusive.0 - 1);
 
-        let mut page_versions = inner
-            .page_versions
-            .ordered_page_version_iter(Some(end_lsn_inclusive));
-
         let mut delta_layers = Vec::new();
 
         if self.start_lsn != end_lsn_inclusive {
@@ -647,7 +644,8 @@ impl InMemoryLayer {
                 self.start_lsn,
                 end_lsn_inclusive,
                 false,
-                page_versions,
+                &inner.page_versions,
+                Some(end_lsn_inclusive),
                 segsizes,
             )?;
             delta_layers.push(delta_layer);
@@ -658,7 +656,11 @@ impl InMemoryLayer {
                 end_lsn_inclusive
             );
         } else {
-            assert!(page_versions.next().is_none());
+            assert!(inner
+                .page_versions
+                .ordered_page_version_iter(None)
+                .next()
+                .is_none());
         }
 
         drop(inner);

pageserver/src/layered_repository/page_versions.rs

@@ -1,13 +1,27 @@
 use std::{collections::HashMap, ops::RangeBounds, slice};
 
+use anyhow::Result;
+
+use std::io::{Read, Write};
+
+use zenith_utils::chunked_buffer::{ChunkedBuffer, ChunkedBufferReader};
 use zenith_utils::{lsn::Lsn, vec_map::VecMap};
 
 use super::storage_layer::PageVersion;
 
-const EMPTY_SLICE: &[(Lsn, PageVersion)] = &[];
+use zenith_utils::bin_ser::LeSer;
+
+const EMPTY_SLICE: &[(Lsn, u64)] = &[];
 
 #[derive(Debug, Default)]
-pub struct PageVersions(HashMap<u32, VecMap<Lsn, PageVersion>>);
+pub struct PageVersions {
+    map: HashMap<u32, VecMap<Lsn, u64>>,
+
+    /// The PageVersion structs are stored in a serialized format in this buffer.
+    /// Each serialized PageVersion is preceded by a 'u32' length field.
+    /// The 'map' stores offsets into this buffer.
+    buffer: ChunkedBuffer,
+}
 
 impl PageVersions {
     pub fn append_or_update_last(
@@ -15,26 +29,35 @@ impl PageVersions {
         blknum: u32,
         lsn: Lsn,
         page_version: PageVersion,
-    ) -> Option<PageVersion> {
-        let map = self.0.entry(blknum).or_insert_with(VecMap::default);
-        map.append_or_update_last(lsn, page_version).unwrap()
+    ) -> Option<u64> {
+        // remember starting position
+        let pos = self.buffer.len();
+
+        // make room for the 'length' field by writing zeros as a placeholder.
+        self.buffer.write_all(&[0u8; 4]).unwrap();
+
+        page_version.ser_into(&mut self.buffer).unwrap();
+
+        // write the 'length' field.
+        let len = self.buffer.len() - pos - 4;
+        let lenbuf = u32::to_ne_bytes(len as u32);
+        self.buffer.write_all_at(&lenbuf, pos).unwrap();
+
+        let map = self.map.entry(blknum).or_insert_with(VecMap::default);
+        map.append_or_update_last(lsn, pos as u64).unwrap()
     }
 
     /// Get all [`PageVersion`]s in a block
-    pub fn get_block_slice(&self, blknum: u32) -> &[(Lsn, PageVersion)] {
-        self.0
+    fn get_block_slice(&self, blknum: u32) -> &[(Lsn, u64)] {
+        self.map
             .get(&blknum)
             .map(VecMap::as_slice)
             .unwrap_or(EMPTY_SLICE)
     }
 
     /// Get a range of [`PageVersions`] in a block
-    pub fn get_block_lsn_range<R: RangeBounds<Lsn>>(
-        &self,
-        blknum: u32,
-        range: R,
-    ) -> &[(Lsn, PageVersion)] {
-        self.0
+    pub fn get_block_lsn_range<R: RangeBounds<Lsn>>(&self, blknum: u32, range: R) -> &[(Lsn, u64)] {
+        self.map
             .get(&blknum)
             .map(|vec_map| vec_map.slice_range(range))
             .unwrap_or(EMPTY_SLICE)
@@ -43,7 +66,7 @@ impl PageVersions {
     /// Iterate through [`PageVersion`]s in (block, lsn) order.
     /// If a [`cutoff_lsn`] is set, only show versions with `lsn < cutoff_lsn`
     pub fn ordered_page_version_iter(&self, cutoff_lsn: Option<Lsn>) -> OrderedPageVersionIter<'_> {
-        let mut ordered_blocks: Vec<u32> = self.0.keys().cloned().collect();
+        let mut ordered_blocks: Vec<u32> = self.map.keys().cloned().collect();
         ordered_blocks.sort_unstable();
 
         let slice = ordered_blocks
@@ -59,6 +82,22 @@ impl PageVersions {
             cur_slice_iter: slice.iter(),
         }
     }
+
+    /// Returns a 'Read' that reads the page version at given offset.
+    pub fn reader(&self, pos: u64) -> Result<std::io::Take<ChunkedBufferReader>, std::io::Error> {
+        // read length
+        let mut lenbuf = [0u8; 4];
+        let mut reader = self.buffer.reader(pos);
+        reader.read_exact(&mut lenbuf)?;
+        let len = u32::from_ne_bytes(lenbuf);
+
+        Ok(reader.take(len as u64))
+    }
+
+    pub fn get_page_version(&self, pos: u64) -> Result<PageVersion> {
+        let mut reader = self.reader(pos)?;
+        Ok(PageVersion::des_from(&mut reader)?)
+    }
 }
 
 pub struct OrderedPageVersionIter<'a> {
@@ -69,7 +108,7 @@ pub struct OrderedPageVersionIter<'a> {
 
     cutoff_lsn: Option<Lsn>,
 
-    cur_slice_iter: slice::Iter<'a, (Lsn, PageVersion)>,
+    cur_slice_iter: slice::Iter<'a, (Lsn, u64)>,
 }
 
 impl OrderedPageVersionIter<'_> {
@@ -83,14 +122,14 @@ impl OrderedPageVersionIter<'_> {
 }
 
 impl<'a> Iterator for OrderedPageVersionIter<'a> {
-    type Item = (u32, Lsn, &'a PageVersion);
+    type Item = (u32, Lsn, u64);
 
     fn next(&mut self) -> Option<Self::Item> {
         loop {
-            if let Some((lsn, page_version)) = self.cur_slice_iter.next() {
+            if let Some((lsn, pos)) = self.cur_slice_iter.next() {
                 if self.is_lsn_before_cutoff(lsn) {
                     let blknum = self.ordered_blocks[self.cur_block_idx];
-                    return Some((blknum, *lsn, page_version));
+                    return Some((blknum, *lsn, *pos));
                 }
             }
 

zenith_utils/src/chunked_buffer.rs

@@ -0,0 +1,185 @@
+use std::cmp::min;
+use std::io::{Read, Write};
+
+const CHUNK_SIZE: usize = 1024;
+
+///
+/// ChunkedBuffer is an expandable byte buffer. You can append to the end of it, and you
+/// can read at any byte position. The advantage over a plain Vec<u8> is that the
+/// buffer consists of smaller chunks, so that a big buffer doesn't require one huge
+/// allocation, and expanding doesn't require copying all the data.
+///
+#[derive(Debug, Default)]
+pub struct ChunkedBuffer {
+    // Clippy considers it unnecessary to have Box here, since the vector is stored
+    // on the heap anyway. But we have our reasons for that: we want each chunk to
+    // be allocated separately, so that we don't require one huge allocation, and so
+    // that expanding the array doesn't require copying all of the data. (Clippy
+    // knows about that and doesn't complain if the array is large enough, 4096
+    // bytes is the default threshold. Our chunk size is smaller than that so that
+    // heuristic doesn't save us.)
+    #[allow(clippy::vec_box)]
+    chunks: Vec<Box<[u8; CHUNK_SIZE]>>,
+    len: usize,
+}
+
+impl ChunkedBuffer {
+    /// Return current length of the buffer, in bytes
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len == 0
+    }
+
+    /// Create a "cursor" for reading from the buffer.
+    pub fn reader(&self, pos: u64) -> ChunkedBufferReader {
+        ChunkedBufferReader {
+            parent: self,
+            pos: pos as usize,
+        }
+    }
+
+    /// Write 'buf' to the given byte position in the buffer.
+    ///
+    /// The buffer is expanded if needed. If 'pos' is past the end of the buffer,
+    /// the "gap" is filled with zeros.
+    ///
+    /// This can stop short, depending on where the chunk boundaries are. But
+    /// always writes at least one byte.
+    pub fn write_at(&mut self, buf: &[u8], pos: usize) -> Result<usize, std::io::Error> {
+        let chunk_idx = pos / CHUNK_SIZE;
+        let chunk_off = pos % CHUNK_SIZE;
+
+        while chunk_idx >= self.chunks.len() {
+            self.chunks.push(Box::new([0; CHUNK_SIZE]));
+        }
+
+        let n = min(CHUNK_SIZE - chunk_off, buf.len());
+
+        let chunk = &mut self.chunks[chunk_idx];
+
+        chunk[chunk_off..(chunk_off + n)].copy_from_slice(&buf[0..n]);
+
+        if pos + n > self.len {
+            self.len = pos + n;
+        }
+        Ok(n)
+    }
+
+    /// Like 'write_at', but doesn't stop until the whole buffer has been written.
+    pub fn write_all_at(&mut self, buf: &[u8], pos: usize) -> Result<(), std::io::Error> {
+        let mut nwritten = 0;
+
+        while nwritten < buf.len() {
+            nwritten += self.write_at(&buf[nwritten..], pos + nwritten)?;
+        }
+        Ok(())
+    }
+}
+
+impl Write for ChunkedBuffer {
+    /// Append to the end of the buffer.
+    ///
+    /// Note how this interacts with 'write_at'. If you use write_at to expand the
+    /// buffer, the "write position" for this function changes too.
+    fn write(&mut self, buf: &[u8]) -> Result<usize, std::io::Error> {
+        self.write_at(buf, self.len)
+    }
+
+    fn flush(&mut self) -> Result<(), std::io::Error> {
+        Ok(())
+    }
+}
+
+pub struct ChunkedBufferReader<'a> {
+    parent: &'a ChunkedBuffer,
+    pos: usize,
+}
+
+impl<'a> Read for ChunkedBufferReader<'a> {
+    fn read(&mut self, dst: &mut [u8]) -> Result<usize, std::io::Error> {
+        if self.pos >= self.parent.len {
+            return Ok(0);
+        }
+
+        let chunk_idx = self.pos / CHUNK_SIZE;
+        let chunk_off = self.pos % CHUNK_SIZE;
+
+        let len = min(
+            min(self.parent.len - self.pos, CHUNK_SIZE - chunk_off),
+            dst.len(),
+        );
+
+        let chunk = &self.parent.chunks[chunk_idx];
+
+        dst[0..len].copy_from_slice(&chunk[chunk_off..(chunk_off + len)]);
+
+        self.pos += len;
+
+        Ok(len)
+    }
+}
+
+#[test]
+fn test_chunked_buffer() -> Result<(), std::io::Error> {
+    let mut chunked_buffer = ChunkedBuffer::default();
+
+    let mut testdata = [0u8; 10000];
+    for (i, b) in testdata.iter_mut().enumerate() {
+        *b = (i % 256) as u8;
+    }
+
+    chunked_buffer.write_all(&testdata)?;
+
+    // Test reader
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader(0);
+    reader.read_exact(&mut resultdata)?;
+    assert_eq!(resultdata, testdata);
+
+    // test zero-length read
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader(0);
+    reader.read_exact(&mut resultdata[0..0])?;
+
+    // test reads around chunk boundaries
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader(0);
+    reader.read_exact(&mut resultdata[0..1])?;
+    reader.read_exact(&mut resultdata[1..CHUNK_SIZE])?;
+    assert_eq!(resultdata[0..CHUNK_SIZE], testdata[0..CHUNK_SIZE]);
+
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader(10);
+    reader.read_exact(&mut resultdata[10..CHUNK_SIZE + 10])?;
+    assert_eq!(
+        resultdata[10..CHUNK_SIZE + 10],
+        testdata[10..CHUNK_SIZE + 10]
+    );
+
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader((CHUNK_SIZE - 1) as u64);
+    reader.read_exact(&mut resultdata[(CHUNK_SIZE - 1)..CHUNK_SIZE])?;
+    assert_eq!(resultdata[CHUNK_SIZE - 1], testdata[CHUNK_SIZE - 1]);
+    reader.read_exact(&mut resultdata[(CHUNK_SIZE)..(CHUNK_SIZE + 1)])?;
+    assert_eq!(resultdata[CHUNK_SIZE], testdata[CHUNK_SIZE]);
+
+    // Read at the end
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader(9900);
+    reader.read_exact(&mut resultdata[9900..10000])?;
+    assert_eq!(resultdata[9900..10000], testdata[9900..10000]);
+
+    // Read past the end
+    let mut resultdata = [0u8; 10001];
+    let mut reader = chunked_buffer.reader(9900);
+    assert!(reader.read_exact(&mut resultdata[9900..10001]).is_err());
+
+    let mut resultdata = [0u8; 10000];
+    let mut reader = chunked_buffer.reader(20000);
+    assert_eq!(reader.read(&mut resultdata)?, 0);
+
+    Ok(())
+}

zenith_utils/src/lib.rs

@@ -11,6 +11,9 @@ pub mod seqwait;
 /// append only ordered map implemented with a Vec
 pub mod vec_map;
 
+/// expandable byte buffer consisting of fixed-size chunks
+pub mod chunked_buffer;
+
 // Async version of SeqWait. Currently unused.
 // pub mod seqwait_async;