storcon: add transactional-ish demotion to secondary

storage_controller/src/service.rs

@@ -5322,28 +5322,30 @@ impl Service {
                         }
                     };
 
-                    if tenant_shard.intent.demote_attached(scheduler, node_id) {
-                        match tenant_shard.schedule(scheduler, &mut schedule_context) {
-                            Err(e) => {
-                                tracing::warn!(
-                                    tenant_id=%tid.tenant_id, shard_id=%tid.shard_slug(),
-                                    "Scheduling error when draining pageserver {} : {e}", node_id
-                                );
-                            }
-                            Ok(()) => {
-                                let scheduled_to = tenant_shard.intent.get_attached();
-                                tracing::info!(
-                                    tenant_id=%tid.tenant_id, shard_id=%tid.shard_slug(),
-                                    "Rescheduled shard while draining node {}: {} -> {:?}",
-                                    node_id,
-                                    node_id,
-                                    scheduled_to
-                                );
+                    match tenant_shard.reschedule_to_secondary(
+                        node_id,
+                        scheduler,
+                        &mut schedule_context,
+                    ) {
+                        Err(e) => {
+                            tracing::warn!(
+                                tenant_id=%tid.tenant_id, shard_id=%tid.shard_slug(),
+                                "Scheduling error when draining pageserver {} : {e}", node_id
+                            );
+                        }
+                        Ok(()) => {
+                            let scheduled_to = tenant_shard.intent.get_attached();
+                            tracing::info!(
+                                tenant_id=%tid.tenant_id, shard_id=%tid.shard_slug(),
+                                "Rescheduled shard while draining node {}: {} -> {:?}",
+                                node_id,
+                                node_id,
+                                scheduled_to
+                            );
 
-                                let waiter = self.maybe_reconcile_shard(tenant_shard, nodes);
-                                if let Some(some) = waiter {
-                                    waiters.push(some);
-                                }
+                            let waiter = self.maybe_reconcile_shard(tenant_shard, nodes);
+                            if let Some(some) = waiter {
+                                waiters.push(some);
                             }
                         }
                     }

storage_controller/src/tenant_shard.rs

@@ -646,6 +646,49 @@ impl TenantShard {
         Ok(())
     }
 
+    /// Attempt to reschedule the tenant shard to one of its secondary locations.
+    /// If no schedulable secondary is found, roll back any changes to the intent state.
+    /// If the `attached_to` node argument doesn't match the intent state, no changes
+    /// are made and the function returns successfully.
+    pub(crate) fn reschedule_to_secondary(
+        &mut self,
+        attached_to: NodeId,
+        scheduler: &mut Scheduler,
+        context: &mut ScheduleContext,
+    ) -> Result<(), ScheduleError> {
+        let original_secondaries = self.intent.get_secondary().clone();
+
+        if self.intent.demote_attached(scheduler, attached_to) {
+            let res = match self.schedule(scheduler, context) {
+                Ok(()) => {
+                    let rescheduled_to_secondary =
+                        self.intent.get_attached().map_or(false, |node| {
+                            original_secondaries.iter().any(|sec| *sec == node)
+                        });
+
+                    if rescheduled_to_secondary {
+                        Ok(())
+                    } else {
+                        Err(ScheduleError::ImpossibleConstraint)
+                    }
+                }
+                Err(e) => Err(e),
+            };
+
+            if res.is_err() {
+                self.intent.set_attached(scheduler, Some(attached_to));
+                self.intent.clear_secondary(scheduler);
+                for sec in original_secondaries {
+                    self.intent.push_secondary(scheduler, sec);
+                }
+            }
+
+            return res;
+        }
+
+        Ok(())
+    }
+
     /// Optimize attachments: if a shard has a secondary location that is preferable to
     /// its primary location based on soft constraints, switch that secondary location
     /// to be attached.