From 10b1ec032ef3083f44b3ee89d484047d094dc615 Mon Sep 17 00:00:00 2001 From: Benjamin Naecker Date: Tue, 24 May 2022 04:27:13 +0000 Subject: [PATCH] Add test for the wrapping next-item query --- nexus/src/db/queries/next_item.rs | 203 ++++++++++++++++++++++++++++-- 1 file changed, 190 insertions(+), 13 deletions(-) diff --git a/nexus/src/db/queries/next_item.rs b/nexus/src/db/queries/next_item.rs index 788ac33b42..d7486e369d 100644 --- a/nexus/src/db/queries/next_item.rs +++ b/nexus/src/db/queries/next_item.rs @@ -72,15 +72,8 @@ use uuid::Uuid; /// existing item. That "retry loop" is converted into a sequential scan in the /// database for an item that does not conflict. /// -/// There are several other potential issues. One is that, we may wish to skip -/// certain reserved items. Second, and slightly more subtle, is that the size -/// of the search space is dependent on the starting point. If we randomly -/// select a base that's pretty close to the end of the search space, then the -/// search would be over a much smaller space than if the base is at the start -/// of the total range. -/// -/// In its most general form, the `NextItem` query accounts for both of these, -/// and looks like: +/// In its most general form, the `NextItem` query accounts for the latter of +/// these, and looks like: /// /// ```sql /// SELECT @@ -133,7 +126,10 @@ use uuid::Uuid; /// This concept of selecting a base and max/min shift is encapsulated in the /// [`ShiftGenerator`] trait. The simplest implementation, /// [`DefaultShiftGenerator`], follows the diagram above, search from `base`, to -/// `base + max_shift`, to `base - min-shift`, to `base - 1`. +/// `base + max_shift`, to `base - min-shift`, to `base - 1`. Other +/// `ShiftGenerator` implementations could provide more complex behavior, such +/// as a deny-list of precluded items, a non-contiguous search space, or other +/// behavior. /// /// Scopes /// ------ @@ -375,11 +371,13 @@ where /// of `self` called `inner`. macro_rules! delegate_query_fragment_impl { ($parent:ty) => { - impl QueryFragment for $parent { + impl ::diesel::query_builder::QueryFragment<::diesel::pg::Pg> + for $parent + { fn walk_ast<'a>( &'a self, - out: AstPass<'_, 'a, Pg>, - ) -> diesel::QueryResult<()> { + out: ::diesel::query_builder::AstPass<'_, 'a, ::diesel::pg::Pg>, + ) -> ::diesel::QueryResult<()> { self.inner.walk_ast(out) } } @@ -470,3 +468,182 @@ impl ShiftGenerator for DefaultShiftGenerator { &self.min_shift } } + +#[cfg(test)] +mod tests { + + use super::DefaultShiftGenerator; + use super::NextItem; + use crate::db; + use async_bb8_diesel::AsyncRunQueryDsl; + use async_bb8_diesel::AsyncSimpleConnection; + use chrono::DateTime; + use chrono::Utc; + use diesel::pg::Pg; + use diesel::query_builder::AstPass; + use diesel::query_builder::QueryFragment; + use diesel::query_builder::QueryId; + use diesel::Column; + use diesel::Insertable; + use diesel::SelectableHelper; + use nexus_test_utils::db::test_setup_database; + use omicron_test_utils::dev; + use std::sync::Arc; + use uuid::Uuid; + + table! { + test_schema.item (id) { + id -> Uuid, + value -> Int4, + time_deleted -> Nullable, + } + } + + async fn setup_test_schema(pool: &db::Pool) { + let connection = pool.pool().get().await.unwrap(); + (*connection) + .batch_execute_async( + "CREATE SCHEMA IF NOT EXISTS test_schema; \ + CREATE TABLE IF NOT EXISTS test_schema.item ( \ + id UUID PRIMARY KEY, \ + value INT4 NOT NULL, \ + time_deleted TIMESTAMPTZ \ + ); \ + TRUNCATE test_schema.item; \ + CREATE UNIQUE INDEX ON test_schema.item (value) WHERE time_deleted IS NULL; \ + ") + .await + .unwrap() + } + + // Describes an item to be allocated with a NextItem query + #[derive(Queryable, Debug, Insertable, Selectable, Clone)] + #[diesel(table_name = item)] + struct Item { + id: Uuid, + value: i32, + time_deleted: Option>, + } + + #[derive(Debug, Clone, Copy)] + struct NextItemQuery { + inner: NextItem, + } + + // These implementations are needed to actually allow inserting the results + // of the `NextItemQuery` itself. + impl NextItemQuery { + fn new(generator: DefaultShiftGenerator) -> Self { + Self { inner: NextItem::new_unscoped(generator) } + } + } + + delegate_query_fragment_impl!(NextItemQuery); + + impl QueryId for NextItemQuery { + type QueryId = (); + const HAS_STATIC_QUERY_ID: bool = false; + } + + impl Insertable for NextItemQuery { + type Values = NextItemQueryValues; + fn values(self) -> Self::Values { + NextItemQueryValues(self) + } + } + + #[derive(Debug, Clone)] + struct NextItemQueryValues(NextItemQuery); + + impl QueryId for NextItemQueryValues { + type QueryId = (); + const HAS_STATIC_QUERY_ID: bool = false; + } + + impl diesel::insertable::CanInsertInSingleQuery for NextItemQueryValues { + fn rows_to_insert(&self) -> Option { + Some(1) + } + } + + impl QueryFragment for NextItemQueryValues { + fn walk_ast<'a>( + &'a self, + mut out: AstPass<'_, 'a, Pg>, + ) -> diesel::QueryResult<()> { + out.push_sql("("); + out.push_identifier(item::dsl::id::NAME)?; + out.push_sql(", "); + out.push_identifier(item::dsl::value::NAME)?; + out.push_sql(", "); + out.push_identifier(item::dsl::time_deleted::NAME)?; + out.push_sql(") VALUES (gen_random_uuid(), ("); + self.0.walk_ast(out.reborrow())?; + out.push_sql("), NULL)"); + Ok(()) + } + } + + // Test that we correctly insert the next available item + #[tokio::test] + async fn test_wrapping_next_item_query() { + // Setup the test database + let logctx = dev::test_setup_log("test_wrapping_next_item_query"); + let log = logctx.log.new(o!()); + let mut db = test_setup_database(&log).await; + let cfg = crate::db::Config { url: db.pg_config().clone() }; + let pool = Arc::new(crate::db::Pool::new(&cfg)); + + // We're going to operate on a separate table, for simplicity. + setup_test_schema(&pool).await; + + // We'll first insert an item at 0. + // + // This generator should start at 0, and then select over the range [0, + // 10], wrapping back to 0. + let generator = + DefaultShiftGenerator { base: 0, max_shift: 10, min_shift: 0 }; + let query = NextItemQuery::new(generator); + let it = diesel::insert_into(item::dsl::item) + .values(query) + .returning(Item::as_returning()) + .get_result_async(pool.pool()) + .await + .unwrap(); + assert_eq!(it.value, 0); + + // Insert the same query again, which should give us 1 now. + let it = diesel::insert_into(item::dsl::item) + .values(query) + .returning(Item::as_returning()) + .get_result_async(pool.pool()) + .await + .unwrap(); + assert_eq!(it.value, 1); + + // Insert 10, and guarantee that we get it back. + let generator = + DefaultShiftGenerator { base: 10, max_shift: 0, min_shift: -10 }; + let query = NextItemQuery::new(generator); + let it = diesel::insert_into(item::dsl::item) + .values(query) + .returning(Item::as_returning()) + .get_result_async(pool.pool()) + .await + .unwrap(); + assert_eq!(it.value, 10); + + // Now, insert the same query again. Since 0, 1, and 10 are all + // allocated, we should wrap around and insert 2. + let it = diesel::insert_into(item::dsl::item) + .values(query) + .returning(Item::as_returning()) + .get_result_async(pool.pool()) + .await + .unwrap(); + assert_eq!(it.value, 2); + + db.cleanup().await.unwrap(); + logctx.cleanup_successful(); + } +}