LCOV - code coverage report
Current view: top level - storage_controller/src - scheduler.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 89.8 % 1081 971
Test Date: 2025-03-12 00:01:28 Functions: 80.0 % 85 68

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::fmt::Debug;
       3              : 
       4              : use http_utils::error::ApiError;
       5              : use itertools::Itertools;
       6              : use pageserver_api::controller_api::AvailabilityZone;
       7              : use pageserver_api::models::PageserverUtilization;
       8              : use serde::Serialize;
       9              : use utils::id::NodeId;
      10              : 
      11              : use crate::metrics::NodeLabelGroup;
      12              : use crate::node::Node;
      13              : use crate::tenant_shard::TenantShard;
      14              : 
      15              : /// Scenarios in which we cannot find a suitable location for a tenant shard
      16              : #[derive(thiserror::Error, Debug)]
      17              : pub enum ScheduleError {
      18              :     #[error("No pageservers found")]
      19              :     NoPageservers,
      20              :     #[error("No pageserver found matching constraint")]
      21              :     ImpossibleConstraint,
      22              : }
      23              : 
      24              : impl From<ScheduleError> for ApiError {
      25            0 :     fn from(value: ScheduleError) -> Self {
      26            0 :         ApiError::Conflict(format!("Scheduling error: {}", value))
      27            0 :     }
      28              : }
      29              : 
      30              : #[derive(Serialize)]
      31              : pub enum MaySchedule {
      32              :     Yes(PageserverUtilization),
      33              :     No,
      34              : }
      35              : 
      36              : #[derive(Serialize)]
      37              : pub(crate) struct SchedulerNode {
      38              :     /// How many shards are currently scheduled on this node, via their [`crate::tenant_shard::IntentState`].
      39              :     shard_count: usize,
      40              :     /// How many shards are currently attached on this node, via their [`crate::tenant_shard::IntentState`].
      41              :     attached_shard_count: usize,
      42              :     /// How many shards have a location on this node (via [`crate::tenant_shard::IntentState`]) _and_ this node
      43              :     /// is in their preferred AZ (i.e. this is their 'home' location)
      44              :     home_shard_count: usize,
      45              :     /// Availability zone id in which the node resides
      46              :     az: AvailabilityZone,
      47              : 
      48              :     /// Whether this node is currently elegible to have new shards scheduled (this is derived
      49              :     /// from a node's availability state and scheduling policy).
      50              :     may_schedule: MaySchedule,
      51              : }
      52              : 
      53              : pub(crate) trait NodeSchedulingScore: Debug + Ord + Copy + Sized {
      54              :     fn generate(
      55              :         node_id: &NodeId,
      56              :         node: &mut SchedulerNode,
      57              :         preferred_az: &Option<AvailabilityZone>,
      58              :         context: &ScheduleContext,
      59              :     ) -> Option<Self>;
      60              : 
      61              :     /// Return a score that drops any components based on node utilization: this is useful
      62              :     /// for finding scores for scheduling optimisation, when we want to avoid rescheduling
      63              :     /// shards due to e.g. disk usage, to avoid flapping.
      64              :     fn for_optimization(&self) -> Self;
      65              : 
      66              :     fn is_overloaded(&self) -> bool;
      67              :     fn node_id(&self) -> NodeId;
      68              : }
      69              : 
      70              : pub(crate) trait ShardTag {
      71              :     type Score: NodeSchedulingScore;
      72              : }
      73              : 
      74              : pub(crate) struct AttachedShardTag {}
      75              : impl ShardTag for AttachedShardTag {
      76              :     type Score = NodeAttachmentSchedulingScore;
      77              : }
      78              : 
      79              : pub(crate) struct SecondaryShardTag {}
      80              : impl ShardTag for SecondaryShardTag {
      81              :     type Score = NodeSecondarySchedulingScore;
      82              : }
      83              : 
      84              : #[derive(PartialEq, Eq, Debug, Clone, Copy)]
      85              : enum AzMatch {
      86              :     Yes,
      87              :     No,
      88              :     Unknown,
      89              : }
      90              : 
      91              : impl AzMatch {
      92        91487 :     fn new(node_az: &AvailabilityZone, shard_preferred_az: Option<&AvailabilityZone>) -> Self {
      93        91315 :         match shard_preferred_az {
      94        91315 :             Some(preferred_az) if preferred_az == node_az => Self::Yes,
      95        52686 :             Some(_preferred_az) => Self::No,
      96          172 :             None => Self::Unknown,
      97              :         }
      98        91487 :     }
      99              : }
     100              : 
     101              : #[derive(PartialEq, Eq, Debug, Clone, Copy)]
     102              : struct AttachmentAzMatch(AzMatch);
     103              : 
     104              : impl Ord for AttachmentAzMatch {
     105        65399 :     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
     106        65399 :         // Lower scores indicate a more suitable node.
     107        65399 :         // Note that we prefer a node for which we don't have
     108        65399 :         // info to a node which we are certain doesn't match the
     109        65399 :         // preferred AZ of the shard.
     110       130798 :         let az_match_score = |az_match: &AzMatch| match az_match {
     111        64280 :             AzMatch::Yes => 0,
     112          162 :             AzMatch::Unknown => 1,
     113        66356 :             AzMatch::No => 2,
     114       130798 :         };
     115              : 
     116        65399 :         az_match_score(&self.0).cmp(&az_match_score(&other.0))
     117        65399 :     }
     118              : }
     119              : 
     120              : impl PartialOrd for AttachmentAzMatch {
     121        65399 :     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
     122        65399 :         Some(self.cmp(other))
     123        65399 :     }
     124              : }
     125              : 
     126              : #[derive(PartialEq, Eq, Debug, Clone, Copy)]
     127              : struct SecondaryAzMatch(AzMatch);
     128              : 
     129              : impl Ord for SecondaryAzMatch {
     130        35455 :     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
     131        35455 :         // Lower scores indicate a more suitable node.
     132        35455 :         // For secondary locations we wish to avoid the preferred AZ
     133        35455 :         // of the shard.
     134        70910 :         let az_match_score = |az_match: &AzMatch| match az_match {
     135        49329 :             AzMatch::No => 0,
     136           24 :             AzMatch::Unknown => 1,
     137        21557 :             AzMatch::Yes => 2,
     138        70910 :         };
     139              : 
     140        35455 :         az_match_score(&self.0).cmp(&az_match_score(&other.0))
     141        35455 :     }
     142              : }
     143              : 
     144              : impl PartialOrd for SecondaryAzMatch {
     145        35451 :     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
     146        35451 :         Some(self.cmp(other))
     147        35451 :     }
     148              : }
     149              : 
     150              : /// Scheduling score of a given node for shard attachments.
     151              : /// Lower scores indicate more suitable nodes.
     152              : /// Ordering is given by member declaration order (top to bottom).
     153              : #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
     154              : pub(crate) struct NodeAttachmentSchedulingScore {
     155              :     /// Flag indicating whether this node matches the preferred AZ
     156              :     /// of the shard. For equal affinity scores, nodes in the matching AZ
     157              :     /// are considered first.
     158              :     az_match: AttachmentAzMatch,
     159              :     /// The number of shards belonging to the tenant currently being
     160              :     /// scheduled that are attached to this node.
     161              :     affinity_score: AffinityScore,
     162              :     /// Utilisation score that combines shard count and disk utilisation
     163              :     utilization_score: u64,
     164              :     /// Total number of shards attached to this node. When nodes have identical utilisation, this
     165              :     /// acts as an anti-affinity between attached shards.
     166              :     total_attached_shard_count: usize,
     167              :     /// Convenience to make selection deterministic in tests and empty systems
     168              :     node_id: NodeId,
     169              : }
     170              : 
     171              : impl NodeSchedulingScore for NodeAttachmentSchedulingScore {
     172        52291 :     fn generate(
     173        52291 :         node_id: &NodeId,
     174        52291 :         node: &mut SchedulerNode,
     175        52291 :         preferred_az: &Option<AvailabilityZone>,
     176        52291 :         context: &ScheduleContext,
     177        52291 :     ) -> Option<Self> {
     178        52291 :         let utilization = match &mut node.may_schedule {
     179        52289 :             MaySchedule::Yes(u) => u,
     180              :             MaySchedule::No => {
     181            2 :                 return None;
     182              :             }
     183              :         };
     184              : 
     185        52289 :         Some(Self {
     186        52289 :             affinity_score: context
     187        52289 :                 .nodes
     188        52289 :                 .get(node_id)
     189        52289 :                 .copied()
     190        52289 :                 .unwrap_or(AffinityScore::FREE),
     191        52289 :             az_match: AttachmentAzMatch(AzMatch::new(&node.az, preferred_az.as_ref())),
     192        52289 :             utilization_score: utilization.cached_score(),
     193        52289 :             total_attached_shard_count: node.attached_shard_count,
     194        52289 :             node_id: *node_id,
     195        52289 :         })
     196        52291 :     }
     197              : 
     198              :     /// For use in scheduling optimisation, where we only want to consider the aspects
     199              :     /// of the score that can only be resolved by moving things (such as inter-shard affinity
     200              :     /// and AZ affinity), and ignore aspects that reflect the total utilization of a node (which
     201              :     /// can fluctuate for other reasons)
     202          118 :     fn for_optimization(&self) -> Self {
     203          118 :         Self {
     204          118 :             utilization_score: 0,
     205          118 :             total_attached_shard_count: 0,
     206          118 :             node_id: NodeId(0),
     207          118 :             ..*self
     208          118 :         }
     209          118 :     }
     210              : 
     211        52170 :     fn is_overloaded(&self) -> bool {
     212        52170 :         PageserverUtilization::is_overloaded(self.utilization_score)
     213        52170 :     }
     214              : 
     215        12880 :     fn node_id(&self) -> NodeId {
     216        12880 :         self.node_id
     217        12880 :     }
     218              : }
     219              : 
     220              : /// Scheduling score of a given node for shard secondaries.
     221              : /// Lower scores indicate more suitable nodes.
     222              : /// Ordering is given by member declaration order (top to bottom).
     223              : #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
     224              : pub(crate) struct NodeSecondarySchedulingScore {
     225              :     /// Flag indicating whether this node matches the preferred AZ
     226              :     /// of the shard. For secondary locations we wish to avoid nodes in.
     227              :     /// the preferred AZ of the shard, since that's where the attached location
     228              :     /// should be scheduled and having the secondary in the same AZ is bad for HA.
     229              :     az_match: SecondaryAzMatch,
     230              :     /// The number of shards belonging to the tenant currently being
     231              :     /// scheduled that are attached to this node.
     232              :     affinity_score: AffinityScore,
     233              :     /// Utilisation score that combines shard count and disk utilisation
     234              :     utilization_score: u64,
     235              :     /// Anti-affinity with other non-home locations: this gives the behavior that secondaries
     236              :     /// will spread out across the nodes in an AZ.
     237              :     total_non_home_shard_count: usize,
     238              :     /// Convenience to make selection deterministic in tests and empty systems
     239              :     node_id: NodeId,
     240              : }
     241              : 
     242              : impl NodeSchedulingScore for NodeSecondarySchedulingScore {
     243        39199 :     fn generate(
     244        39199 :         node_id: &NodeId,
     245        39199 :         node: &mut SchedulerNode,
     246        39199 :         preferred_az: &Option<AvailabilityZone>,
     247        39199 :         context: &ScheduleContext,
     248        39199 :     ) -> Option<Self> {
     249        39199 :         let utilization = match &mut node.may_schedule {
     250        39198 :             MaySchedule::Yes(u) => u,
     251              :             MaySchedule::No => {
     252            1 :                 return None;
     253              :             }
     254              :         };
     255              : 
     256        39198 :         Some(Self {
     257        39198 :             az_match: SecondaryAzMatch(AzMatch::new(&node.az, preferred_az.as_ref())),
     258        39198 :             affinity_score: context
     259        39198 :                 .nodes
     260        39198 :                 .get(node_id)
     261        39198 :                 .copied()
     262        39198 :                 .unwrap_or(AffinityScore::FREE),
     263        39198 :             utilization_score: utilization.cached_score(),
     264        39198 :             total_non_home_shard_count: (node.shard_count - node.home_shard_count),
     265        39198 :             node_id: *node_id,
     266        39198 :         })
     267        39199 :     }
     268              : 
     269           14 :     fn for_optimization(&self) -> Self {
     270           14 :         Self {
     271           14 :             utilization_score: 0,
     272           14 :             total_non_home_shard_count: 0,
     273           14 :             node_id: NodeId(0),
     274           14 :             ..*self
     275           14 :         }
     276           14 :     }
     277              : 
     278        39173 :     fn is_overloaded(&self) -> bool {
     279        39173 :         PageserverUtilization::is_overloaded(self.utilization_score)
     280        39173 :     }
     281              : 
     282        12847 :     fn node_id(&self) -> NodeId {
     283        12847 :         self.node_id
     284        12847 :     }
     285              : }
     286              : 
     287              : impl PartialEq for SchedulerNode {
     288            3 :     fn eq(&self, other: &Self) -> bool {
     289            3 :         let may_schedule_matches = matches!(
     290            3 :             (&self.may_schedule, &other.may_schedule),
     291              :             (MaySchedule::Yes(_), MaySchedule::Yes(_)) | (MaySchedule::No, MaySchedule::No)
     292              :         );
     293              : 
     294            3 :         may_schedule_matches
     295            3 :             && self.shard_count == other.shard_count
     296            3 :             && self.attached_shard_count == other.attached_shard_count
     297            3 :             && self.az == other.az
     298            3 :     }
     299              : }
     300              : 
     301              : impl Eq for SchedulerNode {}
     302              : 
     303              : /// This type is responsible for selecting which node is used when a tenant shard needs to choose a pageserver
     304              : /// on which to run.
     305              : ///
     306              : /// The type has no persistent state of its own: this is all populated at startup.  The Serialize
     307              : /// impl is only for debug dumps.
     308              : #[derive(Serialize)]
     309              : pub(crate) struct Scheduler {
     310              :     nodes: HashMap<NodeId, SchedulerNode>,
     311              : }
     312              : 
     313              : /// Score for soft constraint scheduling: lower scores are preferred to higher scores.
     314              : ///
     315              : /// For example, we may set an affinity score based on the number of shards from the same
     316              : /// tenant already on a node, to implicitly prefer to balance out shards.
     317              : #[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord)]
     318              : pub(crate) struct AffinityScore(pub(crate) usize);
     319              : 
     320              : impl AffinityScore {
     321              :     /// If we have no anti-affinity at all toward a node, this is its score.  It means
     322              :     /// the scheduler has a free choice amongst nodes with this score, and may pick a node
     323              :     /// based on other information such as total utilization.
     324              :     pub(crate) const FREE: Self = Self(0);
     325              : 
     326        25783 :     pub(crate) fn inc(&mut self) {
     327        25783 :         self.0 += 1;
     328        25783 :     }
     329              : 
     330          126 :     pub(crate) fn dec(&mut self) {
     331          126 :         self.0 -= 1;
     332          126 :     }
     333              : }
     334              : 
     335              : impl std::ops::Add for AffinityScore {
     336              :     type Output = Self;
     337              : 
     338            0 :     fn add(self, rhs: Self) -> Self::Output {
     339            0 :         Self(self.0 + rhs.0)
     340            0 :     }
     341              : }
     342              : 
     343              : /// Hint for whether this is a sincere attempt to schedule, or a speculative
     344              : /// check for where we _would_ schedule (done during optimization)
     345              : #[derive(Debug, Clone)]
     346              : pub(crate) enum ScheduleMode {
     347              :     Normal,
     348              :     Speculative,
     349              : }
     350              : 
     351              : impl Default for ScheduleMode {
     352         5334 :     fn default() -> Self {
     353         5334 :         Self::Normal
     354         5334 :     }
     355              : }
     356              : 
     357              : // For carrying state between multiple calls to [`TenantShard::schedule`], e.g. when calling
     358              : // it for many shards in the same tenant.
     359              : #[derive(Debug, Default, Clone)]
     360              : pub(crate) struct ScheduleContext {
     361              :     /// Sparse map of nodes: omitting a node implicitly makes its affinity [`AffinityScore::FREE`]
     362              :     pub(crate) nodes: HashMap<NodeId, AffinityScore>,
     363              : 
     364              :     pub(crate) mode: ScheduleMode,
     365              : }
     366              : 
     367              : impl ScheduleContext {
     368            3 :     pub(crate) fn new(mode: ScheduleMode) -> Self {
     369            3 :         Self {
     370            3 :             nodes: HashMap::new(),
     371            3 :             mode,
     372            3 :         }
     373            3 :     }
     374              : 
     375              :     /// Input is a list of nodes we would like to avoid using again within this context.  The more
     376              :     /// times a node is passed into this call, the less inclined we are to use it.
     377        12900 :     pub(crate) fn avoid(&mut self, nodes: &[NodeId]) {
     378        38683 :         for node_id in nodes {
     379        25783 :             let entry = self.nodes.entry(*node_id).or_insert(AffinityScore::FREE);
     380        25783 :             entry.inc()
     381              :         }
     382        12900 :     }
     383              : 
     384              :     /// Remove `shard`'s contributions to this context.  This is useful when considering scheduling
     385              :     /// this shard afresh, where we don't want it to e.g. experience anti-affinity to its current location.
     386           63 :     pub(crate) fn project_detach(&self, shard: &TenantShard) -> Self {
     387           63 :         let mut new_context = self.clone();
     388              : 
     389           63 :         if let Some(attached) = shard.intent.get_attached() {
     390           60 :             if let Some(score) = new_context.nodes.get_mut(attached) {
     391           60 :                 score.dec();
     392           60 :             }
     393            3 :         }
     394              : 
     395           68 :         for secondary in shard.intent.get_secondary() {
     396           68 :             if let Some(score) = new_context.nodes.get_mut(secondary) {
     397           66 :                 score.dec();
     398           66 :             }
     399              :         }
     400              : 
     401           63 :         new_context
     402           63 :     }
     403              : 
     404              :     /// For test, track the sum of AffinityScore values, which is effectively how many
     405              :     /// attached or secondary locations have been registered with this context.
     406              :     #[cfg(test)]
     407            3 :     pub(crate) fn location_count(&self) -> usize {
     408            7 :         self.nodes.values().map(|i| i.0).sum()
     409            3 :     }
     410              : }
     411              : 
     412              : pub(crate) enum RefCountUpdate<'a> {
     413              :     PromoteSecondary,
     414              :     Attach,
     415              :     Detach,
     416              :     DemoteAttached,
     417              :     AddSecondary,
     418              :     RemoveSecondary,
     419              :     ChangePreferredAzFrom(Option<&'a AvailabilityZone>),
     420              : }
     421              : 
     422              : impl Scheduler {
     423           70 :     pub(crate) fn new<'a>(nodes: impl Iterator<Item = &'a Node>) -> Self {
     424           70 :         let mut scheduler_nodes = HashMap::new();
     425          137 :         for node in nodes {
     426           67 :             scheduler_nodes.insert(
     427           67 :                 node.get_id(),
     428           67 :                 SchedulerNode {
     429           67 :                     shard_count: 0,
     430           67 :                     attached_shard_count: 0,
     431           67 :                     home_shard_count: 0,
     432           67 :                     may_schedule: node.may_schedule(),
     433           67 :                     az: node.get_availability_zone_id().clone(),
     434           67 :                 },
     435           67 :             );
     436           67 :         }
     437              : 
     438           70 :         Self {
     439           70 :             nodes: scheduler_nodes,
     440           70 :         }
     441           70 :     }
     442              : 
     443              :     /// For debug/support: check that our internal statistics are in sync with the state of
     444              :     /// the nodes & tenant shards.
     445              :     ///
     446              :     /// If anything is inconsistent, log details and return an error.
     447            1 :     pub(crate) fn consistency_check<'a>(
     448            1 :         &self,
     449            1 :         nodes: impl Iterator<Item = &'a Node>,
     450            1 :         shards: impl Iterator<Item = &'a TenantShard>,
     451            1 :     ) -> anyhow::Result<()> {
     452            1 :         let mut expect_nodes: HashMap<NodeId, SchedulerNode> = HashMap::new();
     453            4 :         for node in nodes {
     454            3 :             expect_nodes.insert(
     455            3 :                 node.get_id(),
     456            3 :                 SchedulerNode {
     457            3 :                     shard_count: 0,
     458            3 :                     attached_shard_count: 0,
     459            3 :                     home_shard_count: 0,
     460            3 :                     may_schedule: node.may_schedule(),
     461            3 :                     az: node.get_availability_zone_id().clone(),
     462            3 :                 },
     463            3 :             );
     464            3 :         }
     465              : 
     466            2 :         for shard in shards {
     467            1 :             if let Some(node_id) = shard.intent.get_attached() {
     468            1 :                 match expect_nodes.get_mut(node_id) {
     469            1 :                     Some(node) => {
     470            1 :                         node.shard_count += 1;
     471            1 :                         node.attached_shard_count += 1;
     472            1 :                         if Some(&node.az) == shard.preferred_az() {
     473            0 :                             node.home_shard_count += 1;
     474            1 :                         }
     475              :                     }
     476            0 :                     None => anyhow::bail!(
     477            0 :                         "Tenant {} references nonexistent node {}",
     478            0 :                         shard.tenant_shard_id,
     479            0 :                         node_id
     480            0 :                     ),
     481              :                 }
     482            0 :             }
     483              : 
     484            1 :             for node_id in shard.intent.get_secondary() {
     485            1 :                 match expect_nodes.get_mut(node_id) {
     486            1 :                     Some(node) => {
     487            1 :                         node.shard_count += 1;
     488            1 :                         if Some(&node.az) == shard.preferred_az() {
     489            0 :                             node.home_shard_count += 1;
     490            1 :                         }
     491              :                     }
     492            0 :                     None => anyhow::bail!(
     493            0 :                         "Tenant {} references nonexistent node {}",
     494            0 :                         shard.tenant_shard_id,
     495            0 :                         node_id
     496            0 :                     ),
     497              :                 }
     498              :             }
     499              :         }
     500              : 
     501            4 :         for (node_id, expect_node) in &expect_nodes {
     502            3 :             let Some(self_node) = self.nodes.get(node_id) else {
     503            0 :                 anyhow::bail!("Node {node_id} not found in Self")
     504              :             };
     505              : 
     506            3 :             if self_node != expect_node {
     507            0 :                 tracing::error!("Inconsistency detected in scheduling state for node {node_id}");
     508            0 :                 tracing::error!("Expected state: {}", serde_json::to_string(expect_node)?);
     509            0 :                 tracing::error!("Self state: {}", serde_json::to_string(self_node)?);
     510              : 
     511            0 :                 anyhow::bail!("Inconsistent state on {node_id}");
     512            3 :             }
     513              :         }
     514              : 
     515            1 :         if expect_nodes.len() != self.nodes.len() {
     516              :             // We just checked that all the expected nodes are present.  If the lengths don't match,
     517              :             // it means that we have nodes in Self that are unexpected.
     518            0 :             for node_id in self.nodes.keys() {
     519            0 :                 if !expect_nodes.contains_key(node_id) {
     520            0 :                     anyhow::bail!("Node {node_id} found in Self but not in expected nodes");
     521            0 :                 }
     522              :             }
     523            1 :         }
     524              : 
     525            1 :         Ok(())
     526            1 :     }
     527              : 
     528              :     /// Update the reference counts of a node. These reference counts are used to guide scheduling
     529              :     /// decisions, not for memory management: they represent the number of tenant shard whose IntentState
     530              :     /// targets this node and the number of tenants shars whose IntentState is attached to this
     531              :     /// node.
     532              :     ///
     533              :     /// It is an error to call this for a node that is not known to the scheduler (i.e. passed into
     534              :     /// [`Self::new`] or [`Self::node_upsert`])
     535        51432 :     pub(crate) fn update_node_ref_counts(
     536        51432 :         &mut self,
     537        51432 :         node_id: NodeId,
     538        51432 :         preferred_az: Option<&AvailabilityZone>,
     539        51432 :         update: RefCountUpdate,
     540        51432 :     ) {
     541        51432 :         let Some(node) = self.nodes.get_mut(&node_id) else {
     542            0 :             debug_assert!(false);
     543            0 :             tracing::error!("Scheduler missing node {node_id}");
     544            0 :             return;
     545              :         };
     546              : 
     547        51432 :         let is_home_az = Some(&node.az) == preferred_az;
     548        51432 : 
     549        51432 :         match update {
     550            7 :             RefCountUpdate::PromoteSecondary => {
     551            7 :                 node.attached_shard_count += 1;
     552            7 :             }
     553              :             RefCountUpdate::Attach => {
     554        12859 :                 node.shard_count += 1;
     555        12859 :                 node.attached_shard_count += 1;
     556        12859 :                 if is_home_az {
     557        12819 :                     node.home_shard_count += 1;
     558        12819 :                 }
     559              :             }
     560              :             RefCountUpdate::Detach => {
     561        12857 :                 node.shard_count -= 1;
     562        12857 :                 node.attached_shard_count -= 1;
     563        12857 :                 if is_home_az {
     564        12819 :                     node.home_shard_count -= 1;
     565        12819 :                 }
     566              :             }
     567            8 :             RefCountUpdate::DemoteAttached => {
     568            8 :                 node.attached_shard_count -= 1;
     569            8 :             }
     570              :             RefCountUpdate::AddSecondary => {
     571        12848 :                 node.shard_count += 1;
     572        12848 :                 if is_home_az {
     573            6 :                     node.home_shard_count += 1;
     574        12842 :                 }
     575              :             }
     576              :             RefCountUpdate::RemoveSecondary => {
     577        12849 :                 node.shard_count -= 1;
     578        12849 :                 if is_home_az {
     579            5 :                     node.home_shard_count -= 1;
     580        12844 :                 }
     581              :             }
     582            4 :             RefCountUpdate::ChangePreferredAzFrom(old_az) => {
     583            4 :                 if Some(&node.az) == old_az {
     584            2 :                     node.home_shard_count -= 1;
     585            2 :                 }
     586            4 :                 if is_home_az {
     587            1 :                     node.home_shard_count += 1;
     588            3 :                 }
     589              :             }
     590              :         }
     591              : 
     592              :         // Maybe update PageserverUtilization
     593        51432 :         match update {
     594              :             RefCountUpdate::AddSecondary | RefCountUpdate::Attach => {
     595              :                 // Referencing the node: if this takes our shard_count above the utilzation structure's
     596              :                 // shard count, then artifically bump it: this ensures that the scheduler immediately
     597              :                 // recognizes that this node has more work on it, without waiting for the next heartbeat
     598              :                 // to update the utilization.
     599        25707 :                 if let MaySchedule::Yes(utilization) = &mut node.may_schedule {
     600        25707 :                     utilization.adjust_shard_count_max(node.shard_count as u32);
     601        25707 :                 }
     602              :             }
     603              :             RefCountUpdate::PromoteSecondary
     604              :             | RefCountUpdate::Detach
     605              :             | RefCountUpdate::RemoveSecondary
     606              :             | RefCountUpdate::DemoteAttached
     607        25725 :             | RefCountUpdate::ChangePreferredAzFrom(_) => {
     608        25725 :                 // De-referencing the node: leave the utilization's shard_count at a stale higher
     609        25725 :                 // value until some future heartbeat after we have physically removed this shard
     610        25725 :                 // from the node: this prevents the scheduler over-optimistically trying to schedule
     611        25725 :                 // more work onto the node before earlier detaches are done.
     612        25725 :             }
     613              :         }
     614        51432 :     }
     615              : 
     616              :     // Check if the number of shards attached to a given node is lagging below
     617              :     // the cluster average. If that's the case, the node should be filled.
     618            0 :     pub(crate) fn compute_fill_requirement(&self, node_id: NodeId) -> usize {
     619            0 :         let Some(node) = self.nodes.get(&node_id) else {
     620            0 :             debug_assert!(false);
     621            0 :             tracing::error!("Scheduler missing node {node_id}");
     622            0 :             return 0;
     623              :         };
     624            0 :         assert!(!self.nodes.is_empty());
     625            0 :         let expected_attached_shards_per_node = self.expected_attached_shard_count();
     626              : 
     627            0 :         for (node_id, node) in self.nodes.iter() {
     628            0 :             tracing::trace!(%node_id, "attached_shard_count={} shard_count={} expected={}", node.attached_shard_count, node.shard_count, expected_attached_shards_per_node);
     629              :         }
     630              : 
     631            0 :         if node.attached_shard_count < expected_attached_shards_per_node {
     632            0 :             expected_attached_shards_per_node - node.attached_shard_count
     633              :         } else {
     634            0 :             0
     635              :         }
     636            0 :     }
     637              : 
     638            0 :     pub(crate) fn expected_attached_shard_count(&self) -> usize {
     639            0 :         let total_attached_shards: usize =
     640            0 :             self.nodes.values().map(|n| n.attached_shard_count).sum();
     641            0 : 
     642            0 :         assert!(!self.nodes.is_empty());
     643            0 :         total_attached_shards / self.nodes.len()
     644            0 :     }
     645              : 
     646            0 :     pub(crate) fn nodes_by_attached_shard_count(&self) -> Vec<(NodeId, usize)> {
     647            0 :         self.nodes
     648            0 :             .iter()
     649            0 :             .map(|(node_id, stats)| (*node_id, stats.attached_shard_count))
     650            0 :             .sorted_by(|lhs, rhs| Ord::cmp(&lhs.1, &rhs.1).reverse())
     651            0 :             .collect()
     652            0 :     }
     653              : 
     654          215 :     pub(crate) fn node_upsert(&mut self, node: &Node) {
     655              :         use std::collections::hash_map::Entry::*;
     656          215 :         match self.nodes.entry(node.get_id()) {
     657            4 :             Occupied(mut entry) => {
     658            4 :                 // Updates to MaySchedule are how we receive updated PageserverUtilization: adjust these values
     659            4 :                 // to account for any shards scheduled on the controller but not yet visible to the pageserver.
     660            4 :                 let mut may_schedule = node.may_schedule();
     661            4 :                 match &mut may_schedule {
     662            2 :                     MaySchedule::Yes(utilization) => {
     663            2 :                         utilization.adjust_shard_count_max(entry.get().shard_count as u32);
     664            2 :                     }
     665            2 :                     MaySchedule::No => { // Nothing to tweak
     666            2 :                     }
     667              :                 }
     668              : 
     669            4 :                 entry.get_mut().may_schedule = may_schedule;
     670              :             }
     671          211 :             Vacant(entry) => {
     672          211 :                 entry.insert(SchedulerNode {
     673          211 :                     shard_count: 0,
     674          211 :                     attached_shard_count: 0,
     675          211 :                     home_shard_count: 0,
     676          211 :                     may_schedule: node.may_schedule(),
     677          211 :                     az: node.get_availability_zone_id().clone(),
     678          211 :                 });
     679          211 :             }
     680              :         }
     681          215 :     }
     682              : 
     683            0 :     pub(crate) fn node_remove(&mut self, node_id: NodeId) {
     684            0 :         if self.nodes.remove(&node_id).is_none() {
     685            0 :             tracing::warn!(node_id=%node_id, "Removed non-existent node from scheduler");
     686            0 :         }
     687            0 :     }
     688              : 
     689              :     /// Calculate a single node's score, used in optimizer logic to compare specific
     690              :     /// nodes' scores.
     691          146 :     pub(crate) fn compute_node_score<Score>(
     692          146 :         &mut self,
     693          146 :         node_id: NodeId,
     694          146 :         preferred_az: &Option<AvailabilityZone>,
     695          146 :         context: &ScheduleContext,
     696          146 :     ) -> Option<Score>
     697          146 :     where
     698          146 :         Score: NodeSchedulingScore,
     699          146 :     {
     700          146 :         self.nodes
     701          146 :             .get_mut(&node_id)
     702          146 :             .and_then(|node| Score::generate(&node_id, node, preferred_az, context))
     703          146 :     }
     704              : 
     705              :     /// Compute a schedulling score for each node that the scheduler knows of
     706              :     /// minus a set of hard excluded nodes.
     707        25727 :     fn compute_node_scores<Score>(
     708        25727 :         &mut self,
     709        25727 :         hard_exclude: &[NodeId],
     710        25727 :         preferred_az: &Option<AvailabilityZone>,
     711        25727 :         context: &ScheduleContext,
     712        25727 :     ) -> Vec<Score>
     713        25727 :     where
     714        25727 :         Score: NodeSchedulingScore,
     715        25727 :     {
     716        25727 :         self.nodes
     717        25727 :             .iter_mut()
     718       104191 :             .filter_map(|(k, v)| {
     719       104191 :                 if hard_exclude.contains(k) {
     720        12847 :                     None
     721              :                 } else {
     722        91344 :                     Score::generate(k, v, preferred_az, context)
     723              :                 }
     724       104191 :             })
     725        25727 :             .collect()
     726        25727 :     }
     727              : 
     728              :     /// hard_exclude: it is forbidden to use nodes in this list, typically becacuse they
     729              :     /// are already in use by this shard -- we use this to avoid picking the same node
     730              :     /// as both attached and secondary location.  This is a hard constraint: if we cannot
     731              :     /// find any nodes that aren't in this list, then we will return a [`ScheduleError::ImpossibleConstraint`].
     732              :     ///
     733              :     /// context: we prefer to avoid using nodes identified in the context, according
     734              :     /// to their anti-affinity score.  We use this to prefeer to avoid placing shards in
     735              :     /// the same tenant on the same node.  This is a soft constraint: the context will never
     736              :     /// cause us to fail to schedule a shard.
     737        25727 :     pub(crate) fn schedule_shard<Tag: ShardTag>(
     738        25727 :         &mut self,
     739        25727 :         hard_exclude: &[NodeId],
     740        25727 :         preferred_az: &Option<AvailabilityZone>,
     741        25727 :         context: &ScheduleContext,
     742        25727 :     ) -> Result<NodeId, ScheduleError> {
     743        25727 :         if self.nodes.is_empty() {
     744            0 :             return Err(ScheduleError::NoPageservers);
     745        25727 :         }
     746        25727 : 
     747        25727 :         let mut scores =
     748        25727 :             self.compute_node_scores::<Tag::Score>(hard_exclude, preferred_az, context);
     749        25727 : 
     750        25727 :         // Exclude nodes whose utilization is critically high, if there are alternatives available.  This will
     751        25727 :         // cause us to violate affinity rules if it is necessary to avoid critically overloading nodes: for example
     752        25727 :         // we may place shards in the same tenant together on the same pageserver if all other pageservers are
     753        25727 :         // overloaded.
     754        25727 :         let non_overloaded_scores = scores
     755        25727 :             .iter()
     756        91343 :             .filter(|i| !i.is_overloaded())
     757        25727 :             .copied()
     758        25727 :             .collect::<Vec<_>>();
     759        25727 :         if !non_overloaded_scores.is_empty() {
     760        25727 :             scores = non_overloaded_scores;
     761        25727 :         }
     762              : 
     763              :         // Sort the nodes by score. The one with the lowest scores will be the preferred node.
     764              :         // Refer to [`NodeAttachmentSchedulingScore`] for attached locations and
     765              :         // [`NodeSecondarySchedulingScore`] for secondary locations to understand how the nodes
     766              :         // are ranked.
     767        25727 :         scores.sort();
     768        25727 : 
     769        25727 :         if scores.is_empty() {
     770              :             // After applying constraints, no pageservers were left.
     771            0 :             if !matches!(context.mode, ScheduleMode::Speculative) {
     772              :                 // If this was not a speculative attempt, log details to understand why we couldn't
     773              :                 // schedule: this may help an engineer understand if some nodes are marked offline
     774              :                 // in a way that's preventing progress.
     775            0 :                 tracing::info!(
     776            0 :                     "Scheduling failure, while excluding {hard_exclude:?}, node states:"
     777              :                 );
     778            0 :                 for (node_id, node) in &self.nodes {
     779            0 :                     tracing::info!(
     780            0 :                         "Node {node_id}: may_schedule={} shards={}",
     781            0 :                         !matches!(node.may_schedule, MaySchedule::No),
     782              :                         node.shard_count
     783              :                     );
     784              :                 }
     785            0 :             }
     786            0 :             return Err(ScheduleError::ImpossibleConstraint);
     787        25727 :         }
     788        25727 : 
     789        25727 :         // Lowest score wins
     790        25727 :         let node_id = scores.first().unwrap().node_id();
     791              : 
     792        25727 :         if !matches!(context.mode, ScheduleMode::Speculative) {
     793        25727 :             tracing::info!(
     794            0 :                 "scheduler selected node {node_id} (elegible nodes {:?}, hard exclude: {hard_exclude:?}, soft exclude: {context:?}, preferred_az: {:?})",
     795            0 :                 scores.iter().map(|i| i.node_id().0).collect::<Vec<_>>(),
     796              :                 preferred_az,
     797              :             );
     798            0 :         }
     799              : 
     800              :         // Note that we do not update shard count here to reflect the scheduling: that
     801              :         // is IntentState's job when the scheduled location is used.
     802              : 
     803        25727 :         Ok(node_id)
     804        25727 :     }
     805              : 
     806              :     /// Selects any available node. This is suitable for performing background work (e.g. S3
     807              :     /// deletions).
     808            0 :     pub(crate) fn any_available_node(&mut self) -> Result<NodeId, ScheduleError> {
     809            0 :         self.schedule_shard::<AttachedShardTag>(&[], &None, &ScheduleContext::default())
     810            0 :     }
     811              : 
     812              :     /// For choosing which AZ to schedule a new shard into, use this.  It will return the
     813              :     /// AZ with the the lowest number of shards currently scheduled in this AZ as their home
     814              :     /// location.
     815              :     ///
     816              :     /// We use an AZ-wide measure rather than simply selecting the AZ of the least-loaded
     817              :     /// node, because while tenants start out single sharded, when they grow and undergo
     818              :     /// shard-split, they will occupy space on many nodes within an AZ.  It is important
     819              :     /// that we pick the AZ in a way that balances this _future_ load.
     820              :     ///
     821              :     /// Once we've picked an AZ, subsequent scheduling within that AZ will be driven by
     822              :     /// nodes' utilization scores.
     823          303 :     pub(crate) fn get_az_for_new_tenant(&self) -> Option<AvailabilityZone> {
     824          303 :         if self.nodes.is_empty() {
     825            0 :             return None;
     826          303 :         }
     827              : 
     828              :         #[derive(Default)]
     829              :         struct AzScore {
     830              :             home_shard_count: usize,
     831              :             scheduleable: bool,
     832              :         }
     833              : 
     834          303 :         let mut azs: HashMap<&AvailabilityZone, AzScore> = HashMap::new();
     835         1818 :         for node in self.nodes.values() {
     836         1818 :             let az = azs.entry(&node.az).or_default();
     837         1818 :             az.home_shard_count += node.home_shard_count;
     838         1818 :             az.scheduleable |= matches!(node.may_schedule, MaySchedule::Yes(_));
     839              :         }
     840              : 
     841              :         // If any AZs are schedulable, then filter out the non-schedulable ones (i.e. AZs where
     842              :         // all nodes are overloaded or otherwise unschedulable).
     843          303 :         if azs.values().any(|i| i.scheduleable) {
     844          906 :             azs.retain(|_, i| i.scheduleable);
     845          303 :         }
     846              : 
     847              :         // Find the AZ with the lowest number of shards currently allocated
     848          303 :         Some(
     849          303 :             azs.into_iter()
     850          906 :                 .min_by_key(|i| (i.1.home_shard_count, i.0))
     851          303 :                 .unwrap()
     852          303 :                 .0
     853          303 :                 .clone(),
     854          303 :         )
     855          303 :     }
     856              : 
     857            9 :     pub(crate) fn get_node_az(&self, node_id: &NodeId) -> Option<AvailabilityZone> {
     858            9 :         self.nodes.get(node_id).map(|n| n.az.clone())
     859            9 :     }
     860              : 
     861              :     /// For use when choosing a preferred secondary location: filter out nodes that are not
     862              :     /// available, and gather their AZs.
     863        12832 :     pub(crate) fn filter_usable_nodes(
     864        12832 :         &self,
     865        12832 :         nodes: &[NodeId],
     866        12832 :     ) -> Vec<(NodeId, Option<AvailabilityZone>)> {
     867        12832 :         nodes
     868        12832 :             .iter()
     869        12832 :             .filter_map(|node_id| {
     870            3 :                 let node = self
     871            3 :                     .nodes
     872            3 :                     .get(node_id)
     873            3 :                     .expect("Referenced nodes always exist");
     874            3 :                 if matches!(node.may_schedule, MaySchedule::Yes(_)) {
     875            2 :                     Some((*node_id, Some(node.az.clone())))
     876              :                 } else {
     877            1 :                     None
     878              :                 }
     879        12832 :             })
     880        12832 :             .collect()
     881        12832 :     }
     882              : 
     883              :     /// Unit test access to internal state
     884              :     #[cfg(test)]
     885           19 :     pub(crate) fn get_node_shard_count(&self, node_id: NodeId) -> usize {
     886           19 :         self.nodes.get(&node_id).unwrap().shard_count
     887           19 :     }
     888              : 
     889              :     #[cfg(test)]
     890           19 :     pub(crate) fn get_node_attached_shard_count(&self, node_id: NodeId) -> usize {
     891           19 :         self.nodes.get(&node_id).unwrap().attached_shard_count
     892           19 :     }
     893              : 
     894              :     /// Some metrics that we only calculate periodically: this is simpler than
     895              :     /// rigorously updating them on every change.
     896            0 :     pub(crate) fn update_metrics(&self) {
     897            0 :         for (node_id, node) in &self.nodes {
     898            0 :             let node_id_str = format!("{}", node_id);
     899            0 :             let label_group = NodeLabelGroup {
     900            0 :                 az: &node.az.0,
     901            0 :                 node_id: &node_id_str,
     902            0 :             };
     903            0 : 
     904            0 :             crate::metrics::METRICS_REGISTRY
     905            0 :                 .metrics_group
     906            0 :                 .storage_controller_node_shards
     907            0 :                 .set(label_group.clone(), node.shard_count as i64);
     908            0 : 
     909            0 :             crate::metrics::METRICS_REGISTRY
     910            0 :                 .metrics_group
     911            0 :                 .storage_controller_node_attached_shards
     912            0 :                 .set(label_group.clone(), node.attached_shard_count as i64);
     913            0 : 
     914            0 :             crate::metrics::METRICS_REGISTRY
     915            0 :                 .metrics_group
     916            0 :                 .storage_controller_node_home_shards
     917            0 :                 .set(label_group.clone(), node.home_shard_count as i64);
     918            0 :         }
     919            0 :     }
     920              : }
     921              : 
     922              : #[cfg(test)]
     923              : pub(crate) mod test_utils {
     924              : 
     925              :     use std::collections::HashMap;
     926              : 
     927              :     use pageserver_api::controller_api::{AvailabilityZone, NodeAvailability};
     928              :     use pageserver_api::models::utilization::test_utilization;
     929              :     use utils::id::NodeId;
     930              : 
     931              :     use crate::node::Node;
     932              : 
     933              :     /// Test helper: synthesize the requested number of nodes, all in active state.
     934              :     ///
     935              :     /// Node IDs start at one.
     936              :     ///
     937              :     /// The `azs` argument specifies the list of availability zones which will be assigned
     938              :     /// to nodes in round-robin fashion. If empy, a default AZ is assigned.
     939           70 :     pub(crate) fn make_test_nodes(n: u64, azs: &[AvailabilityZone]) -> HashMap<NodeId, Node> {
     940           70 :         let mut az_iter = azs.iter().cycle();
     941           70 : 
     942           70 :         (1..n + 1)
     943          278 :             .map(|i| {
     944          278 :                 (NodeId(i), {
     945          278 :                     let mut node = Node::new(
     946          278 :                         NodeId(i),
     947          278 :                         format!("httphost-{i}"),
     948          278 :                         80 + i as u16,
     949          278 :                         None,
     950          278 :                         format!("pghost-{i}"),
     951          278 :                         5432 + i as u16,
     952          278 :                         az_iter
     953          278 :                             .next()
     954          278 :                             .cloned()
     955          278 :                             .unwrap_or(AvailabilityZone("test-az".to_string())),
     956          278 :                         false,
     957          278 :                     )
     958          278 :                     .unwrap();
     959          278 :                     node.set_availability(NodeAvailability::Active(test_utilization::simple(0, 0)));
     960          278 :                     assert!(node.is_available());
     961          278 :                     node
     962          278 :                 })
     963          278 :             })
     964           70 :             .collect()
     965           70 :     }
     966              : }
     967              : 
     968              : #[cfg(test)]
     969              : mod tests {
     970              :     use pageserver_api::controller_api::NodeAvailability;
     971              :     use pageserver_api::models::utilization::test_utilization;
     972              :     use pageserver_api::shard::ShardIdentity;
     973              :     use utils::id::TenantId;
     974              :     use utils::shard::{ShardCount, ShardNumber, TenantShardId};
     975              : 
     976              :     use super::*;
     977              :     use crate::tenant_shard::IntentState;
     978              :     #[test]
     979            1 :     fn scheduler_basic() -> anyhow::Result<()> {
     980            1 :         let nodes = test_utils::make_test_nodes(2, &[]);
     981            1 : 
     982            1 :         let mut scheduler = Scheduler::new(nodes.values());
     983            1 :         let mut t1_intent = IntentState::new(None);
     984            1 :         let mut t2_intent = IntentState::new(None);
     985            1 : 
     986            1 :         let context = ScheduleContext::default();
     987              : 
     988            1 :         let scheduled = scheduler.schedule_shard::<AttachedShardTag>(&[], &None, &context)?;
     989            1 :         t1_intent.set_attached(&mut scheduler, Some(scheduled));
     990            1 :         let scheduled = scheduler.schedule_shard::<AttachedShardTag>(&[], &None, &context)?;
     991            1 :         t2_intent.set_attached(&mut scheduler, Some(scheduled));
     992            1 : 
     993            1 :         assert_eq!(scheduler.get_node_shard_count(NodeId(1)), 1);
     994            1 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(1)), 1);
     995              : 
     996            1 :         assert_eq!(scheduler.get_node_shard_count(NodeId(2)), 1);
     997            1 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(2)), 1);
     998              : 
     999            1 :         let scheduled = scheduler.schedule_shard::<AttachedShardTag>(
    1000            1 :             &t1_intent.all_pageservers(),
    1001            1 :             &None,
    1002            1 :             &context,
    1003            1 :         )?;
    1004            1 :         t1_intent.push_secondary(&mut scheduler, scheduled);
    1005            1 : 
    1006            1 :         assert_eq!(scheduler.get_node_shard_count(NodeId(1)), 1);
    1007            1 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(1)), 1);
    1008              : 
    1009            1 :         assert_eq!(scheduler.get_node_shard_count(NodeId(2)), 2);
    1010            1 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(2)), 1);
    1011              : 
    1012            1 :         t1_intent.clear(&mut scheduler);
    1013            1 :         assert_eq!(scheduler.get_node_shard_count(NodeId(1)), 0);
    1014            1 :         assert_eq!(scheduler.get_node_shard_count(NodeId(2)), 1);
    1015              : 
    1016            1 :         let total_attached = scheduler.get_node_attached_shard_count(NodeId(1))
    1017            1 :             + scheduler.get_node_attached_shard_count(NodeId(2));
    1018            1 :         assert_eq!(total_attached, 1);
    1019              : 
    1020            1 :         if cfg!(debug_assertions) {
    1021              :             // Dropping an IntentState without clearing it causes a panic in debug mode,
    1022              :             // because we have failed to properly update scheduler shard counts.
    1023            1 :             let result = std::panic::catch_unwind(move || {
    1024            1 :                 drop(t2_intent);
    1025            1 :             });
    1026            1 :             assert!(result.is_err());
    1027              :         } else {
    1028            0 :             t2_intent.clear(&mut scheduler);
    1029            0 : 
    1030            0 :             assert_eq!(scheduler.get_node_shard_count(NodeId(1)), 0);
    1031            0 :             assert_eq!(scheduler.get_node_attached_shard_count(NodeId(1)), 0);
    1032              : 
    1033            0 :             assert_eq!(scheduler.get_node_shard_count(NodeId(2)), 0);
    1034            0 :             assert_eq!(scheduler.get_node_attached_shard_count(NodeId(2)), 0);
    1035              :         }
    1036              : 
    1037            1 :         Ok(())
    1038            1 :     }
    1039              : 
    1040              :     #[test]
    1041              :     /// Test the PageserverUtilization's contribution to scheduling algorithm
    1042            1 :     fn scheduler_utilization() {
    1043            1 :         let mut nodes = test_utils::make_test_nodes(3, &[]);
    1044            1 :         let mut scheduler = Scheduler::new(nodes.values());
    1045            1 : 
    1046            1 :         // Need to keep these alive because they contribute to shard counts via RAII
    1047            1 :         let mut scheduled_intents = Vec::new();
    1048            1 : 
    1049            1 :         let empty_context = ScheduleContext::default();
    1050              : 
    1051           11 :         fn assert_scheduler_chooses(
    1052           11 :             expect_node: NodeId,
    1053           11 :             scheduled_intents: &mut Vec<IntentState>,
    1054           11 :             scheduler: &mut Scheduler,
    1055           11 :             context: &ScheduleContext,
    1056           11 :         ) {
    1057           11 :             let scheduled = scheduler
    1058           11 :                 .schedule_shard::<AttachedShardTag>(&[], &None, context)
    1059           11 :                 .unwrap();
    1060           11 :             let mut intent = IntentState::new(None);
    1061           11 :             intent.set_attached(scheduler, Some(scheduled));
    1062           11 :             scheduled_intents.push(intent);
    1063           11 :             assert_eq!(scheduled, expect_node);
    1064           11 :         }
    1065              : 
    1066              :         // Independent schedule calls onto empty nodes should round-robin, because each node's
    1067              :         // utilization's shard count is updated inline.  The order is determinsitic because when all other factors are
    1068              :         // equal, we order by node ID.
    1069            1 :         assert_scheduler_chooses(
    1070            1 :             NodeId(1),
    1071            1 :             &mut scheduled_intents,
    1072            1 :             &mut scheduler,
    1073            1 :             &empty_context,
    1074            1 :         );
    1075            1 :         assert_scheduler_chooses(
    1076            1 :             NodeId(2),
    1077            1 :             &mut scheduled_intents,
    1078            1 :             &mut scheduler,
    1079            1 :             &empty_context,
    1080            1 :         );
    1081            1 :         assert_scheduler_chooses(
    1082            1 :             NodeId(3),
    1083            1 :             &mut scheduled_intents,
    1084            1 :             &mut scheduler,
    1085            1 :             &empty_context,
    1086            1 :         );
    1087            1 : 
    1088            1 :         // Manually setting utilization higher should cause schedule calls to round-robin the other nodes
    1089            1 :         // which have equal utilization.
    1090            1 :         nodes
    1091            1 :             .get_mut(&NodeId(1))
    1092            1 :             .unwrap()
    1093            1 :             .set_availability(NodeAvailability::Active(test_utilization::simple(
    1094            1 :                 10,
    1095            1 :                 1024 * 1024 * 1024,
    1096            1 :             )));
    1097            1 :         scheduler.node_upsert(nodes.get(&NodeId(1)).unwrap());
    1098            1 : 
    1099            1 :         assert_scheduler_chooses(
    1100            1 :             NodeId(2),
    1101            1 :             &mut scheduled_intents,
    1102            1 :             &mut scheduler,
    1103            1 :             &empty_context,
    1104            1 :         );
    1105            1 :         assert_scheduler_chooses(
    1106            1 :             NodeId(3),
    1107            1 :             &mut scheduled_intents,
    1108            1 :             &mut scheduler,
    1109            1 :             &empty_context,
    1110            1 :         );
    1111            1 :         assert_scheduler_chooses(
    1112            1 :             NodeId(2),
    1113            1 :             &mut scheduled_intents,
    1114            1 :             &mut scheduler,
    1115            1 :             &empty_context,
    1116            1 :         );
    1117            1 :         assert_scheduler_chooses(
    1118            1 :             NodeId(3),
    1119            1 :             &mut scheduled_intents,
    1120            1 :             &mut scheduler,
    1121            1 :             &empty_context,
    1122            1 :         );
    1123            1 : 
    1124            1 :         // The scheduler should prefer nodes with lower affinity score,
    1125            1 :         // even if they have higher utilization (as long as they aren't utilized at >100%)
    1126            1 :         let mut context_prefer_node1 = ScheduleContext::default();
    1127            1 :         context_prefer_node1.avoid(&[NodeId(2), NodeId(3)]);
    1128            1 :         assert_scheduler_chooses(
    1129            1 :             NodeId(1),
    1130            1 :             &mut scheduled_intents,
    1131            1 :             &mut scheduler,
    1132            1 :             &context_prefer_node1,
    1133            1 :         );
    1134            1 :         assert_scheduler_chooses(
    1135            1 :             NodeId(1),
    1136            1 :             &mut scheduled_intents,
    1137            1 :             &mut scheduler,
    1138            1 :             &context_prefer_node1,
    1139            1 :         );
    1140            1 : 
    1141            1 :         // If a node is over-utilized, it will not be used even if affinity scores prefer it
    1142            1 :         nodes
    1143            1 :             .get_mut(&NodeId(1))
    1144            1 :             .unwrap()
    1145            1 :             .set_availability(NodeAvailability::Active(test_utilization::simple(
    1146            1 :                 20000,
    1147            1 :                 1024 * 1024 * 1024,
    1148            1 :             )));
    1149            1 :         scheduler.node_upsert(nodes.get(&NodeId(1)).unwrap());
    1150            1 :         assert_scheduler_chooses(
    1151            1 :             NodeId(2),
    1152            1 :             &mut scheduled_intents,
    1153            1 :             &mut scheduler,
    1154            1 :             &context_prefer_node1,
    1155            1 :         );
    1156            1 :         assert_scheduler_chooses(
    1157            1 :             NodeId(3),
    1158            1 :             &mut scheduled_intents,
    1159            1 :             &mut scheduler,
    1160            1 :             &context_prefer_node1,
    1161            1 :         );
    1162              : 
    1163           12 :         for mut intent in scheduled_intents {
    1164           11 :             intent.clear(&mut scheduler);
    1165           11 :         }
    1166            1 :     }
    1167              : 
    1168              :     #[test]
    1169              :     /// A simple test that showcases AZ-aware scheduling and its interaction with
    1170              :     /// affinity scores.
    1171            1 :     fn az_scheduling() {
    1172            1 :         let az_a_tag = AvailabilityZone("az-a".to_string());
    1173            1 :         let az_b_tag = AvailabilityZone("az-b".to_string());
    1174            1 : 
    1175            1 :         let nodes = test_utils::make_test_nodes(3, &[az_a_tag.clone(), az_b_tag.clone()]);
    1176            1 :         let mut scheduler = Scheduler::new(nodes.values());
    1177            1 : 
    1178            1 :         // Need to keep these alive because they contribute to shard counts via RAII
    1179            1 :         let mut scheduled_intents = Vec::new();
    1180            1 : 
    1181            1 :         let mut context = ScheduleContext::default();
    1182              : 
    1183            4 :         fn assert_scheduler_chooses<Tag: ShardTag>(
    1184            4 :             expect_node: NodeId,
    1185            4 :             preferred_az: Option<AvailabilityZone>,
    1186            4 :             scheduled_intents: &mut Vec<IntentState>,
    1187            4 :             scheduler: &mut Scheduler,
    1188            4 :             context: &mut ScheduleContext,
    1189            4 :         ) {
    1190            4 :             let scheduled = scheduler
    1191            4 :                 .schedule_shard::<Tag>(&[], &preferred_az, context)
    1192            4 :                 .unwrap();
    1193            4 :             let mut intent = IntentState::new(preferred_az.clone());
    1194            4 :             intent.set_attached(scheduler, Some(scheduled));
    1195            4 :             scheduled_intents.push(intent);
    1196            4 :             assert_eq!(scheduled, expect_node);
    1197              : 
    1198            4 :             context.avoid(&[scheduled]);
    1199            4 :         }
    1200              : 
    1201            1 :         assert_scheduler_chooses::<AttachedShardTag>(
    1202            1 :             NodeId(1),
    1203            1 :             Some(az_a_tag.clone()),
    1204            1 :             &mut scheduled_intents,
    1205            1 :             &mut scheduler,
    1206            1 :             &mut context,
    1207            1 :         );
    1208            1 : 
    1209            1 :         // Node 2 and 3 have affinity score equal to 0, but node 3
    1210            1 :         // is in "az-a" so we prefer that.
    1211            1 :         assert_scheduler_chooses::<AttachedShardTag>(
    1212            1 :             NodeId(3),
    1213            1 :             Some(az_a_tag.clone()),
    1214            1 :             &mut scheduled_intents,
    1215            1 :             &mut scheduler,
    1216            1 :             &mut context,
    1217            1 :         );
    1218            1 : 
    1219            1 :         // Node 1 and 3 (az-a) have same affinity score, so prefer the lowest node id.
    1220            1 :         assert_scheduler_chooses::<AttachedShardTag>(
    1221            1 :             NodeId(1),
    1222            1 :             Some(az_a_tag.clone()),
    1223            1 :             &mut scheduled_intents,
    1224            1 :             &mut scheduler,
    1225            1 :             &mut context,
    1226            1 :         );
    1227            1 : 
    1228            1 :         // Avoid nodes in "az-a" for the secondary location.
    1229            1 :         assert_scheduler_chooses::<SecondaryShardTag>(
    1230            1 :             NodeId(2),
    1231            1 :             Some(az_a_tag.clone()),
    1232            1 :             &mut scheduled_intents,
    1233            1 :             &mut scheduler,
    1234            1 :             &mut context,
    1235            1 :         );
    1236              : 
    1237            5 :         for mut intent in scheduled_intents {
    1238            4 :             intent.clear(&mut scheduler);
    1239            4 :         }
    1240            1 :     }
    1241              : 
    1242              :     #[test]
    1243            1 :     fn az_scheduling_for_new_tenant() {
    1244            1 :         let az_a_tag = AvailabilityZone("az-a".to_string());
    1245            1 :         let az_b_tag = AvailabilityZone("az-b".to_string());
    1246            1 :         let nodes = test_utils::make_test_nodes(
    1247            1 :             6,
    1248            1 :             &[
    1249            1 :                 az_a_tag.clone(),
    1250            1 :                 az_a_tag.clone(),
    1251            1 :                 az_a_tag.clone(),
    1252            1 :                 az_b_tag.clone(),
    1253            1 :                 az_b_tag.clone(),
    1254            1 :                 az_b_tag.clone(),
    1255            1 :             ],
    1256            1 :         );
    1257            1 : 
    1258            1 :         let mut scheduler = Scheduler::new(nodes.values());
    1259              : 
    1260              :         /// Force the `home_shard_count` of a node directly: this is the metric used
    1261              :         /// by the scheduler when picking AZs.
    1262            3 :         fn set_shard_count(scheduler: &mut Scheduler, node_id: NodeId, shard_count: usize) {
    1263            3 :             let node = scheduler.nodes.get_mut(&node_id).unwrap();
    1264            3 :             node.home_shard_count = shard_count;
    1265            3 :         }
    1266              : 
    1267              :         // Initial empty state.  Scores are tied, scheduler prefers lower AZ ID.
    1268            1 :         assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_a_tag.clone()));
    1269              : 
    1270              :         // Home shard count is higher in AZ A, so AZ B will be preferred
    1271            1 :         set_shard_count(&mut scheduler, NodeId(1), 10);
    1272            1 :         assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_b_tag.clone()));
    1273              : 
    1274              :         // Total home shard count is higher in AZ B, so we revert to preferring AZ A
    1275            1 :         set_shard_count(&mut scheduler, NodeId(4), 6);
    1276            1 :         set_shard_count(&mut scheduler, NodeId(5), 6);
    1277            1 :         assert_eq!(scheduler.get_az_for_new_tenant(), Some(az_a_tag.clone()));
    1278            1 :     }
    1279              : 
    1280              :     /// Test that when selecting AZs for many new tenants, we get the expected balance across nodes
    1281              :     #[test]
    1282            1 :     fn az_selection_many() {
    1283            1 :         let az_a_tag = AvailabilityZone("az-a".to_string());
    1284            1 :         let az_b_tag = AvailabilityZone("az-b".to_string());
    1285            1 :         let az_c_tag = AvailabilityZone("az-c".to_string());
    1286            1 :         let nodes = test_utils::make_test_nodes(
    1287            1 :             6,
    1288            1 :             &[
    1289            1 :                 az_a_tag.clone(),
    1290            1 :                 az_b_tag.clone(),
    1291            1 :                 az_c_tag.clone(),
    1292            1 :                 az_a_tag.clone(),
    1293            1 :                 az_b_tag.clone(),
    1294            1 :                 az_c_tag.clone(),
    1295            1 :             ],
    1296            1 :         );
    1297            1 : 
    1298            1 :         let mut scheduler = Scheduler::new(nodes.values());
    1299            1 : 
    1300            1 :         // We should get 1/6th of these on each node, give or take a few...
    1301            1 :         let total_tenants = 300;
    1302            1 : 
    1303            1 :         // ...where the 'few' is the number of AZs, because the scheduling will sometimes overshoot
    1304            1 :         // on one AZ before correcting itself.  This is because we select the 'home' AZ based on
    1305            1 :         // an AZ-wide metric, but we select the location for secondaries on a purely node-based
    1306            1 :         // metric (while excluding the home AZ).
    1307            1 :         let grace = 3;
    1308            1 : 
    1309            1 :         let mut scheduled_shards = Vec::new();
    1310          300 :         for _i in 0..total_tenants {
    1311          300 :             let preferred_az = scheduler.get_az_for_new_tenant().unwrap();
    1312          300 : 
    1313          300 :             let mut node_home_counts = scheduler
    1314          300 :                 .nodes
    1315          300 :                 .iter()
    1316         1800 :                 .map(|(node_id, node)| (node_id, node.home_shard_count))
    1317          300 :                 .collect::<Vec<_>>();
    1318         6000 :             node_home_counts.sort_by_key(|i| i.0);
    1319          300 :             eprintln!("Selected {}, vs nodes {:?}", preferred_az, node_home_counts);
    1320          300 : 
    1321          300 :             let tenant_shard_id = TenantShardId {
    1322          300 :                 tenant_id: TenantId::generate(),
    1323          300 :                 shard_number: ShardNumber(0),
    1324          300 :                 shard_count: ShardCount(1),
    1325          300 :             };
    1326          300 : 
    1327          300 :             let shard_identity = ShardIdentity::new(
    1328          300 :                 tenant_shard_id.shard_number,
    1329          300 :                 tenant_shard_id.shard_count,
    1330          300 :                 pageserver_api::shard::ShardStripeSize(1),
    1331          300 :             )
    1332          300 :             .unwrap();
    1333          300 :             let mut shard = TenantShard::new(
    1334          300 :                 tenant_shard_id,
    1335          300 :                 shard_identity,
    1336          300 :                 pageserver_api::controller_api::PlacementPolicy::Attached(1),
    1337          300 :                 Some(preferred_az),
    1338          300 :             );
    1339          300 : 
    1340          300 :             let mut context = ScheduleContext::default();
    1341          300 :             shard.schedule(&mut scheduler, &mut context).unwrap();
    1342          300 :             eprintln!("Scheduled shard at {:?}", shard.intent);
    1343          300 : 
    1344          300 :             scheduled_shards.push(shard);
    1345          300 :         }
    1346              : 
    1347            7 :         for (node_id, node) in &scheduler.nodes {
    1348            6 :             eprintln!(
    1349            6 :                 "Node {}: {} {} {}",
    1350            6 :                 node_id, node.shard_count, node.attached_shard_count, node.home_shard_count
    1351            6 :             );
    1352            6 :         }
    1353              : 
    1354            6 :         for node in scheduler.nodes.values() {
    1355            6 :             assert!((node.home_shard_count as i64 - total_tenants as i64 / 6).abs() < grace);
    1356              :         }
    1357              : 
    1358          301 :         for mut shard in scheduled_shards {
    1359          300 :             shard.intent.clear(&mut scheduler);
    1360          300 :         }
    1361            1 :     }
    1362              : 
    1363              :     #[test]
    1364              :     /// Make sure that when we have an odd number of nodes and an even number of shards, we still
    1365              :     /// get scheduling stability.
    1366            1 :     fn odd_nodes_stability() {
    1367            1 :         let az_a = AvailabilityZone("az-a".to_string());
    1368            1 :         let az_b = AvailabilityZone("az-b".to_string());
    1369            1 : 
    1370            1 :         let nodes = test_utils::make_test_nodes(
    1371            1 :             10,
    1372            1 :             &[
    1373            1 :                 az_a.clone(),
    1374            1 :                 az_a.clone(),
    1375            1 :                 az_a.clone(),
    1376            1 :                 az_a.clone(),
    1377            1 :                 az_a.clone(),
    1378            1 :                 az_b.clone(),
    1379            1 :                 az_b.clone(),
    1380            1 :                 az_b.clone(),
    1381            1 :                 az_b.clone(),
    1382            1 :                 az_b.clone(),
    1383            1 :             ],
    1384            1 :         );
    1385            1 :         let mut scheduler = Scheduler::new(nodes.values());
    1386            1 : 
    1387            1 :         // Need to keep these alive because they contribute to shard counts via RAII
    1388            1 :         let mut scheduled_shards = Vec::new();
    1389            1 : 
    1390            1 :         let mut context = ScheduleContext::default();
    1391              : 
    1392            8 :         fn schedule_shard(
    1393            8 :             tenant_shard_id: TenantShardId,
    1394            8 :             expect_attached: NodeId,
    1395            8 :             expect_secondary: NodeId,
    1396            8 :             scheduled_shards: &mut Vec<TenantShard>,
    1397            8 :             scheduler: &mut Scheduler,
    1398            8 :             preferred_az: Option<AvailabilityZone>,
    1399            8 :             context: &mut ScheduleContext,
    1400            8 :         ) {
    1401            8 :             let shard_identity = ShardIdentity::new(
    1402            8 :                 tenant_shard_id.shard_number,
    1403            8 :                 tenant_shard_id.shard_count,
    1404            8 :                 pageserver_api::shard::ShardStripeSize(1),
    1405            8 :             )
    1406            8 :             .unwrap();
    1407            8 :             let mut shard = TenantShard::new(
    1408            8 :                 tenant_shard_id,
    1409            8 :                 shard_identity,
    1410            8 :                 pageserver_api::controller_api::PlacementPolicy::Attached(1),
    1411            8 :                 preferred_az,
    1412            8 :             );
    1413            8 : 
    1414            8 :             shard.schedule(scheduler, context).unwrap();
    1415            8 : 
    1416            8 :             assert_eq!(shard.intent.get_attached().unwrap(), expect_attached);
    1417            8 :             assert_eq!(
    1418            8 :                 shard.intent.get_secondary().first().unwrap(),
    1419            8 :                 &expect_secondary
    1420            8 :             );
    1421              : 
    1422            8 :             scheduled_shards.push(shard);
    1423            8 :         }
    1424              : 
    1425            1 :         let tenant_id = TenantId::generate();
    1426            1 : 
    1427            1 :         schedule_shard(
    1428            1 :             TenantShardId {
    1429            1 :                 tenant_id,
    1430            1 :                 shard_number: ShardNumber(0),
    1431            1 :                 shard_count: ShardCount(8),
    1432            1 :             },
    1433            1 :             NodeId(1),
    1434            1 :             NodeId(6),
    1435            1 :             &mut scheduled_shards,
    1436            1 :             &mut scheduler,
    1437            1 :             Some(az_a.clone()),
    1438            1 :             &mut context,
    1439            1 :         );
    1440            1 : 
    1441            1 :         schedule_shard(
    1442            1 :             TenantShardId {
    1443            1 :                 tenant_id,
    1444            1 :                 shard_number: ShardNumber(1),
    1445            1 :                 shard_count: ShardCount(8),
    1446            1 :             },
    1447            1 :             NodeId(2),
    1448            1 :             NodeId(7),
    1449            1 :             &mut scheduled_shards,
    1450            1 :             &mut scheduler,
    1451            1 :             Some(az_a.clone()),
    1452            1 :             &mut context,
    1453            1 :         );
    1454            1 : 
    1455            1 :         schedule_shard(
    1456            1 :             TenantShardId {
    1457            1 :                 tenant_id,
    1458            1 :                 shard_number: ShardNumber(2),
    1459            1 :                 shard_count: ShardCount(8),
    1460            1 :             },
    1461            1 :             NodeId(3),
    1462            1 :             NodeId(8),
    1463            1 :             &mut scheduled_shards,
    1464            1 :             &mut scheduler,
    1465            1 :             Some(az_a.clone()),
    1466            1 :             &mut context,
    1467            1 :         );
    1468            1 : 
    1469            1 :         schedule_shard(
    1470            1 :             TenantShardId {
    1471            1 :                 tenant_id,
    1472            1 :                 shard_number: ShardNumber(3),
    1473            1 :                 shard_count: ShardCount(8),
    1474            1 :             },
    1475            1 :             NodeId(4),
    1476            1 :             NodeId(9),
    1477            1 :             &mut scheduled_shards,
    1478            1 :             &mut scheduler,
    1479            1 :             Some(az_a.clone()),
    1480            1 :             &mut context,
    1481            1 :         );
    1482            1 : 
    1483            1 :         schedule_shard(
    1484            1 :             TenantShardId {
    1485            1 :                 tenant_id,
    1486            1 :                 shard_number: ShardNumber(4),
    1487            1 :                 shard_count: ShardCount(8),
    1488            1 :             },
    1489            1 :             NodeId(5),
    1490            1 :             NodeId(10),
    1491            1 :             &mut scheduled_shards,
    1492            1 :             &mut scheduler,
    1493            1 :             Some(az_a.clone()),
    1494            1 :             &mut context,
    1495            1 :         );
    1496            1 : 
    1497            1 :         schedule_shard(
    1498            1 :             TenantShardId {
    1499            1 :                 tenant_id,
    1500            1 :                 shard_number: ShardNumber(5),
    1501            1 :                 shard_count: ShardCount(8),
    1502            1 :             },
    1503            1 :             NodeId(1),
    1504            1 :             NodeId(6),
    1505            1 :             &mut scheduled_shards,
    1506            1 :             &mut scheduler,
    1507            1 :             Some(az_a.clone()),
    1508            1 :             &mut context,
    1509            1 :         );
    1510            1 : 
    1511            1 :         schedule_shard(
    1512            1 :             TenantShardId {
    1513            1 :                 tenant_id,
    1514            1 :                 shard_number: ShardNumber(6),
    1515            1 :                 shard_count: ShardCount(8),
    1516            1 :             },
    1517            1 :             NodeId(2),
    1518            1 :             NodeId(7),
    1519            1 :             &mut scheduled_shards,
    1520            1 :             &mut scheduler,
    1521            1 :             Some(az_a.clone()),
    1522            1 :             &mut context,
    1523            1 :         );
    1524            1 : 
    1525            1 :         schedule_shard(
    1526            1 :             TenantShardId {
    1527            1 :                 tenant_id,
    1528            1 :                 shard_number: ShardNumber(7),
    1529            1 :                 shard_count: ShardCount(8),
    1530            1 :             },
    1531            1 :             NodeId(3),
    1532            1 :             NodeId(8),
    1533            1 :             &mut scheduled_shards,
    1534            1 :             &mut scheduler,
    1535            1 :             Some(az_a.clone()),
    1536            1 :             &mut context,
    1537            1 :         );
    1538              : 
    1539              :         // Assert that the optimizer suggests nochanges, i.e. our initial scheduling was stable.
    1540            9 :         for shard in &scheduled_shards {
    1541            8 :             assert_eq!(shard.optimize_attachment(&mut scheduler, &context), None);
    1542              :         }
    1543              : 
    1544            9 :         for mut shard in scheduled_shards {
    1545            8 :             shard.intent.clear(&mut scheduler);
    1546            8 :         }
    1547            1 :     }
    1548              : 
    1549              :     #[test]
    1550            1 :     fn change_preferred_az() {
    1551            1 :         let az_a = AvailabilityZone("az-a".to_string());
    1552            1 :         let az_b = AvailabilityZone("az-b".to_string());
    1553            1 : 
    1554            1 :         // 2 nodes: 1 az_a and 1 az_b.
    1555            1 :         let nodes = test_utils::make_test_nodes(2, &[az_a.clone(), az_b.clone()]);
    1556            1 :         let mut scheduler = Scheduler::new(nodes.values());
    1557            1 : 
    1558            1 :         let tenant_shard_id = TenantShardId {
    1559            1 :             tenant_id: TenantId::generate(),
    1560            1 :             shard_number: ShardNumber(0),
    1561            1 :             shard_count: ShardCount(1),
    1562            1 :         };
    1563            1 :         let shard_identity = ShardIdentity::new(
    1564            1 :             tenant_shard_id.shard_number,
    1565            1 :             tenant_shard_id.shard_count,
    1566            1 :             pageserver_api::shard::ShardStripeSize(1),
    1567            1 :         )
    1568            1 :         .unwrap();
    1569            1 :         // 1 attached and 1 secondary.
    1570            1 :         let mut shard = TenantShard::new(
    1571            1 :             tenant_shard_id,
    1572            1 :             shard_identity,
    1573            1 :             pageserver_api::controller_api::PlacementPolicy::Attached(1),
    1574            1 :             Some(az_a.clone()),
    1575            1 :         );
    1576            1 : 
    1577            1 :         let mut context = ScheduleContext::default();
    1578            1 :         shard.schedule(&mut scheduler, &mut context).unwrap();
    1579            1 :         eprintln!("Scheduled shard at {:?}", shard.intent);
    1580              : 
    1581            2 :         for node in scheduler.nodes.values() {
    1582              :             // Only 2 nodes, one tenant shard should be scheduled on each of them.
    1583            2 :             assert_eq!(node.shard_count, 1);
    1584            2 :             if node.az == az_a {
    1585            1 :                 assert_eq!(node.home_shard_count, 1);
    1586              :             } else {
    1587            1 :                 assert_eq!(node.home_shard_count, 0);
    1588              :             }
    1589              :         }
    1590              : 
    1591            1 :         shard.set_preferred_az(&mut scheduler, Some(az_b.clone()));
    1592              :         // Home AZ flipped.
    1593            2 :         for node in scheduler.nodes.values() {
    1594            2 :             assert_eq!(node.shard_count, 1);
    1595            2 :             if node.az == az_a {
    1596            1 :                 assert_eq!(node.home_shard_count, 0);
    1597              :             } else {
    1598            1 :                 assert_eq!(node.home_shard_count, 1);
    1599              :             }
    1600              :         }
    1601              : 
    1602            1 :         shard.set_preferred_az(&mut scheduler, None);
    1603              :         // No home AZ.
    1604            2 :         for node in scheduler.nodes.values() {
    1605            2 :             assert_eq!(node.shard_count, 1);
    1606            2 :             assert_eq!(node.home_shard_count, 0);
    1607              :         }
    1608              : 
    1609            1 :         shard.intent.clear(&mut scheduler);
    1610            1 :     }
    1611              : }
        

Generated by: LCOV version 2.1-beta