From 6a267cf9e383082680c424653b791d0ffae7a768 Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Sat, 20 Jul 2024 20:21:31 +0100
Subject: [PATCH] avoid some allocs

---
 .../tenant/storage_layer/inmemory_layer.rs    | 10 +++----
 pageserver/src/tenant/timeline.rs             | 26 +++++++++----------
 2 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/pageserver/src/tenant/storage_layer/inmemory_layer.rs b/pageserver/src/tenant/storage_layer/inmemory_layer.rs
index e7ccd19de8..53fbd8e09e 100644
--- a/pageserver/src/tenant/storage_layer/inmemory_layer.rs
+++ b/pageserver/src/tenant/storage_layer/inmemory_layer.rs
@@ -523,18 +523,16 @@ impl InMemoryLayer {
 
     pub(crate) async fn put_values(
         &self,
-        values: Vec<(Lsn, Key, smallvec::SmallVec<[u8; 256]>)>,
+        mut values: Vec<(Lsn, Key, smallvec::SmallVec<[u8; 256]>, u64)>,
         ctx: &RequestContext,
     ) -> Result<()> {
         let mut inner = self.inner.write().await;
         self.assert_writable();
-        let mut index_updates = Vec::with_capacity(values.len());
-        for (lsn, key, buf) in values {
-            let off = self.put_value_locked2(&mut inner, &buf, ctx).await?;
-            index_updates.push((lsn, key, off));
+        for (_lsn, _key, buf, off) in &mut values {
+            *off = self.put_value_locked2(&mut inner, &buf, ctx).await?;
         }
 
-        for (lsn, key, off) in index_updates {
+        for (lsn, key, _buf, off) in values.into_iter() {
             let vec_map = inner.index.entry(key).or_default();
             let old = vec_map.append_or_update_last(lsn, off).unwrap().0;
             debug_assert!(old.is_none());
diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index fc8a2d5824..264cb5c4ea 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -5998,20 +5998,18 @@ impl<'a> TimelineWriter<'a> {
         let last_lsn = batch.as_slice().last().unwrap().0;
         let mut total_serialized_size = 0;
 
-        let serialized = batch
-            .into_iter()
-            .map(|(l, (k, v))| {
-                // Avoid doing allocations for "small" values.
-                // In the regression test suite, the limit of 256 avoided allocations in 95% of cases:
-                // https://github.com/neondatabase/neon/pull/5056#discussion_r1301975061
-                let mut buf = smallvec::SmallVec::<[u8; 256]>::new();
-                v.ser_into(&mut buf)
-                    .expect("Serialization of Value is infallible");
-                let buf_size: u64 = buf.len().try_into().expect("oversized value buf");
-                total_serialized_size += buf_size;
-                (l, k, buf)
-            })
-            .collect::<Vec<_>>();
+        let mut serialized = Vec::with_capacity(batch.len());
+        for (l, (k, v)) in batch.into_iter() {
+            // Avoid doing allocations for "small" values.
+            // In the regression test suite, the limit of 256 avoided allocations in 95% of cases:
+            // https://github.com/neondatabase/neon/pull/5056#discussion_r1301975061
+            let mut buf = smallvec::SmallVec::<[u8; 256]>::new();
+            v.ser_into(&mut buf)
+                .expect("Serialization of Value is infallible");
+            let buf_size: u64 = buf.len().try_into().expect("oversized value buf");
+            total_serialized_size += buf_size;
+            serialized.push((l, k, buf, 0));
+        }
 
         let action = self.get_open_layer_action(first_lsn, total_serialized_size);
         let layer = self