From 3dcf18429e8606e64bd64bc886313362f7dd60be Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Thu, 8 Jan 2026 13:31:29 -0800 Subject: [PATCH 1/6] Add images to resource accounting --- nexus/db-model/src/schema_versions.rs | 3 +- .../src/virtual_provisioning_resource.rs | 2 + .../virtual_provisioning_collection.rs | 90 +++++ .../virtual_provisioning_collection_update.rs | 330 ++++++++++++++++++ nexus/src/app/sagas/image_create.rs | 119 ++++++- nexus/src/app/sagas/image_delete.rs | 44 +++ nexus/tests/integration_tests/schema.rs | 206 +++++++++++ schema/crdb/dbinit.sql | 2 +- schema/crdb/image-virtual-provisioning/up.sql | 76 ++++ 9 files changed, 856 insertions(+), 16 deletions(-) create mode 100644 schema/crdb/image-virtual-provisioning/up.sql diff --git a/nexus/db-model/src/schema_versions.rs b/nexus/db-model/src/schema_versions.rs index 293b6086c6b..502b6cf1345 100644 --- a/nexus/db-model/src/schema_versions.rs +++ b/nexus/db-model/src/schema_versions.rs @@ -16,7 +16,7 @@ use std::{collections::BTreeMap, sync::LazyLock}; /// /// This must be updated when you change the database schema. Refer to /// schema/crdb/README.adoc in the root of this repository for details. -pub const SCHEMA_VERSION: Version = Version::new(217, 0, 0); +pub const SCHEMA_VERSION: Version = Version::new(218, 0, 0); /// List of all past database schema versions, in *reverse* order /// @@ -28,6 +28,7 @@ static KNOWN_VERSIONS: LazyLock> = LazyLock::new(|| { // | leaving the first copy as an example for the next person. // v // KnownVersion::new(next_int, "unique-dirname-with-the-sql-files"), + KnownVersion::new(218, "image-virtual-provisioning"), KnownVersion::new(217, "multiple-default-ip-pools-per-silo"), KnownVersion::new(216, "add-trust-quorum"), KnownVersion::new(215, "support-up-to-12-disks"), diff --git a/nexus/db-model/src/virtual_provisioning_resource.rs b/nexus/db-model/src/virtual_provisioning_resource.rs index 031d12e268a..347b69f8507 100644 --- a/nexus/db-model/src/virtual_provisioning_resource.rs +++ b/nexus/db-model/src/virtual_provisioning_resource.rs @@ -13,6 +13,7 @@ pub enum ResourceTypeProvisioned { Instance, Disk, Snapshot, + Image, } impl std::fmt::Display for ResourceTypeProvisioned { @@ -21,6 +22,7 @@ impl std::fmt::Display for ResourceTypeProvisioned { ResourceTypeProvisioned::Instance => write!(f, "instance"), ResourceTypeProvisioned::Disk => write!(f, "disk"), ResourceTypeProvisioned::Snapshot => write!(f, "snapshot"), + ResourceTypeProvisioned::Image => write!(f, "image"), } } } diff --git a/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs b/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs index 5cf7f9a9f73..e24ec4c88f4 100644 --- a/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs +++ b/nexus/db-queries/src/db/datastore/virtual_provisioning_collection.rs @@ -24,6 +24,7 @@ use uuid::Uuid; pub enum StorageType { Disk, Snapshot, + Image, } impl From for crate::db::model::ResourceTypeProvisioned { @@ -37,6 +38,9 @@ impl From for crate::db::model::ResourceTypeProvisioned { StorageType::Snapshot => { crate::db::model::ResourceTypeProvisioned::Snapshot } + StorageType::Image => { + crate::db::model::ResourceTypeProvisioned::Image + } } } } @@ -178,6 +182,23 @@ impl DataStore { .await } + pub async fn virtual_provisioning_collection_insert_project_image( + &self, + opctx: &OpContext, + id: Uuid, + project_id: Uuid, + disk_byte_diff: ByteCount, + ) -> Result, Error> { + self.virtual_provisioning_collection_insert_storage( + opctx, + id, + project_id, + disk_byte_diff, + StorageType::Image, + ) + .await + } + /// Transitively updates all provisioned disk provisions from project -> fleet. async fn virtual_provisioning_collection_insert_storage( &self, @@ -236,6 +257,22 @@ impl DataStore { .await } + pub async fn virtual_provisioning_collection_delete_project_image( + &self, + opctx: &OpContext, + id: Uuid, + project_id: Uuid, + disk_byte_diff: ByteCount, + ) -> Result, Error> { + self.virtual_provisioning_collection_delete_storage( + opctx, + id, + project_id, + disk_byte_diff, + ) + .await + } + // Transitively updates all provisioned disk provisions from project -> fleet. async fn virtual_provisioning_collection_delete_storage( &self, @@ -258,6 +295,59 @@ impl DataStore { Ok(provisions) } + /// Insert storage accounting for a silo-scoped image. + /// + /// This updates the Silo and Fleet collections (no Project). + pub async fn virtual_provisioning_collection_insert_silo_image( + &self, + opctx: &OpContext, + id: Uuid, + silo_id: Uuid, + disk_byte_diff: ByteCount, + ) -> Result, Error> { + let provisions = + VirtualProvisioningCollectionUpdate::new_insert_silo_storage( + id, + disk_byte_diff, + silo_id, + StorageType::Image, + ) + .get_results_async(&*self.pool_connection_authorized(opctx).await?) + .await + .map_err(|e| { + crate::db::queries::virtual_provisioning_collection_update::from_diesel(e) + })?; + self.virtual_provisioning_collection_producer + .append_disk_metrics(&provisions)?; + Ok(provisions) + } + + /// Delete storage accounting for a silo-scoped image. + /// + /// This updates the Silo and Fleet collections (no Project). + pub async fn virtual_provisioning_collection_delete_silo_image( + &self, + opctx: &OpContext, + id: Uuid, + silo_id: Uuid, + disk_byte_diff: ByteCount, + ) -> Result, Error> { + let provisions = + VirtualProvisioningCollectionUpdate::new_delete_silo_storage( + id, + disk_byte_diff, + silo_id, + ) + .get_results_async(&*self.pool_connection_authorized(opctx).await?) + .await + .map_err(|e| { + crate::db::queries::virtual_provisioning_collection_update::from_diesel(e) + })?; + self.virtual_provisioning_collection_producer + .append_disk_metrics(&provisions)?; + Ok(provisions) + } + /// Transitively updates all CPU/RAM provisions from project -> fleet. pub async fn virtual_provisioning_collection_insert_instance( &self, diff --git a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs index ff113378424..dc2ba391b19 100644 --- a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs +++ b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs @@ -447,6 +447,336 @@ FROM ) } + /// Insert storage accounting for silo-scoped resources (no project). + /// + /// This updates the Silo and Fleet collections, skipping the Project level. + pub fn new_insert_silo_storage( + id: uuid::Uuid, + disk_byte_diff: ByteCount, + silo_id: uuid::Uuid, + storage_type: crate::db::datastore::StorageType, + ) -> TypedSqlQuery> { + let mut provision = + VirtualProvisioningResource::new(id, storage_type.into()); + provision.virtual_disk_bytes_provisioned = disk_byte_diff; + + Self::apply_silo_update(UpdateKind::InsertStorage(provision), silo_id) + } + + /// Delete storage accounting for silo-scoped resources (no project). + /// + /// This updates the Silo and Fleet collections, skipping the Project level. + pub fn new_delete_silo_storage( + id: uuid::Uuid, + disk_byte_diff: ByteCount, + silo_id: uuid::Uuid, + ) -> TypedSqlQuery> { + Self::apply_silo_update( + UpdateKind::DeleteStorage { id, disk_byte_diff }, + silo_id, + ) + } + + // Similar to apply_update but for silo-scoped resources (no project). + // + // Propagated updates include: + // - Silo + // - Fleet + fn apply_silo_update( + update_kind: UpdateKind, + silo_id: uuid::Uuid, + ) -> TypedSqlQuery> { + let mut query = QueryBuilder::new(); + + // For silo-scoped resources, we don't need to look up the parent silo + // from a project - we use the silo_id directly. + query + .sql( + " +WITH + all_collections + AS ( + (SELECT ", + ) + .param() + .sql( + " AS id) + UNION (SELECT ", + ) + .param() + .sql( + " AS id) + ),", + ) + .bind::(silo_id) + .bind::(*nexus_db_fixed_data::FLEET_ID) + .sql( + " + quotas + AS ( + SELECT + silo_quotas.silo_id, + silo_quotas.cpus, + silo_quotas.memory_bytes AS memory, + silo_quotas.storage_bytes AS storage + FROM + silo_quotas + WHERE + silo_quotas.silo_id = ", + ) + .param() + .sql( + " + ),", + ) + .bind::(silo_id) + .sql( + " + silo_provisioned + AS ( + SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned, + virtual_provisioning_collection.virtual_disk_bytes_provisioned + FROM + virtual_provisioning_collection + WHERE + virtual_provisioning_collection.id = ", + ) + .param() + .sql( + " + ),", + ) + .bind::(silo_id); + + match update_kind.clone() { + UpdateKind::InsertStorage(resource) => query + .sql( + " + do_update + AS ( + SELECT + ( + ( + SELECT count(*) + FROM virtual_provisioning_resource + WHERE virtual_provisioning_resource.id = ", + ) + .param() + .sql( + " + LIMIT 1 + ) + = 0 + AND CAST( + IF( + ( + ", + ) + .param() + .sql( + " = 0 + OR (SELECT quotas.storage FROM quotas LIMIT 1) + >= ( + ( + SELECT + silo_provisioned.virtual_disk_bytes_provisioned + FROM + silo_provisioned + LIMIT + 1 + ) + + ", + ) + .param() + .sql(concatcp!( + " + ) + ), + 'TRUE', + '", + NOT_ENOUGH_STORAGE_SENTINEL, + "' + ) + AS BOOL + ) + ) + AS update + )," + )) + .bind::(resource.id) + .bind::( + resource.virtual_disk_bytes_provisioned, + ) + .bind::( + resource.virtual_disk_bytes_provisioned, + ), + UpdateKind::DeleteStorage { id, .. } => query + .sql( + " + do_update + AS ( + SELECT + ( + SELECT + count(*) + FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = ", + ) + .param() + .sql( + " + LIMIT + 1 + ) = 1 + AS update + ),", + ) + .bind::(id), + // Instance operations are not supported for silo-scoped updates + UpdateKind::InsertInstance(_) + | UpdateKind::DeleteInstance { .. } => { + unreachable!( + "Instance operations are not supported for silo-scoped updates" + ) + } + }; + + match update_kind.clone() { + UpdateKind::InsertStorage(resource) => query + .sql( + " + unused_cte_arm + AS ( + INSERT + INTO + virtual_provisioning_resource + ( + id, + time_modified, + resource_type, + virtual_disk_bytes_provisioned, + cpus_provisioned, + ram_provisioned + ) + VALUES + (", + ) + .param() + .sql(", DEFAULT, ") + .param() + .sql(", ") + .param() + .sql(", ") + .param() + .sql(", ") + .param() + .sql( + ") + ON CONFLICT + DO + NOTHING + RETURNING ", + ) + .sql(AllColumnsOfVirtualResource::with_prefix( + "virtual_provisioning_resource", + )) + .sql("),") + .bind::(resource.id) + .bind::(resource.resource_type) + .bind::( + resource.virtual_disk_bytes_provisioned, + ) + .bind::(resource.cpus_provisioned) + .bind::(resource.ram_provisioned), + UpdateKind::DeleteStorage { id, .. } => query + .sql( + " + unused_cte_arm + AS ( + DELETE FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = ", + ) + .param() + .sql( + " + AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING ", + ) + .sql(AllColumnsOfVirtualResource::with_prefix( + "virtual_provisioning_resource", + )) + .sql("),") + .bind::(id), + // Instance operations are not supported for silo-scoped updates + UpdateKind::InsertInstance(_) + | UpdateKind::DeleteInstance { .. } => { + unreachable!( + "Instance operations are not supported for silo-scoped updates" + ) + } + }; + + query.sql( + " + virtual_provisioning_collection + AS ( + UPDATE + virtual_provisioning_collection + SET", + ); + match update_kind.clone() { + UpdateKind::InsertStorage(resource) => query + .sql( + " + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned + ", + ) + .param() + .bind::( + resource.virtual_disk_bytes_provisioned, + ), + UpdateKind::DeleteStorage { disk_byte_diff, .. } => query + .sql( + " + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned - ", + ) + .param() + .bind::(disk_byte_diff), + // Instance operations are not supported for silo-scoped updates + UpdateKind::InsertInstance(_) + | UpdateKind::DeleteInstance { .. } => { + unreachable!( + "Instance operations are not supported for silo-scoped updates" + ) + } + }; + + query.sql(" + WHERE + virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) + AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING " + ).sql(AllColumnsOfVirtualCollection::with_prefix("virtual_provisioning_collection")).sql(" + ) +SELECT " + ).sql(AllColumnsOfVirtualCollection::with_prefix("virtual_provisioning_collection")).sql(" +FROM + virtual_provisioning_collection +"); + + query.query() + } + pub fn new_insert_instance( id: InstanceUuid, cpus_diff: i64, diff --git a/nexus/src/app/sagas/image_create.rs b/nexus/src/app/sagas/image_create.rs index 60e09b3381c..71e70e376a1 100644 --- a/nexus/src/app/sagas/image_create.rs +++ b/nexus/src/app/sagas/image_create.rs @@ -10,6 +10,7 @@ use crate::app::sagas::declare_saga_actions; use crate::app::{authn, authz, db}; use crate::external_api::params; use nexus_db_lookup::LookupPath; +use nexus_types::identity::Resource; use omicron_common::api::external; use omicron_common::api::external::Error; use omicron_uuid_kinds::GenericUuid; @@ -67,6 +68,10 @@ declare_saga_actions! { + simc_create_image_record - simc_create_image_record_undo } + SPACE_ACCOUNT -> "no_result" { + + simc_account_space + - simc_account_space_undo + } } // image create saga: definition @@ -111,6 +116,7 @@ impl NexusSaga for SagaImageCreate { builder.append(get_source_volume_action()); builder.append(create_image_record_action()); + builder.append(space_account_action()); Ok(builder.build()?) } @@ -321,34 +327,119 @@ async fn simc_create_image_record_undo( let image_id = sagactx.lookup::("image_id")?; + // Use fetch_optional to make this undo idempotent - if the image was + // already deleted (e.g., by a previous undo attempt), we can safely + // skip the delete operation. match ¶ms.image_type { ImageType::Project { .. } => { - let (.., authz_image, db_image) = - LookupPath::new(&opctx, osagactx.datastore()) - .project_image_id(image_id) - .fetch() + let lookup_result = LookupPath::new(&opctx, osagactx.datastore()) + .project_image_id(image_id) + .fetch() + .await; + + if let Ok((.., authz_image, db_image)) = lookup_result { + osagactx + .datastore() + .project_image_delete(&opctx, &authz_image, db_image) .await?; + } + } + ImageType::Silo { .. } => { + let lookup_result = LookupPath::new(&opctx, osagactx.datastore()) + .silo_image_id(image_id) + .fetch() + .await; + + if let Ok((.., authz_image, db_image)) = lookup_result { + osagactx + .datastore() + .silo_image_delete(&opctx, &authz_image, db_image) + .await?; + } + } + } + + Ok(()) +} + +async fn simc_account_space( + sagactx: NexusActionContext, +) -> Result<(), ActionError> { + let osagactx = sagactx.user_data(); + let params = sagactx.saga_params::()?; + let opctx = crate::context::op_context_for_saga_action( + &sagactx, + ¶ms.serialized_authn, + ); + + let created_image = sagactx.lookup::("created_image")?; + + match ¶ms.image_type { + ImageType::Project { authz_project, .. } => { osagactx .datastore() - .project_image_delete(&opctx, &authz_image, db_image) - .await?; + .virtual_provisioning_collection_insert_project_image( + &opctx, + created_image.id(), + authz_project.id(), + created_image.size, + ) + .await + .map_err(ActionError::action_failed)?; + } + ImageType::Silo { authz_silo } => { + osagactx + .datastore() + .virtual_provisioning_collection_insert_silo_image( + &opctx, + created_image.id(), + authz_silo.id(), + created_image.size, + ) + .await + .map_err(ActionError::action_failed)?; } + } + Ok(()) +} - ImageType::Silo { .. } => { - let (.., authz_image, db_image) = - LookupPath::new(&opctx, osagactx.datastore()) - .silo_image_id(image_id) - .fetch() - .await?; +async fn simc_account_space_undo( + sagactx: NexusActionContext, +) -> Result<(), anyhow::Error> { + let osagactx = sagactx.user_data(); + let params = sagactx.saga_params::()?; + let opctx = crate::context::op_context_for_saga_action( + &sagactx, + ¶ms.serialized_authn, + ); + + let created_image = sagactx.lookup::("created_image")?; + match ¶ms.image_type { + ImageType::Project { authz_project, .. } => { osagactx .datastore() - .silo_image_delete(&opctx, &authz_image, db_image) + .virtual_provisioning_collection_delete_project_image( + &opctx, + created_image.id(), + authz_project.id(), + created_image.size, + ) + .await?; + } + ImageType::Silo { authz_silo } => { + osagactx + .datastore() + .virtual_provisioning_collection_delete_silo_image( + &opctx, + created_image.id(), + authz_silo.id(), + created_image.size, + ) .await?; } } - Ok(()) } diff --git a/nexus/src/app/sagas/image_delete.rs b/nexus/src/app/sagas/image_delete.rs index 32b6655abac..1dafdd07e04 100644 --- a/nexus/src/app/sagas/image_delete.rs +++ b/nexus/src/app/sagas/image_delete.rs @@ -6,6 +6,7 @@ use super::{ActionRegistry, NexusActionContext, NexusSaga}; use crate::app::sagas; use crate::app::sagas::declare_saga_actions; use nexus_db_queries::{authn, authz, db}; +use nexus_types::identity::Resource; use omicron_uuid_kinds::VolumeUuid; use serde::Deserialize; use serde::Serialize; @@ -37,6 +38,9 @@ pub(crate) struct Params { declare_saga_actions! { image_delete; + SPACE_ACCOUNT -> "no_result0" { + + sid_account_space + } DELETE_IMAGE_RECORD -> "no_result1" { + sid_delete_image_record } @@ -56,6 +60,7 @@ impl NexusSaga for SagaImageDelete { params: &Self::Params, mut builder: steno::DagBuilder, ) -> Result { + builder.append(space_account_action()); builder.append(delete_image_record_action()); const DELETE_VOLUME_PARAMS: &'static str = "delete_volume_params"; @@ -92,6 +97,45 @@ impl NexusSaga for SagaImageDelete { // image delete saga: action implementations +async fn sid_account_space( + sagactx: NexusActionContext, +) -> Result<(), ActionError> { + let osagactx = sagactx.user_data(); + let params = sagactx.saga_params::()?; + let opctx = crate::context::op_context_for_saga_action( + &sagactx, + ¶ms.serialized_authn, + ); + + match ¶ms.image_param { + ImageParam::Project { image, .. } => { + osagactx + .datastore() + .virtual_provisioning_collection_delete_project_image( + &opctx, + image.id(), + image.project_id, + image.size, + ) + .await + .map_err(ActionError::action_failed)?; + } + ImageParam::Silo { image, .. } => { + osagactx + .datastore() + .virtual_provisioning_collection_delete_silo_image( + &opctx, + image.id(), + image.silo_id, + image.size, + ) + .await + .map_err(ActionError::action_failed)?; + } + } + Ok(()) +} + async fn sid_delete_image_record( sagactx: NexusActionContext, ) -> Result<(), ActionError> { diff --git a/nexus/tests/integration_tests/schema.rs b/nexus/tests/integration_tests/schema.rs index 0b16b3cfca3..9bad5a4b107 100644 --- a/nexus/tests/integration_tests/schema.rs +++ b/nexus/tests/integration_tests/schema.rs @@ -3781,6 +3781,206 @@ mod migration_211 { } } +mod migration_218 { + use super::*; + use pretty_assertions::assert_eq; + + // Randomly-generated test UUIDs + const TEST_SILO_ID: &str = "a1a1a1a1-0000-4000-8000-000000000001"; + const TEST_PROJECT_ID: &str = "b2b2b2b2-0000-4000-8000-000000000002"; + const TEST_PROJECT_IMAGE_ID: &str = "c3c3c3c3-0000-4000-8000-000000000003"; + const TEST_SILO_IMAGE_ID: &str = "d4d4d4d4-0000-4000-8000-000000000004"; + const TEST_VOLUME_ID_1: &str = "e5e5e5e5-0000-4000-8000-000000000005"; + const TEST_VOLUME_ID_2: &str = "f6f6f6f6-0000-4000-8000-000000000006"; + + // 1 GiB and 2 GiB in bytes + const PROJECT_IMAGE_SIZE: i64 = 1073741824; + const SILO_IMAGE_SIZE: i64 = 2147483648; + + async fn before_impl(ctx: &MigrationContext<'_>) { + // Insert test silo with quotas and virtual_provisioning_collection + ctx.client + .batch_execute(&format!( + " + INSERT INTO omicron.public.silo ( + id, name, description, time_created, time_modified, + discoverable, authentication_mode, user_provision_type, + mapped_fleet_roles, rcgen + ) VALUES ( + '{TEST_SILO_ID}', 'migration-218-silo', 'Test silo', + now(), now(), true, 'local', 'jit', '{{}}', 1 + ); + + INSERT INTO omicron.public.silo_quotas ( + silo_id, cpus, memory_bytes, storage_bytes, + time_created, time_modified + ) VALUES ( + '{TEST_SILO_ID}', 100, 1099511627776, 10995116277760, + now(), now() + ); + + INSERT INTO omicron.public.virtual_provisioning_collection ( + id, time_modified, collection_type, + virtual_disk_bytes_provisioned, cpus_provisioned, + ram_provisioned + ) VALUES ( + '{TEST_SILO_ID}', now(), 'Silo', 0, 0, 0 + ); + + INSERT INTO omicron.public.project ( + id, name, description, time_created, time_modified, silo_id, rcgen + ) VALUES ( + '{TEST_PROJECT_ID}', 'migration-218-project', 'Test project', + now(), now(), '{TEST_SILO_ID}', 1 + ); + + INSERT INTO omicron.public.virtual_provisioning_collection ( + id, time_modified, collection_type, + virtual_disk_bytes_provisioned, cpus_provisioned, + ram_provisioned + ) VALUES ( + '{TEST_PROJECT_ID}', now(), 'Project', 0, 0, 0 + ); + + INSERT INTO omicron.public.image ( + id, name, description, time_created, time_modified, + silo_id, project_id, volume_id, os, version, + block_size, size_bytes + ) VALUES ( + '{TEST_PROJECT_IMAGE_ID}', 'project-image', 'Test project image', + now(), now(), '{TEST_SILO_ID}', '{TEST_PROJECT_ID}', + '{TEST_VOLUME_ID_1}', 'linux', '1.0', '512', + {PROJECT_IMAGE_SIZE} + ); + + INSERT INTO omicron.public.image ( + id, name, description, time_created, time_modified, + silo_id, project_id, volume_id, os, version, + block_size, size_bytes + ) VALUES ( + '{TEST_SILO_IMAGE_ID}', 'silo-image', 'Test silo image', + now(), now(), '{TEST_SILO_ID}', NULL, + '{TEST_VOLUME_ID_2}', 'linux', '1.0', '512', + {SILO_IMAGE_SIZE} + ); + " + )) + .await + .expect("inserted pre-migration data"); + } + + async fn after_impl(ctx: &MigrationContext<'_>) { + let project_image_id: Uuid = TEST_PROJECT_IMAGE_ID.parse().unwrap(); + let silo_image_id: Uuid = TEST_SILO_IMAGE_ID.parse().unwrap(); + let project_id: Uuid = TEST_PROJECT_ID.parse().unwrap(); + let silo_id: Uuid = TEST_SILO_ID.parse().unwrap(); + + // Verify virtual_provisioning_resource entries were created for images + let resource_rows = ctx + .client + .query( + " + SELECT id, resource_type, virtual_disk_bytes_provisioned + FROM omicron.public.virtual_provisioning_resource + WHERE id IN ($1, $2) + ORDER BY id + ", + &[&project_image_id, &silo_image_id], + ) + .await + .expect("queried virtual_provisioning_resource"); + + assert_eq!( + resource_rows.len(), + 2, + "Expected 2 virtual_provisioning_resource entries for images" + ); + + for row in &resource_rows { + let resource_type: String = row.get("resource_type"); + assert_eq!( + resource_type, "image", + "Resource type should be 'image'" + ); + } + + // Verify project collection was updated (should have PROJECT_IMAGE_SIZE) + let project_row = ctx + .client + .query_one( + " + SELECT virtual_disk_bytes_provisioned + FROM omicron.public.virtual_provisioning_collection + WHERE id = $1 + ", + &[&project_id], + ) + .await + .expect("queried project collection"); + let project_bytes: i64 = project_row.get(0); + assert_eq!( + project_bytes, PROJECT_IMAGE_SIZE, + "Project should have {} bytes provisioned", + PROJECT_IMAGE_SIZE + ); + + // Verify silo collection was updated + // (should have PROJECT_IMAGE_SIZE + SILO_IMAGE_SIZE) + let silo_row = ctx + .client + .query_one( + " + SELECT virtual_disk_bytes_provisioned + FROM omicron.public.virtual_provisioning_collection + WHERE id = $1 + ", + &[&silo_id], + ) + .await + .expect("queried silo collection"); + let silo_bytes: i64 = silo_row.get(0); + let expected_silo_bytes = PROJECT_IMAGE_SIZE + SILO_IMAGE_SIZE; + assert_eq!( + silo_bytes, expected_silo_bytes, + "Silo should have {} bytes provisioned", + expected_silo_bytes + ); + + // Clean up test data to not affect other tests + ctx.client + .batch_execute(&format!( + " + DELETE FROM omicron.public.virtual_provisioning_resource + WHERE id IN ('{TEST_PROJECT_IMAGE_ID}', '{TEST_SILO_IMAGE_ID}'); + + DELETE FROM omicron.public.image + WHERE id IN ('{TEST_PROJECT_IMAGE_ID}', '{TEST_SILO_IMAGE_ID}'); + + DELETE FROM omicron.public.virtual_provisioning_collection + WHERE id IN ('{TEST_PROJECT_ID}', '{TEST_SILO_ID}'); + + DELETE FROM omicron.public.project WHERE id = '{TEST_PROJECT_ID}'; + DELETE FROM omicron.public.silo_quotas WHERE silo_id = '{TEST_SILO_ID}'; + DELETE FROM omicron.public.silo WHERE id = '{TEST_SILO_ID}'; + " + )) + .await + .expect("cleaned up test data"); + } + + pub(super) fn before<'a>( + ctx: &'a MigrationContext<'a>, + ) -> BoxFuture<'a, ()> { + Box::pin(before_impl(ctx)) + } + + pub(super) fn after<'a>( + ctx: &'a MigrationContext<'a>, + ) -> BoxFuture<'a, ()> { + Box::pin(after_impl(ctx)) + } +} + // Lazily initializes all migration checks. The combination of Rust function // pointers and async makes defining a static table fairly painful, so we're // using lazy initialization instead. @@ -3901,6 +4101,12 @@ fn get_migration_checks() -> BTreeMap { .before(migration_211::before) .after(migration_211::after), ); + map.insert( + Version::new(218, 0, 0), + DataMigrationFns::new() + .before(migration_218::before) + .after(migration_218::after), + ); map } diff --git a/schema/crdb/dbinit.sql b/schema/crdb/dbinit.sql index 29efbfcfb3b..6398fda1d67 100644 --- a/schema/crdb/dbinit.sql +++ b/schema/crdb/dbinit.sql @@ -7655,7 +7655,7 @@ INSERT INTO omicron.public.db_metadata ( version, target_version ) VALUES - (TRUE, NOW(), NOW(), '217.0.0', NULL) + (TRUE, NOW(), NOW(), '218.0.0', NULL) ON CONFLICT DO NOTHING; COMMIT; diff --git a/schema/crdb/image-virtual-provisioning/up.sql b/schema/crdb/image-virtual-provisioning/up.sql new file mode 100644 index 00000000000..960b463f38c --- /dev/null +++ b/schema/crdb/image-virtual-provisioning/up.sql @@ -0,0 +1,76 @@ +SET LOCAL disallow_full_table_scans = OFF; + +-- Insert virtual_provisioning_resource entries for all existing non-deleted images. +-- This tracks each image as a storage-consuming resource. +INSERT INTO omicron.public.virtual_provisioning_resource ( + id, + time_modified, + resource_type, + virtual_disk_bytes_provisioned, + cpus_provisioned, + ram_provisioned +) +SELECT + image.id, + NOW(), + 'image', + image.size_bytes, + 0, + 0 +FROM + omicron.public.image +WHERE + image.time_deleted IS NULL +ON CONFLICT (id) DO NOTHING; + +-- Update project-level collections for project-scoped images. +-- Project-scoped images have a non-null project_id. +UPDATE omicron.public.virtual_provisioning_collection AS vpc +SET + virtual_disk_bytes_provisioned = vpc.virtual_disk_bytes_provisioned + totals.total_size, + time_modified = NOW() +FROM ( + SELECT + project_id, + SUM(size_bytes) AS total_size + FROM + omicron.public.image + WHERE + time_deleted IS NULL AND project_id IS NOT NULL + GROUP BY + project_id +) AS totals +WHERE + vpc.id = totals.project_id; + +-- Update silo-level collections for all images (both project-scoped and silo-scoped). +-- All images belong to a silo, so all images roll up to their silo's collection. +UPDATE omicron.public.virtual_provisioning_collection AS vpc +SET + virtual_disk_bytes_provisioned = vpc.virtual_disk_bytes_provisioned + totals.total_size, + time_modified = NOW() +FROM ( + SELECT + silo_id, + SUM(size_bytes) AS total_size + FROM + omicron.public.image + WHERE + time_deleted IS NULL + GROUP BY + silo_id +) AS totals +WHERE + vpc.id = totals.silo_id; + +-- Update fleet-level collection with total image storage across all silos. +UPDATE omicron.public.virtual_provisioning_collection +SET + virtual_disk_bytes_provisioned = virtual_disk_bytes_provisioned + ( + SELECT COALESCE(SUM(size_bytes), 0) + FROM omicron.public.image + WHERE time_deleted IS NULL + ), + time_modified = NOW() +WHERE + collection_type = 'Fleet'; From 7c03998e17f9c2782efacef627408e0e24910f46 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Thu, 8 Jan 2026 14:31:45 -0800 Subject: [PATCH 2/6] expectorate --- .../virtual_provisioning_collection_update.rs | 112 +++++++++++++ ..._collection_update_delete_silo_storage.sql | 84 ++++++++++ ...collection_update_insert_project_image.sql | 153 ++++++++++++++++++ ..._collection_update_insert_silo_storage.sql | 118 ++++++++++++++ 4 files changed, 467 insertions(+) create mode 100644 nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_silo_storage.sql create mode 100644 nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_project_image.sql create mode 100644 nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_silo_storage.sql diff --git a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs index dc2ba391b19..a70da110943 100644 --- a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs +++ b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs @@ -842,6 +842,25 @@ mod test { ).await; } + #[tokio::test] + async fn expectorate_query_insert_project_image() { + let id = Uuid::nil(); + let project_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + let storage_type = crate::db::datastore::StorageType::Image; + + let query = VirtualProvisioningCollectionUpdate::new_insert_storage( + id, + disk_byte_diff, + project_id, + storage_type, + ); + expectorate_query_contents( + &query, + "tests/output/virtual_provisioning_collection_update_insert_project_image.sql", + ).await; + } + #[tokio::test] async fn expectorate_query_delete_storage() { let id = Uuid::nil(); @@ -860,6 +879,45 @@ mod test { ).await; } + #[tokio::test] + async fn expectorate_query_insert_silo_storage() { + let id = Uuid::nil(); + let silo_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + let storage_type = crate::db::datastore::StorageType::Image; + + let query = + VirtualProvisioningCollectionUpdate::new_insert_silo_storage( + id, + disk_byte_diff, + silo_id, + storage_type, + ); + expectorate_query_contents( + &query, + "tests/output/virtual_provisioning_collection_update_insert_silo_storage.sql", + ).await; + } + + #[tokio::test] + async fn expectorate_query_delete_silo_storage() { + let id = Uuid::nil(); + let silo_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + + let query = + VirtualProvisioningCollectionUpdate::new_delete_silo_storage( + id, + disk_byte_diff, + silo_id, + ); + + expectorate_query_contents( + &query, + "tests/output/virtual_provisioning_collection_update_delete_silo_storage.sql", + ).await; + } + #[tokio::test] async fn expectorate_query_insert_instance() { let id = InstanceUuid::nil(); @@ -949,6 +1007,60 @@ mod test { logctx.cleanup_successful(); } + #[tokio::test] + async fn explain_insert_silo_storage() { + let logctx = dev::test_setup_log("explain_insert_silo_storage"); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = pool.claim().await.unwrap(); + + let id = Uuid::nil(); + let silo_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + let storage_type = crate::db::datastore::StorageType::Image; + + let query = + VirtualProvisioningCollectionUpdate::new_insert_silo_storage( + id, + disk_byte_diff, + silo_id, + storage_type, + ); + let _ = query + .explain_async(&conn) + .await + .expect("Failed to explain query - is it valid SQL?"); + + db.terminate().await; + logctx.cleanup_successful(); + } + + #[tokio::test] + async fn explain_delete_silo_storage() { + let logctx = dev::test_setup_log("explain_delete_silo_storage"); + let db = TestDatabase::new_with_pool(&logctx.log).await; + let pool = db.pool(); + let conn = pool.claim().await.unwrap(); + + let id = Uuid::nil(); + let silo_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + + let query = + VirtualProvisioningCollectionUpdate::new_delete_silo_storage( + id, + disk_byte_diff, + silo_id, + ); + let _ = query + .explain_async(&conn) + .await + .expect("Failed to explain query - is it valid SQL?"); + + db.terminate().await; + logctx.cleanup_successful(); + } + #[tokio::test] async fn explain_insert_instance() { let logctx = dev::test_setup_log("explain_insert_instance"); diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_silo_storage.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_silo_storage.sql new file mode 100644 index 00000000000..97dfe39f8a5 --- /dev/null +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_silo_storage.sql @@ -0,0 +1,84 @@ +WITH + all_collections AS ((SELECT $1 AS id) UNION (SELECT $2 AS id)), + quotas + AS ( + SELECT + silo_quotas.silo_id, + silo_quotas.cpus, + silo_quotas.memory_bytes AS memory, + silo_quotas.storage_bytes AS storage + FROM + silo_quotas + WHERE + silo_quotas.silo_id = $3 + ), + silo_provisioned + AS ( + SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned, + virtual_provisioning_collection.virtual_disk_bytes_provisioned + FROM + virtual_provisioning_collection + WHERE + virtual_provisioning_collection.id = $4 + ), + do_update + AS ( + SELECT + ( + SELECT + count(*) + FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = $5 + LIMIT + 1 + ) + = 1 + AS update + ), + unused_cte_arm + AS ( + DELETE FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = $6 AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING + virtual_provisioning_resource.id, + virtual_provisioning_resource.time_modified, + virtual_provisioning_resource.resource_type, + virtual_provisioning_resource.virtual_disk_bytes_provisioned, + virtual_provisioning_resource.cpus_provisioned, + virtual_provisioning_resource.ram_provisioned + ), + virtual_provisioning_collection + AS ( + UPDATE + virtual_provisioning_collection + SET + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned - $7 + WHERE + virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) + AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING + virtual_provisioning_collection.id, + virtual_provisioning_collection.time_modified, + virtual_provisioning_collection.collection_type, + virtual_provisioning_collection.virtual_disk_bytes_provisioned, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned + ) +SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.time_modified, + virtual_provisioning_collection.collection_type, + virtual_provisioning_collection.virtual_disk_bytes_provisioned, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned +FROM + virtual_provisioning_collection diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_project_image.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_project_image.sql new file mode 100644 index 00000000000..87cd227ed9b --- /dev/null +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_project_image.sql @@ -0,0 +1,153 @@ +WITH + parent_silo AS (SELECT project.silo_id AS id FROM project WHERE project.id = $1), + all_collections + AS ( + ((SELECT $2 AS id) UNION (SELECT parent_silo.id AS id FROM parent_silo)) + UNION (SELECT $3 AS id) + ), + quotas + AS ( + SELECT + silo_quotas.silo_id, + silo_quotas.cpus, + silo_quotas.memory_bytes AS memory, + silo_quotas.storage_bytes AS storage + FROM + silo_quotas INNER JOIN parent_silo ON silo_quotas.silo_id = parent_silo.id + ), + silo_provisioned + AS ( + SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned, + virtual_provisioning_collection.virtual_disk_bytes_provisioned + FROM + virtual_provisioning_collection + INNER JOIN parent_silo ON virtual_provisioning_collection.id = parent_silo.id + ), + do_update + AS ( + SELECT + ( + ( + ( + SELECT + count(*) + FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = $4 + LIMIT + 1 + ) + = 0 + AND CAST( + IF( + ( + $5 = 0 + OR (SELECT quotas.cpus FROM quotas LIMIT 1) + >= ( + (SELECT silo_provisioned.cpus_provisioned FROM silo_provisioned LIMIT 1) + + $6 + ) + ), + 'TRUE', + 'Not enough cpus' + ) + AS BOOL + ) + ) + AND CAST( + IF( + ( + $7 = 0 + OR (SELECT quotas.memory FROM quotas LIMIT 1) + >= ( + (SELECT silo_provisioned.ram_provisioned FROM silo_provisioned LIMIT 1) + $8 + ) + ), + 'TRUE', + 'Not enough memory' + ) + AS BOOL + ) + ) + AND CAST( + IF( + ( + $9 = 0 + OR (SELECT quotas.storage FROM quotas LIMIT 1) + >= ( + ( + SELECT + silo_provisioned.virtual_disk_bytes_provisioned + FROM + silo_provisioned + LIMIT + 1 + ) + + $10 + ) + ), + 'TRUE', + 'Not enough storage' + ) + AS BOOL + ) + AS update + ), + unused_cte_arm + AS ( + INSERT + INTO + virtual_provisioning_resource + ( + id, + time_modified, + resource_type, + virtual_disk_bytes_provisioned, + cpus_provisioned, + ram_provisioned + ) + VALUES + ($11, DEFAULT, $12, $13, $14, $15) + ON CONFLICT + DO + NOTHING + RETURNING + virtual_provisioning_resource.id, + virtual_provisioning_resource.time_modified, + virtual_provisioning_resource.resource_type, + virtual_provisioning_resource.virtual_disk_bytes_provisioned, + virtual_provisioning_resource.cpus_provisioned, + virtual_provisioning_resource.ram_provisioned + ), + virtual_provisioning_collection + AS ( + UPDATE + virtual_provisioning_collection + SET + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned + $16 + WHERE + virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) + AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING + virtual_provisioning_collection.id, + virtual_provisioning_collection.time_modified, + virtual_provisioning_collection.collection_type, + virtual_provisioning_collection.virtual_disk_bytes_provisioned, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned + ) +SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.time_modified, + virtual_provisioning_collection.collection_type, + virtual_provisioning_collection.virtual_disk_bytes_provisioned, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned +FROM + virtual_provisioning_collection diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_silo_storage.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_silo_storage.sql new file mode 100644 index 00000000000..a9e8776ce1e --- /dev/null +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_silo_storage.sql @@ -0,0 +1,118 @@ +WITH + all_collections AS ((SELECT $1 AS id) UNION (SELECT $2 AS id)), + quotas + AS ( + SELECT + silo_quotas.silo_id, + silo_quotas.cpus, + silo_quotas.memory_bytes AS memory, + silo_quotas.storage_bytes AS storage + FROM + silo_quotas + WHERE + silo_quotas.silo_id = $3 + ), + silo_provisioned + AS ( + SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned, + virtual_provisioning_collection.virtual_disk_bytes_provisioned + FROM + virtual_provisioning_collection + WHERE + virtual_provisioning_collection.id = $4 + ), + do_update + AS ( + SELECT + ( + SELECT + count(*) + FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = $5 + LIMIT + 1 + ) + = 0 + AND CAST( + IF( + ( + $6 = 0 + OR (SELECT quotas.storage FROM quotas LIMIT 1) + >= ( + ( + SELECT + silo_provisioned.virtual_disk_bytes_provisioned + FROM + silo_provisioned + LIMIT + 1 + ) + + $7 + ) + ), + 'TRUE', + 'Not enough storage' + ) + AS BOOL + ) + AS update + ), + unused_cte_arm + AS ( + INSERT + INTO + virtual_provisioning_resource + ( + id, + time_modified, + resource_type, + virtual_disk_bytes_provisioned, + cpus_provisioned, + ram_provisioned + ) + VALUES + ($8, DEFAULT, $9, $10, $11, $12) + ON CONFLICT + DO + NOTHING + RETURNING + virtual_provisioning_resource.id, + virtual_provisioning_resource.time_modified, + virtual_provisioning_resource.resource_type, + virtual_provisioning_resource.virtual_disk_bytes_provisioned, + virtual_provisioning_resource.cpus_provisioned, + virtual_provisioning_resource.ram_provisioned + ), + virtual_provisioning_collection + AS ( + UPDATE + virtual_provisioning_collection + SET + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned + $13 + WHERE + virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) + AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING + virtual_provisioning_collection.id, + virtual_provisioning_collection.time_modified, + virtual_provisioning_collection.collection_type, + virtual_provisioning_collection.virtual_disk_bytes_provisioned, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned + ) +SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.time_modified, + virtual_provisioning_collection.collection_type, + virtual_provisioning_collection.virtual_disk_bytes_provisioned, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned +FROM + virtual_provisioning_collection From 7c47803ef8275f98af0f40803d8173c2c36240d8 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Thu, 8 Jan 2026 14:51:46 -0800 Subject: [PATCH 3/6] delete space after deleting record --- nexus/src/app/sagas/image_delete.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/nexus/src/app/sagas/image_delete.rs b/nexus/src/app/sagas/image_delete.rs index 1dafdd07e04..b18c6f42f13 100644 --- a/nexus/src/app/sagas/image_delete.rs +++ b/nexus/src/app/sagas/image_delete.rs @@ -38,12 +38,12 @@ pub(crate) struct Params { declare_saga_actions! { image_delete; - SPACE_ACCOUNT -> "no_result0" { - + sid_account_space - } - DELETE_IMAGE_RECORD -> "no_result1" { + DELETE_IMAGE_RECORD -> "no_result0" { + sid_delete_image_record } + SPACE_ACCOUNT -> "no_result1" { + + sid_account_space + } } #[derive(Debug)] @@ -60,8 +60,8 @@ impl NexusSaga for SagaImageDelete { params: &Self::Params, mut builder: steno::DagBuilder, ) -> Result { - builder.append(space_account_action()); builder.append(delete_image_record_action()); + builder.append(space_account_action()); const DELETE_VOLUME_PARAMS: &'static str = "delete_volume_params"; From fd42157a772b7c80efa9af4368f5a5d1a7b09f9d Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Thu, 8 Jan 2026 15:08:34 -0800 Subject: [PATCH 4/6] comment --- nexus/src/app/sagas/image_create.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nexus/src/app/sagas/image_create.rs b/nexus/src/app/sagas/image_create.rs index 71e70e376a1..95a398877d1 100644 --- a/nexus/src/app/sagas/image_create.rs +++ b/nexus/src/app/sagas/image_create.rs @@ -327,9 +327,9 @@ async fn simc_create_image_record_undo( let image_id = sagactx.lookup::("image_id")?; - // Use fetch_optional to make this undo idempotent - if the image was - // already deleted (e.g., by a previous undo attempt), we can safely - // skip the delete operation. + // Make this undo idempotent by checking if the image still exists. + // If it was already deleted (e.g., by a previous undo attempt), we + // can safely skip the delete operation. match ¶ms.image_type { ImageType::Project { .. } => { let lookup_result = LookupPath::new(&opctx, osagactx.datastore()) From 55cc5d13b4bb87d24467426227d8bc3fcfa9b9d0 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Fri, 9 Jan 2026 09:15:04 -0800 Subject: [PATCH 5/6] e2e test --- nexus/tests/integration_tests/images.rs | 390 ++++++++++++++++++++++++ 1 file changed, 390 insertions(+) diff --git a/nexus/tests/integration_tests/images.rs b/nexus/tests/integration_tests/images.rs index e9f9bfad67d..44f7de4fe02 100644 --- a/nexus/tests/integration_tests/images.rs +++ b/nexus/tests/integration_tests/images.rs @@ -7,6 +7,8 @@ use dropshot::ResultsPage; use http::StatusCode; use http::method::Method; +use nexus_db_queries::context::OpContext; +use nexus_db_queries::db::fixed_data::FLEET_ID; use nexus_db_queries::db::fixed_data::silo::DEFAULT_SILO; use nexus_db_queries::db::fixed_data::silo_user::USER_TEST_UNPRIVILEGED; use nexus_test_utils::http_testing::AuthnMode; @@ -21,6 +23,7 @@ use nexus_types::external_api::shared::SiloRole; use nexus_types::external_api::{params, views}; use nexus_types::identity::Asset; use nexus_types::identity::Resource; +use nexus_types::silo::DEFAULT_SILO_ID; use omicron_common::api::external::Disk; use omicron_common::api::external::{ByteCount, IdentityMetadataCreateParams}; @@ -585,3 +588,390 @@ async fn test_image_deletion_permissions(cptestctx: &ControlPlaneTestContext) { .await .expect("should be able to delete project image as unpriv user!"); } + +/// Test that creating and deleting images correctly updates virtual +/// provisioning at project, silo, and fleet levels. +/// +/// This test creates images from snapshots of blank disks. Since disks, +/// snapshots, and images all contribute to virtual_disk_bytes_provisioned, +/// we track the expected total at each step. +#[nexus_test] +async fn test_image_virtual_provisioning_collection( + cptestctx: &ControlPlaneTestContext, +) { + let client = &cptestctx.external_client; + let nexus = &cptestctx.server.server_context().nexus; + let datastore = nexus.datastore(); + + DiskTest::new(&cptestctx).await; + + let project = create_project(client, PROJECT_NAME).await; + let project_id = project.identity.id; + + let opctx = + OpContext::for_tests(cptestctx.logctx.log.new(o!()), datastore.clone()); + + // Use 1 GiB for all resources (disk, snapshot, image) + let one_gib = ByteCount::from_gibibytes_u32(1); + + // Helper to get provisioning for a collection + async fn get_provisioned_bytes( + datastore: &nexus_db_queries::db::DataStore, + opctx: &OpContext, + id: uuid::Uuid, + ) -> u64 { + datastore + .virtual_provisioning_collection_get(opctx, id) + .await + .unwrap() + .virtual_disk_bytes_provisioned + .to_bytes() + } + + // Initially, all collections should have 0 virtual_disk_bytes_provisioned + assert_eq!(get_provisioned_bytes(&datastore, &opctx, project_id).await, 0); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + 0 + ); + assert_eq!(get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, 0); + + // ========================================================================= + // Create resources for project-scoped image: disk1 -> snapshot1 -> image1 + // ========================================================================= + + // Create disk1 (1 GiB blank disk) + let disks_url = format!("/v1/disks?project={}", PROJECT_NAME); + let _disk1 = NexusRequest::objects_post( + client, + &disks_url, + ¶ms::DiskCreate { + identity: IdentityMetadataCreateParams { + name: "disk1".parse().unwrap(), + description: String::from("first test disk"), + }, + disk_backend: params::DiskBackend::Distributed { + disk_source: params::DiskSource::Blank { + block_size: params::BlockSize::try_from(512).unwrap(), + }, + }, + size: one_gib, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // After disk1: project/silo/fleet all have 1 GiB + let after_disk1 = one_gib.to_bytes(); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_disk1 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_disk1 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_disk1 + ); + + // Create snapshot1 from disk1 + let snapshots_url = format!("/v1/snapshots?project={}", PROJECT_NAME); + let snapshot1: views::Snapshot = NexusRequest::objects_post( + client, + &snapshots_url, + ¶ms::SnapshotCreate { + identity: IdentityMetadataCreateParams { + name: "snapshot1".parse().unwrap(), + description: String::from("first test snapshot"), + }, + disk: "disk1".parse().unwrap(), + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap() + .await; + + // After snapshot1: disk1 + snapshot1 = 2 GiB each + let after_snapshot1 = one_gib.to_bytes() * 2; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_snapshot1 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_snapshot1 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_snapshot1 + ); + + // Create project-scoped image from snapshot1 + let project_images_url = get_project_images_url(PROJECT_NAME); + let project_image = NexusRequest::objects_post( + client, + &project_images_url, + ¶ms::ImageCreate { + identity: IdentityMetadataCreateParams { + name: "project-image".parse().unwrap(), + description: String::from("a project-scoped image"), + }, + os: "testOS".to_string(), + version: "1.0".to_string(), + source: params::ImageSource::Snapshot { id: snapshot1.identity.id }, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // After project_image: disk1 + snapshot1 + image1 = 3 GiB each + let after_project_image = one_gib.to_bytes() * 3; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_project_image + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_project_image + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_project_image + ); + + // ========================================================================= + // Create resources for silo-scoped image: disk2 -> snapshot2 -> image2 + // ========================================================================= + + // Create disk2 (another 1 GiB blank disk) + let _disk2 = NexusRequest::objects_post( + client, + &disks_url, + ¶ms::DiskCreate { + identity: IdentityMetadataCreateParams { + name: "disk2".parse().unwrap(), + description: String::from("second test disk"), + }, + disk_backend: params::DiskBackend::Distributed { + disk_source: params::DiskSource::Blank { + block_size: params::BlockSize::try_from(512).unwrap(), + }, + }, + size: one_gib, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // After disk2: previous + disk2 = 4 GiB + let after_disk2 = one_gib.to_bytes() * 4; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_disk2 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_disk2 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_disk2 + ); + + // Create snapshot2 from disk2 + let snapshot2: views::Snapshot = NexusRequest::objects_post( + client, + &snapshots_url, + ¶ms::SnapshotCreate { + identity: IdentityMetadataCreateParams { + name: "snapshot2".parse().unwrap(), + description: String::from("second test snapshot"), + }, + disk: "disk2".parse().unwrap(), + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap() + .await; + + // After snapshot2: previous + snapshot2 = 5 GiB + let after_snapshot2 = one_gib.to_bytes() * 5; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_snapshot2 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_snapshot2 + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_snapshot2 + ); + + // Create silo-scoped image from snapshot2 + // Note: silo images don't add to project provisioning, only silo and fleet + let silo_images_url = "/v1/images"; + let silo_image = NexusRequest::objects_post( + client, + &silo_images_url, + ¶ms::ImageCreate { + identity: IdentityMetadataCreateParams { + name: "silo-image".parse().unwrap(), + description: String::from("a silo-scoped image"), + }, + os: "testOS".to_string(), + version: "1.0".to_string(), + source: params::ImageSource::Snapshot { id: snapshot2.identity.id }, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // After silo_image: + // - Project still has 5 GiB (disk1 + snapshot1 + image1 + disk2 + snapshot2) + // - Silo/Fleet have 6 GiB (above + silo_image) + let project_after_silo_image = one_gib.to_bytes() * 5; + let silo_after_silo_image = one_gib.to_bytes() * 6; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + project_after_silo_image, + "Project should not include silo-scoped image" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + silo_after_silo_image, + "Silo should include silo-scoped image" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + silo_after_silo_image, + "Fleet should include silo-scoped image" + ); + + // ========================================================================= + // Delete images and verify provisioning decreases correctly + // ========================================================================= + + // Delete the project-scoped image + let project_image_url = format!("/v1/images/{}", project_image.identity.id); + NexusRequest::new( + RequestBuilder::new(client, Method::DELETE, &project_image_url) + .expect_status(Some(StatusCode::NO_CONTENT)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete project image"); + + // After deleting project_image: + // - Project: disk1 + snapshot1 + disk2 + snapshot2 = 4 GiB + // - Silo/Fleet: above + silo_image = 5 GiB + let project_after_delete_project_image = one_gib.to_bytes() * 4; + let silo_after_delete_project_image = one_gib.to_bytes() * 5; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + project_after_delete_project_image + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + silo_after_delete_project_image + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + silo_after_delete_project_image + ); + + // Delete the silo-scoped image + let silo_image_url = format!("/v1/images/{}", silo_image.identity.id); + NexusRequest::new( + RequestBuilder::new(client, Method::DELETE, &silo_image_url) + .expect_status(Some(StatusCode::NO_CONTENT)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete silo image"); + + // After deleting silo_image: + // - All collections: disk1 + snapshot1 + disk2 + snapshot2 = 4 GiB + let after_delete_silo_image = one_gib.to_bytes() * 4; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_delete_silo_image + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_delete_silo_image + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_delete_silo_image + ); + + // ========================================================================= + // Clean up disks and snapshots + // ========================================================================= + + // Delete snapshots first (they depend on disks conceptually, but can be + // deleted independently) + let snapshot1_url = + format!("/v1/snapshots/snapshot1?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &snapshot1_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete snapshot1"); + + let snapshot2_url = + format!("/v1/snapshots/snapshot2?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &snapshot2_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete snapshot2"); + + // After deleting snapshots: disk1 + disk2 = 2 GiB + let after_delete_snapshots = one_gib.to_bytes() * 2; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_delete_snapshots + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_delete_snapshots + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_delete_snapshots + ); + + // Delete disks + let disk1_url = format!("/v1/disks/disk1?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &disk1_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete disk1"); + + let disk2_url = format!("/v1/disks/disk2?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &disk2_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete disk2"); + + // After deleting disks: all collections should be 0 + assert_eq!(get_provisioned_bytes(&datastore, &opctx, project_id).await, 0); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + 0 + ); + assert_eq!(get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, 0); +} From 5b93d80efca9f6dcf1b2be490b3cbfcefa7af19b Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Mon, 12 Jan 2026 14:58:08 -0800 Subject: [PATCH 6/6] feedback --- nexus/db-queries/src/db/datastore/image.rs | 234 +++++++++-- nexus/src/app/sagas/image_create.rs | 38 +- nexus/tests/integration_tests/images.rs | 444 +++++++++++++++++++++ 3 files changed, 661 insertions(+), 55 deletions(-) diff --git a/nexus/db-queries/src/db/datastore/image.rs b/nexus/db-queries/src/db/datastore/image.rs index b48bcf2912b..1edf9827173 100644 --- a/nexus/db-queries/src/db/datastore/image.rs +++ b/nexus/db-queries/src/db/datastore/image.rs @@ -3,16 +3,20 @@ use crate::authz::ApiResource; use crate::context::OpContext; use crate::db::collection_insert::AsyncInsertError; use crate::db::collection_insert::DatastoreCollection; +use crate::db::datastore::StorageType; use crate::db::model::Image; use crate::db::model::Project; use crate::db::model::ProjectImage; use crate::db::model::Silo; use crate::db::model::SiloImage; +use crate::db::model::VirtualProvisioningCollection; use crate::db::pagination::paginated; +use crate::db::queries::virtual_provisioning_collection_update::VirtualProvisioningCollectionUpdate; use async_bb8_diesel::AsyncRunQueryDsl; use chrono::Utc; use diesel::prelude::*; use nexus_db_errors::ErrorHandler; +use nexus_db_errors::OptionalError; use nexus_db_errors::public_error_from_diesel; use nexus_db_model::Name; use nexus_types::identity::Resource; @@ -166,6 +170,13 @@ impl DataStore { Ok(image) } + /// Promotes a project image to a silo image. + /// + /// This operation: + /// 1. Updates the image record to remove the project association + /// 2. Updates virtual provisioning to move accounting from project to silo + /// + /// All operations are performed in a transaction to ensure atomicity. pub async fn project_image_promote( &self, opctx: &OpContext, @@ -176,30 +187,101 @@ impl DataStore { opctx.authorize(authz::Action::CreateChild, authz_silo).await?; opctx.authorize(authz::Action::Modify, authz_project_image).await?; - use nexus_db_schema::schema::image::dsl; - let image = diesel::update(dsl::image) - .filter(dsl::time_deleted.is_null()) - .filter(dsl::id.eq(authz_project_image.id())) - .set(( - dsl::project_id.eq(None::), - dsl::time_modified.eq(Utc::now()), - )) - .returning(Image::as_returning()) - .get_result_async(&*self.pool_connection_authorized(opctx).await?) + let conn = self.pool_connection_authorized(opctx).await?; + let err = OptionalError::new(); + + let image_id = authz_project_image.id(); + let project_id = project_image.project_id; + let silo_id = authz_silo.id(); + let size = project_image.size; + let image_name = project_image.name().as_str().to_string(); + + self.transaction_retry_wrapper("project_image_promote") + .transaction(&conn, |conn| { + let err = err.clone(); + let image_name = image_name.clone(); + async move { + // Delete project-level accounting. + // + // This removes the virtual_provisioning_resource entry and + // decrements project, silo, and fleet collections. + VirtualProvisioningCollectionUpdate::new_delete_storage( + image_id, + size, + project_id, + ) + .get_results_async::(&conn) + .await + .map_err(|e| { + err.bail_retryable_or_else(e, |e| { + crate::db::queries::virtual_provisioning_collection_update::from_diesel(e) + }) + })?; + + // Insert silo-level accounting. + // + // This creates a new virtual_provisioning_resource entry and + // increments silo and fleet collections. + // Net effect: project decremented, silo and fleet unchanged. + VirtualProvisioningCollectionUpdate::new_insert_silo_storage( + image_id, + size, + silo_id, + StorageType::Image, + ) + .get_results_async::(&conn) + .await + .map_err(|e| { + err.bail_retryable_or_else(e, |e| { + crate::db::queries::virtual_provisioning_collection_update::from_diesel(e) + }) + })?; + + // Update the image record to remove project association. + use nexus_db_schema::schema::image::dsl; + let image = diesel::update(dsl::image) + .filter(dsl::time_deleted.is_null()) + .filter(dsl::id.eq(image_id)) + .set(( + dsl::project_id.eq(None::), + dsl::time_modified.eq(Utc::now()), + )) + .returning(Image::as_returning()) + .get_result_async(&conn) + .await + .map_err(|e| { + err.bail_retryable_or_else(e, |e| { + public_error_from_diesel( + e, + ErrorHandler::Conflict( + ResourceType::SiloImage, + image_name.as_str(), + ), + ) + }) + })?; + + Ok(image) + } + }) .await .map_err(|e| { - public_error_from_diesel( - e, - ErrorHandler::Conflict( - ResourceType::SiloImage, - project_image.name().as_str(), - ), - ) - })?; - - Ok(image) + if let Some(err) = err.take() { + return err; + } + public_error_from_diesel(e, ErrorHandler::Server) + }) } + /// Demotes a silo image to a project image. + /// + /// This operation: + /// 1. Updates the image record to add the project association + /// 2. Updates virtual provisioning to move accounting from silo to project + /// + /// All operations are performed in a transaction to ensure atomicity. + /// If the project's quota would be exceeded, the operation fails and + /// rolls back. pub async fn silo_image_demote( &self, opctx: &OpContext, @@ -210,29 +292,98 @@ impl DataStore { opctx.authorize(authz::Action::Modify, authz_silo_image).await?; opctx.authorize(authz::Action::CreateChild, authz_project).await?; - use nexus_db_schema::schema::image::dsl; - let image: Image = diesel::update(dsl::image) - .filter(dsl::time_deleted.is_null()) - .filter(dsl::id.eq(authz_silo_image.id())) - .set(( - dsl::project_id.eq(Some(authz_project.id())), - dsl::time_modified.eq(Utc::now()), - )) - .returning(Image::as_returning()) - .get_result_async(&*self.pool_connection_authorized(opctx).await?) + let conn = self.pool_connection_authorized(opctx).await?; + let err = OptionalError::new(); + + let image_id = authz_silo_image.id(); + let project_id = authz_project.id(); + let silo_id = silo_image.silo_id; + let size = silo_image.size; + let image_name = silo_image.name().as_str().to_string(); + + self.transaction_retry_wrapper("silo_image_demote") + .transaction(&conn, |conn| { + let err = err.clone(); + let image_name = image_name.clone(); + async move { + // Delete silo-level accounting. + // + // This removes the virtual_provisioning_resource entry and + // decrements silo and fleet collections. + VirtualProvisioningCollectionUpdate::new_delete_silo_storage( + image_id, + size, + silo_id, + ) + .get_results_async::(&conn) + .await + .map_err(|e| { + err.bail_retryable_or_else(e, |e| { + crate::db::queries::virtual_provisioning_collection_update::from_diesel(e) + }) + })?; + + // Insert project-level accounting. + // + // This creates a new virtual_provisioning_resource entry and + // increments project, silo, and fleet collections. + // Net effect: project incremented, silo and fleet unchanged. + // This will fail if the silo quota would be exceeded. + VirtualProvisioningCollectionUpdate::new_insert_storage( + image_id, + size, + project_id, + StorageType::Image, + ) + .get_results_async::(&conn) + .await + .map_err(|e| { + err.bail_retryable_or_else(e, |e| { + crate::db::queries::virtual_provisioning_collection_update::from_diesel(e) + }) + })?; + + // Update the image record to add project association. + use nexus_db_schema::schema::image::dsl; + let image: Image = diesel::update(dsl::image) + .filter(dsl::time_deleted.is_null()) + .filter(dsl::id.eq(image_id)) + .set(( + dsl::project_id.eq(Some(project_id)), + dsl::time_modified.eq(Utc::now()), + )) + .returning(Image::as_returning()) + .get_result_async(&conn) + .await + .map_err(|e| { + err.bail_retryable_or_else(e, |e| { + public_error_from_diesel( + e, + ErrorHandler::Conflict( + ResourceType::ProjectImage, + image_name.as_str(), + ), + ) + }) + })?; + + Ok(image) + } + }) .await .map_err(|e| { - public_error_from_diesel( - e, - ErrorHandler::Conflict( - ResourceType::ProjectImage, - silo_image.name().as_str(), - ), - ) - })?; - Ok(image) + if let Some(err) = err.take() { + return err; + } + public_error_from_diesel(e, ErrorHandler::Server) + }) } + /// Deletes a silo image record. + /// + /// Note: Does not update the corresponding accounting for space used by the + /// image. That's the responsibility of the caller (the image deletion + /// saga). pub async fn silo_image_delete( &self, opctx: &OpContext, @@ -243,6 +394,11 @@ impl DataStore { self.image_delete(opctx, image.into()).await } + /// Deletes a project image record. + /// + /// Note: Does not update the corresponding accounting for space used by the + /// image. That's the responsibility of the caller (the image deletion + /// saga). pub async fn project_image_delete( &self, opctx: &OpContext, diff --git a/nexus/src/app/sagas/image_create.rs b/nexus/src/app/sagas/image_create.rs index 95a398877d1..37077ad54dd 100644 --- a/nexus/src/app/sagas/image_create.rs +++ b/nexus/src/app/sagas/image_create.rs @@ -332,30 +332,36 @@ async fn simc_create_image_record_undo( // can safely skip the delete operation. match ¶ms.image_type { ImageType::Project { .. } => { - let lookup_result = LookupPath::new(&opctx, osagactx.datastore()) + match LookupPath::new(&opctx, osagactx.datastore()) .project_image_id(image_id) .fetch() - .await; - - if let Ok((.., authz_image, db_image)) = lookup_result { - osagactx - .datastore() - .project_image_delete(&opctx, &authz_image, db_image) - .await?; + .await + { + Ok((.., authz_image, db_image)) => { + osagactx + .datastore() + .project_image_delete(&opctx, &authz_image, db_image) + .await?; + } + Err(Error::ObjectNotFound { .. }) => {} + Err(e) => return Err(e.into()), } } ImageType::Silo { .. } => { - let lookup_result = LookupPath::new(&opctx, osagactx.datastore()) + match LookupPath::new(&opctx, osagactx.datastore()) .silo_image_id(image_id) .fetch() - .await; - - if let Ok((.., authz_image, db_image)) = lookup_result { - osagactx - .datastore() - .silo_image_delete(&opctx, &authz_image, db_image) - .await?; + .await + { + Ok((.., authz_image, db_image)) => { + osagactx + .datastore() + .silo_image_delete(&opctx, &authz_image, db_image) + .await?; + } + Err(Error::ObjectNotFound { .. }) => {} + Err(e) => return Err(e.into()), } } } diff --git a/nexus/tests/integration_tests/images.rs b/nexus/tests/integration_tests/images.rs index 44f7de4fe02..8da55303a49 100644 --- a/nexus/tests/integration_tests/images.rs +++ b/nexus/tests/integration_tests/images.rs @@ -975,3 +975,447 @@ async fn test_image_virtual_provisioning_collection( ); assert_eq!(get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, 0); } + +/// Test that promoting and demoting images correctly updates virtual +/// provisioning accounting. +/// +/// When a project image is promoted to a silo image: +/// - Project accounting should decrease by the image size +/// - Silo and Fleet accounting should stay the same +/// +/// When a silo image is demoted to a project image: +/// - Project accounting should increase by the image size +/// - Silo and Fleet accounting should stay the same +#[nexus_test] +async fn test_image_promote_demote_virtual_provisioning( + cptestctx: &ControlPlaneTestContext, +) { + let client = &cptestctx.external_client; + let nexus = &cptestctx.server.server_context().nexus; + let datastore = nexus.datastore(); + + DiskTest::new(&cptestctx).await; + + let project = create_project(client, PROJECT_NAME).await; + let project_id = project.identity.id; + + let opctx = + OpContext::for_tests(cptestctx.logctx.log.new(o!()), datastore.clone()); + + let one_gib = ByteCount::from_gibibytes_u32(1); + + // Helper to get provisioning for a collection + async fn get_provisioned_bytes( + datastore: &nexus_db_queries::db::DataStore, + opctx: &OpContext, + id: uuid::Uuid, + ) -> u64 { + datastore + .virtual_provisioning_collection_get(opctx, id) + .await + .unwrap() + .virtual_disk_bytes_provisioned + .to_bytes() + } + + // ========================================================================= + // Create a project image: disk -> snapshot -> image + // ========================================================================= + + // Create disk + let disks_url = format!("/v1/disks?project={}", PROJECT_NAME); + let _disk = NexusRequest::objects_post( + client, + &disks_url, + ¶ms::DiskCreate { + identity: IdentityMetadataCreateParams { + name: "test-disk".parse().unwrap(), + description: String::from("test disk"), + }, + disk_backend: params::DiskBackend::Distributed { + disk_source: params::DiskSource::Blank { + block_size: params::BlockSize::try_from(512).unwrap(), + }, + }, + size: one_gib, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // Create snapshot + let snapshots_url = format!("/v1/snapshots?project={}", PROJECT_NAME); + let snapshot: views::Snapshot = NexusRequest::objects_post( + client, + &snapshots_url, + ¶ms::SnapshotCreate { + identity: IdentityMetadataCreateParams { + name: "test-snapshot".parse().unwrap(), + description: String::from("test snapshot"), + }, + disk: "test-disk".parse().unwrap(), + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap() + .await; + + // Create project image from snapshot + let project_images_url = get_project_images_url(PROJECT_NAME); + let image = NexusRequest::objects_post( + client, + &project_images_url, + ¶ms::ImageCreate { + identity: IdentityMetadataCreateParams { + name: "test-image".parse().unwrap(), + description: String::from("test image for promote/demote"), + }, + os: "testOS".to_string(), + version: "1.0".to_string(), + source: params::ImageSource::Snapshot { id: snapshot.identity.id }, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + let image_id = image.identity.id; + + // At this point we have disk (1 GiB) + snapshot (1 GiB) + image (1 GiB) = 3 GiB + // All in the project, so project/silo/fleet all have 3 GiB + let initial_bytes = one_gib.to_bytes() * 3; + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + initial_bytes, + "Project should have 3 GiB initially" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + initial_bytes, + "Silo should have 3 GiB initially" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + initial_bytes, + "Fleet should have 3 GiB initially" + ); + + // ========================================================================= + // Promote the image from project to silo + // ========================================================================= + + let promote_url = format!("/v1/images/{}/promote", image_id); + NexusRequest::new( + RequestBuilder::new(client, http::Method::POST, &promote_url) + .expect_status(Some(http::StatusCode::ACCEPTED)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // After promote: + // - Project: disk + snapshot = 2 GiB (image moved to silo) + // - Silo/Fleet: disk + snapshot + image = 3 GiB (unchanged) + let project_after_promote = one_gib.to_bytes() * 2; + let silo_after_promote = one_gib.to_bytes() * 3; + + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + project_after_promote, + "Project should decrease by image size after promote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + silo_after_promote, + "Silo should be unchanged after promote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + silo_after_promote, + "Fleet should be unchanged after promote" + ); + + // ========================================================================= + // Demote the image back from silo to project + // ========================================================================= + + let demote_url = + format!("/v1/images/{}/demote?project={}", image_id, PROJECT_NAME); + NexusRequest::new( + RequestBuilder::new(client, http::Method::POST, &demote_url) + .expect_status(Some(http::StatusCode::ACCEPTED)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + // After demote: + // - Project: disk + snapshot + image = 3 GiB (back to original) + // - Silo/Fleet: disk + snapshot + image = 3 GiB (unchanged) + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + initial_bytes, + "Project should increase by image size after demote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + initial_bytes, + "Silo should be unchanged after demote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + initial_bytes, + "Fleet should be unchanged after demote" + ); + + // ========================================================================= + // Clean up + // ========================================================================= + + // Delete image (it's a project image after demote, no project param needed) + let image_url = format!("/v1/images/{}", image_id); + NexusRequest::new( + RequestBuilder::new(client, Method::DELETE, &image_url) + .expect_status(Some(StatusCode::NO_CONTENT)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete image"); + + // Delete snapshot + let snapshot_url = + format!("/v1/snapshots/test-snapshot?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &snapshot_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete snapshot"); + + // Delete disk + let disk_url = format!("/v1/disks/test-disk?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &disk_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete disk"); + + // Verify everything is cleaned up + assert_eq!(get_provisioned_bytes(&datastore, &opctx, project_id).await, 0); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + 0 + ); + assert_eq!(get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, 0); +} + +/// Test that attempting to promote an already-promoted image (silo image) +/// returns an error and doesn't corrupt accounting. +/// +/// Similarly, test that demoting a project image returns an error. +#[nexus_test] +async fn test_image_double_promote_demote_fails_safely( + cptestctx: &ControlPlaneTestContext, +) { + let client = &cptestctx.external_client; + let nexus = &cptestctx.server.server_context().nexus; + let datastore = nexus.datastore(); + + DiskTest::new(&cptestctx).await; + + let project = create_project(client, PROJECT_NAME).await; + let project_id = project.identity.id; + + let opctx = + OpContext::for_tests(cptestctx.logctx.log.new(o!()), datastore.clone()); + + let one_gib = ByteCount::from_gibibytes_u32(1); + + async fn get_provisioned_bytes( + datastore: &nexus_db_queries::db::DataStore, + opctx: &OpContext, + id: uuid::Uuid, + ) -> u64 { + datastore + .virtual_provisioning_collection_get(opctx, id) + .await + .unwrap() + .virtual_disk_bytes_provisioned + .to_bytes() + } + + // Create disk -> snapshot -> project image + let disks_url = format!("/v1/disks?project={}", PROJECT_NAME); + let _disk = NexusRequest::objects_post( + client, + &disks_url, + ¶ms::DiskCreate { + identity: IdentityMetadataCreateParams { + name: "test-disk".parse().unwrap(), + description: String::from("test disk"), + }, + disk_backend: params::DiskBackend::Distributed { + disk_source: params::DiskSource::Blank { + block_size: params::BlockSize::try_from(512).unwrap(), + }, + }, + size: one_gib, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + let snapshots_url = format!("/v1/snapshots?project={}", PROJECT_NAME); + let snapshot: views::Snapshot = NexusRequest::objects_post( + client, + &snapshots_url, + ¶ms::SnapshotCreate { + identity: IdentityMetadataCreateParams { + name: "test-snapshot".parse().unwrap(), + description: String::from("test snapshot"), + }, + disk: "test-disk".parse().unwrap(), + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap() + .await; + + let project_images_url = get_project_images_url(PROJECT_NAME); + let image = NexusRequest::objects_post( + client, + &project_images_url, + ¶ms::ImageCreate { + identity: IdentityMetadataCreateParams { + name: "test-image".parse().unwrap(), + description: String::from("test image"), + }, + os: "testOS".to_string(), + version: "1.0".to_string(), + source: params::ImageSource::Snapshot { id: snapshot.identity.id }, + }, + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + let image_id = image.identity.id; + + // ========================================================================= + // Test: demoting a project image should fail + // ========================================================================= + + let initial_project = + get_provisioned_bytes(&datastore, &opctx, project_id).await; + let initial_silo = + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await; + let initial_fleet = + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await; + + let demote_url = + format!("/v1/images/{}/demote?project={}", image_id, PROJECT_NAME); + NexusRequest::new( + RequestBuilder::new(client, http::Method::POST, &demote_url) + .expect_status(Some(StatusCode::BAD_REQUEST)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("demoting a project image should fail with 400"); + + // Verify accounting is unchanged after failed demote + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + initial_project, + "Project accounting should be unchanged after failed demote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + initial_silo, + "Silo accounting should be unchanged after failed demote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + initial_fleet, + "Fleet accounting should be unchanged after failed demote" + ); + + // ========================================================================= + // Promote the image, then test double-promote + // ========================================================================= + + let promote_url = format!("/v1/images/{}/promote", image_id); + NexusRequest::new( + RequestBuilder::new(client, http::Method::POST, &promote_url) + .expect_status(Some(http::StatusCode::ACCEPTED)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute_and_parse_unwrap::() + .await; + + let after_promote_project = + get_provisioned_bytes(&datastore, &opctx, project_id).await; + let after_promote_silo = + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await; + let after_promote_fleet = + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await; + + // Trying to promote again should fail + NexusRequest::new( + RequestBuilder::new(client, http::Method::POST, &promote_url) + .expect_status(Some(StatusCode::BAD_REQUEST)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("promoting a silo image should fail with 400"); + + // Verify accounting is unchanged after failed double-promote + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, project_id).await, + after_promote_project, + "Project accounting should be unchanged after failed double-promote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, DEFAULT_SILO_ID).await, + after_promote_silo, + "Silo accounting should be unchanged after failed double-promote" + ); + assert_eq!( + get_provisioned_bytes(&datastore, &opctx, *FLEET_ID).await, + after_promote_fleet, + "Fleet accounting should be unchanged after failed double-promote" + ); + + // ========================================================================= + // Clean up + // ========================================================================= + + // Delete image (now a silo image) + let image_url = format!("/v1/images/{}", image_id); + NexusRequest::new( + RequestBuilder::new(client, Method::DELETE, &image_url) + .expect_status(Some(StatusCode::NO_CONTENT)), + ) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete image"); + + let snapshot_url = + format!("/v1/snapshots/test-snapshot?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &snapshot_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete snapshot"); + + let disk_url = format!("/v1/disks/test-disk?project={}", PROJECT_NAME); + NexusRequest::object_delete(client, &disk_url) + .authn_as(AuthnMode::PrivilegedUser) + .execute() + .await + .expect("failed to delete disk"); +}