diff --git a/control_plane/storcon_cli/src/main.rs b/control_plane/storcon_cli/src/main.rs index 2b3861f448..8ab8688622 100644 --- a/control_plane/storcon_cli/src/main.rs +++ b/control_plane/storcon_cli/src/main.rs @@ -920,7 +920,7 @@ async fn main() -> anyhow::Result<()> { storcon_client .dispatch::<(), ()>( Method::PUT, - format!("control/v1/node/{node_id}/drain"), + format!("control/v1/node/{node_id}/drain?graceful=true"), None, ) .await?; diff --git a/storage_controller/src/http.rs b/storage_controller/src/http.rs index 2b1c0db12f..0834b632a5 100644 --- a/storage_controller/src/http.rs +++ b/storage_controller/src/http.rs @@ -47,7 +47,9 @@ use crate::metrics::{ }; use crate::persistence::SafekeeperUpsert; use crate::reconciler::ReconcileError; -use crate::service::{LeadershipStatus, RECONCILE_TIMEOUT, STARTUP_RECONCILE_TIMEOUT, Service}; +use crate::service::{ + LeadershipStatus, NodeDrainMode, RECONCILE_TIMEOUT, STARTUP_RECONCILE_TIMEOUT, Service, +}; /// State available to HTTP request handlers pub struct HttpState { @@ -1000,8 +1002,18 @@ async fn handle_node_drain(req: Request) -> Result, ApiErro let state = get_state(&req); let node_id: NodeId = parse_request_param(&req, "node_id")?; + let graceful: bool = parse_query_param(&req, "graceful")?.unwrap_or(false); - state.service.start_node_drain(node_id).await?; + let node_drain_mode = if graceful { + NodeDrainMode::FullWithReprovisioning + } else { + NodeDrainMode::AttachedOnly + }; + + state + .service + .start_node_drain(node_id, node_drain_mode) + .await?; json_response(StatusCode::ACCEPTED, ()) } diff --git a/storage_controller/src/service.rs b/storage_controller/src/service.rs index 790797bae2..18c81c9bca 100644 --- a/storage_controller/src/service.rs +++ b/storage_controller/src/service.rs @@ -687,6 +687,16 @@ struct ShardMutationLocations { #[derive(Default, Clone)] struct TenantMutationLocations(BTreeMap); +#[allow(dead_code)] +pub(crate) enum NodeDrainMode { + /// Drain only attached locations (primary node itself) + AttachedOnly, + /// Drain attached and secondary locations (but don't create new secondaries) + AttachedAndSecondaries, + /// Drain everything and also provision new secondaries for the newly promoted primaries + FullWithReprovisioning, +} + impl Service { pub fn get_config(&self) -> &Config { &self.config @@ -7548,6 +7558,7 @@ impl Service { pub(crate) async fn start_node_drain( self: &Arc, node_id: NodeId, + _mode: NodeDrainMode, ) -> Result<(), ApiError> { let (ongoing_op, node_available, node_policy, schedulable_nodes_count) = { let locked = self.inner.read().unwrap();