From 93dacd75aa5ae7e3d3a79c13834dd388f0940c1c Mon Sep 17 00:00:00 2001 From: John Spray Date: Mon, 5 Aug 2024 15:10:04 +0000 Subject: [PATCH] Refactor InMemoryLayer put_batch code --- pageserver/src/tenant/ephemeral_file.rs | 19 +++-- .../tenant/storage_layer/inmemory_layer.rs | 75 ++++++++++++++++++- pageserver/src/tenant/timeline.rs | 42 ++--------- 3 files changed, 87 insertions(+), 49 deletions(-) diff --git a/pageserver/src/tenant/ephemeral_file.rs b/pageserver/src/tenant/ephemeral_file.rs index 20c8493d7d..1e6215d23d 100644 --- a/pageserver/src/tenant/ephemeral_file.rs +++ b/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 { 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?; diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs index 5d2213c260..1d8a58a323 100644 --- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs +++ b/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, + + /// 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>) { + 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::::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, - 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; diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index bbf320414f..3cec869c77 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/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::::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.