impr(proxy): introduce proxy_id for cancel key

## Problem

Since #9916 , the chaos code is actively fighting the optimizer: tenants
tend to be attached in their preferred AZ, so most chaos migrations were
moving them to a non-preferred AZ.

## Summary of changes

- When picking migrations, prefer to migrate things _toward_ their
preferred AZ when possible. Then pick shards to move the other way when
necessary.

The resulting behavior should be an alternating "back and forth" where
the chaos code migrates thiings away from home, and then migrates them
back on the next iteration.

The side effect will be that the chaos code actively helps to push
things into their home AZ. That's not contrary to its purpose though: we
mainly just want it to continuously migrate things to exercise
migration+notification code.

.github/workflows/build_and_test.yml

@@ -824,7 +824,7 @@ jobs:
           docker compose -f ./docker-compose/docker-compose.yml down
 
   promote-images-dev:
-    needs: [ check-permissions, tag, vm-compute-node-image ]
+    needs: [ check-permissions, tag, vm-compute-node-image, neon-image ]
     runs-on: ubuntu-22.04
 
     permissions:

libs/safekeeper_api/src/models.rs

@@ -277,3 +277,8 @@ pub struct TimelineTermBumpResponse {
     pub previous_term: u64,
     pub current_term: u64,
 }
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct SafekeeperUtilization {
+    pub timeline_count: u64,
+}

pageserver/src/metrics.rs

@@ -100,6 +100,32 @@ pub(crate) static VEC_READ_NUM_LAYERS_VISITED: Lazy<Histogram> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+pub(crate) static CONCURRENT_INITDBS: Lazy<UIntGauge> = Lazy::new(|| {
+    register_uint_gauge!(
+        "pageserver_concurrent_initdb",
+        "Number of initdb processes running"
+    )
+    .expect("failed to define a metric")
+});
+
+pub(crate) static INITDB_SEMAPHORE_ACQUISITION_TIME: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "pageserver_initdb_semaphore_seconds_global",
+        "Time spent getting a permit from the global initdb semaphore",
+        STORAGE_OP_BUCKETS.into()
+    )
+    .expect("failed to define metric")
+});
+
+pub(crate) static INITDB_RUN_TIME: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "pageserver_initdb_seconds_global",
+        "Time spent performing initdb",
+        STORAGE_OP_BUCKETS.into()
+    )
+    .expect("failed to define metric")
+});
+
 // Metrics collected on operations on the storage repository.
 #[derive(
     Clone, Copy, enum_map::Enum, strum_macros::EnumString, strum_macros::Display, IntoStaticStr,

pageserver/src/tenant.rs

@@ -95,6 +95,9 @@ use crate::deletion_queue::DeletionQueueError;
 use crate::import_datadir;
 use crate::is_uninit_mark;
 use crate::l0_flush::L0FlushGlobalState;
