internal API for fetching MgsUpdateDriver status does not work (#8475)

davepacheco · web-flow · commit 31c8803f39bd · 2025-06-30T15:18:42.000Z
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/nexus/mgs-updates/Cargo.toml b/nexus/mgs-updates/Cargo.toml
@@ -11,6 +11,7 @@ chrono.workspace = true
 futures.workspace = true
 gateway-client.workspace = true
 gateway-types.workspace = true
+iddqd.workspace = true
 id-map.workspace = true
 internal-dns-resolver.workspace = true
 internal-dns-types.workspace = true
diff --git a/nexus/mgs-updates/src/driver.rs b/nexus/mgs-updates/src/driver.rs
@@ -18,6 +18,7 @@ use futures::future::BoxFuture;
 use futures::stream::FuturesUnordered;
 use id_map::IdMap;
 use id_map::IdMappable;
+use iddqd::IdOrdMap;
 use nexus_types::deployment::PendingMgsUpdate;
 use nexus_types::deployment::PendingMgsUpdates;
 use nexus_types::internal_api::views::CompletedAttempt;
@@ -31,7 +32,6 @@ use omicron_uuid_kinds::SpUpdateUuid;
 use qorb::resolver::AllBackends;
 use slog::{error, info, o, warn};
 use slog_error_chain::InlineErrorChain;
-use std::collections::BTreeMap;
 use std::collections::VecDeque;
 use std::sync::Arc;
 use std::time::Duration;
@@ -117,8 +117,8 @@ impl MgsUpdateDriver {
     ) -> MgsUpdateDriver {
         let (status_tx, _) = watch::channel(MgsUpdateDriverStatus {
             recent: VecDeque::with_capacity(N_RECENT_COMPLETIONS),
-            in_progress: BTreeMap::new(),
-            waiting: BTreeMap::new(),
+            in_progress: IdOrdMap::new(),
+            waiting: IdOrdMap::new(),
         });
 
         MgsUpdateDriver {
@@ -346,14 +346,15 @@ impl MgsUpdateDriver {
         // Update status.  We do this before starting the future because it will
         // update this status and expects to find it.
         self.status_tx.send_modify(|driver_status| {
-            driver_status.in_progress.insert(
-                baseboard_id.clone(),
-                InProgressUpdateStatus {
+            driver_status
+                .in_progress
+                .insert_unique(InProgressUpdateStatus {
+                    baseboard_id: baseboard_id.clone(),
                     time_started: in_progress.time_started,
                     status: UpdateAttemptStatus::NotStarted,
                     nattempts_done,
-                },
-            );
+                })
+                .expect("no previous update for this baseboard");
         });
 
         let status_updater = UpdateAttemptStatusUpdater::new(
@@ -444,13 +445,14 @@ impl MgsUpdateDriver {
             assert!(found.is_some());
 
             // Add this item to the list of requests waiting to be retried.
-            driver_status.waiting.insert(
-                baseboard_id.clone(),
-                WaitingStatus {
+            driver_status
+                .waiting
+                .insert_unique(WaitingStatus {
+                    baseboard_id: baseboard_id.clone(),
                     next_attempt_time: status_time_next,
                     nattempts_done,
-                },
-            );
+                })
+                .expect("no previous waiting update for this baseboard");
 
             // Report this recently-completed attempt.
             // This is a ringbuffer of recent attempts.  Make space if we're
@@ -577,7 +579,7 @@ impl UpdateAttemptStatusUpdater {
             // the future that owns it.  The status entry for this future lives
             // in the `in_progress` struct until it completes.  Thus, we should
             // always have a value here.
-            let my_status =
+            let mut my_status =
                 driver_status.in_progress.get_mut(&self.baseboard_id).unwrap();
             my_status.status = new_status;
         });
diff --git a/nexus/mgs-updates/src/test_util/updates.rs b/nexus/mgs-updates/src/test_util/updates.rs
@@ -190,14 +190,15 @@ impl UpdateDescription<'_> {
         let (status_tx, status_rx) =
             watch::channel(MgsUpdateDriverStatus::default());
         status_tx.send_modify(|status| {
-            status.in_progress.insert(
-                baseboard_id.clone(),
-                InProgressUpdateStatus {
+            status
+                .in_progress
+                .insert_unique(InProgressUpdateStatus {
+                    baseboard_id: baseboard_id.clone(),
                     time_started: chrono::Utc::now(),
                     status: UpdateAttemptStatus::NotStarted,
                     nattempts_done: 0,
-                },
-            );
+                })
+                .expect("no value present in object we just created");
         });
         let status_updater =
             UpdateAttemptStatusUpdater::new(status_tx.clone(), baseboard_id);
@@ -316,8 +317,7 @@ impl InProgressAttempt {
             // future itself to finish and then return indicating that it's
             // done.
             let overall_status = self.status_rx.borrow();
-            let maybe_current_status =
-                overall_status.in_progress.values().next();
+            let maybe_current_status = overall_status.in_progress.iter().next();
             if let Some(current_status) = maybe_current_status {
                 if current_status.status == status {
                     debug!(
diff --git a/nexus/types/src/internal_api/views.rs b/nexus/types/src/internal_api/views.rs
@@ -9,6 +9,9 @@ use chrono::SecondsFormat;
 use chrono::Utc;
 use futures::future::ready;
 use futures::stream::StreamExt;
+use iddqd::IdOrdItem;
+use iddqd::IdOrdMap;
+use iddqd::id_upcast;
 use nexus_sled_agent_shared::inventory::ConfigReconcilerInventory;
 use nexus_sled_agent_shared::inventory::ConfigReconcilerInventoryResult;
 use nexus_sled_agent_shared::inventory::OmicronZoneImageSource;
@@ -343,11 +346,13 @@ pub struct Ipv4NatEntryView {
 
 /// Status of ongoing update attempts, recently completed attempts, and update
 /// requests that are waiting for retry.
-#[derive(Clone, Debug, Default, Deserialize, Serialize, JsonSchema)]
+#[derive(
+    Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize, JsonSchema,
+)]
 pub struct MgsUpdateDriverStatus {
     pub recent: VecDeque<CompletedAttempt>,
-    pub in_progress: BTreeMap<Arc<BaseboardId>, InProgressUpdateStatus>,
-    pub waiting: BTreeMap<Arc<BaseboardId>, WaitingStatus>,
+    pub in_progress: IdOrdMap<InProgressUpdateStatus>,
+    pub waiting: IdOrdMap<WaitingStatus>,
 }
 
 impl MgsUpdateDriverStatus {
@@ -384,7 +389,7 @@ impl Display for MgsUpdateDriverStatusDisplay<'_> {
         }
 
         writeln!(f, "\ncurrently in progress:")?;
-        for (baseboard_id, status) in &status.in_progress {
+        for status in &status.in_progress {
             // Ignore units smaller than a millisecond.
             let elapsed = Duration::from_millis(
                 u64::try_from(
@@ -398,19 +403,19 @@ impl Display for MgsUpdateDriverStatusDisplay<'_> {
                 status
                     .time_started
                     .to_rfc3339_opts(SecondsFormat::Millis, true),
-                baseboard_id.serial_number,
+                status.baseboard_id.serial_number,
                 status.status,
                 status.nattempts_done + 1,
                 humantime::format_duration(elapsed),
             )?;
         }
 
         writeln!(f, "\nwaiting for retry:")?;
-        for (baseboard_id, wait_info) in &status.waiting {
+        for wait_info in &status.waiting {
             writeln!(
                 f,
                 "    serial {}: will try again at {} (attempt {})",
-                baseboard_id.serial_number,
+                wait_info.baseboard_id.serial_number,
                 wait_info.next_attempt_time,
                 wait_info.nattempts_done + 1,
             )?;
@@ -421,13 +426,13 @@ impl Display for MgsUpdateDriverStatusDisplay<'_> {
 }
 
 /// externally-exposed status for a completed attempt
-#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema)]
+#[derive(Debug, Clone, Deserialize, Eq, PartialEq, Serialize, JsonSchema)]
 pub struct CompletedAttempt {
     pub time_started: DateTime<Utc>,
     pub time_done: DateTime<Utc>,
     pub elapsed: Duration,
     pub request: PendingMgsUpdate,
-    #[serde(serialize_with = "snake_case_result::serialize")]
+    #[serde(with = "snake_case_result")]
     #[schemars(
         schema_with = "SnakeCaseResult::<UpdateCompletedHow, String>::json_schema"
     )]
@@ -447,13 +452,24 @@ pub enum UpdateCompletedHow {
 }
 
 /// externally-exposed status for each in-progress update
-#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize, JsonSchema)]
 pub struct InProgressUpdateStatus {
+    pub baseboard_id: Arc<BaseboardId>,
     pub time_started: DateTime<Utc>,
     pub status: UpdateAttemptStatus,
     pub nattempts_done: u32,
 }
 
+impl IdOrdItem for InProgressUpdateStatus {
+    type Key<'a> = &'a Arc<BaseboardId>;
+
+    fn key(&self) -> Self::Key<'_> {
+        &self.baseboard_id
+    }
+
+    id_upcast!();
+}
+
 /// status of a single update attempt
 #[derive(
     Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize, JsonSchema,
@@ -471,12 +487,23 @@ pub enum UpdateAttemptStatus {
 }
 
 /// externally-exposed status for waiting updates
-#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize, JsonSchema)]
 pub struct WaitingStatus {
+    pub baseboard_id: Arc<BaseboardId>,
     pub next_attempt_time: DateTime<Utc>,
     pub nattempts_done: u32,
 }
 
+impl IdOrdItem for WaitingStatus {
+    type Key<'a> = &'a Arc<BaseboardId>;
+
+    fn key(&self) -> Self::Key<'_> {
+        &self.baseboard_id
+    }
+
+    id_upcast!();
+}
+
 #[derive(
     Debug,
     Clone,
@@ -594,3 +621,78 @@ impl UpdateStatus {
         ZoneStatusVersion::Unknown
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::CompletedAttempt;
+    use super::InProgressUpdateStatus;
+    use super::MgsUpdateDriverStatus;
+    use super::UpdateCompletedHow;
+    use super::WaitingStatus;
+    use crate::deployment::ExpectedVersion;
+    use crate::deployment::PendingMgsUpdate;
+    use crate::deployment::PendingMgsUpdateDetails;
+    use crate::internal_api::views::UpdateAttemptStatus;
+    use crate::inventory::BaseboardId;
+    use chrono::Utc;
+    use gateway_client::types::SpType;
+    use std::collections::VecDeque;
+    use std::sync::Arc;
+    use std::time::Instant;
+    use tufaceous_artifact::ArtifactHash;
+
+    #[test]
+    fn test_can_serialize_mgs_updates() {
+        let start = Instant::now();
+        let artifact_hash: ArtifactHash =
+            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+                .parse()
+                .unwrap();
+        let baseboard_id = Arc::new(BaseboardId {
+            part_number: String::from("a_port"),
+            serial_number: String::from("a_serial"),
+        });
+        let waiting = WaitingStatus {
+            baseboard_id: baseboard_id.clone(),
+            next_attempt_time: Utc::now(),
+            nattempts_done: 2,
+        };
+        let in_progress = InProgressUpdateStatus {
+            baseboard_id: baseboard_id.clone(),
+            time_started: Utc::now(),
+            status: UpdateAttemptStatus::Updating,
+            nattempts_done: 3,
+        };
+        let completed = CompletedAttempt {
+            time_started: Utc::now(),
+            time_done: Utc::now(),
+            elapsed: start.elapsed(),
+            request: PendingMgsUpdate {
+                baseboard_id: baseboard_id.clone(),
+                sp_type: SpType::Sled,
+                slot_id: 12,
+                details: PendingMgsUpdateDetails::Sp {
+                    expected_active_version: "1.0.0".parse().unwrap(),
+                    expected_inactive_version: ExpectedVersion::NoValidVersion,
+                },
+                artifact_hash,
+                artifact_version: "2.0.0".parse().unwrap(),
+            },
+            result: Ok(UpdateCompletedHow::FoundNoChangesNeeded),
+            nattempts_done: 8,
+        };
+
+        let status = MgsUpdateDriverStatus {
+            recent: VecDeque::from([completed]),
+            in_progress: std::iter::once(in_progress).collect(),
+            waiting: std::iter::once(waiting).collect(),
+        };
+
+        let serialized =
+            serde_json::to_string(&status).expect("failed to serialize value");
+        let deserialized: MgsUpdateDriverStatus =
+            serde_json::from_str(&serialized)
+                .expect("failed to deserialize value");
+        assert_eq!(deserialized, status);
+    }
+}
diff --git a/nexus/types/src/inventory.rs b/nexus/types/src/inventory.rs
@@ -231,6 +231,7 @@ impl Collection {
     Diffable,
     Ord,
     Eq,
+    Hash,
     PartialOrd,
     PartialEq,
     Deserialize,
diff --git a/openapi/nexus-internal.json b/openapi/nexus-internal.json
