okay i really hopep the CTE works

Author: hawkw

clients/nexus-client/src/lib.rs

@@ -136,6 +136,7 @@ impl From<omicron_common::api::internal::nexus::SledInstanceState>
             instance_state: s.instance_state.into(),
             propolis_id: s.propolis_id,
             vmm_state: s.vmm_state.into(),
+            migration_state: None,
         }
     }
 }

common/src/api/internal/nexus.rs

@@ -90,11 +90,13 @@ pub struct SledInstanceState {
     pub migration_state: Option<MigrationRuntimeState>,
 }
 
+/// An update from a sled regarding the state of a migration, indicating the
+/// role of the VMM whose migration state was updated.
 #[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
 pub struct MigrationRuntimeState {
     pub migration_id: Uuid,
-
     pub state: MigrationState,
+    pub role: MigrationRole,
 }
 
 /// The state of an instance's live migration.
@@ -127,6 +129,32 @@ impl fmt::Display for MigrationState {
     }
 }
 
+#[derive(
+    Clone, Copy, Debug, PartialEq, Eq, Deserialize, Serialize, JsonSchema,
+)]
+#[serde(rename_all = "snake_case")]
+pub enum MigrationRole {
+    /// This update concerns the source VMM of a migration.
+    Source,
+    /// This update concerns the target VMM of a migration.
+    Target,
+}
+
+impl MigrationRole {
+    pub fn label(&self) -> &'static str {
+        match self {
+            Self::Source => "source",
+            Self::Target => "target",
+        }
+    }
+}
+
+impl fmt::Display for MigrationRole {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(self.label())
+    }
+}
+
 // Oximeter producer/collector objects.
 
 /// The kind of metric producer this is.

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

@@ -389,6 +389,7 @@ impl DataStore {
             new_instance.clone(),
             *vmm_id,
             new_vmm.clone(),
+            None, // TODO: ELIZA ADD THIS
         );
 
         // The InstanceAndVmmUpdate query handles and indicates failure to find

nexus/db-queries/src/db/queries/instance.rs

@@ -12,8 +12,14 @@ use diesel::sql_types::{Nullable, Uuid as SqlUuid};
 use diesel::{pg::Pg, query_builder::AstPass};
 use diesel::{Column, ExpressionMethods, QueryDsl, RunQueryDsl};
 use nexus_db_model::{
-    schema::{instance::dsl as instance_dsl, vmm::dsl as vmm_dsl},
-    InstanceRuntimeState, VmmRuntimeState,
+    schema::{
+        instance::dsl as instance_dsl, migration::dsl as migration_dsl,
+        vmm::dsl as vmm_dsl,
+    },
+    InstanceRuntimeState, MigrationState, VmmRuntimeState,
+};
+use omicron_common::api::internal::nexus::{
+    MigrationRole, MigrationRuntimeState,
 };
 use uuid::Uuid;
 
@@ -76,6 +82,12 @@ pub struct InstanceAndVmmUpdate {
     vmm_find: Box<dyn QueryFragment<Pg> + Send>,
     instance_update: Box<dyn QueryFragment<Pg> + Send>,
     vmm_update: Box<dyn QueryFragment<Pg> + Send>,
+    migration: Option<MigrationUpdate>,
+}
+
+struct MigrationUpdate {
+    find: Box<dyn QueryFragment<Pg> + Send>,
+    update: Box<dyn QueryFragment<Pg> + Send>,
 }
 
 /// Contains the result of a combined instance-and-VMM update operation.
@@ -89,6 +101,11 @@ pub struct InstanceAndVmmUpdateResult {
     /// `Some(status)` if the target VMM was found; the wrapped `UpdateStatus`
     /// indicates whether the row was updated. `None` if the VMM was not found.
     pub vmm_status: Option<UpdateStatus>,
+
+    /// `Some(status)` if the target migration was found; the wrapped `UpdateStatus`
+    /// indicates whether the row was updated. `None` if the migration was not
+    /// found, or no migration update was performed.
+    pub migration_status: Option<UpdateStatus>,
 }
 
 /// Computes the update status to return from the results of queries that find
