From 35df89dcfe7e8125802bde5de648bd34c5f23cf6 Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Mon, 5 Aug 2024 17:57:41 +0000
Subject: [PATCH] update ingest_bench

---
 pageserver/benches/bench_ingest.rs            | 20 ++++++++++++++++---
 pageserver/src/tenant/storage_layer.rs        |  1 +
 .../tenant/storage_layer/inmemory_layer.rs    |  7 +++----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/pageserver/benches/bench_ingest.rs b/pageserver/benches/bench_ingest.rs
index af2b6934c6..bd0a4be394 100644
--- a/pageserver/benches/bench_ingest.rs
+++ b/pageserver/benches/bench_ingest.rs
@@ -10,7 +10,7 @@ use pageserver::{
     page_cache,
     repository::Value,
     task_mgr::TaskKind,
-    tenant::storage_layer::InMemoryLayer,
+    tenant::storage_layer::{InMemoryLayer, SerializedBatch},
     virtual_file::{self, api::IoEngineKind},
 };
 use pageserver_api::{key::Key, shard::TenantShardId};
@@ -63,12 +63,15 @@ async fn ingest(
 
     let layer = InMemoryLayer::create(conf, timeline_id, tenant_shard_id, lsn, &ctx).await?;
 
-    let data = Value::Image(Bytes::from(vec![0u8; put_size])).ser()?;
+    let value = Value::Image(Bytes::from(vec![0u8; put_size]));
     let ctx = RequestContext::new(
         pageserver::task_mgr::TaskKind::WalReceiverConnectionHandler,
         pageserver::context::DownloadBehavior::Download,
     );
 
+    let batch_pages = 100;
+    let mut batch_values = vec![];
+
     for i in 0..put_count {
         lsn += put_size as u64;
 
@@ -91,7 +94,18 @@ async fn ingest(
             }
         }
 
-        layer.put_value(key, lsn, &data, &ctx).await?;
+        batch_values.push((key, lsn, value.clone()));
+
+        if batch_values.len() >= batch_pages {
+            let batch = SerializedBatch::from_values(vec![(key, lsn, value.clone())]);
+
+            layer.put_batch(batch, &ctx).await?;
+        }
+    }
+    if !batch_values.is_empty() {
+        let batch = SerializedBatch::from_values(vec![(key, lsn, value.clone())]);
+
+        layer.put_batch(batch, &ctx).await?;
     }
     layer.freeze(lsn + 1).await;
 
diff --git a/pageserver/src/tenant/storage_layer.rs b/pageserver/src/tenant/storage_layer.rs
index 59d3e1ce09..56824e1f98 100644
--- a/pageserver/src/tenant/storage_layer.rs
+++ b/pageserver/src/tenant/storage_layer.rs
@@ -29,6 +29,7 @@ use utils::lsn::Lsn;
 pub use delta_layer::{DeltaLayer, DeltaLayerWriter, ValueRef};
 pub use image_layer::{ImageLayer, ImageLayerWriter};
 pub use inmemory_layer::InMemoryLayer;
+pub use inmemory_layer::SerializedBatch;
 pub use layer_desc::{PersistentLayerDesc, PersistentLayerKey};
 pub use layer_name::{DeltaLayerName, ImageLayerName, LayerName};
 
diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs
index a3464b5f02..c6361698a9 100644
--- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs
+++ b/pageserver/src/tenant/storage_layer/inmemory_layer.rs
@@ -454,7 +454,7 @@ impl InMemoryLayer {
     }
 }
 
-pub(crate) struct SerializedBatch {
+pub struct SerializedBatch {
     /// Blobs serialized in EphemeralFile's native format, ready for passing to [`EphemeralFile::write_raw`].
     pub(crate) raw: Vec<u8>,
 
@@ -486,7 +486,7 @@ impl SerializedBatch {
         }
     }
 
-    pub(crate) fn from_values(batch: Vec<(Key, Lsn, Value)>) -> Self {
+    pub fn from_values(batch: Vec<(Key, Lsn, Value)>) -> Self {
         use std::io::Write;
 
         let mut offsets: Vec<(Key, Lsn, u64)> = Vec::new();
@@ -585,9 +585,8 @@ impl InMemoryLayer {
         })
     }
 
-
     // Write path.
-    pub(crate) async fn put_batch(
+    pub async fn put_batch(
         &self,
         serialized_batch: SerializedBatch,
         ctx: &RequestContext,