nexus: provide stable PCI slot assignments for disks (#3408)

gjcolombo · web-flow · commit efc2566433eb · 2023-06-28T20:28:44.000Z
Add a `slot` column to the disk table and modify `instance_attach_disk` to set it when a disk is attached to an instance. Slots are chosen using the `NextItem` query generator that NICs use for their slot assignments. The bulk of this change is new code to combine the `NextItem` fragment with the other SQL statements that mark a disk as attached (changing its state and setting its `attach_instance_id`) in a manner that allows the resulting query fragment to be passed as the body of the SET clause of the UPDATE statement that `instance_attach_disk` passes to the `attach_resource` generic function. Slot assignments remain stable for as long as a disk is attached, but they aren't preserved if a disk is detached and reattached (the algorithm will choose the first available slot, even if the disk's previously assigned slot is available). Tests: - New integration test - On a dev cluster, created an instance with a single boot disk, then attached disks with alphabetically earlier and later names. Verified that the boot disk remained in slot 0 and that booting the instance succeeded (i.e. guest firmware was able to find the boot OS successfully). Also played around with detaching and reattaching disks and making sure via CRDB shell that their slot assignments were as expected. Fixes #3224.
diff --git a/common/src/sql/dbinit.sql b/common/src/sql/dbinit.sql
@@ -947,6 +947,7 @@ CREATE TABLE omicron.public.disk (
      */
     attach_instance_id UUID,
     state_generation INT NOT NULL,
+    slot INT2 CHECK (slot >= 0 AND slot < 8),
     time_state_updated TIMESTAMPTZ NOT NULL,
 
     /* Disk configuration */
diff --git a/nexus/db-model/src/disk.rs b/nexus/db-model/src/disk.rs
@@ -3,7 +3,7 @@
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
 use super::{BlockSize, ByteCount, DiskState, Generation};
-use crate::schema::disk;
+use crate::{schema::disk, unsigned::SqlU8};
 use chrono::{DateTime, Utc};
 use db_macros::Resource;
 use nexus_types::external_api::params;
@@ -44,6 +44,16 @@ pub struct Disk {
     #[diesel(embed)]
     pub runtime_state: DiskRuntimeState,
 
+    /// The PCI slot (within the bank of slots reserved to disks) to which this
+    /// disk should be attached if its attached instance is started, or None
+    /// if there is no such assignment.
+    ///
+    /// Slot assignments are managed entirely in Nexus and aren't modified by
+    /// runtime state changes in the sled agent, so this field is part of the
+    /// "main" disk struct and not the runtime state (even though the attachment
+    /// state and slot assignment will often change together).
+    pub slot: Option<SqlU8>,
+
     /// size of the Disk
     #[diesel(column_name = size_bytes)]
     pub size: ByteCount,
@@ -95,6 +105,7 @@ impl Disk {
             project_id,
             volume_id,
             runtime_state: runtime_initial,
+            slot: None,
             size: params.size.into(),
             block_size,
             create_snapshot_id,
diff --git a/nexus/db-model/src/schema.rs b/nexus/db-model/src/schema.rs
@@ -21,6 +21,7 @@ table! {
         attach_instance_id -> Nullable<Uuid>,
         state_generation -> Int8,
         time_state_updated -> Timestamptz,
+        slot -> Nullable<Int2>,
         size_bytes -> Int8,
         block_size -> crate::BlockSizeEnum,
         origin_snapshot -> Nullable<Uuid>,
diff --git a/nexus/db-model/src/unsigned.rs b/nexus/db-model/src/unsigned.rs
@@ -10,6 +10,57 @@ use diesel::sql_types;
 use serde::{Deserialize, Serialize};
 use std::convert::TryFrom;
 
+/// Representation of a [`u8`] in the database. The database does not support
+/// unsigned types; this adapter converts from the database's INT2 (a 2-byte
+/// signed integer) to an actual u8.
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    AsExpression,
+    Eq,
+    Ord,
+    PartialEq,
+    PartialOrd,
+    FromSqlRow,
+    Serialize,
+    Deserialize,
+)]
+#[diesel(sql_type = sql_types::Int2)]
+#[repr(transparent)]
+pub struct SqlU8(pub u8);
+
+NewtypeFrom! { () pub struct SqlU8(u8); }
+NewtypeDeref! { () pub struct SqlU8(u8); }
+
+impl SqlU8 {
+    pub fn new(value: u8) -> Self {
+        Self(value)
+    }
+}
+
+impl ToSql<sql_types::Int2, Pg> for SqlU8 {
+    fn to_sql<'a>(
+        &'a self,
+        out: &mut serialize::Output<'a, '_, Pg>,
+    ) -> serialize::Result {
+        <i16 as ToSql<sql_types::Int2, Pg>>::to_sql(
+            &i16::from(self.0),
+            &mut out.reborrow(),
+        )
+    }
+}
+
+impl<DB> FromSql<sql_types::Int2, DB> for SqlU8
+where
+    DB: Backend,
+    i16: FromSql<sql_types::Int2, DB>,
+{
+    fn from_sql(bytes: DB::RawValue<'_>) -> deserialize::Result<Self> {
+        u8::try_from(i16::from_sql(bytes)?).map(SqlU8).map_err(|e| e.into())
+    }
+}
+
 /// Representation of a [`u16`] in the database.
 /// We need this because the database does not support unsigned types.
 /// This handles converting from the database's INT4 to the actual u16.
diff --git a/nexus/db-queries/src/db/datastore/disk.rs b/nexus/db-queries/src/db/datastore/disk.rs
@@ -26,6 +26,7 @@ use crate::db::model::Instance;
 use crate::db::model::Name;
 use crate::db::model::Project;
 use crate::db::pagination::paginated;
+use crate::db::queries::disk::DiskSetClauseForAttach;
 use crate::db::update_and_check::UpdateAndCheck;
 use crate::db::update_and_check::UpdateStatus;
 use async_bb8_diesel::AsyncRunQueryDsl;
@@ -183,26 +184,22 @@ impl DataStore {
             db::model::InstanceState(api::external::InstanceState::Stopped),
         ];
 
-        let attached_label =
-            api::external::DiskState::Attached(authz_instance.id()).label();
+        let attach_update = DiskSetClauseForAttach::new(authz_instance.id());
 
-        let (instance, disk) = Instance::attach_resource(
+        let query = Instance::attach_resource(
             authz_instance.id(),
             authz_disk.id(),
-            instance::table
-                .into_boxed()
-                .filter(instance::dsl::state.eq_any(ok_to_attach_instance_states)),
-            disk::table
-                .into_boxed()
-                .filter(disk::dsl::disk_state.eq_any(ok_to_attach_disk_state_labels)),
+            instance::table.into_boxed().filter(
+                instance::dsl::state.eq_any(ok_to_attach_instance_states),
+            ),
+            disk::table.into_boxed().filter(
+                disk::dsl::disk_state.eq_any(ok_to_attach_disk_state_labels),
+            ),
             max_disks,
-            diesel::update(disk::dsl::disk)
-                .set((
-                    disk::dsl::disk_state.eq(attached_label),
-                    disk::dsl::attach_instance_id.eq(authz_instance.id())
-                ))
-        )
-        .attach_and_get_result_async(self.pool_authorized(opctx).await?)
+            diesel::update(disk::dsl::disk).set(attach_update),
+        );
+
+        let (instance, disk) = query.attach_and_get_result_async(self.pool_authorized(opctx).await?)
         .await
         .or_else(|e| {
             match e {
@@ -330,7 +327,8 @@ impl DataStore {
             diesel::update(disk::dsl::disk)
                 .set((
                     disk::dsl::disk_state.eq(detached_label),
-                    disk::dsl::attach_instance_id.eq(Option::<Uuid>::None)
+                    disk::dsl::attach_instance_id.eq(Option::<Uuid>::None),
+                    disk::dsl::slot.eq(Option::<i16>::None)
                 ))
         )
         .detach_and_get_result_async(self.pool_authorized(opctx).await?)
diff --git a/nexus/db-queries/src/db/datastore/instance.rs b/nexus/db-queries/src/db/datastore/instance.rs
@@ -285,6 +285,7 @@ impl DataStore {
             diesel::update(disk::dsl::disk).set((
                 disk::dsl::disk_state.eq(detached_label),
                 disk::dsl::attach_instance_id.eq(Option::<Uuid>::None),
+                disk::dsl::slot.eq(Option::<i16>::None),
             )),
         )
         .detach_and_get_result_async(self.pool_authorized(opctx).await?)
diff --git a/nexus/db-queries/src/db/queries/disk.rs b/nexus/db-queries/src/db/queries/disk.rs
@@ -0,0 +1,168 @@
+// 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/2.0/.
+
+//! Helper queries for working with disks.
+
+use crate::db;
+use crate::db::queries::next_item::{DefaultShiftGenerator, NextItem};
+use diesel::{
+    pg::Pg,
+    query_builder::{AstPass, QueryFragment, QueryId},
+    sql_types, Column, QueryResult,
+};
+use uuid::Uuid;
+
+/// The maximum number of disks that can be attached to an instance.
+//
+// This is defined here for layering reasons: the main Nexus crate depends on
+// the db-queries crate, so the disk-per-instance limit lives here and Nexus
+// proper re-exports it.
+pub const MAX_DISKS_PER_INSTANCE: u32 = 8;
+
+/// A wrapper for the query that selects a PCI slot for a newly-attached disk.
+///
+/// The general idea is to produce a query that left joins a single-column table
+/// containing the sequence [0, 1, 2, ..., MAX_DISKS] against the disk table,
+/// filtering for disks that are attached to the instance of interest and whose
+/// currently-assigned slots match the slots in the sequence. Filtering this
+/// query to just those rows where a slot has no disk attached yields the first
+/// available slot number. If no slots are available, the returned value is
+/// MAX_DISKS, which is not a valid slot number and which will be rejected when
+/// attempting to update the row (the "slot" column has a CHECK clause that
+/// enforces this at the database level).
+///
+/// See the `NextItem` documentation for more details.
+#[derive(Debug, Clone, Copy)]
+struct NextDiskSlot {
+    inner: NextItem<
+        db::schema::disk::table,
+        i16,
+        db::schema::disk::dsl::slot,
+        Uuid,
+        db::schema::disk::dsl::attach_instance_id,
+    >,
+}
+
+impl NextDiskSlot {
+    fn new(instance_id: Uuid) -> Self {
+        let generator = DefaultShiftGenerator {
+            base: 0,
+            max_shift: i64::try_from(MAX_DISKS_PER_INSTANCE).unwrap(),
+            min_shift: 0,
+        };
+        Self { inner: NextItem::new_scoped(generator, instance_id) }
+    }
+}
+
+impl QueryId for NextDiskSlot {
+    type QueryId = ();
+    const HAS_STATIC_QUERY_ID: bool = false;
+}
+
+impl QueryFragment<Pg> for NextDiskSlot {
+    fn walk_ast<'a>(&'a self, mut out: AstPass<'_, 'a, Pg>) -> QueryResult<()> {
+        self.inner.walk_ast(out.reborrow())?;
+        Ok(())
+    }
+}
+
+/// Produces a query fragment containing the SET clause to use in the UPDATE
+/// statement when attaching a disk to an instance.
+///
+/// The expected form of the statement is
+///
+/// ```sql
+/// SET attach_instance_id = instance_id,
+///     disk_state = 'attached',
+///     slot = (SELECT 0 + shift AS slot FROM
+///             (SELECT generate_series(0, 8) AS shift
+///              UNION ALL
+///              SELECT generate_series(0, -1) AS shift)
+///             LEFT OUTER JOIN disk
+///             ON (attach_instance_id, slot, time_deleted IS NULL) =
+///                (instance_id, 0 + shift, TRUE)
+///             WHERE slot IS NULL
+///             LIMIT 1)
+/// ```
+///
+/// This fragment can be passed to an `attach_resource` operation by supplying
+/// it as the argument to a `set`, e.g.
+/// `diesel::update(disk::dsl::disk).set(DiskSetClauseForAttach::new(instance_id))`.
+#[derive(Debug, Clone)]
+pub struct DiskSetClauseForAttach {
+    attach_instance_id: Uuid,
+    next_slot: NextDiskSlot,
+}
+
+impl DiskSetClauseForAttach {
+    pub fn new(instance_id: Uuid) -> Self {
+        Self {
+            attach_instance_id: instance_id,
+            next_slot: NextDiskSlot::new(instance_id),
+        }
+    }
+}
+
+impl QueryId for DiskSetClauseForAttach {
+    type QueryId = ();
+    const HAS_STATIC_QUERY_ID: bool = false;
+}
+
+impl QueryFragment<Pg> for DiskSetClauseForAttach {
+    fn walk_ast<'a>(&'a self, mut out: AstPass<'_, 'a, Pg>) -> QueryResult<()> {
+        let attached_label =
+            omicron_common::api::external::DiskState::Attached(
+                self.attach_instance_id,
+            )
+            .label();
+
+        // Ideally this code would not have to handcraft the entire SET
+        // statement. It would be cleaner to write something like
+        //
+        // diesel::update(disk::dsl::disk).set(
+        //  (disk::dsl::disk::attach_instance_id.eq(attach_instance_id),
+        //   disk::dsl::disk::disk_state.eq(attached_label),
+        //   disk::dsl::disk::slot.eq(next_slot_query))
+        //
+        // where `next_slot_query` is a wrapper around the `NextItem` query
+        // above that yields the slot-choosing SELECT statement. Diesel provides
+        // a `single_value` adapter function that's supposed to do this: given
+        // a query that returns a single column, it appends a LIMIT 1 to the
+        // query and parenthesizes it so that it can be used as the argument to
+        // a SET statement. The sticky bit is the trait bounds: SingleValueDsl
+        // is implemented for SelectQuery + LimitDsl, and while it's easy enough
+        // to impl SelectQuery for NextDiskSlot, implementing LimitDsl is harder
+        // (and seems to be discouraged by the trait docs)--it is natively
+        // implemented for diesel::Table, but here there's no table type, in
+        // part because the left side of the join is synthesized from thin air
+        // by the NextItem query fragment.
+        //
+        // Rather than spar extensively with Diesel over these details, just
+        // implement the whole SET by hand.
+        out.push_identifier(db::schema::disk::dsl::attach_instance_id::NAME)?;
+        out.push_sql(" = ");
+        out.push_bind_param::<sql_types::Uuid, Uuid>(&self.attach_instance_id)?;
+        out.push_sql(", ");
+        out.push_identifier(db::schema::disk::dsl::disk_state::NAME)?;
+        out.push_sql(" = ");
+        out.push_bind_param::<sql_types::Text, str>(attached_label)?;
+        out.push_sql(", ");
+        out.push_identifier(db::schema::disk::dsl::slot::NAME)?;
+        out.push_sql(" = (");
+        self.next_slot.walk_ast(out.reborrow())?;
+        out.push_sql(")");
+
+        Ok(())
+    }
+}
+
+// Required to pass `DiskSetClauseForAttach` to `set`.
+impl diesel::query_builder::AsChangeset for DiskSetClauseForAttach {
+    type Target = db::schema::disk::dsl::disk;
+    type Changeset = Self;
+
+    fn as_changeset(self) -> Self::Changeset {
+        self
+    }
+}
diff --git a/nexus/db-queries/src/db/queries/mod.rs b/nexus/db-queries/src/db/queries/mod.rs
@@ -5,6 +5,7 @@
 //! Specialized queries for inserting database records, usually to maintain
 //! complex invariants that are most accurately expressed in a single query.
 
+pub mod disk;
 pub mod external_ip;
 pub mod ip_pool;
 #[macro_use]
diff --git a/nexus/src/app/instance.rs b/nexus/src/app/instance.rs
@@ -625,12 +625,30 @@ impl super::Nexus {
             .await?;
 
         let mut disk_reqs = vec![];
-        for (i, disk) in disks.iter().enumerate() {
+        for disk in &disks {
+            // Disks that are attached to an instance should always have a slot
+            // assignment, but if for some reason this one doesn't, return an
+            // error instead of taking down the whole process.
+            let slot = match disk.slot {
+                Some(s) => s,
+                None => {
+                    error!(self.log, "attached disk has no PCI slot assignment";
+                       "disk_id" => %disk.id(),
+                       "disk_name" => disk.name().to_string(),
+                       "instance" => ?disk.runtime_state.attach_instance_id);
+
+                    return Err(Error::internal_error(&format!(
+                        "disk {} is attached but has no PCI slot assignment",
+                        disk.id()
+                    )));
+                }
+            };
+
             let volume =
                 self.db_datastore.volume_checkout(disk.volume_id).await?;
             disk_reqs.push(sled_agent_client::types::DiskRequest {
                 name: disk.name().to_string(),
-                slot: sled_agent_client::types::Slot(i as u8),
+                slot: sled_agent_client::types::Slot(slot.0),
                 read_only: false,
                 device: "nvme".to_string(),
                 volume_construction_request: serde_json::from_str(
diff --git a/nexus/src/app/mod.rs b/nexus/src/app/mod.rs
@@ -66,7 +66,7 @@ pub mod sagas;
 // TODO: When referring to API types, we should try to include
 // the prefix unless it is unambiguous.
 
-pub(crate) const MAX_DISKS_PER_INSTANCE: u32 = 8;
+pub(crate) use nexus_db_queries::db::queries::disk::MAX_DISKS_PER_INSTANCE;
 
 pub(crate) const MAX_NICS_PER_INSTANCE: usize = 8;
 
diff --git a/nexus/tests/integration_tests/disks.rs b/nexus/tests/integration_tests/disks.rs

Original file line number	Diff line number	Diff line change
`@@ -285,6 +285,7 @@ impl DataStore {`
`285`	`285`	`diesel::update(disk::dsl::disk).set((`
`286`	`286`	`disk::dsl::disk_state.eq(detached_label),`
`287`	`287`	`disk::dsl::attach_instance_id.eq(Option::<Uuid>::None),`
	`288`	`+ disk::dsl::slot.eq(Option::<i16>::None),`
`288`	`289`	`)),`
`289`	`290`	`)`
`290`	`291`	`.detach_and_get_result_async(self.pool_authorized(opctx).await?)`