Planner/builder split: Move image source decisions to the planner

dev-tools/reconfigurator-cli/src/lib.rs

@@ -23,6 +23,7 @@ use nexus_reconfigurator_planning::system::{SledBuilder, SystemDescription};
 use nexus_reconfigurator_simulation::SimStateBuilder;
 use nexus_reconfigurator_simulation::Simulator;
 use nexus_reconfigurator_simulation::{BlueprintId, SimState};
+use nexus_sled_agent_shared::inventory::ZoneKind;
 use nexus_types::deployment::PlanningInput;
 use nexus_types::deployment::SledFilter;
 use nexus_types::deployment::execution;
@@ -440,9 +441,26 @@ enum BlueprintEditCommands {
     AddNexus {
         /// sled on which to deploy the new instance
         sled_id: SledUuid,
+        /// image source for the new zone
+        ///
+        /// The image source is required if the planning input of the system
+        /// being edited has a TUF repo; otherwise, it will default to the
+        /// install dataset.
+        #[clap(subcommand)]
+        image_source: Option<ImageSourceArgs>,
     },
     /// add a CockroachDB instance to a particular sled
-    AddCockroach { sled_id: SledUuid },
+    AddCockroach {
+        /// sled on which to deploy the new instance
+        sled_id: SledUuid,
+        /// image source for the new zone
+        ///
+        /// The image source is required if the planning input of the system
+        /// being edited has a TUF repo; otherwise, it will default to the
+        /// install dataset.
+        #[clap(subcommand)]
+        image_source: Option<ImageSourceArgs>,
+    },
     /// set the image source for a zone
     SetZoneImage {
         /// id of zone whose image to set
@@ -647,6 +665,24 @@ enum ImageSourceArgs {
     },
 }
 
+fn image_source_unwrap_or(
+    image_source: Option<ImageSourceArgs>,
+    planning_input: &PlanningInput,
+    zone_kind: ZoneKind,
+) -> anyhow::Result<BlueprintZoneImageSource> {
+    if let Some(image_source) = image_source {
+        Ok(image_source.into())
+    } else if planning_input.tuf_repo() == planning_input.old_repo() {
+        planning_input
+            .tuf_repo()
+            .description()
+            .zone_image_source(zone_kind)
+            .context("could not determine image source")
+    } else {
+        bail!("must specify image source for new zone")
+    }
+}
+
 impl From<ImageSourceArgs> for BlueprintZoneImageSource {
     fn from(value: ImageSourceArgs) -> Self {
         match value {
@@ -1198,15 +1234,25 @@ fn cmd_blueprint_edit(
     }
 
     let label = match args.edit_command {
-        BlueprintEditCommands::AddNexus { sled_id } => {
+        BlueprintEditCommands::AddNexus { sled_id, image_source } => {
+            let image_source = image_source_unwrap_or(
+                image_source,
+                &planning_input,
+                ZoneKind::Nexus,
+            )?;
             builder
-                .sled_add_zone_nexus(sled_id)
+                .sled_add_zone_nexus(sled_id, image_source)
                 .context("failed to add Nexus zone")?;
             format!("added Nexus zone to sled {}", sled_id)
         }
-        BlueprintEditCommands::AddCockroach { sled_id } => {
+        BlueprintEditCommands::AddCockroach { sled_id, image_source } => {
+            let image_source = image_source_unwrap_or(
+                image_source,
+                &planning_input,
+                ZoneKind::CockroachDb,
+            )?;
             builder
-                .sled_add_zone_cockroachdb(sled_id)
+                .sled_add_zone_cockroachdb(sled_id, image_source)
                 .context("failed to add CockroachDB zone")?;
             format!("added CockroachDB zone to sled {}", sled_id)
         }

live-tests/tests/test_nexus_add_remove.rs

@@ -18,6 +18,7 @@ use nexus_reconfigurator_planning::blueprint_builder::BlueprintBuilder;
 use nexus_reconfigurator_planning::planner::Planner;
 use nexus_reconfigurator_preparation::PlanningInputFromDb;
 use nexus_sled_agent_shared::inventory::ZoneKind;
+use nexus_types::deployment::BlueprintZoneDisposition;
 use nexus_types::deployment::SledFilter;
 use omicron_common::address::NEXUS_INTERNAL_PORT;
 use omicron_test_utils::dev::poll::CondCheckError;
@@ -54,19 +55,49 @@ async fn test_nexus_add_remove(lc: &LiveTestContext) {
     let nexus = initial_nexus_clients.first().expect("internal Nexus client");
 
     // First, deploy a new Nexus zone to an arbitrary sled.
-    let sled_id = planning_input
+    let commissioned_sled_ids = planning_input
         .all_sled_ids(SledFilter::Commissioned)
-        .next()
-        .expect("any sled id");
+        .collect::<Vec<_>>();
+    let sled_id = *commissioned_sled_ids.first().expect("any sled id");
     let (blueprint1, blueprint2) = blueprint_edit_current_target(
         log,
         &planning_input,
         &collection,
         &nexus,
         &|builder: &mut BlueprintBuilder| {
+            // We have to tell the builder what image source to use for the new
+            // Nexus zone. If we were the planner, we'd check whether we have a
+            // TUF repo (or two) then decide whether to use the image from one
+            // of those or the install dataset. Instead of duplicating all of
+            // that logic, we'll just find an existing Nexus zone and copy its
+            // image source. This should always be right in this context; it
+            // would only be wrong if there are existing Nexus zones with
+            // different image sources, which would only be true in the middle
+            // of an update.
+            let image_source = commissioned_sled_ids
+                .iter()
+                .find_map(|&sled_id| {
+                    builder
+                        .current_sled_zones(
+                            sled_id,
+                            BlueprintZoneDisposition::is_in_service,
+                        )
+                        .find_map(|zone| {
+                            if zone.zone_type.is_nexus() {
+                                Some(zone.image_source.clone())
+                            } else {
+                                None
+                            }
+                        })
+                })
+                .context(
+                    "could not find in-service Nexus in parent blueprint",
+                )?;
+
             builder
-                .sled_add_zone_nexus(sled_id)
+                .sled_add_zone_nexus(sled_id, image_source)
                 .context("adding Nexus zone")?;
+
             Ok(())
         },
     )

nexus/db-queries/src/db/datastore/deployment.rs

@@ -2381,13 +2381,22 @@ mod tests {
 
         // Add zones to our new sled.
         assert_eq!(
-            builder.sled_ensure_zone_ntp(new_sled_id).unwrap(),
+            builder
+                .sled_ensure_zone_ntp(
+                    new_sled_id,
+                    BlueprintZoneImageSource::InstallDataset
+                )
+                .unwrap(),
             Ensure::Added
         );
         for zpool_id in new_sled_zpools.keys() {
             assert_eq!(
                 builder
-                    .sled_ensure_zone_crucible(new_sled_id, *zpool_id)
+                    .sled_ensure_zone_crucible(
+                        new_sled_id,
+                        *zpool_id,
+                        BlueprintZoneImageSource::InstallDataset
+                    )
                     .unwrap(),
                 Ensure::Added
             );

nexus/db-queries/src/db/datastore/vpc.rs

@@ -2883,6 +2883,7 @@ mod tests {
     use nexus_types::deployment::BlueprintTarget;
     use nexus_types::deployment::BlueprintZoneConfig;
     use nexus_types::deployment::BlueprintZoneDisposition;
+    use nexus_types::deployment::BlueprintZoneImageSource;
     use nexus_types::external_api::params;
     use nexus_types::identity::Asset;
     use omicron_common::api::external;
@@ -3238,7 +3239,12 @@ mod tests {
                     .expect("ensured disks");
             }
             builder
-                .sled_add_zone_nexus_with_config(sled_ids[2], false, Vec::new())
+                .sled_add_zone_nexus_with_config(
+                    sled_ids[2],
+                    false,
+                    Vec::new(),
+                    BlueprintZoneImageSource::InstallDataset,
+                )
                 .expect("added nexus to third sled");
             builder.build()
         };
@@ -3307,7 +3313,12 @@ mod tests {
             .expect("created blueprint builder");
             for &sled_id in &sled_ids {
                 builder
-                    .sled_add_zone_nexus_with_config(sled_id, false, Vec::new())
+                    .sled_add_zone_nexus_with_config(
+                        sled_id,
+                        false,
+                        Vec::new(),
+                        BlueprintZoneImageSource::InstallDataset,
+                    )
                     .expect("added nexus to third sled");
             }
             builder.build()

nexus/reconfigurator/execution/src/dns.rs

@@ -1523,7 +1523,12 @@ mod test {
         .unwrap();
         let sled_id =
             blueprint.sleds().next().expect("expected at least one sled");
-        builder.sled_add_zone_nexus(sled_id).unwrap();
+        builder
+            .sled_add_zone_nexus(
+                sled_id,
+                BlueprintZoneImageSource::InstallDataset,
+            )
+            .unwrap();
         let blueprint2 = builder.build();
         eprintln!("blueprint2: {}", blueprint2.display());
         // Figure out the id of the new zone.