Merge pull request #2013 from mozilla-services/feat/post_bsos-STOR-423

pjenvey · web-flow · commit b16ef5064ba9 · 2026-01-21T11:00:17.000-08:00
feat: optimize post_bsos w/ MERGE INTO
diff --git a/syncstorage-postgres/migrations/2025-10-20-155711_create_schema/up.sql b/syncstorage-postgres/migrations/2025-10-20-155711_create_schema/up.sql
@@ -112,3 +112,14 @@ CREATE TABLE batch_bsos (
     ) ON DELETE CASCADE
 );
 
+
+-- solely used for the ease of passing post_bsos data via a single
+-- bind parameter
+CREATE TYPE post_bso AS (
+    bso_id TEXT,
+
+    sortindex INTEGER,
+    payload TEXT,
+
+    ttl BIGINT
+);
diff --git a/syncstorage-postgres/src/db/db_impl.rs b/syncstorage-postgres/src/db/db_impl.rs
@@ -3,13 +3,12 @@ use chrono::{offset::Utc, DateTime, TimeDelta};
 use diesel::{
     delete,
     dsl::{count, max, now, sql},
-    sql_types::{BigInt, Nullable, Timestamptz},
+    sql_types::{Array, BigInt, Integer, Nullable, Timestamptz},
     upsert::excluded,
     ExpressionMethods, IntoSql, OptionalExtension, QueryDsl, SelectableHelper,
 };
 use diesel_async::{AsyncConnection, RunQueryDsl, TransactionManager};
 use futures::TryStreamExt;
-use std::collections::HashMap;
 use syncstorage_db_common::{
     error::DbErrorIntrospect, params, results, util::SyncTimestamp, Db, Sorting, DEFAULT_BSO_TTL,
 };
@@ -18,7 +17,7 @@ use super::{PgDb, TOMBSTONE};
 use crate::{
     bsos_query,
     db::{CollectionLock, PRETOUCH_DT},
-    orm_models::BsoChangeset,
+    orm_models::{sql_types::PostBso, BsoChangeset},
     pool::Conn,
     schema::{bsos, collections, user_collections},
     DbError, DbResult,
@@ -73,6 +72,7 @@ impl Db for PgDb {
         &mut self,
         params: params::LockCollection,
     ) -> DbResult<results::LockCollection> {
+        let user_id = params.user_id.legacy_id as i64;
         let collection_id = self
             .get_collection_id(&params.collection)
             .await
@@ -86,7 +86,6 @@ impl Db for PgDb {
                 }
             })?;
 
-        let user_id = params.user_id.legacy_id as i64;
         let key = (params.user_id, collection_id);
         // If we already have a read or write lock then it's safe to
         // use it as-is.
