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:
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user