Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return richer enum types from datastore functions #6604

Merged
merged 3 commits into from
Sep 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion nexus/db-model/src/region_replacement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ impl std::str::FromStr for RegionReplacementState {
/// | |
/// v |
/// |
/// Completed ---
/// Complete ---
/// ```
///
/// which are captured in the RegionReplacementState enum. Annotated on the
Expand Down
3 changes: 2 additions & 1 deletion nexus/db-queries/src/db/datastore/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ mod rack;
mod region;
mod region_replacement;
mod region_snapshot;
mod region_snapshot_replacement;
pub mod region_snapshot_replacement;
mod role;
mod saga;
mod silo;
Expand Down Expand Up @@ -132,6 +132,7 @@ pub use volume::CrucibleTargets;
pub use volume::ExistingTarget;
pub use volume::ReplacementTarget;
pub use volume::VolumeCheckoutReason;
pub use volume::VolumeReplaceResult;
pub use volume::VolumeReplacementParams;
pub use volume::VolumeToDelete;
pub use volume::VolumeWithTarget;
Expand Down
145 changes: 109 additions & 36 deletions nexus/db-queries/src/db/datastore/region_snapshot_replacement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,13 +23,23 @@ use crate::db::pagination::Paginator;
use crate::db::update_and_check::UpdateAndCheck;
use crate::db::update_and_check::UpdateStatus;
use crate::db::TransactionError;
use crate::transaction_retry::OptionalError;
use async_bb8_diesel::AsyncConnection;
use async_bb8_diesel::AsyncRunQueryDsl;
use diesel::prelude::*;
use omicron_common::api::external::Error;
use uuid::Uuid;

#[must_use]
#[derive(Debug, PartialEq, Eq)]
pub enum InsertStepResult {
/// A new region snapshot replacement step was inserted.
Inserted { step_id: Uuid },

/// A region snapshot replacement step exists already that references this
/// volume id, so no new record is inserted.
AlreadyHandled { existing_step_id: Uuid },
}

impl DataStore {
/// Create and insert a region snapshot replacement request for a
/// RegionSnapshot, returning the ID of the request.
Expand Down Expand Up @@ -614,28 +624,23 @@ impl DataStore {
opctx: &OpContext,
request_id: Uuid,
volume_id: Uuid,
) -> Result<Uuid, Error> {
) -> Result<InsertStepResult, Error> {
let request = RegionSnapshotReplacementStep::new(request_id, volume_id);
let request_id = request.id;

self.insert_region_snapshot_replacement_step(opctx, request).await?;

Ok(request_id)
self.insert_region_snapshot_replacement_step(opctx, request).await
}

pub async fn insert_region_snapshot_replacement_step(
&self,
opctx: &OpContext,
request: RegionSnapshotReplacementStep,
) -> Result<(), Error> {
) -> Result<InsertStepResult, Error> {
let conn = self.pool_connection_authorized(opctx).await?;

let err = OptionalError::new();
self.transaction_retry_wrapper(
"insert_region_snapshot_replacement_step",
)
.transaction(&conn, |conn| {
let err = err.clone();
let request = request.clone();

async move {
Expand All @@ -656,11 +661,37 @@ impl DataStore {
.optional()?;

if let Some(found_record) = maybe_record {
return Err(err.bail(Error::conflict(format!(
"{:?} already referenced in old snapshot volume for \
request {:?}",
request.volume_id, found_record.id,
))));
return Ok(InsertStepResult::AlreadyHandled {
existing_step_id: found_record.id,
});
}

// Skip inserting this record if we found an existing region
// snapshot replacement step for it in a non-complete state.

let maybe_record = dsl::region_snapshot_replacement_step
.filter(dsl::volume_id.eq(request.volume_id))
.get_result_async::<RegionSnapshotReplacementStep>(&conn)
.await
.optional()?;

if let Some(found_record) = maybe_record {
match found_record.replacement_state {
RegionSnapshotReplacementStepState::Complete
| RegionSnapshotReplacementStepState::VolumeDeleted => {
// Ok, we can insert another record with a matching
// volume ID because the volume_repair record would
// have been deleted during the transition to
// Complete.
}

RegionSnapshotReplacementStepState::Requested
| RegionSnapshotReplacementStepState::Running => {
return Ok(InsertStepResult::AlreadyHandled {
existing_step_id: found_record.id,
});
}
}
}

// The region snapshot replacement step saga could invoke a
Expand All @@ -675,22 +706,18 @@ impl DataStore {
.execute_async(&conn)
.await?;

let request_id = request.id;

diesel::insert_into(dsl::region_snapshot_replacement_step)
.values(request)
.execute_async(&conn)
.await?;

Ok(())
Ok(InsertStepResult::Inserted { step_id: request_id })
}
})
.await
.map_err(|e| {
if let Some(err) = err.take() {
return err;
}

public_error_from_diesel(e, ErrorHandler::Server)
})
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
}

pub async fn get_region_snapshot_replacement_step_by_id(
Expand Down Expand Up @@ -1173,10 +1200,12 @@ mod test {
Uuid::new_v4(), // volume id
);

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));
}

assert_eq!(
Expand Down Expand Up @@ -1207,10 +1236,12 @@ mod test {
step.replacement_state =
RegionSnapshotReplacementStepState::Running;

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));
}

assert_eq!(
Expand Down Expand Up @@ -1242,10 +1273,12 @@ mod test {
step.replacement_state =
RegionSnapshotReplacementStepState::VolumeDeleted;

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));
}

