From c78c810118dcba4311e9669b77215be936aa4bcf Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Sat, 29 Jun 2024 21:51:10 +0100
Subject: [PATCH] pageserver: move out common compaction step

---
 pageserver/src/tenant/timeline.rs            | 19 ++++++++++++++++---
 pageserver/src/tenant/timeline/compaction.rs | 17 +++--------------
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index 257a564f3b..3312d05ea6 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -30,7 +30,7 @@ use pageserver_api::{
         InMemoryLayerInfo, LayerMapInfo, LsnLease, TimelineState,
     },
     reltag::BlockNumber,
-    shard::{ShardIdentity, ShardNumber, TenantShardId},
+    shard::{ShardCount, ShardIdentity, ShardNumber, TenantShardId},
 };
 use rand::Rng;
 use serde_with::serde_as;
@@ -1802,9 +1802,22 @@ impl Timeline {
         }
 
         match self.get_compaction_algorithm_settings().kind {
-            CompactionAlgorithm::Tiered => self.compact_tiered(cancel, ctx).await,
-            CompactionAlgorithm::Legacy => self.compact_legacy(cancel, flags, ctx).await,
+            CompactionAlgorithm::Tiered => self.compact_tiered(cancel, ctx).await?,
+            CompactionAlgorithm::Legacy => self.compact_legacy(cancel, flags, ctx).await?,
         }
+
+        if self.shard_identity.count >= ShardCount::new(2) {
+            // Limit the number of layer rewrites to the number of partitions: this means its
+            // runtime should be comparable to a full round of image layer creations, rather than
+            // being potentially much longer.
+            // TODO: make `partitioning` a sync lock: see comment in `repartition()` for why there's no
+            // real async use.
+            let rewrite_max = self.partitioning.try_lock().unwrap().0 .0.parts.len();
+
+            self.compact_shard_ancestors(rewrite_max, ctx).await?;
+        }
+
+        Ok(())
     }
 
     /// Mutate the timeline with a [`TimelineWriter`].
diff --git a/pageserver/src/tenant/timeline/compaction.rs b/pageserver/src/tenant/timeline/compaction.rs
index efaa6144af..f8dc87787a 100644
--- a/pageserver/src/tenant/timeline/compaction.rs
+++ b/pageserver/src/tenant/timeline/compaction.rs
@@ -19,7 +19,7 @@ use enumset::EnumSet;
 use fail::fail_point;
 use itertools::Itertools;
 use pageserver_api::keyspace::ShardedRange;
-use pageserver_api::shard::{ShardCount, ShardIdentity, TenantShardId};
+use pageserver_api::shard::{ShardIdentity, TenantShardId};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, info, info_span, trace, warn, Instrument};
 use utils::id::TimelineId;
@@ -100,7 +100,7 @@ impl Timeline {
         // Define partitioning schema if needed
 
         // FIXME: the match should only cover repartitioning, not the next steps
-        let partition_count = match self
+        match self
             .repartition(
                 self.get_last_record_lsn(),
                 self.get_compaction_target_size(),
@@ -140,7 +140,6 @@ impl Timeline {
                     .await?;
 
                 self.upload_new_image_layers(image_layers)?;
-                partitioning.parts.len()
             }
             Err(err) => {
                 // no partitioning? This is normal, if the timeline was just created
@@ -152,19 +151,9 @@ impl Timeline {
                 if !self.cancel.is_cancelled() {
                     tracing::error!("could not compact, repartitioning keyspace failed: {err:?}");
                 }
-                1
             }
         };
 
-        if self.shard_identity.count >= ShardCount::new(2) {
-            // Limit the number of layer rewrites to the number of partitions: this means its
-            // runtime should be comparable to a full round of image layer creations, rather than
-            // being potentially much longer.
-            let rewrite_max = partition_count;
-
-            self.compact_shard_ancestors(rewrite_max, ctx).await?;
-        }
-
         Ok(())
     }
 
@@ -176,7 +165,7 @@ impl Timeline {
     ///
     /// Note: this phase may read and write many gigabytes of data: use rewrite_max to bound
     /// how much work it will try to do in each compaction pass.
-    async fn compact_shard_ancestors(
+    pub(super) async fn compact_shard_ancestors(
         self: &Arc<Self>,
         rewrite_max: usize,
         ctx: &RequestContext,