From 47553dbaf946de9cdab759c220ec56dffa5f82cd Mon Sep 17 00:00:00 2001
From: Dmitrii Kovalkov <34828390+DimasKovas@users.noreply.github.com>
Date: Sat, 28 Jun 2025 16:59:29 +0400
Subject: [PATCH] neon_local: set timeline_safekeeper_count if we have less
 than 3 safekeepers (#12378)

## Problem
- Closes: https://github.com/neondatabase/neon/issues/12298

## Summary of changes
- Set `timeline_safekeeper_count` in `neon_local` if we have less than 3
safekeepers
- Remove `cfg!(feature = "testing")` code from
`safekeepers_for_new_timeline`
- Change `timeline_safekeeper_count` type to `usize`
---
 control_plane/src/local_env.rs                       | 2 +-
 control_plane/src/storage_controller.rs              | 8 +++++++-
 storage_controller/src/main.rs                       | 4 ++--
 storage_controller/src/service.rs                    | 2 +-
 storage_controller/src/service/safekeeper_service.rs | 9 +--------
 5 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs
index 16cd2d8c08..d0611113e8 100644
--- a/control_plane/src/local_env.rs
+++ b/control_plane/src/local_env.rs
@@ -212,7 +212,7 @@ pub struct NeonStorageControllerConf {
 
     pub use_local_compute_notifications: bool,
 
-    pub timeline_safekeeper_count: Option<i64>,
+    pub timeline_safekeeper_count: Option<usize>,
 
     pub posthog_config: Option<PostHogConfig>,
 
diff --git a/control_plane/src/storage_controller.rs b/control_plane/src/storage_controller.rs
index dea7ae2ccf..bb83a6319c 100644
--- a/control_plane/src/storage_controller.rs
+++ b/control_plane/src/storage_controller.rs
@@ -638,7 +638,13 @@ impl StorageController {
             args.push("--timelines-onto-safekeepers".to_string());
         }
 
-        if let Some(sk_cnt) = self.config.timeline_safekeeper_count {
+        // neon_local is used in test environments where we often have less than 3 safekeepers.
+        if self.config.timeline_safekeeper_count.is_some() || self.env.safekeepers.len() < 3 {
+            let sk_cnt = self
+                .config
+                .timeline_safekeeper_count
+                .unwrap_or(self.env.safekeepers.len());
+
             args.push(format!("--timeline-safekeeper-count={sk_cnt}"));
         }
 
diff --git a/storage_controller/src/main.rs b/storage_controller/src/main.rs
index d1c2858d6f..752262b65e 100644
--- a/storage_controller/src/main.rs
+++ b/storage_controller/src/main.rs
@@ -216,8 +216,8 @@ struct Cli {
     /// Number of safekeepers to choose for a timeline when creating it.
     /// Safekeepers will be choosen from different availability zones.
     /// This option exists primarily for testing purposes.
-    #[arg(long, default_value = "3", value_parser = clap::value_parser!(i64).range(1..))]
-    timeline_safekeeper_count: i64,
+    #[arg(long, default_value = "3", value_parser = clap::builder::RangedU64ValueParser::<usize>::new().range(1..))]
+    timeline_safekeeper_count: usize,
 
     /// When set, actively checks and initiates heatmap downloads/uploads during reconciliation.
     /// This speed up migrations by avoiding the default wait for the heatmap download interval.
diff --git a/storage_controller/src/service.rs b/storage_controller/src/service.rs
index b4dfd01249..19bb0f8671 100644
--- a/storage_controller/src/service.rs
+++ b/storage_controller/src/service.rs
@@ -471,7 +471,7 @@ pub struct Config {
 
     /// Number of safekeepers to choose for a timeline when creating it.
     /// Safekeepers will be choosen from different availability zones.
-    pub timeline_safekeeper_count: i64,
+    pub timeline_safekeeper_count: usize,
 
     /// PostHog integration config
     pub posthog_config: Option<PostHogConfig>,
diff --git a/storage_controller/src/service/safekeeper_service.rs b/storage_controller/src/service/safekeeper_service.rs
index fec81fb661..92d15f3fca 100644
--- a/storage_controller/src/service/safekeeper_service.rs
+++ b/storage_controller/src/service/safekeeper_service.rs
@@ -1,4 +1,3 @@
-use std::cmp::max;
 use std::collections::HashSet;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -654,13 +653,7 @@ impl Service {
             )
         });
         // Number of safekeepers in different AZs we are looking for
-        let mut wanted_count = self.config.timeline_safekeeper_count as usize;
-        // TODO(diko): remove this when `timeline_safekeeper_count` option is in the release
-        // branch and is specified in tests/neon_local config.
-        if cfg!(feature = "testing") && all_safekeepers.len() < wanted_count {
-            // In testing mode, we can have less safekeepers than the config says
-            wanted_count = max(all_safekeepers.len(), 1);
-        }
+        let wanted_count = self.config.timeline_safekeeper_count;
 
         let mut sks = Vec::new();
         let mut azs = HashSet::new();