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:
MechaCat02
2026-07-14 07:18:14 +02:00
parent fc70f5e224
commit 155c5471b7
4 changed files with 170 additions and 33 deletions

View File

@@ -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 {

View File

@@ -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)]

View File

@@ -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 {

View File

@@ -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