diff --git a/crates/manager-core/src/atomic_write.rs b/crates/manager-core/src/atomic_write.rs index 17612cd..8d7cf42 100644 --- a/crates/manager-core/src/atomic_write.rs +++ b/crates/manager-core/src/atomic_write.rs @@ -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, + new: Value, + quota: GroupWriteQuota, + ) -> Result; + + async fn delete(&self, cx: &SdkCallCx, target: GroupKvTarget<'_>) + -> Result; +} + +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 { + 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 { + 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, + new: Value, + quota: GroupWriteQuota, + ) -> Result { + 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 { + // 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 { + self.repo + .count_rows(self.group_id) + .await + .map_err(|e| e.to_string()) + } + + async fn projected_total_bytes(&mut self) -> Result { + 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, + events: Arc, +} + +impl BestEffortGroupKvWriter { + #[must_use] + pub fn new(repo: Arc, events: Arc) -> 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, + new: Value, + quota: GroupWriteQuota, + ) -> Result { + 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 { + 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) + } +} diff --git a/crates/manager-core/src/group_kv_repo.rs b/crates/manager-core/src/group_kv_repo.rs index 659f210..4e1a711 100644 --- a/crates/manager-core/src/group_kv_repo.rs +++ b/crates/manager-core/src/group_kv_repo.rs @@ -132,16 +132,7 @@ impl GroupKvRepo for PostgresGroupKvRepo { collection: &str, key: &str, ) -> Result, 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, GroupKvRepoError> { - let row: Option<(Option,)> = 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, new: serde_json::Value, ) -> Result { - 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, 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 { - 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 { @@ -319,31 +249,193 @@ impl GroupKvRepo for PostgresGroupKvRepo { key: &str, new_value: &serde_json::Value, ) -> Result { - // All three terms use `octet_length(value::text)` (jsonb canonical): the - // group SUM, minus the existing row's bytes (0 if new key), plus the new - // value's bytes. The new value is bound as compact TEXT and re-parsed - // through `::jsonb::text` so PG measures it in the SAME canonical form it - // stores — no Rust/Postgres metric mismatch. - // A `serde_json::Value` always serializes; fall back defensively. - let new_text = serde_json::to_string(new_value).unwrap_or_else(|_| "null".to_string()); - let (n,): (i64,) = sqlx::query_as( - "SELECT ( \ - COALESCE((SELECT SUM(octet_length(value::text)) \ - FROM group_kv_entries WHERE group_id = $1), 0) \ - - COALESCE((SELECT octet_length(value::text) \ - FROM group_kv_entries \ - WHERE group_id = $1 AND collection = $2 AND key = $3), 0) \ - + octet_length($4::jsonb::text) \ - )::BIGINT", + 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, 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, GroupKvRepoError> +where + E: sqlx::PgExecutor<'c>, +{ + let row: Option<(Option,)> = 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, + new: serde_json::Value, +) -> Result +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, 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 +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 +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 { diff --git a/crates/manager-core/src/group_kv_service.rs b/crates/manager-core/src/group_kv_service.rs index be7c5c2..46290ae 100644 --- a/crates/manager-core/src/group_kv_service.rs +++ b/crates/manager-core/src/group_kv_service.rs @@ -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, + /// 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, } 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) -> 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, - ) { - 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 { + /// 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 { 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 { diff --git a/crates/manager-core/src/group_quota.rs b/crates/manager-core/src/group_quota.rs index b4763e9..1280e68 100644 --- a/crates/manager-core/src/group_quota.rs +++ b/crates/manager-core/src/group_quota.rs @@ -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; + + /// 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; +} + +/// 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 { + // 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) +} diff --git a/crates/manager-core/tests/atomic_write.rs b/crates/manager-core/tests/atomic_write.rs index 7b34655..384e1ab 100644 --- a/crates/manager-core/tests/atomic_write.rs +++ b/crates/manager-core/tests/atomic_write.rs @@ -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 { 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"); +} diff --git a/crates/picloud/src/lib.rs b/crates/picloud/src/lib.rs index d85ef8d..2f9ed84 100644 --- a/crates/picloud/src/lib.rs +++ b/crates/picloud/src/lib.rs @@ -195,8 +195,11 @@ pub async fn build_app( authz.clone(), picloud_manager_core::kv_service::kv_max_value_bytes_from_env(), ) - // §11.6 shared-collection triggers: fire `shared = true` triggers. - .with_events(events.clone()), + // The quota reads, the write, and the `shared = true` trigger fan-out all + // run on one connection under a per-group advisory lock: the quota is a + // real bound (concurrent writers to a group serialize), and an emit + // failure rolls the write back. Supersedes `.with_events(...)`. + .with_atomic_writes(pool.clone()), ); // §11.6 shared group collections (docs): group-keyed group_docs store. let group_docs: Arc = Arc::new(