+use crate::metrics::CONCURRENT_INITDBS;
+use crate::metrics::INITDB_RUN_TIME;
+use crate::metrics::INITDB_SEMAPHORE_ACQUISITION_TIME;
 use crate::metrics::TENANT;
 use crate::metrics::{
     remove_tenant_metrics, BROKEN_TENANTS_SET, CIRCUIT_BREAKERS_BROKEN, CIRCUIT_BREAKERS_UNBROKEN,
@@ -5347,8 +5350,17 @@ async fn run_initdb(
         initdb_bin_path, initdb_target_dir, initdb_lib_dir,
     );
 
-    let _permit = INIT_DB_SEMAPHORE.acquire().await;
+    let _permit = {
+        let _timer = INITDB_SEMAPHORE_ACQUISITION_TIME.start_timer();
+        INIT_DB_SEMAPHORE.acquire().await
+    };
 
+    CONCURRENT_INITDBS.inc();
+    scopeguard::defer! {
+        CONCURRENT_INITDBS.dec();
+    }
+
+    let _timer = INITDB_RUN_TIME.start_timer();
     let res = postgres_initdb::do_run_initdb(postgres_initdb::RunInitdbArgs {
         superuser: &conf.superuser,
         locale: &conf.locale,

pageserver/src/tenant/remote_timeline_client.rs

@@ -803,6 +803,12 @@ impl RemoteTimelineClient {
 
         upload_queue.dirty.metadata.apply(update);
 
+        // Defense in depth: if we somehow generated invalid metadata, do not persist it.
+        upload_queue
+            .dirty
+            .validate()
+            .map_err(|e| anyhow::anyhow!(e))?;
+
         self.schedule_index_upload(upload_queue);
 
         Ok(())

pageserver/src/tenant/remote_timeline_client/index.rs

@@ -152,6 +152,21 @@ impl IndexPart {
         };
         is_same_remote_layer_path(name, metadata, name, index_metadata)
     }
+
+    /// Check for invariants in the index: this is useful when uploading an index to ensure that if
+    /// we encounter a bug, we do not persist buggy metadata.
+    pub(crate) fn validate(&self) -> Result<(), String> {
+        if self.import_pgdata.is_none()
+            && self.metadata.ancestor_timeline().is_none()
+            && self.layer_metadata.is_empty()
+        {
+            // Unless we're in the middle of a raw pgdata import, or this is a child timeline,the index must
+            // always have at least one layer.
+            return Err("Index has no ancestor and no layers".to_string());
+        }
+
+        Ok(())
+    }
 }
 
 /// Metadata gathered for each of the layer files.

pageserver/src/tenant/remote_timeline_client/upload.rs

@@ -40,6 +40,10 @@ pub(crate) async fn upload_index_part(
     });
     pausable_failpoint!("before-upload-index-pausable");
 
+    // Safety: refuse to persist invalid index metadata, to mitigate the impact of any bug that produces this
+    // (this should never happen)
+    index_part.validate().map_err(|e| anyhow::anyhow!(e))?;
+
     // FIXME: this error comes too late
     let serialized = index_part.to_json_bytes()?;
     let serialized = Bytes::from(serialized);

pageserver/src/tenant/storage_layer/layer/tests.rs

@@ -1,6 +1,6 @@
 use std::time::UNIX_EPOCH;
 
-use pageserver_api::key::CONTROLFILE_KEY;
+use pageserver_api::key::{Key, CONTROLFILE_KEY};
 use tokio::task::JoinSet;
 use utils::{
     completion::{self, Completion},
@@ -9,7 +9,10 @@ use utils::{
 
 use super::failpoints::{Failpoint, FailpointKind};
 use super::*;
-use crate::{context::DownloadBehavior, tenant::storage_layer::LayerVisibilityHint};
+use crate::{
+    context::DownloadBehavior,
+    tenant::{harness::test_img, storage_layer::LayerVisibilityHint},
+};
 use crate::{task_mgr::TaskKind, tenant::harness::TenantHarness};
 
 /// Used in tests to advance a future to wanted await point, and not futher.
@@ -31,20 +34,51 @@ async fn smoke_test() {
 
     let ctx = RequestContext::new(TaskKind::UnitTest, DownloadBehavior::Download);
 
+    let image_layers = vec![(
+        Lsn(0x40),
+        vec![(
+            Key::from_hex("620000000033333333444444445500000000").unwrap(),
+            test_img("foo"),
+        )],
+    )];
+
+    // Create a test timeline with one real layer, and one synthetic test layer.  The synthetic
+    // one is only there so that we can GC the real one without leaving the timeline's metadata
+    // empty, which is an illegal state (see [`IndexPart::validate`]).
     let timeline = tenant
-        .create_test_timeline(TimelineId::generate(), Lsn(0x10), 14, &ctx)
+        .create_test_timeline_with_layers(
+            TimelineId::generate(),
+            Lsn(0x10),
+            14,
+            &ctx,
+            Default::default(),
+            image_layers,
+            Lsn(0x100),
+        )
         .await
         .unwrap();
 
-    let layer = {
+    // Grab one of the timeline's layers to exercise in the test, and the other layer that is just
+    // there to avoid the timeline being illegally empty
+    let (layer, dummy_layer) = {
         let mut layers = {
             let layers = timeline.layers.read().await;
             layers.likely_resident_layers().cloned().collect::<Vec<_>>()
         };
 
-        assert_eq!(layers.len(), 1);
+        assert_eq!(layers.len(), 2);
 
-        layers.swap_remove(0)
+        layers.sort_by_key(|l| l.layer_desc().get_key_range().start);
+        let synthetic_layer = layers.pop().unwrap();
+        let real_layer = layers.pop().unwrap();
+        tracing::info!(
+            "real_layer={:?} ({}), synthetic_layer={:?} ({})",
+            real_layer,
+            real_layer.layer_desc().file_size,
+            synthetic_layer,
+            synthetic_layer.layer_desc().file_size
+        );
+        (real_layer, synthetic_layer)
     };
 
     // all layers created at pageserver are like `layer`, initialized with strong
@@ -173,10 +207,13 @@ async fn smoke_test() {
 
     let rtc = &timeline.remote_client;
 
+    // Simulate GC removing our test layer.
     {
-        let layers = &[layer];
         let mut g = timeline.layers.write().await;
+
+        let layers = &[layer];
         g.open_mut().unwrap().finish_gc_timeline(layers);
+
         // this just updates the remote_physical_size for demonstration purposes
         rtc.schedule_gc_update(layers).unwrap();
     }
@@ -191,7 +228,10 @@ async fn smoke_test() {
 
     rtc.wait_completion().await.unwrap();
 
-    assert_eq!(rtc.get_remote_physical_size(), 0);
+    assert_eq!(
+        rtc.get_remote_physical_size(),
+        dummy_layer.metadata().file_size
+    );
     assert_eq!(0, LAYER_IMPL_METRICS.inits_cancelled.get())
 }
 

pageserver/src/tenant/timeline.rs

@@ -5678,9 +5678,17 @@ impl Timeline {
         info!("force created image layer {}", image_layer.local_path());
         {
             let mut guard = self.layers.write().await;
-            guard.open_mut().unwrap().force_insert_layer(image_layer);
+            guard
+                .open_mut()
+                .unwrap()
+                .force_insert_layer(image_layer.clone());
         }
 
+        // Update remote_timeline_client state to reflect existence of this layer
+        self.remote_client
+            .schedule_layer_file_upload(image_layer)
+            .unwrap();
+
         Ok(())
     }
 
@@ -5731,9 +5739,17 @@ impl Timeline {
         info!("force created delta layer {}", delta_layer.local_path());
         {
             let mut guard = self.layers.write().await;
-            guard.open_mut().unwrap().force_insert_layer(delta_layer);
+            guard
+                .open_mut()
+                .unwrap()
+                .force_insert_layer(delta_layer.clone());
         }
 
+        // Update remote_timeline_client state to reflect existence of this layer
+        self.remote_client
+            .schedule_layer_file_upload(delta_layer)
+            .unwrap();
+
         Ok(())
     }
 

proxy/src/bin/local_proxy.rs

@@ -14,7 +14,7 @@ use proxy::auth::backend::local::{LocalBackend, JWKS_ROLE_MAP};
 use proxy::auth::{self};
 use proxy::cancellation::CancellationHandlerMain;
 use proxy::config::{
-    self, AuthenticationConfig, ComputeConfig, HttpConfig, ProxyConfig, RetryConfig,
+    self, obfuscated_proxy_id, AuthenticationConfig, ComputeConfig, HttpConfig, ProxyConfig, RetryConfig
 };
 use proxy::control_plane::locks::ApiLocks;
 use proxy::control_plane::messages::{EndpointJwksResponse, JwksSettings};
@@ -218,6 +218,7 @@ async fn main() -> anyhow::Result<()> {
             proxy::metrics::CancellationSource::Local,
         )),
         endpoint_rate_limiter,
+        obfuscated_proxy_id(std::process::id(), "local"),
     );
 
     match futures::future::select(pin!(maintenance_tasks.join_next()), pin!(task)).await {

proxy/src/bin/pg_sni_router.rs

@@ -10,6 +10,7 @@ use clap::Arg;
 use futures::future::Either;
 use futures::TryFutureExt;
 use itertools::Itertools;
+use proxy::config::obfuscated_proxy_id;
 use proxy::context::RequestContext;
 use proxy::metrics::{Metrics, ThreadPoolMetrics};
 use proxy::protocol2::ConnectionInfo;
@@ -185,6 +186,7 @@ async fn task_main(
                     },
                     proxy::metrics::Protocol::SniRouter,
                     "sni",
+                    obfuscated_proxy_id(std::process::id(), "sni-router"), // just a shim for context
                 );
                 handle_client(ctx, dest_suffix, tls_config, tls_server_end_point, socket).await
             }

proxy/src/bin/proxy.rs

@@ -10,7 +10,7 @@ use proxy::auth::backend::{AuthRateLimiter, ConsoleRedirectBackend, MaybeOwned};
 use proxy::cancellation::{CancelMap, CancellationHandler};
 use proxy::config::{
     self, remote_storage_from_toml, AuthenticationConfig, CacheOptions, ComputeConfig, HttpConfig,
-    ProjectInfoCacheOptions, ProxyConfig, ProxyProtocolV2,
+    ProjectInfoCacheOptions, ProxyConfig, ProxyProtocolV2, obfuscated_proxy_id,
 };
 use proxy::context::parquet::ParquetUploadArgs;
 use proxy::http::health_server::AppMetrics;
@@ -396,6 +396,8 @@ async fn main() -> anyhow::Result<()> {
         None => None,
     };
 
+    let proxy_id: u16 = obfuscated_proxy_id(std::process::id(), &args.region);
+
     let cancellation_handler = Arc::new(CancellationHandler::<
         Option<Arc<Mutex<RedisPublisherClient>>>,
     >::new(
@@ -437,6 +439,7 @@ async fn main() -> anyhow::Result<()> {
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
                     endpoint_rate_limiter.clone(),
+                    proxy_id,
                 ));
             }
 
@@ -448,6 +451,7 @@ async fn main() -> anyhow::Result<()> {
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
                     endpoint_rate_limiter.clone(),
+                    proxy_id,
                 ));
             }
         }
