fabi/PiCloud
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

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

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

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

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

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

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

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

173
crates/manager-core/tests/atomic_write.rs
View File

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