Refactor InMemoryLayer put_batch code

pageserver/src/tenant/ephemeral_file.rs

@@ -80,6 +80,7 @@ impl EphemeralFile {
     }
 
     #[cfg(test)]
+    // This is a test helper: outside of tests, we are always written do via a pre-serialized batch.
     pub(crate) async fn write_blob(
         &mut self,
         srcbuf: &[u8],
@@ -87,17 +88,15 @@ impl EphemeralFile {
     ) -> Result<u64, io::Error> {
         let pos = self.rw.bytes_written();
 
-        // Write the length field
-        if srcbuf.len() < 0x80 {
-            // short one-byte length header
-            let len_buf = [srcbuf.len() as u8];
+        let mut len_bytes = std::io::Cursor::new(Vec::new());
+        crate::tenant::storage_layer::inmemory_layer::SerializedBatch::write_blob_length(
+            srcbuf.len(),
+            &mut len_bytes,
+        );
+        let len_bytes = len_bytes.into_inner();
 
-            self.rw.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.rw.write_all_borrowed(&len_buf, ctx).await?;
-        }
+        // Write the length field
+        self.rw.write_all_borrowed(&len_bytes, ctx).await?;
 
         // Write the payload
         self.rw.write_all_borrowed(srcbuf, ctx).await?;

pageserver/src/tenant/storage_layer/inmemory_layer.rs

@@ -416,6 +416,74 @@ impl InMemoryLayer {
     }
 }
 
+pub(crate) struct SerializedBatch {
+    /// Blobs serialized in EphemeralFile's native format, ready for passing to [`EphemeralFile::write_raw`].
+    pub(crate) raw: Vec<u8>,
+
+    /// Index of values in [`Self::raw`], using offsets relative to the start of the buffer.
+    pub(crate) offsets: Vec<(Key, Lsn, u64)>,
+
+    /// The highest LSN of any value in the batch
+    pub(crate) max_lsn: Lsn,
+}
+
+impl SerializedBatch {
+    /// Write a blob length in the internal format of the EphemeralFile
+    pub(crate) fn write_blob_length(len: usize, cursor: &mut std::io::Cursor<Vec<u8>>) {
+        use std::io::Write;
+
+        if len < 0x80 {
+            // short one-byte length header
+            let len_buf = [len as u8];
+
+            cursor
+                .write_all(&len_buf)
+                .expect("Writing to Vec is infallible");
+        } else {
+            let mut len_buf = u32::to_be_bytes(len as u32);
+            len_buf[0] |= 0x80;
+            cursor
+                .write_all(&len_buf)
+                .expect("Writing to Vec is infallible");
+        }
+    }
+
+    pub(crate) fn from_values(batch: Vec<(Key, Lsn, Value)>) -> Self {
+        use std::io::Write;
+
+        let mut offsets: Vec<(Key, Lsn, u64)> = Vec::new();
+        let mut cursor = std::io::Cursor::new(Vec::<u8>::new());
+        let mut max_lsn: Lsn = Lsn(0);
+        let mut value_buf = smallvec::SmallVec::<[u8; 256]>::new();
+        for (key, lsn, val) in batch {
+            let relative_off = cursor.position();
+
+            value_buf.clear();
+            val.ser_into(&mut value_buf)
+                .expect("Value serialization is infallible");
+            Self::write_blob_length(value_buf.len(), &mut cursor);
+
+            cursor
+                .write_all(&value_buf)
+                .expect("Writing to Vec is infallible");
+
+            // We can't write straight into the buffer, because the InMemoryLayer file format requires
+            // the size to come before the value.  However... we could probably calculate the size before
+            // actually serializing the value
+            //val.ser_into(&mut cursor)?;
+
+            offsets.push((key, lsn, relative_off));
+            max_lsn = std::cmp::max(max_lsn, lsn);
+        }
+
+        Self {
+            raw: cursor.into_inner(),
+            offsets,
+            max_lsn,
+        }
+    }
+}
+
 fn inmem_layer_display(mut f: impl Write, start_lsn: Lsn, end_lsn: Lsn) -> std::fmt::Result {
     write!(f, "inmem-{:016X}-{:016X}", start_lsn.0, end_lsn.0)
 }
@@ -483,8 +551,7 @@ impl InMemoryLayer {
     // Write path.
     pub(crate) async fn put_batch(
         &self,
-        buf: Vec<u8>,
-        values: Vec<(Key, Lsn, u64)>,
+        serialized_batch: SerializedBatch,
         ctx: &RequestContext,
     ) -> Result<()> {
         let mut inner = self.inner.write().await;
@@ -494,7 +561,7 @@ impl InMemoryLayer {
             inner
                 .file
                 .write_raw(
-                    &buf,
+                    &serialized_batch.raw,
                     &RequestContextBuilder::extend(ctx)
                         .page_content_kind(PageContentKind::InMemoryLayer)
                         .build(),
@@ -502,7 +569,7 @@ impl InMemoryLayer {
                 .await?
         };
 
-        for (key, lsn, relative_off) in values {
+        for (key, lsn, relative_off) in serialized_batch.offsets {
             let off = base_off + relative_off;
             let vec_map = inner.index.entry(key).or_default();
             let old = vec_map.append_or_update_last(lsn, off).unwrap().0;

pageserver/src/tenant/timeline.rs

@@ -49,6 +49,7 @@ use utils::{
     sync::gate::{Gate, GateGuard},
 };
 
+use std::pin::pin;
 use std::sync::atomic::Ordering as AtomicOrdering;
 use std::sync::{Arc, Mutex, RwLock, Weak};
 use std::time::{Duration, Instant, SystemTime};
@@ -65,7 +66,6 @@ use std::{
     collections::btree_map::Entry,
     ops::{Deref, Range},
 };
-use std::{io::Write, pin::pin};
 
 use crate::{
     aux_file::AuxFileSizeEstimator,
@@ -139,7 +139,7 @@ use self::layer_manager::LayerManager;
 use self::logical_size::LogicalSize;
 use self::walreceiver::{WalReceiver, WalReceiverConf};
 
-use super::{config::TenantConf, upload_queue::NotInitialized};
+use super::{config::TenantConf, storage_layer::inmemory_layer, upload_queue::NotInitialized};
 use super::{debug_assert_current_span_has_tenant_and_timeline_id, AttachedTenantConf};
 use super::{remote_timeline_client::index::IndexPart, storage_layer::LayerFringe};
 use super::{
@@ -6042,44 +6042,16 @@ impl<'a> TimelineWriter<'a> {
             return Ok(());
         }
 
-        let mut values: Vec<(Key, Lsn, u64)> = Vec::new();
-        let mut cursor = std::io::Cursor::new(Vec::<u8>::new());
-        let mut batch_max_lsn: Lsn = Lsn(0);
-        let mut value_buf = smallvec::SmallVec::<[u8; 256]>::new();
-        for (key, lsn, val) in batch {
-            let relative_off = cursor.position();
-
-            value_buf.clear();
-            val.ser_into(&mut value_buf)?;
-            if value_buf.len() < 0x80 {
-                // short one-byte length header
-                let len_buf = [value_buf.len() as u8];
-
-                cursor.write_all(&len_buf)?;
-            } else {
-                let mut len_buf = u32::to_be_bytes(value_buf.len() as u32);
-                len_buf[0] |= 0x80;
-                cursor.write_all(&len_buf)?;
-            }
-            cursor.write_all(&value_buf)?;
-
-            // We can't write straight into the buffer, because the InMemoryLayer file format requires
-            // the size to come before the value.  However... we could probably calculate the size before
-            // actually serializing the value
-            //val.ser_into(&mut cursor)?;
-
-            values.push((key, lsn, relative_off));
-            batch_max_lsn = std::cmp::max(batch_max_lsn, lsn);
-        }
-
-        let buf = cursor.into_inner();
-        let buf_size: u64 = buf.len() as u64;
+        let serialized_batch = inmemory_layer::SerializedBatch::from_values(batch);
+        let batch_max_lsn = serialized_batch.max_lsn;
+        let buf_size: u64 = serialized_batch.raw.len() as u64;
 
         let action = self.get_open_layer_action(batch_max_lsn, buf_size);
         let layer = self
             .handle_open_layer_action(batch_max_lsn, action, ctx)
             .await?;
-        let res = layer.put_batch(buf, values, ctx).await;
+
+        let res = layer.put_batch(serialized_batch, ctx).await;
 
         if res.is_ok() {
             // Update the current size only when the entire write was ok.