feat: allow quota enforcement for all db backends

chenba · chenba · commit ea3422c10633 · 2026-06-10T12:41:17.000-05:00
This commit:
 - implements storage quotas for MySQL and Postgres
 - fixes a bug in Spanner's batch quota calculation
diff --git a/syncstorage-db/src/tests/batch.rs b/syncstorage-db/src/tests/batch.rs
@@ -420,6 +420,53 @@ async fn test_append_upsert_overwrites_same_batch_bso_id() -> Result<(), DbError
     .await
 }
 
+#[tokio::test]
+async fn test_quota_pending_batch_accumulates() -> Result<(), DbError> {
+    let settings = Settings::test_settings();
+    if !settings.syncstorage.enable_quota {
+        debug!("[test] Skipping test");
+        return Ok(());
+    }
+
+    with_test_transaction(
+        settings.syncstorage,
+        async |db: &mut dyn Db<Error = DbError>| {
+            let uid = 9001;
+            let coll = "clients";
+
+            let chunk = 3838;
+            let payload = "z".repeat(chunk);
+            // hits quota on second append.
+            db.set_quota(true, chunk + chunk / 2, true);
+
+            let new_batch = db.create_batch(cb(uid, coll, vec![])).await?;
+
+            db.append_to_batch(ab(
+                uid,
+                coll,
+                new_batch.clone(),
+                vec![postbso("b0", Some(&payload), None, None)],
+            ))
+            .await?;
+
+            let result = db
+                .append_to_batch(ab(
+                    uid,
+                    coll,
+                    new_batch.clone(),
+                    vec![postbso("b1", Some(&payload), None, None)],
+                ))
+                .await;
+            assert!(
+                result.as_ref().err().is_some_and(|e| e.is_quota()),
+                "expected an over-quota error, got {result:?}"
+            );
+            Ok(())
+        },
+    )
+    .await
+}
+
 #[tokio::test]
 async fn test_commit_batch_partial_overlap() -> Result<(), DbError> {
     let settings = Settings::test_settings().syncstorage;
diff --git a/syncstorage-db/src/tests/db.rs b/syncstorage-db/src/tests/db.rs
@@ -651,14 +651,14 @@ async fn get_collection_usage() -> Result<(), DbError> {
             let quota = db
                 .get_quota_usage(params::GetQuotaUsage {
                     user_id: hid(uid),
-                    collection: "ignored".to_owned(),
+                    collection: "bookmarks".to_owned(),
                 })
                 .await?;
             assert_eq!(
                 &(quota.total_bytes as i64),
                 expected.get("bookmarks").unwrap()
             );
-            assert_eq!(quota.count, 5); // 3 collections, 5 records
+            assert_eq!(quota.count, 5); // 5 records in the collection
         }
         Ok(())
     })
diff --git a/syncstorage-mysql/src/db/batch_impl.rs b/syncstorage-mysql/src/db/batch_impl.rs
@@ -1,5 +1,5 @@
 use base64::Engine;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 use async_trait::async_trait;
 use diesel::{
@@ -69,11 +69,25 @@ impl BatchDb for MysqlDb {
             }
         }
 
-        do_append(self, batch_id, params.user_id, collection_id, params.bsos).await?;
-        Ok(results::CreateBatch {
+        let batch = results::CreateBatch {
             id: encode_id(batch_id),
-            size: None,
-        })
+            // `size` is the committed collection size from `check_quota`.
+            size: self
+                .check_quota(&params.user_id, &params.collection)
+                .await?,
+        };
+
+        do_append(
+            self,
+            params.user_id,
+            collection_id,
+            batch.clone(),
+            params.bsos,
+            &params.collection,
+        )
+        .await?;
+
+        Ok(batch)
     }
 
     async fn validate_batch(&mut self, params: params::ValidateBatch) -> DbResult<bool> {
@@ -110,9 +124,24 @@ impl BatchDb for MysqlDb {
             return Err(DbError::batch_not_found());
         }
 
-        let batch_id = decode_id(&params.batch.id)?;
         let collection_id = self._get_collection_id(&params.collection).await?;
-        do_append(self, batch_id, params.user_id, collection_id, params.bsos).await?;
+
+        // `batch.size` is the committed collection size from `check_quota`.
+        let mut batch = params.batch;
+        batch.size = self
+            .check_quota(&params.user_id, &params.collection)
+            .await?;
+
+        do_append(
+            self,
+            params.user_id,
+            collection_id,
+            batch,
+            params.bsos,
+            &params.collection,
+        )
+        .await?;
+
         Ok(())
     }
 
