diff --git a/Cargo.lock b/Cargo.lock index 652de41921..34984f018e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4838,6 +4838,7 @@ dependencies = [ "serde", "slog-scope", "syncserver-common", + "syncstorage-db-common", "syncstorage-settings", "temp-env", "tokenserver-settings", diff --git a/config/local.example.spanner.toml b/config/local.example.spanner.toml index 34ab6488c8..2c5b102c1c 100644 --- a/config/local.example.spanner.toml +++ b/config/local.example.spanner.toml @@ -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 diff --git a/config/local.example.toml b/config/local.example.toml index 2318a6d2bb..702dfede19 100644 --- a/config/local.example.toml +++ b/config/local.example.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 diff --git a/docs/src/config.md b/docs/src/config.md index c2dfad5c6e..a6b67e945b 100644 --- a/docs/src/config.md +++ b/docs/src/config.md @@ -72,15 +72,16 @@ The following configuration options are available. | SYNC_SYNCSTORAGE__LIMITS__MAX_REQUEST_BYTES | 2,625,536 | Max Content-Length for requests | | SYNC_SYNCSTORAGE__LIMITS__MAX_TOTAL_BYTES | 262,144,000 | Max BSO payload size per batch | | SYNC_SYNCSTORAGE__LIMITS__MAX_TOTAL_RECORDS | 10,000 | Max BSO count per batch | -| SYNC_SYNCSTORAGE__LIMITS__MAX_QUOTA_LIMIT | 2,147,483,648 | Max storage quota per user (2 GB) | +| SYNC_SYNCSTORAGE__LIMITS__MAX_QUOTA_LIMIT | 2,147,483,648 | Account-wide storage quota for collections without their own per-collection quota (2 GB) | +| 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 | | --- | --- | --- | | SYNC_SYNCSTORAGE__ENABLED | true | Enable syncstorage service | -| SYNC_SYNCSTORAGE__ENABLE_QUOTA | false | Enable quota tracking (Spanner only) | -| SYNC_SYNCSTORAGE__ENFORCE_QUOTA | false | Enforce quota limits (Spanner only) | +| SYNC_SYNCSTORAGE__ENABLE_QUOTA | false | Enable quota tracking and configured per-collection quota limits | +| SYNC_SYNCSTORAGE__ENFORCE_QUOTA | false | Enforce account-wide and per-collection quota limits | | SYNC_SYNCSTORAGE__GLEAN_ENABLED | true | Enable Glean telemetry | | SYNC_SYNCSTORAGE__LBHEARTBEAT_TTL | None | Load balancer heartbeat period in seconds | | SYNC_SYNCSTORAGE__LBHEARTBEAT_TTL_JITTER | 25 | Jitter percentage for the load balancer heartbeat period | diff --git a/docs/src/syncstorage/syncstorage-spanner-db.md b/docs/src/syncstorage/syncstorage-spanner-db.md index 2c49e810cd..d985e5fc80 100644 --- a/docs/src/syncstorage/syncstorage-spanner-db.md +++ b/docs/src/syncstorage/syncstorage-spanner-db.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 diff --git a/syncserver-settings/Cargo.toml b/syncserver-settings/Cargo.toml index d26e6c61f5..b2337e31fa 100644 --- a/syncserver-settings/Cargo.toml +++ b/syncserver-settings/Cargo.toml @@ -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" diff --git a/syncserver-settings/src/lib.rs b/syncserver-settings/src/lib.rs index 1feb487c11..ed75136b52 100644 --- a/syncserver-settings/src/lib.rs +++ b/syncserver-settings/src/lib.rs @@ -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. diff --git a/syncserver/src/server/mod.rs b/syncserver/src/server/mod.rs index 3c85561920..9076bd9493 100644 --- a/syncserver/src/server/mod.rs +++ b/syncserver/src/server/mod.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 diff --git a/syncserver/src/server/test.rs b/syncserver/src/server/test.rs index 936b7ada53..947ba2f462 100644 --- a/syncserver/src/server/test.rs +++ b/syncserver/src/server/test.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()); +} diff --git a/syncserver/src/web/handlers.rs b/syncserver/src/web/handlers.rs index 424bc1897d..4e1e5d83c0 100644 --- a/syncserver/src/web/handlers.rs +++ b/syncserver/src/web/handlers.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()); diff --git a/syncstorage-db-common/src/lib.rs b/syncstorage-db-common/src/lib.rs index f41bd497f1..632a5f847c 100644 --- a/syncstorage-db-common/src/lib.rs +++ b/syncstorage-db-common/src/lib.rs @@ -20,21 +20,23 @@ lazy_static! { /// common collection names. This is the canonical list of such /// names. Non-standard collections will be allocated IDs starting /// from the highest ID in this collection. - pub static ref STD_COLLS: Vec<(i32, &'static str)> = { + /// + /// Each entry is `(id, name, per-collection byte limit)`. + pub static ref STD_COLLS: Vec<(i32, &'static str, Option)> = { vec![ - (1, "clients"), - (2, "crypto"), - (3, "forms"), - (4, "history"), - (5, "keys"), - (6, "meta"), - (7, "bookmarks"), - (8, "prefs"), - (9, "tabs"), - (10, "passwords"), - (11, "addons"), - (12, "addresses"), - (13, "creditcards"), + (1, "clients", None), + (2, "crypto", None), + (3, "forms", None), + (4, "history", None), + (5, "keys", None), + (6, "meta", None), + (7, "bookmarks", None), + (8, "prefs", None), + (9, "tabs", None), + (10, "passwords", None), + (11, "addons", None), + (12, "addresses", None), + (13, "creditcards", None), ] }; } @@ -233,6 +235,9 @@ pub trait Db: BatchDb { #[cfg(debug_assertions)] fn set_quota(&mut self, enabled: bool, limit: usize, enforce: bool); + + #[cfg(debug_assertions)] + fn set_collection_limits(&mut self, _limits: std::collections::HashMap) {} } #[async_trait(?Send)] diff --git a/syncstorage-db/src/lib.rs b/syncstorage-db/src/lib.rs index 376fe2856b..b956ecf0cd 100644 --- a/syncstorage-db/src/lib.rs +++ b/syncstorage-db/src/lib.rs @@ -33,7 +33,7 @@ pub use syncserver_db_common::GetPoolStatus; pub use syncstorage_db_common::error::DbErrorIntrospect; pub use syncstorage_db_common::{ - Db, DbPool, Sorting, UserIdentifier, params, results, + Db, DbPool, STD_COLLS, Sorting, UserIdentifier, params, results, util::{SyncTimestamp, to_rfc3339}, }; diff --git a/syncstorage-db/src/tests/batch.rs b/syncstorage-db/src/tests/batch.rs index 3511acf867..69c50eadfa 100644 --- a/syncstorage-db/src/tests/batch.rs +++ b/syncstorage-db/src/tests/batch.rs @@ -253,6 +253,63 @@ async fn quota_test_append_batch() -> Result<(), DbError> { .await } +#[tokio::test] +async fn test_quota_per_collection_batch() -> Result<(), DbError> { + let mut settings = Settings::test_settings().syncstorage; + settings.enable_quota = true; + settings.enforce_quota = true; + + with_test_transaction(settings, async |db: &mut dyn Db| { + let uid = 1; + let limited_coll = "tabs"; + let limited_id = 9; + let unlimited_coll = "clients"; + let limit = 256; + + db.set_quota(true, limit * 1000, true); + db.set_collection_limits(std::collections::HashMap::from([(limited_id, limit)])); + + let filler = (0..limit - 8).map(|_| "#").collect::>().concat(); + + let b0 = db + .create_batch(cb( + uid, + limited_coll, + vec![postbso("t0", Some(&filler), None, None)], + )) + .await?; + let batch = db + .get_batch(gb(uid, limited_coll, b0.id.clone())) + .await? + .unwrap(); + db.commit_batch(params::CommitBatch { + user_id: hid(uid), + collection: limited_coll.to_owned(), + batch, + }) + .await?; + + let over = db + .create_batch(cb( + uid, + limited_coll, + vec![postbso("t1", Some(&filler), None, None)], + )) + .await; + assert!(over.is_err(),); + + // account-wide limit not affected + db.create_batch(cb( + uid, + unlimited_coll, + vec![postbso("c0", Some(&filler), None, None)], + )) + .await?; + Ok(()) + }) + .await +} + #[tokio::test] async fn test_append_async_w_null() -> Result<(), DbError> { let settings = Settings::test_settings().syncstorage; diff --git a/syncstorage-db/src/tests/db.rs b/syncstorage-db/src/tests/db.rs index 4ae0c48044..41bf150e92 100644 --- a/syncstorage-db/src/tests/db.rs +++ b/syncstorage-db/src/tests/db.rs @@ -704,6 +704,65 @@ async fn test_quota() -> Result<(), DbError> { .await } +#[tokio::test] +async fn per_collection_quota() -> Result<(), DbError> { + let mut settings = Settings::test_settings().syncstorage; + settings.enable_quota = true; + settings.enforce_quota = true; + + with_test_transaction(settings, async |db: &mut dyn Db| { + let uid = *UID; + let limited_coll = "tabs"; + let limited_id = 9; + let unlimited_coll = "clients"; + let size = 256; + let limit = size * 2; + + db.set_quota(true, size * 1000, true); + db.set_collection_limits(HashMap::from([(limited_id, limit)])); + + let random = rng() + .sample_iter(&Alphanumeric) + .take(size) + .collect::>(); + let payload = String::from_utf8_lossy(&random); + + // fill to limit + db.put_bso(pbso(uid, limited_coll, "t0", Some(&payload), None, None)) + .await?; + db.put_bso(pbso(uid, limited_coll, "t1", Some(&payload), None, None)) + .await?; + + let replace_over = db + .put_bso(pbso(uid, limited_coll, "t1", Some(&payload), None, None)) + .await; + assert!(replace_over.is_err()); + + let new_bso_over = db + .put_bso(pbso(uid, limited_coll, "t2", Some(&payload), None, None)) + .await; + assert!(new_bso_over.is_err()); + + db.put_bso(pbso(uid, unlimited_coll, "c0", Some(&payload), None, None)) + .await?; + + // quota disabled, no per-collection limit + db.set_quota(false, 0, false); + db.put_bso(pbso(uid, limited_coll, "t2", Some(&payload), None, None)) + .await?; + let disabled_usage = db + .get_quota_usage(params::GetQuotaUsage { + user_id: hid(uid), + collection: limited_coll.to_owned(), + }) + .await?; + assert_eq!(disabled_usage.total_bytes, 0); + assert_eq!(disabled_usage.count, 0); + Ok(()) + }) + .await +} + #[tokio::test] async fn get_collection_counts() -> Result<(), DbError> { with_test_transaction(None, async |db: &mut dyn Db| { diff --git a/syncstorage-mysql/src/db/batch_impl.rs b/syncstorage-mysql/src/db/batch_impl.rs index f5df34cac5..8f2983c573 100644 --- 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::{HashMap, HashSet}; +use std::collections::HashSet; use async_trait::async_trait; use diesel::{ @@ -221,7 +221,7 @@ impl BatchDb for MysqlDb { pub async fn do_append( db: &mut MysqlDb, user_id: UserIdentifier, - _collection_id: i32, + collection_id: i32, batch: results::CreateBatch, bsos: Vec, collection: &str, @@ -248,17 +248,14 @@ pub async fn do_append( .sum(); let incoming_ids: Vec = 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 ({} bytes)", projected_total; - "collection" => collection); - } + let ceiling = db + .get_collection_limit(collection_id) + .unwrap_or(db.quota.size); + let projected_total = size + .saturating_add(pending_size) + .saturating_add(running_size); + if projected_total >= ceiling { + db.emit_at_limit(collection, projected_total)?; } } diff --git a/syncstorage-mysql/src/db/db_impl.rs b/syncstorage-mysql/src/db/db_impl.rs index 796b414ecb..2f2596f0d5 100644 --- a/syncstorage-mysql/src/db/db_impl.rs +++ b/syncstorage-mysql/src/db/db_impl.rs @@ -194,96 +194,8 @@ impl Db for MysqlDb { */ let collection_id = self.get_or_create_collection_id(&bso.collection).await?; - let user_id: u64 = bso.user_id.legacy_id; - let timestamp = self.session.timestamp.as_i64(); - if self.quota.enabled { - let usage = self - .get_quota_usage(params::GetQuotaUsage { - user_id: bso.user_id.clone(), - collection: bso.collection.clone(), - }) - .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 ({} bytes)", usage.total_bytes; "collection"=>bso.collection.clone()); - } - } - } - - 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 q = format!( - r#" - INSERT INTO bso ({user_id}, {collection_id}, id, sortindex, payload, {modified}, {expiry}) - VALUES (?, ?, ?, ?, ?, ?, ?) - ON DUPLICATE KEY UPDATE - {user_id} = VALUES({user_id}), - {collection_id} = VALUES({collection_id}), - id = VALUES(id) - "#, - user_id = USER_ID, - modified = MODIFIED, - collection_id = COLLECTION_ID, - expiry = EXPIRY - ); - let q = format!( - "{}{}", - q, - if bso.sortindex.is_some() { - ", sortindex = VALUES(sortindex)" - } else { - "" - }, - ); - let q = format!( - "{}{}", - q, - if bso.payload.is_some() { - ", payload = VALUES(payload)" - } else { - "" - }, - ); - let q = format!( - "{}{}", - q, - if bso.ttl.is_some() { - format!(", {expiry} = VALUES({expiry})", expiry = EXPIRY) - } else { - "".to_owned() - }, - ); - let q = format!( - "{}{}", - q, - if bso.payload.is_some() || bso.sortindex.is_some() { - format!(", {modified} = VALUES({modified})", modified = MODIFIED) - } else { - "".to_owned() - }, - ); - sql_query(q) - .bind::(user_id as i64) // XXX: - .bind::(&collection_id) - .bind::(&bso.id) - .bind::, _>(sortindex) - .bind::(payload) - .bind::(timestamp) - .bind::(timestamp + (i64::from(ttl) * 1000)) // remember: this is in millis - .execute(&mut self.conn) - .await?; - self.update_collection(params::UpdateCollection { - user_id: bso.user_id, - collection_id, - collection: bso.collection, - }) - .await + self.check_quota(&bso.user_id, &bso.collection).await?; + self.write_bso(bso, collection_id).await } async fn get_bsos(&mut self, params: params::GetBsos) -> DbResult { @@ -516,16 +428,21 @@ impl Db for MysqlDb { let collection_id = self.get_or_create_collection_id(&input.collection).await?; let modified = self.session.timestamp; + self.check_quota(&input.user_id, &input.collection).await?; + for pbso in input.bsos { - self.put_bso(params::PutBso { - user_id: input.user_id.clone(), - collection: input.collection.clone(), - id: pbso.id.clone(), - payload: pbso.payload, - payload_link: pbso.payload_link, - sortindex: pbso.sortindex, - ttl: pbso.ttl, - }) + self.write_bso( + params::PutBso { + user_id: input.user_id.clone(), + collection: input.collection.clone(), + id: pbso.id.clone(), + payload: pbso.payload, + payload_link: pbso.payload_link, + sortindex: pbso.sortindex, + ttl: pbso.ttl, + }, + collection_id, + ) .await?; } self.update_collection(params::UpdateCollection { @@ -686,13 +603,23 @@ impl Db for MysqlDb { &mut self, params: params::GetQuotaUsage, ) -> DbResult { + if !self.quota.enabled { + return Ok(results::GetQuotaUsage::default()); + } + let uid = params.user_id.legacy_id as i64; - let (total_bytes, count): (i64, i32) = user_collections::table + let mut query = user_collections::table .select(( sql::("COALESCE(SUM(COALESCE(total_bytes, 0)), 0)"), sql::("COALESCE(SUM(COALESCE(count, 0)), 0)"), )) .filter(user_collections::user_id.eq(uid)) + .into_boxed(); + let limited_ids = self.get_quota_limited_collection_ids(); + if !limited_ids.is_empty() { + query = query.filter(user_collections::collection_id.ne_all(limited_ids)); + } + let (total_bytes, count): (i64, i32) = query .get_result(&mut self.conn) .await .optional()? @@ -779,6 +706,92 @@ impl Db for MysqlDb { enforced, } } + + #[cfg(debug_assertions)] + fn set_collection_limits(&mut self, limits: std::collections::HashMap) { + self.coll_limits = std::sync::Arc::new(limits); + } +} + +impl MysqlDb { + /// Write a single BSO row; callers should run `check_quota` first. + async fn write_bso( + &mut self, + bso: params::PutBso, + collection_id: i32, + ) -> DbResult { + let user_id: u64 = bso.user_id.legacy_id; + let timestamp = self.session.timestamp.as_i64(); + 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 q = format!( + r#" + INSERT INTO bso ({user_id}, {collection_id}, id, sortindex, payload, {modified}, {expiry}) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON DUPLICATE KEY UPDATE + {user_id} = VALUES({user_id}), + {collection_id} = VALUES({collection_id}), + id = VALUES(id) + "#, + user_id = USER_ID, + modified = MODIFIED, + collection_id = COLLECTION_ID, + expiry = EXPIRY + ); + let q = format!( + "{}{}", + q, + if bso.sortindex.is_some() { + ", sortindex = VALUES(sortindex)" + } else { + "" + }, + ); + let q = format!( + "{}{}", + q, + if bso.payload.is_some() { + ", payload = VALUES(payload)" + } else { + "" + }, + ); + let q = format!( + "{}{}", + q, + if bso.ttl.is_some() { + format!(", {expiry} = VALUES({expiry})", expiry = EXPIRY) + } else { + "".to_owned() + }, + ); + let q = format!( + "{}{}", + q, + if bso.payload.is_some() || bso.sortindex.is_some() { + format!(", {modified} = VALUES({modified})", modified = MODIFIED) + } else { + "".to_owned() + }, + ); + sql_query(q) + .bind::(user_id as i64) // XXX: + .bind::(&collection_id) + .bind::(&bso.id) + .bind::, _>(sortindex) + .bind::(payload) + .bind::(timestamp) + .bind::(timestamp + (i64::from(ttl) * 1000)) // remember: this is in millis + .execute(&mut self.conn) + .await?; + self.update_collection(params::UpdateCollection { + user_id: bso.user_id, + collection_id, + collection: bso.collection, + }) + .await + } } #[derive(Debug, QueryableByName)] diff --git a/syncstorage-mysql/src/db/mod.rs b/syncstorage-mysql/src/db/mod.rs index 4e25472f6f..d9b44aeb8c 100644 --- a/syncstorage-mysql/src/db/mod.rs +++ b/syncstorage-mysql/src/db/mod.rs @@ -18,7 +18,7 @@ use crate::{ DbError, DbResult, pool::{CollectionCache, Conn}, }; -use schema::{bso, collections, last_insert_id}; +use schema::{bso, collections, last_insert_id, user_collections}; mod batch_impl; mod db_impl; @@ -63,6 +63,8 @@ pub struct MysqlDb { session: MysqlDbSession, /// Pool level cache of collection_ids and their names coll_cache: Arc, + /// Resolved per-collection limits. + coll_limits: Arc>, metrics: Metrics, quota: Quota, } @@ -72,6 +74,7 @@ impl fmt::Debug for MysqlDb { f.debug_struct("MysqlDb") .field("session", &self.session) .field("coll_cache", &self.coll_cache) + .field("coll_limits", &self.coll_limits) .field("metrics", &self.metrics) .field("quota", &self.quota) .finish() @@ -84,11 +87,13 @@ impl MysqlDb { coll_cache: Arc, metrics: &Metrics, quota: &Quota, + coll_limits: Arc>, ) -> Self { MysqlDb { conn, session: Default::default(), coll_cache, + coll_limits, metrics: metrics.clone(), quota: *quota, } @@ -238,6 +243,49 @@ impl MysqlDb { }) } + pub(super) fn get_quota_limited_collection_ids(&self) -> Vec { + self.coll_limits.keys().copied().collect() + } + + pub(super) fn get_collection_limit(&self, collection_id: i32) -> Option { + self.coll_limits.get(&collection_id).copied() + } + + pub(super) fn emit_at_limit(&self, collection: &str, total_bytes: usize) -> DbResult<()> { + 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()); + } + warn!("Quota at limit for user ({total_bytes} bytes)"; "collection" => collection); + Ok(()) + } + + /// Single-collection usage to check a collection against its limit. + async fn get_collection_quota_usage( + &mut self, + user_id: &UserIdentifier, + collection_id: i32, + ) -> DbResult { + let uid = user_id.legacy_id as i64; + let (total_bytes, count): (i64, i32) = user_collections::table + .select(( + sql::("COALESCE(SUM(COALESCE(total_bytes, 0)), 0)"), + sql::("COALESCE(SUM(COALESCE(count, 0)), 0)"), + )) + .filter(user_collections::user_id.eq(uid)) + .filter(user_collections::collection_id.eq(collection_id)) + .get_result(&mut self.conn) + .await + .optional()? + .unwrap_or_default(); + Ok(results::GetQuotaUsage { + total_bytes: total_bytes as usize, + count, + }) + } + async fn check_quota( &mut self, user_id: &UserIdentifier, @@ -246,23 +294,35 @@ impl MysqlDb { 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 ({} bytes)", usage.total_bytes; - "collection" => collection); - } + + let collection_id = match self._get_collection_id(collection).await { + Ok(id) => Some(id), + Err(e) if e.is_collection_not_found() => None, + Err(e) => return Err(e), + }; + + let (usage, ceiling) = match collection_id + .and_then(|id| self.get_collection_limit(id).map(|limit| (id, limit))) + { + Some((collection_id, limit)) => ( + self.get_collection_quota_usage(user_id, collection_id) + .await?, + limit, + ), + None => ( + self.get_quota_usage(params::GetQuotaUsage { + user_id: user_id.clone(), + collection: collection.to_owned(), + }) + .await?, + self.quota.size, + ), + }; + + if usage.total_bytes >= ceiling { + self.emit_at_limit(collection, usage.total_bytes)?; } + Ok(Some(usage.total_bytes)) } } diff --git a/syncstorage-mysql/src/pool.rs b/syncstorage-mysql/src/pool.rs index 150aeeaa12..ae4fe5674a 100644 --- a/syncstorage-mysql/src/pool.rs +++ b/syncstorage-mysql/src/pool.rs @@ -39,6 +39,7 @@ pub struct MysqlDbPool { /// Thread Pool for running synchronous db calls /// In-memory cache of collection_ids and their names coll_cache: Arc, + coll_limits: Arc>, metrics: Metrics, quota: Quota, @@ -90,6 +91,7 @@ impl MysqlDbPool { Ok(Self { pool, coll_cache: Default::default(), + coll_limits: Arc::new(settings.limits.resolved_collection_limits(&STD_COLLS)), metrics: metrics.clone(), quota: Quota { size: settings.limits.max_quota_limit as usize, @@ -124,6 +126,7 @@ impl MysqlDbPool { Arc::clone(&self.coll_cache), &self.metrics, &self.quota, + Arc::clone(&self.coll_limits), )) } } @@ -229,13 +232,13 @@ impl Default for CollectionCache { by_name: RwLock::new( STD_COLLS .iter() - .map(|(k, v)| ((*v).to_owned(), *k)) + .map(|(k, v, _)| ((*v).to_owned(), *k)) .collect(), ), by_id: RwLock::new( STD_COLLS .iter() - .map(|(k, v)| (*k, (*v).to_owned())) + .map(|(k, v, _)| (*k, (*v).to_owned())) .collect(), ), } diff --git a/syncstorage-postgres/src/db/batch_impl.rs b/syncstorage-postgres/src/db/batch_impl.rs index 828d54bb75..048e1790cc 100644 --- a/syncstorage-postgres/src/db/batch_impl.rs +++ b/syncstorage-postgres/src/db/batch_impl.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; - use async_trait::async_trait; use diesel::{ self, ExpressionMethods, OptionalExtension, QueryDsl, delete, @@ -225,17 +223,14 @@ pub async fn do_append( let incoming_ids: Vec = 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 ({} bytes)", projected_total; - "collection" => collection); - } + let ceiling = db + .get_collection_limit(collection_id) + .unwrap_or(db.quota.size); + let projected_total = size + .saturating_add(pending_size) + .saturating_add(running_size); + if projected_total >= ceiling { + db.emit_at_limit(collection, projected_total)?; } } diff --git a/syncstorage-postgres/src/db/db_impl.rs b/syncstorage-postgres/src/db/db_impl.rs index d57e8632ff..d361e85431 100644 --- a/syncstorage-postgres/src/db/db_impl.rs +++ b/syncstorage-postgres/src/db/db_impl.rs @@ -265,12 +265,22 @@ impl Db for PgDb { &mut self, params: params::GetQuotaUsage, ) -> DbResult { - let (total_bytes, count): (i64, i64) = user_collections::table + if !self.quota.enabled { + return Ok(results::GetQuotaUsage::default()); + } + + let mut query = user_collections::table .select(( sql::("COALESCE(SUM(COALESCE(total_bytes, 0)), 0)::BIGINT"), sql::("COALESCE(SUM(COALESCE(count, 0)), 0)::BIGINT"), )) .filter(user_collections::user_id.eq(params.user_id.legacy_id as i64)) + .into_boxed(); + let limited_ids = self.get_quota_limited_collection_ids(); + if !limited_ids.is_empty() { + query = query.filter(user_collections::collection_id.ne_all(limited_ids)); + } + let (total_bytes, count): (i64, i64) = query .get_result(&mut self.conn) .await .optional()? @@ -586,6 +596,11 @@ impl Db for PgDb { enforced, } } + + #[cfg(debug_assertions)] + fn set_collection_limits(&mut self, limits: std::collections::HashMap) { + self.coll_limits = std::sync::Arc::new(limits); + } } #[derive(Debug, Queryable, Selectable)] diff --git a/syncstorage-postgres/src/db/mod.rs b/syncstorage-postgres/src/db/mod.rs index dabe154ead..06090b43d0 100644 --- a/syncstorage-postgres/src/db/mod.rs +++ b/syncstorage-postgres/src/db/mod.rs @@ -46,6 +46,8 @@ pub struct PgDb { session: PgDbSession, /// Pool level cache of collection_ids and their names. coll_cache: Arc, + /// Resolved per-collection limits. + coll_limits: Arc>, /// Configured quota, with defined size, enabled, and enforced attributes. metrics: Metrics, /// Configured quota, with defined size, enabled, and enforced attributes. @@ -57,6 +59,7 @@ impl fmt::Debug for PgDb { f.debug_struct("PgDb") .field("session", &self.session) .field("coll_cache", &self.coll_cache) + .field("coll_limits", &self.coll_limits) .field("metrics", &self.metrics) .field("quota", &self.quota) .finish() @@ -86,11 +89,13 @@ impl PgDb { coll_cache: Arc, metrics: &Metrics, quota: &Quota, + coll_limits: Arc>, ) -> Self { PgDb { conn, session: Default::default(), coll_cache, + coll_limits, metrics: metrics.clone(), quota: *quota, } @@ -235,23 +240,35 @@ impl PgDb { 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 ({} bytes)", usage.total_bytes; - "collection" => collection); - } + + let collection_id = match self._get_collection_id(collection).await { + Ok(id) => Some(id), + Err(e) if e.is_collection_not_found() => None, + Err(e) => return Err(e), + }; + + let (usage, ceiling) = match collection_id + .and_then(|id| self.get_collection_limit(id).map(|limit| (id, limit))) + { + Some((collection_id, limit)) => ( + self.get_collection_quota_usage(user_id, collection_id) + .await?, + limit, + ), + None => ( + self.get_quota_usage(params::GetQuotaUsage { + user_id: user_id.clone(), + collection: collection.to_owned(), + }) + .await?, + self.quota.size, + ), + }; + + if usage.total_bytes >= ceiling { + self.emit_at_limit(collection, usage.total_bytes)?; } + Ok(Some(usage.total_bytes)) } @@ -279,6 +296,47 @@ impl PgDb { }) } + pub(super) fn get_quota_limited_collection_ids(&self) -> Vec { + self.coll_limits.keys().copied().collect() + } + + pub(super) fn get_collection_limit(&self, collection_id: i32) -> Option { + self.coll_limits.get(&collection_id).copied() + } + + pub(super) fn emit_at_limit(&self, collection: &str, total_bytes: usize) -> DbResult<()> { + 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()); + } + warn!("Quota at limit for user ({total_bytes} bytes)"; "collection" => collection); + Ok(()) + } + + async fn get_collection_quota_usage( + &mut self, + user_id: &UserIdentifier, + collection_id: i32, + ) -> DbResult { + let (total_bytes, count): (i64, i64) = user_collections::table + .select(( + sql::("COALESCE(SUM(COALESCE(total_bytes, 0)), 0)::BIGINT"), + sql::("COALESCE(SUM(COALESCE(count, 0)), 0)::BIGINT"), + )) + .filter(user_collections::user_id.eq(user_id.legacy_id as i64)) + .filter(user_collections::collection_id.eq(collection_id)) + .get_result(&mut self.conn) + .await + .optional()? + .unwrap_or_default(); + Ok(results::GetQuotaUsage { + total_bytes: total_bytes as usize, + count: count as i32, + }) + } + pub fn checked_timestamp(&self) -> DbResult { self.session .timestamp diff --git a/syncstorage-postgres/src/pool.rs b/syncstorage-postgres/src/pool.rs index da0ca1f803..a012aec56e 100644 --- a/syncstorage-postgres/src/pool.rs +++ b/syncstorage-postgres/src/pool.rs @@ -43,6 +43,8 @@ pub struct PgDbPool { metrics: Metrics, /// Configured quota, with defined size, enabled, and enforced attributes. quota: Quota, + /// Resolved per-collection byte limits keyed by collection id. + coll_limits: Arc>, } impl PgDbPool { @@ -98,6 +100,7 @@ impl PgDbPool { enabled: settings.enable_quota, enforced: settings.enforce_quota, }, + coll_limits: Arc::new(settings.limits.resolved_collection_limits(&STD_COLLS)), }) } @@ -125,6 +128,7 @@ impl PgDbPool { Arc::clone(&self.coll_cache), &self.metrics, &self.quota, + Arc::clone(&self.coll_limits), )) } } @@ -225,13 +229,13 @@ impl Default for CollectionCache { by_name: RwLock::new( STD_COLLS .iter() - .map(|(k, v)| ((*v).to_owned(), *k)) + .map(|(k, v, _)| ((*v).to_owned(), *k)) .collect(), ), by_id: RwLock::new( STD_COLLS .iter() - .map(|(k, v)| (*k, (*v).to_owned())) + .map(|(k, v, _)| (*k, (*v).to_owned())) .collect(), ), } diff --git a/syncstorage-settings/src/lib.rs b/syncstorage-settings/src/lib.rs index 13d0630a3a..9d87efdff5 100644 --- a/syncstorage-settings/src/lib.rs +++ b/syncstorage-settings/src/lib.rs @@ -2,6 +2,7 @@ use std::{ cmp::min, + collections::HashMap, time::{Duration, Instant}, }; @@ -96,12 +97,11 @@ pub struct Settings { /// StatsD metrics label prefix for syncstorage. Default: "syncstorage". pub statsd_label: String, - /// Track per-user storage quota. Spanner-only; force-disabled on other - /// backends. Default: false. + /// Track per-user storage quota and configured per-collection quota + /// limits. Default: false. pub enable_quota: bool, - /// Reject writes that exceed `limits.max_quota_limit`. Requires - /// `enable_quota`. Spanner-only; force-disabled on other backends. - /// Default: false. + /// Reject writes that exceed `limits.max_quota_limit` or configured + /// per-collection quota limits. Requires `enable_quota`. Default: false. pub enforce_quota: bool, /// Whether Glean telemetry metric emission is enabled. @@ -221,6 +221,58 @@ pub struct ServerLimits { /// Maximum BSO count across a batch upload. pub max_total_records: u32, pub max_quota_limit: u64, + + /// Optional per-collection limits, keyed by collection name. Overrides the defaults carried in + /// `STD_COLLS`. + #[serde(default, skip_serializing)] + pub collections: HashMap, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct CollectionLimits { + pub max_quota_limit: usize, +} + +impl ServerLimits { + /// Resolve the per-collection limits. + pub fn resolved_collection_limits( + &self, + std_colls: &[(i32, &'static str, Option)], + ) -> HashMap { + self.resolve_collection_limits(std_colls) + .map(|(id, _, limit)| (id, limit)) + .collect() + } + + /// Resolved limits by collection name. + pub fn resolved_collection_limits_by_name( + &self, + std_colls: &[(i32, &'static str, Option)], + ) -> HashMap { + self.resolve_collection_limits(std_colls) + .map(|(_, name, limit)| { + ( + name.to_owned(), + CollectionLimits { + max_quota_limit: limit, + }, + ) + }) + .collect() + } + + fn resolve_collection_limits<'a>( + &'a self, + std_colls: &'a [(i32, &'static str, Option)], + ) -> impl Iterator + 'a { + std_colls.iter().filter_map(|(id, name, default)| { + self.collections + .get(*name) + .copied() + .or(*default) + .map(|limit| (*id, *name, limit)) + }) + } } impl Default for ServerLimits { @@ -234,6 +286,66 @@ impl Default for ServerLimits { max_total_bytes: DEFAULT_MAX_TOTAL_BYTES, max_total_records: DEFAULT_MAX_TOTAL_RECORDS, max_quota_limit: DEFAULT_MAX_QUOTA_LIMIT, + collections: HashMap::new(), } } } + +#[cfg(test)] +mod tests { + use super::*; + + const STD_COLLS: &[(i32, &str, Option)] = &[ + (4, "history", None), + (9, "tabs", Some(5_242_880)), + (10, "passwords", None), + ]; + + fn limits_with(collections: &[(&str, usize)]) -> ServerLimits { + ServerLimits { + collections: collections + .iter() + .map(|&(n, v)| (n.to_owned(), v)) + .collect(), + ..Default::default() + } + } + + #[test] + fn resolver_settings_override_std_colls() { + let limits = limits_with(&[("tabs", 100), ("history", 200)]); + let resolved = limits.resolved_collection_limits(STD_COLLS); + assert_eq!(resolved.get(&9), Some(&100)); // tabs + assert_eq!(resolved.get(&4), Some(&200)); // history + assert_eq!(resolved.get(&10), None); // passwords, no override + } + + #[test] + fn resolver_falls_back_to_std_default() { + let limits = limits_with(&[]); + let resolved = limits.resolved_collection_limits(STD_COLLS); + assert_eq!(resolved.get(&9), Some(&5_242_880)); // tabs + assert_eq!(resolved.get(&4), None); // history + assert_eq!(resolved.len(), 1); // only tabs + } + + #[test] + fn resolver_by_name_matches_by_id() { + let limits = limits_with(&[("history", 200)]); + let by_name = limits.resolved_collection_limits_by_name(STD_COLLS); + assert_eq!( + by_name.get("history"), + Some(&CollectionLimits { + max_quota_limit: 200 + }) + ); + assert_eq!( + by_name.get("tabs"), + Some(&CollectionLimits { + max_quota_limit: 5_242_880 + }) + ); + assert_eq!(by_name.get("passwords"), None); + assert_eq!(by_name.len(), 2); + } +} diff --git a/syncstorage-spanner/src/db/batch_impl.rs b/syncstorage-spanner/src/db/batch_impl.rs index ad4f6c00bc..18515f82e2 100644 --- a/syncstorage-spanner/src/db/batch_impl.rs +++ b/syncstorage-spanner/src/db/batch_impl.rs @@ -294,13 +294,14 @@ pub async fn do_append( { let pending_size = pending_batch_size(db, &user_id, collection_id, &batch.id, &incoming_ids).await?; - let projected_total = size + pending_size + running_size; - if projected_total >= db.quota.size { - if db.quota.enforced { - return Err(db.quota_error(collection)); - } else { - warn!("Quota at limit for user ({} bytes)", projected_total; "collection"=>collection); - } + let ceiling = db + .get_collection_limit(collection_id) + .unwrap_or(db.quota.size); + let projected_total = size + .saturating_add(pending_size) + .saturating_add(running_size); + if projected_total >= ceiling { + db.emit_at_limit(collection, projected_total)?; } } diff --git a/syncstorage-spanner/src/db/db_impl.rs b/syncstorage-spanner/src/db/db_impl.rs index 730660ff95..d2011760ea 100644 --- a/syncstorage-spanner/src/db/db_impl.rs +++ b/syncstorage-spanner/src/db/db_impl.rs @@ -411,40 +411,29 @@ impl Db for SpannerDb { if !self.quota.enabled { return Ok(results::GetQuotaUsage::default()); } - let check_sql = "SELECT COALESCE(SUM(COALESCE(total_bytes, 0)), 0), \ + // account-wide usage + let mut check_sql = "SELECT COALESCE(SUM(COALESCE(total_bytes, 0)), 0), \ COALESCE(SUM(COALESCE(count, 0)), 0) FROM user_collections WHERE fxa_uid = @fxa_uid - AND fxa_kid = @fxa_kid"; + AND fxa_kid = @fxa_kid" + .to_owned(); + let mut limited_ids = self.get_quota_limited_collection_ids(); + if !limited_ids.is_empty() { + limited_ids.sort_unstable(); + let ids = limited_ids + .iter() + .map(|id| id.to_string()) + .collect::>() + .join(", "); + check_sql.push_str(&format!(" AND collection_id NOT IN ({ids})")); + } let (sqlparams, sqlparam_types) = params! { "fxa_uid" => params.user_id.fxa_uid.clone(), "fxa_kid" => params.user_id.fxa_kid.clone(), }; - let result = self - .sql(check_sql) - .await? - .params(sqlparams) - .param_types(sqlparam_types) - .execute(&self.conn)? - .one_or_none() - .await?; - if let Some(result) = result { - let total_bytes = if self.quota.enabled { - result[0] - .get_string_value() - .parse::() - .map_err(|e| DbError::integrity(e.to_string()))? - } else { - 0 - }; - let count = result[1] - .get_string_value() - .parse::() - .map_err(|e| DbError::integrity(e.to_string()))?; - Ok(results::GetQuotaUsage { total_bytes, count }) - } else { - Ok(results::GetQuotaUsage::default()) - } + self.query_quota_usage(&check_sql, sqlparams, sqlparam_types) + .await } async fn delete_storage(&mut self, user_id: params::DeleteStorage) -> DbResult<()> { @@ -812,15 +801,17 @@ impl Db for SpannerDb { let user_id = params.user_id; let collection_id = self.get_or_create_collection_id(¶ms.collection).await?; - if !params.for_batch { - self.check_quota(&user_id, ¶ms.collection).await?; - } - let load_size: usize = params .bsos .iter() .map(|bso| bso.payload.as_ref().map_or(0, String::len)) .sum(); + + // batch commits enforce quota in do_append + if !params.for_batch { + self.check_quota(&user_id, ¶ms.collection).await?; + } + if load_size > MAX_SPANNER_LOAD_SIZE { self.metrics.clone().incr("error.tooMuchData"); trace!( @@ -915,6 +906,11 @@ impl Db for SpannerDb { enforced, }; } + + #[cfg(debug_assertions)] + fn set_collection_limits(&mut self, limits: std::collections::HashMap) { + self.coll_limits = std::sync::Arc::new(limits); + } } fn sync_timestamp_from_rfc3339(val: &str) -> Result { diff --git a/syncstorage-spanner/src/db/mod.rs b/syncstorage-spanner/src/db/mod.rs index cc28a6ff1f..80d78e89be 100644 --- a/syncstorage-spanner/src/db/mod.rs +++ b/syncstorage-spanner/src/db/mod.rs @@ -13,7 +13,7 @@ use syncserver_common::Metrics; use syncstorage_db_common::{ DEFAULT_BSO_TTL, Db, FIRST_CUSTOM_COLLECTION_ID, Sorting, UserIdentifier, error::DbErrorIntrospect, - params, + params, results, util::{SyncTimestamp, to_rfc3339}, }; use syncstorage_settings::Quota; @@ -50,6 +50,8 @@ pub struct SpannerDb { /// Pool level cache of collection_ids and their names coll_cache: Arc, + /// Resolved per-collection limits by id + pub coll_limits: Arc>, pub metrics: Metrics, pub quota: Quota, @@ -89,11 +91,13 @@ impl SpannerDb { coll_cache: Arc, metrics: &Metrics, quota: Quota, + coll_limits: Arc>, ) -> Self { SpannerDb { conn, session: Default::default(), coll_cache, + coll_limits, metrics: metrics.clone(), quota, } @@ -530,12 +534,23 @@ impl SpannerDb { */ } - pub fn quota_error(&self, collection: &str) -> DbError { - // return the over quota error. + pub(super) fn get_quota_limited_collection_ids(&self) -> Vec { + self.coll_limits.keys().copied().collect() + } + + pub(super) fn get_collection_limit(&self, collection_id: i32) -> Option { + self.coll_limits.get(&collection_id).copied() + } + + pub(super) fn emit_at_limit(&self, collection: &str, total_bytes: usize) -> DbResult<()> { let mut tags = HashMap::default(); tags.insert("collection".to_owned(), collection.to_owned()); self.metrics.incr_with_tags("storage.quota.at_limit", tags); - DbError::quota() + if self.quota.enforced { + return Err(DbError::quota()); + } + warn!("Quota at limit for user ({total_bytes} bytes)"; "collection" => collection); + Ok(()) } pub(super) async fn check_quota( @@ -547,22 +562,87 @@ impl SpannerDb { 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 { - if self.quota.enforced { - return Err(self.quota_error(collection)); - } else { - warn!("Quota at limit for user ({} bytes)", usage.total_bytes; "collection"=>collection); - } + let collection_id = match self._get_collection_id(collection).await { + Ok(id) => Some(id), + Err(e) if e.is_collection_not_found() => None, + Err(e) => return Err(e), + }; + let collection_limit = collection_id.and_then(|id| self.coll_limits.get(&id).copied()); + let (usage, ceiling) = match (collection_id, collection_limit) { + (Some(collection_id), Some(limit)) => ( + self.get_collection_quota_usage(user_id, collection_id) + .await?, + limit, + ), + _ => ( + self.get_quota_usage(params::GetQuotaUsage { + user_id: user_id.clone(), + collection: collection.to_owned(), + }) + .await?, + self.quota.size, + ), + }; + + if usage.total_bytes >= ceiling { + self.emit_at_limit(collection, usage.total_bytes)?; } + Ok(Some(usage.total_bytes)) } + pub(super) async fn query_quota_usage( + &mut self, + sql: &str, + sqlparams: HashMap, + sqlparam_types: HashMap, + ) -> DbResult { + let result = self + .sql(sql) + .await? + .params(sqlparams) + .param_types(sqlparam_types) + .execute(&self.conn)? + .one_or_none() + .await?; + if let Some(result) = result { + let total_bytes = result[0] + .get_string_value() + .parse::() + .map_err(|e| DbError::integrity(e.to_string()))?; + let count = result[1] + .get_string_value() + .parse::() + .map_err(|e| DbError::integrity(e.to_string()))?; + Ok(results::GetQuotaUsage { total_bytes, count }) + } else { + Ok(results::GetQuotaUsage::default()) + } + } + + pub(super) async fn get_collection_quota_usage( + &mut self, + user_id: &UserIdentifier, + collection_id: i32, + ) -> DbResult { + if !self.quota.enabled { + return Ok(results::GetQuotaUsage::default()); + } + let check_sql = "SELECT COALESCE(SUM(COALESCE(total_bytes, 0)), 0), \ + COALESCE(SUM(COALESCE(count, 0)), 0) + FROM user_collections + WHERE fxa_uid = @fxa_uid + AND fxa_kid = @fxa_kid + AND collection_id = @collection_id"; + let (sqlparams, sqlparam_types) = params! { + "fxa_uid" => user_id.fxa_uid.clone(), + "fxa_kid" => user_id.fxa_kid.clone(), + "collection_id" => collection_id, + }; + self.query_quota_usage(check_sql, sqlparams, sqlparam_types) + .await + } + /// Write a bso using an `INSERT OR UPDATE`. async fn put_bso_dml( &mut self, diff --git a/syncstorage-spanner/src/pool.rs b/syncstorage-spanner/src/pool.rs index 3fa821d90b..6a72ee0871 100644 --- a/syncstorage-spanner/src/pool.rs +++ b/syncstorage-spanner/src/pool.rs @@ -17,6 +17,8 @@ pub struct SpannerDbPool { pool: deadpool::managed::Pool, /// In-memory cache of collection_ids and their names coll_cache: Arc, + /// Resolved per-collection limits + coll_limits: Arc>, metrics: Metrics, quota: Quota, @@ -69,6 +71,7 @@ impl SpannerDbPool { enabled: settings.enable_quota, enforced: settings.enforce_quota, }, + coll_limits: Arc::new(settings.limits.resolved_collection_limits(&STD_COLLS)), }) } @@ -85,6 +88,7 @@ impl SpannerDbPool { Arc::clone(&self.coll_cache), &self.metrics, self.quota, + Arc::clone(&self.coll_limits), )) } @@ -203,13 +207,13 @@ impl Default for CollectionCache { by_name: RwLock::new( STD_COLLS .iter() - .map(|(k, v)| ((*v).to_owned(), *k)) + .map(|(k, v, _)| ((*v).to_owned(), *k)) .collect(), ), by_id: RwLock::new( STD_COLLS .iter() - .map(|(k, v)| (*k, (*v).to_owned())) + .map(|(k, v, _)| (*k, (*v).to_owned())) .collect(), ), }