feat(rack): make maintenance requests atomic and idempotent

crates/api-core/src/handlers/component_manager.rs

@@ -68,7 +68,7 @@ fn unsupported_from_json_firmware_versions(target: &str) -> Status {
     ))
 }
 
-fn component_manager_error_to_status(err: ComponentManagerError) -> Status {
+pub(crate) fn component_manager_error_to_status(err: ComponentManagerError) -> Status {
     match err {
         ComponentManagerError::Unavailable(msg) => Status::unavailable(msg),
         ComponentManagerError::NotFound(msg) => Status::not_found(msg),

crates/api-core/src/handlers/rack.rs

@@ -22,11 +22,14 @@ use ::rpc::forge::{self as rpc, HealthReportEntry};
 use carbide_rack::firmware_object::{
     rack_maintenance_access_token_key, rms_access_token_or_noauth,
 };
-use carbide_secrets::credentials::Credentials;
 use carbide_uuid::machine::MachineId;
 use carbide_uuid::power_shelf::PowerShelfId;
 use carbide_uuid::rack::RackId;
 use carbide_uuid::switch::SwitchId;
+use component_manager::component_manager::{
+    RackMaintenanceAccessToken, RackMaintenanceEligibility, RackMaintenanceRequestOutcome,
+    request_rack_maintenance_via_state_controller,
+};
 use db::{
     ObjectColumnFilter, WithTransaction, machine as db_machine, power_shelf as db_power_shelf,
     rack as db_rack, switch as db_switch,
@@ -41,6 +44,7 @@ use tonic::{Request, Response, Status};
 use crate::CarbideError;
 use crate::api::{Api, log_request_data, log_request_data_redacted};
 use crate::auth::AuthContext;
+use crate::handlers::component_manager::component_manager_error_to_status;
 
 pub async fn get_rack(
     api: &Api,
@@ -242,6 +246,8 @@ pub async fn admin_force_delete_rack(
         .into());
     }
 
+    let maintenance_request_id = rack_list[0].config.maintenance_request_id;
+
     db_rack::final_delete(&mut txn, &rack_id)
         .await
         .map_err(CarbideError::from)?;
@@ -250,7 +256,10 @@ pub async fn admin_force_delete_rack(
 
     if let Err(error) = api
         .credential_manager
-        .delete_credentials(&rack_maintenance_access_token_key(&rack_id))
+        .delete_credentials(&rack_maintenance_access_token_key(
+            &rack_id,
+            maintenance_request_id,
+        ))
         .await
     {
         tracing::warn!(
@@ -766,58 +775,44 @@ pub(crate) async fn on_demand_rack_maintenance(
         }
     }
 
-    let access_token_stored = maintenance_access_token.is_some();
-    if let Some(token) = maintenance_access_token {
-        api.credential_manager
-            .set_credentials(
-                &rack_maintenance_access_token_key(&rack_id),
-                &Credentials::UsernamePassword {
-                    username: "access_token".into(),
-                    password: token,
-                },
-            )
-            .await
-            .map_err(|error| CarbideError::Internal {
-                message: format!("failed to store rack maintenance access token: {error}"),
-            })?;
-    }
-
-    let mut updated_config = rack.config.clone();
-    updated_config.maintenance_requested = Some(scope);
+    let maintenance_access_token =
+        maintenance_access_token.map(|token| RackMaintenanceAccessToken {
+            credential_manager: api.credential_manager.as_ref(),
+            token,
+        });
 
-    let db_result: Result<(), Status> = async {
-        let mut txn = api.txn_begin().await?;
-        db_rack::update(&mut txn, &rack_id, &updated_config).await?;
-        if updated_config
-            .maintenance_requested
-            .as_ref()
-            .is_some_and(|scope| {
-                scope.should_run(&MaintenanceActivity::FirmwareUpgrade {
-                    firmware_version: None,
-                    components: vec![],
-                    force_update: false,
-                })
-            })
-        {
-            db_rack::update_firmware_upgrade_job(txn.as_mut(), &rack_id, None).await?;
-        }
-        txn.commit().await?;
-        Ok(())
-    }
+    let schedule_result = request_rack_maintenance_via_state_controller(
+        &api.database_connection,
+        &rack_id,
+        scope,
+        RackMaintenanceEligibility::AllowError,
+        maintenance_access_token,
+    )
     .await;
-    if let Err(status) = db_result {
-        if access_token_stored
-            && let Err(error) = api
-                .credential_manager
-                .delete_credentials(&rack_maintenance_access_token_key(&rack_id))
-                .await
-        {
-            tracing::warn!(
-                rack_id = %rack_id,
-                error = %error,
-                "failed to delete rack maintenance access token after DB error",
-            );
-        }
+
+    let scheduling_error = match schedule_result {
+        Ok(
+            RackMaintenanceRequestOutcome::Scheduled { .. }
+            | RackMaintenanceRequestOutcome::AlreadyPending { .. },
+        ) => None,
+        Ok(RackMaintenanceRequestOutcome::Busy) => Some(
+            CarbideError::InvalidArgument(format!(
+                "On-demand maintenance for rack {} is already scheduled.",
+                rack_id,
+            ))
+            .into(),
+        ),
+        Ok(RackMaintenanceRequestOutcome::Deferred { state }) => Some(
+            CarbideError::InvalidArgument(format!(
+                "Rack {} is not in Ready or Error state (current: {:?}). Maintenance can only be requested when the rack is Ready or in Error.",
+                rack_id, state,
+            ))
+            .into(),
+        ),
+        Err(error) => Some(component_manager_error_to_status(error)),
+    };
+
+    if let Some(status) = scheduling_error {
         return Err(status);
     }
 

crates/api-core/src/tests/rack_state_controller/handler.rs

@@ -36,7 +36,8 @@ use model::rack::{
     ConfigureNmxClusterCertificateState, ConfigureNmxClusterState, FirmwareUpgradeDeviceStatus,
     FirmwareUpgradeJob, FirmwareUpgradeState, MaintenanceActivity, MaintenanceScope,
     NvosUpdateState, NvosUpdateSwitchStatus, Rack, RackConfig, RackFirmwareUpgradeState,
-    RackMaintenanceState, RackPowerState, RackState, RackValidationState,
+    RackMaintenanceRequestStatus, RackMaintenanceState, RackPowerState, RackState,
+    RackValidationState,
 };
 use model::rack_type::{
     RackCapabilitiesSet, RackCapabilityCompute, RackCapabilityPowerShelf, RackCapabilitySwitch,
@@ -341,6 +342,49 @@ pub(crate) fn new_rack_id() -> RackId {
     RackId::new(uuid::Uuid::new_v4().to_string())
 }
 
+async fn create_ready_rack_with_maintenance_request(
+    pool: &sqlx::PgPool,
+    scope: MaintenanceScope,
+) -> Result<(RackId, uuid::Uuid), Box<dyn std::error::Error>> {
+    let rack_id = new_rack_id();
+    let request_id = uuid::Uuid::new_v4();
+    let config = RackConfig {
+        maintenance_requested: Some(scope.clone()),
+        maintenance_request_id: Some(request_id),
+        ..Default::default()
+    };
+    let mut txn = pool.begin().await?;
+    let rack = db_rack::create(
+        txn.as_mut(),
+        &rack_id,
+        Some(&RackProfileId::new("Empty")),
+        &config,
+        None,
+    )
+    .await?;
+    db::rack_maintenance_request::insert_preparing(
+        txn.as_mut(),
+        request_id,
+        &rack_id,
+        &scope,
+        false,
+    )
+    .await?;
+    assert!(db::rack_maintenance_request::mark_ready(txn.as_mut(), request_id).await?);
+    assert!(
+        db_rack::try_update_controller_state(
+            txn.as_mut(),
+            &rack_id,
+            rack.controller_state.version,
+            rack.controller_state.version.increment(),
+            &RackState::Ready,
+        )
+        .await?
+    );
+    txn.commit().await?;
+    Ok((rack_id, request_id))
+}
+
 async fn create_ready_rack_with_switch(
     env: &TestEnv,
     pool: &sqlx::PgPool,
@@ -1181,6 +1225,147 @@ async fn test_error_state_does_nothing(
     Ok(())
 }
 
+#[crate::sqlx_test]
+async fn test_first_class_maintenance_request_runs_and_completes(
+    pool: sqlx::PgPool,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let env = create_test_env_with_overrides(pool.clone(), TestEnvOverrides::default()).await;
+    let scope = MaintenanceScope {
+        activities: vec![MaintenanceActivity::ConfigureNmxCluster],
+        ..Default::default()
+    };
+    let (rack_id, request_id) = create_ready_rack_with_maintenance_request(&pool, scope).await?;
+    let mut rack = get_db_rack(env.db_reader().as_mut(), &rack_id).await;
+
+    let handler = RackStateHandler::default();
+    let mut services = env.rack_state_handler_services();
+    let mut metrics = RackMetrics::default();
+    let mut db_writes = DbWriteBatch::default();
+    let mut ctx = StateHandlerContext::<RackStateHandlerContextObjects> {
+        services: &mut services,
+        metrics: &mut metrics,
+        pending_db_writes: &mut db_writes,
+    };
+
+    let mut outcome = handler
+        .handle_object_state(&rack_id, &mut rack, &RackState::Ready, &mut ctx)
+        .await?;
+    assert!(matches!(
+        &outcome,
+        StateHandlerOutcome::Transition {
+            next_state: RackState::Maintenance { .. },
+            ..
+        }
+    ));
+    if let Some(txn) = outcome.take_transaction() {
+        txn.commit().await?;
+    }
+
+    let mut conn = pool.acquire().await?;
+    let request = db::rack_maintenance_request::find_by_id(conn.as_mut(), request_id)
+        .await?
+        .unwrap();
+    assert_eq!(request.status, RackMaintenanceRequestStatus::Running);
+    assert!(request.started.is_some());
+    drop(conn);
+
+    let completed_state = RackState::Maintenance {
+        maintenance_state: RackMaintenanceState::Completed,
+    };
+    let mut outcome = handler
+        .handle_object_state(&rack_id, &mut rack, &completed_state, &mut ctx)
+        .await?;
+    assert!(matches!(
+        &outcome,
+        StateHandlerOutcome::Transition {
+            next_state: RackState::Validating {
+                validating_state: RackValidationState::Pending,
+            },
+            ..
+        }
+    ));
+    if let Some(txn) = outcome.take_transaction() {
+        txn.commit().await?;
+    }
+
+    let mut conn = pool.acquire().await?;
+    let request = db::rack_maintenance_request::find_by_id(conn.as_mut(), request_id)
+        .await?
+        .unwrap();
+    assert_eq!(request.status, RackMaintenanceRequestStatus::Completed);
+    assert!(request.completed.is_some());
+    let rack = get_db_rack(conn.as_mut(), &rack_id).await;
+    assert!(rack.config.maintenance_requested.is_none());
+    assert!(rack.config.maintenance_request_id.is_none());
+
+    Ok(())
+}
+
+#[crate::sqlx_test]
+async fn test_first_class_maintenance_request_fails_with_maintenance(
+    pool: sqlx::PgPool,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let env = create_test_env_with_overrides(pool.clone(), TestEnvOverrides::default()).await;
+    let scope = MaintenanceScope {
+        activities: vec![MaintenanceActivity::FirmwareUpgrade {
+            firmware_version: None,
+            components: vec![],
+            force_update: false,
+        }],
+        ..Default::default()
+    };
+    let (rack_id, request_id) = create_ready_rack_with_maintenance_request(&pool, scope).await?;
+    let mut rack = get_db_rack(env.db_reader().as_mut(), &rack_id).await;
+
+    let handler = RackStateHandler::default();
+    let mut services = env.rack_state_handler_services();
+    let mut metrics = RackMetrics::default();
+    let mut db_writes = DbWriteBatch::default();
+    let mut ctx = StateHandlerContext::<RackStateHandlerContextObjects> {
+        services: &mut services,
+        metrics: &mut metrics,
+        pending_db_writes: &mut db_writes,
+    };
+
+    let mut outcome = handler
+        .handle_object_state(&rack_id, &mut rack, &RackState::Ready, &mut ctx)
+        .await?;
+    if let Some(txn) = outcome.take_transaction() {
+        txn.commit().await?;
+    }
+
+    let firmware_state = RackState::Maintenance {
+        maintenance_state: RackMaintenanceState::FirmwareUpgrade {
+            rack_firmware_upgrade: FirmwareUpgradeState::Start,
+        },
+    };
+    let mut outcome = handler
+        .handle_object_state(&rack_id, &mut rack, &firmware_state, &mut ctx)
+        .await?;
+    assert!(matches!(
+        &outcome,
+        StateHandlerOutcome::Transition {
+            next_state: RackState::Error { .. },
+            ..
+        }
+    ));
+    if let Some(txn) = outcome.take_transaction() {
+        txn.commit().await?;
+    }
+
+    let mut conn = pool.acquire().await?;
+    let request = db::rack_maintenance_request::find_by_id(conn.as_mut(), request_id)
+        .await?
+        .unwrap();
+    assert_eq!(request.status, RackMaintenanceRequestStatus::Failed);
+    assert!(request.error_message.is_some());
+    let rack = get_db_rack(conn.as_mut(), &rack_id).await;
+    assert!(rack.config.maintenance_requested.is_none());
+    assert!(rack.config.maintenance_request_id.is_none());
+
+    Ok(())
+}
+
 /// test_maintenance_completed_transitions_to_validation verifies that
 /// Maintenance::Completed transitions to Validation(Pending).
 #[crate::sqlx_test]