@@ -191,10 +220,11 @@ impl BatchDb for MysqlDb {
 
 pub async fn do_append(
     db: &mut MysqlDb,
-    batch_id: i64,
     user_id: UserIdentifier,
     _collection_id: i32,
+    batch: results::CreateBatch,
     bsos: Vec<params::PostCollectionBso>,
+    collection: &str,
 ) -> DbResult<()> {
     fn exist_idx(user_id: u64, batch_id: i64, bso_id: &str) -> String {
         // Construct something that matches the key for batch_upload_items
@@ -206,6 +236,32 @@ pub async fn do_append(
         )
     }
 
+    let batch_id = decode_id(&batch.id)?;
+    let uid = user_id.legacy_id as i64;
+
+    if db.quota.enabled
+        && let Some(size) = batch.size
+    {
+        let running_size: usize = bsos
+            .iter()
+            .filter_map(|bso| bso.payload.as_ref().map(|p| p.len()))
+            .sum();
+        let incoming_ids: Vec<String> = bsos.iter().map(|bso| bso.id.clone()).collect();
+        let pending_size = pending_batch_size(db, uid, batch_id, &incoming_ids).await?;
+        let projected_total = size + pending_size + running_size;
+        if projected_total >= db.quota.size {
+            let mut tags = HashMap::default();
+            tags.insert("collection".to_owned(), collection.to_owned());
+            db.metrics.incr_with_tags("storage.quota.at_limit", tags);
+            if db.quota.enforced {
+                return Err(DbError::quota());
+            } else {
+                warn!("Quota at limit for user's collection ({} bytes)", projected_total;
+                      "collection" => collection);
+            }
+        }
+    }
+
     // It's possible for the list of items to contain a duplicate key entry.
     // This means that we can't really call `ON DUPLICATE` here, because that's
     // more about inserting one item at a time. (e.g. it works great if the
@@ -281,6 +337,25 @@ pub async fn do_append(
     Ok(())
 }
 
+/// Sum the pending payload bytes, excluding any ids in `incoming_ids`.
+async fn pending_batch_size(
+    db: &mut MysqlDb,
+    user_id: i64,
+    batch_id: i64,
+    incoming_ids: &[String],
+) -> DbResult<usize> {
+    let current_batch_size: i64 = batch_upload_items::table
+        .select(sql::<BigInt>("COALESCE(SUM(payload_size), 0)"))
+        .filter(batch_upload_items::user_id.eq(user_id))
+        .filter(batch_upload_items::batch_id.eq(batch_id))
+        .filter(batch_upload_items::id.ne_all(incoming_ids))
+        .get_result(&mut db.conn)
+        .await
+        .optional()?
+        .unwrap_or_default();
+    Ok(current_batch_size.max(0) as usize)
+}
+
 pub fn validate_batch_id(id: &str) -> DbResult<()> {
     decode_id(id).map(|_| ())
 }
diff --git a/syncstorage-mysql/src/db/mod.rs b/syncstorage-mysql/src/db/mod.rs
@@ -9,7 +9,7 @@ use diesel::{
 use diesel_async::RunQueryDsl;
 use syncserver_common::Metrics;
 use syncstorage_db_common::{
-    FIRST_CUSTOM_COLLECTION_ID, UserIdentifier, error::DbErrorIntrospect, results,
+    Db, FIRST_CUSTOM_COLLECTION_ID, UserIdentifier, error::DbErrorIntrospect, params, results,
     util::SyncTimestamp,
 };
 use syncstorage_settings::Quota;
@@ -237,6 +237,34 @@ impl MysqlDb {
             count,
         })
     }
+
+    async fn check_quota(
+        &mut self,
+        user_id: &UserIdentifier,
+        collection: &str,
+    ) -> 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(),
+            })
+            .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))
+    }
 }
 
 #[allow(dead_code)] // Not really dead, Rust can't see the use above
