fix(quota): make the per-group ceiling a bound, not a hint

Audit #8. `GroupKvServiceImpl` read the group's usage (`count_rows` /
`projected_total_bytes`) on one pooled connection and wrote on another. N
concurrent writers therefore each observed the same pre-write total, each
concluded it had room, and together sailed past the ceiling. The code
called this out as "best-effort ... quotas are safety rails, not exact
accounting" — but the overshoot is not small. With a ceiling of 5 and 20
concurrent writers, 19 rows land.

Route group-KV mutations through `atomic_write::GroupKvWriter`, whose
Postgres impl runs the quota reads, the row write, and the shared-trigger
fan-out on ONE connection in ONE transaction — and, crucially, takes a
per-group `pg_advisory_xact_lock` first.

The lock is the fix, not the transaction. Putting the check inside the
writing tx is necessary but NOT sufficient: under READ COMMITTED each
transaction's `COUNT(*)`/`SUM(...)` still sees a snapshot without the other
writers' uncommitted rows, so they all still pass. Serializing the
check-then-write per group is what makes a writer see its predecessor's row.
The lock key is namespaced per kind (kv/docs/…) so writes drawing on
different ceilings don't contend. A delete takes no lock — it only frees
space.

The quota POLICY (row ceiling, projected-bytes ceiling, and the
upper-bound fast path that skips the O(n) SUM scan for a group nowhere near
its cap) moves to one place, `group_quota::check_group_write`, over a small
`GroupUsage` trait. Both backends — the transactional one reading through
`&mut *tx`, and the in-memory one the unit tests use — share it, so there is
exactly one copy of the decision.

`set_if` gains a correctness nicety on the way: the row ceiling now keys on
whether the write ADDS a row, so a compare-and-swap against an absent key
with a `Some` precondition (which cannot swap, and so consumes nothing) is
no longer charged for one.

tests/atomic_write.rs pins both ceilings against 20/16 concurrent writers.
Both tests fail without the lock (19 rows stored against a ceiling of 5).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-07-14 19:37:11 +02:00
parent 01e4e5a8f5
commit b086713e3d
6 changed files with 951 additions and 242 deletions

View File

