From 849c05b0fcb2ad02c313db5e62e9224f93739b37 Mon Sep 17 00:00:00 2001
From: Alex Chi Z <chi@neon.tech>
Date: Fri, 17 May 2024 14:16:49 -0400
Subject: [PATCH] fix test case

Signed-off-by: Alex Chi Z <chi@neon.tech>
---
 pageserver/src/tenant.rs          | 55 +++++++++++++++++++------------
 pageserver/src/tenant/timeline.rs |  5 ++-
 2 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs
index 61761e86dce5..0c1d94d36d10 100644
--- a/pageserver/src/tenant.rs
+++ b/pageserver/src/tenant.rs
@@ -5732,23 +5732,11 @@ mod tests {
             lsn: Lsn,
             ctx: &RequestContext,
         ) -> anyhow::Result<(BTreeMap<Key, Result<Bytes, PageReconstructError>>, usize)> {
-            let begin_files_accessed = ctx
-                .vectored_access_delta_file_cnt
-                .load(std::sync::atomic::Ordering::SeqCst);
+            let mut reconstruct_state = ValuesReconstructState::default();
             let res = tline
-                .get_vectored_impl(
-                    keyspace.clone(),
-                    lsn,
-                    ValuesReconstructState::default(),
-                    ctx,
-                )
+                .get_vectored_impl(keyspace.clone(), lsn, &mut reconstruct_state, ctx)
                 .await?;
-            Ok((
-                res,
-                ctx.vectored_access_delta_file_cnt
-                    .load(std::sync::atomic::Ordering::SeqCst)
-                    - begin_files_accessed,
-            ))
+            Ok((res, reconstruct_state.get_delta_layers_visited() as usize))
         }
 
         #[allow(clippy::needless_range_loop)]
@@ -5827,6 +5815,8 @@ mod tests {
             .create_test_timeline(TIMELINE_ID, Lsn(0x10), DEFAULT_PG_VERSION, &ctx)
             .await?;
 
+        let cancel = CancellationToken::new();
+
         let base_key = Key::from_hex("000000000033333333444444445500000000").unwrap();
         let base_key_child = Key::from_hex("000000000033333333444444445500000001").unwrap();
         let base_key_nonexist = Key::from_hex("000000000033333333444444445500000002").unwrap();
@@ -5841,7 +5831,7 @@ mod tests {
             writer.finish_write(lsn);
             drop(writer);
 
-            tline.freeze_and_flush().await?; // this will create an image layer
+            tline.freeze_and_flush().await?; // this will create a image layer
         }
 
         let child = tenant
@@ -5864,7 +5854,30 @@ mod tests {
             writer.finish_write(lsn);
             drop(writer);
 
-            child.freeze_and_flush().await?; // this will create an image layer
+            child.freeze_and_flush().await?; // this will create a delta
+
+            {
+                // update the partitioning to include the test key space, otherwise they
+                // will be dropped by image layer creation
+                let mut guard = child.partitioning.lock().await;
+                let ((partitioning, _), partition_lsn) = &mut *guard;
+                partitioning
+                    .parts
+                    .push(KeySpace::single(base_key..base_key_nonexist)); // exclude the nonexist key
+                *partition_lsn = lsn;
+            }
+
+            child
+                .compact(
+                    &cancel,
+                    {
+                        let mut set = EnumSet::empty();
+                        set.insert(CompactFlags::ForceImageLayerCreation);
+                        set
+                    },
+                    &ctx,
+                )
+                .await?; // force create an image layer for the keys
         }
 
         async fn get_vectored_impl_wrapper(
@@ -5873,12 +5886,12 @@ mod tests {
             lsn: Lsn,
             ctx: &RequestContext,
         ) -> Result<Option<Bytes>, GetVectoredError> {
-            let reconstruct_state = ValuesReconstructState::new();
+            let mut reconstruct_state = ValuesReconstructState::new();
             let mut res = tline
                 .get_vectored_impl(
                     KeySpace::single(key..key.next()),
                     lsn,
-                    reconstruct_state,
+                    &mut reconstruct_state,
                     ctx,
                 )
                 .await?;
@@ -6019,12 +6032,12 @@ mod tests {
             lsn: Lsn,
             ctx: &RequestContext,
         ) -> Result<Option<Bytes>, GetVectoredError> {
-            let reconstruct_state = ValuesReconstructState::new();
+            let mut reconstruct_state = ValuesReconstructState::new();
             let mut res = tline
                 .get_vectored_impl(
                     KeySpace::single(key..key.next()),
                     lsn,
-                    reconstruct_state,
+                    &mut reconstruct_state,
                     ctx,
                 )
                 .await?;
diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index c9ef6e8fd5ba..8d27c21219b5 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -349,7 +349,7 @@ pub struct Timeline {
     pub initdb_lsn: Lsn,
 
     /// When did we last calculate the partitioning?
-    partitioning: tokio::sync::Mutex<((KeyPartitioning, SparseKeyPartitioning), Lsn)>,
+    pub(crate) partitioning: tokio::sync::Mutex<((KeyPartitioning, SparseKeyPartitioning), Lsn)>,
 
     /// Configuration: how often should the partitioning be recalculated.
     repartition_threshold: u64,
@@ -4161,6 +4161,8 @@ impl Timeline {
         }; // no partitioning for metadata keys for now
         *partitioning_guard = ((dense_partitioning, sparse_partitioning), lsn);
 
+        info!("repartitioning!");
+
         Ok((partitioning_guard.0.clone(), partitioning_guard.1))
     }
 
@@ -4230,6 +4232,7 @@ impl Timeline {
 
         let mut key_request_accum = KeySpaceAccum::new();
         for range in &partition.ranges {
+            info!("partition: {}..{}", range.start, range.end);
             let mut key = range.start;
             while key < range.end {
                 // Decide whether to retain this key: usually we do, but sharded tenants may