feat: optimize postgres get_or_create_collection_id (#2028)

pjenvey · web-flow · commit 935b468cb077 · 2026-01-29T10:46:26.000-08:00
making it (and mysql's) match spanner's, also fixing/simplifying their create_collection test method and fix mysql/postgres transactions to begin first and foremost Closes STOR-465 Closes STOR-371
diff --git a/syncstorage-mysql/src/db/db_impl.rs b/syncstorage-mysql/src/db/db_impl.rs
@@ -39,6 +39,8 @@ impl Db for MysqlDb {
     /// than explicit locking, but our ops team have expressed concerns about
     /// the efficiency of that approach at scale.
     async fn lock_for_read(&mut self, params: params::LockCollection) -> DbResult<()> {
+        // Lock the db
+        self.begin(false).await?;
         let collection_id = self
             .get_collection_id(&params.collection)
             .await
@@ -60,8 +62,6 @@ impl Db for MysqlDb {
             return Ok(());
         }
 
-        // Lock the db
-        self.begin(false).await?;
         let modified = user_collections::table
             .select(user_collections::modified)
             .filter(user_collections::user_id.eq(user_id))
@@ -82,6 +82,8 @@ impl Db for MysqlDb {
     }
 
     async fn lock_for_write(&mut self, params: params::LockCollection) -> DbResult<()> {
+        // Lock the db
+        self.begin(true).await?;
         let collection_id = self.get_or_create_collection_id(&params.collection).await?;
         let user_id = params.user_id.legacy_id as i64;
         let key = (params.user_id, collection_id);
@@ -92,8 +94,6 @@ impl Db for MysqlDb {
             ));
         }
 
-        // Lock the db
-        self.begin(true).await?;
         let modified = user_collections::table
             .select(user_collections::modified)
             .filter(user_collections::user_id.eq(user_id))
@@ -754,7 +754,7 @@ impl Db for MysqlDb {
 
     #[cfg(debug_assertions)]
     async fn create_collection(&mut self, name: &str) -> Result<i32, Self::Error> {
-        self.get_or_create_collection_id(name).await
+        self._create_collection(name).await
     }
 
     #[cfg(debug_assertions)]
diff --git a/syncstorage-mysql/src/db/mod.rs b/syncstorage-mysql/src/db/mod.rs
@@ -8,14 +8,17 @@ use diesel::{
 };
 use diesel_async::RunQueryDsl;
 use syncserver_common::Metrics;
-use syncstorage_db_common::{results, util::SyncTimestamp, UserIdentifier};
+use syncstorage_db_common::{
+    error::DbErrorIntrospect, results, util::SyncTimestamp, Db, UserIdentifier,
+    FIRST_CUSTOM_COLLECTION_ID,
+};
 use syncstorage_settings::Quota;
 
 use crate::{
     pool::{CollectionCache, Conn},
     DbError, DbResult,
 };
-use schema::{bso, collections};
+use schema::{bso, collections, last_insert_id};
 
 mod batch_impl;
 mod db_impl;
@@ -110,45 +113,36 @@ impl MysqlDb {
     }
 
     pub(super) async fn get_or_create_collection_id(&mut self, name: &str) -> DbResult<i32> {
-        if let Some(id) = self.coll_cache.get_id(name)? {
-            return Ok(id);
+        match self.get_collection_id(name).await {
+            Err(e) if e.is_collection_not_found() => self._create_collection(name).await,
+            result => result,
         }
+    }
 
-        diesel::insert_or_ignore_into(collections::table)
+    async fn _create_collection(&mut self, name: &str) -> DbResult<i32> {
+        if !cfg!(debug_assertions) && !self.session.in_write_transaction {
+            return Err(DbError::internal(
+                "Can't escalate read-lock to write-lock".to_owned(),
+            ));
+        }
+        diesel::insert_into(collections::table)
             .values(collections::name.eq(name))
             .execute(&mut self.conn)
             .await?;
-
-        let id = collections::table
-            .select(collections::id)
-            .filter(collections::name.eq(name))
-            .first(&mut self.conn)
-            .await?;
-
-        if !self.session.in_write_transaction {
-            self.coll_cache.put(id, name.to_owned())?;
+        let collection_id = diesel::select(last_insert_id())
+            .first::<u64>(&mut self.conn)
+            .await?
+            .try_into()
+            .map_err(|e| {
+                DbError::internal(format!("Couldn't convert last_insert_id() to i32: {e}"))
+            })?;
+        if collection_id < FIRST_CUSTOM_COLLECTION_ID {
+            // DDL ensures this should never occur
+            return Err(DbError::internal(
+                "create_collection < {FIRST_CUSTOM_COLLECTION_ID}".to_owned(),
+            ));
         }
-
-        Ok(id)
-    }
-
-    async fn _get_collection_name(&mut self, id: i32) -> DbResult<String> {
-        let name = if let Some(name) = self.coll_cache.get_name(id)? {
-            name
-        } else {
-            sql_query(
-                "SELECT name
-                   FROM collections
-                  WHERE id = ?",
-            )
-            .bind::<Integer, _>(&id)
-            .get_result::<NameResult>(&mut self.conn)
-            .await
-            .optional()?
-            .ok_or_else(DbError::collection_not_found)?
-            .name
-        };
-        Ok(name)
+        Ok(collection_id)
     }
 
     async fn map_collection_names<T>(
diff --git a/syncstorage-mysql/src/db/schema.rs b/syncstorage-mysql/src/db/schema.rs
@@ -69,3 +69,8 @@ allow_tables_to_appear_in_same_query!(
     collections,
     user_collections,
 );
+
+define_sql_function! {
+    /// MySQL's LAST_INSERT_ID()
+    fn last_insert_id() -> Unsigned<BigInt>
+}
diff --git a/syncstorage-postgres/src/db/db_impl.rs b/syncstorage-postgres/src/db/db_impl.rs
@@ -72,6 +72,7 @@ impl Db for PgDb {
         &mut self,
         params: params::LockCollection,
     ) -> DbResult<results::LockCollection> {
+        self.begin(false).await?;
         let user_id = params.user_id.legacy_id as i64;
         let collection_id = self
             .get_collection_id(&params.collection)
@@ -95,8 +96,6 @@ impl Db for PgDb {
 
         // `FOR SHARE`
         // Obtains shared lock, allowing multiple transactions to read rows simultaneously.
-        self.begin(false).await?;
-
         let modified = user_collections::table
             .select(user_collections::modified)
             .filter(user_collections::user_id.eq(user_id))
@@ -116,6 +115,7 @@ impl Db for PgDb {
     }
 
     async fn lock_for_write(&mut self, params: params::LockCollection) -> DbResult<()> {
+        self.begin(true).await?;
         let user_id = params.user_id.legacy_id as i64;
         let collection_id = self.get_or_create_collection_id(&params.collection).await?;
         let key = (params.user_id, collection_id);
@@ -130,7 +130,6 @@ impl Db for PgDb {
         // Acquires exclusive lock on select rows, prohibits other transactions from modifying
         // until complete.
         let nowtz = now.into_sql::<Timestamptz>();
-        self.begin(true).await?;
         let row = user_collections::table
             .select((user_collections::modified, nowtz))
             .filter(user_collections::user_id.eq(user_id))
@@ -549,7 +548,7 @@ impl Db for PgDb {
 
     #[cfg(debug_assertions)]
     async fn create_collection(&mut self, name: &str) -> DbResult<i32> {
-        self.get_or_create_collection_id(name).await
+        self._create_collection(name).await
     }
 
     #[cfg(debug_assertions)]
diff --git a/syncstorage-postgres/src/db/mod.rs b/syncstorage-postgres/src/db/mod.rs
@@ -10,7 +10,8 @@ use std::{collections::HashMap, fmt, sync::Arc};
 
 use syncserver_common::Metrics;
 use syncstorage_db_common::{
-    diesel::DbError, params, results, util::SyncTimestamp, Db, UserIdentifier,
+    diesel::DbError, error::DbErrorIntrospect, params, results, util::SyncTimestamp, Db,
+    UserIdentifier, FIRST_CUSTOM_COLLECTION_ID,
 };
 use syncstorage_settings::Quota;
 
@@ -98,30 +99,31 @@ impl PgDb {
     /// Gets the provided collection by name and creates it if not present.
     /// Checks collection cache first to see if matching collection stored.
     /// Uses logic to not make change sif there is a conflict during insert.
-    pub(super) async fn get_or_create_collection_id(&mut self, name: &str) -> DbResult<i32> {
-        if let Some(id) = self.coll_cache.get_id(name)? {
-            return Ok(id);
+    async fn get_or_create_collection_id(&mut self, name: &str) -> DbResult<i32> {
+        match self.get_collection_id(name).await {
+            Err(e) if e.is_collection_not_found() => self._create_collection(name).await,
+            result => result,
         }
+    }
 
-        // Postgres specific ON CONFLICT DO NOTHING statement.
-        // https://docs.diesel.rs/2.0.x/diesel/query_builder/struct.InsertStatement.html#method.on_conflict_do_nothing
-        diesel::insert_into(collections::table)
+    async fn _create_collection(&mut self, name: &str) -> DbResult<i32> {
+        if !cfg!(debug_assertions) && !self.session.in_write_transaction {
+            return Err(DbError::internal(
+                "Can't escalate read-lock to write-lock".to_owned(),
+            ));
+        }
+        let collection_id = diesel::insert_into(collections::table)
             .values(collections::name.eq(name))
-            .on_conflict_do_nothing()
-            .execute(&mut self.conn)
-            .await?;
-
-        let id = collections::table
-            .select(collections::collection_id)
-            .filter(collections::name.eq(name))
-            .first(&mut self.conn)
+            .returning(collections::collection_id)
+            .get_result(&mut self.conn)
             .await?;
-
-        if !self.session.in_write_transaction {
-            self.coll_cache.put(id, name.to_owned())?;
+        if collection_id < FIRST_CUSTOM_COLLECTION_ID {
+            // DDL ensures this should never occur
+            return Err(DbError::internal(
+                "create_collection < {FIRST_CUSTOM_COLLECTION_ID}".to_owned(),
+            ));
         }
-
-        Ok(id)
+        Ok(collection_id)
     }
 
     /// Given a set of collection_ids, return a HashMap of collection_id's
diff --git a/syncstorage-spanner/src/db/mod.rs b/syncstorage-spanner/src/db/mod.rs
@@ -113,7 +113,7 @@ impl SpannerDb {
 
     async fn get_or_create_collection_id(&mut self, name: &str) -> DbResult<i32> {
         match self.get_collection_id(name).await {
-            Err(err) if err.is_collection_not_found() => self._create_collection(name).await,
+            Err(e) if e.is_collection_not_found() => self._create_collection(name).await,
             result => result,
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -113,7 +113,7 @@ impl SpannerDb {`
`113`	`113`
`114`	`114`	`async fn get_or_create_collection_id(&mut self, name: &str) -> DbResult<i32> {`
`115`	`115`	`match self.get_collection_id(name).await {`
`116`		`- Err(err) if err.is_collection_not_found() => self._create_collection(name).await,`
	`116`	`+ Err(e) if e.is_collection_not_found() => self._create_collection(name).await,`
`117`	`117`	`result => result,`
`118`	`118`	`}`
`119`	`119`	`}`