(5/5) [nexus] Implement Affinity/Anti-Affinity Groups in external API

common/src/api/external/http_pagination.rs

@@ -312,6 +312,19 @@ pub type PaginatedByNameOrId<Selector = ()> = PaginationParams<
 pub type PageSelectorByNameOrId<Selector = ()> =
     PageSelector<ScanByNameOrId<Selector>, NameOrId>;
 
+pub fn id_pagination<'a, Selector>(
+    pag_params: &'a DataPageParams<Uuid>,
+    scan_params: &'a ScanById<Selector>,
+) -> Result<PaginatedBy<'a>, HttpError>
+where
+    Selector:
+        Clone + Debug + DeserializeOwned + JsonSchema + PartialEq + Serialize,
+{
+    match scan_params.sort_by {
+        IdSortMode::IdAscending => Ok(PaginatedBy::Id(pag_params.clone())),
+    }
+}
+
 pub fn name_or_id_pagination<'a, Selector>(
     pag_params: &'a DataPageParams<NameOrId>,
     scan_params: &'a ScanByNameOrId<Selector>,

common/src/api/external/mod.rs

@@ -982,6 +982,10 @@ impl JsonSchema for Hostname {
 pub enum ResourceType {
     AddressLot,
     AddressLotBlock,
+    AffinityGroup,
+    AffinityGroupMember,
+    AntiAffinityGroup,
+    AntiAffinityGroupMember,
     AllowList,
     BackgroundTask,
     BgpConfig,
@@ -1312,6 +1316,56 @@ pub enum InstanceAutoRestartPolicy {
     BestEffort,
 }
 
+// AFFINITY GROUPS
+
+#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum AffinityPolicy {
+    /// If the affinity request cannot be satisfied, allow it anyway.
+    ///
+    /// This enables a "best-effort" attempt to satisfy the affinity policy.
+    Allow,
+
+    /// If the affinity request cannot be satisfied, fail explicitly.
+    Fail,
+}
+
+/// Describes the scope of affinity for the purposes of co-location.
+#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum FailureDomain {
+    /// Instances are considered co-located if they are on the same sled
+    Sled,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
+#[serde(tag = "type", content = "value", rename_all = "snake_case")]
+pub enum AffinityGroupMember {
+    Instance(Uuid),
+}
+
+impl SimpleIdentity for AffinityGroupMember {
+    fn id(&self) -> Uuid {
+        match self {
+            AffinityGroupMember::Instance(id) => *id,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
+#[serde(tag = "type", content = "value", rename_all = "snake_case")]
+pub enum AntiAffinityGroupMember {
+    Instance(Uuid),
+}
+
+impl SimpleIdentity for AntiAffinityGroupMember {
+    fn id(&self) -> Uuid {
+        match self {
+            AntiAffinityGroupMember::Instance(id) => *id,
+        }
+    }
+}
+
 // DISKS
 
 /// View of a Disk

dev-tools/omdb/src/bin/omdb/db.rs

@@ -6191,6 +6191,7 @@ async fn cmd_db_vmm_info(
                     rss_ram: ByteCount(rss),
                     reservoir_ram: ByteCount(reservoir),
                 },
+            instance_id: _,
         } = resource;
         const SLED_ID: &'static str = "sled ID";
         const THREADS: &'static str = "hardware threads";

nexus/auth/src/authz/api_resources.rs

@@ -713,6 +713,22 @@ authz_resource! {
     polar_snippet = InProject,
 }
 
+authz_resource! {
+    name = "AffinityGroup",
+    parent = "Project",
+    primary_key = Uuid,
+    roles_allowed = false,
+    polar_snippet = InProject,
+}
+
+authz_resource! {
+    name = "AntiAffinityGroup",
+    parent = "Project",
+    primary_key = Uuid,
+    roles_allowed = false,
+    polar_snippet = InProject,
+}
+
 authz_resource! {
     name = "InstanceNetworkInterface",
     parent = "Instance",

nexus/auth/src/authz/oso_generic.rs

@@ -125,6 +125,8 @@ pub fn make_omicron_oso(log: &slog::Logger) -> Result<OsoInit, anyhow::Error> {
         Disk::init(),
         Snapshot::init(),
         ProjectImage::init(),
+        AffinityGroup::init(),
+        AntiAffinityGroup::init(),
         Instance::init(),
         IpPool::init(),
         InstanceNetworkInterface::init(),

nexus/db-model/src/affinity.rs

@@ -0,0 +1,260 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at https://mozilla.org/MPL/5.0/.
+
+// Copyright 2024 Oxide Computer Company
+
+//! Database representation of affinity and anti-affinity groups
+
+use super::impl_enum_type;
+use super::Name;
+use crate::schema::affinity_group;
+use crate::schema::affinity_group_instance_membership;
+use crate::schema::anti_affinity_group;
+use crate::schema::anti_affinity_group_instance_membership;
+use crate::typed_uuid::DbTypedUuid;
+use chrono::{DateTime, Utc};
+use db_macros::Resource;
+use nexus_types::external_api::params;
+use nexus_types::external_api::views;
+use omicron_common::api::external;
+use omicron_common::api::external::IdentityMetadata;
+use omicron_uuid_kinds::AffinityGroupKind;
+use omicron_uuid_kinds::AffinityGroupUuid;
+use omicron_uuid_kinds::AntiAffinityGroupKind;
+use omicron_uuid_kinds::AntiAffinityGroupUuid;
+use omicron_uuid_kinds::GenericUuid;
+use omicron_uuid_kinds::InstanceKind;
+use omicron_uuid_kinds::InstanceUuid;
+use uuid::Uuid;
+
+impl_enum_type!(
+    #[derive(SqlType, Debug, QueryId)]
+    #[diesel(postgres_type(name = "affinity_policy", schema = "public"))]
+    pub struct AffinityPolicyEnum;
+
+    #[derive(Clone, Copy, Debug, AsExpression, FromSqlRow, PartialEq, Eq, Ord, PartialOrd)]
+    #[diesel(sql_type = AffinityPolicyEnum)]
+    pub enum AffinityPolicy;
+
+    // Enum values
+    Fail => b"fail"
+    Allow => b"allow"
+);
+
+impl From<AffinityPolicy> for external::AffinityPolicy {
+    fn from(policy: AffinityPolicy) -> Self {
+        match policy {
+            AffinityPolicy::Fail => Self::Fail,
+            AffinityPolicy::Allow => Self::Allow,
+        }
+    }
+}
+
+impl From<external::AffinityPolicy> for AffinityPolicy {
+    fn from(policy: external::AffinityPolicy) -> Self {
+        match policy {
+            external::AffinityPolicy::Fail => Self::Fail,
+            external::AffinityPolicy::Allow => Self::Allow,
+        }
+    }
+}
+
+impl_enum_type!(
+    #[derive(SqlType, Debug)]
+    #[diesel(postgres_type(name = "failure_domain", schema = "public"))]
+    pub struct FailureDomainEnum;
+
+    #[derive(Clone, Copy, Debug, AsExpression, FromSqlRow, PartialEq)]
+    #[diesel(sql_type = FailureDomainEnum)]
+    pub enum FailureDomain;
+
+    // Enum values
+    Sled => b"sled"
+);
+
+impl From<FailureDomain> for external::FailureDomain {
+    fn from(domain: FailureDomain) -> Self {
+        match domain {
+            FailureDomain::Sled => Self::Sled,
+        }
+    }
+}
+
+impl From<external::FailureDomain> for FailureDomain {
+    fn from(domain: external::FailureDomain) -> Self {
+        match domain {
+            external::FailureDomain::Sled => Self::Sled,
+        }
+    }
+}
+
+#[derive(
+    Queryable, Insertable, Clone, Debug, Resource, Selectable, PartialEq,
+)]
+#[diesel(table_name = affinity_group)]
+pub struct AffinityGroup {
+    #[diesel(embed)]
+    pub identity: AffinityGroupIdentity,
+    pub project_id: Uuid,
+    pub policy: AffinityPolicy,
+    pub failure_domain: FailureDomain,
+}
+
+impl AffinityGroup {
+    pub fn new(project_id: Uuid, params: params::AffinityGroupCreate) -> Self {
+        Self {
+            identity: AffinityGroupIdentity::new(
+                Uuid::new_v4(),
+                params.identity,
+            ),
+            project_id,
+            policy: params.policy.into(),
+            failure_domain: params.failure_domain.into(),
+        }
+    }
+}
+
+impl From<AffinityGroup> for views::AffinityGroup {
+    fn from(group: AffinityGroup) -> Self {
+        let identity = IdentityMetadata {
+            id: group.identity.id,
+            name: group.identity.name.into(),
+            description: group.identity.description,
+            time_created: group.identity.time_created,
+            time_modified: group.identity.time_modified,
+        };
+        Self {
+            identity,
+            policy: group.policy.into(),
+            failure_domain: group.failure_domain.into(),
+        }
+    }
+}
+
+/// Describes a set of updates for the [`AffinityGroup`] model.
+#[derive(AsChangeset)]
+#[diesel(table_name = affinity_group)]
+pub struct AffinityGroupUpdate {
+    pub name: Option<Name>,
+    pub description: Option<String>,
+    pub time_modified: DateTime<Utc>,
+}
+
+impl From<params::AffinityGroupUpdate> for AffinityGroupUpdate {
+    fn from(params: params::AffinityGroupUpdate) -> Self {
+        Self {
+            name: params.identity.name.map(Name),
+            description: params.identity.description,
+            time_modified: Utc::now(),
+        }
+    }
+}
+
+#[derive(
+    Queryable, Insertable, Clone, Debug, Resource, Selectable, PartialEq,
+)]
+#[diesel(table_name = anti_affinity_group)]
+pub struct AntiAffinityGroup {
+    #[diesel(embed)]
+    identity: AntiAffinityGroupIdentity,
+    pub project_id: Uuid,
+    pub policy: AffinityPolicy,
+    pub failure_domain: FailureDomain,
+}
+
+impl AntiAffinityGroup {
+    pub fn new(
+        project_id: Uuid,
+        params: params::AntiAffinityGroupCreate,
+    ) -> Self {
+        Self {
+            identity: AntiAffinityGroupIdentity::new(
+                Uuid::new_v4(),
+                params.identity,
+            ),
+            project_id,
+            policy: params.policy.into(),
+            failure_domain: params.failure_domain.into(),
+        }
+    }
+}
+
+impl From<AntiAffinityGroup> for views::AntiAffinityGroup {
+    fn from(group: AntiAffinityGroup) -> Self {
+        let identity = IdentityMetadata {
+            id: group.identity.id,
+            name: group.identity.name.into(),
+            description: group.identity.description,
+            time_created: group.identity.time_created,
+            time_modified: group.identity.time_modified,
+        };
+        Self {
+            identity,
+            policy: group.policy.into(),
+            failure_domain: group.failure_domain.into(),
+        }
+    }
+}
+
+/// Describes a set of updates for the [`AntiAffinityGroup`] model.
+#[derive(AsChangeset)]
+#[diesel(table_name = anti_affinity_group)]
+pub struct AntiAffinityGroupUpdate {
+    pub name: Option<Name>,
+    pub description: Option<String>,
+    pub time_modified: DateTime<Utc>,
+}
+
+impl From<params::AntiAffinityGroupUpdate> for AntiAffinityGroupUpdate {
+    fn from(params: params::AntiAffinityGroupUpdate) -> Self {
+        Self {
+            name: params.identity.name.map(Name),
+            description: params.identity.description,
+            time_modified: Utc::now(),
+        }
+    }
+}
+
+#[derive(Queryable, Insertable, Clone, Debug, Selectable)]
+#[diesel(table_name = affinity_group_instance_membership)]
+pub struct AffinityGroupInstanceMembership {
+    pub group_id: DbTypedUuid<AffinityGroupKind>,
+    pub instance_id: DbTypedUuid<InstanceKind>,
+}
+
+impl AffinityGroupInstanceMembership {
+    pub fn new(group_id: AffinityGroupUuid, instance_id: InstanceUuid) -> Self {
+        Self { group_id: group_id.into(), instance_id: instance_id.into() }
+    }
+}
+
+impl From<AffinityGroupInstanceMembership> for external::AffinityGroupMember {
+    fn from(member: AffinityGroupInstanceMembership) -> Self {
+        Self::Instance(member.instance_id.into_untyped_uuid())
+    }
+}
+
+#[derive(Queryable, Insertable, Clone, Debug, Selectable)]
+#[diesel(table_name = anti_affinity_group_instance_membership)]
+pub struct AntiAffinityGroupInstanceMembership {
+    pub group_id: DbTypedUuid<AntiAffinityGroupKind>,
+    pub instance_id: DbTypedUuid<InstanceKind>,
+}
+
+impl AntiAffinityGroupInstanceMembership {
+    pub fn new(
+        group_id: AntiAffinityGroupUuid,
+        instance_id: InstanceUuid,
+    ) -> Self {
+        Self { group_id: group_id.into(), instance_id: instance_id.into() }
+    }
+}
+
+impl From<AntiAffinityGroupInstanceMembership>
+    for external::AntiAffinityGroupMember
+{
+    fn from(member: AntiAffinityGroupInstanceMembership) -> Self {
+        Self::Instance(member.instance_id.into_untyped_uuid())
+    }
+}