Reconfigurator: Add cockroachdb zones as needed (#5797)

This builds on #5788 to add support to the planner to add new CRDB zones
to blueprints (if the current count is below the policy's target count).
No changes were needed on the execution side, and the CRDB zones already
bring themselves up as part of the existing cluster by looking up other
node names in DNS.

A big chunk of the diff comes from expectorate output - our simulated
test system was producing sleds that all had the same physical disk and
zpool UUIDs, which messed with the test I wrote for the builder's zpool
allocation. I tweaked the `TypedUuidRng` that the sled uses to include
the `sled_id` (which itself comes from a "parent" `TypedUuidRng`) as the
second seed argument. If that seems unreasonable, I am very open to
other fixes!

common/src/address.rs

@@ -31,6 +31,12 @@ pub const MIN_PORT: u16 = u16::MIN;
 /// Reconfigurator (to know whether to add new Nexus zones)
 pub const NEXUS_REDUNDANCY: usize = 3;
 
+/// The amount of redundancy for CockroachDb services.
+///
+/// This is used by both RSS (to distribute the initial set of services) and the
+/// Reconfigurator (to know whether to add new crdb zones)
+pub const COCKROACHDB_REDUNDANCY: usize = 5;
+
 /// The amount of redundancy for internal DNS servers.
 ///
 /// Must be less than or equal to MAX_DNS_REDUNDANCY.

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

@@ -24,6 +24,7 @@ use nexus_types::deployment::BlueprintZoneFilter;
 use nexus_types::deployment::OmicronZoneNic;
 use nexus_types::deployment::PlanningInput;
 use nexus_types::deployment::SledFilter;
+use nexus_types::deployment::ZoneKind;
 use nexus_types::deployment::{Blueprint, UnstableReconfiguratorState};
 use nexus_types::internal_api::params::DnsConfigParams;
 use nexus_types::inventory::Collection;
@@ -744,7 +745,8 @@ fn cmd_blueprint_edit(
 
     let label = match args.edit_command {
         BlueprintEditCommands::AddNexus { sled_id } => {
-            let current = builder.sled_num_nexus_zones(sled_id);
+            let current =
+                builder.sled_num_zones_of_kind(sled_id, ZoneKind::Nexus);
             let added = builder
                 .sled_ensure_zone_multiple_nexus(sled_id, current + 1)
                 .context("failed to add Nexus zone")?;

nexus/reconfigurator/execution/src/dns.rs

@@ -493,6 +493,7 @@ mod test {
     use omicron_common::address::get_switch_zone_address;
     use omicron_common::address::IpRange;
     use omicron_common::address::Ipv6Subnet;
+    use omicron_common::address::COCKROACHDB_REDUNDANCY;
     use omicron_common::address::NEXUS_REDUNDANCY;
     use omicron_common::address::RACK_PREFIX;
     use omicron_common::address::SLED_PREFIX;
@@ -501,6 +502,7 @@ mod test {
     use omicron_test_utils::dev::test_setup_log;
     use omicron_uuid_kinds::ExternalIpUuid;
     use omicron_uuid_kinds::OmicronZoneUuid;
+    use sled_agent_client::ZoneKind;
     use std::collections::BTreeMap;
     use std::collections::BTreeSet;
     use std::collections::HashMap;
@@ -1235,6 +1237,7 @@ mod test {
                 external_ip_rows: &[],
                 service_nic_rows: &[],
                 target_nexus_zone_count: NEXUS_REDUNDANCY,
+                target_cockroachdb_zone_count: COCKROACHDB_REDUNDANCY,
                 target_cockroachdb_cluster_version:
                     CockroachDbClusterVersion::POLICY,
                 log,
@@ -1260,7 +1263,7 @@ mod test {
         .unwrap();
         let sled_id =
             blueprint.sleds().next().expect("expected at least one sled");
-        let nalready = builder.sled_num_nexus_zones(sled_id);
+        let nalready = builder.sled_num_zones_of_kind(sled_id, ZoneKind::Nexus);
         let rv = builder
             .sled_ensure_zone_multiple_nexus(sled_id, nalready + 1)
             .unwrap();

nexus/reconfigurator/planning/src/blueprint_builder/builder.rs

@@ -6,6 +6,7 @@
 
 use crate::ip_allocator::IpAllocator;
 use crate::planner::zone_needs_expungement;
+use crate::planner::DiscretionaryOmicronZone;
 use crate::planner::ZoneExpungeReason;
 use anyhow::anyhow;
 use internal_dns::config::Host;
@@ -29,6 +30,7 @@ use nexus_types::deployment::PlanningInput;
 use nexus_types::deployment::SledDetails;
 use nexus_types::deployment::SledFilter;
 use nexus_types::deployment::SledResources;
+use nexus_types::deployment::ZpoolFilter;
 use nexus_types::deployment::ZpoolName;
 use nexus_types::external_api::views::SledState;
 use omicron_common::address::get_internal_dns_server_addresses;
@@ -50,6 +52,7 @@ use omicron_uuid_kinds::SledUuid;
 use omicron_uuid_kinds::ZpoolUuid;
 use rand::rngs::StdRng;
 use rand::SeedableRng;
+use sled_agent_client::ZoneKind;
 use slog::debug;
 use slog::error;
 use slog::info;
@@ -78,6 +81,10 @@ use super::zones::BuilderZonesConfig;
 pub enum Error {
     #[error("sled {sled_id}: ran out of available addresses for sled")]
     OutOfAddresses { sled_id: SledUuid },
+    #[error(
+        "sled {sled_id}: no available zpools for additional {kind:?} zones"
+    )]
+    NoAvailableZpool { sled_id: SledUuid, kind: ZoneKind },
     #[error("no Nexus zones exist in parent blueprint")]
     NoNexusZonesInParentBlueprint,
     #[error("no external service IP addresses are available")]
@@ -722,15 +729,19 @@ impl<'a> BlueprintBuilder<'a> {
         Ok(Ensure::Added)
     }
 
-    /// Return the number of Nexus zones that would be configured to run on the
-    /// given sled if this builder generated a blueprint
+    /// Return the number of zones of a given kind that would be configured to
+    /// run on the given sled if this builder generated a blueprint.
     ///
     /// This value may change before a blueprint is actually generated if
     /// further changes are made to the builder.
-    pub fn sled_num_nexus_zones(&self, sled_id: SledUuid) -> usize {
+    pub fn sled_num_zones_of_kind(
+        &self,
+        sled_id: SledUuid,
+        kind: ZoneKind,
+    ) -> usize {
         self.zones
             .current_sled_zones(sled_id)
-            .filter(|(z, _)| z.zone_type.is_nexus())
+            .filter(|(z, _)| z.zone_type.kind() == kind)
             .count()
     }
 
@@ -777,7 +788,7 @@ impl<'a> BlueprintBuilder<'a> {
         external_dns_servers: Vec<IpAddr>,
     ) -> Result<EnsureMultiple, Error> {
         // How many Nexus zones do we need to add?
-        let nexus_count = self.sled_num_nexus_zones(sled_id);
+        let nexus_count = self.sled_num_zones_of_kind(sled_id, ZoneKind::Nexus);
         let num_nexus_to_add = match desired_zone_count.checked_sub(nexus_count)
         {
             Some(0) => return Ok(EnsureMultiple::NotNeeded),
@@ -850,6 +861,51 @@ impl<'a> BlueprintBuilder<'a> {
         self.cockroachdb_setting_preserve_downgrade = version;
     }
 
+    pub fn sled_ensure_zone_multiple_cockroachdb(
+        &mut self,
+        sled_id: SledUuid,
+        desired_zone_count: usize,
+    ) -> Result<EnsureMultiple, Error> {
+        // How many CRDB zones do we need to add?
+        let crdb_count =
+            self.sled_num_zones_of_kind(sled_id, ZoneKind::CockroachDb);
+        let num_crdb_to_add = match desired_zone_count.checked_sub(crdb_count) {
+            Some(0) => return Ok(EnsureMultiple::NotNeeded),
+            Some(n) => n,
+            None => {
+                return Err(Error::Planner(anyhow!(
+                    "removing a CockroachDb zone not yet supported \
+                     (sled {sled_id} has {crdb_count}; \
+                     planner wants {desired_zone_count})"
+                )));
+            }
+        };
+        for _ in 0..num_crdb_to_add {
+            let zone_id = self.rng.zone_rng.next();
+            let underlay_ip = self.sled_alloc_ip(sled_id)?;
+            let pool_name = self.sled_alloc_zpool(
+                sled_id,
+                DiscretionaryOmicronZone::CockroachDb,
+            )?;
+            let port = omicron_common::address::COCKROACH_PORT;
+            let address = SocketAddrV6::new(underlay_ip, port, 0, 0);
+            let zone = BlueprintZoneConfig {
+                disposition: BlueprintZoneDisposition::InService,
+                id: zone_id,
+                underlay_address: underlay_ip,
+                zone_type: BlueprintZoneType::CockroachDb(
+                    blueprint_zone_type::CockroachDb {
+                        address,
+                        dataset: OmicronZoneDataset { pool_name },
+                    },
+                ),
+            };
+            self.sled_add_zone(sled_id, zone)?;
+        }
+
+        Ok(EnsureMultiple::Changed { added: num_crdb_to_add, removed: 0 })
+    }
+
     fn sled_add_zone(
         &mut self,
         sled_id: SledUuid,
@@ -906,6 +962,46 @@ impl<'a> BlueprintBuilder<'a> {
         allocator.alloc().ok_or(Error::OutOfAddresses { sled_id })
     }
 
+    fn sled_alloc_zpool(
+        &mut self,
+        sled_id: SledUuid,
+        kind: DiscretionaryOmicronZone,
+    ) -> Result<ZpoolName, Error> {
+        let resources = self.sled_resources(sled_id)?;
+
+        // We refuse to choose a zpool for a zone of a given `kind` if this
+        // sled already has a zone of that kind on the same zpool. Build up a
+        // set of invalid zpools for this sled/kind pair.
+        let mut skip_zpools = BTreeSet::new();
+        for zone_config in self.current_sled_zones(sled_id) {
+            match (kind, &zone_config.zone_type) {
+                (
+                    DiscretionaryOmicronZone::Nexus,
+                    BlueprintZoneType::Nexus(_),
+                ) => {
+                    // TODO handle this case once we track transient datasets
+                }
+                (
+                    DiscretionaryOmicronZone::CockroachDb,
+                    BlueprintZoneType::CockroachDb(crdb),
+                ) => {
+                    skip_zpools.insert(&crdb.dataset.pool_name);
+                }
+                (DiscretionaryOmicronZone::Nexus, _)
+                | (DiscretionaryOmicronZone::CockroachDb, _) => (),
+            }
+        }
+
+        for &zpool_id in resources.all_zpools(ZpoolFilter::InService) {
+            let zpool_name = ZpoolName::new_external(zpool_id);
+            if !skip_zpools.contains(&zpool_name) {
+                return Ok(zpool_name);
+            }
+        }
+
+        Err(Error::NoAvailableZpool { sled_id, kind: kind.into() })
+    }
+
     fn sled_resources(
         &self,
         sled_id: SledUuid,
@@ -1179,6 +1275,7 @@ pub mod test {
 
     /// Checks various conditions that should be true for all blueprints
     pub fn verify_blueprint(blueprint: &Blueprint) {
+        // There should be no duplicate underlay IPs.
         let mut underlay_ips: BTreeMap<Ipv6Addr, &BlueprintZoneConfig> =
             BTreeMap::new();
         for (_, zone) in blueprint.all_omicron_zones(BlueprintZoneFilter::All) {
@@ -1196,6 +1293,33 @@ pub mod test {
                 );
             }
         }
+
+        // On any given zpool, we should have at most one zone of any given
+        // kind.
+        let mut kinds_by_zpool: BTreeMap<
+            ZpoolUuid,
+            BTreeMap<ZoneKind, OmicronZoneUuid>,
+        > = BTreeMap::new();
+        for (_, zone) in blueprint.all_omicron_zones(BlueprintZoneFilter::All) {
+            if let Some(dataset) = zone.zone_type.dataset() {
+                let kind = zone.zone_type.kind();
+                if let Some(previous) = kinds_by_zpool
+                    .entry(dataset.pool_name.id())
+                    .or_default()
+                    .insert(kind, zone.id)
+                {
+                    panic!(
+                        "zpool {} has two zones of kind {kind:?}: {} and {}\
+                            \n\n\
+                            blueprint: {}",
+                        dataset.pool_name,
+                        zone.id,
+                        previous,
+                        blueprint.display(),
+                    );
+                }
+            }
+        }
     }
 
     #[test]
@@ -1881,4 +2005,96 @@ pub mod test {
 
         logctx.cleanup_successful();
     }
+
+    #[test]
+    fn test_ensure_cockroachdb() {
+        static TEST_NAME: &str = "blueprint_builder_test_ensure_cockroachdb";
+        let logctx = test_setup_log(TEST_NAME);
+
+        // Discard the example blueprint and start with an empty one.
+        let (_, input, _) = example(&logctx.log, TEST_NAME, DEFAULT_N_SLEDS);
+        let input = {
+            // Clear out the external networking records from `input`, since
+            // we're building an empty blueprint.
+            let mut builder = input.into_builder();
+            *builder.network_resources_mut() =
+                OmicronZoneNetworkResources::new();
+            builder.build()
+        };
+        let parent = BlueprintBuilder::build_empty_with_sleds_seeded(
+            input.all_sled_ids(SledFilter::Commissioned),
+            "test",
+            TEST_NAME,
+        );
+
+        // Pick an arbitrary sled.
+        let (target_sled_id, sled_resources) = input
+            .all_sled_resources(SledFilter::InService)
+            .next()
+            .expect("at least one sled");
+
+        // It should have multiple zpools.
+        let num_sled_zpools = sled_resources.zpools.len();
+        assert!(
+            num_sled_zpools > 1,
+            "expected more than 1 zpool, got {num_sled_zpools}"
+        );
+
+        // We should be able to ask for a CRDB zone per zpool.
+        let mut builder = BlueprintBuilder::new_based_on(
+            &logctx.log,
+            &parent,
+            &input,
+            "test",
+        )
+        .expect("constructed builder");
+        let ensure_result = builder
+            .sled_ensure_zone_multiple_cockroachdb(
+                target_sled_id,
+                num_sled_zpools,
+            )
+            .expect("ensured multiple CRDB zones");
+        assert_eq!(
+            ensure_result,
+            EnsureMultiple::Changed { added: num_sled_zpools, removed: 0 }
+        );
+
+        let blueprint = builder.build();
+        verify_blueprint(&blueprint);
+        assert_eq!(
+            blueprint
+                .all_omicron_zones(BlueprintZoneFilter::ShouldBeRunning)
+                .filter(|(sled_id, z)| {
+                    *sled_id == target_sled_id
+                        && z.zone_type.kind() == ZoneKind::CockroachDb
+                })
+                .count(),
+            num_sled_zpools
+        );
+
+        // If we instead ask for one more zone than there are zpools, we should
+        // get a zpool allocation error.
+        let mut builder = BlueprintBuilder::new_based_on(
+            &logctx.log,
+            &parent,
+            &input,
+            "test",
+        )
+        .expect("constructed builder");
+        let ensure_error = builder
+            .sled_ensure_zone_multiple_cockroachdb(
+                target_sled_id,
+                num_sled_zpools + 1,
+            )
+            .expect_err("failed to create too many CRDB zones");
+        match ensure_error {
+            Error::NoAvailableZpool { sled_id, kind } => {
+                assert_eq!(target_sled_id, sled_id);
+                assert_eq!(kind, ZoneKind::CockroachDb);
+            }
+            _ => panic!("unexpected error {ensure_error}"),
+        }
+
+        logctx.cleanup_successful();
+    }
 }