feat: add optional per-collection storage limits

This commit:
 - adds a third field to the elements of STD_COLLS for per-collection
   storage limit
 - adds an optional `limits.collections` map in settings that overrides
   the default limits from STD_COLLS
 - enforces the per-collection limits separately from account-wide usage

Author: chenba

Cargo.lock

@@ -59,9 +59,9 @@ syncstorage.database_url = "spanner://projects/SAMPLE_GCP_PROJECT/instances/SAMP
 # Max size of the syncstorage DB connection pool. Default: 10.
 # syncstorage.database_pool_max_size = 10
 
-# Quota tracking/enforcement (Spanner-only). Both default to false. Enable
-# tracking to record per-user usage; additionally enable enforcement to reject
-# writes over `limits.max_quota_limit` (default 2 GB).
+# Quota tracking/enforcement. Both default to false. Enable tracking to record
+# per-user usage; additionally enable enforcement to reject writes over
+# `limits.max_quota_limit` (default 2GB) or configured per-collection limits.
 syncstorage.enable_quota = 0
 # syncstorage.enforce_quota = 0
 

config/local.example.spanner.toml

@@ -49,9 +49,11 @@ syncstorage.database_url = "mysql://sample_user:sample_password@localhost/syncst
 # Max size of the syncstorage DB connection pool. Default: 10.
 # syncstorage.database_pool_max_size = 10
 
-# Quota tracking/enforcement are Spanner-only features and are force-disabled
-# for non-Spanner backends, so they have no effect on the MySQL build.
+# Quota tracking/enforcement. It is disabled by default; enable tracking to
+# record per-user usage and, optionally, enforcement to reject writes over the
+# configured quota limits.
 syncstorage.enable_quota = 0
+# syncstorage.enforce_quota = 0
 
 # Request/batch limits. Defaults shown in docs/src/config.md (Syncstorage
 # Limits). `max_total_records` is also used to cap MySQL `get_bsos` result

config/local.example.toml

@@ -72,15 +72,16 @@ The following configuration options are available.
 | <span id="SYNC_SYNCSTORAGE__LIMITS__MAX_REQUEST_BYTES"></span>SYNC_SYNCSTORAGE__LIMITS__MAX_REQUEST_BYTES | 2,625,536 | Max Content-Length for requests |
 | <span id="SYNC_SYNCSTORAGE__LIMITS__MAX_TOTAL_BYTES"></span>SYNC_SYNCSTORAGE__LIMITS__MAX_TOTAL_BYTES | 262,144,000 | Max BSO payload size per batch |
 | <span id="SYNC_SYNCSTORAGE__LIMITS__MAX_TOTAL_RECORDS"></span>SYNC_SYNCSTORAGE__LIMITS__MAX_TOTAL_RECORDS | 10,000 | Max BSO count per batch |
-| <span id="SYNC_SYNCSTORAGE__LIMITS__MAX_QUOTA_LIMIT"></span>SYNC_SYNCSTORAGE__LIMITS__MAX_QUOTA_LIMIT | 2,147,483,648 | Max storage quota per user (2 GB) |
+| <span id="SYNC_SYNCSTORAGE__LIMITS__MAX_QUOTA_LIMIT"></span>SYNC_SYNCSTORAGE__LIMITS__MAX_QUOTA_LIMIT | 2,147,483,648 | Account-wide storage quota for collections without their own per-collection quota (2 GB) |
+| <span id="SYNC_SYNCSTORAGE__LIMITS__COLLECTIONS__COLL_NAME"></span>SYNC_SYNCSTORAGE__LIMITS__COLLECTIONS__COLL_NAME | unset | Optional per-collection storage quota override, in bytes. Replace `COLL_NAME` with collection name. |
 
 ### Syncstorage Features
 
 | Env Var | Default Value | Description |
 | --- | --- | --- |
 | <span id="SYNC_SYNCSTORAGE__ENABLED"></span>SYNC_SYNCSTORAGE__ENABLED | true | Enable syncstorage service |
