fix(nvlink): refresh NVL domain UUIDs for machines

crates/nvlink-manager/src/lib.rs

@@ -145,6 +145,7 @@ fn domain_uuid_from_nmx_c_hello(
         .as_ref()
         .and_then(|header| uuid::Uuid::parse_str(&header.domain_uuid).ok())
         .map(NvLinkDomainId::from)
+        .filter(|domain_uuid| *domain_uuid != NvLinkDomainId::nil())
         .ok_or_else(|| {
             NvLinkManagerError::internal(format!(
                 "Failed to parse domain UUID from NMX-C hello response: {hello:?}"
@@ -183,9 +184,8 @@ fn build_machine_nvlink_info_from_nmx_c_hello(
 ) -> MachineNvLinkInfo {
     if let Some(existing) = existing {
         let mut info = existing.clone();
-        if info.domain_uuid == NvLinkDomainId::nil() {
-            info.domain_uuid = domain_uuid;
-        }
+        info.domain_uuid = domain_uuid;
+
         if info.chassis_serial.trim().is_empty() {
             info.chassis_serial = chassis_serial.to_string();
         }
@@ -218,11 +218,11 @@ fn build_machine_nvlink_info_from_nmx_c_hello(
     }
 }
 
-/// Populates missing `machines.status.nvlink_info` entries (or nil `domain_uuid`) using NMX-C hello.
+/// Updates missing or stale machine NVLink domains using a validated NMX-C Hello.
 fn populate_machine_nvlink_info_if_needed(
     machine_nvlink_info: &mut HashMap<MachineId, Option<MachineNvLinkInfo>>,
     managed_host_snapshots: &HashMap<MachineId, ManagedHostStateSnapshot>,
-    chassis_serial: &str,
+    snapshot_chassis_serial: Option<&str>,
     machine_ids: &[MachineId],
     domain_uuid: NvLinkDomainId,
 ) -> Vec<(MachineId, MachineNvLinkInfo)> {
@@ -233,14 +233,22 @@ fn populate_machine_nvlink_info_if_needed(
             .and_then(|info| info.as_ref());
         let needs_update = match existing {
             None => true,
-            Some(info) => {
-                info.domain_uuid == NvLinkDomainId::nil() || info.chassis_serial.trim().is_empty()
-            }
+            Some(info) => info.domain_uuid != domain_uuid || info.chassis_serial.trim().is_empty(),
         };
         if !needs_update {
             continue;
         }
 
+        // Fall back to persisted chassis inventory so a missing snapshot serial does
+        // not block refreshing the machine's NVLink domain UUID.
+        let Some(chassis_serial) = snapshot_chassis_serial
+            .or_else(|| existing.map(|info| info.chassis_serial.as_str()))
+            .map(str::trim)
+            .filter(|serial| !serial.is_empty())
+        else {
+            continue;
+        };
+
         let snapshot = managed_host_snapshots.get(machine_id);
         let updated = build_machine_nvlink_info_from_nmx_c_hello(
             existing,
@@ -1548,21 +1556,26 @@ impl NvlPartitionMonitor {
             managed_host_snapshots_domain.keys().copied().collect();
         let mut nvlink_info_db_updates = Vec::new();
         for snapshot in snapshots {
-            // The chassis serial number should always be available from hardware_info.
-            let Some(chassis_serial) = managed_host_chassis_serial(snapshot) else {
+            let machine_id = snapshot.host_snapshot.id;
+            let snapshot_chassis_serial = managed_host_chassis_serial(snapshot);
+            let has_persisted_chassis_serial = machine_nvlink_info
+                .get(&machine_id)
+                .and_then(|info| info.as_ref())
+                .is_some_and(|info| !info.chassis_serial.trim().is_empty());
+            if snapshot_chassis_serial.is_none() && !has_persisted_chassis_serial {
                 tracing::warn!(
                     group_id,
                     group_type = group_type_label,
-                    machine_id = %snapshot.host_snapshot.id,
+                    %machine_id,
                     "Skipping nvlink_info population; chassis serial unavailable"
                 );
                 continue;
-            };
+            }
             nvlink_info_db_updates.extend(populate_machine_nvlink_info_if_needed(
                 machine_nvlink_info,
                 all_managed_host_snapshots,
-                &chassis_serial,
-                &[snapshot.host_snapshot.id],
+                snapshot_chassis_serial.as_deref(),
+                &[machine_id],
                 domain_uuid,
             ));
         }
@@ -2950,6 +2963,119 @@ fn record_domain_health(
     }
 }
 
+#[cfg(test)]
+mod domain_uuid_observation_tests {
+    use carbide_test_support::{Check, check_values, value_scenarios};
+    use carbide_uuid::machine::{MachineIdSource, MachineType};
+    use libnmxc::nmxc_model::{ServerHeader, ServerHello};
+
+    use super::*;
+
+    #[test]
+    fn hello_domain_uuid_accepts_only_valid_non_nil_values() {
+        value_scenarios!(
+            run = |domain_uuid| {
+                domain_uuid_from_nmx_c_hello(&ServerHello {
+                    server_header: Some(ServerHeader {
+                        domain_uuid: domain_uuid.to_string(),
+                        ..Default::default()
+                    }),
+                    ..Default::default()
+                })
+                .is_ok()
+            };
+
+            "accepted" {
+                "11111111-1111-1111-1111-111111111111" => true,
+            }
+
+            "rejected" {
+                "not-a-uuid" => false,
+                "00000000-0000-0000-0000-000000000000" => false,
+            }
+
+        );
+    }
+
+    #[test]
+    fn valid_domain_observations_update_only_changed_machine_domains()
+    -> Result<(), Box<dyn std::error::Error>> {
+        let machine_id = MachineId::new(MachineIdSource::Tpm, [1; 32], MachineType::Host);
+        let old_domain: NvLinkDomainId = "11111111-1111-1111-1111-111111111111".parse()?;
+        let observed_domain: NvLinkDomainId = "22222222-2222-2222-2222-222222222222".parse()?;
+
+        let gpu = NvLinkGpu {
+            tray_index: 1,
+            slot_id: 2,
+            device_id: 3,
+            guid: 4,
+        };
+
+        let old_info = MachineNvLinkInfo {
+            domain_uuid: old_domain,
+            chassis_serial: "CHASSIS-A".to_string(),
+            gpus: vec![gpu],
+        };
+
+        let observed_info = MachineNvLinkInfo {
+            domain_uuid: observed_domain,
+            ..old_info.clone()
+        };
+
+        let initial_info = MachineNvLinkInfo {
+            domain_uuid: observed_domain,
+            chassis_serial: "CHASSIS-A".to_string(),
+            gpus: Vec::new(),
+        };
+
+        check_values(
+            [
+                Check {
+                    scenario: "initial population",
+                    input: (None, Some("CHASSIS-A")),
+                    expect: (1, initial_info),
+                },
+                Check {
+                    scenario: "changed UUID",
+                    input: (Some(old_info.clone()), Some("CHASSIS-A")),
+                    expect: (1, observed_info.clone()),
+                },
+                Check {
+                    scenario: "changed UUID with persisted chassis serial only",
+                    input: (Some(old_info), None),
+                    expect: (1, observed_info.clone()),
+                },
+                Check {
+                    scenario: "identical persisted UUID after restart",
+                    input: (Some(observed_info.clone()), Some("CHASSIS-A")),
+                    expect: (0, observed_info),
+                },
+            ],
+            |(existing, snapshot_chassis_serial)| {
+                let mut machine_nvlink_info = HashMap::from([(machine_id, existing)]);
+
+                let updates = populate_machine_nvlink_info_if_needed(
+                    &mut machine_nvlink_info,
+                    &HashMap::new(),
+                    snapshot_chassis_serial,
+                    &[machine_id],
+                    observed_domain,
+                );
+
+                (
+                    updates.len(),
+                    machine_nvlink_info
+                        .remove(&machine_id)
+                        .flatten()
+                        .expect("observation populates machine NVLink info"),
+                )
+            },
+        );
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod partition_classification_tests {
     use libnmxc::nmxc_model::PartitionInfo;