Change expected error message for test_physical_replication_config_mismatch_max_locks_per_transaction

Cargo.lock

@@ -1605,32 +1605,6 @@ dependencies = [
  "typenum",
 ]
 
-[[package]]
-name = "curve25519-dalek"
-version = "4.1.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be"
-dependencies = [
- "cfg-if",
- "cpufeatures",
- "curve25519-dalek-derive",
- "digest",
- "fiat-crypto",
- "rustc_version",
- "subtle",
-]
-
-[[package]]
-name = "curve25519-dalek-derive"
-version = "0.1.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.90",
-]
-
 [[package]]
 name = "darling"
 version = "0.20.1"
@@ -1901,28 +1875,6 @@ dependencies = [
  "spki 0.7.3",
 ]
 
-[[package]]
-name = "ed25519"
-version = "2.2.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53"
-dependencies = [
- "signature 2.2.0",
-]
-
-[[package]]
-name = "ed25519-dalek"
-version = "2.1.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4a3daa8e81a3963a60642bcc1f90a670680bd4a77535faa384e9d1c79d620871"
-dependencies = [
- "curve25519-dalek",
- "ed25519",
- "rand_core 0.6.4",
- "sha2",
- "subtle",
-]
-
 [[package]]
 name = "either"
 version = "1.8.1"
@@ -2161,12 +2113,6 @@ dependencies = [
  "subtle",
 ]
 
-[[package]]
-name = "fiat-crypto"
-version = "0.2.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d"
-
 [[package]]
 name = "filetime"
 version = "0.2.22"
@@ -4478,8 +4424,6 @@ dependencies = [
  "bytes",
  "fallible-iterator",
  "postgres-protocol",
- "serde",
- "serde_json",
 ]
 
 [[package]]
