From 58bf0ab3eca0fb1981426c980fc944fdcd42abbe Mon Sep 17 00:00:00 2001 From: MechaCat02 Date: Tue, 14 Jul 2026 19:45:26 +0200 Subject: [PATCH] fix(docs): make docs writes transactional and the group quota a bound MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Audit #6 and #8, applied to the docs store — the same two bugs KV had, in the same two places. Per-app docs (#6): create/update/delete wrote the row, then emitted best-effort. An outbox failure left a committed doc whose trigger never fired. They now go through `atomic_write::DocsWriter`, whose Postgres impl writes and fans out on one connection in one transaction. Group docs (#8): the service read `count_rows`/`projected_total_bytes` on pooled connections and wrote on another, so concurrent creators each saw the same pre-write count and together overshot the ceiling. `PostgresGroupDocsWriter` takes a per-group advisory lock — on its OWN `docs`-namespaced key, so it serializes against other docs writers to that group but not against KV writers, which draw on a separate ceiling. The bespoke `check_total_bytes` (with its own copy of the upper-bound fast path) is gone; group docs now shares `group_quota::check_group_write` with group KV, so the row ceiling, the projected-bytes ceiling, and the fast path that skips the O(n) SUM scan have exactly one implementation between them. tests/atomic_write.rs gains the docs row-quota race (16 concurrent creators against a ceiling of 4). Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/manager-core/src/atomic_write.rs | 547 +++++++++++++++++- crates/manager-core/src/docs_repo.rs | 154 +++-- crates/manager-core/src/docs_service.rs | 98 +--- crates/manager-core/src/group_docs_repo.rs | 220 ++++--- crates/manager-core/src/group_docs_service.rs | 202 +++---- crates/manager-core/tests/atomic_write.rs | 71 ++- crates/picloud/src/lib.rs | 18 +- 7 files changed, 959 insertions(+), 351 deletions(-) diff --git a/crates/manager-core/src/atomic_write.rs b/crates/manager-core/src/atomic_write.rs index 8d7cf42..836d76b 100644 --- a/crates/manager-core/src/atomic_write.rs +++ b/crates/manager-core/src/atomic_write.rs @@ -41,11 +41,14 @@ use std::sync::Arc; use async_trait::async_trait; use picloud_shared::{ - GroupId, GroupKvError, KvError, SdkCallCx, ServiceEvent, ServiceEventEmitter, + DocId, DocsError, GroupDocsError, GroupId, GroupKvError, KvError, SdkCallCx, ServiceEvent, + ServiceEventEmitter, }; use serde_json::Value; use sqlx::{PgConnection, PgPool}; +use crate::docs_repo::{self, DocsRepo}; +use crate::group_docs_repo::{self, GroupDocsRepo}; use crate::group_kv_repo::{self, GroupKvRepo}; use crate::group_quota::{check_group_write, GroupUsage, GroupWriteQuota, QuotaOutcome}; use crate::kv_repo::{self, KvRepo}; @@ -716,3 +719,545 @@ impl GroupKvWriter for BestEffortGroupKvWriter { Ok(deleted) } } + +// ---------------------------------------------------------------------------- +// Docs (per-app) +// ---------------------------------------------------------------------------- + +/// A docs mutation's `ServiceEvent`. The emitter turns `old_payload` into the +/// `prev_data` change-data-capture surface handlers see. +fn docs_event( + op: &'static str, + collection: &str, + id: DocId, + payload: Option, + old_payload: Option, +) -> ServiceEvent { + ServiceEvent { + source: "docs", + op, + collection: Some(collection.to_string()), + key: Some(id.to_string()), + payload, + old_payload, + } +} + +fn docs_tx_err(e: &sqlx::Error) -> DocsError { + DocsError::Backend(format!("database error: {e}")) +} + +/// The mutating half of the docs service: the write AND the trigger fan-out it +/// produces. Reads stay on `DocsRepo`. +#[async_trait] +pub trait DocsWriter: Send + Sync { + async fn create( + &self, + cx: &SdkCallCx, + collection: &str, + data: Value, + ) -> Result; + + /// `Ok(false)` when the doc does not exist (the caller maps that to + /// `NotFound`); no event is emitted in that case. + async fn update( + &self, + cx: &SdkCallCx, + collection: &str, + id: DocId, + data: Value, + ) -> Result; + + /// `Ok(false)` when the doc was already absent. + async fn delete(&self, cx: &SdkCallCx, collection: &str, id: DocId) -> Result; +} + +pub struct PostgresDocsWriter { + pool: PgPool, +} + +impl PostgresDocsWriter { + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { pool } + } +} + +#[async_trait] +impl DocsWriter for PostgresDocsWriter { + async fn create( + &self, + cx: &SdkCallCx, + collection: &str, + data: Value, + ) -> Result { + let mut tx = self.pool.begin().await.map_err(|e| docs_tx_err(&e))?; + let row = docs_repo::create_on(&mut *tx, cx.app_id, collection, data.clone()).await?; + emit_on( + &mut tx, + cx, + &docs_event("create", collection, row.id, Some(data), None), + ) + .await + .map_err(|e| DocsError::Backend(format!("event emit: {e}")))?; + tx.commit().await.map_err(|e| docs_tx_err(&e))?; + Ok(row.id) + } + + async fn update( + &self, + cx: &SdkCallCx, + collection: &str, + id: DocId, + data: Value, + ) -> Result { + let mut tx = self.pool.begin().await.map_err(|e| docs_tx_err(&e))?; + let previous = + docs_repo::update_on(&mut *tx, cx.app_id, collection, id, data.clone()).await?; + let Some(prev) = previous else { + return Ok(false); + }; + emit_on( + &mut tx, + cx, + &docs_event("update", collection, id, Some(data), Some(prev)), + ) + .await + .map_err(|e| DocsError::Backend(format!("event emit: {e}")))?; + tx.commit().await.map_err(|e| docs_tx_err(&e))?; + Ok(true) + } + + async fn delete(&self, cx: &SdkCallCx, collection: &str, id: DocId) -> Result { + let mut tx = self.pool.begin().await.map_err(|e| docs_tx_err(&e))?; + let previous = docs_repo::delete_on(&mut *tx, cx.app_id, collection, id).await?; + let Some(prev) = previous else { + return Ok(false); + }; + emit_on( + &mut tx, + cx, + &docs_event("delete", collection, id, None, Some(prev)), + ) + .await + .map_err(|e| DocsError::Backend(format!("event emit: {e}")))?; + tx.commit().await.map_err(|e| docs_tx_err(&e))?; + Ok(true) + } +} + +pub struct BestEffortDocsWriter { + repo: Arc, + events: Arc, +} + +impl BestEffortDocsWriter { + #[must_use] + pub fn new(repo: Arc, events: Arc) -> Self { + Self { repo, events } + } +} + +#[async_trait] +impl DocsWriter for BestEffortDocsWriter { + async fn create( + &self, + cx: &SdkCallCx, + collection: &str, + data: Value, + ) -> Result { + let row = self + .repo + .create(cx.app_id, collection, data.clone()) + .await?; + best_effort_emit( + &*self.events, + cx, + docs_event("create", collection, row.id, Some(data), None), + ) + .await; + Ok(row.id) + } + + async fn update( + &self, + cx: &SdkCallCx, + collection: &str, + id: DocId, + data: Value, + ) -> Result { + let Some(prev) = self + .repo + .update(cx.app_id, collection, id, data.clone()) + .await? + else { + return Ok(false); + }; + best_effort_emit( + &*self.events, + cx, + docs_event("update", collection, id, Some(data), Some(prev)), + ) + .await; + Ok(true) + } + + async fn delete(&self, cx: &SdkCallCx, collection: &str, id: DocId) -> Result { + let Some(prev) = self.repo.delete(cx.app_id, collection, id).await? else { + return Ok(false); + }; + best_effort_emit( + &*self.events, + cx, + docs_event("delete", collection, id, None, Some(prev)), + ) + .await; + Ok(true) + } +} + +// ---------------------------------------------------------------------------- +// §11.6 group (shared-collection) docs +// +// Same shape as group KV: the quota reads, the write, and the shared-trigger +// fan-out on one connection under a per-group advisory lock. +// ---------------------------------------------------------------------------- + +/// Where a group-docs mutation lands. `id` is `None` for a create (the row's id +/// is minted by the write). +#[derive(Clone, Copy)] +pub struct GroupDocsTarget<'a> { + pub group_id: GroupId, + pub collection: &'a str, +} + +#[async_trait] +pub trait GroupDocsWriter: Send + Sync { + async fn create( + &self, + cx: &SdkCallCx, + target: GroupDocsTarget<'_>, + data: Value, + quota: GroupWriteQuota, + ) -> Result; + + /// `Ok(false)` when the doc does not exist. + async fn update( + &self, + cx: &SdkCallCx, + target: GroupDocsTarget<'_>, + id: DocId, + data: Value, + quota: GroupWriteQuota, + ) -> Result; + + async fn delete( + &self, + cx: &SdkCallCx, + target: GroupDocsTarget<'_>, + id: DocId, + ) -> Result; +} + +fn group_docs_tx_err(e: &sqlx::Error) -> GroupDocsError { + GroupDocsError::Backend(format!("database error: {e}")) +} + +fn docs_quota_gate(outcome: QuotaOutcome) -> Result<(), GroupDocsError> { + match outcome { + QuotaOutcome::Allowed => Ok(()), + QuotaOutcome::RowsExceeded { limit, actual } => Err(GroupDocsError::QuotaExceeded { + limit: usize::try_from(limit).unwrap_or(usize::MAX), + actual: usize::try_from(actual).unwrap_or(usize::MAX), + }), + QuotaOutcome::BytesExceeded { limit, actual } => { + Err(GroupDocsError::TotalBytesQuotaExceeded { limit, actual }) + } + } +} + +/// Group-docs usage, read on the WRITING transaction's connection. +struct TxGroupDocsUsage<'a> { + conn: &'a mut PgConnection, + group_id: GroupId, + /// `Some(id)` when this write replaces an existing doc, so its bytes are + /// subtracted from the projection; `None` on create. + replacing: Option, + new_data: &'a Value, +} + +#[async_trait] +impl GroupUsage for TxGroupDocsUsage<'_> { + async fn count_rows(&mut self) -> Result { + group_docs_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_docs_repo::projected_total_bytes_on( + &mut *self.conn, + self.group_id, + self.replacing, + self.new_data, + ) + .await + .map_err(|e| e.to_string()) + } +} + +pub struct PostgresGroupDocsWriter { + pool: PgPool, +} + +impl PostgresGroupDocsWriter { + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { pool } + } +} + +#[async_trait] +impl GroupDocsWriter for PostgresGroupDocsWriter { + async fn create( + &self, + cx: &SdkCallCx, + t: GroupDocsTarget<'_>, + data: Value, + quota: GroupWriteQuota, + ) -> Result { + let mut tx = self.pool.begin().await.map_err(|e| group_docs_tx_err(&e))?; + lock_group_quota(&mut tx, t.group_id, "docs") + .await + .map_err(|e| group_docs_tx_err(&e))?; + + let mut usage = TxGroupDocsUsage { + conn: &mut tx, + group_id: t.group_id, + replacing: None, + new_data: &data, + }; + let outcome = check_group_write(quota, true, &mut usage) + .await + .map_err(GroupDocsError::Backend)?; + docs_quota_gate(outcome)?; + + let created = + group_docs_repo::create_on(&mut *tx, t.group_id, t.collection, data.clone()).await?; + emit_shared_on( + &mut tx, + cx, + t.group_id, + &docs_event("create", t.collection, created.id, Some(data), None), + ) + .await + .map_err(|e| GroupDocsError::Backend(format!("event emit: {e}")))?; + tx.commit().await.map_err(|e| group_docs_tx_err(&e))?; + Ok(created.id) + } + + async fn update( + &self, + cx: &SdkCallCx, + t: GroupDocsTarget<'_>, + id: DocId, + data: Value, + quota: GroupWriteQuota, + ) -> Result { + let mut tx = self.pool.begin().await.map_err(|e| group_docs_tx_err(&e))?; + lock_group_quota(&mut tx, t.group_id, "docs") + .await + .map_err(|e| group_docs_tx_err(&e))?; + + // An update replaces a row rather than adding one, so it is exempt from + // the row ceiling; the byte projection subtracts the replaced doc in SQL. + // A missing doc subtracts 0 and the write below no-ops to `None`. + let mut usage = TxGroupDocsUsage { + conn: &mut tx, + group_id: t.group_id, + replacing: Some(id), + new_data: &data, + }; + let outcome = check_group_write(quota, false, &mut usage) + .await + .map_err(GroupDocsError::Backend)?; + docs_quota_gate(outcome)?; + + let previous = + group_docs_repo::update_on(&mut *tx, t.group_id, t.collection, id, data.clone()) + .await?; + let Some(prev) = previous else { + return Ok(false); + }; + emit_shared_on( + &mut tx, + cx, + t.group_id, + &docs_event("update", t.collection, id, Some(data), Some(prev)), + ) + .await + .map_err(|e| GroupDocsError::Backend(format!("event emit: {e}")))?; + tx.commit().await.map_err(|e| group_docs_tx_err(&e))?; + Ok(true) + } + + async fn delete( + &self, + cx: &SdkCallCx, + t: GroupDocsTarget<'_>, + id: DocId, + ) -> 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_docs_tx_err(&e))?; + let previous = group_docs_repo::delete_on(&mut *tx, t.group_id, t.collection, id).await?; + let Some(prev) = previous else { + return Ok(false); + }; + emit_shared_on( + &mut tx, + cx, + t.group_id, + &docs_event("delete", t.collection, id, None, Some(prev)), + ) + .await + .map_err(|e| GroupDocsError::Backend(format!("event emit: {e}")))?; + tx.commit().await.map_err(|e| group_docs_tx_err(&e))?; + Ok(true) + } +} + +// ---- Best-effort group-docs writer (in-memory tests) ------------------------ + +struct RepoGroupDocsUsage<'a> { + repo: &'a dyn GroupDocsRepo, + group_id: GroupId, + replacing: Option, + new_data: &'a Value, +} + +#[async_trait] +impl GroupUsage for RepoGroupDocsUsage<'_> { + 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.replacing, self.new_data) + .await + .map_err(|e| e.to_string()) + } +} + +pub struct BestEffortGroupDocsWriter { + repo: Arc, + events: Arc, +} + +impl BestEffortGroupDocsWriter { + #[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, + group_id: GroupId, + adds_row: bool, + replacing: Option, + new_data: &Value, + quota: GroupWriteQuota, + ) -> Result<(), GroupDocsError> { + let mut usage = RepoGroupDocsUsage { + repo: &*self.repo, + group_id, + replacing, + new_data, + }; + let outcome = check_group_write(quota, adds_row, &mut usage) + .await + .map_err(GroupDocsError::Backend)?; + docs_quota_gate(outcome) + } +} + +#[async_trait] +impl GroupDocsWriter for BestEffortGroupDocsWriter { + async fn create( + &self, + cx: &SdkCallCx, + t: GroupDocsTarget<'_>, + data: Value, + quota: GroupWriteQuota, + ) -> Result { + self.check(t.group_id, true, None, &data, quota).await?; + let created = self + .repo + .create(t.group_id, t.collection, data.clone()) + .await?; + self.emit( + cx, + t.group_id, + docs_event("create", t.collection, created.id, Some(data), None), + ) + .await; + Ok(created.id) + } + + async fn update( + &self, + cx: &SdkCallCx, + t: GroupDocsTarget<'_>, + id: DocId, + data: Value, + quota: GroupWriteQuota, + ) -> Result { + self.check(t.group_id, false, Some(id), &data, quota) + .await?; + let Some(prev) = self + .repo + .update(t.group_id, t.collection, id, data.clone()) + .await? + else { + return Ok(false); + }; + self.emit( + cx, + t.group_id, + docs_event("update", t.collection, id, Some(data), Some(prev)), + ) + .await; + Ok(true) + } + + async fn delete( + &self, + cx: &SdkCallCx, + t: GroupDocsTarget<'_>, + id: DocId, + ) -> Result { + let Some(prev) = self.repo.delete(t.group_id, t.collection, id).await? else { + return Ok(false); + }; + self.emit( + cx, + t.group_id, + docs_event("delete", t.collection, id, None, Some(prev)), + ) + .await; + Ok(true) + } +} diff --git a/crates/manager-core/src/docs_repo.rs b/crates/manager-core/src/docs_repo.rs index eeb3264..80efb5b 100644 --- a/crates/manager-core/src/docs_repo.rs +++ b/crates/manager-core/src/docs_repo.rs @@ -118,24 +118,7 @@ impl DocsRepo for PostgresDocsRepo { collection: &str, data: Value, ) -> Result { - let id = Uuid::new_v4(); - let row: (DateTime, DateTime) = sqlx::query_as( - "INSERT INTO docs (app_id, collection, id, data) \ - VALUES ($1, $2, $3, $4) \ - RETURNING created_at, updated_at", - ) - .bind(app_id.into_inner()) - .bind(collection) - .bind(id) - .bind(&data) - .fetch_one(&self.pool) - .await?; - Ok(DocRow { - id, - data, - created_at: row.0, - updated_at: row.1, - }) + create_on(&self.pool, app_id, collection, data).await } async fn get( @@ -179,34 +162,7 @@ impl DocsRepo for PostgresDocsRepo { id: DocId, data: Value, ) -> Result, DocsRepoError> { - // Same CTE shape as KV's set ([kv_repo.rs:101-132]): SELECT the - // previous data before the UPDATE so the service can emit - // `prev_data` in the update ServiceEvent. Single statement, no - // explicit transaction. Inherits KV's last-writer-wins race - // under concurrent writers; documented as a known limitation - // for v1.1.2. - let row: Option<(Option,)> = sqlx::query_as( - "WITH prev AS ( \ - SELECT data FROM docs \ - WHERE app_id = $1 AND collection = $2 AND id = $3 \ - ), \ - updated AS ( \ - UPDATE docs SET data = $4, updated_at = NOW() \ - WHERE app_id = $1 AND collection = $2 AND id = $3 \ - RETURNING 1 \ - ) \ - SELECT (SELECT data FROM prev) FROM updated", - ) - .bind(app_id.into_inner()) - .bind(collection) - .bind(id) - .bind(&data) - .fetch_optional(&self.pool) - .await?; - // `row` is None when the UPDATE matched no rows (missing doc); - // Some((Some(prev),)) on success. `data` is JSONB NOT NULL so - // the inner Option is always Some when prev exists. - Ok(row.and_then(|(v,)| v)) + update_on(&self.pool, app_id, collection, id, data).await } async fn delete( @@ -215,17 +171,7 @@ impl DocsRepo for PostgresDocsRepo { collection: &str, id: DocId, ) -> Result, DocsRepoError> { - let row: Option<(Value,)> = sqlx::query_as( - "DELETE FROM docs \ - WHERE app_id = $1 AND collection = $2 AND id = $3 \ - RETURNING data", - ) - .bind(app_id.into_inner()) - .bind(collection) - .bind(id) - .fetch_optional(&self.pool) - .await?; - Ok(row.map(|(v,)| v)) + delete_on(&self.pool, app_id, collection, id).await } async fn list( @@ -448,6 +394,100 @@ fn value_to_text(v: &Value) -> Option { // pin the cross-app isolation invariant at the SQL level. // ---------------------------------------------------------------------------- +// ---------------------------------------------------------------------------- +// Connection-scoped mutations +// +// Generic over the executor so the same SQL serves the pooled trait methods +// above and `crate::atomic_write`, which runs the write and the trigger fan-out +// it produces on ONE connection inside ONE transaction. +// ---------------------------------------------------------------------------- + +pub(crate) async fn create_on<'c, E>( + exec: E, + app_id: AppId, + collection: &str, + data: Value, +) -> Result +where + E: sqlx::PgExecutor<'c>, +{ + let id = Uuid::new_v4(); + let row: (DateTime, DateTime) = sqlx::query_as( + "INSERT INTO docs (app_id, collection, id, data) \ + VALUES ($1, $2, $3, $4) \ + RETURNING created_at, updated_at", + ) + .bind(app_id.into_inner()) + .bind(collection) + .bind(id) + .bind(&data) + .fetch_one(exec) + .await?; + Ok(DocRow { + id, + data, + created_at: row.0, + updated_at: row.1, + }) +} + +/// Returns the previous data (for the update event), `None` if the doc is +/// missing. The CTE captures the prior data alongside the UPDATE so the service +/// can emit `prev_data` without a second round trip. +pub(crate) async fn update_on<'c, E>( + exec: E, + app_id: AppId, + collection: &str, + id: DocId, + data: Value, +) -> Result, DocsRepoError> +where + E: sqlx::PgExecutor<'c>, +{ + let row: Option<(Option,)> = sqlx::query_as( + "WITH prev AS ( \ + SELECT data FROM docs \ + WHERE app_id = $1 AND collection = $2 AND id = $3 \ + ), \ + updated AS ( \ + UPDATE docs SET data = $4, updated_at = NOW() \ + WHERE app_id = $1 AND collection = $2 AND id = $3 \ + RETURNING 1 \ + ) \ + SELECT (SELECT data FROM prev) FROM updated", + ) + .bind(app_id.into_inner()) + .bind(collection) + .bind(id) + .bind(&data) + .fetch_optional(exec) + .await?; + Ok(row.and_then(|(v,)| v)) +} + +/// Returns the deleted doc's data if it existed, `None` if no such doc. +pub(crate) async fn delete_on<'c, E>( + exec: E, + app_id: AppId, + collection: &str, + id: DocId, +) -> Result, DocsRepoError> +where + E: sqlx::PgExecutor<'c>, +{ + let row: Option<(Value,)> = sqlx::query_as( + "DELETE FROM docs \ + WHERE app_id = $1 AND collection = $2 AND id = $3 \ + RETURNING data", + ) + .bind(app_id.into_inner()) + .bind(collection) + .bind(id) + .fetch_optional(exec) + .await?; + Ok(row.map(|(v,)| v)) +} + #[cfg(test)] mod sql_shape_tests { use super::*; diff --git a/crates/manager-core/src/docs_service.rs b/crates/manager-core/src/docs_service.rs index 42308e2..7e30223 100644 --- a/crates/manager-core/src/docs_service.rs +++ b/crates/manager-core/src/docs_service.rs @@ -26,10 +26,10 @@ use std::sync::Arc; use async_trait::async_trait; use picloud_shared::{ - DocId, DocRow, DocsError, DocsListPage, DocsService, SdkCallCx, ServiceEvent, - ServiceEventEmitter, + DocId, DocRow, DocsError, DocsListPage, DocsService, SdkCallCx, ServiceEventEmitter, }; +use crate::atomic_write::{BestEffortDocsWriter, DocsWriter, PostgresDocsWriter}; use crate::authz::{self, AuthzRepo, Capability}; use crate::docs_filter::{parse_filter, FilterParseError}; use crate::docs_repo::{DocsRepo, DocsRepoError}; @@ -55,9 +55,11 @@ pub fn docs_max_value_bytes_from_env() -> usize { } pub struct DocsServiceImpl { + /// Reads only. Mutations go through `writer`, which owns the write AND the + /// trigger fan-out so the two can share a transaction. repo: Arc, authz: Arc, - events: Arc, + writer: Arc, max_value_bytes: usize, } @@ -79,13 +81,23 @@ impl DocsServiceImpl { max_value_bytes: usize, ) -> Self { Self { + writer: Arc::new(BestEffortDocsWriter::new(repo.clone(), events)), repo, authz, - events, max_value_bytes, } } + /// Swap the best-effort writer for the transactional one: the write and its + /// trigger fan-out then commit together, so an outbox failure rolls the + /// write back instead of silently losing the event. The host always calls + /// this; the in-memory unit tests do not. + #[must_use] + pub fn with_atomic_writes(mut self, pool: sqlx::PgPool) -> Self { + self.writer = Arc::new(PostgresDocsWriter::new(pool)); + self + } + fn check_data_size(&self, data: &serde_json::Value) -> Result<(), DocsError> { let encoded_len = serde_json::to_vec(data) .map(|v| v.len()) @@ -163,30 +175,7 @@ impl DocsService for DocsServiceImpl { validate_data(&data)?; self.check_data_size(&data)?; self.check_write(cx).await?; - let row = self - .repo - .create(cx.app_id, collection, data.clone()) - .await?; - // Best-effort emit — a failed emit logs but does not roll back - // the write (mirrors KV's pattern). - if let Err(e) = self - .events - .emit( - cx, - ServiceEvent { - source: "docs", - op: "create", - collection: Some(collection.to_string()), - key: Some(row.id.to_string()), - payload: Some(data), - old_payload: None, - }, - ) - .await - { - tracing::error!(error = %e, source = "docs", op = "create", event_emit_failure = true, "event emit failed"); - } - Ok(row.id) + self.writer.create(cx, collection, data).await } async fn get( @@ -241,60 +230,17 @@ impl DocsService for DocsServiceImpl { validate_data(&data)?; self.check_data_size(&data)?; self.check_write(cx).await?; - let previous = self - .repo - .update(cx.app_id, collection, id, data.clone()) - .await?; - match previous { - Some(prev) => { - if let Err(e) = self - .events - .emit( - cx, - ServiceEvent { - source: "docs", - op: "update", - collection: Some(collection.to_string()), - key: Some(id.to_string()), - payload: Some(data), - old_payload: Some(prev), - }, - ) - .await - { - tracing::error!(error = %e, source = "docs", op = "update", event_emit_failure = true, "event emit failed"); - } - Ok(()) - } - None => Err(DocsError::NotFound), + if self.writer.update(cx, collection, id, data).await? { + Ok(()) + } else { + Err(DocsError::NotFound) } } async fn delete(&self, cx: &SdkCallCx, collection: &str, id: DocId) -> Result { validate_collection(collection)?; self.check_write(cx).await?; - let previous = self.repo.delete(cx.app_id, collection, id).await?; - let was_present = previous.is_some(); - if let Some(prev) = previous { - if let Err(e) = self - .events - .emit( - cx, - ServiceEvent { - source: "docs", - op: "delete", - collection: Some(collection.to_string()), - key: Some(id.to_string()), - payload: None, - old_payload: Some(prev), - }, - ) - .await - { - tracing::error!(error = %e, source = "docs", op = "delete", event_emit_failure = true, "event emit failed"); - } - } - Ok(was_present) + self.writer.delete(cx, collection, id).await } async fn list( diff --git a/crates/manager-core/src/group_docs_repo.rs b/crates/manager-core/src/group_docs_repo.rs index 5ecb828..4c9ac8e 100644 --- a/crates/manager-core/src/group_docs_repo.rs +++ b/crates/manager-core/src/group_docs_repo.rs @@ -136,24 +136,7 @@ impl GroupDocsRepo for PostgresGroupDocsRepo { collection: &str, data: Value, ) -> Result { - let id = Uuid::new_v4(); - let row: (DateTime, DateTime) = sqlx::query_as( - "INSERT INTO group_docs (group_id, collection, id, data) \ - VALUES ($1, $2, $3, $4) \ - RETURNING created_at, updated_at", - ) - .bind(group_id.into_inner()) - .bind(collection) - .bind(id) - .bind(&data) - .fetch_one(&self.pool) - .await?; - Ok(DocRow { - id, - data, - created_at: row.0, - updated_at: row.1, - }) + create_on(&self.pool, group_id, collection, data).await } async fn get( @@ -206,25 +189,7 @@ impl GroupDocsRepo for PostgresGroupDocsRepo { id: DocId, data: Value, ) -> Result, GroupDocsRepoError> { - let row: Option<(Option,)> = sqlx::query_as( - "WITH prev AS ( \ - SELECT data FROM group_docs \ - WHERE group_id = $1 AND collection = $2 AND id = $3 \ - ), \ - updated AS ( \ - UPDATE group_docs SET data = $4, updated_at = NOW() \ - WHERE group_id = $1 AND collection = $2 AND id = $3 \ - RETURNING 1 \ - ) \ - SELECT (SELECT data FROM prev) FROM updated", - ) - .bind(group_id.into_inner()) - .bind(collection) - .bind(id) - .bind(&data) - .fetch_optional(&self.pool) - .await?; - Ok(row.and_then(|(v,)| v)) + update_on(&self.pool, group_id, collection, id, data).await } async fn delete( @@ -233,17 +198,7 @@ impl GroupDocsRepo for PostgresGroupDocsRepo { collection: &str, id: DocId, ) -> Result, GroupDocsRepoError> { - let row: Option<(Value,)> = sqlx::query_as( - "DELETE FROM group_docs \ - WHERE group_id = $1 AND collection = $2 AND id = $3 \ - RETURNING data", - ) - .bind(group_id.into_inner()) - .bind(collection) - .bind(id) - .fetch_optional(&self.pool) - .await?; - Ok(row.map(|(v,)| v)) + delete_on(&self.pool, group_id, collection, id).await } async fn list( @@ -299,11 +254,7 @@ impl GroupDocsRepo for PostgresGroupDocsRepo { } async fn count_rows(&self, group_id: GroupId) -> Result { - let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_docs 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 { @@ -323,29 +274,152 @@ impl GroupDocsRepo for PostgresGroupDocsRepo { replacing: Option, new_data: &serde_json::Value, ) -> Result { - // All three terms use `octet_length(data::text)` (jsonb canonical). The - // new value is bound as compact TEXT and re-parsed through - // `::jsonb::text` so PG measures it in the SAME canonical form it stores. - // `replacing = None` (create) → the subtrahend row never matches → 0. - let new_text = serde_json::to_string(new_data).unwrap_or_else(|_| "null".to_string()); - let (n,): (i64,) = sqlx::query_as( - "SELECT ( \ - COALESCE((SELECT SUM(octet_length(data::text)) \ - FROM group_docs WHERE group_id = $1), 0) \ - - COALESCE((SELECT octet_length(data::text) \ - FROM group_docs WHERE group_id = $1 AND id = $2), 0) \ - + octet_length($3::jsonb::text) \ - )::BIGINT", - ) - .bind(group_id.into_inner()) - .bind(replacing) - .bind(new_text) - .fetch_one(&self.pool) - .await?; - Ok(u64::try_from(n).unwrap_or(0)) + projected_total_bytes_on(&self.pool, group_id, replacing, new_data).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 (audit #6/#8). +// ---------------------------------------------------------------------------- + +pub(crate) async fn create_on<'c, E>( + exec: E, + group_id: GroupId, + collection: &str, + data: Value, +) -> Result +where + E: sqlx::PgExecutor<'c>, +{ + let id = Uuid::new_v4(); + let row: (DateTime, DateTime) = sqlx::query_as( + "INSERT INTO group_docs (group_id, collection, id, data) \ + VALUES ($1, $2, $3, $4) \ + RETURNING created_at, updated_at", + ) + .bind(group_id.into_inner()) + .bind(collection) + .bind(id) + .bind(&data) + .fetch_one(exec) + .await?; + Ok(DocRow { + id, + data, + created_at: row.0, + updated_at: row.1, + }) +} + +/// Returns the previous data (for the update event), `None` if the doc is absent. +pub(crate) async fn update_on<'c, E>( + exec: E, + group_id: GroupId, + collection: &str, + id: DocId, + data: Value, +) -> Result, GroupDocsRepoError> +where + E: sqlx::PgExecutor<'c>, +{ + let row: Option<(Option,)> = sqlx::query_as( + "WITH prev AS ( \ + SELECT data FROM group_docs \ + WHERE group_id = $1 AND collection = $2 AND id = $3 \ + ), \ + updated AS ( \ + UPDATE group_docs SET data = $4, updated_at = NOW() \ + WHERE group_id = $1 AND collection = $2 AND id = $3 \ + RETURNING 1 \ + ) \ + SELECT (SELECT data FROM prev) FROM updated", + ) + .bind(group_id.into_inner()) + .bind(collection) + .bind(id) + .bind(&data) + .fetch_optional(exec) + .await?; + Ok(row.and_then(|(v,)| v)) +} + +pub(crate) async fn delete_on<'c, E>( + exec: E, + group_id: GroupId, + collection: &str, + id: DocId, +) -> Result, GroupDocsRepoError> +where + E: sqlx::PgExecutor<'c>, +{ + let row: Option<(Value,)> = sqlx::query_as( + "DELETE FROM group_docs \ + WHERE group_id = $1 AND collection = $2 AND id = $3 \ + RETURNING data", + ) + .bind(group_id.into_inner()) + .bind(collection) + .bind(id) + .fetch_optional(exec) + .await?; + Ok(row.map(|(v,)| v)) +} + +/// Total doc count across the group's shared-docs 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_docs WHERE group_id = $1") + .bind(group_id.into_inner()) + .fetch_one(exec) + .await?; + Ok(u64::try_from(n).unwrap_or(0)) +} + +/// §11.6 M4 quota: the PROJECTED total stored bytes for the group AFTER writing +/// `new_data` — current SUM, minus the replaced doc's bytes (`replacing = +/// Some(id)` on update, `None` on create), plus the new data's. Computed +/// entirely in SQL so all three terms use the SAME canonical metric +/// (`octet_length(data::text)`). +pub(crate) async fn projected_total_bytes_on<'c, E>( + exec: E, + group_id: GroupId, + replacing: Option, + new_data: &serde_json::Value, +) -> Result +where + E: sqlx::PgExecutor<'c>, +{ + // The new data is bound as compact TEXT and re-parsed through `::jsonb::text` + // so PG measures it in the same canonical form it stores. `replacing = None` + // (create) → the subtrahend row never matches → 0. + let new_text = serde_json::to_string(new_data).unwrap_or_else(|_| "null".to_string()); + let (n,): (i64,) = sqlx::query_as( + "SELECT ( \ + COALESCE((SELECT SUM(octet_length(data::text)) \ + FROM group_docs WHERE group_id = $1), 0) \ + - COALESCE((SELECT octet_length(data::text) \ + FROM group_docs WHERE group_id = $1 AND id = $2), 0) \ + + octet_length($3::jsonb::text) \ + )::BIGINT", + ) + .bind(group_id.into_inner()) + .bind(replacing) + .bind(new_text) + .fetch_one(exec) + .await?; + Ok(u64::try_from(n).unwrap_or(0)) +} + fn encode_cursor(last_id: &Uuid) -> String { URL_SAFE_NO_PAD.encode(last_id.as_bytes()) } diff --git a/crates/manager-core/src/group_docs_service.rs b/crates/manager-core/src/group_docs_service.rs index f1ab1cc..2b3c557 100644 --- a/crates/manager-core/src/group_docs_service.rs +++ b/crates/manager-core/src/group_docs_service.rs @@ -12,14 +12,18 @@ use std::sync::Arc; use async_trait::async_trait; use picloud_shared::{ DocId, DocRow, DocsListPage, GroupDocsError, GroupDocsService, GroupId, NoopEventEmitter, - SdkCallCx, ServiceEvent, ServiceEventEmitter, + SdkCallCx, ServiceEventEmitter, }; +use crate::atomic_write::{ + BestEffortGroupDocsWriter, GroupDocsTarget, GroupDocsWriter, PostgresGroupDocsWriter, +}; use crate::authz::{self, AuthzRepo, Capability}; use crate::docs_filter::{parse_filter, FilterParseError}; use crate::docs_service::docs_max_value_bytes_from_env; use crate::group_collection_repo::GroupCollectionResolver; use crate::group_docs_repo::{GroupDocsRepo, GroupDocsRepoError}; +use crate::group_quota::GroupWriteQuota; /// The registry `kind` this service resolves. const KIND_DOCS: &str = "docs"; @@ -36,8 +40,10 @@ pub struct GroupDocsServiceImpl { /// shared-docs collections (`PICLOUD_GROUP_DOCS_MAX_TOTAL_BYTES`). Checked on /// the projected total so a same/smaller update near the cap is still allowed. max_total_bytes: u64, - /// §11.6: fires `shared = true` docs triggers on a shared-collection write. - 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 GroupDocsServiceImpl { @@ -50,10 +56,21 @@ impl GroupDocsServiceImpl { Self::with_max_value_bytes(repo, resolver, authz, docs_max_value_bytes_from_env()) } - /// Wire the event emitter (§11.6 shared-collection triggers). + /// Wire the event emitter (§11.6 shared-collection triggers) on the + /// best-effort writer. Superseded by [`Self::with_atomic_writes`]. #[must_use] pub fn with_events(mut self, events: Arc) -> Self { - self.events = events; + self.writer = Arc::new(BestEffortGroupDocsWriter::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, so the quota is a real bound and an emit failure rolls the + /// write back. Supersedes [`Self::with_events`]. + #[must_use] + pub fn with_atomic_writes(mut self, pool: sqlx::PgPool) -> Self { + self.writer = Arc::new(PostgresGroupDocsWriter::new(pool)); self } @@ -65,7 +82,10 @@ impl GroupDocsServiceImpl { max_value_bytes: usize, ) -> Self { Self { - events: Arc::new(NoopEventEmitter), + writer: Arc::new(BestEffortGroupDocsWriter::new( + repo.clone(), + Arc::new(NoopEventEmitter), + )), max_rows: crate::group_quota::group_docs_max_rows_from_env(), max_total_bytes: crate::group_quota::group_docs_max_total_bytes_from_env(), repo, @@ -82,75 +102,16 @@ impl GroupDocsServiceImpl { self } - /// §11.6 M4: reject a write whose PROJECTED total (old doc's bytes subtracted, - /// new doc's added — `old_len = 0` for a create) would exceed the group's - /// total-bytes ceiling. `new_len` is the JSON-encoded size of the incoming - /// data (already validated ≤ per-doc cap). `rows_after` is the doc count the - /// collection would hold post-write. - /// - /// Fast path (matches the KV service): every doc is ≤ `max_value_bytes`, so - /// the collection holds at most `rows_after * max_value_bytes`. When that - /// exact upper bound fits under the ceiling the write can't exceed it — skip - /// the O(collection-size) `SUM(octet_length(...))` scan. It runs only near cap. - async fn check_total_bytes( - &self, - group_id: GroupId, - replacing: Option, - new_data: &serde_json::Value, - rows_after: u64, - ) -> Result<(), GroupDocsError> { - let upper_bound = rows_after.saturating_mul(self.max_value_bytes as u64); - if upper_bound <= self.max_total_bytes { - return Ok(()); + /// The ceilings a shared-collection write must fit under. The row/byte + /// policy itself lives in `group_quota::check_group_write`, shared with KV. + fn quota(&self) -> GroupWriteQuota { + GroupWriteQuota { + max_rows: self.max_rows, + max_total_bytes: self.max_total_bytes, + max_value_bytes: self.max_value_bytes, } - // Consistent-metric projection (canonical `octet_length(data::text)` for - // the current SUM, the replaced doc, and the new data) — no Rust-compact - // vs DB-canonical drift. - let projected = self - .repo - .projected_total_bytes(group_id, replacing, new_data) - .await?; - if projected > self.max_total_bytes { - return Err(GroupDocsError::TotalBytesQuotaExceeded { - limit: self.max_total_bytes, - actual: projected, - }); - } - Ok(()) } - /// §11.6: fire `shared = true` docs triggers on the owning group. - /// Best-effort; an emit failure is logged, never surfaced. - #[allow(clippy::too_many_arguments)] // op/collection/id + new + old payloads - async fn emit_shared( - &self, - cx: &SdkCallCx, - group_id: GroupId, - op: &'static str, - collection: &str, - id: &str, - payload: Option, - old_payload: Option, - ) { - crate::group_collection_repo::best_effort_emit_shared( - &*self.events, - cx, - group_id, - ServiceEvent { - source: "docs", - op, - collection: Some(collection.to_string()), - key: Some(id.to_string()), - payload, - old_payload, - }, - ) - .await; - } - - /// The structural boundary: resolve the collection to its owning group - /// (kind=`docs`) on the calling app's chain. `CollectionNotShared` when no - /// ancestor group declares it. async fn owning_group( &self, cx: &SdkCallCx, @@ -238,29 +199,17 @@ impl GroupDocsService for GroupDocsServiceImpl { validate_data(&data)?; self.check_data_size(&data)?; self.check_write(cx, group_id).await?; - // §11.6 quota: a new doc must fit under the group's row ceiling. - let count = self.repo.count_rows(group_id).await?; - if count >= self.max_rows { - return Err(GroupDocsError::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: a create only adds bytes (nothing replaced), one row. - self.check_total_bytes(group_id, None, &data, count + 1) - .await?; - let created = self.repo.create(group_id, collection, data.clone()).await?; - self.emit_shared( - cx, - group_id, - "create", - collection, - &created.id.to_string(), - Some(data), - None, - ) - .await; - Ok(created.id) + self.writer + .create( + cx, + GroupDocsTarget { + group_id, + collection, + }, + data, + self.quota(), + ) + .await } async fn get( @@ -313,34 +262,23 @@ impl GroupDocsService for GroupDocsServiceImpl { validate_data(&data)?; self.check_data_size(&data)?; self.check_write(cx, group_id).await?; - // §11.6 M4 byte quota: check the projected total before writing. The - // projection subtracts the replaced doc's bytes in SQL, so there is no - // need to fetch the old doc here (a missing doc subtracts 0, and the - // update below no-ops to `None`). - // An update leaves the row count unchanged; the cheap COUNT lets the byte - // check skip the text SUM when the collection is comfortably under cap. - let count = self.repo.count_rows(group_id).await?; - self.check_total_bytes(group_id, Some(id), &data, count) - .await?; - match self - .repo - .update(group_id, collection, id, data.clone()) + if self + .writer + .update( + cx, + GroupDocsTarget { + group_id, + collection, + }, + id, + data, + self.quota(), + ) .await? { - Some(prev) => { - self.emit_shared( - cx, - group_id, - "update", - collection, - &id.to_string(), - Some(data), - Some(prev), - ) - .await; - Ok(()) - } - None => Err(GroupDocsError::NotFound), + Ok(()) + } else { + Err(GroupDocsError::NotFound) } } @@ -352,22 +290,16 @@ impl GroupDocsService for GroupDocsServiceImpl { ) -> Result { let group_id = self.owning_group(cx, collection).await?; self.check_write(cx, group_id).await?; - match self.repo.delete(group_id, collection, id).await? { - Some(prev) => { - self.emit_shared( - cx, + self.writer + .delete( + cx, + GroupDocsTarget { group_id, - "delete", collection, - &id.to_string(), - None, - Some(prev), - ) - .await; - Ok(true) - } - None => Ok(false), - } + }, + id, + ) + .await } async fn list( diff --git a/crates/manager-core/tests/atomic_write.rs b/crates/manager-core/tests/atomic_write.rs index 384e1ab..60d0027 100644 --- a/crates/manager-core/tests/atomic_write.rs +++ b/crates/manager-core/tests/atomic_write.rs @@ -16,10 +16,13 @@ use std::sync::Arc; use picloud_manager_core::atomic_write::{ - GroupKvTarget, GroupKvWriter, KvWriter, PostgresGroupKvWriter, PostgresKvWriter, + GroupDocsTarget, GroupDocsWriter, GroupKvTarget, GroupKvWriter, KvWriter, + PostgresGroupDocsWriter, PostgresGroupKvWriter, PostgresKvWriter, }; use picloud_manager_core::group_quota::GroupWriteQuota; -use picloud_shared::{AppId, ExecutionId, GroupId, GroupKvError, RequestId, ScriptId, SdkCallCx}; +use picloud_shared::{ + AppId, ExecutionId, GroupDocsError, GroupId, GroupKvError, RequestId, ScriptId, SdkCallCx, +}; use sqlx::postgres::PgPoolOptions; use sqlx::PgPool; use uuid::Uuid; @@ -436,3 +439,67 @@ async fn concurrent_writers_cannot_push_a_group_past_its_byte_quota() { .await .expect("cleanup"); } + +/// Group DOCS carried the identical check-then-write race as group KV, and the +/// identical fix (a per-group advisory lock, on its own `docs`-namespaced key so +/// it doesn't contend with KV writes). +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_writers_cannot_push_a_group_past_its_docs_row_quota() { + const MAX_DOCS: u64 = 4; + + let Some(pool) = pool_or_skip().await else { + return; + }; + let group = mk_group(&pool).await; + let writer = Arc::new(PostgresGroupDocsWriter::new(pool.clone())); + let quota = GroupWriteQuota { + max_rows: MAX_DOCS, + max_total_bytes: u64::MAX, + max_value_bytes: 256 * 1024, + }; + + let app = Uuid::new_v4(); + let mut set = tokio::task::JoinSet::new(); + for i in 0..16 { + let w = Arc::clone(&writer); + set.spawn(async move { + let cx = cx(app); + w.create( + &cx, + GroupDocsTarget { + group_id: GroupId::from(group), + collection: "articles", + }, + serde_json::json!({ "n": i }), + quota, + ) + .await + }); + } + let mut ok = 0; + while let Some(r) = set.join_next().await { + match r.expect("task") { + Ok(_) => ok += 1, + Err(GroupDocsError::QuotaExceeded { .. }) => {} + Err(e) => panic!("unexpected error: {e}"), + } + } + + let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM group_docs WHERE group_id = $1") + .bind(group) + .fetch_one(&pool) + .await + .expect("count"); + assert_eq!( + n, + i64::try_from(MAX_DOCS).unwrap(), + "the group must hold EXACTLY its doc ceiling" + ); + assert_eq!(ok, usize::try_from(MAX_DOCS).unwrap()); + + 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 2f9ed84..0570bba 100644 --- a/crates/picloud/src/lib.rs +++ b/crates/picloud/src/lib.rs @@ -209,14 +209,18 @@ pub async fn build_app( authz.clone(), picloud_manager_core::docs_service::docs_max_value_bytes_from_env(), ) - .with_events(events.clone()), + // Quota reads + write + shared fan-out in one advisory-locked tx. + .with_atomic_writes(pool.clone()), + ); + let docs: Arc = Arc::new( + DocsServiceImpl::with_max_value_bytes( + docs_repo, + authz.clone(), + events.clone(), + picloud_manager_core::docs_service::docs_max_value_bytes_from_env(), + ) + .with_atomic_writes(pool.clone()), ); - let docs: Arc = Arc::new(DocsServiceImpl::with_max_value_bytes( - docs_repo, - authz.clone(), - events.clone(), - picloud_manager_core::docs_service::docs_max_value_bytes_from_env(), - )); let dl_service: Arc = Arc::new(PostgresDeadLetterService::new( dl_repo.clone(), outbox_repo.clone(),