@@ -459,6 +463,7 @@ async fn main() -> anyhow::Result<()> {
                     proxy_listener,
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
+                    proxy_id,
                 ));
             }
         }

proxy/src/cancellation.rs

@@ -80,15 +80,24 @@ impl ReportableError for CancelError {
 
 impl<P: CancellationPublisher> CancellationHandler<P> {
     /// Run async action within an ephemeral session identified by [`CancelKeyData`].
-    pub(crate) fn get_session(self: Arc<Self>) -> Session<P> {
+    pub(crate) fn get_session(self: Arc<Self>, proxy_id: u16) -> Session<P> {
         // we intentionally generate a random "backend pid" and "secret key" here.
         // we use the corresponding u64 as an identifier for the
         // actual endpoint+pid+secret for postgres/pgbouncer.
         //
         // if we forwarded the backend_pid from postgres to the client, there would be a lot
         // of overlap between our computes as most pids are small (~100).
+
         let key = loop {
-            let key = rand::random();
+            let key_rand: u64 = rand::random::<u64>() & 0x0000_ffff_ffff_ffff;
+
+            let backend_pid = ((proxy_id as u32) << 16) | ((key_rand >> 32) as u32) as u32;
+            let cancel_key = (key_rand as u32) as i32;
+
+            let key = CancelKeyData {
+                backend_pid: (backend_pid as i32),
+                cancel_key,
+            };
 
             // Random key collisions are unlikely to happen here, but they're still possible,
             // which is why we have to take care not to rewrite an existing key.
@@ -451,7 +460,7 @@ mod tests {
             CancellationSource::FromRedis,
         ));
 
-        let session = cancellation_handler.clone().get_session();
+        let session = cancellation_handler.clone().get_session(123);
         assert!(cancellation_handler.contains(&session));
         drop(session);
         // Check that the session has been dropped.

proxy/src/config.rs

@@ -16,6 +16,8 @@ use crate::serverless::GlobalConnPoolOptions;
 pub use crate::tls::server_config::{configure_tls, TlsConfig};
 use crate::types::Host;
 
+use sha2::{Digest, Sha256};
+
 pub struct ProxyConfig {
     pub tls_config: Option<TlsConfig>,
     pub metric_collection: Option<MetricCollectionConfig>,
@@ -416,6 +418,23 @@ impl FromStr for ConcurrencyLockOptions {
     }
 }
 
+fn map_u32_to_u8(value: u32) -> u8 {
+    ((value * 31 + 17) % 255) as u8
+}
+
+pub fn obfuscated_proxy_id(process_id: u32, region_id: &str) -> u16 {
+    let process_id = map_u32_to_u8(process_id);
+    let hash_region_id =  Sha256::digest(region_id.as_bytes());
+
+    const BASE: u64 = 257;
+    let combined_region = hash_region_id.iter().enumerate().fold(0u64, |acc, (i, &byte)| {
+        (acc + (byte as u64 * BASE.pow(i as u32))) % 255
+    });
+
+    let combined_region = (combined_region % 255) as u8;
+    ((combined_region as u16) * 257 + (process_id as u16)) % 65535
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -511,4 +530,12 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_proxy_id_obfuscation() {
+        let process_id = 123;
+        let region_id = "us-west-2";
+        let proxy_id = obfuscated_proxy_id(process_id, region_id);
+        assert_eq!(proxy_id, 0x1f7b);
+    }
 }

proxy/src/console_redirect_proxy.rs

@@ -25,6 +25,7 @@ pub async fn task_main(
     listener: tokio::net::TcpListener,
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandlerMain>,
+    proxy_id: u16,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -89,6 +90,7 @@ pub async fn task_main(
                 peer_addr,
                 crate::metrics::Protocol::Tcp,
                 &config.region,
+                proxy_id,
             );
 
             let res = handle_client(
@@ -222,7 +224,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
 
     node.cancel_closure
         .set_ip_allowlist(ip_allowlist.unwrap_or_default());
-    let session = cancellation_handler.get_session();
+    let session = cancellation_handler.get_session(ctx.proxy_id());
     prepare_client_connection(&node, &session, &mut stream).await?;
 
     // Before proxy passing, forward to compute whatever data is left in the

proxy/src/context/mod.rs

@@ -46,6 +46,7 @@ struct RequestContextInner {
     pub(crate) protocol: Protocol,
     first_packet: chrono::DateTime<Utc>,
     region: &'static str,
+    pub(crate) proxy_id: u16, // for generating cancel keys per region/process
     pub(crate) span: Span,
 
     // filled in as they are discovered
@@ -92,6 +93,7 @@ impl Clone for RequestContext {
             protocol: inner.protocol,
             first_packet: inner.first_packet,
             region: inner.region,
+            proxy_id: inner.proxy_id,
             span: info_span!("background_task"),
 
             project: inner.project,
@@ -124,6 +126,7 @@ impl RequestContext {
         conn_info: ConnectionInfo,
         protocol: Protocol,
         region: &'static str,
+        proxy_id: u16,
     ) -> Self {
         // TODO: be careful with long lived spans
         let span = info_span!(
@@ -141,6 +144,7 @@ impl RequestContext {
             protocol,
             first_packet: Utc::now(),
             region,
+            proxy_id,
             span,
 
             project: None,
@@ -172,7 +176,7 @@ impl RequestContext {
         let ip = IpAddr::from([127, 0, 0, 1]);
         let addr = SocketAddr::new(ip, 5432);
         let conn_info = ConnectionInfo { addr, extra: None };
-        RequestContext::new(Uuid::now_v7(), conn_info, Protocol::Tcp, "test")
+        RequestContext::new(Uuid::now_v7(), conn_info, Protocol::Tcp, "test", 1)
     }
 
     pub(crate) fn console_application_name(&self) -> String {
@@ -334,6 +338,10 @@ impl RequestContext {
             .latency_timer
             .success();
     }
+
+    pub(crate) fn proxy_id(&self) -> u16 {
+        self.0.try_lock().expect("should not deadlock").proxy_id
+    }
 }
 
 pub(crate) struct LatencyTimerPause<'a> {

proxy/src/proxy/mod.rs

@@ -59,6 +59,7 @@ pub async fn task_main(
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandlerMain>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+    proxy_id: u16,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -124,6 +125,7 @@ pub async fn task_main(
                 conn_info,
                 crate::metrics::Protocol::Tcp,
                 &config.region,
+                proxy_id,
             );
 
             let res = handle_client(
@@ -358,7 +360,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
 
     node.cancel_closure
         .set_ip_allowlist(ip_allowlist.unwrap_or_default());
-    let session = cancellation_handler.get_session();
+    let session = cancellation_handler.get_session(ctx.proxy_id());
     prepare_client_connection(&node, &session, &mut stream).await?;
 
     // Before proxy passing, forward to compute whatever data is left in the

proxy/src/serverless/mod.rs

@@ -63,6 +63,7 @@ pub async fn task_main(
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandlerMain>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+    proxy_id: u16,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("websocket server has shut down");
@@ -198,6 +199,7 @@ pub async fn task_main(
                     conn,
                     conn_info,
                     session_id,
+                    proxy_id,
                 ))
                 .await;
             }
@@ -324,6 +326,7 @@ async fn connection_handler(
     conn: AsyncRW,
     conn_info: ConnectionInfo,
     session_id: uuid::Uuid,
+    proxy_id: u16,
 ) {
     let session_id = AtomicTake::new(session_id);
 
@@ -371,6 +374,7 @@ async fn connection_handler(
                     http_request_token,
                     endpoint_rate_limiter.clone(),
                     cancellations,
+                    proxy_id,
                 )
                 .in_current_span()
                 .map_ok_or_else(api_error_into_response, |r| r),
@@ -419,6 +423,7 @@ async fn request_handler(
     http_cancellation_token: CancellationToken,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellations: TaskTracker,
+    proxy_id: u16,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
     let host = request
         .headers()
@@ -436,6 +441,7 @@ async fn request_handler(
             conn_info,
             crate::metrics::Protocol::Ws,
             &config.region,
+            proxy_id,
         );
 
         let span = ctx.span();
@@ -473,6 +479,7 @@ async fn request_handler(
             conn_info,
             crate::metrics::Protocol::Http,
             &config.region,
+            proxy_id,
         );
         let span = ctx.span();
 

safekeeper/client/src/mgmt_api.rs

@@ -102,6 +102,11 @@ impl Client {
         self.get(&uri).await
     }
 
+    pub async fn utilization(&self) -> Result<reqwest::Response> {
+        let uri = format!("{}/v1/utilization/", self.mgmt_api_endpoint);
+        self.get(&uri).await
+    }
+
     async fn get<U: IntoUrl>(&self, uri: U) -> Result<reqwest::Response> {
         self.request(Method::GET, uri, ()).await
     }

safekeeper/src/http/routes.rs

@@ -127,6 +127,13 @@ async fn timeline_create_handler(mut request: Request<Body>) -> Result<Response<
     json_response(StatusCode::OK, ())
 }
 
+async fn utilization_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
+    check_permission(&request, None)?;
+    let global_timelines = get_global_timelines(&request);
+    let utilization = global_timelines.get_timeline_counts();
+    json_response(StatusCode::OK, utilization)
+}
+
 /// List all (not deleted) timelines.
 /// Note: it is possible to do the same with debug_dump.
 async fn timeline_list_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
@@ -620,6 +627,7 @@ pub fn make_router(
                 failpoints_handler(r, cancel).await
             })
         })
+        .get("/v1/uzilization", |r| request_span(r, utilization_handler))
         .delete("/v1/tenant/:tenant_id", |r| {
             request_span(r, tenant_delete_handler)
         })

safekeeper/src/timelines_global_map.rs

@@ -13,6 +13,7 @@ use anyhow::{bail, Context, Result};
 use camino::Utf8PathBuf;
 use camino_tempfile::Utf8TempDir;
 use safekeeper_api::membership::Configuration;
+use safekeeper_api::models::SafekeeperUtilization;
 use safekeeper_api::ServerInfo;
 use serde::Serialize;
 use std::collections::HashMap;
@@ -416,6 +417,20 @@ impl GlobalTimelines {
             .collect()
     }
 
+    /// Returns statistics about timeline counts
+    pub fn get_timeline_counts(&self) -> SafekeeperUtilization {
+        let global_lock = self.state.lock().unwrap();
+        let timeline_count = global_lock
+            .timelines
+            .values()
+            .filter(|t| match t {
+                GlobalMapTimeline::CreationInProgress => false,
+                GlobalMapTimeline::Timeline(t) => !t.is_cancelled(),
+            })
+            .count() as u64;
+        SafekeeperUtilization { timeline_count }
+    }
+
     /// Returns all timelines belonging to a given tenant. Used for deleting all timelines of a tenant,
     /// and that's why it can return cancelled timelines, to retry deleting them.
     fn get_all_for_tenant(&self, tenant_id: TenantId) -> Vec<Arc<Timeline>> {

storage_controller/src/service/chaos_injector.rs

@@ -1,11 +1,17 @@
-use std::{sync::Arc, time::Duration};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+    time::Duration,
+};
 
 use pageserver_api::controller_api::ShardSchedulingPolicy;
 use rand::seq::SliceRandom;
 use rand::thread_rng;
 use tokio_util::sync::CancellationToken;
+use utils::id::NodeId;
+use utils::shard::TenantShardId;
 
-use super::Service;
+use super::{Node, Scheduler, Service, TenantShard};
 
 pub struct ChaosInjector {
     service: Arc<Service>,
@@ -35,50 +41,86 @@ impl ChaosInjector {
         }
     }
 
+    /// If a shard has a secondary and attached location, then re-assign the secondary to be
+    /// attached and the attached to be secondary.
+    ///
+    /// Only modifies tenants if they're in Active scheduling policy.
+    fn maybe_migrate_to_secondary(
+        &self,
+        tenant_shard_id: TenantShardId,
+        nodes: &Arc<HashMap<NodeId, Node>>,
+        tenants: &mut BTreeMap<TenantShardId, TenantShard>,
+        scheduler: &mut Scheduler,
+    ) {
+        let shard = tenants
+            .get_mut(&tenant_shard_id)
+            .expect("Held lock between choosing ID and this get");
+
+        if !matches!(shard.get_scheduling_policy(), ShardSchedulingPolicy::Active) {
+            // Skip non-active scheduling policies, so that a shard with a policy like Pause can
+            // be pinned without being disrupted by us.
+            tracing::info!(
+                "Skipping shard {tenant_shard_id}: scheduling policy is {:?}",
+                shard.get_scheduling_policy()
+            );
+            return;
+        }
+
+        // Pick a secondary to promote
+        let Some(new_location) = shard
+            .intent
+            .get_secondary()
+            .choose(&mut thread_rng())
+            .cloned()
+        else {
+            tracing::info!(
+                "Skipping shard {tenant_shard_id}: no secondary location, can't migrate"
+            );
+            return;
+        };
+
+        let Some(old_location) = *shard.intent.get_attached() else {
+            tracing::info!("Skipping shard {tenant_shard_id}: currently has no attached location");
+            return;
+        };
+
+        tracing::info!("Injecting chaos: migrate {tenant_shard_id} {old_location}->{new_location}");
+
+        shard.intent.demote_attached(scheduler, old_location);
+        shard.intent.promote_attached(scheduler, new_location);
+        self.service.maybe_reconcile_shard(shard, nodes);
+    }
+
     async fn inject_chaos(&mut self) {
         // Pick some shards to interfere with
         let batch_size = 128;
         let mut inner = self.service.inner.write().unwrap();
         let (nodes, tenants, scheduler) = inner.parts_mut();
         let tenant_ids = tenants.keys().cloned().collect::<Vec<_>>();
-        let victims = tenant_ids.choose_multiple(&mut thread_rng(), batch_size);
 
-        for victim in victims {
-            let shard = tenants
-                .get_mut(victim)
-                .expect("Held lock between choosing ID and this get");
-
-            if !matches!(shard.get_scheduling_policy(), ShardSchedulingPolicy::Active) {
-                // Skip non-active scheduling policies, so that a shard with a policy like Pause can
-                // be pinned without being disrupted by us.
-                tracing::info!(
-                    "Skipping shard {victim}: scheduling policy is {:?}",
-                    shard.get_scheduling_policy()
-                );
-                continue;
+        // Prefer to migrate tenants that are currently outside their home AZ.  This avoids the chaos injector
+        // continuously pushing tenants outside their home AZ: instead, we'll tend to cycle between picking some
+        // random tenants to move, and then on next chaos iteration moving them back, then picking some new
+        // random tenants on the next iteration.
+        let mut victims = Vec::with_capacity(batch_size);
+        for shard in tenants.values() {
+            if shard.is_attached_outside_preferred_az(nodes) {
+                victims.push(shard.tenant_shard_id);
             }
 
-            // Pick a secondary to promote
-            let Some(new_location) = shard
-                .intent
-                .get_secondary()
-                .choose(&mut thread_rng())
-                .cloned()
-            else {
-                tracing::info!("Skipping shard {victim}: no secondary location, can't migrate");
-                continue;
-            };
+            if victims.len() >= batch_size {
+                break;
+            }
+        }
 
-            let Some(old_location) = *shard.intent.get_attached() else {
-                tracing::info!("Skipping shard {victim}: currently has no attached location");
-                continue;
-            };
+        let choose_random = batch_size.saturating_sub(victims.len());
+        tracing::info!("Injecting chaos: found {} shards to migrate back to home AZ, picking {choose_random} random shards to migrate", victims.len());
 
-            tracing::info!("Injecting chaos: migrate {victim} {old_location}->{new_location}");
+        let random_victims = tenant_ids.choose_multiple(&mut thread_rng(), choose_random);
+        victims.extend(random_victims);
 
-            shard.intent.demote_attached(scheduler, old_location);
-            shard.intent.promote_attached(scheduler, new_location);
-            self.service.maybe_reconcile_shard(shard, nodes);
+        for victim in victims {
+            self.maybe_migrate_to_secondary(victim, nodes, tenants, scheduler);
         }
     }
 }

storage_controller/src/tenant_shard.rs

@@ -1793,6 +1793,23 @@ impl TenantShard {
             }
         }
     }
+
+    /// Returns true if the tenant shard is attached to a node that is outside the preferred AZ.
+    ///
+    /// If the shard does not have a preferred AZ, returns false.
+    pub(crate) fn is_attached_outside_preferred_az(&self, nodes: &HashMap<NodeId, Node>) -> bool {
+        self.intent
+            .get_attached()
+            .map(|node_id| {
+                Some(
+                    nodes
+                        .get(&node_id)
+                        .expect("referenced node exists")
+                        .get_availability_zone_id(),
+                ) == self.intent.preferred_az_id.as_ref()
+            })
+            .unwrap_or(false)
+    }
 }
 
 impl Drop for TenantShard {