diff --git a/src/db/mock.rs b/src/db/mock.rs index b9b52c62fe..85aa7ed5b9 100644 --- a/src/db/mock.rs +++ b/src/db/mock.rs @@ -121,7 +121,7 @@ impl<'a> Db<'a> for MockDb { fn clear_coll_cache(&self) {} #[cfg(test)] - fn set_quota(&mut self, _: bool, _: usize) {} + fn set_quota(&mut self, _: bool, _: usize, _: bool) {} } unsafe impl Send for MockDb {} diff --git a/src/db/mod.rs b/src/db/mod.rs index 4584405a71..ea942a0c1b 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -239,7 +239,7 @@ pub trait Db<'a>: Debug + 'a { fn clear_coll_cache(&self); #[cfg(test)] - fn set_quota(&mut self, enabled: bool, limit: usize); + fn set_quota(&mut self, enabled: bool, limit: usize, enforce: bool); } impl<'a> Clone for Box> { diff --git a/src/db/mysql/models.rs b/src/db/mysql/models.rs index 43b6e3b9c1..57296c54f5 100644 --- a/src/db/mysql/models.rs +++ b/src/db/mysql/models.rs @@ -31,6 +31,7 @@ use crate::db::{ Db, DbFuture, Sorting, }; use crate::server::metrics::Metrics; +use crate::settings::Quota; use crate::web::extractors::{BsoQueryParams, HawkIdentifier}; use crate::web::tags::Tags; @@ -86,8 +87,7 @@ pub struct MysqlDb { coll_cache: Arc, pub metrics: Metrics, - pub quota: usize, - pub quota_enabled: bool, + pub quota: Quota, } /// Despite the db conn structs being !Sync (see Arc above) we @@ -123,8 +123,7 @@ impl MysqlDb { conn: Conn, coll_cache: Arc, metrics: &Metrics, - quota: &usize, - quota_enabled: bool, + quota: &Quota, ) -> Self { let inner = MysqlDbInner { #[cfg(not(test))] @@ -138,7 +137,6 @@ impl MysqlDb { coll_cache, metrics: metrics.clone(), quota: *quota, - quota_enabled, } } @@ -398,18 +396,23 @@ impl MysqlDb { let collection_id = self.get_or_create_collection_id(&bso.collection)?; let user_id: u64 = bso.user_id.legacy_id; let timestamp = self.timestamp().as_i64(); - if self.quota_enabled { + if self.quota.enabled { let usage = self.get_quota_usage_sync(params::GetQuotaUsage { user_id: HawkIdentifier::new_legacy(user_id), collection: bso.collection.clone(), collection_id, })?; - if usage.total_bytes >= self.quota as usize { + if usage.total_bytes >= self.quota.size as usize { let mut tags = Tags::default(); - tags.tags.insert("collection".to_owned(), bso.collection); + tags.tags + .insert("collection".to_owned(), bso.collection.clone()); self.metrics .incr_with_tags("storage.quota.at_limit", Some(tags)); - return Err(DbErrorKind::Quota.into()); + if self.quota.enforced { + return Err(DbErrorKind::Quota.into()); + } else { + warn!("Quota at limit for user's collection ({} bytes)", usage.total_bytes; "collection"=>bso.collection.clone()); + } } } @@ -829,7 +832,7 @@ impl MysqlDb { user_id: u32, collection_id: i32, ) -> Result { - let quota = if self.quota_enabled { + let quota = if self.quota.enabled { self.calc_quota_usage_sync(user_id, collection_id)? } else { results::GetQuotaUsage { @@ -1114,9 +1117,12 @@ impl<'a> Db<'a> for MysqlDb { } #[cfg(test)] - fn set_quota(&mut self, enabled: bool, limit: usize) { - self.quota = limit; - self.quota_enabled = enabled; + fn set_quota(&mut self, enabled: bool, limit: usize, enforced: bool) { + self.quota = Quota { + size: limit, + enabled, + enforced, + } } } diff --git a/src/db/mysql/pool.rs b/src/db/mysql/pool.rs index dff66417e4..3738c25a43 100644 --- a/src/db/mysql/pool.rs +++ b/src/db/mysql/pool.rs @@ -26,7 +26,7 @@ use crate::db::{ }; use crate::error::{ApiError, ApiResult}; use crate::server::metrics::Metrics; -use crate::settings::Settings; +use crate::settings::{Quota, Settings}; embed_migrations!(); @@ -54,8 +54,7 @@ pub struct MysqlDbPool { coll_cache: Arc, metrics: Metrics, - quota: usize, - quota_enabled: bool, + quota: Quota, } impl MysqlDbPool { @@ -84,8 +83,11 @@ impl MysqlDbPool { pool: builder.build(manager)?, coll_cache: Default::default(), metrics: metrics.clone(), - quota: settings.limits.max_quota_limit as usize, - quota_enabled: settings.enable_quota, + quota: Quota { + size: settings.limits.max_quota_limit as usize, + enabled: settings.enable_quota, + enforced: settings.enforce_quota, + }, }) } @@ -95,7 +97,6 @@ impl MysqlDbPool { Arc::clone(&self.coll_cache), &self.metrics, &self.quota, - self.quota_enabled, )) } } diff --git a/src/db/spanner/batch.rs b/src/db/spanner/batch.rs index a5c8f2e66f..41ccd275fc 100644 --- a/src/db/spanner/batch.rs +++ b/src/db/spanner/batch.rs @@ -363,10 +363,14 @@ pub async fn do_append_async( }; } - if db.quota_enabled { + if db.quota.enabled { if let Some(size) = batch.size { - if size + running_size >= (db.quota as usize) { - return Err(db.quota_error(collection)); + if size + running_size >= db.quota.size { + if db.quota.enforced { + return Err(db.quota_error(collection)); + } else { + warn!("Quota at limit for user's collection ({} bytes)", size + running_size; "collection"=>collection); + } } } } @@ -510,7 +514,7 @@ async fn pretouch_collection_async( .await?; if result.is_none() { sqlparams.insert("modified".to_owned(), as_value(PRETOUCH_TS.to_owned())); - let sql = if db.quota_enabled { + let sql = if db.quota.enabled { "INSERT INTO user_collections (fxa_uid, fxa_kid, collection_id, modified, count, total_bytes) VALUES (@fxa_uid, @fxa_kid, @collection_id, @modified, 0, 0)" } else { diff --git a/src/db/spanner/models.rs b/src/db/spanner/models.rs index c9dff91646..af330b6c3f 100644 --- a/src/db/spanner/models.rs +++ b/src/db/spanner/models.rs @@ -31,6 +31,7 @@ use crate::{ Db, DbFuture, Sorting, FIRST_CUSTOM_COLLECTION_ID, }, server::metrics::Metrics, + settings::Quota, web::{ extractors::{BsoQueryParams, HawkIdentifier, Offset}, tags::Tags, @@ -92,8 +93,7 @@ pub struct SpannerDb { coll_cache: Arc, pub metrics: Metrics, - pub quota: usize, - pub quota_enabled: bool, + pub quota: Quota, } pub struct SpannerDbInner { @@ -121,8 +121,7 @@ impl SpannerDb { conn: Conn, coll_cache: Arc, metrics: &Metrics, - quota: usize, - quota_enabled: bool, + quota: Quota, ) -> Self { let inner = SpannerDbInner { conn, @@ -133,7 +132,6 @@ impl SpannerDb { coll_cache, metrics: metrics.clone(), quota, - quota_enabled, } } @@ -801,7 +799,7 @@ impl SpannerDb { &self, params: params::GetQuotaUsage, ) -> Result { - if !self.quota_enabled { + if !self.quota.enabled { return Ok(results::GetQuotaUsage::default()); } let check_sql = "SELECT COALESCE(total_bytes,0), COALESCE(count,0) @@ -820,7 +818,7 @@ impl SpannerDb { .one_or_none() .await?; if let Some(result) = result { - let total_bytes = if self.quota_enabled { + let total_bytes = if self.quota.enabled { result[0] .get_string_value() .parse::() @@ -862,7 +860,7 @@ impl SpannerDb { self.metrics .clone() .start_timer("storage.quota.update_existing_totals", None); - let calc_sql = if self.quota_enabled { + let calc_sql = if self.quota.enabled { "SELECT SUM(BYTE_LENGTH(payload)), COUNT(*) FROM bsos WHERE fxa_uid = @fxa_uid @@ -891,7 +889,7 @@ impl SpannerDb { // Update the user_collections table to reflect current numbers. // If there are BSOs, there are user_collections (or else something // really bad already happened.) - if self.quota_enabled { + if self.quota.enabled { sqlparams.insert( "total_bytes".to_owned(), as_value(result[0].take_string_value()), @@ -934,7 +932,7 @@ impl SpannerDb { .await?; if result.is_none() { // No collections, so insert what we've got. - if self.quota_enabled { + if self.quota.enabled { "INSERT INTO user_collections (fxa_uid, fxa_kid, collection_id, modified, total_bytes, count) VALUES (@fxa_uid, @fxa_kid, @collection_id, @modified, 0, 0)" } else { @@ -944,7 +942,7 @@ impl SpannerDb { } else { // there are collections, best modify what's there. // NOTE, tombstone is a single collection id, it would have been created above. - if self.quota_enabled { + if self.quota.enabled { "UPDATE user_collections SET modified=@modified, total_bytes=0, count=0 WHERE fxa_uid=@fxa_uid AND fxa_kid=@fxa_kid AND collection_id=@collection_id" } else { @@ -1119,7 +1117,7 @@ impl SpannerDb { self.metrics .clone() .start_timer("storage.quota.init_totals", Some(tags)); - let update_sql = if self.quota_enabled { + let update_sql = if self.quota.enabled { "INSERT INTO user_collections (fxa_uid, fxa_kid, collection_id, modified, count, total_bytes) VALUES (@fxa_uid, @fxa_kid, @collection_id, @modified, 0, 0)" } else { @@ -1630,7 +1628,7 @@ impl SpannerDb { collection_id: i32, ) -> Result> { // duplicate quota trap in test func below. - if !self.quota_enabled { + if !self.quota.enabled { return Ok(None); } let usage = self @@ -1640,8 +1638,12 @@ impl SpannerDb { collection_id, }) .await?; - if usage.total_bytes >= self.quota { - return Err(self.quota_error(collection)); + if usage.total_bytes >= self.quota.size { + if self.quota.enforced { + return Err(self.quota_error(collection)); + } else { + warn!("Quota at limit for user's collection: ({} bytes)", usage.total_bytes; "collection"=>collection); + } } Ok(Some(usage.total_bytes as usize)) } @@ -2096,8 +2098,11 @@ impl<'a> Db<'a> for SpannerDb { } #[cfg(test)] - fn set_quota(&mut self, enabled: bool, limit: usize) { - self.quota_enabled = enabled; - self.quota = limit; + fn set_quota(&mut self, enabled: bool, limit: usize, enforced: bool) { + self.quota = Quota { + size: limit, + enabled, + enforced, + }; } } diff --git a/src/db/spanner/pool.rs b/src/db/spanner/pool.rs index 994ca22c8f..5be8bb14f4 100644 --- a/src/db/spanner/pool.rs +++ b/src/db/spanner/pool.rs @@ -10,7 +10,7 @@ use std::{ use super::models::Result; use crate::db::{error::DbError, results, Db, DbPool, STD_COLLS}; use crate::server::metrics::Metrics; -use crate::settings::Settings; +use crate::settings::{Quota, Settings}; use super::manager::{SpannerSession, SpannerSessionManager}; use super::models::SpannerDb; @@ -37,8 +37,7 @@ pub struct SpannerDbPool { coll_cache: Arc, metrics: Metrics, - quota: usize, - quota_enabled: bool, + quota: Quota, } impl SpannerDbPool { @@ -58,8 +57,11 @@ impl SpannerDbPool { pool, coll_cache: Default::default(), metrics: metrics.clone(), - quota: settings.limits.max_quota_limit as usize, - quota_enabled: settings.enable_quota, + quota: Quota { + size: settings.limits.max_quota_limit as usize, + enabled: settings.enable_quota, + enforced: settings.enforce_quota, + }, }) } @@ -75,7 +77,6 @@ impl SpannerDbPool { Arc::clone(&self.coll_cache), &self.metrics, self.quota, - self.quota_enabled, )) } } diff --git a/src/db/tests/batch.rs b/src/db/tests/batch.rs index 8dabb75c5a..2b334a57e2 100644 --- a/src/db/tests/batch.rs +++ b/src/db/tests/batch.rs @@ -173,7 +173,7 @@ async fn quota_test_create_batch() -> Result<()> { let limit = 300; settings.limits.max_quota_limit = limit; - let pool = db_pool(Some(settings)).await?; + let pool = db_pool(Some(settings.clone())).await?; let db = test_db(pool.as_ref()).await?; let uid = 1; @@ -193,7 +193,12 @@ async fn quota_test_create_batch() -> Result<()> { }) .await?; - assert!(db.create_batch(cb(uid, coll, bsos2)).await.is_err()); + let result = db.create_batch(cb(uid, coll, bsos2)).await; + if settings.enforce_quota { + assert!(result.is_err()); + } else { + assert!(result.is_ok()); + } Ok(()) } @@ -210,7 +215,7 @@ async fn quota_test_append_batch() -> Result<()> { let limit = 300; settings.limits.max_quota_limit = limit; - let pool = db_pool(Some(settings)).await?; + let pool = db_pool(Some(settings.clone())).await?; let db = test_db(pool.as_ref()).await?; let uid = 1; @@ -234,11 +239,10 @@ async fn quota_test_append_batch() -> Result<()> { }) .await?; let id2 = db.create_batch(cb(uid, coll, bsos2)).await?; - assert!(db - .append_to_batch(ab(uid, coll, id2.clone(), bsos3)) - .await - .is_err()); - + let result = db.append_to_batch(ab(uid, coll, id2.clone(), bsos3)).await; + if settings.enforce_quota { + assert!(result.is_err()) + } Ok(()) } diff --git a/src/db/tests/db.rs b/src/db/tests/db.rs index 39e76c74b4..68ccc08b9f 100644 --- a/src/db/tests/db.rs +++ b/src/db/tests/db.rs @@ -681,7 +681,7 @@ async fn test_quota() -> Result<()> { .sample_iter(&Alphanumeric) .take(size) .collect::(); - db.set_quota(false, 0); + db.set_quota(false, 0, false); // These should work db.put_bso(pbso(uid, coll, "100", Some(&payload), None, None)) @@ -689,7 +689,7 @@ async fn test_quota() -> Result<()> { db.put_bso(pbso(uid, coll, "101", Some(&payload), None, None)) .await?; - db.set_quota(true, size * 2); + db.set_quota(true, size * 2, true); // Allow the put, but calculate the quota db.put_bso(pbso(uid, coll, "102", Some(&payload), None, None)) diff --git a/src/settings.rs b/src/settings.rs index a74cfdcbcb..39e6061e48 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -25,6 +25,13 @@ static DEFAULT_MAX_TOTAL_RECORDS: u32 = 100 * DEFAULT_MAX_POST_RECORDS; static DEFAULT_MAX_QUOTA_LIMIT: u32 = 2 * GIGABYTE; static PREFIX: &str = "sync"; +#[derive(Clone, Debug, Default, Copy)] +pub struct Quota { + pub size: usize, + pub enabled: bool, + pub enforced: bool, +} + #[derive(Clone, Debug, Deserialize)] pub struct Settings { pub debug: bool, @@ -53,6 +60,7 @@ pub struct Settings { pub statsd_label: String, pub enable_quota: bool, + pub enforce_quota: bool, } impl Default for Settings { @@ -74,6 +82,7 @@ impl Default for Settings { statsd_label: "syncstorage".to_string(), human_logs: false, enable_quota: false, + enforce_quota: false, } } } @@ -115,6 +124,7 @@ impl Settings { s.set_default("statsd_port", 8125)?; s.set_default("statsd_label", "syncstorage")?; s.set_default("enable_quota", false)?; + s.set_default("enforce_quota", false)?; // Merge the config file if supplied if let Some(config_filename) = filename { @@ -149,13 +159,13 @@ impl Settings { env::set_var("ACTIX_THREADPOOL", database_pool_max_size.to_string()); } } - // No quotas for stand alone servers - s.limits.max_quota_limit = 0; - s.enable_quota = false; } if s.limits.max_quota_limit == 0 { s.enable_quota = false } + if s.enforce_quota { + s.enable_quota = true + } s } Err(e) => match e { @@ -194,7 +204,11 @@ impl Settings { /// A simple banner for display of certain settings at startup pub fn banner(&self) -> String { let quota = if self.enable_quota { - format!("Quota: {} bytes", self.limits.max_quota_limit) + format!( + "Quota: {} bytes ({}enforced)", + self.limits.max_quota_limit, + if !self.enforce_quota { "un" } else { "" } + ) } else { "No quota".to_owned() };