@@ -117,8 +116,8 @@ impl Db for PgDb {
     }
 
     async fn lock_for_write(&mut self, params: params::LockCollection) -> DbResult<()> {
-        let collection_id = self.get_or_create_collection_id(&params.collection).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);
 
         if let Some(CollectionLock::Read) = self.session.coll_locks.get(&key) {
@@ -353,38 +352,11 @@ impl Db for PgDb {
         Ok(results::GetBsoIds { items, offset })
     }
 
-    async fn post_bsos(&mut self, params: params::PostBsos) -> DbResult<results::PostBsos> {
-        let collection_id = self.get_or_create_collection_id(&params.collection).await?;
-        let modified = self.checked_timestamp()?;
-
-        self.ensure_user_collection(params.user_id.legacy_id as i64, collection_id)
-            .await?;
-        for pbso in params.bsos {
-            self.put_bso(params::PutBso {
-                user_id: params.user_id.clone(),
-                collection: params.collection.clone(),
-                id: pbso.id.clone(),
-                payload: pbso.payload,
-                sortindex: pbso.sortindex,
-                ttl: pbso.ttl,
-            })
-            .await?;
-        }
-        self.update_collection(params::UpdateCollection {
-            user_id: params.user_id,
-            collection_id,
-            collection: params.collection,
-        })
-        .await?;
-
-        Ok(modified)
-    }
-
     async fn delete_bso(&mut self, params: params::DeleteBso) -> DbResult<results::DeleteBso> {
-        let user_id = params.user_id.legacy_id;
+        let user_id = params.user_id.legacy_id as i64;
         let collection_id = self.get_collection_id(&params.collection).await?;
         let affected_rows = delete(bsos::table)
-            .filter(bsos::user_id.eq(user_id as i64))
+            .filter(bsos::user_id.eq(user_id))
             .filter(bsos::collection_id.eq(&collection_id))
             .filter(bsos::bso_id.eq(params.id))
             .filter(bsos::expiry.gt(now))
@@ -437,65 +409,31 @@ impl Db for PgDb {
     }
 
     async fn put_bso(&mut self, bso: params::PutBso) -> DbResult<results::PutBso> {
+        let user_id = bso.user_id.legacy_id as i64;
         let collection_id = self.get_or_create_collection_id(&bso.collection).await?;
-        let user_id: u64 = bso.user_id.legacy_id;
-        if self.quota.enabled {
-            let usage = self
-                .get_quota_usage(params::GetQuotaUsage {
-                    user_id: bso.user_id.clone(),
-                    collection: bso.collection.clone(),
-                    collection_id,
-                })
-                .await?;
-            if usage.total_bytes >= self.quota.size {
-                let mut tags = HashMap::default();
-                tags.insert("collection".to_owned(), bso.collection.clone());
-                self.metrics.incr_with_tags("storage.quota.at_limit", tags);
-                if self.quota.enforced {
-                    return Err(DbError::quota());
-                } else {
-                    warn!("Quota at limit for user's collection ({} bytes)", usage.total_bytes; "collection"=>bso.collection.clone());
-                }
-            }
-        }
+
+        self.check_quota(&bso.user_id, &bso.collection, collection_id)
+            .await?;
 
         let payload = bso.payload.as_deref().unwrap_or_default();
         let sortindex = bso.sortindex;
-        let ttl = bso.ttl.map_or(DEFAULT_BSO_TTL, |ttl| ttl);
+        let ttl = bso.ttl.unwrap_or(DEFAULT_BSO_TTL);
 
         let modified = self.checked_timestamp()?.as_datetime()?;
         // Expiry originally required millisecond conversion
         let expiry = modified + TimeDelta::seconds(ttl as i64);
-
         // The changeset utilizes Diesel's `AsChangeset` trait.
         // This allows selective updates of fields if and only if they are `Some(<T>)`
         let changeset = BsoChangeset {
-            sortindex: if bso.sortindex.is_some() {
-                sortindex // sortindex is already an Option of type `Option<i32>`
-            } else {
-                None
-            },
-            payload: if bso.payload.is_some() {
-                Some(payload)
-            } else {
-                None
-            },
-            expiry: if bso.ttl.is_some() {
-                Some(expiry)
-            } else {
-                None
-            },
-            modified: if bso.payload.is_some() || bso.sortindex.is_some() {
-                Some(modified)
-            } else {
-                None
-            },
+            sortindex: bso.sortindex,
+            payload: bso.payload.as_deref(),
+            modified: (bso.payload.is_some() || bso.sortindex.is_some()).then_some(modified),
+            expiry: bso.ttl.map(|_| expiry),
         };
-        self.ensure_user_collection(user_id as i64, collection_id)
-            .await?;
+        self.ensure_user_collection(user_id, collection_id).await?;
         diesel::insert_into(bsos::table)
             .values((
-                bsos::user_id.eq(user_id as i64),
+                bsos::user_id.eq(user_id),
                 bsos::collection_id.eq(&collection_id),
                 bsos::bso_id.eq(&bso.id),
                 bsos::sortindex.eq(sortindex),
@@ -517,6 +455,35 @@ impl Db for PgDb {
         .await
     }
 
+    async fn post_bsos(&mut self, params: params::PostBsos) -> DbResult<results::PostBsos> {
+        let user_id = params.user_id.legacy_id as i64;
+        let collection_id = self.get_or_create_collection_id(&params.collection).await?;
+        self.check_quota(&params.user_id, &params.collection, collection_id)
+            .await?;
+        self.ensure_user_collection(user_id, collection_id).await?;
+
+        // Rendering a VALUES statement for MERGE INTO here is painful so we
+        // pass the bsos in a single bind param as an Array of a named
+        // composite type (post_bso[]). The composite type must be explicitly
+        // named/declared as postgres disallows unnamed/anonymous composite
+        // types to be used as bind parameters
+        diesel::sql_query(include_str!("post_bsos.sql"))
+            .bind::<BigInt, _>(user_id)
+            .bind::<Integer, _>(collection_id)
+            .bind::<Array<PostBso>, _>(params.bsos)
+            .bind::<Timestamptz, _>(self.checked_timestamp()?.as_datetime()?)
+            .bind::<BigInt, _>(DEFAULT_BSO_TTL as i64)
+            .execute(&mut self.conn)
+            .await?;
+
+        self.update_collection(params::UpdateCollection {
+            user_id: params.user_id,
+            collection_id,
+            collection: params.collection,
+        })
+        .await
+    }
+
     async fn get_collection_id(&mut self, name: &str) -> DbResult<results::GetCollectionId> {
         if let Some(id) = self.coll_cache.get_id(name)? {
             return Ok(id);
@@ -548,9 +515,9 @@ impl Db for PgDb {
         &mut self,
         params: params::UpdateCollection,
     ) -> DbResult<SyncTimestamp> {
+        let user_id = params.user_id.legacy_id as i64;
         let quota = if self.quota.enabled {
-            self.calc_quota_usage(params.user_id.legacy_id as i64, params.collection_id)
-                .await?
+            self.calc_quota_usage(user_id, params.collection_id).await?
         } else {
             results::GetQuotaUsage {
                 count: 0,
@@ -562,7 +529,7 @@ impl Db for PgDb {
 
         diesel::insert_into(user_collections::table)
             .values((
-                user_collections::user_id.eq(params.user_id.legacy_id as i64),
+                user_collections::user_id.eq(user_id),
                 user_collections::collection_id.eq(params.collection_id),
                 user_collections::modified.eq(modified.as_datetime()?),
                 user_collections::count.eq(quota.count as i64),
diff --git a/syncstorage-postgres/src/db/mod.rs b/syncstorage-postgres/src/db/mod.rs
@@ -9,8 +9,9 @@ use diesel_async::RunQueryDsl;
 use std::{collections::HashMap, fmt, sync::Arc};
 
 use syncserver_common::Metrics;
-use syncstorage_db_common::diesel::DbError;
-use syncstorage_db_common::{results, util::SyncTimestamp, UserIdentifier};
+use syncstorage_db_common::{
+    diesel::DbError, params, results, util::SyncTimestamp, Db, UserIdentifier,
+};
 use syncstorage_settings::Quota;
 
 use super::schema::{bsos, collections, user_collections};
@@ -205,6 +206,36 @@ impl PgDb {
         Ok(())
     }
 
+    async fn check_quota(
+        &mut self,
+        user_id: &UserIdentifier,
+        collection: &str,
+        collection_id: i32,
+    ) -> DbResult<Option<usize>> {
+        if !self.quota.enabled {
+            return Ok(None);
+        }
+        let usage = self
+            .get_quota_usage(params::GetQuotaUsage {
+                user_id: user_id.clone(),
+                collection: collection.to_owned(),
+                collection_id,
+            })
+            .await?;
+        if usage.total_bytes >= self.quota.size {
+            let mut tags = HashMap::default();
+            tags.insert("collection".to_owned(), collection.to_owned());
+            self.metrics.incr_with_tags("storage.quota.at_limit", tags);
+            if self.quota.enforced {
+                return Err(DbError::quota());
+            } else {
+                warn!("Quota at limit for user's collection: ({} bytes)", usage.total_bytes;
+                      "collection" => collection);
+            }
+        }
+        Ok(Some(usage.total_bytes))
+    }
+
     // perform a heavier weight quota calculation
     async fn calc_quota_usage(
         &mut self,
diff --git a/syncstorage-postgres/src/db/post_bsos.sql b/syncstorage-postgres/src/db/post_bsos.sql
@@ -0,0 +1,33 @@
+           MERGE INTO bsos
+                USING UNNEST($3) AS post
+                   ON bsos.user_id = $1
+                  AND bsos.collection_id = $2
+                  AND bsos.bso_id = post.bso_id
+    WHEN MATCHED THEN
+           UPDATE SET
+                      sortindex = COALESCE(post.sortindex, bsos.sortindex),
+                      payload = COALESCE(post.payload, bsos.payload),
+                      modified = COALESCE(
+                         CASE
+                         WHEN post.payload is NOT NULL OR post.sortindex IS NOT NULL THEN $4
+                         ELSE NULL
+                         END,
+                         bsos.modified
+                     ),
+                      expiry = COALESCE(
+                          CASE
+                          WHEN post.ttl IS NOT NULL THEN $4 + (post.ttl || ' seconds')::INTERVAL
+                          ELSE NULL
+                          END,
+                          bsos.expiry
+                      )
+WHEN NOT MATCHED THEN
+               INSERT (user_id, collection_id, bso_id, sortindex, payload, modified, expiry)
+               VALUES ($1,
+                       $2,
+                       post.bso_id,
+                       post.sortindex,
+                       COALESCE(post.payload, ''),
+                       $4,
+                       $4 + (COALESCE(post.ttl, $5) || ' seconds')::INTERVAL
+                      )
diff --git a/syncstorage-postgres/src/orm_models.rs b/syncstorage-postgres/src/orm_models.rs
@@ -1,8 +1,8 @@
 use chrono::{offset::Utc, DateTime};
+use diesel::{AsChangeset, Identifiable, Queryable};
 use uuid::Uuid;
 
 use crate::schema::{batch_bsos, batches, bsos, collections, user_collections};
-use diesel::{AsChangeset, Identifiable, Queryable};
 
 #[derive(Queryable, Debug, Identifiable)]
 #[diesel(primary_key(user_id, collection_id, batch_id, batch_bso_id))]
@@ -18,16 +18,17 @@ pub struct BatchBso {
 
 #[derive(Queryable, Debug, Identifiable)]
 #[diesel(primary_key(user_id, collection_id, batch_id))]
-#[diesel(table_name=batches)]
+#[diesel(table_name = batches)]
 pub struct Batch {
     pub user_id: i64,
     pub collection_id: i32,
     pub batch_id: Uuid,
     pub expiry: DateTime<Utc>,
 }
 
-#[derive(Queryable, Debug, Identifiable)]
+#[derive(Queryable, Debug, Identifiable, Insertable)]
 #[diesel(primary_key(user_id, collection_id, bso_id))]
+#[diesel(table_name = bsos)]
 pub struct Bso {
     pub user_id: i64,
     pub collection_id: i32,
@@ -64,3 +65,30 @@ pub struct UserCollection {
     pub count: Option<i64>,
     pub total_bytes: Option<i64>,
 }
+
+pub mod sql_types {
+    use diesel::{
+        pg::Pg,
+        serialize::{self, Output, ToSql, WriteTuple},
+        sql_types::{BigInt, Integer, Nullable, Text},
+    };
+    use syncstorage_db_common::params;
+
+    #[derive(SqlType, QueryId)]
+    #[diesel(postgres_type(name = "post_bso"))]
+    pub struct PostBso;
+
+    impl ToSql<PostBso, Pg> for params::PostCollectionBso {
+        fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result {
+            WriteTuple::<(Text, Nullable<Integer>, Nullable<Text>, Nullable<BigInt>)>::write_tuple(
+                &(
+                    &self.id,
+                    &self.sortindex,
+                    &self.payload,
+                    &self.ttl.map(|ttl| ttl as i64),
+                ),
+                &mut out.reborrow(),
+            )
+        }
+    }
+}