@@ -40,12 +40,16 @@
use std::sync::Arc;
use async_trait::async_trait;
use picloud_shared::{KvError, SdkCallCx, ServiceEvent, ServiceEventEmitter};
use picloud_shared::{
GroupId, GroupKvError, KvError, SdkCallCx, ServiceEvent, ServiceEventEmitter,
};
use serde_json::Value;
use sqlx::PgPool;
use sqlx::{PgConnection, PgPool};
use crate::group_kv_repo::{self, GroupKvRepo};
use crate::group_quota::{check_group_write, GroupUsage, GroupWriteQuota, QuotaOutcome};
use crate::kv_repo::{self, KvRepo};
use crate::outbox_event_emitter::emit_on;
use crate::outbox_event_emitter::{emit_on, emit_shared_on};
/// Map a `sqlx` failure from the transaction plumbing itself (begin/commit).
fn tx_err(e: &sqlx::Error) -> KvError {
@@ -295,3 +299,420 @@ impl KvWriter for BestEffortKvWriter {
Ok(previous)
}
}
// ----------------------------------------------------------------------------
// §11.6 group (shared-collection) KV
//
// Adds a second reason to hold a transaction: the QUOTA. The service used to
// read the group's usage on one pooled connection and write on another, so two
// concurrent writers both observed the pre-write total, both passed the check,
// and together pushed the group past its ceiling (audit #8). Reading usage
// inside the writing transaction is necessary but not sufficient — under READ
// COMMITTED each transaction's aggregate still sees a snapshot without the
// other's uncommitted row. So the transaction ALSO takes a per-group advisory
// lock, which serializes check-then-write for that group and makes the ceiling
// an actual bound rather than a hint.
// ----------------------------------------------------------------------------
/// Where a group-KV mutation lands.
#[derive(Clone, Copy)]
pub struct GroupKvTarget<'a> {
pub group_id: GroupId,
pub collection: &'a str,
pub key: &'a str,
}
/// The mutating half of the group-KV service: the quota decision, the row write,
/// and the shared-trigger fan-out — which must all agree, so they all commit
/// together.
#[async_trait]
pub trait GroupKvWriter: Send + Sync {
async fn set(
&self,
cx: &SdkCallCx,
target: GroupKvTarget<'_>,
value: Value,
quota: GroupWriteQuota,
) -> Result<(), GroupKvError>;
async fn set_if(
&self,
cx: &SdkCallCx,
target: GroupKvTarget<'_>,
expected: Option<Value>,
new: Value,
quota: GroupWriteQuota,
) -> Result<bool, GroupKvError>;
async fn delete(&self, cx: &SdkCallCx, target: GroupKvTarget<'_>)
-> Result<bool, GroupKvError>;
}
fn group_tx_err(e: &sqlx::Error) -> GroupKvError {
GroupKvError::Backend(format!("database error: {e}"))
}
/// Turn the shared quota decision into the service's error type.
fn quota_gate(outcome: QuotaOutcome) -> Result<(), GroupKvError> {
match outcome {
QuotaOutcome::Allowed => Ok(()),
QuotaOutcome::RowsExceeded { limit, actual } => Err(GroupKvError::QuotaExceeded {
limit: usize::try_from(limit).unwrap_or(usize::MAX),
actual: usize::try_from(actual).unwrap_or(usize::MAX),
}),
QuotaOutcome::BytesExceeded { limit, actual } => {
Err(GroupKvError::TotalBytesQuotaExceeded { limit, actual })
}
}
}
/// Stable `pg_advisory_xact_lock` key for a group's shared-collection quota
/// accounting. Held for the life of the writing transaction, so concurrent
/// writers to the same group queue up behind each other for the
/// check-then-write and one of them sees the other's row.
///
/// The key is namespaced per KIND (`kv` / `docs` / …): a KV write and a docs
/// write draw on separate ceilings, so serializing them against each other would
/// be pure contention for no safety.
fn group_quota_lock_key(group_id: GroupId, kind: &str) -> i64 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut h = DefaultHasher::new();
"group_quota".hash(&mut h);
kind.hash(&mut h);
group_id.into_inner().hash(&mut h);
// Postgres advisory locks key on bigint; reinterpret the u64 as i64.
#[allow(clippy::cast_possible_wrap)]
let key = h.finish() as i64;
key
}
async fn lock_group_quota(
conn: &mut PgConnection,
group_id: GroupId,
kind: &str,
) -> Result<(), sqlx::Error> {
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(group_quota_lock_key(group_id, kind))
.execute(conn)
.await?;
Ok(())
}
/// Group-KV usage, read on the WRITING transaction's connection.
struct TxGroupKvUsage<'a> {
conn: &'a mut PgConnection,
group_id: GroupId,
collection: &'a str,
key: &'a str,
new_value: &'a Value,
}
#[async_trait]
impl GroupUsage for TxGroupKvUsage<'_> {
async fn count_rows(&mut self) -> Result<u64, String> {
group_kv_repo::count_rows_on(&mut *self.conn, self.group_id)
.await
.map_err(|e| e.to_string())
}
async fn projected_total_bytes(&mut self) -> Result<u64, String> {
group_kv_repo::projected_total_bytes_on(
&mut *self.conn,
self.group_id,
self.collection,
self.key,
self.new_value,
)
.await
.map_err(|e| e.to_string())
}
}
pub struct PostgresGroupKvWriter {
pool: PgPool,
}
impl PostgresGroupKvWriter {
#[must_use]
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
}
#[async_trait]
impl GroupKvWriter for PostgresGroupKvWriter {
async fn set(
&self,
cx: &SdkCallCx,
t: GroupKvTarget<'_>,
value: Value,
quota: GroupWriteQuota,
) -> Result<(), GroupKvError> {
let mut tx = self.pool.begin().await.map_err(|e| group_tx_err(&e))?;
lock_group_quota(&mut tx, t.group_id, "kv")
.await
.map_err(|e| group_tx_err(&e))?;
let existed = group_kv_repo::get_on(&mut *tx, t.group_id, t.collection, t.key)
.await?
.is_some();
let mut usage = TxGroupKvUsage {
conn: &mut tx,
group_id: t.group_id,
collection: t.collection,
key: t.key,
new_value: &value,
};
let outcome = check_group_write(quota, !existed, &mut usage)
.await
.map_err(GroupKvError::Backend)?;
quota_gate(outcome)?;
group_kv_repo::set_on(&mut *tx, t.group_id, t.collection, t.key, value.clone()).await?;
let op = if existed { "update" } else { "insert" };
emit_shared_on(
&mut tx,
cx,
t.group_id,
&kv_event(op, t.collection, t.key, Some(value), None),
)
.await
.map_err(|e| GroupKvError::Backend(format!("event emit: {e}")))?;
tx.commit().await.map_err(|e| group_tx_err(&e))?;
Ok(())
}
async fn set_if(
&self,
cx: &SdkCallCx,
t: GroupKvTarget<'_>,
expected: Option<Value>,
new: Value,
quota: GroupWriteQuota,
) -> Result<bool, GroupKvError> {
let mut tx = self.pool.begin().await.map_err(|e| group_tx_err(&e))?;
lock_group_quota(&mut tx, t.group_id, "kv")
.await
.map_err(|e| group_tx_err(&e))?;
let existed = group_kv_repo::get_on(&mut *tx, t.group_id, t.collection, t.key)
.await?
.is_some();
// A swap only ADDS a row when the precondition was "absent" and the key
// really is absent. With a `Some` precondition against an absent key the
// swap cannot happen at all, so it consumes nothing.
let adds_row = expected.is_none() && !existed;
let mut usage = TxGroupKvUsage {
conn: &mut tx,
group_id: t.group_id,
collection: t.collection,
key: t.key,
new_value: &new,
};
let outcome = check_group_write(quota, adds_row, &mut usage)
.await
.map_err(GroupKvError::Backend)?;
quota_gate(outcome)?;
let op = if expected.is_some() {
"update"
} else {
"insert"
};
let swapped = group_kv_repo::set_if_on(
&mut *tx,
t.group_id,
t.collection,
t.key,
expected,
new.clone(),
)
.await?;
if swapped {
emit_shared_on(
&mut tx,
cx,
t.group_id,
&kv_event(op, t.collection, t.key, Some(new), None),
)
.await
.map_err(|e| GroupKvError::Backend(format!("event emit: {e}")))?;
}
tx.commit().await.map_err(|e| group_tx_err(&e))?;
Ok(swapped)
}
async fn delete(&self, cx: &SdkCallCx, t: GroupKvTarget<'_>) -> Result<bool, GroupKvError> {
// No quota to check (a delete only frees space) and so no lock to take.
let mut tx = self.pool.begin().await.map_err(|e| group_tx_err(&e))?;
let deleted = group_kv_repo::delete_on(&mut *tx, t.group_id, t.collection, t.key)
.await?
.is_some();
if deleted {
emit_shared_on(
&mut tx,
cx,
t.group_id,
&kv_event("delete", t.collection, t.key, None, None),
)
.await
.map_err(|e| GroupKvError::Backend(format!("event emit: {e}")))?;
}
tx.commit().await.map_err(|e| group_tx_err(&e))?;
Ok(deleted)
}
}
// ---- Best-effort group-KV writer (in-memory tests) --------------------------
/// Group-KV usage read through the repo trait — one pooled call each, i.e. NOT
/// serialized against a concurrent writer. Only the in-memory unit tests run
/// on this; the host uses `PostgresGroupKvWriter`.
struct RepoGroupKvUsage<'a> {
repo: &'a dyn GroupKvRepo,
group_id: GroupId,
collection: &'a str,
key: &'a str,
new_value: &'a Value,
}
#[async_trait]
impl GroupUsage for RepoGroupKvUsage<'_> {
async fn count_rows(&mut self) -> Result<u64, String> {
self.repo
.count_rows(self.group_id)
.await
.map_err(|e| e.to_string())
}
async fn projected_total_bytes(&mut self) -> Result<u64, String> {
self.repo
.projected_total_bytes(self.group_id, self.collection, self.key, self.new_value)
.await
.map_err(|e| e.to_string())
}
}
pub struct BestEffortGroupKvWriter {
repo: Arc<dyn GroupKvRepo>,
events: Arc<dyn ServiceEventEmitter>,
}
impl BestEffortGroupKvWriter {
#[must_use]
pub fn new(repo: Arc<dyn GroupKvRepo>, events: Arc<dyn ServiceEventEmitter>) -> Self {
Self { repo, events }
}
async fn emit(&self, cx: &SdkCallCx, group_id: GroupId, event: ServiceEvent) {
let (source, op) = (event.source, event.op);
if let Err(e) = self.events.emit_shared(cx, group_id, event).await {
tracing::error!(
error = %e, source, op, event_emit_failure = true,
"shared event emit failed — the write committed but its triggers will not fire"
);
}
}
async fn check(
&self,
t: GroupKvTarget<'_>,
adds_row: bool,
new_value: &Value,
quota: GroupWriteQuota,
) -> Result<(), GroupKvError> {
let mut usage = RepoGroupKvUsage {
repo: &*self.repo,
group_id: t.group_id,
collection: t.collection,
key: t.key,
new_value,
};
let outcome = check_group_write(quota, adds_row, &mut usage)
.await
.map_err(GroupKvError::Backend)?;
quota_gate(outcome)
}
}
#[async_trait]
impl GroupKvWriter for BestEffortGroupKvWriter {
async fn set(
&self,
cx: &SdkCallCx,
t: GroupKvTarget<'_>,
value: Value,
quota: GroupWriteQuota,
) -> Result<(), GroupKvError> {
let existed = self
.repo
.get(t.group_id, t.collection, t.key)
.await?
.is_some();
self.check(t, !existed, &value, quota).await?;
self.repo
.set(t.group_id, t.collection, t.key, value.clone())
.await?;
let op = if existed { "update" } else { "insert" };
self.emit(
cx,
t.group_id,
kv_event(op, t.collection, t.key, Some(value), None),
)
.await;
Ok(())
}
async fn set_if(
&self,
cx: &SdkCallCx,
t: GroupKvTarget<'_>,
expected: Option<Value>,
new: Value,
quota: GroupWriteQuota,
) -> Result<bool, GroupKvError> {
let existed = self
.repo
.get(t.group_id, t.collection, t.key)
.await?
.is_some();
let adds_row = expected.is_none() && !existed;
self.check(t, adds_row, &new, quota).await?;
let op = if expected.is_some() {
"update"
} else {
"insert"
};
let swapped = self
.repo
.set_if(t.group_id, t.collection, t.key, expected, new.clone())
.await?;
if swapped {
self.emit(
cx,
t.group_id,
kv_event(op, t.collection, t.key, Some(new), None),
)
.await;
}
Ok(swapped)
}
async fn delete(&self, cx: &SdkCallCx, t: GroupKvTarget<'_>) -> Result<bool, GroupKvError> {
let deleted = self
.repo
.delete(t.group_id, t.collection, t.key)
.await?
.is_some();
if deleted {
self.emit(
cx,
t.group_id,
kv_event("delete", t.collection, t.key, None, None),
)
.await;
}
Ok(deleted)
}
}