-| <span id="SYNC_SYNCSTORAGE__ENABLE_QUOTA"></span>SYNC_SYNCSTORAGE__ENABLE_QUOTA | false | Enable quota tracking (Spanner only) |
-| <span id="SYNC_SYNCSTORAGE__ENFORCE_QUOTA"></span>SYNC_SYNCSTORAGE__ENFORCE_QUOTA | false | Enforce quota limits (Spanner only) |
+| <span id="SYNC_SYNCSTORAGE__ENABLE_QUOTA"></span>SYNC_SYNCSTORAGE__ENABLE_QUOTA | false | Enable quota tracking and configured per-collection quota limits |
+| <span id="SYNC_SYNCSTORAGE__ENFORCE_QUOTA"></span>SYNC_SYNCSTORAGE__ENFORCE_QUOTA | false | Enforce account-wide and per-collection quota limits |
 | <span id="SYNC_SYNCSTORAGE__GLEAN_ENABLED"></span>SYNC_SYNCSTORAGE__GLEAN_ENABLED | true | Enable Glean telemetry |
 | <span id="SYNC_SYNCSTORAGE__LBHEARTBEAT_TTL"></span>SYNC_SYNCSTORAGE__LBHEARTBEAT_TTL | None | Load balancer heartbeat period in seconds |
 | <span id="SYNC_SYNCSTORAGE__LBHEARTBEAT_TTL_JITTER"></span>SYNC_SYNCSTORAGE__LBHEARTBEAT_TTL_JITTER | 25 | Jitter percentage for the load balancer heartbeat period |

docs/src/config.md

@@ -48,16 +48,14 @@ Set under `[syncstorage.limits]` in TOML, or as `SYNC_SYNCSTORAGE__LIMITS__*` en
 | `max_request_bytes`        | 2,625,536 (≈ 2.5 MB + 4 KB)              | HTTP `Content-Length` ceiling also enforced upstream of the API (e.g., nginx).                                                                                  |
 | `max_total_bytes`          | 250 MB declared, clamped                 | Combined batch payload size. **`Settings::normalize()` clamps this to `MAX_SPANNER_LOAD_SIZE` (100 MB) for Spanner deployments.**                                |
 | `max_total_records`        | 10,000 default; 9,984 in Spanner prod    | BSOs per batch. The Spanner ceiling is driven by the per-commit mutation budget, see [below](#batch-commit-mutation-budget).                                    |
-| `max_quota_limit`          | 2 GB                                     | Per-collection quota; only enforced when `enforce_quota = true`.                                                                                                  |
+| `max_quota_limit`          | 2 GB                                     | Account-wide quota for collections without their own per-collection limit; only enforced when `enforce_quota = true`.                                             |
 
 ### Quota toggles
 
 | Setting          | Default | Effect                                                                                                                                                              |
 | ---------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `enable_quota`   | false   | Enables `count` / `total_bytes` tracking on `user_collections`. Adds 6 mutations per batch commit (Steps 1 & 4 in the budget below).                                |
-| `enforce_quota`  | false   | When true, returns 403 once a user's collection exceeds `max_quota_limit`. When false but `enable_quota` is true, the server only logs a warning.                   |
-
-`Settings::normalize()` forces `max_quota_limit = 0` and disables both quota flags for non-Spanner deployments, quota is a Spanner-only feature in this codebase.
+| `enable_quota`   | false   | Enables `count` / `total_bytes` tracking on `user_collections` and activates configured per-collection quota limits. Adds 6 mutations per batch commit (Steps 1 & 4 in the budget below). |
+| `enforce_quota`  | false   | When true, returns 403 once a user's account-wide or per-collection usage reaches its configured quota. When false but `enable_quota` is true, the server only logs a warning. |
 
 ### Compile-time constants
 

docs/src/syncstorage/syncstorage-spanner-db.md

@@ -12,6 +12,7 @@ slog-scope.workspace = true
 config = "0.15"
 num_cpus = "1"
 syncserver-common = { path = "../syncserver-common" }
+syncstorage-db-common = { path = "../syncstorage-db-common" }
 syncstorage-settings = { path = "../syncstorage-settings" }
 tokenserver-settings = { path = "../tokenserver-settings" }
 url = "2.5"

syncserver-settings/Cargo.toml

@@ -7,6 +7,7 @@ use syncserver_common::{
     X_LAST_MODIFIED, X_VERIFY_CODE, X_WEAVE_BYTES, X_WEAVE_NEXT_OFFSET, X_WEAVE_RECORDS,
     X_WEAVE_TIMESTAMP, X_WEAVE_TOTAL_BYTES, X_WEAVE_TOTAL_RECORDS,
 };
+use syncstorage_db_common::STD_COLLS;
 use syncstorage_settings::Settings as SyncstorageSettings;
 use tokenserver_settings::Settings as TokenserverSettings;
 use url::Url;
@@ -168,6 +169,16 @@ impl Settings {
             ));
         }
 