assert_eq!(
Expand Down Expand Up @@ -1288,18 +1321,26 @@ mod test {
RegionSnapshotReplacementStep::new(Uuid::new_v4(), volume_id);
let first_request_id = step.id;

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));

let step =
RegionSnapshotReplacementStep::new(Uuid::new_v4(), volume_id);

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step.clone())
.await
.unwrap_err();
.unwrap();

let InsertStepResult::AlreadyHandled { existing_step_id } = result
else {
panic!("wrong return type: {result:?}");
};
assert_eq!(existing_step_id, first_request_id);

// Ensure that transitioning the first step to running doesn't change
// things.
Expand All @@ -1315,10 +1356,16 @@ mod test {
.await
.unwrap();

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step.clone())
.await
.unwrap_err();
.unwrap();

let InsertStepResult::AlreadyHandled { existing_step_id } = result
else {
panic!("wrong return type: {result:?}");
};
assert_eq!(existing_step_id, first_request_id);

// Ensure that transitioning the first step to complete means another
// can be added.
Expand All @@ -1333,11 +1380,16 @@ mod test {
.await
.unwrap();

datastore
let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step.clone())
.await
.unwrap();

let InsertStepResult::Inserted { step_id } = result else {
panic!("wrong return type: {result:?}");
};
assert_eq!(step_id, step.id);

// Ensure that transitioning the first step to volume deleted still
// works.

Expand Down Expand Up @@ -1388,19 +1440,25 @@ mod test {
let mut step =
RegionSnapshotReplacementStep::new(request_id, Uuid::new_v4());
step.replacement_state = RegionSnapshotReplacementStepState::Complete;
datastore

let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));

let mut step =
RegionSnapshotReplacementStep::new(request_id, Uuid::new_v4());
step.replacement_state = RegionSnapshotReplacementStepState::Complete;
datastore

let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));

assert_eq!(
2,
datastore
Expand Down Expand Up @@ -1435,19 +1493,32 @@ mod test {
RegionSnapshotReplacementStep::new(request_id, volume_id);
step.replacement_state = RegionSnapshotReplacementStepState::Complete;
step.old_snapshot_volume_id = Some(old_snapshot_volume_id);
datastore

let first_step_id = step.id;

let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap();

assert!(matches!(result, InsertStepResult::Inserted { .. }));

let step = RegionSnapshotReplacementStep::new(
request_id,
old_snapshot_volume_id,
);
datastore

let result = datastore
.insert_region_snapshot_replacement_step(&opctx, step)
.await
.unwrap_err();
.unwrap();

assert_eq!(
result,
InsertStepResult::AlreadyHandled {
existing_step_id: first_step_id
}
);

db.cleanup().await.unwrap();
logctx.cleanup_successful();
Expand All @@ -1468,13 +1539,15 @@ mod test {
let volume_id = Uuid::new_v4();

let request = RegionReplacement::new(Uuid::new_v4(), volume_id);

datastore
.insert_region_replacement_request(&opctx, request)
.await
.unwrap();

let request =
RegionSnapshotReplacementStep::new(Uuid::new_v4(), volume_id);

datastore
.insert_region_snapshot_replacement_step(&opctx, request)
.await
Expand Down
Loading
Loading