View File

@@ -132,16 +132,7 @@ impl GroupKvRepo for PostgresGroupKvRepo {
collection: &str,
key: &str,
) -> Result<Option<serde_json::Value>, GroupKvRepoError> {
let row: Option<(serde_json::Value,)> = sqlx::query_as(
"SELECT value FROM group_kv_entries \
WHERE group_id = $1 AND collection = $2 AND key = $3",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|(v,)| v))
get_on(&self.pool, group_id, collection, key).await
}
async fn set(
@@ -151,27 +142,7 @@ impl GroupKvRepo for PostgresGroupKvRepo {
key: &str,
value: serde_json::Value,
) -> Result<Option<serde_json::Value>, GroupKvRepoError> {
let row: Option<(Option<serde_json::Value>,)> = sqlx::query_as(
"WITH prev AS (\
SELECT value FROM group_kv_entries \
WHERE group_id = $1 AND collection = $2 AND key = $3\
), \
upserted AS (\
INSERT INTO group_kv_entries (group_id, collection, key, value) \
VALUES ($1, $2, $3, $4) \
ON CONFLICT (group_id, collection, key) DO UPDATE \
SET value = EXCLUDED.value, updated_at = NOW() \
RETURNING 1\
) \
SELECT (SELECT value FROM prev) FROM upserted",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.bind(value)
.fetch_optional(&self.pool)
.await?;
Ok(row.and_then(|(v,)| v))
set_on(&self.pool, group_id, collection, key, value).await
}
async fn set_if(
@@ -182,33 +153,7 @@ impl GroupKvRepo for PostgresGroupKvRepo {
expected: Option<serde_json::Value>,
new: serde_json::Value,
) -> Result<bool, GroupKvRepoError> {
let affected = match expected {
Some(exp) => sqlx::query(
"UPDATE group_kv_entries SET value = $4, updated_at = NOW() \
WHERE group_id = $1 AND collection = $2 AND key = $3 AND value = $5",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.bind(new)
.bind(exp)
.execute(&self.pool)
.await?
.rows_affected(),
None => sqlx::query(
"INSERT INTO group_kv_entries (group_id, collection, key, value) \
VALUES ($1, $2, $3, $4) \
ON CONFLICT (group_id, collection, key) DO NOTHING",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.bind(new)
.execute(&self.pool)
.await?
.rows_affected(),
};
Ok(affected == 1)
set_if_on(&self.pool, group_id, collection, key, expected, new).await
}
async fn delete(
@@ -217,17 +162,7 @@ impl GroupKvRepo for PostgresGroupKvRepo {
collection: &str,
key: &str,
) -> Result<Option<serde_json::Value>, GroupKvRepoError> {
let row: Option<(serde_json::Value,)> = sqlx::query_as(
"DELETE FROM group_kv_entries \
WHERE group_id = $1 AND collection = $2 AND key = $3 \
RETURNING value",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|(v,)| v))
delete_on(&self.pool, group_id, collection, key).await
}
async fn has(
@@ -293,12 +228,7 @@ impl GroupKvRepo for PostgresGroupKvRepo {
}
async fn count_rows(&self, group_id: GroupId) -> Result<u64, GroupKvRepoError> {
let (n,): (i64,) =
sqlx::query_as("SELECT COUNT(*) FROM group_kv_entries WHERE group_id = $1")
.bind(group_id.into_inner())
.fetch_one(&self.pool)
.await?;
Ok(u64::try_from(n).unwrap_or(0))
count_rows_on(&self.pool, group_id).await
}
async fn total_bytes(&self, group_id: GroupId) -> Result<u64, GroupKvRepoError> {
@@ -319,31 +249,193 @@ impl GroupKvRepo for PostgresGroupKvRepo {
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",
projected_total_bytes_on(&self.pool, group_id, collection, key, new_value).await
}
}
// ----------------------------------------------------------------------------
// Connection-scoped operations
//
// Generic over the executor so the same SQL serves the pooled trait methods
// above and `crate::atomic_write`, which runs the per-group advisory lock, the
// quota reads, the write, and the shared-trigger fan-out on ONE connection
// inside ONE transaction. The quota reads in particular MUST run on the writing
// transaction's connection — that is what makes check-then-write atomic rather
// than a race (audit #8).
// ----------------------------------------------------------------------------
pub(crate) async fn get_on<'c, E>(
exec: E,
group_id: GroupId,
collection: &str,
key: &str,
) -> Result<Option<serde_json::Value>, GroupKvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
let row: Option<(serde_json::Value,)> = sqlx::query_as(
"SELECT value FROM group_kv_entries \
WHERE group_id = $1 AND collection = $2 AND key = $3",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.fetch_optional(exec)
.await?;
Ok(row.map(|(v,)| v))
}
/// Upsert. Returns the previous value (`None` = this was an insert).
pub(crate) async fn set_on<'c, E>(
exec: E,
group_id: GroupId,
collection: &str,
key: &str,
value: serde_json::Value,
) -> Result<Option<serde_json::Value>, GroupKvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
let row: Option<(Option<serde_json::Value>,)> = sqlx::query_as(
"WITH prev AS (\
SELECT value FROM group_kv_entries \
WHERE group_id = $1 AND collection = $2 AND key = $3\
), \
upserted AS (\
INSERT INTO group_kv_entries (group_id, collection, key, value) \
VALUES ($1, $2, $3, $4) \
ON CONFLICT (group_id, collection, key) DO UPDATE \
SET value = EXCLUDED.value, updated_at = NOW() \
RETURNING 1\
) \
SELECT (SELECT value FROM prev) FROM upserted",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.bind(value)
.fetch_optional(exec)
.await?;
Ok(row.and_then(|(v,)| v))
}
/// Compare-and-swap. Writes `new` iff the current value equals `expected`
/// (`expected = None` → iff the key is ABSENT). Returns whether it wrote.
pub(crate) async fn set_if_on<'c, E>(
exec: E,
group_id: GroupId,
collection: &str,
key: &str,
expected: Option<serde_json::Value>,
new: serde_json::Value,
) -> Result<bool, GroupKvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
let affected = match expected {
Some(exp) => sqlx::query(
"UPDATE group_kv_entries SET value = $4, updated_at = NOW() \
WHERE group_id = $1 AND collection = $2 AND key = $3 AND value = $5",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.bind(new_text)
.fetch_one(&self.pool)
.bind(new)
.bind(exp)
.execute(exec)
.await?
.rows_affected(),
None => sqlx::query(
"INSERT INTO group_kv_entries (group_id, collection, key, value) \
VALUES ($1, $2, $3, $4) \
ON CONFLICT (group_id, collection, key) DO NOTHING",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.bind(new)
.execute(exec)
.await?
.rows_affected(),
};
Ok(affected == 1)
}
/// Returns the deleted value if present, `None` if the row didn't exist.
pub(crate) async fn delete_on<'c, E>(
exec: E,
group_id: GroupId,
collection: &str,
key: &str,
) -> Result<Option<serde_json::Value>, GroupKvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
let row: Option<(serde_json::Value,)> = sqlx::query_as(
"DELETE FROM group_kv_entries \
WHERE group_id = $1 AND collection = $2 AND key = $3 \
RETURNING value",
)
.bind(group_id.into_inner())
.bind(collection)
.bind(key)
.fetch_optional(exec)
.await?;
Ok(row.map(|(v,)| v))
}
/// Total key count across the group's shared-KV collections (§11.6 row quota).
pub(crate) async fn count_rows_on<'c, E>(
exec: E,
group_id: GroupId,
) -> Result<u64, GroupKvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_kv_entries WHERE group_id = $1")
.bind(group_id.into_inner())
.fetch_one(exec)
.await?;
Ok(u64::try_from(n).unwrap_or(0))
}
Ok(u64::try_from(n).unwrap_or(0))
}
/// §11.6 M4 quota: the PROJECTED total stored bytes for the group AFTER writing
/// `new_value` at `(collection, key)` — current SUM, minus the existing row's
/// bytes (0 if the key is new), plus the new value's bytes. Computed entirely in
/// SQL so all three terms use the SAME canonical metric
/// (`octet_length(value::text)`); mixing the DB's canonical SUM with a
/// `serde_json` compact length drifts the ceiling permissive.
pub(crate) async fn projected_total_bytes_on<'c, E>(
exec: E,
group_id: GroupId,
collection: &str,
key: &str,
new_value: &serde_json::Value,
) -> Result<u64, GroupKvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
// 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.
// 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(exec)
.await?;
Ok(u64::try_from(n).unwrap_or(0))
}
fn encode_cursor(last_key: &str) -> String {

View File

@@ -22,13 +22,17 @@ use std::sync::Arc;
use async_trait::async_trait;
use picloud_shared::{
GroupId, GroupKvError, GroupKvService, KvListPage, NoopEventEmitter, SdkCallCx, ServiceEvent,
GroupId, GroupKvError, GroupKvService, KvListPage, NoopEventEmitter, SdkCallCx,
ServiceEventEmitter,
};
use crate::atomic_write::{
BestEffortGroupKvWriter, GroupKvTarget, GroupKvWriter, PostgresGroupKvWriter,
};
use crate::authz::{self, AuthzRepo, Capability};
use crate::group_collection_repo::GroupCollectionResolver;
use crate::group_kv_repo::{GroupKvRepo, GroupKvRepoError};
use crate::group_quota::GroupWriteQuota;
use crate::kv_service::kv_max_value_bytes_from_env;
/// The registry `kind` this service resolves.
@@ -48,10 +52,10 @@ pub struct GroupKvServiceImpl {
/// projected total (old value subtracted, new added) so a same/smaller update
/// near the cap is still allowed.
max_total_bytes: u64,
/// §11.6: fires `shared = true` triggers on a shared-collection write.
/// Defaults to the noop emitter; the host wires the outbox emitter via
/// [`Self::with_events`].
events: Arc<dyn ServiceEventEmitter>,
/// Mutations: the quota decision, the row write, and the shared-trigger
/// fan-out, which for the host all commit in one advisory-locked
/// transaction. Reads stay on `repo`.
writer: Arc<dyn GroupKvWriter>,
}
impl GroupKvServiceImpl {
@@ -64,11 +68,26 @@ impl GroupKvServiceImpl {
Self::with_max_value_bytes(repo, resolver, authz, kv_max_value_bytes_from_env())
}
/// Wire the event emitter (§11.6 shared-collection triggers). Without this
/// the service uses the noop emitter and shared writes fire nothing.
/// Wire the event emitter (§11.6 shared-collection triggers) on the
/// best-effort writer. Without this — or [`Self::with_atomic_writes`], which
/// supersedes it — shared writes fire nothing.
#[must_use]
pub fn with_events(mut self, events: Arc<dyn ServiceEventEmitter>) -> Self {
self.events = events;
self.writer = Arc::new(BestEffortGroupKvWriter::new(self.repo.clone(), events));
self
}
/// Use the transactional writer: the quota reads, the row write, and the
/// shared-trigger fan-out all run on ONE connection under a per-group
/// advisory lock. That makes the quota an actual bound (concurrent writers
/// to a group serialize, so one sees the other's row) and makes an emit
/// failure roll the write back instead of silently losing the event.
///
/// Supersedes [`Self::with_events`] — the transactional writer resolves and
/// writes the fan-out itself and needs no injected emitter.
#[must_use]
pub fn with_atomic_writes(mut self, pool: sqlx::PgPool) -> Self {
self.writer = Arc::new(PostgresGroupKvWriter::new(pool));
self
}
@@ -88,7 +107,10 @@ impl GroupKvServiceImpl {
max_value_bytes: usize,
) -> Self {
Self {
events: Arc::new(NoopEventEmitter),
writer: Arc::new(BestEffortGroupKvWriter::new(
repo.clone(),
Arc::new(NoopEventEmitter),
)),
max_rows: crate::group_quota::group_kv_max_rows_from_env(),
max_total_bytes: crate::group_quota::group_kv_max_total_bytes_from_env(),
repo,
@@ -123,39 +145,19 @@ impl GroupKvServiceImpl {
.ok_or_else(|| GroupKvError::CollectionNotShared(collection.to_string()))
}
/// §11.6: fire `shared = true` kv triggers on the owning group. Best-effort
/// — an emit failure is logged, never surfaced (the write already
/// committed), matching the per-app `KvServiceImpl` behaviour.
async fn emit_shared(
&self,
cx: &SdkCallCx,
group_id: GroupId,
op: &'static str,
collection: &str,
key: &str,
payload: Option<serde_json::Value>,
) {
crate::group_collection_repo::best_effort_emit_shared(
&*self.events,
cx,
group_id,
ServiceEvent {
source: "kv",
op,
collection: Some(collection.to_string()),
key: Some(key.to_string()),
payload,
old_payload: None,
},
)
.await;
/// The ceilings a shared-collection write must fit under.
fn quota(&self) -> GroupWriteQuota {
GroupWriteQuota {
max_rows: self.max_rows,
max_total_bytes: self.max_total_bytes,
max_value_bytes: self.max_value_bytes,
}
}
/// Encode `value` and enforce the per-value byte cap, returning its encoded
/// 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> {
/// Enforce the per-value byte cap. The per-GROUP byte quota is measured
/// canonically in SQL inside the writing transaction — see
/// `group_quota::check_group_write`.
fn check_value_size(&self, value: &serde_json::Value) -> Result<(), GroupKvError> {
let encoded_len = serde_json::to_vec(value)
.map(|v| v.len())
.map_err(|e| GroupKvError::Backend(format!("encode value: {e}")))?;
@@ -165,7 +167,7 @@ impl GroupKvServiceImpl {
actual: encoded_len,
});
}
Ok(encoded_len)
Ok(())
}
async fn check_read(&self, cx: &SdkCallCx, group_id: GroupId) -> Result<(), GroupKvError> {
@@ -222,61 +224,18 @@ impl GroupKvService for GroupKvServiceImpl {
let group_id = self.owning_group(cx, collection).await?;
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
// are fire-and-forget and quotas are safety rails, not exact accounting).
let existing = self.repo.get(group_id, collection, key).await?;
let existed = existing.is_some();
// One cheap COUNT backs BOTH the row quota AND the byte-quota fast path
// below (an index-only count, far lighter than the per-value text SUM).
let count = self.repo.count_rows(group_id).await?;
// §11.6 row quota: a NEW key must fit under the group's row ceiling; an
// update of an existing key is net-zero and exempt.
if !existed && count >= self.max_rows {
return Err(GroupKvError::QuotaExceeded {
limit: usize::try_from(self.max_rows).unwrap_or(usize::MAX),
actual: usize::try_from(count).unwrap_or(usize::MAX),
});
}
// §11.6 M4 byte quota: the PROJECTED total (old value's bytes subtracted,
// new value's added) must fit under the group's total-bytes ceiling, so a
// same-or-smaller update near the cap is still allowed.
//
// Fast path: every stored value is capped at `max_value_bytes`, so the
// whole collection can hold at most `rows_after * max_value_bytes`. When
// that exact upper bound already fits under the ceiling there is no way to
// exceed it — skip the O(collection-size) `SUM(octet_length(...))` scan
// entirely. The scan runs only for a group actually near its byte cap.
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 {
// 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,
actual: projected,
});
}
}
self.repo
.set(group_id, collection, key, value.clone())
.await?;
self.emit_shared(
cx,
group_id,
if existed { "update" } else { "insert" },
collection,
key,
Some(value),
)
.await;
Ok(())
self.writer
.set(
cx,
GroupKvTarget {
group_id,
collection,
key,
},
value,
self.quota(),
)
.await
}
async fn set_if(
@@ -290,43 +249,19 @@ impl GroupKvService for GroupKvServiceImpl {
let group_id = self.owning_group(cx, collection).await?;
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)
// must fit the row ceiling; the byte ceiling uses the projected total.
let existing = self.repo.get(group_id, collection, key).await?;
if expected.is_none() && existing.is_none() {
let count = self.repo.count_rows(group_id).await?;
if count >= self.max_rows {
return Err(GroupKvError::QuotaExceeded {
limit: usize::try_from(self.max_rows).unwrap_or(usize::MAX),
actual: usize::try_from(count).unwrap_or(usize::MAX),
});
}
}
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,
actual: projected,
});
}
let op = if expected.is_some() {
"update"
} else {
"insert"
};
let swapped = self
.repo
.set_if(group_id, collection, key, expected, new.clone())
.await?;
if swapped {
self.emit_shared(cx, group_id, op, collection, key, Some(new))
.await;
}
Ok(swapped)
self.writer
.set_if(
cx,
GroupKvTarget {
group_id,
collection,
key,
},
expected,
new,
self.quota(),
)
.await
}
async fn delete(
@@ -337,12 +272,16 @@ impl GroupKvService for GroupKvServiceImpl {
) -> Result<bool, GroupKvError> {
let group_id = self.owning_group(cx, collection).await?;
self.check_write(cx, group_id).await?;
let deleted = self.repo.delete(group_id, collection, key).await?.is_some();
if deleted {
self.emit_shared(cx, group_id, "delete", collection, key, None)
.await;
}
Ok(deleted)
self.writer
.delete(
cx,
GroupKvTarget {
group_id,
collection,
key,
},
)
.await
}
async fn has(&self, cx: &SdkCallCx, collection: &str, key: &str) -> Result<bool, GroupKvError> {

View File

@@ -67,3 +67,90 @@ pub fn group_files_max_total_bytes_from_env() -> u64 {
DEFAULT_GROUP_FILES_MAX_TOTAL_BYTES,
)
}
// ----------------------------------------------------------------------------
// The row/byte decision, shared by every backend
// ----------------------------------------------------------------------------
/// The ceilings one shared-collection write must fit under.
#[derive(Debug, Clone, Copy)]
pub struct GroupWriteQuota {
pub max_rows: u64,
pub max_total_bytes: u64,
/// The per-value cap. Not a ceiling of its own here — it is what makes the
/// fast path below sound.
pub max_value_bytes: usize,
}
/// The two usage reads a quota decision needs. Implemented over the writing
/// transaction's connection (the real path — see `crate::atomic_write`) and over
/// the in-memory test store, so the POLICY below has one home regardless of
/// backend.
///
/// **The reads must come from the same transaction as the write they gate.**
/// Reading usage on one connection and writing on another is the check-then-write
/// race that lets N concurrent writers each observe the pre-write total, each
/// pass, and together blow the ceiling.
#[async_trait::async_trait]
pub trait GroupUsage: Send {
/// Current row count for the group, across its collections of this kind.
async fn count_rows(&mut self) -> Result<u64, String>;
/// Total stored bytes for the group AFTER this write — the current total,
/// minus the bytes of the row being replaced (if any), plus the new value's.
async fn projected_total_bytes(&mut self) -> Result<u64, String>;
}
/// What the ceilings say about a pending write.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuotaOutcome {
Allowed,
RowsExceeded { limit: u64, actual: u64 },
BytesExceeded { limit: u64, actual: u64 },
}
/// Decide whether a shared-collection write fits. `adds_row` = the write will
/// ADD a row rather than replace one (a new key, or a doc create) — a
/// compare-and-swap against an absent key with a `Some` precondition adds
/// nothing, because it cannot swap at all.
///
/// # Errors
/// Propagates a backend failure from either usage read, as a string.
pub async fn check_group_write(
q: GroupWriteQuota,
adds_row: bool,
usage: &mut dyn GroupUsage,
) -> Result<QuotaOutcome, String> {
// One cheap COUNT backs BOTH the row ceiling and the byte fast path below.
let count = usage.count_rows().await?;
// Row ceiling: only a write that adds a row consumes one. An update is
// net-zero and exempt.
if adds_row && count >= q.max_rows {
return Ok(QuotaOutcome::RowsExceeded {
limit: q.max_rows,
actual: count,
});
}
// Byte ceiling, on the PROJECTED total — so an update that is the same size
// or smaller still goes through when the group is sitting near its cap.
//
// Fast path: every stored value is capped at `max_value_bytes`, so the group
// can hold at most `rows_after * max_value_bytes`. When even that upper bound
// fits under the ceiling, no write can breach it — skip the
// O(collection-size) `SUM(octet_length(...))` scan entirely. The scan runs
// only for a group actually near its byte cap.
let rows_after = if adds_row { count + 1 } else { count };
let upper_bound = rows_after.saturating_mul(q.max_value_bytes as u64);
if upper_bound > q.max_total_bytes {
let projected = usage.projected_total_bytes().await?;
if projected > q.max_total_bytes {
return Ok(QuotaOutcome::BytesExceeded {
limit: q.max_total_bytes,
actual: projected,
});
}
}
Ok(QuotaOutcome::Allowed)
}