+        // Per-collection limits may only target standard collections in STD_COLLS.
+        for name in self.syncstorage.limits.collections.keys() {
+            if !STD_COLLS.iter().any(|(_, std_name, _)| std_name == name) {
+                return Err(ConfigError::Message(format!(
+                    "Invalid SYNC_SYNCSTORAGE__LIMITS__COLLECTIONS: \
+                     \"{name}\" is not a standard collection name"
+                )));
+            }
+        }
+
         if let Some(init_node_url) = &self.tokenserver.init_node_url {
             let url = Url::parse(init_node_url).map_err(|e| {
                 ConfigError::Message(format!("Invalid SYNC_TOKENSERVER__INIT_NODE_URL: {e}"))
@@ -410,6 +421,49 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_unknown_collection_limit_name_fails_fast() {
+        temp_env::with_vars(
+            [
+                (
+                    "SYNC_SYNCSTORAGE__DATABASE_URL",
+                    Some(TEST_SYNCSTORAGE_DATABASE_URL),
+                ),
+                ("SYNC_TOKENSERVER__DATABASE_URL", None),
+                ("SYNC_TOKENSERVER__ENABLED", None),
+                ("SYNC_SYNCSTORAGE__LIMITS__COLLECTIONS__BATS", Some("9001")),
+            ],
+            || {
+                let err = Settings::with_env_and_config_file(None)
+                    .expect_err("an unknown collection limit name should fail validation");
+                assert!(err.to_string().contains("not a standard collection name"));
+            },
+        );
+    }
+
+    #[test]
+    fn test_known_collection_limit_name_validates() {
+        temp_env::with_vars(
+            [
+                (
+                    "SYNC_SYNCSTORAGE__DATABASE_URL",
+                    Some(TEST_SYNCSTORAGE_DATABASE_URL),
+                ),
+                ("SYNC_TOKENSERVER__DATABASE_URL", None),
+                ("SYNC_TOKENSERVER__ENABLED", None),
+                ("SYNC_SYNCSTORAGE__LIMITS__COLLECTIONS__TABS", Some("1024")),
+            ],
+            || {
+                let settings = Settings::with_env_and_config_file(None)
+                    .expect("a standard collection limit name should validate");
+                assert_eq!(
+                    settings.syncstorage.limits.collections.get("tabs"),
+                    Some(&1024)
+                );
+            },
+        );
+    }
+
     #[test]
     fn test_disabled_tokenserver_does_not_require_database_url() {
         // A disabled Tokenserver should not require a DATABASE_URL.

syncserver-settings/src/lib.rs

@@ -21,7 +21,7 @@ use syncserver_common::{
 };
 use syncserver_db_common::GetPoolStatus;
 use syncserver_settings::Settings;
-use syncstorage_db::{DbError, DbPool, DbPoolImpl};
+use syncstorage_db::{DbError, DbPool, DbPoolImpl, STD_COLLS};
 use syncstorage_settings::{Deadman, ServerLimits};
 use tokio::{sync::RwLock, time};
 use utoipa::OpenApi;
@@ -398,11 +398,10 @@ impl Server {
         let gcs_endpoint = settings.syncstorage.gcs_endpoint.clone();
         let worker_thread_count =
             calculate_worker_max_blocking_threads(settings.worker_max_blocking_threads);
+        let quota_enabled = settings.syncstorage.enable_quota;
+        let limits_json = build_limits_json(&settings.syncstorage.limits, quota_enabled);
         let limits = Arc::new(settings.syncstorage.limits);
-        let limits_json =
-            serde_json::to_string(&*limits).expect("ServerLimits failed to serialize");
         let secrets = Arc::new(settings.master_secret);
-        let quota_enabled = settings.syncstorage.enable_quota;
         let actix_keep_alive = settings.actix_keep_alive;
         let tokenserver_state = if settings.tokenserver.enabled {
             let mut state = tokenserver::ServerState::from_settings(
@@ -529,6 +528,29 @@ fn calculate_worker_max_blocking_threads(count: usize) -> usize {
     std::cmp::max(count / parallelism, 1)
 }
 
+/// Build the `info/configuration` client-limits JSON.
+///
+/// When `quota_enabled` and at least one collection resolves to a limit, injects a `collections`
+/// section (name -> `{max_quota_limit}`).
+pub(crate) fn build_limits_json(limits: &ServerLimits, quota_enabled: bool) -> String {
+    let collection_limits = if quota_enabled {
+        limits.resolved_collection_limits_by_name(&STD_COLLS)
+    } else {
+        Default::default()
+    };
+    let mut limits_value = serde_json::to_value(limits).expect("ServerLimits failed to serialize");
+    if !collection_limits.is_empty()
+        && let Some(obj) = limits_value.as_object_mut()
+    {
+        obj.insert(
+            "collections".to_owned(),
+            serde_json::to_value(&collection_limits)
+                .expect("collection limits failed to serialize"),
+        );
+    }
+    serde_json::to_string(&limits_value).expect("ServerLimits failed to serialize")
+}
+
 fn build_cors(settings: &Settings) -> Cors {
     // Followed by the "official middleware" so they run first.
     // actix is getting increasingly tighter about CORS headers. Our server is

syncserver/src/server/mod.rs

@@ -979,3 +979,36 @@ async fn put_bso_offloads_to_gcs() {
         resp.response()
     );
 }
+
+fn limits_with_collection(name: &str, limit: usize) -> ServerLimits {
+    let mut limits = ServerLimits::default();
+    limits.collections.insert(name.to_owned(), limit);
+    limits
+}
+
+#[::core::prelude::v1::test]
+fn limits_json_advertises_collections_when_quota_enabled() {
+    let limits = limits_with_collection("tabs", 5_242_880);
+    let json: Value = serde_json::from_str(&build_limits_json(&limits, true)).unwrap();
+    assert!(json["collections"]["tabs"].is_object());
+    assert_eq!(
+        json["collections"]["tabs"]["max_quota_limit"],
+        json!(5_242_880)
+    );
+}
+
+#[::core::prelude::v1::test]
+fn limits_json_omits_collections_when_quota_disabled() {
+    // quota is off, so the section is omitted
+    let limits = limits_with_collection("tabs", 5_242_880);
+    let json: Value = serde_json::from_str(&build_limits_json(&limits, false)).unwrap();
+    assert!(json.get("collections").is_none());
+}
+
+#[::core::prelude::v1::test]
+fn limits_json_omits_collections_when_none_limited() {
+    // quota enabled but nothing limited
+    let json: Value =
+        serde_json::from_str(&build_limits_json(&ServerLimits::default(), true)).unwrap();
+    assert!(json.get("collections").is_none());
+}

syncserver/src/server/test.rs

@@ -509,19 +509,9 @@ pub async fn post_collection_batch(
             .await?;
 
         if is_valid {
-            let usage = db
-                .get_quota_usage(params::GetQuotaUsage {
-                    user_id: coll.user_id.clone(),
-                    collection: coll.collection.clone(),
-                })
-                .await?;
             CreateBatch {
                 id: id.clone(),
-                size: if coll.quota_enabled {
-                    Some(usage.total_bytes)
-                } else {
-                    None
-                },
+                size: None, // will be seeded with value from `check_quota` later
             }
         } else {
             return Err(ApiErrorKind::Db(DbError::batch_not_found()).into());