LCOV - code coverage report
Current view: top level - storage_controller/src - tenant_shard.rs (source / functions) Coverage Total Hit
Test: c789ec21f6053d4c25d2419c4a34ed298d5f69f5.info Lines: 66.2 % 863 571
Test Date: 2024-06-20 08:12:09 Functions: 50.0 % 74 37

            Line data    Source code
       1              : use std::{
       2              :     collections::{HashMap, HashSet},
       3              :     sync::Arc,
       4              :     time::Duration,
       5              : };
       6              : 
       7              : use crate::{
       8              :     metrics::{self, ReconcileCompleteLabelGroup, ReconcileOutcome},
       9              :     persistence::TenantShardPersistence,
      10              :     reconciler::ReconcileUnits,
      11              :     scheduler::{AffinityScore, MaySchedule, RefCountUpdate, ScheduleContext},
      12              : };
      13              : use pageserver_api::controller_api::{
      14              :     NodeSchedulingPolicy, PlacementPolicy, ShardSchedulingPolicy,
      15              : };
      16              : use pageserver_api::{
      17              :     models::{LocationConfig, LocationConfigMode, TenantConfig},
      18              :     shard::{ShardIdentity, TenantShardId},
      19              : };
      20              : use serde::Serialize;
      21              : use tokio::task::JoinHandle;
      22              : use tokio_util::sync::CancellationToken;
      23              : use tracing::{instrument, Instrument};
      24              : use utils::{
      25              :     generation::Generation,
      26              :     id::NodeId,
      27              :     seqwait::{SeqWait, SeqWaitError},
      28              :     sync::gate::GateGuard,
      29              : };
      30              : 
      31              : use crate::{
      32              :     compute_hook::ComputeHook,
      33              :     node::Node,
      34              :     persistence::{split_state::SplitState, Persistence},
      35              :     reconciler::{
      36              :         attached_location_conf, secondary_location_conf, ReconcileError, Reconciler, TargetState,
      37              :     },
      38              :     scheduler::{ScheduleError, Scheduler},
      39              :     service, Sequence,
      40              : };
      41              : 
      42              : /// Serialization helper
      43            0 : fn read_last_error<S, T>(v: &std::sync::Mutex<Option<T>>, serializer: S) -> Result<S::Ok, S::Error>
      44            0 : where
      45            0 :     S: serde::ser::Serializer,
      46            0 :     T: std::fmt::Display,
      47            0 : {
      48            0 :     serializer.collect_str(
      49            0 :         &v.lock()
      50            0 :             .unwrap()
      51            0 :             .as_ref()
      52            0 :             .map(|e| format!("{e}"))
      53            0 :             .unwrap_or("".to_string()),
      54            0 :     )
      55            0 : }
      56              : 
      57              : /// In-memory state for a particular tenant shard.
      58              : ///
      59              : /// This struct implement Serialize for debugging purposes, but is _not_ persisted
      60              : /// itself: see [`crate::persistence`] for the subset of tenant shard state that is persisted.
      61            0 : #[derive(Serialize)]
      62              : pub(crate) struct TenantShard {
      63              :     pub(crate) tenant_shard_id: TenantShardId,
      64              : 
      65              :     pub(crate) shard: ShardIdentity,
      66              : 
      67              :     // Runtime only: sequence used to coordinate when updating this object while
      68              :     // with background reconcilers may be running.  A reconciler runs to a particular
      69              :     // sequence.
      70              :     pub(crate) sequence: Sequence,
      71              : 
      72              :     // Latest generation number: next time we attach, increment this
      73              :     // and use the incremented number when attaching.
      74              :     //
      75              :     // None represents an incompletely onboarded tenant via the [`Service::location_config`]
      76              :     // API, where this tenant may only run in PlacementPolicy::Secondary.
      77              :     pub(crate) generation: Option<Generation>,
      78              : 
      79              :     // High level description of how the tenant should be set up.  Provided
      80              :     // externally.
      81              :     pub(crate) policy: PlacementPolicy,
      82              : 
      83              :     // Low level description of exactly which pageservers should fulfil
      84              :     // which role.  Generated by `Self::schedule`.
      85              :     pub(crate) intent: IntentState,
      86              : 
      87              :     // Low level description of how the tenant is configured on pageservers:
      88              :     // if this does not match `Self::intent` then the tenant needs reconciliation
      89              :     // with `Self::reconcile`.
      90              :     pub(crate) observed: ObservedState,
      91              : 
      92              :     // Tenant configuration, passed through opaquely to the pageserver.  Identical
      93              :     // for all shards in a tenant.
      94              :     pub(crate) config: TenantConfig,
      95              : 
      96              :     /// If a reconcile task is currently in flight, it may be joined here (it is
      97              :     /// only safe to join if either the result has been received or the reconciler's
      98              :     /// cancellation token has been fired)
      99              :     #[serde(skip)]
     100              :     pub(crate) reconciler: Option<ReconcilerHandle>,
     101              : 
     102              :     /// If a tenant is being split, then all shards with that TenantId will have a
     103              :     /// SplitState set, this acts as a guard against other operations such as background
     104              :     /// reconciliation, and timeline creation.
     105              :     pub(crate) splitting: SplitState,
     106              : 
     107              :     /// If a tenant was enqueued for later reconcile due to hitting concurrency limit, this flag
     108              :     /// is set. This flag is cleared when the tenant is popped off the delay queue.
     109              :     pub(crate) delayed_reconcile: bool,
     110              : 
     111              :     /// Optionally wait for reconciliation to complete up to a particular
     112              :     /// sequence number.
     113              :     #[serde(skip)]
     114              :     pub(crate) waiter: std::sync::Arc<SeqWait<Sequence, Sequence>>,
     115              : 
     116              :     /// Indicates sequence number for which we have encountered an error reconciling.  If
     117              :     /// this advances ahead of [`Self::waiter`] then a reconciliation error has occurred,
     118              :     /// and callers should stop waiting for `waiter` and propagate the error.
     119              :     #[serde(skip)]
     120              :     pub(crate) error_waiter: std::sync::Arc<SeqWait<Sequence, Sequence>>,
     121              : 
     122              :     /// The most recent error from a reconcile on this tenant.  This is a nested Arc
     123              :     /// because:
     124              :     ///  - ReconcileWaiters need to Arc-clone the overall object to read it later
     125              :     ///  - ReconcileWaitError needs to use an `Arc<ReconcileError>` because we can construct
     126              :     ///    many waiters for one shard, and the underlying error types are not Clone.
     127              :     /// TODO: generalize to an array of recent events
     128              :     /// TOOD: use a ArcSwap instead of mutex for faster reads?
     129              :     #[serde(serialize_with = "read_last_error")]
     130              :     pub(crate) last_error: std::sync::Arc<std::sync::Mutex<Option<Arc<ReconcileError>>>>,
     131              : 
     132              :     /// If we have a pending compute notification that for some reason we weren't able to send,
     133              :     /// set this to true. If this is set, calls to [`Self::get_reconcile_needed`] will return Yes
     134              :     /// and trigger a Reconciler run.  This is the mechanism by which compute notifications are included in the scope
     135              :     /// of state that we publish externally in an eventually consistent way.
     136              :     pub(crate) pending_compute_notification: bool,
     137              : 
     138              :     // Support/debug tool: if something is going wrong or flapping with scheduling, this may
     139              :     // be set to a non-active state to avoid making changes while the issue is fixed.
     140              :     scheduling_policy: ShardSchedulingPolicy,
     141              : }
     142              : 
     143              : #[derive(Default, Clone, Debug, Serialize)]
     144              : pub(crate) struct IntentState {
     145              :     attached: Option<NodeId>,
     146              :     secondary: Vec<NodeId>,
     147              : }
     148              : 
     149              : impl IntentState {
     150            4 :     pub(crate) fn new() -> Self {
     151            4 :         Self {
     152            4 :             attached: None,
     153            4 :             secondary: vec![],
     154            4 :         }
     155            4 :     }
     156            0 :     pub(crate) fn single(scheduler: &mut Scheduler, node_id: Option<NodeId>) -> Self {
     157            0 :         if let Some(node_id) = node_id {
     158            0 :             scheduler.update_node_ref_counts(node_id, RefCountUpdate::Attach);
     159            0 :         }
     160            0 :         Self {
     161            0 :             attached: node_id,
     162            0 :             secondary: vec![],
     163            0 :         }
     164            0 :     }
     165              : 
     166           26 :     pub(crate) fn set_attached(&mut self, scheduler: &mut Scheduler, new_attached: Option<NodeId>) {
     167           26 :         if self.attached != new_attached {
     168           26 :             if let Some(old_attached) = self.attached.take() {
     169            0 :                 scheduler.update_node_ref_counts(old_attached, RefCountUpdate::Detach);
     170           26 :             }
     171           26 :             if let Some(new_attached) = &new_attached {
     172           26 :                 scheduler.update_node_ref_counts(*new_attached, RefCountUpdate::Attach);
     173           26 :             }
     174           26 :             self.attached = new_attached;
     175            0 :         }
     176           26 :     }
     177              : 
     178              :     /// Like set_attached, but the node is from [`Self::secondary`].  This swaps the node from
     179              :     /// secondary to attached while maintaining the scheduler's reference counts.
     180           14 :     pub(crate) fn promote_attached(
     181           14 :         &mut self,
     182           14 :         scheduler: &mut Scheduler,
     183           14 :         promote_secondary: NodeId,
     184           14 :     ) {
     185           14 :         // If we call this with a node that isn't in secondary, it would cause incorrect
     186           14 :         // scheduler reference counting, since we assume the node is already referenced as a secondary.
     187           14 :         debug_assert!(self.secondary.contains(&promote_secondary));
     188              : 
     189           28 :         self.secondary.retain(|n| n != &promote_secondary);
     190           14 : 
     191           14 :         let demoted = self.attached;
     192           14 :         self.attached = Some(promote_secondary);
     193           14 : 
     194           14 :         scheduler.update_node_ref_counts(promote_secondary, RefCountUpdate::PromoteSecondary);
     195           14 :         if let Some(demoted) = demoted {
     196            0 :             scheduler.update_node_ref_counts(demoted, RefCountUpdate::DemoteAttached);
     197           14 :         }
     198           14 :     }
     199              : 
     200           34 :     pub(crate) fn push_secondary(&mut self, scheduler: &mut Scheduler, new_secondary: NodeId) {
     201           34 :         debug_assert!(!self.secondary.contains(&new_secondary));
     202           34 :         scheduler.update_node_ref_counts(new_secondary, RefCountUpdate::AddSecondary);
     203           34 :         self.secondary.push(new_secondary);
     204           34 :     }
     205              : 
     206              :     /// It is legal to call this with a node that is not currently a secondary: that is a no-op
     207           10 :     pub(crate) fn remove_secondary(&mut self, scheduler: &mut Scheduler, node_id: NodeId) {
     208           10 :         let index = self.secondary.iter().position(|n| *n == node_id);
     209           10 :         if let Some(index) = index {
     210           10 :             scheduler.update_node_ref_counts(node_id, RefCountUpdate::RemoveSecondary);
     211           10 :             self.secondary.remove(index);
     212           10 :         }
     213           10 :     }
     214              : 
     215           24 :     pub(crate) fn clear_secondary(&mut self, scheduler: &mut Scheduler) {
     216           24 :         for secondary in self.secondary.drain(..) {
     217           24 :             scheduler.update_node_ref_counts(secondary, RefCountUpdate::RemoveSecondary);
     218           24 :         }
     219           24 :     }
     220              : 
     221              :     /// Remove the last secondary node from the list of secondaries
     222            0 :     pub(crate) fn pop_secondary(&mut self, scheduler: &mut Scheduler) {
     223            0 :         if let Some(node_id) = self.secondary.pop() {
     224            0 :             scheduler.update_node_ref_counts(node_id, RefCountUpdate::RemoveSecondary);
     225            0 :         }
     226            0 :     }
     227              : 
     228           24 :     pub(crate) fn clear(&mut self, scheduler: &mut Scheduler) {
     229           24 :         if let Some(old_attached) = self.attached.take() {
     230           24 :             scheduler.update_node_ref_counts(old_attached, RefCountUpdate::Detach);
     231           24 :         }
     232              : 
     233           24 :         self.clear_secondary(scheduler);
     234           24 :     }
     235              : 
     236          172 :     pub(crate) fn all_pageservers(&self) -> Vec<NodeId> {
     237          172 :         let mut result = Vec::new();
     238          172 :         if let Some(p) = self.attached {
     239          168 :             result.push(p)
     240            4 :         }
     241              : 
     242          172 :         result.extend(self.secondary.iter().copied());
     243          172 : 
     244          172 :         result
     245          172 :     }
     246              : 
     247          144 :     pub(crate) fn get_attached(&self) -> &Option<NodeId> {
     248          144 :         &self.attached
     249          144 :     }
     250              : 
     251           38 :     pub(crate) fn get_secondary(&self) -> &Vec<NodeId> {
     252           38 :         &self.secondary
     253           38 :     }
     254              : 
     255              :     /// If the node is in use as the attached location, demote it into
     256              :     /// the list of secondary locations.  This is used when a node goes offline,
     257              :     /// and we want to use a different node for attachment, but not permanently
     258              :     /// forget the location on the offline node.
     259              :     ///
     260              :     /// Returns true if a change was made
     261           14 :     pub(crate) fn demote_attached(&mut self, scheduler: &mut Scheduler, node_id: NodeId) -> bool {
     262           14 :         if self.attached == Some(node_id) {
     263           14 :             self.attached = None;
     264           14 :             self.secondary.push(node_id);
     265           14 :             scheduler.update_node_ref_counts(node_id, RefCountUpdate::DemoteAttached);
     266           14 :             true
     267              :         } else {
     268            0 :             false
     269              :         }
     270           14 :     }
     271              : }
     272              : 
     273              : impl Drop for IntentState {
     274           26 :     fn drop(&mut self) {
     275           26 :         // Must clear before dropping, to avoid leaving stale refcounts in the Scheduler.
     276           26 :         // We do not check this while panicking, to avoid polluting unit test failures or
     277           26 :         // other assertions with this assertion's output.  It's still wrong to leak these,
     278           26 :         // but if we already have a panic then we don't need to independently flag this case.
     279           26 :         if !(std::thread::panicking()) {
     280           26 :             debug_assert!(self.attached.is_none() && self.secondary.is_empty());
     281            0 :         }
     282           24 :     }
     283              : }
     284              : 
     285              : #[derive(Default, Clone, Serialize)]
     286              : pub(crate) struct ObservedState {
     287              :     pub(crate) locations: HashMap<NodeId, ObservedStateLocation>,
     288              : }
     289              : 
     290              : /// Our latest knowledge of how this tenant is configured in the outside world.
     291              : ///
     292              : /// Meaning:
     293              : ///     * No instance of this type exists for a node: we are certain that we have nothing configured on that
     294              : ///       node for this shard.
     295              : ///     * Instance exists with conf==None: we *might* have some state on that node, but we don't know
     296              : ///       what it is (e.g. we failed partway through configuring it)
     297              : ///     * Instance exists with conf==Some: this tells us what we last successfully configured on this node,
     298              : ///       and that configuration will still be present unless something external interfered.
     299              : #[derive(Clone, Serialize)]
     300              : pub(crate) struct ObservedStateLocation {
     301              :     /// If None, it means we do not know the status of this shard's location on this node, but
     302              :     /// we know that we might have some state on this node.
     303              :     pub(crate) conf: Option<LocationConfig>,
     304              : }
     305              : pub(crate) struct ReconcilerWaiter {
     306              :     // For observability purposes, remember the ID of the shard we're
     307              :     // waiting for.
     308              :     pub(crate) tenant_shard_id: TenantShardId,
     309              : 
     310              :     seq_wait: std::sync::Arc<SeqWait<Sequence, Sequence>>,
     311              :     error_seq_wait: std::sync::Arc<SeqWait<Sequence, Sequence>>,
     312              :     error: std::sync::Arc<std::sync::Mutex<Option<Arc<ReconcileError>>>>,
     313              :     seq: Sequence,
     314              : }
     315              : 
     316              : pub(crate) enum ReconcilerStatus {
     317              :     Done,
     318              :     Failed,
     319              :     InProgress,
     320              : }
     321              : 
     322            0 : #[derive(thiserror::Error, Debug)]
     323              : pub(crate) enum ReconcileWaitError {
     324              :     #[error("Timeout waiting for shard {0}")]
     325              :     Timeout(TenantShardId),
     326              :     #[error("shutting down")]
     327              :     Shutdown,
     328              :     #[error("Reconcile error on shard {0}: {1}")]
     329              :     Failed(TenantShardId, Arc<ReconcileError>),
     330              : }
     331              : 
     332              : #[derive(Eq, PartialEq, Debug)]
     333              : pub(crate) struct ReplaceSecondary {
     334              :     old_node_id: NodeId,
     335              :     new_node_id: NodeId,
     336              : }
     337              : 
     338              : #[derive(Eq, PartialEq, Debug)]
     339              : pub(crate) struct MigrateAttachment {
     340              :     pub(crate) old_attached_node_id: NodeId,
     341              :     pub(crate) new_attached_node_id: NodeId,
     342              : }
     343              : 
     344              : #[derive(Eq, PartialEq, Debug)]
     345              : pub(crate) enum ScheduleOptimizationAction {
     346              :     // Replace one of our secondary locations with a different node
     347              :     ReplaceSecondary(ReplaceSecondary),
     348              :     // Migrate attachment to an existing secondary location
     349              :     MigrateAttachment(MigrateAttachment),
     350              : }
     351              : 
     352              : #[derive(Eq, PartialEq, Debug)]
     353              : pub(crate) struct ScheduleOptimization {
     354              :     // What was the reconcile sequence when we generated this optimization?  The optimization
     355              :     // should only be applied if the shard's sequence is still at this value, in case other changes
     356              :     // happened between planning the optimization and applying it.
     357              :     sequence: Sequence,
     358              : 
     359              :     pub(crate) action: ScheduleOptimizationAction,
     360              : }
     361              : 
     362              : impl ReconcilerWaiter {
     363            0 :     pub(crate) async fn wait_timeout(&self, timeout: Duration) -> Result<(), ReconcileWaitError> {
     364              :         tokio::select! {
     365              :             result = self.seq_wait.wait_for_timeout(self.seq, timeout)=> {
     366            0 :                 result.map_err(|e| match e {
     367            0 :                     SeqWaitError::Timeout => ReconcileWaitError::Timeout(self.tenant_shard_id),
     368            0 :                     SeqWaitError::Shutdown => ReconcileWaitError::Shutdown
     369            0 :                 })?;
     370              :             },
     371              :             result = self.error_seq_wait.wait_for(self.seq) => {
     372            0 :                 result.map_err(|e| match e {
     373            0 :                     SeqWaitError::Shutdown => ReconcileWaitError::Shutdown,
     374            0 :                     SeqWaitError::Timeout => unreachable!()
     375            0 :                 })?;
     376              : 
     377              :                 return Err(ReconcileWaitError::Failed(self.tenant_shard_id,
     378              :                     self.error.lock().unwrap().clone().expect("If error_seq_wait was advanced error was set").clone()))
     379              :             }
     380              :         }
     381              : 
     382            0 :         Ok(())
     383            0 :     }
     384              : 
     385            0 :     pub(crate) fn get_status(&self) -> ReconcilerStatus {
     386            0 :         if self.seq_wait.would_wait_for(self.seq).is_err() {
     387            0 :             ReconcilerStatus::Done
     388            0 :         } else if self.error_seq_wait.would_wait_for(self.seq).is_err() {
     389            0 :             ReconcilerStatus::Failed
     390              :         } else {
     391            0 :             ReconcilerStatus::InProgress
     392              :         }
     393            0 :     }
     394              : }
     395              : 
     396              : /// Having spawned a reconciler task, the tenant shard's state will carry enough
     397              : /// information to optionally cancel & await it later.
     398              : pub(crate) struct ReconcilerHandle {
     399              :     sequence: Sequence,
     400              :     handle: JoinHandle<()>,
     401              :     cancel: CancellationToken,
     402              : }
     403              : 
     404              : pub(crate) enum ReconcileNeeded {
     405              :     /// shard either doesn't need reconciliation, or is forbidden from spawning a reconciler
     406              :     /// in its current state (e.g. shard split in progress, or ShardSchedulingPolicy forbids it)
     407              :     No,
     408              :     /// shard has a reconciler running, and its intent hasn't changed since that one was
     409              :     /// spawned: wait for the existing reconciler rather than spawning a new one.
     410              :     WaitExisting(ReconcilerWaiter),
     411              :     /// shard needs reconciliation: call into [`TenantShard::spawn_reconciler`]
     412              :     Yes,
     413              : }
     414              : 
     415              : /// When a reconcile task completes, it sends this result object
     416              : /// to be applied to the primary TenantShard.
     417              : pub(crate) struct ReconcileResult {
     418              :     pub(crate) sequence: Sequence,
     419              :     /// On errors, `observed` should be treated as an incompleted description
     420              :     /// of state (i.e. any nodes present in the result should override nodes
     421              :     /// present in the parent tenant state, but any unmentioned nodes should
     422              :     /// not be removed from parent tenant state)
     423              :     pub(crate) result: Result<(), ReconcileError>,
     424              : 
     425              :     pub(crate) tenant_shard_id: TenantShardId,
     426              :     pub(crate) generation: Option<Generation>,
     427              :     pub(crate) observed: ObservedState,
     428              : 
     429              :     /// Set [`TenantShard::pending_compute_notification`] from this flag
     430              :     pub(crate) pending_compute_notification: bool,
     431              : }
     432              : 
     433              : impl ObservedState {
     434            0 :     pub(crate) fn new() -> Self {
     435            0 :         Self {
     436            0 :             locations: HashMap::new(),
     437            0 :         }
     438            0 :     }
     439              : }
     440              : 
     441              : impl TenantShard {
     442           22 :     pub(crate) fn new(
     443           22 :         tenant_shard_id: TenantShardId,
     444           22 :         shard: ShardIdentity,
     445           22 :         policy: PlacementPolicy,
     446           22 :     ) -> Self {
     447           22 :         Self {
     448           22 :             tenant_shard_id,
     449           22 :             policy,
     450           22 :             intent: IntentState::default(),
     451           22 :             generation: Some(Generation::new(0)),
     452           22 :             shard,
     453           22 :             observed: ObservedState::default(),
     454           22 :             config: TenantConfig::default(),
     455           22 :             reconciler: None,
     456           22 :             splitting: SplitState::Idle,
     457           22 :             sequence: Sequence(1),
     458           22 :             delayed_reconcile: false,
     459           22 :             waiter: Arc::new(SeqWait::new(Sequence(0))),
     460           22 :             error_waiter: Arc::new(SeqWait::new(Sequence(0))),
     461           22 :             last_error: Arc::default(),
     462           22 :             pending_compute_notification: false,
     463           22 :             scheduling_policy: ShardSchedulingPolicy::default(),
     464           22 :         }
     465           22 :     }
     466              : 
     467              :     /// For use on startup when learning state from pageservers: generate my [`IntentState`] from my
     468              :     /// [`ObservedState`], even if it violates my [`PlacementPolicy`].  Call [`Self::schedule`] next,
     469              :     /// to get an intent state that complies with placement policy.  The overall goal is to do scheduling
     470              :     /// in a way that makes use of any configured locations that already exist in the outside world.
     471            2 :     pub(crate) fn intent_from_observed(&mut self, scheduler: &mut Scheduler) {
     472            2 :         // Choose an attached location by filtering observed locations, and then sorting to get the highest
     473            2 :         // generation
     474            2 :         let mut attached_locs = self
     475            2 :             .observed
     476            2 :             .locations
     477            2 :             .iter()
     478            4 :             .filter_map(|(node_id, l)| {
     479            4 :                 if let Some(conf) = &l.conf {
     480            4 :                     if conf.mode == LocationConfigMode::AttachedMulti
     481            2 :                         || conf.mode == LocationConfigMode::AttachedSingle
     482            2 :                         || conf.mode == LocationConfigMode::AttachedStale
     483              :                     {
     484            4 :                         Some((node_id, conf.generation))
     485              :                     } else {
     486            0 :                         None
     487              :                     }
     488              :                 } else {
     489            0 :                     None
     490              :                 }
     491            4 :             })
     492            2 :             .collect::<Vec<_>>();
     493            2 : 
     494            4 :         attached_locs.sort_by_key(|i| i.1);
     495            2 :         if let Some((node_id, _gen)) = attached_locs.into_iter().last() {
     496            2 :             self.intent.set_attached(scheduler, Some(*node_id));
     497            2 :         }
     498              : 
     499              :         // All remaining observed locations generate secondary intents.  This includes None
     500              :         // observations, as these may well have some local content on disk that is usable (this
     501              :         // is an edge case that might occur if we restarted during a migration or other change)
     502              :         //
     503              :         // We may leave intent.attached empty if we didn't find any attached locations: [`Self::schedule`]
     504              :         // will take care of promoting one of these secondaries to be attached.
     505            4 :         self.observed.locations.keys().for_each(|node_id| {
     506            4 :             if Some(*node_id) != self.intent.attached {
     507            2 :                 self.intent.push_secondary(scheduler, *node_id);
     508            2 :             }
     509            4 :         });
     510            2 :     }
     511              : 
     512              :     /// Part of [`Self::schedule`] that is used to choose exactly one node to act as the
     513              :     /// attached pageserver for a shard.
     514              :     ///
     515              :     /// Returns whether we modified it, and the NodeId selected.
     516           14 :     fn schedule_attached(
     517           14 :         &mut self,
     518           14 :         scheduler: &mut Scheduler,
     519           14 :         context: &ScheduleContext,
     520           14 :     ) -> Result<(bool, NodeId), ScheduleError> {
     521              :         // No work to do if we already have an attached tenant
     522           14 :         if let Some(node_id) = self.intent.attached {
     523            0 :             return Ok((false, node_id));
     524           14 :         }
     525              : 
     526           14 :         if let Some(promote_secondary) = scheduler.node_preferred(&self.intent.secondary) {
     527              :             // Promote a secondary
     528            2 :             tracing::debug!("Promoted secondary {} to attached", promote_secondary);
     529            2 :             self.intent.promote_attached(scheduler, promote_secondary);
     530            2 :             Ok((true, promote_secondary))
     531              :         } else {
     532              :             // Pick a fresh node: either we had no secondaries or none were schedulable
     533           12 :             let node_id = scheduler.schedule_shard(&self.intent.secondary, context)?;
     534           12 :             tracing::debug!("Selected {} as attached", node_id);
     535           12 :             self.intent.set_attached(scheduler, Some(node_id));
     536           12 :             Ok((true, node_id))
     537              :         }
     538           14 :     }
     539              : 
     540           16 :     pub(crate) fn schedule(
     541           16 :         &mut self,
     542           16 :         scheduler: &mut Scheduler,
     543           16 :         context: &mut ScheduleContext,
     544           16 :     ) -> Result<(), ScheduleError> {
     545           16 :         let r = self.do_schedule(scheduler, context);
     546           16 : 
     547           16 :         context.avoid(&self.intent.all_pageservers());
     548           16 :         if let Some(attached) = self.intent.get_attached() {
     549           14 :             context.push_attached(*attached);
     550           14 :         }
     551              : 
     552           16 :         r
     553           16 :     }
     554              : 
     555           16 :     pub(crate) fn do_schedule(
     556           16 :         &mut self,
     557           16 :         scheduler: &mut Scheduler,
     558           16 :         context: &ScheduleContext,
     559           16 :     ) -> Result<(), ScheduleError> {
     560           16 :         // TODO: before scheduling new nodes, check if any existing content in
     561           16 :         // self.intent refers to pageservers that are offline, and pick other
     562           16 :         // pageservers if so.
     563           16 : 
     564           16 :         // TODO: respect the splitting bit on tenants: if they are currently splitting then we may not
     565           16 :         // change their attach location.
     566           16 : 
     567           16 :         match self.scheduling_policy {
     568           14 :             ShardSchedulingPolicy::Active | ShardSchedulingPolicy::Essential => {}
     569              :             ShardSchedulingPolicy::Pause | ShardSchedulingPolicy::Stop => {
     570              :                 // Warn to make it obvious why other things aren't happening/working, if we skip scheduling
     571            2 :                 tracing::warn!(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug(),
     572            0 :                     "Scheduling is disabled by policy {:?}", self.scheduling_policy);
     573            2 :                 return Ok(());
     574              :             }
     575              :         }
     576              : 
     577              :         // Build the set of pageservers already in use by this tenant, to avoid scheduling
     578              :         // more work on the same pageservers we're already using.
     579           14 :         let mut modified = false;
     580           14 : 
     581           14 :         // Add/remove nodes to fulfil policy
     582           14 :         use PlacementPolicy::*;
     583           14 :         match self.policy {
     584           14 :             Attached(secondary_count) => {
     585           14 :                 let retain_secondaries = if self.intent.attached.is_none()
     586           14 :                     && scheduler.node_preferred(&self.intent.secondary).is_some()
     587              :                 {
     588              :                     // If we have no attached, and one of the secondaries is elegible to be promoted, retain
     589              :                     // one more secondary than we usually would, as one of them will become attached futher down this function.
     590            2 :                     secondary_count + 1
     591              :                 } else {
     592           12 :                     secondary_count
     593              :                 };
     594              : 
     595           14 :                 while self.intent.secondary.len() > retain_secondaries {
     596            0 :                     // We have no particular preference for one secondary location over another: just
     597            0 :                     // arbitrarily drop from the end
     598            0 :                     self.intent.pop_secondary(scheduler);
     599            0 :                     modified = true;
     600            0 :                 }
     601              : 
     602              :                 // Should have exactly one attached, and N secondaries
     603           14 :                 let (modified_attached, attached_node_id) =
     604           14 :                     self.schedule_attached(scheduler, context)?;
     605           14 :                 modified |= modified_attached;
     606           14 : 
     607           14 :                 let mut used_pageservers = vec![attached_node_id];
     608           26 :                 while self.intent.secondary.len() < secondary_count {
     609           12 :                     let node_id = scheduler.schedule_shard(&used_pageservers, context)?;
     610           12 :                     self.intent.push_secondary(scheduler, node_id);
     611           12 :                     used_pageservers.push(node_id);
     612           12 :                     modified = true;
     613              :                 }
     614              :             }
     615              :             Secondary => {
     616            0 :                 if let Some(node_id) = self.intent.get_attached() {
     617            0 :                     // Populate secondary by demoting the attached node
     618            0 :                     self.intent.demote_attached(scheduler, *node_id);
     619            0 :                     modified = true;
     620            0 :                 } else if self.intent.secondary.is_empty() {
     621            0 :                     // Populate secondary by scheduling a fresh node
     622            0 :                     let node_id = scheduler.schedule_shard(&[], context)?;
     623            0 :                     self.intent.push_secondary(scheduler, node_id);
     624            0 :                     modified = true;
     625            0 :                 }
     626            0 :                 while self.intent.secondary.len() > 1 {
     627            0 :                     // We have no particular preference for one secondary location over another: just
     628            0 :                     // arbitrarily drop from the end
     629            0 :                     self.intent.pop_secondary(scheduler);
     630            0 :                     modified = true;
     631            0 :                 }
     632              :             }
     633              :             Detached => {
     634              :                 // Never add locations in this mode
     635            0 :                 if self.intent.get_attached().is_some() || !self.intent.get_secondary().is_empty() {
     636            0 :                     self.intent.clear(scheduler);
     637            0 :                     modified = true;
     638            0 :                 }
     639              :             }
     640              :         }
     641              : 
     642           14 :         if modified {
     643           14 :             self.sequence.0 += 1;
     644           14 :         }
     645              : 
     646           14 :         Ok(())
     647           16 :     }
     648              : 
     649              :     /// Optimize attachments: if a shard has a secondary location that is preferable to
     650              :     /// its primary location based on soft constraints, switch that secondary location
     651              :     /// to be attached.
     652           40 :     #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug()))]
     653              :     pub(crate) fn optimize_attachment(
     654              :         &self,
     655              :         nodes: &HashMap<NodeId, Node>,
     656              :         schedule_context: &ScheduleContext,
     657              :     ) -> Option<ScheduleOptimization> {
     658              :         let attached = (*self.intent.get_attached())?;
     659              :         if self.intent.secondary.is_empty() {
     660              :             // We can only do useful work if we have both attached and secondary locations: this
     661              :             // function doesn't schedule new locations, only swaps between attached and secondaries.
     662              :             return None;
     663              :         }
     664              : 
     665              :         let current_affinity_score = schedule_context.get_node_affinity(attached);
     666              :         let current_attachment_count = schedule_context.get_node_attachments(attached);
     667              : 
     668              :         // Generate score for each node, dropping any un-schedulable nodes.
     669              :         let all_pageservers = self.intent.all_pageservers();
     670              :         let mut scores = all_pageservers
     671              :             .iter()
     672           80 :             .flat_map(|node_id| {
     673           80 :                 let node = nodes.get(node_id);
     674           80 :                 if node.is_none() {
     675            0 :                     None
     676           80 :                 } else if matches!(
     677           80 :                     node.unwrap().get_scheduling(),
     678              :                     NodeSchedulingPolicy::Filling
     679              :                 ) {
     680              :                     // If the node is currently filling, don't count it as a candidate to avoid,
     681              :                     // racing with the background fill.
     682            0 :                     None
     683           80 :                 } else if matches!(node.unwrap().may_schedule(), MaySchedule::No) {
     684            0 :                     None
     685              :                 } else {
     686           80 :                     let affinity_score = schedule_context.get_node_affinity(*node_id);
     687           80 :                     let attachment_count = schedule_context.get_node_attachments(*node_id);
     688           80 :                     Some((*node_id, affinity_score, attachment_count))
     689              :                 }
     690           80 :             })
     691              :             .collect::<Vec<_>>();
     692              : 
     693              :         // Sort precedence:
     694              :         //  1st - prefer nodes with the lowest total affinity score
     695              :         //  2nd - prefer nodes with the lowest number of attachments in this context
     696              :         //  3rd - if all else is equal, sort by node ID for determinism in tests.
     697           80 :         scores.sort_by_key(|i| (i.1, i.2, i.0));
     698              : 
     699              :         if let Some((preferred_node, preferred_affinity_score, preferred_attachment_count)) =
     700              :             scores.first()
     701              :         {
     702              :             if attached != *preferred_node {
     703              :                 // The best alternative must be more than 1 better than us, otherwise we could end
     704              :                 // up flapping back next time we're called (e.g. there's no point migrating from
     705              :                 // a location with score 1 to a score zero, because on next location the situation
     706              :                 // would be the same, but in reverse).
     707              :                 if current_affinity_score > *preferred_affinity_score + AffinityScore(1)
     708              :                     || current_attachment_count > *preferred_attachment_count + 1
     709              :                 {
     710              :                     tracing::info!(
     711              :                         "Identified optimization: migrate attachment {attached}->{preferred_node} (secondaries {:?})",
     712              :                         self.intent.get_secondary()
     713              :                     );
     714              :                     return Some(ScheduleOptimization {
     715              :                         sequence: self.sequence,
     716              :                         action: ScheduleOptimizationAction::MigrateAttachment(MigrateAttachment {
     717              :                             old_attached_node_id: attached,
     718              :                             new_attached_node_id: *preferred_node,
     719              :                         }),
     720              :                     });
     721              :                 }
     722              :             } else {
     723              :                 tracing::debug!(
     724              :                     "Node {} is already preferred (score {:?})",
     725              :                     preferred_node,
     726              :                     preferred_affinity_score
     727              :                 );
     728              :             }
     729              :         }
     730              : 
     731              :         // Fall-through: we didn't find an optimization
     732              :         None
     733              :     }
     734              : 
     735           30 :     #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug()))]
     736              :     pub(crate) fn optimize_secondary(
     737              :         &self,
     738              :         scheduler: &Scheduler,
     739              :         schedule_context: &ScheduleContext,
     740              :     ) -> Option<ScheduleOptimization> {
     741              :         if self.intent.secondary.is_empty() {
     742              :             // We can only do useful work if we have both attached and secondary locations: this
     743              :             // function doesn't schedule new locations, only swaps between attached and secondaries.
     744              :             return None;
     745              :         }
     746              : 
     747              :         for secondary in self.intent.get_secondary() {
     748              :             let Some(affinity_score) = schedule_context.nodes.get(secondary) else {
     749              :                 // We're already on a node unaffected any affinity constraints,
     750              :                 // so we won't change it.
     751              :                 continue;
     752              :             };
     753              : 
     754              :             // Let the scheduler suggest a node, where it would put us if we were scheduling afresh
     755              :             // This implicitly limits the choice to nodes that are available, and prefers nodes
     756              :             // with lower utilization.
     757              :             let Ok(candidate_node) =
     758              :                 scheduler.schedule_shard(&self.intent.all_pageservers(), schedule_context)
     759              :             else {
     760              :                 // A scheduling error means we have no possible candidate replacements
     761              :                 continue;
     762              :             };
     763              : 
     764              :             let candidate_affinity_score = schedule_context
     765              :                 .nodes
     766              :                 .get(&candidate_node)
     767              :                 .unwrap_or(&AffinityScore::FREE);
     768              : 
     769              :             // The best alternative must be more than 1 better than us, otherwise we could end
     770              :             // up flapping back next time we're called.
     771              :             if *candidate_affinity_score + AffinityScore(1) < *affinity_score {
     772              :                 // If some other node is available and has a lower score than this node, then
     773              :                 // that other node is a good place to migrate to.
     774              :                 tracing::info!(
     775              :                     "Identified optimization: replace secondary {secondary}->{candidate_node} (current secondaries {:?})",
     776              :                     self.intent.get_secondary()
     777              :                 );
     778              :                 return Some(ScheduleOptimization {
     779              :                     sequence: self.sequence,
     780              :                     action: ScheduleOptimizationAction::ReplaceSecondary(ReplaceSecondary {
     781              :                         old_node_id: *secondary,
     782              :                         new_node_id: candidate_node,
     783              :                     }),
     784              :                 });
     785              :             }
     786              :         }
     787              : 
     788              :         None
     789              :     }
     790              : 
     791              :     /// Return true if the optimization was really applied: it will not be applied if the optimization's
     792              :     /// sequence is behind this tenant shard's
     793           22 :     pub(crate) fn apply_optimization(
     794           22 :         &mut self,
     795           22 :         scheduler: &mut Scheduler,
     796           22 :         optimization: ScheduleOptimization,
     797           22 :     ) -> bool {
     798           22 :         if optimization.sequence != self.sequence {
     799            0 :             return false;
     800           22 :         }
     801           22 : 
     802           22 :         metrics::METRICS_REGISTRY
     803           22 :             .metrics_group
     804           22 :             .storage_controller_schedule_optimization
     805           22 :             .inc();
     806           22 : 
     807           22 :         match optimization.action {
     808              :             ScheduleOptimizationAction::MigrateAttachment(MigrateAttachment {
     809           12 :                 old_attached_node_id,
     810           12 :                 new_attached_node_id,
     811           12 :             }) => {
     812           12 :                 self.intent.demote_attached(scheduler, old_attached_node_id);
     813           12 :                 self.intent
     814           12 :                     .promote_attached(scheduler, new_attached_node_id);
     815           12 :             }
     816              :             ScheduleOptimizationAction::ReplaceSecondary(ReplaceSecondary {
     817           10 :                 old_node_id,
     818           10 :                 new_node_id,
     819           10 :             }) => {
     820           10 :                 self.intent.remove_secondary(scheduler, old_node_id);
     821           10 :                 self.intent.push_secondary(scheduler, new_node_id);
     822           10 :             }
     823              :         }
     824              : 
     825           22 :         true
     826           22 :     }
     827              : 
     828              :     /// Query whether the tenant's observed state for attached node matches its intent state, and if so,
     829              :     /// yield the node ID.  This is appropriate for emitting compute hook notifications: we are checking that
     830              :     /// the node in question is not only where we intend to attach, but that the tenant is indeed already attached there.
     831              :     ///
     832              :     /// Reconciliation may still be needed for other aspects of state such as secondaries (see [`Self::dirty`]): this
     833              :     /// funciton should not be used to decide whether to reconcile.
     834            0 :     pub(crate) fn stably_attached(&self) -> Option<NodeId> {
     835            0 :         if let Some(attach_intent) = self.intent.attached {
     836            0 :             match self.observed.locations.get(&attach_intent) {
     837            0 :                 Some(loc) => match &loc.conf {
     838            0 :                     Some(conf) => match conf.mode {
     839              :                         LocationConfigMode::AttachedMulti
     840              :                         | LocationConfigMode::AttachedSingle
     841              :                         | LocationConfigMode::AttachedStale => {
     842              :                             // Our intent and observed state agree that this node is in an attached state.
     843            0 :                             Some(attach_intent)
     844              :                         }
     845              :                         // Our observed config is not an attached state
     846            0 :                         _ => None,
     847              :                     },
     848              :                     // Our observed state is None, i.e. in flux
     849            0 :                     None => None,
     850              :                 },
     851              :                 // We have no observed state for this node
     852            0 :                 None => None,
     853              :             }
     854              :         } else {
     855              :             // Our intent is not to attach
     856            0 :             None
     857              :         }
     858            0 :     }
     859              : 
     860            0 :     fn dirty(&self, nodes: &Arc<HashMap<NodeId, Node>>) -> bool {
     861            0 :         let mut dirty_nodes = HashSet::new();
     862              : 
     863            0 :         if let Some(node_id) = self.intent.attached {
     864              :             // Maybe panic: it is a severe bug if we try to attach while generation is null.
     865            0 :             let generation = self
     866            0 :                 .generation
     867            0 :                 .expect("Attempted to enter attached state without a generation");
     868            0 : 
     869            0 :             let wanted_conf = attached_location_conf(
     870            0 :                 generation,
     871            0 :                 &self.shard,
     872            0 :                 &self.config,
     873            0 :                 !self.intent.secondary.is_empty(),
     874            0 :             );
     875            0 :             match self.observed.locations.get(&node_id) {
     876            0 :                 Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => {}
     877            0 :                 Some(_) | None => {
     878            0 :                     dirty_nodes.insert(node_id);
     879            0 :                 }
     880              :             }
     881            0 :         }
     882              : 
     883            0 :         for node_id in &self.intent.secondary {
     884            0 :             let wanted_conf = secondary_location_conf(&self.shard, &self.config);
     885            0 :             match self.observed.locations.get(node_id) {
     886            0 :                 Some(conf) if conf.conf.as_ref() == Some(&wanted_conf) => {}
     887            0 :                 Some(_) | None => {
     888            0 :                     dirty_nodes.insert(*node_id);
     889            0 :                 }
     890              :             }
     891              :         }
     892              : 
     893            0 :         for node_id in self.observed.locations.keys() {
     894            0 :             if self.intent.attached != Some(*node_id) && !self.intent.secondary.contains(node_id) {
     895            0 :                 // We have observed state that isn't part of our intent: need to clean it up.
     896            0 :                 dirty_nodes.insert(*node_id);
     897            0 :             }
     898              :         }
     899              : 
     900            0 :         dirty_nodes.retain(|node_id| {
     901            0 :             nodes
     902            0 :                 .get(node_id)
     903            0 :                 .map(|n| n.is_available())
     904            0 :                 .unwrap_or(false)
     905            0 :         });
     906            0 : 
     907            0 :         !dirty_nodes.is_empty()
     908            0 :     }
     909              : 
     910              :     #[allow(clippy::too_many_arguments)]
     911            0 :     #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug()))]
     912              :     pub(crate) fn get_reconcile_needed(
     913              :         &mut self,
     914              :         pageservers: &Arc<HashMap<NodeId, Node>>,
     915              :     ) -> ReconcileNeeded {
     916              :         // If there are any ambiguous observed states, and the nodes they refer to are available,
     917              :         // we should reconcile to clean them up.
     918              :         let mut dirty_observed = false;
     919              :         for (node_id, observed_loc) in &self.observed.locations {
     920              :             let node = pageservers
     921              :                 .get(node_id)
     922              :                 .expect("Nodes may not be removed while referenced");
     923              :             if observed_loc.conf.is_none() && node.is_available() {
     924              :                 dirty_observed = true;
     925              :                 break;
     926              :             }
     927              :         }
     928              : 
     929              :         let active_nodes_dirty = self.dirty(pageservers);
     930              : 
     931              :         // Even if there is no pageserver work to be done, if we have a pending notification to computes,
     932              :         // wake up a reconciler to send it.
     933              :         let do_reconcile =
     934              :             active_nodes_dirty || dirty_observed || self.pending_compute_notification;
     935              : 
     936              :         if !do_reconcile {
     937              :             tracing::debug!("Not dirty, no reconciliation needed.");
     938              :             return ReconcileNeeded::No;
     939              :         }
     940              : 
     941              :         // If we are currently splitting, then never start a reconciler task: the splitting logic
     942              :         // requires that shards are not interfered with while it runs. Do this check here rather than
     943              :         // up top, so that we only log this message if we would otherwise have done a reconciliation.
     944              :         if !matches!(self.splitting, SplitState::Idle) {
     945              :             tracing::info!("Refusing to reconcile, splitting in progress");
     946              :             return ReconcileNeeded::No;
     947              :         }
     948              : 
     949              :         // Reconcile already in flight for the current sequence?
     950              :         if let Some(handle) = &self.reconciler {
     951              :             if handle.sequence == self.sequence {
     952              :                 tracing::info!(
     953              :                     "Reconciliation already in progress for sequence {:?}",
     954              :                     self.sequence,
     955              :                 );
     956              :                 return ReconcileNeeded::WaitExisting(ReconcilerWaiter {
     957              :                     tenant_shard_id: self.tenant_shard_id,
     958              :                     seq_wait: self.waiter.clone(),
     959              :                     error_seq_wait: self.error_waiter.clone(),
     960              :                     error: self.last_error.clone(),
     961              :                     seq: self.sequence,
     962              :                 });
     963              :             }
     964              :         }
     965              : 
     966              :         // Pre-checks done: finally check whether we may actually do the work
     967              :         match self.scheduling_policy {
     968              :             ShardSchedulingPolicy::Active
     969              :             | ShardSchedulingPolicy::Essential
     970              :             | ShardSchedulingPolicy::Pause => {}
     971              :             ShardSchedulingPolicy::Stop => {
     972              :                 // We only reach this point if there is work to do and we're going to skip
     973              :                 // doing it: warn it obvious why this tenant isn't doing what it ought to.
     974              :                 tracing::warn!("Skipping reconcile for policy {:?}", self.scheduling_policy);
     975              :                 return ReconcileNeeded::No;
     976              :             }
     977              :         }
     978              : 
     979              :         ReconcileNeeded::Yes
     980              :     }
     981              : 
     982              :     /// Ensure the sequence number is set to a value where waiting for this value will make us wait
     983              :     /// for the next reconcile: i.e. it is ahead of all completed or running reconcilers.
     984              :     ///
     985              :     /// Constructing a ReconcilerWaiter with the resulting sequence number gives the property
     986              :     /// that the waiter will not complete until some future Reconciler is constructed and run.
     987            0 :     fn ensure_sequence_ahead(&mut self) {
     988            0 :         // Find the highest sequence for which a Reconciler has previously run or is currently
     989            0 :         // running
     990            0 :         let max_seen = std::cmp::max(
     991            0 :             self.reconciler
     992            0 :                 .as_ref()
     993            0 :                 .map(|r| r.sequence)
     994            0 :                 .unwrap_or(Sequence(0)),
     995            0 :             std::cmp::max(self.waiter.load(), self.error_waiter.load()),
     996            0 :         );
     997            0 : 
     998            0 :         if self.sequence <= max_seen {
     999            0 :             self.sequence = max_seen.next();
    1000            0 :         }
    1001            0 :     }
    1002              : 
    1003              :     /// Create a waiter that will wait for some future Reconciler that hasn't been spawned yet.
    1004              :     ///
    1005              :     /// This is appropriate when you can't spawn a reconciler (e.g. due to resource limits), but
    1006              :     /// you would like to wait on the next reconciler that gets spawned in the background.
    1007            0 :     pub(crate) fn future_reconcile_waiter(&mut self) -> ReconcilerWaiter {
    1008            0 :         self.ensure_sequence_ahead();
    1009            0 : 
    1010            0 :         ReconcilerWaiter {
    1011            0 :             tenant_shard_id: self.tenant_shard_id,
    1012            0 :             seq_wait: self.waiter.clone(),
    1013            0 :             error_seq_wait: self.error_waiter.clone(),
    1014            0 :             error: self.last_error.clone(),
    1015            0 :             seq: self.sequence,
    1016            0 :         }
    1017            0 :     }
    1018              : 
    1019              :     #[allow(clippy::too_many_arguments)]
    1020            0 :     #[instrument(skip_all, fields(tenant_id=%self.tenant_shard_id.tenant_id, shard_id=%self.tenant_shard_id.shard_slug()))]
    1021              :     pub(crate) fn spawn_reconciler(
    1022              :         &mut self,
    1023              :         result_tx: &tokio::sync::mpsc::UnboundedSender<ReconcileResult>,
    1024              :         pageservers: &Arc<HashMap<NodeId, Node>>,
    1025              :         compute_hook: &Arc<ComputeHook>,
    1026              :         service_config: &service::Config,
    1027              :         persistence: &Arc<Persistence>,
    1028              :         units: ReconcileUnits,
    1029              :         gate_guard: GateGuard,
    1030              :         cancel: &CancellationToken,
    1031              :     ) -> Option<ReconcilerWaiter> {
    1032              :         // Reconcile in flight for a stale sequence?  Our sequence's task will wait for it before
    1033              :         // doing our sequence's work.
    1034              :         let old_handle = self.reconciler.take();
    1035              : 
    1036              :         // Build list of nodes from which the reconciler should detach
    1037              :         let mut detach = Vec::new();
    1038              :         for node_id in self.observed.locations.keys() {
    1039              :             if self.intent.get_attached() != &Some(*node_id)
    1040              :                 && !self.intent.secondary.contains(node_id)
    1041              :             {
    1042              :                 detach.push(
    1043              :                     pageservers
    1044              :                         .get(node_id)
    1045              :                         .expect("Intent references non-existent pageserver")
    1046              :                         .clone(),
    1047              :                 )
    1048              :             }
    1049              :         }
    1050              : 
    1051              :         // Advance the sequence before spawning a reconciler, so that sequence waiters
    1052              :         // can distinguish between before+after the reconcile completes.
    1053              :         self.ensure_sequence_ahead();
    1054              : 
    1055              :         let reconciler_cancel = cancel.child_token();
    1056              :         let reconciler_intent = TargetState::from_intent(pageservers, &self.intent);
    1057              :         let mut reconciler = Reconciler {
    1058              :             tenant_shard_id: self.tenant_shard_id,
    1059              :             shard: self.shard,
    1060              :             generation: self.generation,
    1061              :             intent: reconciler_intent,
    1062              :             detach,
    1063              :             config: self.config.clone(),
    1064              :             observed: self.observed.clone(),
    1065              :             compute_hook: compute_hook.clone(),
    1066              :             service_config: service_config.clone(),
    1067              :             _gate_guard: gate_guard,
    1068              :             _resource_units: units,
    1069              :             cancel: reconciler_cancel.clone(),
    1070              :             persistence: persistence.clone(),
    1071              :             compute_notify_failure: false,
    1072              :         };
    1073              : 
    1074              :         let reconcile_seq = self.sequence;
    1075              : 
    1076              :         tracing::info!(seq=%reconcile_seq, "Spawning Reconciler for sequence {}", self.sequence);
    1077              :         let must_notify = self.pending_compute_notification;
    1078              :         let reconciler_span = tracing::info_span!(parent: None, "reconciler", seq=%reconcile_seq,
    1079              :                                                         tenant_id=%reconciler.tenant_shard_id.tenant_id,
    1080              :                                                         shard_id=%reconciler.tenant_shard_id.shard_slug());
    1081              :         metrics::METRICS_REGISTRY
    1082              :             .metrics_group
    1083              :             .storage_controller_reconcile_spawn
    1084              :             .inc();
    1085              :         let result_tx = result_tx.clone();
    1086              :         let join_handle = tokio::task::spawn(
    1087            0 :             async move {
    1088              :                 // Wait for any previous reconcile task to complete before we start
    1089            0 :                 if let Some(old_handle) = old_handle {
    1090            0 :                     old_handle.cancel.cancel();
    1091            0 :                     if let Err(e) = old_handle.handle.await {
    1092              :                         // We can't do much with this other than log it: the task is done, so
    1093              :                         // we may proceed with our work.
    1094            0 :                         tracing::error!("Unexpected join error waiting for reconcile task: {e}");
    1095            0 :                     }
    1096            0 :                 }
    1097              : 
    1098              :                 // Early check for cancellation before doing any work
    1099              :                 // TODO: wrap all remote API operations in cancellation check
    1100              :                 // as well.
    1101            0 :                 if reconciler.cancel.is_cancelled() {
    1102            0 :                     metrics::METRICS_REGISTRY
    1103            0 :                         .metrics_group
    1104            0 :                         .storage_controller_reconcile_complete
    1105            0 :                         .inc(ReconcileCompleteLabelGroup {
    1106            0 :                             status: ReconcileOutcome::Cancel,
    1107            0 :                         });
    1108            0 :                     return;
    1109            0 :                 }
    1110              : 
    1111              :                 // Attempt to make observed state match intent state
    1112            0 :                 let result = reconciler.reconcile().await;
    1113              : 
    1114              :                 // If we know we had a pending compute notification from some previous action, send a notification irrespective
    1115              :                 // of whether the above reconcile() did any work
    1116            0 :                 if result.is_ok() && must_notify {
    1117              :                     // If this fails we will send the need to retry in [`ReconcileResult::pending_compute_notification`]
    1118            0 :                     reconciler.compute_notify().await.ok();
    1119            0 :                 }
    1120              : 
    1121              :                 // Update result counter
    1122            0 :                 let outcome_label = match &result {
    1123            0 :                     Ok(_) => ReconcileOutcome::Success,
    1124            0 :                     Err(ReconcileError::Cancel) => ReconcileOutcome::Cancel,
    1125            0 :                     Err(_) => ReconcileOutcome::Error,
    1126              :                 };
    1127              : 
    1128            0 :                 metrics::METRICS_REGISTRY
    1129            0 :                     .metrics_group
    1130            0 :                     .storage_controller_reconcile_complete
    1131            0 :                     .inc(ReconcileCompleteLabelGroup {
    1132            0 :                         status: outcome_label,
    1133            0 :                     });
    1134            0 : 
    1135            0 :                 // Constructing result implicitly drops Reconciler, freeing any ReconcileUnits before the Service might
    1136            0 :                 // try and schedule more work in response to our result.
    1137            0 :                 let result = ReconcileResult {
    1138            0 :                     sequence: reconcile_seq,
    1139            0 :                     result,
    1140            0 :                     tenant_shard_id: reconciler.tenant_shard_id,
    1141            0 :                     generation: reconciler.generation,
    1142            0 :                     observed: reconciler.observed,
    1143            0 :                     pending_compute_notification: reconciler.compute_notify_failure,
    1144            0 :                 };
    1145            0 : 
    1146            0 :                 result_tx.send(result).ok();
    1147            0 :             }
    1148              :             .instrument(reconciler_span),
    1149              :         );
    1150              : 
    1151              :         self.reconciler = Some(ReconcilerHandle {
    1152              :             sequence: self.sequence,
    1153              :             handle: join_handle,
    1154              :             cancel: reconciler_cancel,
    1155              :         });
    1156              : 
    1157              :         Some(ReconcilerWaiter {
    1158              :             tenant_shard_id: self.tenant_shard_id,
    1159              :             seq_wait: self.waiter.clone(),
    1160              :             error_seq_wait: self.error_waiter.clone(),
    1161              :             error: self.last_error.clone(),
    1162              :             seq: self.sequence,
    1163              :         })
    1164              :     }
    1165              : 
    1166              :     /// Get a waiter for any reconciliation in flight, but do not start reconciliation
    1167              :     /// if it is not already running
    1168            0 :     pub(crate) fn get_waiter(&self) -> Option<ReconcilerWaiter> {
    1169            0 :         if self.reconciler.is_some() {
    1170            0 :             Some(ReconcilerWaiter {
    1171            0 :                 tenant_shard_id: self.tenant_shard_id,
    1172            0 :                 seq_wait: self.waiter.clone(),
    1173            0 :                 error_seq_wait: self.error_waiter.clone(),
    1174            0 :                 error: self.last_error.clone(),
    1175            0 :                 seq: self.sequence,
    1176            0 :             })
    1177              :         } else {
    1178            0 :             None
    1179              :         }
    1180            0 :     }
    1181              : 
    1182              :     /// Called when a ReconcileResult has been emitted and the service is updating
    1183              :     /// our state: if the result is from a sequence >= my ReconcileHandle, then drop
    1184              :     /// the handle to indicate there is no longer a reconciliation in progress.
    1185            0 :     pub(crate) fn reconcile_complete(&mut self, sequence: Sequence) {
    1186            0 :         if let Some(reconcile_handle) = &self.reconciler {
    1187            0 :             if reconcile_handle.sequence <= sequence {
    1188            0 :                 self.reconciler = None;
    1189            0 :             }
    1190            0 :         }
    1191            0 :     }
    1192              : 
    1193              :     // If we had any state at all referring to this node ID, drop it.  Does not
    1194              :     // attempt to reschedule.
    1195            0 :     pub(crate) fn deref_node(&mut self, node_id: NodeId) {
    1196            0 :         if self.intent.attached == Some(node_id) {
    1197            0 :             self.intent.attached = None;
    1198            0 :         }
    1199              : 
    1200            0 :         self.intent.secondary.retain(|n| n != &node_id);
    1201            0 : 
    1202            0 :         self.observed.locations.remove(&node_id);
    1203            0 : 
    1204            0 :         debug_assert!(!self.intent.all_pageservers().contains(&node_id));
    1205            0 :     }
    1206              : 
    1207            0 :     pub(crate) fn set_scheduling_policy(&mut self, p: ShardSchedulingPolicy) {
    1208            0 :         self.scheduling_policy = p;
    1209            0 :     }
    1210              : 
    1211            0 :     pub(crate) fn get_scheduling_policy(&self) -> &ShardSchedulingPolicy {
    1212            0 :         &self.scheduling_policy
    1213            0 :     }
    1214              : 
    1215            0 :     pub(crate) fn set_last_error(&mut self, sequence: Sequence, error: ReconcileError) {
    1216            0 :         // Ordering: always set last_error before advancing sequence, so that sequence
    1217            0 :         // waiters are guaranteed to see a Some value when they see an error.
    1218            0 :         *(self.last_error.lock().unwrap()) = Some(Arc::new(error));
    1219            0 :         self.error_waiter.advance(sequence);
    1220            0 :     }
    1221              : 
    1222            0 :     pub(crate) fn from_persistent(
    1223            0 :         tsp: TenantShardPersistence,
    1224            0 :         intent: IntentState,
    1225            0 :     ) -> anyhow::Result<Self> {
    1226            0 :         let tenant_shard_id = tsp.get_tenant_shard_id()?;
    1227            0 :         let shard_identity = tsp.get_shard_identity()?;
    1228              : 
    1229            0 :         Ok(Self {
    1230            0 :             tenant_shard_id,
    1231            0 :             shard: shard_identity,
    1232            0 :             sequence: Sequence::initial(),
    1233            0 :             generation: tsp.generation.map(|g| Generation::new(g as u32)),
    1234            0 :             policy: serde_json::from_str(&tsp.placement_policy).unwrap(),
    1235            0 :             intent,
    1236            0 :             observed: ObservedState::new(),
    1237            0 :             config: serde_json::from_str(&tsp.config).unwrap(),
    1238            0 :             reconciler: None,
    1239            0 :             splitting: tsp.splitting,
    1240            0 :             waiter: Arc::new(SeqWait::new(Sequence::initial())),
    1241            0 :             error_waiter: Arc::new(SeqWait::new(Sequence::initial())),
    1242            0 :             last_error: Arc::default(),
    1243            0 :             pending_compute_notification: false,
    1244            0 :             delayed_reconcile: false,
    1245            0 :             scheduling_policy: serde_json::from_str(&tsp.scheduling_policy).unwrap(),
    1246            0 :         })
    1247            0 :     }
    1248              : 
    1249            0 :     pub(crate) fn to_persistent(&self) -> TenantShardPersistence {
    1250            0 :         TenantShardPersistence {
    1251            0 :             tenant_id: self.tenant_shard_id.tenant_id.to_string(),
    1252            0 :             shard_number: self.tenant_shard_id.shard_number.0 as i32,
    1253            0 :             shard_count: self.tenant_shard_id.shard_count.literal() as i32,
    1254            0 :             shard_stripe_size: self.shard.stripe_size.0 as i32,
    1255            0 :             generation: self.generation.map(|g| g.into().unwrap_or(0) as i32),
    1256            0 :             generation_pageserver: self.intent.get_attached().map(|n| n.0 as i64),
    1257            0 :             placement_policy: serde_json::to_string(&self.policy).unwrap(),
    1258            0 :             config: serde_json::to_string(&self.config).unwrap(),
    1259            0 :             splitting: SplitState::default(),
    1260            0 :             scheduling_policy: serde_json::to_string(&self.scheduling_policy).unwrap(),
    1261            0 :         }
    1262            0 :     }
    1263              : }
    1264              : 
    1265              : #[cfg(test)]
    1266              : pub(crate) mod tests {
    1267              :     use pageserver_api::{
    1268              :         controller_api::NodeAvailability,
    1269              :         shard::{ShardCount, ShardNumber},
    1270              :     };
    1271              :     use utils::id::TenantId;
    1272              : 
    1273              :     use crate::scheduler::test_utils::make_test_nodes;
    1274              : 
    1275              :     use super::*;
    1276              : 
    1277           14 :     fn make_test_tenant_shard(policy: PlacementPolicy) -> TenantShard {
    1278           14 :         let tenant_id = TenantId::generate();
    1279           14 :         let shard_number = ShardNumber(0);
    1280           14 :         let shard_count = ShardCount::new(1);
    1281           14 : 
    1282           14 :         let tenant_shard_id = TenantShardId {
    1283           14 :             tenant_id,
    1284           14 :             shard_number,
    1285           14 :             shard_count,
    1286           14 :         };
    1287           14 :         TenantShard::new(
    1288           14 :             tenant_shard_id,
    1289           14 :             ShardIdentity::new(
    1290           14 :                 shard_number,
    1291           14 :                 shard_count,
    1292           14 :                 pageserver_api::shard::ShardStripeSize(32768),
    1293           14 :             )
    1294           14 :             .unwrap(),
    1295           14 :             policy,
    1296           14 :         )
    1297           14 :     }
    1298              : 
    1299            2 :     fn make_test_tenant(policy: PlacementPolicy, shard_count: ShardCount) -> Vec<TenantShard> {
    1300            2 :         let tenant_id = TenantId::generate();
    1301            2 : 
    1302            2 :         (0..shard_count.count())
    1303            8 :             .map(|i| {
    1304            8 :                 let shard_number = ShardNumber(i);
    1305            8 : 
    1306            8 :                 let tenant_shard_id = TenantShardId {
    1307            8 :                     tenant_id,
    1308            8 :                     shard_number,
    1309            8 :                     shard_count,
    1310            8 :                 };
    1311            8 :                 TenantShard::new(
    1312            8 :                     tenant_shard_id,
    1313            8 :                     ShardIdentity::new(
    1314            8 :                         shard_number,
    1315            8 :                         shard_count,
    1316            8 :                         pageserver_api::shard::ShardStripeSize(32768),
    1317            8 :                     )
    1318            8 :                     .unwrap(),
    1319            8 :                     policy.clone(),
    1320            8 :                 )
    1321            8 :             })
    1322            2 :             .collect()
    1323            2 :     }
    1324              : 
    1325              :     /// Test the scheduling behaviors used when a tenant configured for HA is subject
    1326              :     /// to nodes being marked offline.
    1327              :     #[test]
    1328            2 :     fn tenant_ha_scheduling() -> anyhow::Result<()> {
    1329            2 :         // Start with three nodes.  Our tenant will only use two.  The third one is
    1330            2 :         // expected to remain unused.
    1331            2 :         let mut nodes = make_test_nodes(3);
    1332            2 : 
    1333            2 :         let mut scheduler = Scheduler::new(nodes.values());
    1334            2 :         let mut context = ScheduleContext::default();
    1335            2 : 
    1336            2 :         let mut tenant_shard = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1337            2 :         tenant_shard
    1338            2 :             .schedule(&mut scheduler, &mut context)
    1339            2 :             .expect("we have enough nodes, scheduling should work");
    1340            2 : 
    1341            2 :         // Expect to initially be schedule on to different nodes
    1342            2 :         assert_eq!(tenant_shard.intent.secondary.len(), 1);
    1343            2 :         assert!(tenant_shard.intent.attached.is_some());
    1344              : 
    1345            2 :         let attached_node_id = tenant_shard.intent.attached.unwrap();
    1346            2 :         let secondary_node_id = *tenant_shard.intent.secondary.iter().last().unwrap();
    1347            2 :         assert_ne!(attached_node_id, secondary_node_id);
    1348              : 
    1349              :         // Notifying the attached node is offline should demote it to a secondary
    1350            2 :         let changed = tenant_shard
    1351            2 :             .intent
    1352            2 :             .demote_attached(&mut scheduler, attached_node_id);
    1353            2 :         assert!(changed);
    1354            2 :         assert!(tenant_shard.intent.attached.is_none());
    1355            2 :         assert_eq!(tenant_shard.intent.secondary.len(), 2);
    1356              : 
    1357              :         // Update the scheduler state to indicate the node is offline
    1358            2 :         nodes
    1359            2 :             .get_mut(&attached_node_id)
    1360            2 :             .unwrap()
    1361            2 :             .set_availability(NodeAvailability::Offline);
    1362            2 :         scheduler.node_upsert(nodes.get(&attached_node_id).unwrap());
    1363            2 : 
    1364            2 :         // Scheduling the node should promote the still-available secondary node to attached
    1365            2 :         tenant_shard
    1366            2 :             .schedule(&mut scheduler, &mut context)
    1367            2 :             .expect("active nodes are available");
    1368            2 :         assert_eq!(tenant_shard.intent.attached.unwrap(), secondary_node_id);
    1369              : 
    1370              :         // The original attached node should have been retained as a secondary
    1371            2 :         assert_eq!(
    1372            2 :             *tenant_shard.intent.secondary.iter().last().unwrap(),
    1373            2 :             attached_node_id
    1374            2 :         );
    1375              : 
    1376            2 :         tenant_shard.intent.clear(&mut scheduler);
    1377            2 : 
    1378            2 :         Ok(())
    1379            2 :     }
    1380              : 
    1381              :     #[test]
    1382            2 :     fn intent_from_observed() -> anyhow::Result<()> {
    1383            2 :         let nodes = make_test_nodes(3);
    1384            2 :         let mut scheduler = Scheduler::new(nodes.values());
    1385            2 : 
    1386            2 :         let mut tenant_shard = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1387            2 : 
    1388            2 :         tenant_shard.observed.locations.insert(
    1389            2 :             NodeId(3),
    1390            2 :             ObservedStateLocation {
    1391            2 :                 conf: Some(LocationConfig {
    1392            2 :                     mode: LocationConfigMode::AttachedMulti,
    1393            2 :                     generation: Some(2),
    1394            2 :                     secondary_conf: None,
    1395            2 :                     shard_number: tenant_shard.shard.number.0,
    1396            2 :                     shard_count: tenant_shard.shard.count.literal(),
    1397            2 :                     shard_stripe_size: tenant_shard.shard.stripe_size.0,
    1398            2 :                     tenant_conf: TenantConfig::default(),
    1399            2 :                 }),
    1400            2 :             },
    1401            2 :         );
    1402            2 : 
    1403            2 :         tenant_shard.observed.locations.insert(
    1404            2 :             NodeId(2),
    1405            2 :             ObservedStateLocation {
    1406            2 :                 conf: Some(LocationConfig {
    1407            2 :                     mode: LocationConfigMode::AttachedStale,
    1408            2 :                     generation: Some(1),
    1409            2 :                     secondary_conf: None,
    1410            2 :                     shard_number: tenant_shard.shard.number.0,
    1411            2 :                     shard_count: tenant_shard.shard.count.literal(),
    1412            2 :                     shard_stripe_size: tenant_shard.shard.stripe_size.0,
    1413            2 :                     tenant_conf: TenantConfig::default(),
    1414            2 :                 }),
    1415            2 :             },
    1416            2 :         );
    1417            2 : 
    1418            2 :         tenant_shard.intent_from_observed(&mut scheduler);
    1419            2 : 
    1420            2 :         // The highest generationed attached location gets used as attached
    1421            2 :         assert_eq!(tenant_shard.intent.attached, Some(NodeId(3)));
    1422              :         // Other locations get used as secondary
    1423            2 :         assert_eq!(tenant_shard.intent.secondary, vec![NodeId(2)]);
    1424              : 
    1425            2 :         scheduler.consistency_check(nodes.values(), [&tenant_shard].into_iter())?;
    1426              : 
    1427            2 :         tenant_shard.intent.clear(&mut scheduler);
    1428            2 :         Ok(())
    1429            2 :     }
    1430              : 
    1431              :     #[test]
    1432            2 :     fn scheduling_mode() -> anyhow::Result<()> {
    1433            2 :         let nodes = make_test_nodes(3);
    1434            2 :         let mut scheduler = Scheduler::new(nodes.values());
    1435            2 : 
    1436            2 :         let mut tenant_shard = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1437            2 : 
    1438            2 :         // In pause mode, schedule() shouldn't do anything
    1439            2 :         tenant_shard.scheduling_policy = ShardSchedulingPolicy::Pause;
    1440            2 :         assert!(tenant_shard
    1441            2 :             .schedule(&mut scheduler, &mut ScheduleContext::default())
    1442            2 :             .is_ok());
    1443            2 :         assert!(tenant_shard.intent.all_pageservers().is_empty());
    1444              : 
    1445              :         // In active mode, schedule() works
    1446            2 :         tenant_shard.scheduling_policy = ShardSchedulingPolicy::Active;
    1447            2 :         assert!(tenant_shard
    1448            2 :             .schedule(&mut scheduler, &mut ScheduleContext::default())
    1449            2 :             .is_ok());
    1450            2 :         assert!(!tenant_shard.intent.all_pageservers().is_empty());
    1451              : 
    1452            2 :         tenant_shard.intent.clear(&mut scheduler);
    1453            2 :         Ok(())
    1454            2 :     }
    1455              : 
    1456              :     #[test]
    1457            2 :     fn optimize_attachment() -> anyhow::Result<()> {
    1458            2 :         let nodes = make_test_nodes(3);
    1459            2 :         let mut scheduler = Scheduler::new(nodes.values());
    1460            2 : 
    1461            2 :         let mut shard_a = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1462            2 :         let mut shard_b = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1463            2 : 
    1464            2 :         // Initially: both nodes attached on shard 1, and both have secondary locations
    1465            2 :         // on different nodes.
    1466            2 :         shard_a.intent.set_attached(&mut scheduler, Some(NodeId(1)));
    1467            2 :         shard_a.intent.push_secondary(&mut scheduler, NodeId(2));
    1468            2 :         shard_b.intent.set_attached(&mut scheduler, Some(NodeId(1)));
    1469            2 :         shard_b.intent.push_secondary(&mut scheduler, NodeId(3));
    1470            2 : 
    1471            2 :         let mut schedule_context = ScheduleContext::default();
    1472            2 :         schedule_context.avoid(&shard_a.intent.all_pageservers());
    1473            2 :         schedule_context.push_attached(shard_a.intent.get_attached().unwrap());
    1474            2 :         schedule_context.avoid(&shard_b.intent.all_pageservers());
    1475            2 :         schedule_context.push_attached(shard_b.intent.get_attached().unwrap());
    1476            2 : 
    1477            2 :         let optimization_a = shard_a.optimize_attachment(&nodes, &schedule_context);
    1478            2 : 
    1479            2 :         // Either shard should recognize that it has the option to switch to a secondary location where there
    1480            2 :         // would be no other shards from the same tenant, and request to do so.
    1481            2 :         assert_eq!(
    1482            2 :             optimization_a,
    1483            2 :             Some(ScheduleOptimization {
    1484            2 :                 sequence: shard_a.sequence,
    1485            2 :                 action: ScheduleOptimizationAction::MigrateAttachment(MigrateAttachment {
    1486            2 :                     old_attached_node_id: NodeId(1),
    1487            2 :                     new_attached_node_id: NodeId(2)
    1488            2 :                 })
    1489            2 :             })
    1490            2 :         );
    1491              : 
    1492              :         // Note that these optimizing two shards in the same tenant with the same ScheduleContext is
    1493              :         // mutually exclusive (the optimization of one invalidates the stats) -- it is the responsibility
    1494              :         // of [`Service::optimize_all`] to avoid trying
    1495              :         // to do optimizations for multiple shards in the same tenant at the same time.  Generating
    1496              :         // both optimizations is just done for test purposes
    1497            2 :         let optimization_b = shard_b.optimize_attachment(&nodes, &schedule_context);
    1498            2 :         assert_eq!(
    1499            2 :             optimization_b,
    1500            2 :             Some(ScheduleOptimization {
    1501            2 :                 sequence: shard_b.sequence,
    1502            2 :                 action: ScheduleOptimizationAction::MigrateAttachment(MigrateAttachment {
    1503            2 :                     old_attached_node_id: NodeId(1),
    1504            2 :                     new_attached_node_id: NodeId(3)
    1505            2 :                 })
    1506            2 :             })
    1507            2 :         );
    1508              : 
    1509              :         // Applying these optimizations should result in the end state proposed
    1510            2 :         shard_a.apply_optimization(&mut scheduler, optimization_a.unwrap());
    1511            2 :         assert_eq!(shard_a.intent.get_attached(), &Some(NodeId(2)));
    1512            2 :         assert_eq!(shard_a.intent.get_secondary(), &vec![NodeId(1)]);
    1513            2 :         shard_b.apply_optimization(&mut scheduler, optimization_b.unwrap());
    1514            2 :         assert_eq!(shard_b.intent.get_attached(), &Some(NodeId(3)));
    1515            2 :         assert_eq!(shard_b.intent.get_secondary(), &vec![NodeId(1)]);
    1516              : 
    1517            2 :         shard_a.intent.clear(&mut scheduler);
    1518            2 :         shard_b.intent.clear(&mut scheduler);
    1519            2 : 
    1520            2 :         Ok(())
    1521            2 :     }
    1522              : 
    1523              :     #[test]
    1524            2 :     fn optimize_secondary() -> anyhow::Result<()> {
    1525            2 :         let nodes = make_test_nodes(4);
    1526            2 :         let mut scheduler = Scheduler::new(nodes.values());
    1527            2 : 
    1528            2 :         let mut shard_a = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1529            2 :         let mut shard_b = make_test_tenant_shard(PlacementPolicy::Attached(1));
    1530            2 : 
    1531            2 :         // Initially: both nodes attached on shard 1, and both have secondary locations
    1532            2 :         // on different nodes.
    1533            2 :         shard_a.intent.set_attached(&mut scheduler, Some(NodeId(1)));
    1534            2 :         shard_a.intent.push_secondary(&mut scheduler, NodeId(3));
    1535            2 :         shard_b.intent.set_attached(&mut scheduler, Some(NodeId(2)));
    1536            2 :         shard_b.intent.push_secondary(&mut scheduler, NodeId(3));
    1537            2 : 
    1538            2 :         let mut schedule_context = ScheduleContext::default();
    1539            2 :         schedule_context.avoid(&shard_a.intent.all_pageservers());
    1540            2 :         schedule_context.push_attached(shard_a.intent.get_attached().unwrap());
    1541            2 :         schedule_context.avoid(&shard_b.intent.all_pageservers());
    1542            2 :         schedule_context.push_attached(shard_b.intent.get_attached().unwrap());
    1543            2 : 
    1544            2 :         let optimization_a = shard_a.optimize_secondary(&scheduler, &schedule_context);
    1545            2 : 
    1546            2 :         // Since there is a node with no locations available, the node with two locations for the
    1547            2 :         // same tenant should generate an optimization to move one away
    1548            2 :         assert_eq!(
    1549            2 :             optimization_a,
    1550            2 :             Some(ScheduleOptimization {
    1551            2 :                 sequence: shard_a.sequence,
    1552            2 :                 action: ScheduleOptimizationAction::ReplaceSecondary(ReplaceSecondary {
    1553            2 :                     old_node_id: NodeId(3),
    1554            2 :                     new_node_id: NodeId(4)
    1555            2 :                 })
    1556            2 :             })
    1557            2 :         );
    1558              : 
    1559            2 :         shard_a.apply_optimization(&mut scheduler, optimization_a.unwrap());
    1560            2 :         assert_eq!(shard_a.intent.get_attached(), &Some(NodeId(1)));
    1561            2 :         assert_eq!(shard_a.intent.get_secondary(), &vec![NodeId(4)]);
    1562              : 
    1563            2 :         shard_a.intent.clear(&mut scheduler);
    1564            2 :         shard_b.intent.clear(&mut scheduler);
    1565            2 : 
    1566            2 :         Ok(())
    1567            2 :     }
    1568              : 
    1569              :     // Optimize til quiescent: this emulates what Service::optimize_all does, when
    1570              :     // called repeatedly in the background.
    1571            2 :     fn optimize_til_idle(
    1572            2 :         nodes: &HashMap<NodeId, Node>,
    1573            2 :         scheduler: &mut Scheduler,
    1574            2 :         shards: &mut [TenantShard],
    1575            2 :     ) {
    1576            2 :         let mut loop_n = 0;
    1577              :         loop {
    1578           18 :             let mut schedule_context = ScheduleContext::default();
    1579           18 :             let mut any_changed = false;
    1580              : 
    1581           72 :             for shard in shards.iter() {
    1582           72 :                 schedule_context.avoid(&shard.intent.all_pageservers());
    1583           72 :                 if let Some(attached) = shard.intent.get_attached() {
    1584           72 :                     schedule_context.push_attached(*attached);
    1585           72 :                 }
    1586              :             }
    1587              : 
    1588           36 :             for shard in shards.iter_mut() {
    1589           36 :                 let optimization = shard.optimize_attachment(nodes, &schedule_context);
    1590           36 :                 if let Some(optimization) = optimization {
    1591            8 :                     shard.apply_optimization(scheduler, optimization);
    1592            8 :                     any_changed = true;
    1593            8 :                     break;
    1594           28 :                 }
    1595           28 : 
    1596           28 :                 let optimization = shard.optimize_secondary(scheduler, &schedule_context);
    1597           28 :                 if let Some(optimization) = optimization {
    1598            8 :                     shard.apply_optimization(scheduler, optimization);
    1599            8 :                     any_changed = true;
    1600            8 :                     break;
    1601           20 :                 }
    1602              :             }
    1603              : 
    1604           18 :             if !any_changed {
    1605            2 :                 break;
    1606           16 :             }
    1607           16 : 
    1608           16 :             // Assert no infinite loop
    1609           16 :             loop_n += 1;
    1610           16 :             assert!(loop_n < 1000);
    1611              :         }
    1612            2 :     }
    1613              : 
    1614              :     /// Test the balancing behavior of shard scheduling: that it achieves a balance, and
    1615              :     /// that it converges.
    1616              :     #[test]
    1617            2 :     fn optimize_add_nodes() -> anyhow::Result<()> {
    1618            2 :         let nodes = make_test_nodes(4);
    1619            2 : 
    1620            2 :         // Only show the scheduler a couple of nodes
    1621            2 :         let mut scheduler = Scheduler::new([].iter());
    1622            2 :         scheduler.node_upsert(nodes.get(&NodeId(1)).unwrap());
    1623            2 :         scheduler.node_upsert(nodes.get(&NodeId(2)).unwrap());
    1624            2 : 
    1625            2 :         let mut shards = make_test_tenant(PlacementPolicy::Attached(1), ShardCount::new(4));
    1626            2 :         let mut schedule_context = ScheduleContext::default();
    1627           10 :         for shard in &mut shards {
    1628            8 :             assert!(shard
    1629            8 :                 .schedule(&mut scheduler, &mut schedule_context)
    1630            8 :                 .is_ok());
    1631              :         }
    1632              : 
    1633              :         // We should see equal number of locations on the two nodes.
    1634            2 :         assert_eq!(scheduler.get_node_shard_count(NodeId(1)), 4);
    1635              :         // Scheduling does not consider the number of attachments picking the initial
    1636              :         // pageserver to attach to (hence the assertion that all primaries are on the
    1637              :         // same node)
    1638              :         // TODO: Tweak the scheduling to evenly distribute attachments for new shards.
    1639            2 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(1)), 4);
    1640              : 
    1641            2 :         assert_eq!(scheduler.get_node_shard_count(NodeId(2)), 4);
    1642            2 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(2)), 0);
    1643              : 
    1644              :         // Add another two nodes: we should see the shards spread out when their optimize
    1645              :         // methods are called
    1646            2 :         scheduler.node_upsert(nodes.get(&NodeId(3)).unwrap());
    1647            2 :         scheduler.node_upsert(nodes.get(&NodeId(4)).unwrap());
    1648            2 :         optimize_til_idle(&nodes, &mut scheduler, &mut shards);
    1649            2 : 
    1650            2 :         assert_eq!(scheduler.get_node_shard_count(NodeId(1)), 2);
    1651            2 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(1)), 1);
    1652              : 
    1653            2 :         assert_eq!(scheduler.get_node_shard_count(NodeId(2)), 2);
    1654            2 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(2)), 1);
    1655              : 
    1656            2 :         assert_eq!(scheduler.get_node_shard_count(NodeId(3)), 2);
    1657            2 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(3)), 1);
    1658              : 
    1659            2 :         assert_eq!(scheduler.get_node_shard_count(NodeId(4)), 2);
    1660            2 :         assert_eq!(scheduler.get_node_attached_shard_count(NodeId(4)), 1);
    1661              : 
    1662            8 :         for shard in shards.iter_mut() {
    1663            8 :             shard.intent.clear(&mut scheduler);
    1664            8 :         }
    1665              : 
    1666            2 :         Ok(())
    1667            2 :     }
    1668              : }
        

Generated by: LCOV version 2.1-beta