View File

@@ -13,8 +13,13 @@
//! running in parallel against the same database. Skips when `DATABASE_URL` is
//! unset.
use picloud_manager_core::atomic_write::{KvWriter, PostgresKvWriter};
use picloud_shared::{AppId, ExecutionId, RequestId, ScriptId, SdkCallCx};
use std::sync::Arc;
use picloud_manager_core::atomic_write::{
GroupKvTarget, GroupKvWriter, KvWriter, PostgresGroupKvWriter, PostgresKvWriter,
};
use picloud_manager_core::group_quota::GroupWriteQuota;
use picloud_shared::{AppId, ExecutionId, GroupId, GroupKvError, RequestId, ScriptId, SdkCallCx};
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use uuid::Uuid;
@@ -25,7 +30,7 @@ async fn pool_or_skip() -> Option<PgPool> {
return None;
};
let pool = PgPoolOptions::new()
.max_connections(3)
.max_connections(24)
.connect(&url)
.await
.expect("connect");
@@ -269,3 +274,165 @@ async fn a_failed_fan_out_rolls_a_delete_back() {
cleanup(&pool, f.app).await;
}
// ----------------------------------------------------------------------------
// Audit fix #8: the per-group quota must be a BOUND, not a hint.
// ----------------------------------------------------------------------------
async fn mk_group(pool: &PgPool) -> Uuid {
let (g,): (Uuid,) =
sqlx::query_as("INSERT INTO groups (slug, name) VALUES ($1, $1) RETURNING id")
.bind(format!("q-grp-{}", Uuid::new_v4().simple()))
.fetch_one(pool)
.await
.expect("group");
g
}
async fn group_kv_count(pool: &PgPool, group: Uuid) -> i64 {
let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_kv_entries WHERE group_id = $1")
.bind(group)
.fetch_one(pool)
.await
.expect("count");
n
}
/// The service read the group's usage on one pooled connection and wrote on
/// another, so N concurrent writers each saw the same pre-write count, each
/// decided it had room, and together blew past the ceiling. The writer now holds
/// a per-group advisory lock across the check AND the write, so they serialize
/// and each sees its predecessor's row.
///
/// Note this is NOT fixed by merely putting the check in the same transaction:
/// under READ COMMITTED each transaction's `COUNT(*)` still sees a snapshot
/// without the others' uncommitted rows. The lock is what does the work.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_writers_cannot_push_a_group_past_its_row_quota() {
const MAX_ROWS: u64 = 5;
let Some(pool) = pool_or_skip().await else {
return;
};
let group = mk_group(&pool).await;
let writer = Arc::new(PostgresGroupKvWriter::new(pool.clone()));
let quota = GroupWriteQuota {
max_rows: MAX_ROWS,
max_total_bytes: u64::MAX,
max_value_bytes: 256 * 1024,
};
// 20 writers race to insert 20 DISTINCT new keys into an empty group whose
// ceiling is 5. Exactly 5 may win.
let app = Uuid::new_v4(); // cx.app_id is not used by the group writer's SQL
let mut set = tokio::task::JoinSet::new();
for i in 0..20 {
let w = Arc::clone(&writer);
set.spawn(async move {
let cx = cx(app);
w.set(
&cx,
GroupKvTarget {
group_id: GroupId::from(group),
collection: "shared",
key: &format!("k{i}"),
},
serde_json::json!(i),
quota,
)
.await
});
}
let mut ok = 0;
let mut denied = 0;
while let Some(r) = set.join_next().await {
match r.expect("task") {
Ok(()) => ok += 1,
Err(GroupKvError::QuotaExceeded { .. }) => denied += 1,
Err(e) => panic!("unexpected error: {e}"),
}
}
assert_eq!(
group_kv_count(&pool, group).await,
i64::try_from(MAX_ROWS).unwrap(),
"the group must hold EXACTLY its ceiling — a check-then-write race would overshoot"
);
let max_rows = usize::try_from(MAX_ROWS).unwrap();
assert_eq!(ok, max_rows, "exactly {MAX_ROWS} writers may win");
assert_eq!(denied, 20 - max_rows, "the rest are refused");
sqlx::query("DELETE FROM groups WHERE id = $1")
.bind(group)
.execute(&pool)
.await
.expect("cleanup");
}
/// The byte ceiling has the same race, and the same fix.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_writers_cannot_push_a_group_past_its_byte_quota() {
let Some(pool) = pool_or_skip().await else {
return;
};
let group = mk_group(&pool).await;
let writer = Arc::new(PostgresGroupKvWriter::new(pool.clone()));
// A ~100-byte value, and a ceiling that fits about 4 of them. `max_value_bytes`
// is set low enough that the upper-bound fast path can't short-circuit the
// scan (rows_after * max_value_bytes must exceed the ceiling immediately).
let value = serde_json::json!("x".repeat(100));
let quota = GroupWriteQuota {
max_total_bytes: 450,
max_rows: u64::MAX,
max_value_bytes: 200,
};
let app = Uuid::new_v4();
let mut set = tokio::task::JoinSet::new();
for i in 0..16 {
let w = Arc::clone(&writer);
let v = value.clone();
set.spawn(async move {
let cx = cx(app);
w.set(
&cx,
GroupKvTarget {
group_id: GroupId::from(group),
collection: "shared",
key: &format!("k{i}"),
},
v,
quota,
)
.await
});
}
while let Some(r) = set.join_next().await {
match r.expect("task") {
Ok(()) | Err(GroupKvError::TotalBytesQuotaExceeded { .. }) => {}
Err(e) => panic!("unexpected error: {e}"),
}
}
let (bytes,): (i64,) = sqlx::query_as(
"SELECT COALESCE(SUM(octet_length(value::text)), 0)::BIGINT \
FROM group_kv_entries WHERE group_id = $1",
)
.bind(group)
.fetch_one(&pool)
.await
.expect("bytes");
assert!(
bytes <= 450,
"stored bytes ({bytes}) must not exceed the 450-byte ceiling — \
concurrent writers must not each see the same pre-write total"
);
assert!(bytes > 0, "some writers should have succeeded");
sqlx::query("DELETE FROM groups WHERE id = $1")
.bind(group)
.execute(&pool)
.await
.expect("cleanup");
}