@@ -135,6 +152,7 @@ impl InstanceAndVmmUpdate {
         new_instance_runtime_state: InstanceRuntimeState,
         vmm_id: Uuid,
         new_vmm_runtime_state: VmmRuntimeState,
+        migration: Option<MigrationRuntimeState>,
     ) -> Self {
         let instance_find = Box::new(
             instance_dsl::instance
@@ -165,24 +183,92 @@ impl InstanceAndVmmUpdate {
                 .set(new_vmm_runtime_state),
         );
 
-        Self { instance_find, vmm_find, instance_update, vmm_update }
+        let migration = migration.map(
+            |MigrationRuntimeState { role, migration_id, state }| {
+                let state = MigrationState::from(state);
+                match role {
+                    MigrationRole::Target => {
+                        let find = Box::new(
+                            migration_dsl::migration
+                                .filter(migration_dsl::id.eq(migration_id))
+                                .filter(
+                                    migration_dsl::target_propolis_id
+                                        .eq(vmm_id),
+                                )
+                                .select(migration_dsl::id),
+                        );
+                        let update = Box::new(
+                            diesel::update(migration_dsl::migration)
+                                .filter(migration_dsl::id.eq(migration_id))
+                                .filter(
+                                    migration_dsl::target_propolis_id
+                                        .eq(vmm_id),
+                                )
+                                .set(migration_dsl::target_state.eq(state)),
+                        );
+                        MigrationUpdate { find, update }
+                    }
+                    MigrationRole::Source => {
+                        let find = Box::new(
+                            migration_dsl::migration
+                                .filter(migration_dsl::id.eq(migration_id))
+                                .filter(
+                                    migration_dsl::source_propolis_id
+                                        .eq(vmm_id),
+                                )
+                                .select(migration_dsl::id),
+                        );
+                        let update = Box::new(
+                            diesel::update(migration_dsl::migration)
+                                .filter(migration_dsl::id.eq(migration_id))
+                                .filter(
+                                    migration_dsl::source_propolis_id
+                                        .eq(vmm_id),
+                                )
+                                .set(migration_dsl::source_state.eq(state)),
+                        );
+                        MigrationUpdate { find, update }
+                    }
+                }
+            },
+        );
+
+        Self { instance_find, vmm_find, instance_update, vmm_update, migration }
     }
 
     pub async fn execute_and_check(
         self,
         conn: &(impl async_bb8_diesel::AsyncConnection<DbConnection> + Sync),
     ) -> Result<InstanceAndVmmUpdateResult, DieselError> {
-        let (vmm_found, vmm_updated, instance_found, instance_updated) =
-            self.get_result_async::<(Option<Uuid>,
-                                     Option<Uuid>,
-                                     Option<Uuid>,
-                                     Option<Uuid>)>(conn).await?;
+        let (
+            vmm_found,
+            vmm_updated,
+            instance_found,
+            instance_updated,
+            migration_found,
+            migration_updated,
+        ) = self
+            .get_result_async::<(
+                Option<Uuid>,
+                Option<Uuid>,
+                Option<Uuid>,
+                Option<Uuid>,
+                Option<Uuid>,
+                Option<Uuid>,
+            )>(conn)
+            .await?;
 
         let instance_status =
             compute_update_status(instance_found, instance_updated);
         let vmm_status = compute_update_status(vmm_found, vmm_updated);
+        let migration_status =
+            compute_update_status(migration_found, migration_updated);
 
-        Ok(InstanceAndVmmUpdateResult { instance_status, vmm_status })
+        Ok(InstanceAndVmmUpdateResult {
+            instance_status,
+            vmm_status,
+            migration_status,
+        })
     }
 }
 
@@ -197,6 +283,8 @@ impl Query for InstanceAndVmmUpdate {
         Nullable<SqlUuid>,
         Nullable<SqlUuid>,
         Nullable<SqlUuid>,
+        Nullable<SqlUuid>,
+        Nullable<SqlUuid>,
     );
 }
 
