From 1745fe5c65fbcc9422bc267ef778193e564bfc0c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arpad=20M=C3=BCller?= <arpad@neon.tech>
Date: Mon, 27 Jan 2025 12:57:44 +0100
Subject: [PATCH] Add SafekeeperGeneration

---
 libs/utils/src/generation.rs      | 24 ++++++++++++++++++++++++
 storage_controller/src/service.rs |  8 ++++----
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/libs/utils/src/generation.rs b/libs/utils/src/generation.rs
index 44565ee6a2..255f3b0dc3 100644
--- a/libs/utils/src/generation.rs
+++ b/libs/utils/src/generation.rs
@@ -144,6 +144,30 @@ impl Debug for Generation {
     }
 }
 
+/// Like tenant generations, but for safekeepers
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+pub struct SafekeeperGeneration(u32);
+
+impl SafekeeperGeneration {
+    pub const fn new(v: u32) -> Self {
+        Self(v)
+    }
+
+    #[track_caller]
+    pub fn previous(&self) -> Option<Self> {
+        Some(Self(self.0.checked_sub(1)?))
+    }
+
+    #[track_caller]
+    pub fn next(&self) -> Self {
+        Self(self.0 + 1)
+    }
+
+    pub fn into_inner(self) -> u32 {
+        self.0
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/storage_controller/src/service.rs b/storage_controller/src/service.rs
index 98613ff641..127f9611dd 100644
--- a/storage_controller/src/service.rs
+++ b/storage_controller/src/service.rs
@@ -85,7 +85,7 @@ use utils::{
     backoff,
     completion::Barrier,
     failpoint_support,
-    generation::Generation,
+    generation::{Generation, SafekeeperGeneration},
     http::error::ApiError,
     id::{NodeId, TenantId, TimelineId},
     logging::SecretString,
@@ -3564,7 +3564,7 @@ impl Service {
         tenant_id: TenantId,
         timeline_info: &TimelineInfo,
         create_mode: models::TimelineCreateRequestMode,
-    ) -> Result<(u32, Vec<NodeId>), ApiError> {
+    ) -> Result<(SafekeeperGeneration, Vec<NodeId>), ApiError> {
         let timeline_id = timeline_info.timeline_id;
         let pg_version = timeline_info.pg_version;
         let start_lsn = match create_mode {
@@ -3613,7 +3613,7 @@ impl Service {
             &sk_persistences,
         )
         .await?;
-        Ok((0, sks))
+        Ok((SafekeeperGeneration::new(0), sks))
     }
 
     pub(crate) async fn tenant_timeline_create(
@@ -3646,7 +3646,7 @@ impl Service {
             let res = self
                 .tenant_timeline_create_safekeepers(tenant_id, &timeline_info, create_mode)
                 .await?;
-            (Some(res.0), Some(res.1))
+            (Some(res.0.into_inner()), Some(res.1))
         } else {
             (None, None)
         };