@@ -4801,7 +4745,6 @@ dependencies = [
  "consumption_metrics",
  "dashmap 5.5.0",
  "ecdsa 0.16.9",
- "ed25519-dalek",
  "env_logger 0.10.2",
  "fallible-iterator",
  "flate2",
@@ -8024,7 +7967,6 @@ dependencies = [
  "num-traits",
  "once_cell",
  "parquet",
- "postgres-types",
  "prettyplease",
  "proc-macro2",
  "prost",
@@ -8051,7 +7993,6 @@ dependencies = [
  "time",
  "time-macros",
  "tokio",
- "tokio-postgres",
  "tokio-rustls 0.26.0",
  "tokio-stream",
  "tokio-util",

build-tools.Dockerfile

@@ -115,7 +115,7 @@ RUN set -e \
 
 # Keep the version the same as in compute/compute-node.Dockerfile and
 # test_runner/regress/test_compute_metrics.py.
-ENV SQL_EXPORTER_VERSION=0.17.0
+ENV SQL_EXPORTER_VERSION=0.16.0
 RUN curl -fsSL \
     "https://github.com/burningalchemist/sql_exporter/releases/download/${SQL_EXPORTER_VERSION}/sql_exporter-${SQL_EXPORTER_VERSION}.linux-$(case "$(uname -m)" in x86_64) echo amd64;; aarch64) echo arm64;; esac).tar.gz" \
     --output sql_exporter.tar.gz \

compute/compute-node.Dockerfile

@@ -170,6 +170,7 @@ RUN case "${PG_VERSION}" in \
     wget https://download.osgeo.org/postgis/source/postgis-${POSTGIS_VERSION}.tar.gz -O postgis.tar.gz && \
     echo "${POSTGIS_CHECKSUM} postgis.tar.gz" | sha256sum --check && \
     mkdir postgis-src && cd postgis-src && tar xzf ../postgis.tar.gz --strip-components=1 -C . && \
+    find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /before.txt &&\
     ./autogen.sh && \
     ./configure --with-sfcgal=/usr/local/bin/sfcgal-config && \
     make -j $(getconf _NPROCESSORS_ONLN) && \
@@ -219,7 +220,11 @@ RUN case "${PG_VERSION}" in \
     cmake -GNinja -DCMAKE_BUILD_TYPE=Release .. && \
     ninja -j $(getconf _NPROCESSORS_ONLN) && \
     ninja -j $(getconf _NPROCESSORS_ONLN) install && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrouting.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrouting.control && \
+    find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /after.txt &&\
+    cp /usr/local/pgsql/share/extension/pgrouting.control /extensions/postgis && \
+    sort -o /before.txt /before.txt && sort -o /after.txt /after.txt && \
+    comm -13 /before.txt /after.txt | tar --directory=/usr/local/pgsql --zstd -cf /extensions/postgis.tar.zst -T -
 
 #########################################################################################
 #
@@ -837,8 +842,13 @@ RUN case "${PG_VERSION}" in "v17") \
     wget  https://github.com/neondatabase/postgresql_anonymizer/archive/refs/tags/neon_1.1.1.tar.gz -O pg_anon.tar.gz && \
     echo "321ea8d5c1648880aafde850a2c576e4a9e7b9933a34ce272efc839328999fa9  pg_anon.tar.gz" | sha256sum --check && \
     mkdir pg_anon-src && cd pg_anon-src && tar xzf ../pg_anon.tar.gz --strip-components=1 -C . && \
+    find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /before.txt &&\
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/anon.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/anon.control && \
+    find /usr/local/pgsql -type f | sed 's|^/usr/local/pgsql/||' > /after.txt &&\
+    mkdir -p /extensions/anon && cp /usr/local/pgsql/share/extension/anon.control /extensions/anon && \
+    sort -o /before.txt /before.txt && sort -o /after.txt /after.txt && \
+    comm -13 /before.txt /after.txt | tar --directory=/usr/local/pgsql --zstd -cf /extensions/anon.tar.zst -T -
 
 #########################################################################################
 #
@@ -1055,8 +1065,8 @@ ARG PG_VERSION
 # NOTE: local_proxy depends on the version of pg_session_jwt
 # Do not update without approve from proxy team
 # Make sure the version is reflected in proxy/src/serverless/local_conn_pool.rs
-RUN wget https://github.com/neondatabase/pg_session_jwt/archive/refs/tags/v0.2.0.tar.gz -O pg_session_jwt.tar.gz && \
-    echo "5ace028e591f2e000ca10afa5b1ca62203ebff014c2907c0ec3b29c36f28a1bb pg_session_jwt.tar.gz" | sha256sum --check && \
+RUN wget https://github.com/neondatabase/pg_session_jwt/archive/refs/tags/v0.1.2-v17.tar.gz -O pg_session_jwt.tar.gz && \
+    echo "c8ecbed9cb8c6441bce5134a176002b043018adf9d05a08e457dda233090a86e pg_session_jwt.tar.gz" | sha256sum --check && \
     mkdir pg_session_jwt-src && cd pg_session_jwt-src && tar xzf ../pg_session_jwt.tar.gz --strip-components=1 -C . && \
     sed -i 's/pgrx = "0.12.6"/pgrx = { version = "=0.12.6", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     cargo pgrx install --release
@@ -1274,11 +1284,11 @@ RUN set -e \
 #
 #########################################################################################
 
-FROM quay.io/prometheuscommunity/postgres-exporter:v0.16.0 AS postgres-exporter
+FROM quay.io/prometheuscommunity/postgres-exporter:v0.12.1 AS postgres-exporter
 
 # Keep the version the same as in build-tools.Dockerfile and
 # test_runner/regress/test_compute_metrics.py.
-FROM burningalchemist/sql_exporter:0.17.0 AS sql-exporter
+FROM burningalchemist/sql_exporter:0.16.0 AS sql-exporter
 
 #########################################################################################
 #

control_plane/src/storage_controller.rs

@@ -822,7 +822,10 @@ impl StorageController {
         self.dispatch(
             Method::PUT,
             format!("control/v1/tenant/{tenant_shard_id}/migrate"),
-            Some(TenantShardMigrateRequest { node_id }),
+            Some(TenantShardMigrateRequest {
+                tenant_shard_id,
+                node_id,
+            }),
         )
         .await
     }

control_plane/storcon_cli/src/main.rs

@@ -112,13 +112,6 @@ enum Command {
         #[arg(long)]
         node: NodeId,
     },
-    /// Migrate the secondary location for a tenant shard to a specific pageserver.
-    TenantShardMigrateSecondary {
-        #[arg(long)]
-        tenant_shard_id: TenantShardId,
-        #[arg(long)]
-        node: NodeId,
-    },
     /// Cancel any ongoing reconciliation for this shard
     TenantShardCancelReconcile {
         #[arg(long)]
@@ -547,7 +540,10 @@ async fn main() -> anyhow::Result<()> {
             tenant_shard_id,
             node,
         } => {
-            let req = TenantShardMigrateRequest { node_id: node };
+            let req = TenantShardMigrateRequest {
+                tenant_shard_id,
+                node_id: node,
+            };
 
             storcon_client
                 .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
@@ -557,20 +553,6 @@ async fn main() -> anyhow::Result<()> {
                 )
                 .await?;
         }
-        Command::TenantShardMigrateSecondary {
-            tenant_shard_id,
-            node,
-        } => {
-            let req = TenantShardMigrateRequest { node_id: node };
-
-            storcon_client
-                .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
-                    Method::PUT,
-                    format!("control/v1/tenant/{tenant_shard_id}/migrate_secondary"),
-                    Some(req),
-                )
-                .await?;
-        }
         Command::TenantShardCancelReconcile { tenant_shard_id } => {
             storcon_client
                 .dispatch::<(), ()>(
@@ -933,7 +915,10 @@ async fn main() -> anyhow::Result<()> {
                             .dispatch::<TenantShardMigrateRequest, TenantShardMigrateResponse>(
                                 Method::PUT,
                                 format!("control/v1/tenant/{}/migrate", mv.tenant_shard_id),
-                                Some(TenantShardMigrateRequest { node_id: mv.to }),
+                                Some(TenantShardMigrateRequest {
+                                    tenant_shard_id: mv.tenant_shard_id,
+                                    node_id: mv.to,
+                                }),
                             )
                             .await
                             .map_err(|e| (mv.tenant_shard_id, mv.from, mv.to, e))

libs/pageserver_api/src/controller_api.rs

@@ -179,6 +179,7 @@ pub struct TenantDescribeResponseShard {
 /// specifies some constraints, e.g. asking it to get off particular node(s)
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TenantShardMigrateRequest {
+    pub tenant_shard_id: TenantShardId,
     pub node_id: NodeId,
 }
 
@@ -367,16 +368,6 @@ pub enum PlacementPolicy {
     Detached,
 }
 
-impl PlacementPolicy {
-    pub fn want_secondaries(&self) -> usize {
-        match self {
-            PlacementPolicy::Attached(secondary_count) => *secondary_count,
-            PlacementPolicy::Secondary => 1,
-            PlacementPolicy::Detached => 0,
-        }
-    }
-}
-
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TenantShardMigrateResponse {}
 

libs/pageserver_api/src/key.rs

@@ -706,7 +706,7 @@ pub fn repl_origin_key_range() -> Range<Key> {
 /// Non inherited range for vectored get.
 pub const NON_INHERITED_RANGE: Range<Key> = AUX_FILES_KEY..AUX_FILES_KEY.next();
 /// Sparse keyspace range for vectored get. Missing key error will be ignored for this range.
-pub const SPARSE_RANGE: Range<Key> = Key::metadata_key_range();
+pub const NON_INHERITED_SPARSE_RANGE: Range<Key> = Key::metadata_key_range();
 
 impl Key {
     // AUX_FILES currently stores only data for logical replication (slots etc), and
@@ -714,42 +714,7 @@ impl Key {
     // switch (and generally it likely should be optional), so ignore these.
     #[inline(always)]
     pub fn is_inherited_key(self) -> bool {
-        if self.is_sparse() {
-            self.is_inherited_sparse_key()
-        } else {
-            !NON_INHERITED_RANGE.contains(&self)
-        }
-    }
-
-    #[inline(always)]
-    pub fn is_sparse(self) -> bool {
-        self.field1 >= METADATA_KEY_BEGIN_PREFIX && self.field1 < METADATA_KEY_END_PREFIX
-    }
-
-    /// Check if the key belongs to the inherited keyspace.
-    fn is_inherited_sparse_key(self) -> bool {
-        debug_assert!(self.is_sparse());
-        self.field1 == RELATION_SIZE_PREFIX
-    }
-
-    pub fn sparse_non_inherited_keyspace() -> Range<Key> {
-        // The two keys are adjacent; if we will have non-adjancent keys in the future, we should return a keyspace
-        debug_assert_eq!(AUX_KEY_PREFIX + 1, REPL_ORIGIN_KEY_PREFIX);
-        Key {
-            field1: AUX_KEY_PREFIX,
-            field2: 0,
-            field3: 0,
-            field4: 0,
-            field5: 0,
-            field6: 0,
-        }..Key {
-            field1: REPL_ORIGIN_KEY_PREFIX + 1,
-            field2: 0,
-            field3: 0,
-            field4: 0,
-            field5: 0,
-            field6: 0,
-        }
+        !NON_INHERITED_RANGE.contains(&self) && !NON_INHERITED_SPARSE_RANGE.contains(&self)
     }
 
     #[inline(always)]

pageserver/src/tenant.rs

@@ -5682,7 +5682,7 @@ mod tests {
     use bytes::{Bytes, BytesMut};
     use hex_literal::hex;
     use itertools::Itertools;
-    use pageserver_api::key::{Key, AUX_KEY_PREFIX, NON_INHERITED_RANGE, RELATION_SIZE_PREFIX};
+    use pageserver_api::key::{Key, AUX_KEY_PREFIX, NON_INHERITED_RANGE};
     use pageserver_api::keyspace::KeySpace;
     use pageserver_api::models::{CompactionAlgorithm, CompactionAlgorithmSettings};
     use pageserver_api::value::Value;
@@ -7741,18 +7741,7 @@ mod tests {
         let base_key = Key::from_hex("620000000033333333444444445500000000").unwrap();
         let base_key_child = Key::from_hex("620000000033333333444444445500000001").unwrap();
         let base_key_nonexist = Key::from_hex("620000000033333333444444445500000002").unwrap();
-        let base_key_overwrite = Key::from_hex("620000000033333333444444445500000003").unwrap();
-
-        let base_inherited_key = Key::from_hex("610000000033333333444444445500000000").unwrap();
-        let base_inherited_key_child =
-            Key::from_hex("610000000033333333444444445500000001").unwrap();
-        let base_inherited_key_nonexist =
-            Key::from_hex("610000000033333333444444445500000002").unwrap();
-        let base_inherited_key_overwrite =
-            Key::from_hex("610000000033333333444444445500000003").unwrap();
-
         assert_eq!(base_key.field1, AUX_KEY_PREFIX); // in case someone accidentally changed the prefix...
-        assert_eq!(base_inherited_key.field1, RELATION_SIZE_PREFIX);
 
         let tline = tenant
             .create_test_timeline_with_layers(
@@ -7761,18 +7750,7 @@ mod tests {
                 DEFAULT_PG_VERSION,
                 &ctx,
                 Vec::new(), // delta layers
-                vec![(
-                    Lsn(0x20),
-                    vec![
-                        (base_inherited_key, test_img("metadata inherited key 1")),
-                        (
-                            base_inherited_key_overwrite,
-                            test_img("metadata key overwrite 1a"),
-                        ),
-                        (base_key, test_img("metadata key 1")),
-                        (base_key_overwrite, test_img("metadata key overwrite 1b")),
-                    ],
-                )], // image layers
+                vec![(Lsn(0x20), vec![(base_key, test_img("metadata key 1"))])], // image layers
                 Lsn(0x20), // it's fine to not advance LSN to 0x30 while using 0x30 to get below because `get_vectored_impl` does not wait for LSN
             )
             .await?;
@@ -7786,18 +7764,7 @@ mod tests {
                 Vec::new(), // delta layers
                 vec![(
                     Lsn(0x30),
-                    vec![
-                        (
-                            base_inherited_key_child,
-                            test_img("metadata inherited key 2"),
-                        ),
-                        (
-                            base_inherited_key_overwrite,
-                            test_img("metadata key overwrite 2a"),
-                        ),
-                        (base_key_child, test_img("metadata key 2")),
-                        (base_key_overwrite, test_img("metadata key overwrite 2b")),
-                    ],
+                    vec![(base_key_child, test_img("metadata key 2"))],
                 )], // image layers
                 Lsn(0x30),
             )
@@ -7819,26 +7786,6 @@ mod tests {
             get_vectored_impl_wrapper(&tline, base_key_nonexist, lsn, &ctx).await?,
             None
         );
-        assert_eq!(
-            get_vectored_impl_wrapper(&tline, base_key_overwrite, lsn, &ctx).await?,
-            Some(test_img("metadata key overwrite 1b"))
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&tline, base_inherited_key, lsn, &ctx).await?,
-            Some(test_img("metadata inherited key 1"))
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&tline, base_inherited_key_child, lsn, &ctx).await?,
-            None
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&tline, base_inherited_key_nonexist, lsn, &ctx).await?,
-            None
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&tline, base_inherited_key_overwrite, lsn, &ctx).await?,
-            Some(test_img("metadata key overwrite 1a"))
-        );
 
         // test vectored get on child timeline
         assert_eq!(
@@ -7853,82 +7800,6 @@ mod tests {
             get_vectored_impl_wrapper(&child, base_key_nonexist, lsn, &ctx).await?,
             None
         );
-        assert_eq!(
-            get_vectored_impl_wrapper(&child, base_inherited_key, lsn, &ctx).await?,
-            Some(test_img("metadata inherited key 1"))
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&child, base_inherited_key_child, lsn, &ctx).await?,
-            Some(test_img("metadata inherited key 2"))
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&child, base_inherited_key_nonexist, lsn, &ctx).await?,
-            None
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&child, base_key_overwrite, lsn, &ctx).await?,
-            Some(test_img("metadata key overwrite 2b"))
-        );
-        assert_eq!(
-            get_vectored_impl_wrapper(&child, base_inherited_key_overwrite, lsn, &ctx).await?,
-            Some(test_img("metadata key overwrite 2a"))
-        );
-
-        // test vectored scan on parent timeline
-        let mut reconstruct_state = ValuesReconstructState::new();
-        let res = tline
-            .get_vectored_impl(
-                KeySpace::single(Key::metadata_key_range()),
-                lsn,
-                &mut reconstruct_state,
-                &ctx,
-            )
-            .await?;
-
-        assert_eq!(
-            res.into_iter()
-                .map(|(k, v)| (k, v.unwrap()))
-                .collect::<Vec<_>>(),
-            vec![
-                (base_inherited_key, test_img("metadata inherited key 1")),
-                (
-                    base_inherited_key_overwrite,
-                    test_img("metadata key overwrite 1a")
-                ),
-                (base_key, test_img("metadata key 1")),
-                (base_key_overwrite, test_img("metadata key overwrite 1b")),
-            ]
-        );
-
-        // test vectored scan on child timeline
-        let mut reconstruct_state = ValuesReconstructState::new();
-        let res = child
-            .get_vectored_impl(
-                KeySpace::single(Key::metadata_key_range()),
-                lsn,
-                &mut reconstruct_state,
-                &ctx,
-            )
-            .await?;
-
-        assert_eq!(
-            res.into_iter()
-                .map(|(k, v)| (k, v.unwrap()))
-                .collect::<Vec<_>>(),
-            vec![
-                (base_inherited_key, test_img("metadata inherited key 1")),
-                (
-                    base_inherited_key_child,
-                    test_img("metadata inherited key 2")
-                ),
-                (
-                    base_inherited_key_overwrite,
-                    test_img("metadata key overwrite 2a")
-                ),
-                (base_key_child, test_img("metadata key 2")),
-                (base_key_overwrite, test_img("metadata key overwrite 2b")),
-            ]
-        );
 
         Ok(())
     }

pageserver/src/tenant/storage_layer.rs

@@ -12,7 +12,7 @@ pub mod merge_iterator;
 
 use crate::context::{AccessStatsBehavior, RequestContext};
 use bytes::Bytes;
-use pageserver_api::key::Key;
+use pageserver_api::key::{Key, NON_INHERITED_SPARSE_RANGE};
 use pageserver_api::keyspace::{KeySpace, KeySpaceRandomAccum};
 use pageserver_api::record::NeonWalRecord;
 use pageserver_api::value::Value;
@@ -209,7 +209,7 @@ impl ValuesReconstructState {
             .keys
             .entry(*key)
             .or_insert(Ok(VectoredValueReconstructState::default()));
-        let is_sparse_key = key.is_sparse();
+        let is_sparse_key = NON_INHERITED_SPARSE_RANGE.contains(key);
         if let Ok(state) = state {
             let key_done = match state.situation {
                 ValueReconstructSituation::Complete => {

pageserver/src/tenant/timeline.rs

@@ -27,7 +27,7 @@ use pageserver_api::{
     config::tenant_conf_defaults::DEFAULT_COMPACTION_THRESHOLD,
     key::{
         KEY_SIZE, METADATA_KEY_BEGIN_PREFIX, METADATA_KEY_END_PREFIX, NON_INHERITED_RANGE,
-        SPARSE_RANGE,
+        NON_INHERITED_SPARSE_RANGE,
     },
     keyspace::{KeySpaceAccum, KeySpaceRandomAccum, SparseKeyPartitioning},
     models::{
@@ -3221,7 +3221,7 @@ impl Timeline {
             // We don't return a blanket [`GetVectoredError::MissingKey`] to avoid
             // stalling compaction.
             keyspace.remove_overlapping_with(&KeySpace {
-                ranges: vec![NON_INHERITED_RANGE, Key::sparse_non_inherited_keyspace()],
+                ranges: vec![NON_INHERITED_RANGE, NON_INHERITED_SPARSE_RANGE],
             });
 
             // Keyspace is fully retrieved
@@ -3242,11 +3242,7 @@ impl Timeline {
             // keys from `keyspace`, we expect there to be no overlap between it and the image covered key
             // space. If that's not the case, we had at least one key encounter a gap in the image layer
             // and stop the search as a result of that.
-            let mut removed = keyspace.remove_overlapping_with(&image_covered_keyspace);
-            // Do not fire missing key error for sparse keys.
-            removed.remove_overlapping_with(&KeySpace {
-                ranges: vec![SPARSE_RANGE],
-            });
+            let removed = keyspace.remove_overlapping_with(&image_covered_keyspace);
             if !removed.is_empty() {
                 break Some(removed);
             }
@@ -3261,21 +3257,6 @@ impl Timeline {
             timeline = &*timeline_owned;
         };
 
-        // Remove sparse keys from the keyspace so that it doesn't fire errors.
-        let missing_keyspace = if let Some(missing_keyspace) = missing_keyspace {
-            let mut missing_keyspace = missing_keyspace;
-            missing_keyspace.remove_overlapping_with(&KeySpace {
-                ranges: vec![SPARSE_RANGE],
-            });
-            if missing_keyspace.is_empty() {
-                None
-            } else {
-                Some(missing_keyspace)
-            }
-        } else {
-            None
-        };
-
         if let Some(missing_keyspace) = missing_keyspace {
             return Err(GetVectoredError::MissingKey(MissingKeyError {
                 key: missing_keyspace.start().unwrap(), /* better if we can store the full keyspace */

proxy/Cargo.toml

@@ -106,7 +106,6 @@ jose-jwk = { version = "0.1.2", features = ["p256", "p384", "rsa"] }
 signature = "2"
 ecdsa = "0.16"
 p256 = { version = "0.13", features = ["jwk"] }
-ed25519-dalek = { version = "2", default-features = false, features = ["rand_core"] }
 rsa = "0.9"
 
 workspace_hack.workspace = true
@@ -121,4 +120,4 @@ rcgen.workspace = true
 rstest.workspace = true
 walkdir.workspace = true
 rand_distr = "0.4"
-tokio-postgres = { workspace = true, features = ["with-serde_json-1"] }
+tokio-postgres.workspace = true

proxy/src/auth/backend/console_redirect.rs

@@ -94,7 +94,7 @@ impl BackendIpAllowlist for ConsoleRedirectBackend {
         self.api
             .get_allowed_ips_and_secret(ctx, user_info)
             .await
-            .map(|(ips, _)| ips.0.clone())
+            .map(|(ips, _)| ips.as_ref().clone())
             .map_err(|e| e.into())
     }
 }

proxy/src/auth/backend/jwt.rs

@@ -44,18 +44,6 @@ pub(crate) trait FetchAuthRules: Clone + Send + Sync + 'static {
     ) -> impl Future<Output = Result<Vec<AuthRule>, FetchAuthRulesError>> + Send;
 }
 
-#[derive(Clone)]
-pub(crate) struct StaticAuthRules(pub Vec<AuthRule>);
-impl FetchAuthRules for StaticAuthRules {
-    async fn fetch_auth_rules(
-        &self,
-        _ctx: &RequestContext,
-        _endpoint: EndpointId,
-    ) -> Result<Vec<AuthRule>, FetchAuthRulesError> {
-        Ok(self.0.clone())
-    }
-}
-
 #[derive(Error, Debug)]
 pub(crate) enum FetchAuthRulesError {
     #[error(transparent)]
@@ -65,7 +53,7 @@ pub(crate) enum FetchAuthRulesError {
     RoleJwksNotConfigured,
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone)]
 pub(crate) struct AuthRule {
     pub(crate) id: String,
     pub(crate) jwks_url: url::Url,

proxy/src/auth/backend/mod.rs

@@ -10,7 +10,6 @@ use std::sync::Arc;
 pub use console_redirect::ConsoleRedirectBackend;
 pub(crate) use console_redirect::ConsoleRedirectError;
 use ipnet::{Ipv4Net, Ipv6Net};
-use jwt::{JwkCache, StaticAuthRules};
 use local::LocalBackend;
 use postgres_client::config::AuthKeys;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -260,7 +259,6 @@ pub(crate) trait BackendIpAllowlist {
 /// Here, we choose the appropriate auth flow based on circumstances.
 ///
 /// All authentication flows will emit an AuthenticationOk message if successful.
-#[allow(clippy::too_many_arguments)]
 async fn auth_quirks(
     ctx: &RequestContext,
     api: &impl control_plane::ControlPlaneApi,
@@ -269,7 +267,6 @@ async fn auth_quirks(
     allow_cleartext: bool,
     config: &'static AuthenticationConfig,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-    jwks_cache: Arc<JwkCache>,
 ) -> auth::Result<(ComputeCredentials, Option<Vec<IpPattern>>)> {
     // If there's no project so far, that entails that client doesn't
     // support SNI or other means of passing the endpoint (project) name.
@@ -285,54 +282,11 @@ async fn auth_quirks(
     };
 
     debug!("fetching user's authentication info");
-    let (x, maybe_secret) = api.get_allowed_ips_and_secret(ctx, &info).await?;
-    let (allowed_ips, auth_rules) = &**x;
-
-    // we expect a jwt in the options field
-    if !auth_rules.is_empty() {
-        match info.options.get("jwt") {
-            Some(jwt) => {
-                let creds = jwks_cache
-                    .check_jwt(
-                        ctx,
-                        info.endpoint.clone(),
-                        &info.user,
-                        &StaticAuthRules(auth_rules.clone()),
-                        &jwt,
-                    )
-                    .await?;
-                let token = match creds {
-                    ComputeCredentialKeys::JwtPayload(payload) => {
-                        serde_json::from_slice::<serde_json::Value>(&payload)
-                            .expect("jwt payload is valid json")
-                    }
-                    _ => unreachable!(),
-                };
-
-                // the token has a required IP claim.
-                if let Some(expected_ip) = token.get("ip") {
-                    // todo: don't panic here, obviously.
-                    let allowed_ips: Vec<IpPattern> = expected_ip
-                        .as_str()
-                        .expect("jwt should not have an invalid IP claim")
-                        .split(',')
-                        .map(|s| s.parse().expect("jwt should not have an invalid IP claim"))
-                        .collect();
-
-                    if !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips) {
-                        return Err(auth::AuthError::ip_address_not_allowed(ctx.peer_addr()));
-                    }
-                }
-            }
-            None => {
-                return Err(AuthError::bad_auth_method("needs jwt"));
-            }
-        }
-    }
+    let (allowed_ips, maybe_secret) = api.get_allowed_ips_and_secret(ctx, &info).await?;
 
     // check allowed list
     if config.ip_allowlist_check_enabled
-        && !check_peer_addr_is_in_list(&ctx.peer_addr(), allowed_ips)
+        && !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips)
     {
         return Err(auth::AuthError::ip_address_not_allowed(ctx.peer_addr()));
     }
@@ -372,7 +326,7 @@ async fn auth_quirks(
     )
     .await
     {
-        Ok(keys) => Ok((keys, Some(allowed_ips.clone()))),
+        Ok(keys) => Ok((keys, Some(allowed_ips.as_ref().clone()))),
         Err(e) => {
             if e.is_password_failed() {
                 // The password could have been changed, so we invalidate the cache.
@@ -442,7 +396,6 @@ impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
         allow_cleartext: bool,
         config: &'static AuthenticationConfig,
         endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-        jwks_cache: Arc<JwkCache>,
     ) -> auth::Result<(Backend<'a, ComputeCredentials>, Option<Vec<IpPattern>>)> {
         let res = match self {
             Self::ControlPlane(api, user_info) => {
@@ -460,7 +413,6 @@ impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
                     allow_cleartext,
                     config,
                     endpoint_rate_limiter,
-                    jwks_cache,
                 )
                 .await?;
                 Ok((Backend::ControlPlane(api, credentials), ip_allowlist))
@@ -495,7 +447,7 @@ impl Backend<'_, ComputeUserInfo> {
             Self::ControlPlane(api, user_info) => {
                 api.get_allowed_ips_and_secret(ctx, user_info).await
             }
-            Self::Local(_) => Ok((Cached::new_uncached(Arc::new((vec![], vec![]))), None)),
+            Self::Local(_) => Ok((Cached::new_uncached(Arc::new(vec![])), None)),
         }
     }
 }
@@ -509,11 +461,11 @@ impl BackendIpAllowlist for Backend<'_, ()> {
     ) -> auth::Result<Vec<auth::IpPattern>> {
         let auth_data = match self {
             Self::ControlPlane(api, ()) => api.get_allowed_ips_and_secret(ctx, user_info).await,
-            Self::Local(_) => Ok((Cached::new_uncached(Arc::new((vec![], vec![]))), None)),
+            Self::Local(_) => Ok((Cached::new_uncached(Arc::new(vec![])), None)),
         };
 
         auth_data
-            .map(|(ips, _)| ips.0.clone())
+            .map(|(ips, _)| ips.as_ref().clone())
             .map_err(|e| e.into())
     }
 }
@@ -591,7 +543,7 @@ mod tests {
             control_plane::errors::GetAuthInfoError,
         > {
             Ok((
-                CachedAllowedIps::new_uncached(Arc::new((self.ips.clone(), vec![]))),
+                CachedAllowedIps::new_uncached(Arc::new(self.ips.clone())),
                 Some(CachedRoleSecret::new_uncached(Some(self.secret.clone()))),
             ))
         }
@@ -751,7 +703,6 @@ mod tests {
             false,
             &CONFIG,
             endpoint_rate_limiter,
-            Arc::new(JwkCache::default()),
         )
         .await
         .unwrap();
@@ -807,7 +758,6 @@ mod tests {
             true,
             &CONFIG,
             endpoint_rate_limiter,
-            Arc::new(JwkCache::default()),
         )
         .await
         .unwrap();
@@ -861,7 +811,6 @@ mod tests {
             true,
             &CONFIG,
             endpoint_rate_limiter,
-            Arc::new(JwkCache::default()),
         )
         .await
         .unwrap();

proxy/src/cache/project_info.rs

@@ -1,5 +1,3 @@
-#![allow(clippy::type_complexity)]
-
 use std::collections::HashSet;
 use std::convert::Infallible;
 use std::sync::atomic::AtomicU64;
@@ -15,7 +13,6 @@ use tokio::time::Instant;
 use tracing::{debug, info};
 
 use super::{Cache, Cached};
-use crate::auth::backend::jwt::AuthRule;
 use crate::auth::IpPattern;
 use crate::config::ProjectInfoCacheOptions;
 use crate::control_plane::AuthSecret;
@@ -53,7 +50,7 @@ impl<T> From<T> for Entry<T> {
 #[derive(Default)]
 struct EndpointInfo {
     secret: std::collections::HashMap<RoleNameInt, Entry<Option<AuthSecret>>>,
-    allowed_ips: Option<Entry<Arc<(Vec<IpPattern>, Vec<AuthRule>)>>>,
+    allowed_ips: Option<Entry<Arc<Vec<IpPattern>>>>,
 }
 
 impl EndpointInfo {
@@ -84,7 +81,7 @@ impl EndpointInfo {
         &self,
         valid_since: Instant,
         ignore_cache_since: Option<Instant>,
-    ) -> Option<(Arc<(Vec<IpPattern>, Vec<AuthRule>)>, bool)> {
+    ) -> Option<(Arc<Vec<IpPattern>>, bool)> {
         if let Some(allowed_ips) = &self.allowed_ips {
             if valid_since < allowed_ips.created_at {
                 return Some((
@@ -214,7 +211,7 @@ impl ProjectInfoCacheImpl {
     pub(crate) fn get_allowed_ips(
         &self,
         endpoint_id: &EndpointId,
-    ) -> Option<Cached<&Self, Arc<(Vec<IpPattern>, Vec<AuthRule>)>>> {
+    ) -> Option<Cached<&Self, Arc<Vec<IpPattern>>>> {
         let endpoint_id = EndpointIdInt::get(endpoint_id)?;
         let (valid_since, ignore_cache_since) = self.get_cache_times();
         let endpoint_info = self.cache.get(&endpoint_id)?;
@@ -250,7 +247,7 @@ impl ProjectInfoCacheImpl {
         &self,
         project_id: ProjectIdInt,
         endpoint_id: EndpointIdInt,
-        allowed_ips: Arc<(Vec<IpPattern>, Vec<AuthRule>)>,
+        allowed_ips: Arc<Vec<IpPattern>>,
     ) {
         if self.cache.len() >= self.config.size {
             // If there are too many entries, wait until the next gc cycle.
@@ -389,10 +386,10 @@ mod tests {
         let user2: RoleName = "user2".into();
         let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
         let secret2 = None;
-        let allowed_ips = Arc::new((
-            vec!["127.0.0.1".parse().unwrap(), "127.0.0.2".parse().unwrap()],
-            vec![],
-        ));
+        let allowed_ips = Arc::new(vec![
+            "127.0.0.1".parse().unwrap(),
+            "127.0.0.2".parse().unwrap(),
+        ]);
         cache.insert_role_secret(
             (&project_id).into(),
             (&endpoint_id).into(),
@@ -460,10 +457,10 @@ mod tests {
         let user2: RoleName = "user2".into();
         let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
         let secret2 = Some(AuthSecret::Scram(ServerSecret::mock([2; 32])));
-        let allowed_ips = Arc::new((
-            vec!["127.0.0.1".parse().unwrap(), "127.0.0.2".parse().unwrap()],
-            vec![],
-        ));
+        let allowed_ips = Arc::new(vec![
+            "127.0.0.1".parse().unwrap(),
+            "127.0.0.2".parse().unwrap(),
+        ]);
         cache.insert_role_secret(
             (&project_id).into(),
             (&endpoint_id).into(),
@@ -523,10 +520,10 @@ mod tests {
         let user2: RoleName = "user2".into();
         let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
         let secret2 = Some(AuthSecret::Scram(ServerSecret::mock([2; 32])));
-        let allowed_ips = Arc::new((
-            vec!["127.0.0.1".parse().unwrap(), "127.0.0.2".parse().unwrap()],
-            vec![],
-        ));
+        let allowed_ips = Arc::new(vec![
+            "127.0.0.1".parse().unwrap(),
+            "127.0.0.2".parse().unwrap(),
+        ]);
         cache.insert_role_secret(
             (&project_id).into(),
             (&endpoint_id).into(),

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -153,22 +153,9 @@ impl NeonControlPlaneClient {
                 .proxy
                 .allowed_ips_number
                 .observe(allowed_ips.len() as f64);
-
-            let auth_rules = body
-                .jwks
-                .into_iter()
-                .map(|jwks| AuthRule {
-                    id: jwks.id,
-                    jwks_url: jwks.jwks_url,
-                    audience: jwks.jwt_audience,
-                    role_names: jwks.role_names,
-                })
-                .collect();
-
             Ok(AuthInfo {
                 secret,
                 allowed_ips,
-                auth_rules,
                 project_id: body.project_id,
             })
         }
@@ -323,7 +310,7 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
             self.caches.project_info.insert_allowed_ips(
                 project_id,
                 normalized_ep_int,
-                Arc::new((auth_info.allowed_ips, auth_info.auth_rules)),
+                Arc::new(auth_info.allowed_ips),
             );
             ctx.set_project_id(project_id);
         }
@@ -349,7 +336,7 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
             .allowed_ips_cache_misses
             .inc(CacheOutcome::Miss);
         let auth_info = self.do_get_auth_info(ctx, user_info).await?;
-        let allowed_ips = Arc::new((auth_info.allowed_ips, auth_info.auth_rules));
+        let allowed_ips = Arc::new(auth_info.allowed_ips);
         let user = &user_info.user;
         if let Some(project_id) = auth_info.project_id {
             let normalized_ep_int = normalized_ep.into();

proxy/src/control_plane/client/mock.rs

@@ -5,9 +5,8 @@ use std::sync::Arc;
 
 use futures::TryFutureExt;
 use thiserror::Error;
-use tokio_postgres::types::Json;
 use tokio_postgres::Client;
-use tracing::{error, info, warn};
+use tracing::{error, info, info_span, warn, Instrument};
 
 use crate::auth::backend::jwt::AuthRule;
 use crate::auth::backend::ComputeUserInfo;
@@ -18,7 +17,7 @@ use crate::control_plane::client::{CachedAllowedIps, CachedRoleSecret};
 use crate::control_plane::errors::{
     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
 };
-use crate::control_plane::messages::{JwksSettings, MetricsAuxInfo};
+use crate::control_plane::messages::MetricsAuxInfo;
 use crate::control_plane::{AuthInfo, AuthSecret, CachedNodeInfo, NodeInfo};
 use crate::error::io_error;
 use crate::intern::RoleNameInt;
@@ -66,70 +65,61 @@ impl MockControlPlane {
         &self,
         user_info: &ComputeUserInfo,
     ) -> Result<AuthInfo, GetAuthInfoError> {
-        // Perhaps we could persist this connection, but then we'd have to
-        // write more code for reopening it if it got closed, which doesn't
-        // seem worth it.
-        let (client, connection) =
-            tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?;
+        let (secret, allowed_ips) = async {
+            // Perhaps we could persist this connection, but then we'd have to
+            // write more code for reopening it if it got closed, which doesn't
+            // seem worth it.
+            let (client, connection) =
+                tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?;
 
-        tokio::spawn(connection);
+            tokio::spawn(connection);
 
-        let secret = if let Some(entry) = get_execute_postgres_query(
-            &client,
-            "select rolpassword from pg_catalog.pg_authid where rolname = $1",
-            &[&&*user_info.user],
-            "rolpassword",
-        )
-        .await?
-        {
-            info!("got a secret: {entry}"); // safe since it's not a prod scenario
-            let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram);
-            secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5))
-        } else {
-            warn!("user '{}' does not exist", user_info.user);
-            None
-        };
+            let secret = if let Some(entry) = get_execute_postgres_query(
+                &client,
+                "select rolpassword from pg_catalog.pg_authid where rolname = $1",
+                &[&&*user_info.user],
+                "rolpassword",
+            )
+            .await?
+            {
+                info!("got a secret: {entry}"); // safe since it's not a prod scenario
+                let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram);
+                secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5))
+            } else {
+                warn!("user '{}' does not exist", user_info.user);
+                None
+            };
 
-        let (allowed_ips, auth_rules) = if self.ip_allowlist_check_enabled {
-            let row = client.query_opt("select allowed_ips, jwks from neon_control_plane.endpoints where endpoint_id = $1", &[&user_info.endpoint.as_str()]).await?;
-            match row {
-                Some(row) => {
-                    let allowed_ips: String = row
-                        .try_get("allowed_ips")
-                        .map_err(MockApiError::PasswordNotSet)?;
-                    let jwks: Json<Vec<JwksSettings>> =
-                        row.try_get("jwks").map_err(MockApiError::PasswordNotSet)?;
-
-                    info!("got allowed_ips: {allowed_ips}");
-                    let allowed_ips = allowed_ips
-                        .split(',')
-                        .map(|s| {
-                            IpPattern::from_str(s).expect("mocked ip pattern should be correct")
-                        })
-                        .collect();
-
-                    let auth_rules = jwks
-                        .0
-                        .into_iter()
-                        .map(|jwks| AuthRule {
-                            id: jwks.id,
-                            jwks_url: jwks.jwks_url,
-                            audience: jwks.jwt_audience,
-                            role_names: jwks.role_names,
-                        })
-                        .collect();
-
-                    (allowed_ips, auth_rules)
+            let allowed_ips = if self.ip_allowlist_check_enabled {
+                match get_execute_postgres_query(
+                    &client,
+                    "select allowed_ips from neon_control_plane.endpoints where endpoint_id = $1",
+                    &[&user_info.endpoint.as_str()],
+                    "allowed_ips",
+                )
+                .await?
+                {
+                    Some(s) => {
+                        info!("got allowed_ips: {s}");
+                        s.split(',')
+                            .map(|s| {
+                                IpPattern::from_str(s).expect("mocked ip pattern should be correct")
+                            })
+                            .collect()
+                    }
+                    None => vec![],
                 }
-                None => (vec![], vec![]),
-            }
-        } else {
-            (vec![], vec![])
-        };
+            } else {
+                vec![]
+            };
 
+            Ok((secret, allowed_ips))
+        }
+        .inspect_err(|e: &GetAuthInfoError| tracing::error!("{e}"))
+        .instrument(info_span!("postgres", url = self.endpoint.as_str()))
+        .await?;
         Ok(AuthInfo {
             secret,
-            auth_rules,
             allowed_ips,
             project_id: None,
         })
@@ -213,7 +203,7 @@ async fn get_execute_postgres_query(
 }
 
 impl super::ControlPlaneApi for MockControlPlane {
-    #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))]
+    #[tracing::instrument(skip_all)]
     async fn get_role_secret(
         &self,
         _ctx: &RequestContext,
@@ -224,20 +214,19 @@ impl super::ControlPlaneApi for MockControlPlane {
         ))
     }
 
-    #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))]
     async fn get_allowed_ips_and_secret(
         &self,
         _ctx: &RequestContext,
         user_info: &ComputeUserInfo,
     ) -> Result<(CachedAllowedIps, Option<CachedRoleSecret>), GetAuthInfoError> {
-        let res = self.do_get_auth_info(user_info).await?;
         Ok((
-            Cached::new_uncached(Arc::new((res.allowed_ips, res.auth_rules))),
+            Cached::new_uncached(Arc::new(
+                self.do_get_auth_info(user_info).await?.allowed_ips,
+            )),
             None,
         ))
     }
 
-    #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))]
     async fn get_endpoint_jwks(
         &self,
         _ctx: &RequestContext,
@@ -246,7 +235,7 @@ impl super::ControlPlaneApi for MockControlPlane {
         self.do_get_endpoint_jwks(endpoint).await
     }
 
-    #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))]
+    #[tracing::instrument(skip_all)]
     async fn wake_compute(
         &self,
         _ctx: &RequestContext,

proxy/src/control_plane/messages.rs

@@ -229,8 +229,6 @@ pub(crate) struct GetEndpointAccessControl {
     pub(crate) allowed_ips: Option<Vec<IpPattern>>,
     pub(crate) project_id: Option<ProjectIdInt>,
     pub(crate) allowed_vpc_endpoint_ids: Option<Vec<EndpointIdInt>>,
-    #[serde(default)]
-    pub(crate) jwks: Vec<JwksSettings>,
 }
 
 /// Response which holds compute node's `host:port` pair.

proxy/src/control_plane/mod.rs

@@ -54,7 +54,6 @@ pub(crate) struct AuthInfo {
     pub(crate) allowed_ips: Vec<IpPattern>,
     /// Project ID. This is used for cache invalidation.
     pub(crate) project_id: Option<ProjectIdInt>,
-    pub(crate) auth_rules: Vec<AuthRule>,
 }
 
 /// Info for establishing a connection to a compute node.
@@ -100,8 +99,7 @@ pub(crate) type NodeInfoCache =
     TimedLru<EndpointCacheKey, Result<NodeInfo, Box<ControlPlaneErrorMessage>>>;
 pub(crate) type CachedNodeInfo = Cached<&'static NodeInfoCache, NodeInfo>;
 pub(crate) type CachedRoleSecret = Cached<&'static ProjectInfoCacheImpl, Option<AuthSecret>>;
-pub(crate) type CachedAllowedIps =
-    Cached<&'static ProjectInfoCacheImpl, Arc<(Vec<IpPattern>, Vec<AuthRule>)>>;
+pub(crate) type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc<Vec<IpPattern>>>;
 
 /// This will allocate per each call, but the http requests alone
 /// already require a few allocations, so it should be fine.

proxy/src/proxy/mod.rs

@@ -23,7 +23,6 @@ use tracing::{debug, error, info, warn, Instrument};
 
 use self::connect_compute::{connect_to_compute, TcpMechanism};
 use self::passthrough::ProxyPassthrough;
-use crate::auth::backend::jwt::JwkCache;
 use crate::cancellation::{self, CancellationHandlerMain, CancellationHandlerMainInternal};
 use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
 use crate::context::RequestContext;
@@ -72,8 +71,6 @@ pub async fn task_main(
     let connections = tokio_util::task::task_tracker::TaskTracker::new();
     let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
 
-    let jwks_cache = Arc::new(JwkCache::default());
-
     while let Some(accept_result) =
         run_until_cancelled(listener.accept(), &cancellation_token).await
     {
@@ -87,7 +84,6 @@ pub async fn task_main(
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
         let cancellations = cancellations.clone();
-        let jwks_cache = jwks_cache.clone();
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
         let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
@@ -140,7 +136,6 @@ pub async fn task_main(
                 endpoint_rate_limiter2,
                 conn_gauge,
                 cancellations,
-                jwks_cache,
             )
             .instrument(ctx.span())
             .boxed()
@@ -254,7 +249,6 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     conn_gauge: NumClientConnectionsGuard<'static>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
-    jwks_cache: Arc<JwkCache>,
 ) -> Result<Option<ProxyPassthrough<CancellationHandlerMainInternal, S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -325,7 +319,6 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
             mode.allow_cleartext(),
             &config.authentication_config,
             endpoint_rate_limiter,
-            jwks_cache,
         )
         .await
     {

proxy/src/serverless/backend.rs

@@ -3,9 +3,9 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ed25519_dalek::SigningKey;
 use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
-use jose_jwk::jose_b64;
+use p256::ecdsa::SigningKey;
+use p256::elliptic_curve::JwkEcKey;
 use rand::rngs::OsRng;
 use tokio::net::{lookup_host, TcpStream};
 use tracing::field::display;
@@ -57,10 +57,9 @@ impl PoolingBackend {
 
         let user_info = user_info.clone();
         let backend = self.auth_backend.as_ref().map(|()| user_info.clone());
-        let (x, maybe_secret) = backend.get_allowed_ips_and_secret(ctx).await?;
-        let (allowed_ips, _) = &**x;
+        let (allowed_ips, maybe_secret) = backend.get_allowed_ips_and_secret(ctx).await?;
         if self.config.authentication_config.ip_allowlist_check_enabled
-            && !check_peer_addr_is_in_list(&ctx.peer_addr(), allowed_ips)
+            && !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips)
         {
             return Err(AuthError::ip_address_not_allowed(ctx.peer_addr()));
         }
@@ -355,15 +354,9 @@ impl PoolingBackend {
     }
 }
 
-fn create_random_jwk() -> (SigningKey, jose_jwk::Key) {
-    let key = SigningKey::generate(&mut OsRng);
-
-    let jwk = jose_jwk::Key::Okp(jose_jwk::Okp {
-        crv: jose_jwk::OkpCurves::Ed25519,
-        x: jose_b64::serde::Bytes::from(key.verifying_key().to_bytes().to_vec()),
-        d: None,
-    });
-
+fn create_random_jwk() -> (SigningKey, JwkEcKey) {
+    let key = SigningKey::random(&mut OsRng);
+    let jwk = p256::PublicKey::from(key.verifying_key()).to_jwk();
     (key, jwk)
 }
 

proxy/src/serverless/local_conn_pool.rs

@@ -16,16 +16,17 @@ use std::sync::Arc;
 use std::task::{ready, Poll};
 use std::time::Duration;
 
-use ed25519_dalek::{Signature, Signer, SigningKey};
 use futures::future::poll_fn;
 use futures::Future;
 use indexmap::IndexMap;
 use jose_jwk::jose_b64::base64ct::{Base64UrlUnpadded, Encoding};
+use p256::ecdsa::{Signature, SigningKey};
 use parking_lot::RwLock;
 use postgres_client::tls::NoTlsStream;
 use postgres_client::types::ToSql;
 use postgres_client::AsyncMessage;
 use serde_json::value::RawValue;
+use signature::Signer;
 use tokio::net::TcpStream;
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
@@ -41,7 +42,7 @@ use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo};
 use crate::metrics::Metrics;
 
 pub(crate) const EXT_NAME: &str = "pg_session_jwt";
-pub(crate) const EXT_VERSION: &str = "0.2.0";
+pub(crate) const EXT_VERSION: &str = "0.1.2";
 pub(crate) const EXT_SCHEMA: &str = "auth";
 
 #[derive(Clone)]
@@ -338,8 +339,8 @@ fn sign_jwt(sk: &SigningKey, payload: &[u8]) -> String {
     let cap = jwt.capacity();
 
     // we only need an empty header with the alg specified.
-    // base64url(r#"{"alg":"EdDSA"}"#) == "eyJhbGciOiJFZERTQSJ9"
-    jwt.push_str("eyJhbGciOiJFZERTQSJ9.");
+    // base64url(r#"{"alg":"ES256"}"#) == "eyJhbGciOiJFUzI1NiJ9"
+    jwt.push_str("eyJhbGciOiJFUzI1NiJ9.");
 
     // encode the jwt payload in-place
     base64::encode_config_buf(payload, base64::URL_SAFE_NO_PAD, &mut jwt);
@@ -365,14 +366,14 @@ fn sign_jwt(sk: &SigningKey, payload: &[u8]) -> String {
 #[cfg(test)]
 #[expect(clippy::unwrap_used)]
 mod tests {
-    use ed25519_dalek::SigningKey;
+    use p256::ecdsa::SigningKey;
     use typed_json::json;
 
     use super::resign_jwt;
 
     #[test]
     fn jwt_token_snapshot() {
-        let key = SigningKey::from_bytes(&[1; 32]);
+        let key = SigningKey::from_bytes(&[1; 32].into()).unwrap();
         let data =
             json!({"foo":"bar","jti":"foo\nbar","nested":{"jti":"tricky nesting"}}).to_string();
 
@@ -380,17 +381,12 @@ mod tests {
 
         // To validate the JWT, copy the JWT string and paste it into https://jwt.io/.
         // In the public-key box, paste the following jwk public key
-        // `{"kty":"OKP","crv":"Ed25519","x":"iojj3XQJ8ZX9UtstPLpdcspnCb8dlBIb83SIAbQPb1w"}`
-        // Note - jwt.io doesn't support EdDSA :(
-        // https://github.com/jsonwebtoken/jsonwebtoken.github.io/issues/509
+        // `{"kty":"EC","crv":"P-256","x":"b_A7lJJBzh2t1DUZ5pYOCoW0GmmgXDKBA6orzhWUyhY","y":"PE91OlW_AdxT9sCwx-7ni0DG_30lqW4igrmJzvccFEo"}`
 
-        // let jwk = jose_jwk::Key::Okp(jose_jwk::Okp {
-        //     crv: jose_jwk::OkpCurves::Ed25519,
-        //     x: jose_jwk::jose_b64::serde::Bytes::from(key.verifying_key().to_bytes().to_vec()),
-        //     d: None,
-        // });
-        // println!("{}", serde_json::to_string(&jwk).unwrap());
+        // let pub_key = p256::ecdsa::VerifyingKey::from(&key);
+        // let pub_key = p256::PublicKey::from(pub_key);
+        // println!("{}", pub_key.to_jwk_string());
 
-        assert_eq!(jwt, "eyJhbGciOiJFZERTQSJ9.eyJmb28iOiJiYXIiLCJqdGkiOjIsIm5lc3RlZCI6eyJqdGkiOiJ0cmlja3kgbmVzdGluZyJ9fQ.Cvyc2By33KI0f0obystwdy8PN111L3Sc9_Mr2CU3XshtSqSdxuRxNEZGbb_RvyJf2IzheC_s7aBZ-jLeQ9N0Bg");
+        assert_eq!(jwt, "eyJhbGciOiJFUzI1NiJ9.eyJmb28iOiJiYXIiLCJqdGkiOjIsIm5lc3RlZCI6eyJqdGkiOiJ0cmlja3kgbmVzdGluZyJ9fQ.pYf0LxoJ8sDgpmsYOgrbNecOSipnPBEGwnZzB-JhW2cONrKlqRsgXwK8_cOsyolGy-hTTe8GXbWTl_UdpF5RyA");
     }
 }

proxy/src/serverless/mod.rs

@@ -43,7 +43,6 @@ use tokio_util::task::TaskTracker;
 use tracing::{info, warn, Instrument};
 use utils::http::error::ApiError;
 
-use crate::auth::backend::jwt::JwkCache;
 use crate::cancellation::CancellationHandlerMain;
 use crate::config::{ProxyConfig, ProxyProtocolV2};
 use crate::context::RequestContext;
@@ -332,8 +331,6 @@ async fn connection_handler(
     let http_cancellation_token = CancellationToken::new();
     let _cancel_connection = http_cancellation_token.clone().drop_guard();
 
-    let jwks_cache = Arc::new(JwkCache::default());
-
     let conn_info2 = conn_info.clone();
     let server = Builder::new(TokioExecutor::new());
     let conn = server.serve_connection_with_upgrades(
@@ -374,7 +371,6 @@ async fn connection_handler(
                     http_request_token,
                     endpoint_rate_limiter.clone(),
                     cancellations,
-                    jwks_cache.clone(),
                 )
                 .in_current_span()
                 .map_ok_or_else(api_error_into_response, |r| r),
@@ -423,7 +419,6 @@ async fn request_handler(
     http_cancellation_token: CancellationToken,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellations: TaskTracker,
-    jwks_cache: Arc<JwkCache>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
     let host = request
         .headers()
@@ -461,7 +456,6 @@ async fn request_handler(
                     endpoint_rate_limiter,
                     host,
                     cancellations,
-                    jwks_cache,
                 )
                 .await
                 {

proxy/src/serverless/websocket.rs

@@ -12,7 +12,6 @@ use pin_project_lite::pin_project;
 use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
 use tracing::warn;
 
-use crate::auth::backend::jwt::JwkCache;
 use crate::cancellation::CancellationHandlerMain;
 use crate::config::ProxyConfig;
 use crate::context::RequestContext;
@@ -134,7 +133,6 @@ pub(crate) async fn serve_websocket(
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     hostname: Option<String>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
-    jwks_cache: Arc<JwkCache>,
 ) -> anyhow::Result<()> {
     let websocket = websocket.await?;
     let websocket = WebSocketServer::after_handshake(TokioIo::new(websocket));
@@ -154,7 +152,6 @@ pub(crate) async fn serve_websocket(
         endpoint_rate_limiter,
         conn_gauge,
         cancellations,
-        jwks_cache,
     ))
     .await;
 

storage_controller/Cargo.toml

@@ -55,4 +55,4 @@ r2d2 = { version = "0.8.10" }
 utils = { path = "../libs/utils/" }
 metrics = { path = "../libs/metrics/" }
 control_plane = { path = "../control_plane" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+workspace_hack = { version = "0.1", path = "../workspace_hack" }

storage_controller/src/compute_hook.rs

@@ -124,10 +124,7 @@ impl ComputeHookTenant {
                 if let Some(shard_idx) = shard_idx {
                     sharded.shards.remove(shard_idx);
                 } else {
-                    // This is a valid but niche case, where the tenant was previously attached
-                    // as a Secondary location and then detached, so has no previously notified
-                    // state.
-                    tracing::info!("Shard not found while handling detach")
+                    tracing::warn!("Shard not found while handling detach")
                 }
             }
             ComputeHookTenant::Unsharded(_) => {
@@ -764,10 +761,7 @@ impl ComputeHook {
         let mut state_locked = self.state.lock().unwrap();
         match state_locked.entry(tenant_shard_id.tenant_id) {
             Entry::Vacant(_) => {
-                // This is a valid but niche case, where the tenant was previously attached
-                // as a Secondary location and then detached, so has no previously notified
-                // state.
-                tracing::info!("Compute hook tenant not found for detach");
+                tracing::warn!("Compute hook tenant not found for detach");
             }
             Entry::Occupied(mut e) => {
                 let sharded = e.get().is_sharded();

storage_controller/src/drain_utils.rs

@@ -112,7 +112,7 @@ impl TenantShardDrain {
             }
         }
 
-        match tenant_shard.preferred_secondary(scheduler) {
+        match scheduler.node_preferred(tenant_shard.intent.get_secondary()) {
             Some(node) => Some(node),
             None => {
                 tracing::warn!(

storage_controller/src/http.rs

@@ -690,8 +690,7 @@ async fn handle_node_list(req: Request<Body>) -> Result<Response<Body>, ApiError
     };
 
     let state = get_state(&req);
-    let mut nodes = state.service.node_list().await?;
-    nodes.sort_by_key(|n| n.get_id());
+    let nodes = state.service.node_list().await?;
     let api_nodes = nodes.into_iter().map(|n| n.describe()).collect::<Vec<_>>();
 
     json_response(StatusCode::OK, api_nodes)
@@ -1006,29 +1005,6 @@ async fn handle_tenant_shard_migrate(
     )
 }
 
-async fn handle_tenant_shard_migrate_secondary(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let mut req = match maybe_forward(req).await {
-        ForwardOutcome::Forwarded(res) => {
-            return res;
-        }
-        ForwardOutcome::NotForwarded(req) => req,
-    };
-
-    let tenant_shard_id: TenantShardId = parse_request_param(&req, "tenant_shard_id")?;
-    let migrate_req = json_request::<TenantShardMigrateRequest>(&mut req).await?;
-    json_response(
-        StatusCode::OK,
-        service
-            .tenant_shard_migrate_secondary(tenant_shard_id, migrate_req)
-            .await?,
-    )
-}
-
 async fn handle_tenant_shard_cancel_reconcile(
     service: Arc<Service>,
     req: Request<Body>,
@@ -1879,16 +1855,6 @@ pub fn make_router(
                 RequestName("control_v1_tenant_migrate"),
             )
         })
-        .put(
-            "/control/v1/tenant/:tenant_shard_id/migrate_secondary",
-            |r| {
-                tenant_service_handler(
-                    r,
-                    handle_tenant_shard_migrate_secondary,
-                    RequestName("control_v1_tenant_migrate_secondary"),
-                )
-            },
-        )
         .put(
             "/control/v1/tenant/:tenant_shard_id/cancel_reconcile",
             |r| {

storage_controller/src/reconciler.rs

@@ -826,21 +826,7 @@ impl Reconciler {
             if self.cancel.is_cancelled() {
                 return Err(ReconcileError::Cancel);
             }
-            // We only try to configure secondary locations if the node is available.  This does
-            // not stop us succeeding with the reconcile, because our core goal is to make the
-            // shard _available_ (the attached location), and configuring secondary locations
-            // can be done lazily when the node becomes available (via background reconciliation).
-            if node.is_available() {
-                self.location_config(&node, conf, None, false).await?;
-            } else {
-                // If the node is unavailable, we skip and consider the reconciliation successful: this
-                // is a common case where a pageserver is marked unavailable: we demote a location on
-                // that unavailable pageserver to secondary.
-                tracing::info!("Skipping configuring secondary location {node}, it is unavailable");
-                self.observed
-                    .locations
-                    .insert(node.get_id(), ObservedStateLocation { conf: None });
-            }
+            self.location_config(&node, conf, None, false).await?;
         }
 
         // The condition below identifies a detach. We must have no attached intent and

storage_controller/src/scheduler.rs

@@ -32,9 +32,6 @@ pub(crate) struct SchedulerNode {
     shard_count: usize,
     /// How many shards are currently attached on this node, via their [`crate::tenant_shard::IntentState`].
     attached_shard_count: usize,
-    /// How many shards have a location on this node (via [`crate::tenant_shard::IntentState`]) _and_ this node
-    /// is in their preferred AZ (i.e. this is their 'home' location)
-    home_shard_count: usize,
     /// Availability zone id in which the node resides
     az: AvailabilityZone,
 
@@ -50,12 +47,6 @@ pub(crate) trait NodeSchedulingScore: Debug + Ord + Copy + Sized {
         preferred_az: &Option<AvailabilityZone>,
         context: &ScheduleContext,
     ) -> Option<Self>;
-
-    /// Return a score that drops any components based on node utilization: this is useful
-    /// for finding scores for scheduling optimisation, when we want to avoid rescheduling
-    /// shards due to e.g. disk usage, to avoid flapping.
-    fn for_optimization(&self) -> Self;
-
     fn is_overloaded(&self) -> bool;
     fn node_id(&self) -> NodeId;
 }
@@ -145,13 +136,17 @@ impl PartialOrd for SecondaryAzMatch {
 /// Ordering is given by member declaration order (top to bottom).
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
 pub(crate) struct NodeAttachmentSchedulingScore {
+    /// The number of shards belonging to the tenant currently being
+    /// scheduled that are attached to this node.
+    affinity_score: AffinityScore,
     /// Flag indicating whether this node matches the preferred AZ
     /// of the shard. For equal affinity scores, nodes in the matching AZ
     /// are considered first.
     az_match: AttachmentAzMatch,
-    /// The number of shards belonging to the tenant currently being
-    /// scheduled that are attached to this node.
-    affinity_score: AffinityScore,
+    /// Size of [`ScheduleContext::attached_nodes`] for the current node.
+    /// This normally tracks the number of attached shards belonging to the
+    /// tenant being scheduled that are already on this node.
+    attached_shards_in_context: usize,
     /// Utilisation score that combines shard count and disk utilisation
     utilization_score: u64,
     /// Total number of shards attached to this node. When nodes have identical utilisation, this
@@ -182,25 +177,13 @@ impl NodeSchedulingScore for NodeAttachmentSchedulingScore {
                 .copied()
                 .unwrap_or(AffinityScore::FREE),
             az_match: AttachmentAzMatch(AzMatch::new(&node.az, preferred_az.as_ref())),
+            attached_shards_in_context: context.attached_nodes.get(node_id).copied().unwrap_or(0),
             utilization_score: utilization.cached_score(),
             total_attached_shard_count: node.attached_shard_count,
             node_id: *node_id,
         })
     }
 
-    /// For use in scheduling optimisation, where we only want to consider the aspects
-    /// of the score that can only be resolved by moving things (such as inter-shard affinity
-    /// and AZ affinity), and ignore aspects that reflect the total utilization of a node (which
-    /// can fluctuate for other reasons)
-    fn for_optimization(&self) -> Self {
-        Self {
-            utilization_score: 0,
-            total_attached_shard_count: 0,
-            node_id: NodeId(0),
-            ..*self
-        }
-    }
-
     fn is_overloaded(&self) -> bool {
         PageserverUtilization::is_overloaded(self.utilization_score)
     }
@@ -225,9 +208,9 @@ pub(crate) struct NodeSecondarySchedulingScore {
     affinity_score: AffinityScore,
     /// Utilisation score that combines shard count and disk utilisation
     utilization_score: u64,
-    /// Anti-affinity with other non-home locations: this gives the behavior that secondaries
-    /// will spread out across the nodes in an AZ.
-    total_non_home_shard_count: usize,
+    /// Total number of shards attached to this node. When nodes have identical utilisation, this
+    /// acts as an anti-affinity between attached shards.
+    total_attached_shard_count: usize,
     /// Convenience to make selection deterministic in tests and empty systems
     node_id: NodeId,
 }
@@ -254,20 +237,11 @@ impl NodeSchedulingScore for NodeSecondarySchedulingScore {
                 .copied()
                 .unwrap_or(AffinityScore::FREE),
             utilization_score: utilization.cached_score(),
-            total_non_home_shard_count: (node.shard_count - node.home_shard_count),
+            total_attached_shard_count: node.attached_shard_count,
             node_id: *node_id,
         })
     }
 
-    fn for_optimization(&self) -> Self {
-        Self {
-            utilization_score: 0,
-            total_non_home_shard_count: 0,
-            node_id: NodeId(0),
-            ..*self
-        }
-    }
-
     fn is_overloaded(&self) -> bool {
         PageserverUtilization::is_overloaded(self.utilization_score)
     }
@@ -319,10 +293,6 @@ impl AffinityScore {
     pub(crate) fn inc(&mut self) {
         self.0 += 1;
     }
-
-    pub(crate) fn dec(&mut self) {
-        self.0 -= 1;
-    }
 }
 
 impl std::ops::Add for AffinityScore {
@@ -354,6 +324,9 @@ pub(crate) struct ScheduleContext {
     /// Sparse map of nodes: omitting a node implicitly makes its affinity [`AffinityScore::FREE`]
     pub(crate) nodes: HashMap<NodeId, AffinityScore>,
 
+    /// Specifically how many _attached_ locations are on each node
+    pub(crate) attached_nodes: HashMap<NodeId, usize>,
+
     pub(crate) mode: ScheduleMode,
 }
 
@@ -361,6 +334,7 @@ impl ScheduleContext {
     pub(crate) fn new(mode: ScheduleMode) -> Self {
         Self {
             nodes: HashMap::new(),
+            attached_nodes: HashMap::new(),
             mode,
         }
     }
@@ -374,31 +348,25 @@ impl ScheduleContext {
         }
     }
 
-    /// Remove `shard`'s contributions to this context.  This is useful when considering scheduling
-    /// this shard afresh, where we don't want it to e.g. experience anti-affinity to its current location.
-    pub(crate) fn project_detach(&self, shard: &TenantShard) -> Self {
-        let mut new_context = self.clone();
-
-        if let Some(attached) = shard.intent.get_attached() {
-            if let Some(score) = new_context.nodes.get_mut(attached) {
-                score.dec();
-            }
-        }
-
-        for secondary in shard.intent.get_secondary() {
-            if let Some(score) = new_context.nodes.get_mut(secondary) {
-                score.dec();
-            }
-        }
-
-        new_context
+    pub(crate) fn push_attached(&mut self, node_id: NodeId) {
+        let entry = self.attached_nodes.entry(node_id).or_default();
+        *entry += 1;
+    }
+
+    pub(crate) fn get_node_affinity(&self, node_id: NodeId) -> AffinityScore {
+        self.nodes
+            .get(&node_id)
+            .copied()
+            .unwrap_or(AffinityScore::FREE)
+    }
+
+    pub(crate) fn get_node_attachments(&self, node_id: NodeId) -> usize {
+        self.attached_nodes.get(&node_id).copied().unwrap_or(0)
     }
 
-    /// For test, track the sum of AffinityScore values, which is effectively how many
-    /// attached or secondary locations have been registered with this context.
     #[cfg(test)]
-    pub(crate) fn location_count(&self) -> usize {
-        self.nodes.values().map(|i| i.0).sum()
+    pub(crate) fn attach_count(&self) -> usize {
+        self.attached_nodes.values().sum()
     }
 }
 
@@ -420,7 +388,6 @@ impl Scheduler {
                 SchedulerNode {
                     shard_count: 0,
                     attached_shard_count: 0,
-                    home_shard_count: 0,
                     may_schedule: node.may_schedule(),
                     az: node.get_availability_zone_id().clone(),
                 },
@@ -448,7 +415,6 @@ impl Scheduler {
                 SchedulerNode {
                     shard_count: 0,
                     attached_shard_count: 0,
-                    home_shard_count: 0,
                     may_schedule: node.may_schedule(),
                     az: node.get_availability_zone_id().clone(),
                 },
@@ -461,9 +427,6 @@ impl Scheduler {
                     Some(node) => {
                         node.shard_count += 1;
                         node.attached_shard_count += 1;
-                        if Some(&node.az) == shard.preferred_az() {
-                            node.home_shard_count += 1;
-                        }
                     }
                     None => anyhow::bail!(
                         "Tenant {} references nonexistent node {}",
@@ -475,12 +438,7 @@ impl Scheduler {
 
             for node_id in shard.intent.get_secondary() {
                 match expect_nodes.get_mut(node_id) {
-                    Some(node) => {
-                        node.shard_count += 1;
-                        if Some(&node.az) == shard.preferred_az() {
-                            node.home_shard_count += 1;
-                        }
-                    }
+                    Some(node) => node.shard_count += 1,
                     None => anyhow::bail!(
                         "Tenant {} references nonexistent node {}",
                         shard.tenant_shard_id,
@@ -524,20 +482,13 @@ impl Scheduler {
     ///
     /// It is an error to call this for a node that is not known to the scheduler (i.e. passed into
     /// [`Self::new`] or [`Self::node_upsert`])
-    pub(crate) fn update_node_ref_counts(
-        &mut self,
-        node_id: NodeId,
-        preferred_az: Option<&AvailabilityZone>,
-        update: RefCountUpdate,
-    ) {
+    pub(crate) fn update_node_ref_counts(&mut self, node_id: NodeId, update: RefCountUpdate) {
         let Some(node) = self.nodes.get_mut(&node_id) else {
             debug_assert!(false);
             tracing::error!("Scheduler missing node {node_id}");
             return;
         };
 
-        let is_home_az = Some(&node.az) == preferred_az;
-
         match update {
             RefCountUpdate::PromoteSecondary => {
                 node.attached_shard_count += 1;
@@ -545,31 +496,19 @@ impl Scheduler {
             RefCountUpdate::Attach => {
                 node.shard_count += 1;
                 node.attached_shard_count += 1;
-                if is_home_az {
-                    node.home_shard_count += 1;
-                }
             }
             RefCountUpdate::Detach => {
                 node.shard_count -= 1;
                 node.attached_shard_count -= 1;
-                if is_home_az {
-                    node.home_shard_count -= 1;
-                }
             }
             RefCountUpdate::DemoteAttached => {
                 node.attached_shard_count -= 1;
             }
             RefCountUpdate::AddSecondary => {
                 node.shard_count += 1;
-                if is_home_az {
-                    node.home_shard_count += 1;
-                }
             }
             RefCountUpdate::RemoveSecondary => {
                 node.shard_count -= 1;
-                if is_home_az {
-                    node.home_shard_count -= 1;
-                }
             }
         }
 
@@ -655,7 +594,6 @@ impl Scheduler {
                 entry.insert(SchedulerNode {
                     shard_count: 0,
                     attached_shard_count: 0,
-                    home_shard_count: 0,
                     may_schedule: node.may_schedule(),
                     az: node.get_availability_zone_id().clone(),
                 });
@@ -669,20 +607,33 @@ impl Scheduler {
         }
     }
 
-    /// Calculate a single node's score, used in optimizer logic to compare specific
-    /// nodes' scores.
-    pub(crate) fn compute_node_score<Score>(
-        &mut self,
-        node_id: NodeId,
-        preferred_az: &Option<AvailabilityZone>,
-        context: &ScheduleContext,
-    ) -> Option<Score>
-    where
-        Score: NodeSchedulingScore,
-    {
-        self.nodes
-            .get_mut(&node_id)
-            .and_then(|node| Score::generate(&node_id, node, preferred_az, context))
+    /// Where we have several nodes to choose from, for example when picking a secondary location
+    /// to promote to an attached location, this method may be used to pick the best choice based
+    /// on the scheduler's knowledge of utilization and availability.
+    ///
+    /// If the input is empty, or all the nodes are not elegible for scheduling, return None: the
+    /// caller can pick a node some other way.
+    pub(crate) fn node_preferred(&self, nodes: &[NodeId]) -> Option<NodeId> {
+        if nodes.is_empty() {
+            return None;
+        }
+
+        // TODO: When the utilization score returned by the pageserver becomes meaningful,
+        // schedule based on that instead of the shard count.
+        let node = nodes
+            .iter()
+            .map(|node_id| {
+                let may_schedule = self
+                    .nodes
+                    .get(node_id)
+                    .map(|n| !matches!(n.may_schedule, MaySchedule::No))
+                    .unwrap_or(false);
+                (*node_id, may_schedule)
+            })
+            .max_by_key(|(_n, may_schedule)| *may_schedule);
+
+        // If even the preferred node has may_schedule==false, return None
+        node.and_then(|(node_id, may_schedule)| if may_schedule { Some(node_id) } else { None })
     }
 
     /// Compute a schedulling score for each node that the scheduler knows of
@@ -776,7 +727,7 @@ impl Scheduler {
             tracing::info!(
             "scheduler selected node {node_id} (elegible nodes {:?}, hard exclude: {hard_exclude:?}, soft exclude: {context:?})",
             scores.iter().map(|i| i.node_id().0).collect::<Vec<_>>()
-       );
+        );
         }
 
         // Note that we do not update shard count here to reflect the scheduling: that
@@ -792,74 +743,47 @@ impl Scheduler {
     }
 
     /// For choosing which AZ to schedule a new shard into, use this.  It will return the
-    /// AZ with the the lowest number of shards currently scheduled in this AZ as their home
-    /// location.
+    /// AZ with the lowest median utilization.
     ///
     /// We use an AZ-wide measure rather than simply selecting the AZ of the least-loaded
     /// node, because while tenants start out single sharded, when they grow and undergo
-    /// shard-split, they will occupy space on many nodes within an AZ.  It is important
-    /// that we pick the AZ in a way that balances this _future_ load.
+    /// shard-split, they will occupy space on many nodes within an AZ.
     ///
-    /// Once we've picked an AZ, subsequent scheduling within that AZ will be driven by
-    /// nodes' utilization scores.
+    /// We use median rather than total free space or mean utilization, because
+    /// we wish to avoid preferring AZs that have low-load nodes resulting from
+    /// recent replacements.
+    ///
+    /// The practical result is that we will pick an AZ based on its median node, and
+    /// then actually _schedule_ the new shard onto the lowest-loaded node in that AZ.
     pub(crate) fn get_az_for_new_tenant(&self) -> Option<AvailabilityZone> {
         if self.nodes.is_empty() {
             return None;
         }
 
-        #[derive(Default)]
-        struct AzScore {
-            home_shard_count: usize,
-            scheduleable: bool,
+        let mut scores_by_az = HashMap::new();
+        for (node_id, node) in &self.nodes {
+            let az_scores = scores_by_az.entry(&node.az).or_insert_with(Vec::new);
+            let score = match &node.may_schedule {
+                MaySchedule::Yes(utilization) => utilization.score(),
+                MaySchedule::No => PageserverUtilization::full().score(),
+            };
+            az_scores.push((node_id, node, score));
         }
 
-        let mut azs: HashMap<&AvailabilityZone, AzScore> = HashMap::new();
-        for node in self.nodes.values() {
-            let az = azs.entry(&node.az).or_default();
-            az.home_shard_count += node.home_shard_count;
-            az.scheduleable |= matches!(node.may_schedule, MaySchedule::Yes(_));
+        // Sort by utilization.  Also include the node ID to break ties.
+        for scores in scores_by_az.values_mut() {
+            scores.sort_by_key(|i| (i.2, i.0));
         }
 
-        // If any AZs are schedulable, then filter out the non-schedulable ones (i.e. AZs where
-        // all nodes are overloaded or otherwise unschedulable).
-        if azs.values().any(|i| i.scheduleable) {
-            azs.retain(|_, i| i.scheduleable);
-        }
-
-        // Find the AZ with the lowest number of shards currently allocated
-        Some(
-            azs.into_iter()
-                .min_by_key(|i| (i.1.home_shard_count, i.0))
-                .unwrap()
-                .0
-                .clone(),
-        )
-    }
-
-    pub(crate) fn get_node_az(&self, node_id: &NodeId) -> Option<AvailabilityZone> {
-        self.nodes.get(node_id).map(|n| n.az.clone())
-    }
-
-    /// For use when choosing a preferred secondary location: filter out nodes that are not
-    /// available, and gather their AZs.
-    pub(crate) fn filter_usable_nodes(
-        &self,
-        nodes: &[NodeId],
-    ) -> Vec<(NodeId, Option<AvailabilityZone>)> {
-        nodes
+        let mut median_by_az = scores_by_az
             .iter()
-            .filter_map(|node_id| {
-                let node = self
-                    .nodes
-                    .get(node_id)
-                    .expect("Referenced nodes always exist");
-                if matches!(node.may_schedule, MaySchedule::Yes(_)) {
-                    Some((*node_id, Some(node.az.clone())))
-                } else {
-                    None
-                }
-            })
-            .collect()
+            .map(|(az, nodes)| (*az, nodes.get(nodes.len() / 2).unwrap().2))
+            .collect::<Vec<_>>();
+        // Sort by utilization.  Also include the AZ to break ties.
+        median_by_az.sort_by_key(|i| (i.1, i.0));
+
+        // Return the AZ with the lowest median utilization
+        Some(median_by_az.first().unwrap().0.clone())
     }
 
     /// Unit test access to internal state
@@ -919,14 +843,7 @@ pub(crate) mod test_utils {
 
 #[cfg(test)]
 mod tests {
-    use pageserver_api::{
-        controller_api::NodeAvailability, models::utilization::test_utilization,
-        shard::ShardIdentity,
-    };
-    use utils::{
-        id::TenantId,
-        shard::{ShardCount, ShardNumber, TenantShardId},
-    };
+    use pageserver_api::{controller_api::NodeAvailability, models::utilization::test_utilization};
 
     use super::*;
 
@@ -936,8 +853,8 @@ mod tests {
         let nodes = test_utils::make_test_nodes(2, &[]);
 
         let mut scheduler = Scheduler::new(nodes.values());
-        let mut t1_intent = IntentState::new(None);
-        let mut t2_intent = IntentState::new(None);
+        let mut t1_intent = IntentState::new();
+        let mut t2_intent = IntentState::new();
 
         let context = ScheduleContext::default();
 
@@ -1013,7 +930,7 @@ mod tests {
             let scheduled = scheduler
                 .schedule_shard::<AttachedShardTag>(&[], &None, context)
                 .unwrap();
-            let mut intent = IntentState::new(None);
+            let mut intent = IntentState::new();
             intent.set_attached(scheduler, Some(scheduled));
             scheduled_intents.push(intent);
             assert_eq!(scheduled, expect_node);
@@ -1146,7 +1063,7 @@ mod tests {
             let scheduled = scheduler
                 .schedule_shard::<Tag>(&[], &preferred_az, context)
                 .unwrap();
-            let mut intent = IntentState::new(preferred_az.clone());
+            let mut intent = IntentState::new();
             intent.set_attached(scheduler, Some(scheduled));
             scheduled_intents.push(intent);
             assert_eq!(scheduled, expect_node);
@@ -1172,9 +1089,9 @@ mod tests {
             &mut context,
         );
 
-        // Node 1 and 3 (az-a) have same affinity score, so prefer the lowest node id.
+        // Node 2 is not in "az-a", but it has the lowest affinity so we prefer that.
         assert_scheduler_chooses::<AttachedShardTag>(
-            NodeId(1),
+            NodeId(2),
             Some(az_a_tag.clone()),
             &mut scheduled_intents,
             &mut scheduler,
@@ -1190,6 +1107,26 @@ mod tests {
             &mut context,
         );
 
+        // Avoid nodes in "az-b" for the secondary location.
+        // Nodes 1 and 3 are identically loaded, so prefer the lowest node id.
+        assert_scheduler_chooses::<SecondaryShardTag>(
+            NodeId(1),
+            Some(az_b_tag.clone()),
+            &mut scheduled_intents,
+            &mut scheduler,
+            &mut context,
+        );
+
+        // Avoid nodes in "az-b" for the secondary location.
+        // Node 3 has lower affinity score than 1, so prefer that.
+        assert_scheduler_chooses::<SecondaryShardTag>(
+            NodeId(3),
+            Some(az_b_tag.clone()),
+            &mut scheduled_intents,
+            &mut scheduler,
+            &mut context,
+        );
+
         for mut intent in scheduled_intents {
             intent.clear(&mut scheduler);
         }
@@ -1213,292 +1150,34 @@ mod tests {
 
         let mut scheduler = Scheduler::new(nodes.values());
 
-        /// Force the `home_shard_count` of a node directly: this is the metric used
-        /// by the scheduler when picking AZs.
-        fn set_shard_count(scheduler: &mut Scheduler, node_id: NodeId, shard_count: usize) {
-            let node = scheduler.nodes.get_mut(&node_id).unwrap();
-            node.home_shard_count = shard_count;
+        /// Force the utilization of a node in Scheduler's state to a particular
+        /// number of bytes used.
+        fn set_utilization(scheduler: &mut Scheduler, node_id: NodeId, shard_count: u32) {
+            let mut node = Node::new(
+                node_id,
+                "".to_string(),
+                0,
+                "".to_string(),
+                0,
+                scheduler.nodes.get(&node_id).unwrap().az.clone(),
+            );
+            node.set_availability(NodeAvailability::Active(test_utilization::simple(
+                shard_count,
+                0,
+            )));
+            scheduler.node_upsert(&node);
         }
 
         // Initial empty state.  Scores are tied, scheduler prefers lower AZ ID.
         assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_a_tag.clone()));
 
-        // Home shard count is higher in AZ A, so AZ B will be preferred
-        set_shard_count(&mut scheduler, NodeId(1), 10);
-        assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_b_tag.clone()));
-
-        // Total home shard count is higher in AZ B, so we revert to preferring AZ A
-        set_shard_count(&mut scheduler, NodeId(4), 6);
-        set_shard_count(&mut scheduler, NodeId(5), 6);
+        // Put some utilization on one node in AZ A: this should change nothing, as the median hasn't changed
+        set_utilization(&mut scheduler, NodeId(1), 1000000);
         assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_a_tag.clone()));
-    }
 
-    /// Test that when selecting AZs for many new tenants, we get the expected balance across nodes
-    #[test]
-    fn az_selection_many() {
-        let az_a_tag = AvailabilityZone("az-a".to_string());
-        let az_b_tag = AvailabilityZone("az-b".to_string());
-        let az_c_tag = AvailabilityZone("az-c".to_string());
-        let nodes = test_utils::make_test_nodes(
-            6,
-            &[
-                az_a_tag.clone(),
-                az_b_tag.clone(),
-                az_c_tag.clone(),
-                az_a_tag.clone(),
-                az_b_tag.clone(),
-                az_c_tag.clone(),
-            ],
-        );
-
-        let mut scheduler = Scheduler::new(nodes.values());
-
-        // We should get 1/6th of these on each node, give or take a few...
-        let total_tenants = 300;
-
-        // ...where the 'few' is the number of AZs, because the scheduling will sometimes overshoot
-        // on one AZ before correcting itself.  This is because we select the 'home' AZ based on
-        // an AZ-wide metric, but we select the location for secondaries on a purely node-based
-        // metric (while excluding the home AZ).
-        let grace = 3;
-
-        let mut scheduled_shards = Vec::new();
-        for _i in 0..total_tenants {
-            let preferred_az = scheduler.get_az_for_new_tenant().unwrap();
-
-            let mut node_home_counts = scheduler
-                .nodes
-                .iter()
-                .map(|(node_id, node)| (node_id, node.home_shard_count))
-                .collect::<Vec<_>>();
-            node_home_counts.sort_by_key(|i| i.0);
-            eprintln!("Selected {}, vs nodes {:?}", preferred_az, node_home_counts);
-
-            let tenant_shard_id = TenantShardId {
-                tenant_id: TenantId::generate(),
-                shard_number: ShardNumber(0),
-                shard_count: ShardCount(1),
-            };
-
-            let shard_identity = ShardIdentity::new(
-                tenant_shard_id.shard_number,
-                tenant_shard_id.shard_count,
-                pageserver_api::shard::ShardStripeSize(1),
-            )
-            .unwrap();
-            let mut shard = TenantShard::new(
-                tenant_shard_id,
-                shard_identity,
-                pageserver_api::controller_api::PlacementPolicy::Attached(1),
-                Some(preferred_az),
-            );
-
-            let mut context = ScheduleContext::default();
-            shard.schedule(&mut scheduler, &mut context).unwrap();
-            eprintln!("Scheduled shard at {:?}", shard.intent);
-
-            scheduled_shards.push(shard);
-        }
-
-        for (node_id, node) in &scheduler.nodes {
-            eprintln!(
-                "Node {}: {} {} {}",
-                node_id, node.shard_count, node.attached_shard_count, node.home_shard_count
-            );
-        }
-
-        for node in scheduler.nodes.values() {
-            assert!((node.home_shard_count as i64 - total_tenants as i64 / 6).abs() < grace);
-        }
-
-        for mut shard in scheduled_shards {
-            shard.intent.clear(&mut scheduler);
-        }
-    }
-
-    #[test]
-    /// Make sure that when we have an odd number of nodes and an even number of shards, we still
-    /// get scheduling stability.
-    fn odd_nodes_stability() {
-        let az_a = AvailabilityZone("az-a".to_string());
-        let az_b = AvailabilityZone("az-b".to_string());
-
-        let nodes = test_utils::make_test_nodes(
-            10,
-            &[
-                az_a.clone(),
-                az_a.clone(),
-                az_a.clone(),
-                az_a.clone(),
-                az_a.clone(),
-                az_b.clone(),
-                az_b.clone(),
-                az_b.clone(),
-                az_b.clone(),
-                az_b.clone(),
-            ],
-        );
-        let mut scheduler = Scheduler::new(nodes.values());
-
-        // Need to keep these alive because they contribute to shard counts via RAII
-        let mut scheduled_shards = Vec::new();
-
-        let mut context = ScheduleContext::default();
-
-        fn schedule_shard(
-            tenant_shard_id: TenantShardId,
-            expect_attached: NodeId,
-            expect_secondary: NodeId,
-            scheduled_shards: &mut Vec<TenantShard>,
-            scheduler: &mut Scheduler,
-            preferred_az: Option<AvailabilityZone>,
-            context: &mut ScheduleContext,
-        ) {
-            let shard_identity = ShardIdentity::new(
-                tenant_shard_id.shard_number,
-                tenant_shard_id.shard_count,
-                pageserver_api::shard::ShardStripeSize(1),
-            )
-            .unwrap();
-            let mut shard = TenantShard::new(
-                tenant_shard_id,
-                shard_identity,
-                pageserver_api::controller_api::PlacementPolicy::Attached(1),
-                preferred_az,
-            );
-
-            shard.schedule(scheduler, context).unwrap();
-
-            assert_eq!(shard.intent.get_attached().unwrap(), expect_attached);
-            assert_eq!(
-                shard.intent.get_secondary().first().unwrap(),
-                &expect_secondary
-            );
-
-            scheduled_shards.push(shard);
-        }
-
-        let tenant_id = TenantId::generate();
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(0),
-                shard_count: ShardCount(8),
-            },
-            NodeId(1),
-            NodeId(6),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(1),
-                shard_count: ShardCount(8),
-            },
-            NodeId(2),
-            NodeId(7),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(2),
-                shard_count: ShardCount(8),
-            },
-            NodeId(3),
-            NodeId(8),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(3),
-                shard_count: ShardCount(8),
-            },
-            NodeId(4),
-            NodeId(9),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(4),
-                shard_count: ShardCount(8),
-            },
-            NodeId(5),
-            NodeId(10),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(5),
-                shard_count: ShardCount(8),
-            },
-            NodeId(1),
-            NodeId(6),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(6),
-                shard_count: ShardCount(8),
-            },
-            NodeId(2),
-            NodeId(7),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        schedule_shard(
-            TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(7),
-                shard_count: ShardCount(8),
-            },
-            NodeId(3),
-            NodeId(8),
-            &mut scheduled_shards,
-            &mut scheduler,
-            Some(az_a.clone()),
-            &mut context,
-        );
-
-        // Assert that the optimizer suggests nochanges, i.e. our initial scheduling was stable.
-        for shard in &scheduled_shards {
-            assert_eq!(shard.optimize_attachment(&mut scheduler, &context), None);
-        }
-
-        for mut shard in scheduled_shards {
-            shard.intent.clear(&mut scheduler);
-        }
+        // Put some utilization on a second node in AZ A: now the median has changed, so the scheduler
+        // should prefer the other AZ.
+        set_utilization(&mut scheduler, NodeId(2), 1000000);
+        assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_b_tag.clone()));
     }
 }

storage_controller/src/service.rs

@@ -1404,11 +1404,7 @@ impl Service {
 
             // We will populate intent properly later in [`Self::startup_reconcile`], initially populate
             // it with what we can infer: the node for which a generation was most recently issued.
-            let mut intent = IntentState::new(
-                tsp.preferred_az_id
-                    .as_ref()
-                    .map(|az| AvailabilityZone(az.clone())),
-            );
+            let mut intent = IntentState::new();
             if let Some(generation_pageserver) = tsp.generation_pageserver.map(|n| NodeId(n as u64))
             {
                 if nodes.contains_key(&generation_pageserver) {
@@ -2478,30 +2474,19 @@ impl Service {
         tenant_id: TenantId,
         _guard: &TracingExclusiveGuard<TenantOperations>,
     ) -> Result<(), ApiError> {
-        // Check if the tenant is present in memory, and select an AZ to use when loading
-        // if we will load it.
-        let load_in_az = {
+        let present_in_memory = {
             let locked = self.inner.read().unwrap();
-            let existing = locked
+            locked
                 .tenants
                 .range(TenantShardId::tenant_range(tenant_id))
-                .next();
-
-            // If the tenant is not present in memory, we expect to load it from database,
-            // so let's figure out what AZ to load it into while we have self.inner locked.
-            if existing.is_none() {
-                locked
-                    .scheduler
-                    .get_az_for_new_tenant()
-                    .ok_or(ApiError::BadRequest(anyhow::anyhow!(
-                        "No AZ with nodes found to load tenant"
-                    )))?
-            } else {
-                // We already have this tenant in memory
-                return Ok(());
-            }
+                .next()
+                .is_some()
         };
 
+        if present_in_memory {
+            return Ok(());
+        }
+
         let tenant_shards = self.persistence.load_tenant(tenant_id).await?;
         if tenant_shards.is_empty() {
             return Err(ApiError::NotFound(
@@ -2509,20 +2494,8 @@ impl Service {
             ));
         }
 
-        // Update the persistent shards with the AZ that we are about to apply to in-memory state
-        self.persistence
-            .set_tenant_shard_preferred_azs(
-                tenant_shards
-                    .iter()
-                    .map(|t| {
-                        (
-                            t.get_tenant_shard_id().expect("Corrupt shard in database"),
-                            load_in_az.clone(),
-                        )
-                    })
-                    .collect(),
-            )
-            .await?;
+        // TODO: choose a fresh AZ to use for this tenant when un-detaching: there definitely isn't a running
+        // compute, so no benefit to making AZ sticky across detaches.
 
         let mut locked = self.inner.write().unwrap();
         tracing::info!(
@@ -2532,7 +2505,7 @@ impl Service {
         );
 
         locked.tenants.extend(tenant_shards.into_iter().map(|p| {
-            let intent = IntentState::new(Some(load_in_az.clone()));
+            let intent = IntentState::new();
             let shard =
                 TenantShard::from_persistent(p, intent).expect("Corrupt shard row in database");
 
@@ -4263,22 +4236,6 @@ impl Service {
                 }
 
                 tracing::info!("Restoring parent shard {tenant_shard_id}");
-
-                // Drop any intents that refer to unavailable nodes, to enable this abort to proceed even
-                // if the original attachment location is offline.
-                if let Some(node_id) = shard.intent.get_attached() {
-                    if !nodes.get(node_id).unwrap().is_available() {
-                        tracing::info!("Demoting attached intent for {tenant_shard_id} on unavailable node {node_id}");
-                        shard.intent.demote_attached(scheduler, *node_id);
-                    }
-                }
-                for node_id in shard.intent.get_secondary().clone() {
-                    if !nodes.get(&node_id).unwrap().is_available() {
-                        tracing::info!("Dropping secondary intent for {tenant_shard_id} on unavailable node {node_id}");
-                        shard.intent.remove_secondary(scheduler, node_id);
-                    }
-                }
-
                 shard.splitting = SplitState::Idle;
                 if let Err(e) = shard.schedule(scheduler, &mut ScheduleContext::default()) {
                     // If this shard can't be scheduled now (perhaps due to offline nodes or
@@ -4432,13 +4389,15 @@ impl Service {
 
                     let mut child_state =
                         TenantShard::new(child, child_shard, policy.clone(), preferred_az.clone());
-                    child_state.intent =
-                        IntentState::single(scheduler, Some(pageserver), preferred_az.clone());
+                    child_state.intent = IntentState::single(scheduler, Some(pageserver));
                     child_state.observed = ObservedState {
                         locations: child_observed,
                     };
                     child_state.generation = Some(generation);
                     child_state.config = config.clone();
+                    if let Some(preferred_az) = &preferred_az {
+                        child_state.set_preferred_az(preferred_az.clone());
+                    }
 
                     // The child's TenantShard::splitting is intentionally left at the default value of Idle,
                     // as at this point in the split process we have succeeded and this part is infallible:
@@ -5055,8 +5014,6 @@ impl Service {
                         // If our new attached node was a secondary, it no longer should be.
                         shard.intent.remove_secondary(scheduler, migrate_req.node_id);
 
-                        shard.intent.set_attached(scheduler, Some(migrate_req.node_id));
-
                         // If we were already attached to something, demote that to a secondary
                         if let Some(old_attached) = old_attached {
                             if n > 0 {
@@ -5068,6 +5025,8 @@ impl Service {
                                 shard.intent.push_secondary(scheduler, old_attached);
                             }
                         }
+
+                        shard.intent.set_attached(scheduler, Some(migrate_req.node_id));
                     }
                     PlacementPolicy::Secondary => {
                         shard.intent.clear(scheduler);
@@ -5096,69 +5055,6 @@ impl Service {
         Ok(TenantShardMigrateResponse {})
     }
 
-    pub(crate) async fn tenant_shard_migrate_secondary(
-        &self,
-        tenant_shard_id: TenantShardId,
-        migrate_req: TenantShardMigrateRequest,
-    ) -> Result<TenantShardMigrateResponse, ApiError> {
-        let waiter = {
-            let mut locked = self.inner.write().unwrap();
-            let (nodes, tenants, scheduler) = locked.parts_mut();
-
-            let Some(node) = nodes.get(&migrate_req.node_id) else {
-                return Err(ApiError::BadRequest(anyhow::anyhow!(
-                    "Node {} not found",
-                    migrate_req.node_id
-                )));
-            };
-
-            if !node.is_available() {
-                // Warn but proceed: the caller may intend to manually adjust the placement of
-                // a shard even if the node is down, e.g. if intervening during an incident.
-                tracing::warn!("Migrating to unavailable node {node}");
-            }
-
-            let Some(shard) = tenants.get_mut(&tenant_shard_id) else {
-                return Err(ApiError::NotFound(
-                    anyhow::anyhow!("Tenant shard not found").into(),
-                ));
-            };
-
-            if shard.intent.get_secondary().len() == 1
-                && shard.intent.get_secondary()[0] == migrate_req.node_id
-            {
-                tracing::info!(
-                    "Migrating secondary to {node}: intent is unchanged {:?}",
-                    shard.intent
-                );
-            } else if shard.intent.get_attached() == &Some(migrate_req.node_id) {
-                tracing::info!("Migrating secondary to {node}: already attached where we were asked to create a secondary");
-            } else {
-                let old_secondaries = shard.intent.get_secondary().clone();
-                for secondary in old_secondaries {
-                    shard.intent.remove_secondary(scheduler, secondary);
-                }
-
-                shard.intent.push_secondary(scheduler, migrate_req.node_id);
-                shard.sequence = shard.sequence.next();
-                tracing::info!(
-                    "Migrating secondary to {node}: new intent {:?}",
-                    shard.intent
-                );
-            }
-
-            self.maybe_reconcile_shard(shard, nodes)
-        };
-
-        if let Some(waiter) = waiter {
-            waiter.wait_timeout(RECONCILE_TIMEOUT).await?;
-        } else {
-            tracing::info!("Migration is a no-op");
-        }
-
-        Ok(TenantShardMigrateResponse {})
-    }
-
     /// 'cancel' in this context means cancel any ongoing reconcile
     pub(crate) async fn tenant_shard_cancel_reconcile(
         &self,
@@ -5753,7 +5649,7 @@ impl Service {
             register_req.listen_http_port,
             register_req.listen_pg_addr,
             register_req.listen_pg_port,
-            register_req.availability_zone_id.clone(),
+            register_req.availability_zone_id,
         );
 
         // TODO: idempotency if the node already exists in the database
@@ -5773,9 +5669,8 @@ impl Service {
             .set(locked.nodes.len() as i64);
 
         tracing::info!(
-            "Registered pageserver {} ({}), now have {} pageservers",
+            "Registered pageserver {}, now have {} pageservers",
             register_req.node_id,
-            register_req.availability_zone_id,
             locked.nodes.len()
         );
         Ok(())
@@ -6509,7 +6404,6 @@ impl Service {
                 // Shard was dropped between planning and execution;
                 continue;
             };
-            tracing::info!("Applying optimization: {optimization:?}");
             if shard.apply_optimization(scheduler, optimization) {
                 optimizations_applied += 1;
                 if self.maybe_reconcile_shard(shard, nodes).is_some() {
@@ -6540,13 +6434,7 @@ impl Service {
 
         let mut work = Vec::new();
         let mut locked = self.inner.write().unwrap();
-        let (_nodes, tenants, scheduler) = locked.parts_mut();
-
-        // We are going to plan a bunch of optimisations before applying any of them, so the
-        // utilisation stats on nodes will be effectively stale for the >1st optimisation we
-        // generate.  To avoid this causing unstable migrations/flapping, it's important that the
-        // code in TenantShard for finding optimisations uses [`NodeAttachmentSchedulingScore::disregard_utilization`]
-        // to ignore the utilisation component of the score.
+        let (nodes, tenants, scheduler) = locked.parts_mut();
 
         for (_tenant_id, schedule_context, shards) in
             TenantShardContextIterator::new(tenants, ScheduleMode::Speculative)
@@ -6577,28 +6465,13 @@ impl Service {
                     continue;
                 }
 
-                // Fast path: we may quickly identify shards that don't have any possible optimisations
-                if !shard.maybe_optimizable(scheduler, &schedule_context) {
-                    if cfg!(feature = "testing") {
-                        // Check that maybe_optimizable doesn't disagree with the actual optimization functions.
-                        // Only do this in testing builds because it is not a correctness-critical check, so we shouldn't
-                        // panic in prod if we hit this, or spend cycles on it in prod.
-                        assert!(shard
-                            .optimize_attachment(scheduler, &schedule_context)
-                            .is_none());
-                        assert!(shard
-                            .optimize_secondary(scheduler, &schedule_context)
-                            .is_none());
-                    }
-                    continue;
-                }
-
+                // TODO: optimization calculations are relatively expensive: create some fast-path for
+                // the common idle case (avoiding the search on tenants that we have recently checked)
                 if let Some(optimization) =
-                    // If idle, maybe optimize attachments: if a shard has a secondary location that is preferable to
+                    // If idle, maybe ptimize attachments: if a shard has a secondary location that is preferable to
                     // its primary location based on soft constraints, cut it over.
-                    shard.optimize_attachment(scheduler, &schedule_context)
+                    shard.optimize_attachment(nodes, &schedule_context)
                 {
-                    tracing::info!(tenant_shard_id=%shard.tenant_shard_id, "Identified optimization for attachment: {optimization:?}");
                     work.push((shard.tenant_shard_id, optimization));
                     break;
                 } else if let Some(optimization) =
@@ -6608,7 +6481,6 @@ impl Service {
                     // in the same tenant with secondary locations on the node where they originally split.
                     shard.optimize_secondary(scheduler, &schedule_context)
                 {
-                    tracing::info!(tenant_shard_id=%shard.tenant_shard_id, "Identified optimization for secondary: {optimization:?}");
                     work.push((shard.tenant_shard_id, optimization));
                     break;
                 }
@@ -6657,10 +6529,8 @@ impl Service {
                         }
                     }
                 }
-                ScheduleOptimizationAction::ReplaceSecondary(_)
-                | ScheduleOptimizationAction::CreateSecondary(_)
-                | ScheduleOptimizationAction::RemoveSecondary(_) => {
-                    // No extra checks needed to manage secondaries: this does not interrupt client access
+                ScheduleOptimizationAction::ReplaceSecondary(_) => {
+                    // No extra checks needed to replace a secondary: this does not interrupt client access
                     validated_work.push((tenant_shard_id, optimization))
                 }
             };
@@ -6732,35 +6602,26 @@ impl Service {
     /// we have this helper to move things along faster.
     #[cfg(feature = "testing")]
     async fn kick_secondary_download(&self, tenant_shard_id: TenantShardId) {
-        let (attached_node, secondaries) = {
+        let (attached_node, secondary_node) = {
             let locked = self.inner.read().unwrap();
             let Some(shard) = locked.tenants.get(&tenant_shard_id) else {
-                tracing::warn!(
-                    "Skipping kick of secondary download for {tenant_shard_id}: not found"
-                );
                 return;
             };
-
-            let Some(attached) = shard.intent.get_attached() else {
-                tracing::warn!(
-                    "Skipping kick of secondary download for {tenant_shard_id}: no attached"
-                );
+            let (Some(attached), Some(secondary)) = (
+                shard.intent.get_attached(),
+                shard.intent.get_secondary().first(),
+            ) else {
                 return;
             };
-
-            let secondaries = shard
-                .intent
-                .get_secondary()
-                .iter()
-                .map(|n| locked.nodes.get(n).unwrap().clone())
-                .collect::<Vec<_>>();
-
-            (locked.nodes.get(attached).unwrap().clone(), secondaries)
+            (
+                locked.nodes.get(attached).unwrap().clone(),
+                locked.nodes.get(secondary).unwrap().clone(),
+            )
         };
 
         // Make remote API calls to upload + download heatmaps: we ignore errors because this is just
         // a 'kick' to let scheduling optimisation run more promptly.
-        match attached_node
+        attached_node
             .with_client_retries(
                 |client| async move { client.tenant_heatmap_upload(tenant_shard_id).await },
                 &self.config.jwt_token,
@@ -6769,57 +6630,22 @@ impl Service {
                 SHORT_RECONCILE_TIMEOUT,
                 &self.cancel,
             )
-            .await
-        {
-            Some(Err(e)) => {
-                tracing::info!(
-                    "Failed to upload heatmap from {attached_node} for {tenant_shard_id}: {e}"
-                );
-            }
-            None => {
-                tracing::info!(
-                    "Cancelled while uploading heatmap from {attached_node} for {tenant_shard_id}"
-                );
-            }
-            Some(Ok(_)) => {
-                tracing::info!(
-                    "Successfully uploaded heatmap from {attached_node} for {tenant_shard_id}"
-                );
-            }
-        }
+            .await;
 
-        for secondary_node in secondaries {
-            match secondary_node
-                .with_client_retries(
-                    |client| async move {
-                        client
-                            .tenant_secondary_download(
-                                tenant_shard_id,
-                                Some(Duration::from_secs(1)),
-                            )
-                            .await
-                    },
-                    &self.config.jwt_token,
-                    3,
-                    10,
-                    SHORT_RECONCILE_TIMEOUT,
-                    &self.cancel,
-                )
-                .await
-            {
-                Some(Err(e)) => {
-                    tracing::info!(
-                "Failed to download heatmap from {secondary_node} for {tenant_shard_id}: {e}"
-            );
-                }
-                None => {
-                    tracing::info!("Cancelled while downloading heatmap from {secondary_node} for {tenant_shard_id}");
-                }
-                Some(Ok(progress)) => {
-                    tracing::info!("Successfully downloaded heatmap from {secondary_node} for {tenant_shard_id}: {progress:?}");
-                }
-            }
-        }
+        secondary_node
+            .with_client_retries(
+                |client| async move {
+                    client
+                        .tenant_secondary_download(tenant_shard_id, Some(Duration::from_secs(1)))
+                        .await
+                },
+                &self.config.jwt_token,
+                3,
+                10,
+                SHORT_RECONCILE_TIMEOUT,
+                &self.cancel,
+            )
+            .await;
     }
 
     /// Look for shards which are oversized and in need of splitting
@@ -7255,15 +7081,9 @@ impl Service {
     fn fill_node_plan(&self, node_id: NodeId) -> Vec<TenantShardId> {
         let mut locked = self.inner.write().unwrap();
         let fill_requirement = locked.scheduler.compute_fill_requirement(node_id);
-        let (nodes, tenants, _scheduler) = locked.parts_mut();
 
-        let node_az = nodes
-            .get(&node_id)
-            .expect("Node must exist")
-            .get_availability_zone_id()
-            .clone();
-
-        let mut tids_by_node = tenants
+        let mut tids_by_node = locked
+            .tenants
             .iter_mut()
             .filter_map(|(tid, tenant_shard)| {
                 if !matches!(
@@ -7276,25 +7096,6 @@ impl Service {
                     return None;
                 }
 
-                // AZ check: when filling nodes after a restart, our intent is to move _back_ the
-                // shards which belong on this node, not to promote shards whose scheduling preference
-                // would be on their currently attached node.  So will avoid promoting shards whose
-                // home AZ doesn't match the AZ of the node we're filling.
-                match tenant_shard.preferred_az() {
-                    None => {
-                        // Shard doesn't have an AZ preference: it is elegible to be moved.
-                    }
-                    Some(az) if az == &node_az => {
-                        // This shard's home AZ is equal to the node we're filling: it is
-                        // elegible to be moved: fall through;
-                    }
-                    Some(_) => {
-                        // This shard's home AZ is somewhere other than the node we're filling:
-                        // do not include it in the fill plan.
-                        return None;
-                    }
-                }
-
                 if tenant_shard.intent.get_secondary().contains(&node_id) {
                     if let Some(primary) = tenant_shard.intent.get_attached() {
                         return Some((*primary, *tid));

storage_controller/src/service/context_iterator.rs

@@ -43,6 +43,9 @@ impl<'a> Iterator for TenantShardContextIterator<'a> {
 
             // Accumulate the schedule context for all the shards in a tenant
             schedule_context.avoid(&shard.intent.all_pageservers());
+            if let Some(attached) = shard.intent.get_attached() {
+                schedule_context.push_attached(*attached);
+            }
             tenant_shards.push(shard);
 
             if tenant_shard_id.shard_number.0 == tenant_shard_id.shard_count.count() - 1 {
@@ -112,7 +115,7 @@ mod tests {
         assert_eq!(tenant_id, t1_id);
         assert_eq!(shards[0].tenant_shard_id.shard_number, ShardNumber(0));
         assert_eq!(shards.len(), 1);
-        assert_eq!(context.location_count(), 2);
+        assert_eq!(context.attach_count(), 1);
 
         let (tenant_id, context, shards) = iter.next().unwrap();
         assert_eq!(tenant_id, t2_id);
@@ -121,13 +124,13 @@ mod tests {
         assert_eq!(shards[2].tenant_shard_id.shard_number, ShardNumber(2));
         assert_eq!(shards[3].tenant_shard_id.shard_number, ShardNumber(3));
         assert_eq!(shards.len(), 4);
-        assert_eq!(context.location_count(), 8);
+        assert_eq!(context.attach_count(), 4);
 
         let (tenant_id, context, shards) = iter.next().unwrap();
         assert_eq!(tenant_id, t3_id);
         assert_eq!(shards[0].tenant_shard_id.shard_number, ShardNumber(0));
         assert_eq!(shards.len(), 1);
-        assert_eq!(context.location_count(), 2);
+        assert_eq!(context.attach_count(), 1);
 
         for shard in tenants.values_mut() {
             shard.intent.clear(&mut scheduler);

test_runner/fixtures/neon_fixtures.py

@@ -3645,7 +3645,7 @@ def static_proxy(
     vanilla_pg.safe_psql("create user proxy with login superuser password 'password'")
     vanilla_pg.safe_psql("CREATE SCHEMA IF NOT EXISTS neon_control_plane")
     vanilla_pg.safe_psql(
-        "CREATE TABLE neon_control_plane.endpoints (endpoint_id VARCHAR(255) PRIMARY KEY, allowed_ips VARCHAR(255), jwks jsonb default '[]'::jsonb)"
+        "CREATE TABLE neon_control_plane.endpoints (endpoint_id VARCHAR(255) PRIMARY KEY, allowed_ips VARCHAR(255))"
     )
 
     proxy_port = port_distributor.get_port()

test_runner/performance/test_sharding_autosplit.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import concurrent.futures
 import re
-import threading
 from pathlib import Path
 
 import pytest
@@ -189,20 +188,7 @@ def test_sharding_autosplit(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin):
 
         check_pgbench_output(out_path)
 
-    stop_pump = threading.Event()
-
-    def pump_controller():
-        # Run a background loop to force the storage controller to run its
-        # background work faster than it otherwise would: this helps
-        # us:
-        #  A) to create a test that runs in a shorter time
-        #  B) to create a test that is more intensive by doing the shard migrations
-        #     after splits happen more rapidly.
-        while not stop_pump.is_set():
-            env.storage_controller.reconcile_all()
-            stop_pump.wait(0.1)
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=tenant_count + 1) as pgbench_threads:
+    with concurrent.futures.ThreadPoolExecutor(max_workers=tenant_count) as pgbench_threads:
         pgbench_futs = []
         for tenant_state in tenants.values():
             fut = pgbench_threads.submit(run_pgbench_init, tenant_state.endpoint)
@@ -212,8 +198,6 @@ def test_sharding_autosplit(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin):
         for fut in pgbench_futs:
             fut.result()
 
-        pump_fut = pgbench_threads.submit(pump_controller)
-
         pgbench_futs = []
         for tenant_state in tenants.values():
             fut = pgbench_threads.submit(run_pgbench_main, tenant_state.endpoint)
@@ -223,9 +207,6 @@ def test_sharding_autosplit(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin):
         for fut in pgbench_futs:
             fut.result()
 
-        stop_pump.set()
-        pump_fut.result()
-
     def assert_all_split():
         for tenant_id in tenants.keys():
             shards = tenant_get_shards(env, tenant_id)

test_runner/performance/test_storage_controller_scale.py

@@ -13,13 +13,11 @@ from fixtures.log_helper import log
 from fixtures.neon_fixtures import (
     NeonEnv,
     NeonEnvBuilder,
-    NeonPageserver,
     PageserverAvailability,
     PageserverSchedulingPolicy,
 )
 from fixtures.pageserver.http import PageserverApiException, PageserverHttpClient
 from fixtures.pg_version import PgVersion
-from fixtures.utils import wait_until
 
 
 def get_consistent_node_shard_counts(env: NeonEnv, total_shards) -> defaultdict[str, int]:
@@ -87,12 +85,8 @@ def test_storage_controller_many_tenants(
     )
 
     AZS = ["alpha", "bravo", "charlie"]
-
-    def az_selector(node_id):
-        return f"az-{AZS[(node_id - 1) % len(AZS)]}"
-
     neon_env_builder.pageserver_config_override = lambda ps_cfg: ps_cfg.update(
-        {"availability_zone": az_selector(ps_cfg["id"])}
+        {"availability_zone": f"az-{AZS[ps_cfg['id'] % len(AZS)]}"}
     )
 
     # A small sleep on each call into the notify hook, to simulate the latency of doing a database write
@@ -174,31 +168,6 @@ def test_storage_controller_many_tenants(
         log.info(f"Resident memory: {rss} ({ rss / total_shards} per shard)")
         assert rss < expect_memory_per_shard * total_shards
 
-    def assert_all_tenants_scheduled_in_home_az():
-        for tenant_id in tenant_ids:
-            desc = env.storage_controller.tenant_describe(tenant_id)
-            preferred_az = None
-            for shard in desc["shards"]:
-                # All shards in a tenant should have the same preferred AZ
-                if preferred_az is None:
-                    preferred_az = shard["preferred_az_id"]
-                else:
-                    assert preferred_az == shard["preferred_az_id"]
-
-                # Attachment should be in the preferred AZ
-                assert shard["preferred_az_id"] == az_selector(
-                    shard["node_attached"]
-                ), f"Shard {shard['tenant_shard_id']} not in {shard['preferred_az_id']}"
-
-                # Secondary locations should not be in the preferred AZ
-                for node_secondary in shard["node_secondary"]:
-                    assert (
-                        shard["preferred_az_id"] != az_selector(node_secondary)
-                    ), f"Shard {shard['tenant_shard_id']} secondary should be in {shard['preferred_az_id']}"
-
-                # There should only be one secondary location (i.e. no migrations in flight)
-                assert len(shard["node_secondary"]) == 1
-
     # Issue more concurrent operations than the storage controller's reconciler concurrency semaphore
     # permits, to ensure that we are exercising stressing that.
     api_concurrency = 135
@@ -273,22 +242,6 @@ def test_storage_controller_many_tenants(
             f"Created {len(tenants_with_timelines)} timelines in {time.time() - t1}, {len(tenants_with_timelines) / (time.time() - t1)}/s"
         )
 
-    # Check initial scheduling
-    assert_all_tenants_scheduled_in_home_az()
-    az_attached_counts: defaultdict[str, int] = defaultdict(int)
-    az_secondary_counts: defaultdict[str, int] = defaultdict(int)
-    node_attached_counts: defaultdict[str, int] = defaultdict(int)
-    for tenant_id in tenants.keys():
-        desc = env.storage_controller.tenant_describe(tenant_id)
-        for shard in desc["shards"]:
-            az_attached_counts[az_selector(shard["node_attached"])] += 1
-            node_attached_counts[shard["node_attached"]] += 1
-            for node_secondary in shard["node_secondary"]:
-                az_secondary_counts[az_selector(node_secondary)] += 1
-
-    log.info(f"Initial node attached counts: {node_attached_counts}")
-    log.info(f"Initial AZ shard counts: {az_attached_counts}, {az_secondary_counts}")
-
     # Plan operations: ensure each tenant with a timeline gets at least
     # one of each operation type.  Then add other tenants to make up the
     # numbers.
@@ -497,77 +450,11 @@ def test_storage_controller_many_tenants(
         env.storage_controller.reconcile_until_idle(max_interval=0.1, timeout_secs=120)
         env.storage_controller.consistency_check()
 
-    # Since we did `reconcile_until_idle` during the above loop, the system should be left in
-    # an optimally scheduled state.  Validate that this includes all the tenants being scheduled
-    # in their home AZ.
-    assert_all_tenants_scheduled_in_home_az()
-
     # Consistency check is safe here: restarting pageservers should not have caused any Reconcilers to spawn,
     # as they were not offline long enough to trigger any scheduling changes.
     env.storage_controller.consistency_check()
     check_memory()
 
-    # Simulate loss of an AZ
-    victim_az = "az-alpha"
-    killed_pageservers = []
-    for ps in env.pageservers:
-        if az_selector(ps.id) == victim_az:
-            ps.stop(immediate=True)
-            killed_pageservers.append(ps)
-            log.info(f"Killed pageserver {ps.id}")
-
-    assert killed_pageservers
-
-    # Wait for the controller to notice the pageservers are dead
-    def assert_pageservers_availability(
-        pageservers: list[NeonPageserver], expected_availability: PageserverAvailability
-    ):
-        nodes = env.storage_controller.nodes()
-        checked_any = False
-        node_ids = [ps.id for ps in pageservers]
-        for node in nodes:
-            if node["id"] in node_ids:
-                checked_any = True
-                assert (
-                    node["availability"] == expected_availability
-                ), f"Node {node['id']} is not {expected_availability} yet: {node['availability']}"
-
-        assert checked_any
-
-    wait_until(
-        lambda: assert_pageservers_availability(killed_pageservers, PageserverAvailability.OFFLINE),
-        timeout=60,
-    )
-
-    # Let the controller finish all its rescheduling
-    env.storage_controller.reconcile_until_idle(max_interval=0.1, timeout_secs=120)
-
-    # Check that all the tenants are rescheduled to the remaining pageservers
-    for tenant_id in tenant_ids:
-        desc = env.storage_controller.tenant_describe(tenant_id)
-        for shard in desc["shards"]:
-            # Attachment should be outside the AZ where we killed the pageservers
-            assert (
-                az_selector(shard["node_attached"]) != victim_az
-            ), f"Shard {shard['tenant_shard_id']} still in {victim_az} (node {shard['node_attached']})"
-
-    # Bring back the pageservers
-    for ps in killed_pageservers:
-        ps.start()
-
-    wait_until(
-        lambda: assert_pageservers_availability(killed_pageservers, PageserverAvailability.ACTIVE),
-        timeout=60,
-    )
-
-    # A very long timeout is required: we will be migrating all the tenants on all the pageservers
-    # in the region that we just restored.  Assume it'll take up to twice as long as it took to fill
-    # a single node
-    env.storage_controller.reconcile_until_idle(
-        max_interval=0.1, timeout_secs=DRAIN_FILL_TIMEOUT * 4
-    )
-    assert_all_tenants_scheduled_in_home_az()
-
     # Stop the storage controller before tearing down fixtures, because it otherwise might log
     # errors trying to call our `ComputeReconfigure`.
     env.storage_controller.stop()

test_runner/regress/test_compute_metrics.py

@@ -219,7 +219,7 @@ if SQL_EXPORTER is None:
             #
             # The "host" network mode allows sql_exporter to talk to the
             # endpoint which is running on the host.
-            super().__init__("docker.io/burningalchemist/sql_exporter:0.17.0", network_mode="host")
+            super().__init__("docker.io/burningalchemist/sql_exporter:0.16.0", network_mode="host")
 
             self.__logs_dir = logs_dir
             self.__port = port
@@ -252,7 +252,7 @@ if SQL_EXPORTER is None:
             log.info("Waiting for sql_exporter to be ready")
             wait_for_logs(
                 self,
-                rf'msg="Listening on" address=\[::\]:{self.__port}',
+                rf'level=info msg="Listening on" address=\[::\]:{self.__port}',
                 timeout=5,
             )
 
@@ -344,7 +344,10 @@ else:
                         time.sleep(0.5)
                         continue
 
-                    if f'msg="Listening on" address=[::]:{self._sql_exporter_port}' in line:
+                    if (
+                        f'level=info msg="Listening on" address=[::]:{self._sql_exporter_port}'
+                        in line
+                    ):
                         break
 
         @override

test_runner/regress/test_physical_replication.py

@@ -187,7 +187,7 @@ def test_physical_replication_config_mismatch_too_many_known_xids(neon_simple_en
         origin=primary,
         endpoint_id="secondary",
         config_lines=[
-            "max_connections=5",
+            "max_connections=2",
             "autovacuum_max_workers=1",
             "max_worker_processes=5",
             "max_wal_senders=1",
@@ -260,4 +260,4 @@ def test_physical_replication_config_mismatch_max_locks_per_transaction(neon_sim
         secondary.stop()
 
     log.info(f"Replica crashed with {e}")
-    assert secondary.log_contains("You might need to increase")
+    assert secondary.log_contains("out of shared memory")

test_runner/regress/test_sharding.py

@@ -520,18 +520,14 @@ def test_sharding_split_smoke(
     shard_count = 2
     # Shard count we split into
     split_shard_count = 4
-    # In preferred AZ & other AZ we will end up with one shard per pageserver
-    neon_env_builder.num_pageservers = split_shard_count * 2
+    # We will have 2 shards per pageserver once done (including secondaries)
+    neon_env_builder.num_pageservers = split_shard_count
 
     # Two AZs
     def assign_az(ps_cfg):
         az = f"az-{(ps_cfg['id'] - 1) % 2}"
         ps_cfg["availability_zone"] = az
 
-        # We will run more pageservers than tests usually do, so give them tiny page caches
-        # in case we're on a test node under memory pressure.
-        ps_cfg["page_cache_size"] = 128
-
     neon_env_builder.pageserver_config_override = assign_az
 
     # 1MiB stripes: enable getting some meaningful data distribution without
@@ -683,8 +679,8 @@ def test_sharding_split_smoke(
     # - shard_count reconciles for the original setup of the tenant
     # - shard_count reconciles for detaching the original secondary locations during split
     # - split_shard_count reconciles during shard splitting, for setting up secondaries.
-    # - split_shard_count/2 reconciles to migrate shards to their temporary secondaries
-    expect_reconciles = shard_count * 2 + split_shard_count + 3 * (split_shard_count / 2)
+    # - split_shard_count/2 of the child shards will need to fail over to their secondaries (since we have 8 shards and 4 pageservers, only 4 will move)
+    expect_reconciles = shard_count * 2 + split_shard_count + split_shard_count / 2
 
     reconcile_ok = env.storage_controller.get_metric_value(
         "storage_controller_reconcile_complete_total", filter={"status": "ok"}
@@ -749,14 +745,10 @@ def test_sharding_split_smoke(
     # dominated by shard count.
     log.info(f"total: {total}")
     assert total == {
-        1: 1,
-        2: 1,
-        3: 1,
-        4: 1,
-        5: 1,
-        6: 1,
-        7: 1,
-        8: 1,
+        1: 2,
+        2: 2,
+        3: 2,
+        4: 2,
     }
 
     # The controller is not required to lay out the attached locations in any particular way, but
@@ -1395,7 +1387,13 @@ def test_sharding_split_failures(
                 else:
                     attached_count += 1
 
-        assert secondary_count == initial_shard_count
+        if exclude_ps_id is not None:
+            # For a node failure case, we expect there to be a secondary location
+            # scheduled on the offline node, so expect one fewer secondary in total
+            assert secondary_count == initial_shard_count - 1
+        else:
+            assert secondary_count == initial_shard_count
+
         assert attached_count == initial_shard_count
 
     def assert_split_done(exclude_ps_id: int | None = None) -> None:

test_runner/regress/test_storage_controller.py

@@ -1052,7 +1052,7 @@ def test_storage_controller_debug_apis(neon_env_builder: NeonEnvBuilder):
     that just hits the endpoints to check that they don't bitrot.
     """
 
-    neon_env_builder.num_pageservers = 3
+    neon_env_builder.num_pageservers = 2
     env = neon_env_builder.init_start()
 
     tenant_id = TenantId.generate()
@@ -1077,7 +1077,7 @@ def test_storage_controller_debug_apis(neon_env_builder: NeonEnvBuilder):
         "GET", f"{env.storage_controller_api}/debug/v1/scheduler"
     )
     # Two nodes, in a dict of node_id->node
-    assert len(response.json()["nodes"]) == 3
+    assert len(response.json()["nodes"]) == 2
     assert sum(v["shard_count"] for v in response.json()["nodes"].values()) == 3
     assert all(v["may_schedule"] for v in response.json()["nodes"].values())
 
@@ -1088,25 +1088,13 @@ def test_storage_controller_debug_apis(neon_env_builder: NeonEnvBuilder):
         headers=env.storage_controller.headers(TokenScope.ADMIN),
     )
 
-    # Secondary migration API: superficial check that it migrates
-    secondary_dest = env.pageservers[2].id
-    env.storage_controller.request(
-        "PUT",
-        f"{env.storage_controller_api}/control/v1/tenant/{tenant_id}-0002/migrate_secondary",
-        headers=env.storage_controller.headers(TokenScope.ADMIN),
-        json={"tenant_shard_id": f"{tenant_id}-0002", "node_id": secondary_dest},
-    )
-    assert env.storage_controller.tenant_describe(tenant_id)["shards"][0]["node_secondary"] == [
-        secondary_dest
-    ]
-
     # Node unclean drop API
     response = env.storage_controller.request(
         "POST",
         f"{env.storage_controller_api}/debug/v1/node/{env.pageservers[1].id}/drop",
         headers=env.storage_controller.headers(TokenScope.ADMIN),
     )
-    assert len(env.storage_controller.node_list()) == 2
+    assert len(env.storage_controller.node_list()) == 1
 
     # Tenant unclean drop API
     response = env.storage_controller.request(
@@ -1824,13 +1812,7 @@ def test_storcon_cli(neon_env_builder: NeonEnvBuilder):
     """
     output_dir = neon_env_builder.test_output_dir
     shard_count = 4
-    neon_env_builder.num_pageservers = 2
-    env = neon_env_builder.init_configs()
-    env.start()
-
-    tenant_id = TenantId.generate()
-    env.create_tenant(tenant_id, placement_policy='{"Attached":1}', shard_count=shard_count)
-
+    env = neon_env_builder.init_start(initial_tenant_shard_count=shard_count)
     base_args = [env.neon_binpath / "storcon_cli", "--api", env.storage_controller_api]
 
     def storcon_cli(args):
@@ -1859,7 +1841,7 @@ def test_storcon_cli(neon_env_builder: NeonEnvBuilder):
     # List nodes
     node_lines = storcon_cli(["nodes"])
     # Table header, footer, and one line of data
-    assert len(node_lines) == 7
+    assert len(node_lines) == 5
     assert "localhost" in node_lines[3]
 
     # Pause scheduling onto a node
@@ -1877,21 +1859,10 @@ def test_storcon_cli(neon_env_builder: NeonEnvBuilder):
     storcon_cli(["node-configure", "--node-id", "1", "--availability", "offline"])
     assert "Offline" in storcon_cli(["nodes"])[3]
 
-    # Restore node, verify status changes in CLI output
-    env.pageservers[0].start()
-
-    def is_online():
-        assert "Offline" not in storcon_cli(["nodes"])
-
-    wait_until(is_online)
-
-    # Let everything stabilize after node failure to avoid interfering with subsequent steps
-    env.storage_controller.reconcile_until_idle(timeout_secs=10)
-
     # List tenants
     tenant_lines = storcon_cli(["tenants"])
     assert len(tenant_lines) == 5
-    assert str(tenant_id) in tenant_lines[3]
+    assert str(env.initial_tenant) in tenant_lines[3]
 
     # Setting scheduling policies intentionally result in warnings, they're for rare use.
     env.storage_controller.allowed_errors.extend(
@@ -1899,58 +1870,23 @@ def test_storcon_cli(neon_env_builder: NeonEnvBuilder):
     )
 
     # Describe a tenant
-    tenant_lines = storcon_cli(["tenant-describe", "--tenant-id", str(tenant_id)])
+    tenant_lines = storcon_cli(["tenant-describe", "--tenant-id", str(env.initial_tenant)])
     assert len(tenant_lines) >= 3 + shard_count * 2
-    assert str(tenant_id) in tenant_lines[0]
-
-    # Migrate an attached location
-    def other_ps_id(current_ps_id):
-        return (
-            env.pageservers[0].id
-            if current_ps_id == env.pageservers[1].id
-            else env.pageservers[1].id
-        )
-
-    storcon_cli(
-        [
-            "tenant-shard-migrate",
-            "--tenant-shard-id",
-            f"{tenant_id}-0004",
-            "--node",
-            str(
-                other_ps_id(
-                    env.storage_controller.tenant_describe(tenant_id)["shards"][0]["node_attached"]
-                )
-            ),
-        ]
-    )
-
-    # Migrate a secondary location
-    storcon_cli(
-        [
-            "tenant-shard-migrate-secondary",
-            "--tenant-shard-id",
-            f"{tenant_id}-0004",
-            "--node",
-            str(
-                other_ps_id(
-                    env.storage_controller.tenant_describe(tenant_id)["shards"][0][
-                        "node_secondary"
-                    ][0]
-                )
-            ),
-        ]
-    )
+    assert str(env.initial_tenant) in tenant_lines[0]
 
     # Pause changes on a tenant
-    storcon_cli(["tenant-policy", "--tenant-id", str(tenant_id), "--scheduling", "stop"])
+    storcon_cli(["tenant-policy", "--tenant-id", str(env.initial_tenant), "--scheduling", "stop"])
     assert "Stop" in storcon_cli(["tenants"])[3]
 
     # Cancel ongoing reconcile on a tenant
-    storcon_cli(["tenant-shard-cancel-reconcile", "--tenant-shard-id", f"{tenant_id}-0104"])
+    storcon_cli(
+        ["tenant-shard-cancel-reconcile", "--tenant-shard-id", f"{env.initial_tenant}-0104"]
+    )
 
     # Change a tenant's placement
-    storcon_cli(["tenant-policy", "--tenant-id", str(tenant_id), "--placement", "secondary"])
+    storcon_cli(
+        ["tenant-policy", "--tenant-id", str(env.initial_tenant), "--placement", "secondary"]
+    )
     assert "Secondary" in storcon_cli(["tenants"])[3]
 
     # Modify a tenant's config
@@ -1958,7 +1894,7 @@ def test_storcon_cli(neon_env_builder: NeonEnvBuilder):
         [
             "patch-tenant-config",
             "--tenant-id",
-            str(tenant_id),
+            str(env.initial_tenant),
             "--config",
             json.dumps({"pitr_interval": "1m"}),
         ]
@@ -3213,12 +3149,11 @@ def eq_safekeeper_records(a: dict[str, Any], b: dict[str, Any]) -> bool:
 @run_only_on_default_postgres("this is like a 'unit test' against storcon db")
 def test_shard_preferred_azs(neon_env_builder: NeonEnvBuilder):
     def assign_az(ps_cfg):
-        az = f"az-{ps_cfg['id'] % 2}"
-        log.info("Assigned AZ {az}")
+        az = f"az-{ps_cfg['id']}"
         ps_cfg["availability_zone"] = az
 
     neon_env_builder.pageserver_config_override = assign_az
-    neon_env_builder.num_pageservers = 4
+    neon_env_builder.num_pageservers = 2
     env = neon_env_builder.init_configs()
     env.start()
 
@@ -3233,14 +3168,8 @@ def test_shard_preferred_azs(neon_env_builder: NeonEnvBuilder):
 
         assert shards[0]["preferred_az_id"] == expected_az
 
-    # When all other schedule scoring parameters are equal, tenants should round-robin on AZs
-    assert env.storage_controller.tenant_describe(tids[0])["shards"][0]["preferred_az_id"] == "az-0"
-    assert env.storage_controller.tenant_describe(tids[1])["shards"][0]["preferred_az_id"] == "az-1"
-    assert env.storage_controller.tenant_describe(tids[2])["shards"][0]["preferred_az_id"] == "az-0"
-
-    # Try modifying preferred AZ
     updated = env.storage_controller.set_preferred_azs(
-        {TenantShardId(tid, 0, 0): "az-0" for tid in tids}
+        {TenantShardId(tid, 0, 0): "foo" for tid in tids}
     )
 
     assert set(updated) == set([TenantShardId(tid, 0, 0) for tid in tids])
@@ -3248,24 +3177,29 @@ def test_shard_preferred_azs(neon_env_builder: NeonEnvBuilder):
     for tid in tids:
         shards = env.storage_controller.tenant_describe(tid)["shards"]
         assert len(shards) == 1
-        assert shards[0]["preferred_az_id"] == "az-0"
+        assert shards[0]["preferred_az_id"] == "foo"
 
-    # Having modified preferred AZ, we should get moved there
-    env.storage_controller.reconcile_until_idle(max_interval=0.1)
-    for tid in tids:
-        shard = env.storage_controller.tenant_describe(tid)["shards"][0]
-        attached_to = shard["node_attached"]
-        attached_in_az = env.get_pageserver(attached_to).az_id
-        assert shard["preferred_az_id"] == attached_in_az == "az-0"
+    # Generate a layer to avoid shard split handling on ps from tripping
+    # up on debug assert.
+    timeline_id = TimelineId.generate()
+    env.create_timeline("bar", tids[0], timeline_id)
+
+    workload = Workload(env, tids[0], timeline_id, branch_name="bar")
+    workload.init()
+    workload.write_rows(256)
+    workload.validate()
 
     env.storage_controller.tenant_shard_split(tids[0], shard_count=2)
-    env.storage_controller.reconcile_until_idle(max_interval=0.1)
     shards = env.storage_controller.tenant_describe(tids[0])["shards"]
     assert len(shards) == 2
     for shard in shards:
         attached_to = shard["node_attached"]
-        attached_in_az = env.get_pageserver(attached_to).az_id
-        assert shard["preferred_az_id"] == attached_in_az == "az-0"
+        expected_az = env.get_pageserver(attached_to).az_id
+
+        # The scheduling optimization logic is not yet AZ-aware, so doesn't succeed
+        # in putting the tenant shards in the preferred AZ.
+        # To be fixed in https://github.com/neondatabase/neon/pull/9916
+        # assert shard["preferred_az_id"] == expected_az
 
 
 @run_only_on_default_postgres("Postgres version makes no difference here")

vendor/revisions.json

@@ -1,7 +1,7 @@
 {
   "v17": [
     "17.2",
-    "0f8da73ed08d4fc4ee58cccea008c75bfb20baa8"
+    "9c9e9a78a93aebec2f6a2f54644442d35ffa245c"
   ],
   "v16": [
     "16.6",

workspace_hack/Cargo.toml

@@ -65,7 +65,6 @@ num-rational = { version = "0.4", default-features = false, features = ["num-big
 num-traits = { version = "0.2", features = ["i128", "libm"] }
 once_cell = { version = "1" }
 parquet = { version = "53", default-features = false, features = ["zstd"] }
-postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", branch = "neon", default-features = false, features = ["with-serde_json-1"] }
 prost = { version = "0.13", features = ["no-recursion-limit", "prost-derive"] }
 rand = { version = "0.8", features = ["small_rng"] }
 regex = { version = "1" }
@@ -87,7 +86,6 @@ tikv-jemalloc-ctl = { version = "0.6", features = ["stats", "use_std"] }
 tikv-jemalloc-sys = { version = "0.6", features = ["profiling", "stats", "unprefixed_malloc_on_supported_platforms"] }
 time = { version = "0.3", features = ["macros", "serde-well-known"] }
 tokio = { version = "1", features = ["full", "test-util"] }
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", branch = "neon", features = ["with-serde_json-1"] }
 tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "ring", "tls12"] }
 tokio-stream = { version = "0.1", features = ["net"] }
 tokio-util = { version = "0.7", features = ["codec", "compat", "io", "rt"] }