@@ -212,6 +300,12 @@ impl QueryFragment<Pg> for InstanceAndVmmUpdate {
         self.vmm_find.walk_ast(out.reborrow())?;
         out.push_sql(") AS id), ");
 
+        if let Some(MigrationUpdate { ref find, .. }) = self.migration {
+            out.push_sql("migration_found AS (SELECT (");
+            find.walk_ast(out.reborrow())?;
+            out.push_sql(") AS id), ");
+        }
+
         out.push_sql("instance_updated AS (");
         self.instance_update.walk_ast(out.reborrow())?;
         out.push_sql(" RETURNING id), ");
@@ -220,6 +314,12 @@ impl QueryFragment<Pg> for InstanceAndVmmUpdate {
         self.vmm_update.walk_ast(out.reborrow())?;
         out.push_sql(" RETURNING id), ");
 
+        if let Some(MigrationUpdate { ref update, .. }) = self.migration {
+            out.push_sql("migration_updated AS (");
+            update.walk_ast(out.reborrow())?;
+            out.push_sql(" RETURNING id), ");
+        }
+
         out.push_sql("vmm_result AS (");
         out.push_sql("SELECT vmm_found.");
         out.push_identifier(vmm_dsl::id::NAME)?;
@@ -246,9 +346,33 @@ impl QueryFragment<Pg> for InstanceAndVmmUpdate {
         out.push_identifier(instance_dsl::id::NAME)?;
         out.push_sql(") ");
 
+        if self.migration.is_some() {
+            out.push_sql("migration_result AS (");
+            out.push_sql("SELECT migration_found.");
+            out.push_identifier(migration_dsl::id::NAME)?;
+            out.push_sql(" AS found, migration_updated.");
+            out.push_identifier(migration_dsl::id::NAME)?;
+            out.push_sql(" AS updated");
+            out.push_sql(
+                " FROM migration_found LEFT JOIN migration_updated ON migration_found.",
+            );
+            out.push_identifier(migration_dsl::id::NAME)?;
+            out.push_sql(" = migration_updated.");
+            out.push_identifier(migration_dsl::id::NAME)?;
+            out.push_sql(") ");
+        }
+
         out.push_sql("SELECT vmm_result.found, vmm_result.updated, ");
-        out.push_sql("instance_result.found, instance_result.updated ");
-        out.push_sql("FROM vmm_result, instance_result;");
+        out.push_sql("instance_result.found, instance_result.updated, ");
+        if self.migration.is_some() {
+            out.push_sql("migration_result.found, migration_result.updated");
+        } else {
+            out.push_sql("NULL, NULL");
+        }
+        out.push_sql("FROM vmm_result, instance_result");
+        if self.migration.is_some() {
+            out.push_sql(", migration_result");
+        }
 
         Ok(())
     }

openapi/nexus-internal.json

@@ -3503,6 +3503,71 @@
         "minLength": 5,
         "maxLength": 17
       },
+      "MigrationRole": {
+        "oneOf": [
+          {
+            "description": "This update concerns the source VMM of a migration.",
+            "type": "string",
+            "enum": [
+              "source"
+            ]
+          },
+          {
+            "description": "This update concerns the target VMM of a migration.",
+            "type": "string",
+            "enum": [
+              "target"
+            ]
+          }
+        ]
+      },
+      "MigrationRuntimeState": {
+        "description": "An update from a sled regarding the state of a migration, indicating the role of the VMM whose migration state was updated.",
+        "type": "object",
+        "properties": {
+          "migration_id": {
+            "type": "string",
+            "format": "uuid"
+          },
+          "role": {
+            "$ref": "#/components/schemas/MigrationRole"
+          },
+          "state": {
+            "$ref": "#/components/schemas/MigrationState"
+          }
+        },
+        "required": [
+          "migration_id",
+          "role",
+          "state"
+        ]
+      },
+      "MigrationState": {
+        "description": "The state of an instance's live migration.",
+        "oneOf": [
+          {
+            "description": "The migration is in progress.",
+            "type": "string",
+            "enum": [
+              "in_progress"
+            ]
+          },
+          {
+            "description": "The migration has failed.",
+            "type": "string",
+            "enum": [
+              "failed"
+            ]
+          },
+          {
+            "description": "The migration has completed.",
+            "type": "string",
+            "enum": [
+              "completed"
+            ]
+          }
+        ]
+      },
       "Name": {
         "title": "A name unique within the parent collection",
         "description": "Names must begin with a lower case ASCII letter, be composed exclusively of lowercase ASCII, uppercase ASCII, numbers, and '-', and may not end with a '-'. Names cannot be a UUID though they may contain a UUID.",
@@ -4647,6 +4712,15 @@
               }
             ]
           },
+          "migration_state": {
+            "nullable": true,
+            "description": "The current state of any in-progress migration for this instance, as understood by this sled.",
+            "allOf": [
+              {
+                "$ref": "#/components/schemas/MigrationRuntimeState"
+              }
+            ]
+          },
           "propolis_id": {
             "description": "The ID of the VMM whose state is being reported.",
             "type": "string",