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

@@ -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>,
 

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}"));
         }
 

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.

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>,

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();