diff --git a/crates/manager-core/src/group_docs_repo.rs b/crates/manager-core/src/group_docs_repo.rs index 1f13143..5ecb828 100644 --- a/crates/manager-core/src/group_docs_repo.rs +++ b/crates/manager-core/src/group_docs_repo.rs @@ -95,6 +95,23 @@ pub trait GroupDocsRepo: Send + Sync { let _ = group_id; Ok(0) } + + /// §11.6 M4 quota: the PROJECTED total stored bytes for the group AFTER + /// writing `new_data` — current SUM, minus the existing doc's bytes + /// (`replacing = Some(id)` on update; `None` on create), plus the new data's + /// bytes. Computed entirely in SQL so ALL three terms use the SAME canonical + /// metric (`octet_length(data::text)`); the Rust-side subtract/add previously + /// mixed the DB canonical SUM with a `serde_json` compact length, drifting the + /// ceiling permissive. Default `Ok(0)`. + async fn projected_total_bytes( + &self, + group_id: GroupId, + replacing: Option, + new_data: &serde_json::Value, + ) -> Result { + let _ = (group_id, replacing, new_data); + Ok(0) + } } pub struct PostgresGroupDocsRepo { @@ -299,6 +316,34 @@ impl GroupDocsRepo for PostgresGroupDocsRepo { .await?; Ok(u64::try_from(n).unwrap_or(0)) } + + async fn projected_total_bytes( + &self, + group_id: GroupId, + replacing: Option, + new_data: &serde_json::Value, + ) -> Result { + // All three terms use `octet_length(data::text)` (jsonb canonical). The + // new value is bound as compact TEXT and re-parsed through + // `::jsonb::text` so PG measures it in the SAME canonical form it stores. + // `replacing = None` (create) → the subtrahend row never matches → 0. + let new_text = serde_json::to_string(new_data).unwrap_or_else(|_| "null".to_string()); + let (n,): (i64,) = sqlx::query_as( + "SELECT ( \ + COALESCE((SELECT SUM(octet_length(data::text)) \ + FROM group_docs WHERE group_id = $1), 0) \ + - COALESCE((SELECT octet_length(data::text) \ + FROM group_docs WHERE group_id = $1 AND id = $2), 0) \ + + octet_length($3::jsonb::text) \ + )::BIGINT", + ) + .bind(group_id.into_inner()) + .bind(replacing) + .bind(new_text) + .fetch_one(&self.pool) + .await?; + Ok(u64::try_from(n).unwrap_or(0)) + } } fn encode_cursor(last_id: &Uuid) -> String { diff --git a/crates/manager-core/src/group_docs_service.rs b/crates/manager-core/src/group_docs_service.rs index efccde4..f1ab1cc 100644 --- a/crates/manager-core/src/group_docs_service.rs +++ b/crates/manager-core/src/group_docs_service.rs @@ -95,16 +95,21 @@ impl GroupDocsServiceImpl { async fn check_total_bytes( &self, group_id: GroupId, - new_len: u64, - old_len: u64, + replacing: Option, + new_data: &serde_json::Value, rows_after: u64, ) -> Result<(), GroupDocsError> { let upper_bound = rows_after.saturating_mul(self.max_value_bytes as u64); if upper_bound <= self.max_total_bytes { return Ok(()); } - let used = self.repo.total_bytes(group_id).await?; - let projected = used.saturating_sub(old_len) + new_len; + // Consistent-metric projection (canonical `octet_length(data::text)` for + // the current SUM, the replaced doc, and the new data) — no Rust-compact + // vs DB-canonical drift. + let projected = self + .repo + .projected_total_bytes(group_id, replacing, new_data) + .await?; if projected > self.max_total_bytes { return Err(GroupDocsError::TotalBytesQuotaExceeded { limit: self.max_total_bytes, @@ -241,9 +246,8 @@ impl GroupDocsService for GroupDocsServiceImpl { actual: usize::try_from(count).unwrap_or(usize::MAX), }); } - // §11.6 M4 byte quota: a create only adds bytes (old_len = 0), one row. - let new_len = serde_json::to_vec(&data).map_or(0u64, |b| b.len() as u64); - self.check_total_bytes(group_id, new_len, 0, count + 1) + // §11.6 M4 byte quota: a create only adds bytes (nothing replaced), one row. + self.check_total_bytes(group_id, None, &data, count + 1) .await?; let created = self.repo.create(group_id, collection, data.clone()).await?; self.emit_shared( @@ -309,20 +313,14 @@ impl GroupDocsService for GroupDocsServiceImpl { validate_data(&data)?; self.check_data_size(&data)?; self.check_write(cx, group_id).await?; - // §11.6 M4 byte quota: check the projected total before writing. Fetch the - // current doc for its old byte size (a missing doc → old_len 0, and the + // §11.6 M4 byte quota: check the projected total before writing. The + // projection subtracts the replaced doc's bytes in SQL, so there is no + // need to fetch the old doc here (a missing doc subtracts 0, and the // update below no-ops to `None`). - let old_len = self - .repo - .get(group_id, collection, id) - .await? - .and_then(|d| serde_json::to_vec(&d.data).ok()) - .map_or(0u64, |b| b.len() as u64); - let new_len = serde_json::to_vec(&data).map_or(0u64, |b| b.len() as u64); // An update leaves the row count unchanged; the cheap COUNT lets the byte // check skip the text SUM when the collection is comfortably under cap. let count = self.repo.count_rows(group_id).await?; - self.check_total_bytes(group_id, new_len, old_len, count) + self.check_total_bytes(group_id, Some(id), &data, count) .await?; match self .repo @@ -516,6 +514,27 @@ mod tests { .map(|v| serde_json::to_vec(v).map_or(0, |b| b.len() as u64)) .sum()) } + + /// Same projection the Postgres impl computes in SQL, but with this + /// mock's own (compact) metric: every doc in the group EXCEPT the one + /// being replaced, plus the new data — i.e. `used - old + new`. + async fn projected_total_bytes( + &self, + group_id: GroupId, + replacing: Option, + new_data: &serde_json::Value, + ) -> Result { + let data = self.data.lock().await; + let others: u64 = data + .iter() + .filter(|((g, _), _)| *g == group_id) + .flat_map(|(_, docs)| docs.iter()) + .filter(|(id, _)| Some(**id) != replacing) + .map(|(_, v)| serde_json::to_vec(v).map_or(0, |b| b.len() as u64)) + .sum(); + let new_len = serde_json::to_vec(new_data).map_or(0, |b| b.len() as u64); + Ok(others + new_len) + } } #[derive(Default)] diff --git a/crates/manager-core/src/group_kv_repo.rs b/crates/manager-core/src/group_kv_repo.rs index 139225c..659f210 100644 --- a/crates/manager-core/src/group_kv_repo.rs +++ b/crates/manager-core/src/group_kv_repo.rs @@ -88,6 +88,25 @@ pub trait GroupKvRepo: Send + Sync { let _ = group_id; Ok(0) } + + /// §11.6 M4 quota: the PROJECTED total stored bytes for the group AFTER + /// writing `new_value` at `(collection, key)` — i.e. current SUM minus the + /// existing row's bytes (if any) plus the new value's bytes. Computed + /// entirely in SQL so ALL three terms use the SAME canonical metric + /// (`octet_length(value::text)`); the Rust-side subtract/add previously + /// mixed the DB canonical SUM with a `serde_json` compact length, which + /// drifted the ceiling permissive. Default `Ok(0)` (non-Postgres impls skip + /// the byte quota). + async fn projected_total_bytes( + &self, + group_id: GroupId, + collection: &str, + key: &str, + new_value: &serde_json::Value, + ) -> Result { + let _ = (group_id, collection, key, new_value); + Ok(0) + } } pub struct PostgresGroupKvRepo { @@ -292,6 +311,39 @@ impl GroupKvRepo for PostgresGroupKvRepo { .await?; Ok(u64::try_from(n).unwrap_or(0)) } + + async fn projected_total_bytes( + &self, + group_id: GroupId, + collection: &str, + key: &str, + new_value: &serde_json::Value, + ) -> Result { + // All three terms use `octet_length(value::text)` (jsonb canonical): the + // group SUM, minus the existing row's bytes (0 if new key), plus the new + // value's bytes. The new value is bound as compact TEXT and re-parsed + // through `::jsonb::text` so PG measures it in the SAME canonical form it + // stores — no Rust/Postgres metric mismatch. + // A `serde_json::Value` always serializes; fall back defensively. + let new_text = serde_json::to_string(new_value).unwrap_or_else(|_| "null".to_string()); + let (n,): (i64,) = sqlx::query_as( + "SELECT ( \ + COALESCE((SELECT SUM(octet_length(value::text)) \ + FROM group_kv_entries WHERE group_id = $1), 0) \ + - COALESCE((SELECT octet_length(value::text) \ + FROM group_kv_entries \ + WHERE group_id = $1 AND collection = $2 AND key = $3), 0) \ + + octet_length($4::jsonb::text) \ + )::BIGINT", + ) + .bind(group_id.into_inner()) + .bind(collection) + .bind(key) + .bind(new_text) + .fetch_one(&self.pool) + .await?; + Ok(u64::try_from(n).unwrap_or(0)) + } } fn encode_cursor(last_key: &str) -> String { diff --git a/crates/manager-core/src/group_kv_service.rs b/crates/manager-core/src/group_kv_service.rs index 38f4a02..be7c5c2 100644 --- a/crates/manager-core/src/group_kv_service.rs +++ b/crates/manager-core/src/group_kv_service.rs @@ -152,8 +152,9 @@ impl GroupKvServiceImpl { } /// Encode `value` and enforce the per-value byte cap, returning its encoded - /// length (both callers need it for the byte-quota projection). Mirrors the - /// sibling `group_docs_service::check_data_size`. + /// length. (The per-group byte quota measures canonically in SQL via + /// `projected_total_bytes`, so it no longer consumes this length.) Mirrors + /// the sibling `group_docs_service::check_data_size`. fn check_value_size(&self, value: &serde_json::Value) -> Result { let encoded_len = serde_json::to_vec(value) .map(|v| v.len()) @@ -219,7 +220,7 @@ impl GroupKvService for GroupKvServiceImpl { value: serde_json::Value, ) -> Result<(), GroupKvError> { let group_id = self.owning_group(cx, collection).await?; - let encoded_len = self.check_value_size(&value)?; + self.check_value_size(&value)?; self.check_write(cx, group_id).await?; // Determine insert vs update for the event op + quota deltas (best-effort, // like the per-app path — the get+set is non-transactional but triggers @@ -249,12 +250,13 @@ impl GroupKvService for GroupKvServiceImpl { let rows_after = if existed { count } else { count + 1 }; let upper_bound = rows_after.saturating_mul(self.max_value_bytes as u64); if upper_bound > self.max_total_bytes { - let old_len = existing - .as_ref() - .and_then(|v| serde_json::to_vec(v).ok()) - .map_or(0u64, |b| b.len() as u64); - let used = self.repo.total_bytes(group_id).await?; - let projected = used.saturating_sub(old_len) + encoded_len as u64; + // Consistent-metric projection (canonical `octet_length(value::text)` + // for the current SUM, the old row, and the new value) — no + // Rust-compact vs DB-canonical drift. + let projected = self + .repo + .projected_total_bytes(group_id, collection, key, &value) + .await?; if projected > self.max_total_bytes { return Err(GroupKvError::TotalBytesQuotaExceeded { limit: self.max_total_bytes, @@ -286,7 +288,7 @@ impl GroupKvService for GroupKvServiceImpl { new: serde_json::Value, ) -> Result { let group_id = self.owning_group(cx, collection).await?; - let encoded_len = self.check_value_size(&new)?; + self.check_value_size(&new)?; self.check_write(cx, group_id).await?; // Same quota rules as `set` (best-effort pre-check; the swap itself is // atomic in the repo). A swap that inserts a NEW key (expected absent) @@ -301,12 +303,10 @@ impl GroupKvService for GroupKvServiceImpl { }); } } - let old_len = existing - .as_ref() - .and_then(|v| serde_json::to_vec(v).ok()) - .map_or(0u64, |b| b.len() as u64); - let used = self.repo.total_bytes(group_id).await?; - let projected = used.saturating_sub(old_len) + encoded_len as u64; + let projected = self + .repo + .projected_total_bytes(group_id, collection, key, &new) + .await?; if projected > self.max_total_bytes { return Err(GroupKvError::TotalBytesQuotaExceeded { limit: self.max_total_bytes, @@ -495,6 +495,27 @@ mod tests { .map(|(_, v)| serde_json::to_vec(v).map_or(0, |b| b.len() as u64)) .sum()) } + + /// Same projection the Postgres impl computes in SQL, but with this + /// mock's own (compact) metric: every OTHER row in the group, plus the + /// new value — i.e. `used - old + new`, measured consistently. + async fn projected_total_bytes( + &self, + group_id: GroupId, + collection: &str, + key: &str, + new_value: &serde_json::Value, + ) -> Result { + let data = self.data.lock().await; + let target = (group_id, collection.to_string(), key.to_string()); + let others: u64 = data + .iter() + .filter(|(k, _)| k.0 == group_id && **k != target) + .map(|(_, v)| serde_json::to_vec(v).map_or(0, |b| b.len() as u64)) + .sum(); + let new_len = serde_json::to_vec(new_value).map_or(0, |b| b.len() as u64); + Ok(others + new_len) + } } /// Maps `(app_id, lowercased name) -> owning group`. Anything absent is