fix(quota): measure the per-group byte projection with one consistent metric
The §11.6 M4 byte quota mixed two incompatible byte measures: `used` came from Postgres as `octet_length(value::text)` (jsonb CANONICAL text — whitespace after `:` and `,`, keys normalized), while `old_len`/`new_len` were computed in Rust with `serde_json::to_vec` (COMPACT, no spaces). The projection `used - old_len + new_len` therefore subtracted and added a *smaller* measure than what is actually stored, so a group could be admitted over its ceiling — the cap drifted permissive. The two forms are not reconcilable in Rust (jsonb also reorders/dedups keys), so the projection has to be measured by Postgres. Add `projected_total_bytes` to `GroupKvRepo` (keyed by collection+key) and `GroupDocsRepo` (keyed by the replaced doc id, `None` on create). Each computes `SUM(canonical) - existing_row(canonical) + new_value(canonical)` in ONE query, binding the new value as compact TEXT and re-parsing it through `::jsonb::text` so PG measures it in exactly the form it stores. All three terms now share one metric. The services call it instead of the mixed Rust math; the near-cap fast-path (`rows_after * max_value_bytes <= ceiling` → skip the SUM) is kept, so the common case still does no extra work. The docs `update` path no longer needs its old-doc fetch (the subtraction happens in SQL), removing a read. The in-memory test repos implement the same projection with their own (compact) metric, so the quota unit tests keep their exact byte arithmetic. No schema change. Note: the check-then-write TOCTOU (concurrent writers can each pass and collectively overshoot) is NOT addressed here — it needs a transaction spanning the check and the write, the same missing infrastructure as the transactional outbox. Tracked separately. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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<DocId>,
|
||||
new_data: &serde_json::Value,
|
||||
) -> Result<u64, GroupDocsRepoError> {
|
||||
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<DocId>,
|
||||
new_data: &serde_json::Value,
|
||||
) -> Result<u64, GroupDocsRepoError> {
|
||||
// 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 {
|
||||
|
||||
@@ -95,16 +95,21 @@ impl GroupDocsServiceImpl {
|
||||
async fn check_total_bytes(
|
||||
&self,
|
||||
group_id: GroupId,
|
||||
new_len: u64,
|
||||
old_len: u64,
|
||||
replacing: Option<DocId>,
|
||||
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<DocId>,
|
||||
new_data: &serde_json::Value,
|
||||
) -> Result<u64, GroupDocsRepoError> {
|
||||
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)]
|
||||
|
||||
@@ -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<u64, GroupKvRepoError> {
|
||||
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<u64, GroupKvRepoError> {
|
||||
// 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 {
|
||||
|
||||
@@ -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<usize, GroupKvError> {
|
||||
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<bool, GroupKvError> {
|
||||
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<u64, GroupKvRepoError> {
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user