diff --git a/syncstorage-postgres/src/db/batch_impl.rs b/syncstorage-postgres/src/db/batch_impl.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use async_trait::async_trait;
 use diesel::{
     self, ExpressionMethods, OptionalExtension, QueryDsl, delete,
@@ -45,7 +47,10 @@ impl BatchDb for PgDb {
 
         let batch = results::CreateBatch {
             id: batch_id.to_string(),
-            size: None,
+            // `size` is the committed collection size from `check_quota`.
+            size: self
+                .check_quota(&params.user_id, &params.collection)
+                .await?,
         };
 
         do_append(
@@ -54,6 +59,7 @@ impl BatchDb for PgDb {
             collection_id,
             batch.clone(),
             params.bsos,
+            &params.collection,
         )
         .await?;
 
@@ -100,12 +106,19 @@ impl BatchDb for PgDb {
 
         let collection_id = self.get_or_create_collection_id(&params.collection).await?;
 
+        // `batch.size` is the committed collection size from `check_quota`.
+        let mut batch = params.batch;
+        batch.size = self
+            .check_quota(&params.user_id, &params.collection)
+            .await?;
+
         do_append(
             self,
             params.user_id,
             collection_id,
-            params.batch,
+            batch,
             params.bsos,
+            &params.collection,
         )
         .await?;
 
@@ -132,13 +145,7 @@ impl BatchDb for PgDb {
         let user_id = params.user_id.legacy_id as i64;
         let collection_id = self.get_or_create_collection_id(&params.collection).await?;
 
-        let timestamp = self
-            .update_collection(params::UpdateCollection {
-                user_id: params.user_id.clone(),
-                collection_id,
-                collection: params.collection.clone(),
-            })
-            .await?;
+        let timestamp = self.checked_timestamp()?;
         let default_ttl_seconds = DEFAULT_BSO_TTL as i64;
         let ts_datetime = timestamp.as_datetime()?;
 
@@ -151,6 +158,14 @@ impl BatchDb for PgDb {
             .execute(&mut self.conn)
             .await?;
 
+        let timestamp = self
+            .update_collection(params::UpdateCollection {
+                user_id: params.user_id.clone(),
+                collection_id,
+                collection: params.collection.clone(),
+            })
+            .await?;
+
         self.delete_batch(params::DeleteBatch {
             user_id: params.user_id,
             collection: params.collection,
@@ -194,14 +209,39 @@ pub async fn do_append(
     collection_id: i32,
     batch: results::CreateBatch,
     bsos: Vec<params::PostCollectionBso>,
+    collection: &str,
 ) -> DbResult<()> {
     let batch_id = Uuid::parse_str(&batch.id)
         .map_err(|e| DbError::internal(format!("Invalid batch_id in batch: {}", e)))?;
+    let user_id_i64 = user_id.legacy_id as i64;
+
+    if db.quota.enabled
+        && let Some(size) = batch.size
+    {
+        let running_size: usize = bsos
+            .iter()
+            .filter_map(|bso| bso.payload.as_ref().map(|p| p.len()))
+            .sum();
+        let incoming_ids: Vec<String> = bsos.iter().map(|bso| bso.id.clone()).collect();
+        let pending_size =
+            pending_batch_size(db, user_id_i64, collection_id, &batch_id, &incoming_ids).await?;
+        let projected_total = size + pending_size + running_size;
+        if projected_total >= db.quota.size {
+            let mut tags = HashMap::default();
+            tags.insert("collection".to_owned(), collection.to_owned());
+            db.metrics.incr_with_tags("storage.quota.at_limit", tags);
+            if db.quota.enforced {
+                return Err(DbError::quota());
+            } else {
+                warn!("Quota at limit for user's collection ({} bytes)", projected_total;
+                      "collection" => collection);
+            }
+        }
+    }
 
     for bso in bsos {
         let ttl = bso.ttl.map(|t| t as i64);
         let sortindex = bso.sortindex;
-        let user_id_i64 = user_id.legacy_id as i64;
 
         sql_query(
             "INSERT INTO batch_bsos (user_id, collection_id, batch_id, batch_bso_id, sortindex, payload, ttl)
@@ -225,6 +265,29 @@ pub async fn do_append(
     Ok(())
 }
 
+/// Sum the pending payload bytes, excluding any ids in `incoming_ids`.
+async fn pending_batch_size(
+    db: &mut PgDb,
+    user_id: i64,
+    collection_id: i32,
+    batch_id: &Uuid,
+    incoming_ids: &[String],
+) -> DbResult<usize> {
+    let current_batch_size: i64 = batch_bsos::table
+        .select(sql::<BigInt>(
+            "COALESCE(SUM(LENGTH(COALESCE(payload, ''))), 0)::BIGINT",
+        ))
+        .filter(batch_bsos::user_id.eq(user_id))
+        .filter(batch_bsos::collection_id.eq(collection_id))
+        .filter(batch_bsos::batch_id.eq(batch_id))
+        .filter(batch_bsos::batch_bso_id.ne_all(incoming_ids))
+        .get_result(&mut db.conn)
+        .await
+        .optional()?
+        .unwrap_or_default();
+    Ok(current_batch_size.max(0) as usize)
+}
+
 pub fn validate_batch_id(id: &str) -> DbResult<Uuid> {
     Uuid::parse_str(id).map_err(|e| DbError::internal(format!("Invalid batch_id: {}", e)))
 }
diff --git a/syncstorage-settings/src/lib.rs b/syncstorage-settings/src/lib.rs
diff --git a/syncstorage-spanner/src/db/batch_impl.rs b/syncstorage-spanner/src/db/batch_impl.rs
