Skip to content

Commit

Permalink
okay i really hopep the CTE works
Browse files Browse the repository at this point in the history
  • Loading branch information
hawkw committed Jun 5, 2024
1 parent 05945f2 commit ef871e2
Show file tree
Hide file tree
Showing 6 changed files with 314 additions and 12 deletions.
1 change: 1 addition & 0 deletions clients/nexus-client/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}
}
Expand Down
30 changes: 29 additions & 1 deletion common/src/api/internal/nexus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions nexus/db-queries/src/db/datastore/instance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
146 changes: 135 additions & 11 deletions nexus/db-queries/src/db/queries/instance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
})
}
}

Expand All @@ -197,6 +283,8 @@ impl Query for InstanceAndVmmUpdate {
Nullable<SqlUuid>,
Nullable<SqlUuid>,
Nullable<SqlUuid>,
Nullable<SqlUuid>,
Nullable<SqlUuid>,
);
}

Expand All @@ -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), ");
Expand All @@ -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)?;
Expand All @@ -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(())
}
Expand Down
74 changes: 74 additions & 0 deletions openapi/nexus-internal.json
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down Expand Up @@ -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",
Expand Down
Loading

0 comments on commit ef871e2

Please sign in to comment.