diff --git a/crates/manager-core/src/atomic_write.rs b/crates/manager-core/src/atomic_write.rs index 087a45f..8e649d9 100644 --- a/crates/manager-core/src/atomic_write.rs +++ b/crates/manager-core/src/atomic_write.rs @@ -135,15 +135,21 @@ impl PostgresKvWriter { Self { pool, max_rows } } - /// Refuse a key-adding write once the app is at its ceiling. Runs on the - /// writing transaction's connection. - async fn check_rows( + /// Refuse a key-adding write once the app is over its ceiling. + /// + /// Called AFTER the insert, so the count already includes the new key — hence + /// `>` rather than `>=`. Checking afterwards lets us reuse the previous value + /// the upsert already returns, instead of paying a second SELECT just to learn + /// whether the write adds a row; the transaction rolls the insert back on + /// refusal, so the outcome is identical. `kv::set` is the hottest write path in + /// the system and does not need a redundant round-trip. + async fn check_rows_after_insert( &self, conn: &mut PgConnection, app_id: picloud_shared::AppId, ) -> Result<(), KvError> { let count = kv_repo::count_rows_on(&mut *conn, app_id).await?; - if count >= self.max_rows { + if count > self.max_rows { return Err(KvError::QuotaExceeded { limit: self.max_rows, actual: count, @@ -163,14 +169,13 @@ impl KvWriter for PostgresKvWriter { value: Value, ) -> Result, KvError> { let mut tx = self.pool.begin().await.map_err(|e| tx_err(&e))?; - // A NEW key consumes a row against the app's ceiling; an update does not. - if kv_repo::get_on(&mut *tx, cx.app_id, collection, key) - .await? - .is_none() - { - self.check_rows(&mut tx, cx.app_id).await?; - } let previous = kv_repo::set_on(&mut *tx, cx.app_id, collection, key, value.clone()).await?; + // `previous.is_none()` means the upsert INSERTED, i.e. this write consumed + // a row against the app's ceiling. An update is net-zero and pays nothing — + // not even the COUNT. + if previous.is_none() { + self.check_rows_after_insert(&mut tx, cx.app_id).await?; + } let op = if previous.is_some() { "update" } else { @@ -200,18 +205,16 @@ impl KvWriter for PostgresKvWriter { "insert" }; let mut tx = self.pool.begin().await.map_err(|e| tx_err(&e))?; - // A swap only ADDS a key 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. - if expected.is_none() - && kv_repo::get_on(&mut *tx, cx.app_id, collection, key) - .await? - .is_none() - { - self.check_rows(&mut tx, cx.app_id).await?; - } + // A swap only ADDS a key when the precondition was "absent" AND it actually + // inserted. With a `Some` precondition it is an UPDATE (net-zero rows), and + // against an absent key it cannot swap at all — either way, nothing to + // charge. So the ceiling is checked only on an insert that really happened. + let inserted = expected.is_none(); let swapped = kv_repo::set_if_on(&mut *tx, cx.app_id, collection, key, expected, new.clone()).await?; + if swapped && inserted { + self.check_rows_after_insert(&mut tx, cx.app_id).await?; + } if swapped { let event = kv_event(op, collection, key, Some(new), None); emit_on(&mut tx, cx, &event) @@ -245,6 +248,15 @@ impl KvWriter for PostgresKvWriter { // Best-effort (repo trait + emitter trait) // ---------------------------------------------------------------------------- +/// A `files_event` that failed to serialize. Never silently dropped — that is the +/// exact class of loss this branch exists to remove. +fn log_event_build_failure(op: &str, e: &str) { + tracing::error!( + error = %e, source = "files", op, event_emit_failure = true, + "could not build the file event — the write committed but its triggers will not fire" + ); +} + /// Emit, downgrading a failure to an error-level log. The write has already /// committed by this point, so there is nothing useful to return to the caller — /// which is exactly the durability hole `PostgresKvWriter` exists to close. @@ -1447,15 +1459,21 @@ impl FilesWriter for PostgresFilesWriter { let owner = files_repo::app_owner_dir(cx.app_id); let id = Uuid::new_v4(); let size = i64::try_from(new.data.len()).unwrap_or(i64::MAX); - // The bytes go down first — they cannot be part of the transaction. - let checksum = files_repo::write_atomic_at(&self.root, &owner, collection, id, &new.data)?; let committed = async { let mut tx = self.pool.begin().await.map_err(|e| files_tx_err(&e))?; lock_app_quota(&mut tx, cx.app_id, "files") .await .map_err(|e| files_tx_err(&e))?; + // Quota FIRST, before a single byte reaches the disk. Writing the blob + // and only then refusing would let an over-quota caller drive unbounded + // write+unlink churn — the ceiling would bound storage but not I/O. self.check_bytes(&mut tx, cx.app_id, None, size).await?; + // The bytes cannot join the transaction, so they go down inside it, + // under the lock: a rollback below unlinks them, and nothing else can + // reference them in the meantime. + let checksum = + files_repo::write_atomic_at(&self.root, &owner, collection, id, &new.data)?; let meta = files_repo::insert_meta_on( &mut *tx, cx.app_id, @@ -1502,22 +1520,22 @@ impl FilesWriter for PostgresFilesWriter { return Ok(false); }; let size = i64::try_from(upd.data.len()).unwrap_or(i64::MAX); - let checksum = files_repo::write_atomic_at(&self.root, &owner, collection, id, &upd.data)?; let mut tx = self.pool.begin().await.map_err(|e| files_tx_err(&e))?; lock_app_quota(&mut tx, cx.app_id, "files") .await .map_err(|e| files_tx_err(&e))?; // Growing an existing blob consumes disk too — an update that skipped the - // ceiling would be a free bypass (exactly the group-files bug). The - // projection subtracts the blob being replaced, so a same-size-or-smaller - // update near the cap still goes through. - if let Err(e) = self.check_bytes(&mut tx, cx.app_id, Some(id), size).await { - // The new bytes are already on disk but no row will reference them. - drop(tx); - files_repo::unlink_blob(&self.root, &owner, collection, id); - return Err(e); - } + // ceiling would be a free bypass. The projection subtracts the blob being + // replaced, so a same-size-or-smaller update near the cap still goes through. + // + // **This MUST precede the blob write.** `write_atomic_at` renames over the + // FINAL path, destroying the previous bytes. Refusing afterwards would + // leave the row pointing at a file that is either gone (if we unlink) or + // has the wrong checksum (if we don't) — i.e. a user would permanently + // lose a file merely by exceeding a quota. + self.check_bytes(&mut tx, cx.app_id, Some(id), size).await?; + let checksum = files_repo::write_atomic_at(&self.root, &owner, collection, id, &upd.data)?; let updated = files_repo::update_meta_on( &mut *tx, cx.app_id, @@ -1591,8 +1609,9 @@ impl FilesWriter for BestEffortFilesWriter { new: NewFile, ) -> Result { let meta = self.repo.create(cx.app_id, collection, new).await?; - if let Ok(event) = files_event("create", collection, &meta, None) { - best_effort_emit(&*self.events, cx, event).await; + match files_event("create", collection, &meta, None) { + Ok(event) => best_effort_emit(&*self.events, cx, event).await, + Err(e) => log_event_build_failure("create", &e), } Ok(meta.id) } @@ -1921,8 +1940,9 @@ impl GroupFilesWriter for BestEffortGroupFilesWriter { }); } let meta = self.repo.create(group_id, collection, new).await?; - if let Ok(event) = files_event("create", collection, &meta, None) { - self.emit(cx, group_id, event).await; + match files_event("create", collection, &meta, None) { + Ok(event) => self.emit(cx, group_id, event).await, + Err(e) => log_event_build_failure("create", &e), } Ok(meta.id) } @@ -1934,8 +1954,22 @@ impl GroupFilesWriter for BestEffortGroupFilesWriter { collection: &str, id: Uuid, upd: FileUpdate, - _quota: GroupFilesQuota, + quota: GroupFilesQuota, ) -> Result { + // Mirror the transactional writer's ceiling, so a unit test can catch its + // removal. (Not serialized — that is what the Postgres writer is for.) + if let Some(prev) = self.repo.head(group_id, collection, id).await? { + let used = self.repo.total_bytes(group_id).await?; + let projected = used + .saturating_sub(prev.size) + .saturating_add(upd.data.len() as u64); + if projected > quota.max_total_bytes { + return Err(GroupFilesError::QuotaExceeded { + limit: quota.max_total_bytes, + actual: projected, + }); + } + } let Some(FileUpdated { new, prev }) = self.repo.update(group_id, collection, id, upd).await? else { diff --git a/crates/manager-core/src/files_repo.rs b/crates/manager-core/src/files_repo.rs index 421b7b0..c8e669e 100644 --- a/crates/manager-core/src/files_repo.rs +++ b/crates/manager-core/src/files_repo.rs @@ -209,11 +209,6 @@ impl FsFilesRepo { } Ok(()) } - - fn final_path(&self, app_id: AppId, collection: &str, id: Uuid) -> PathBuf { - final_path_at(&self.config.root, &app_owner_dir(app_id), collection, id) - } - fn write_atomic( &self, app_id: AppId, @@ -468,42 +463,14 @@ impl FilesRepo for FsFilesRepo { id: Uuid, ) -> Result, FilesRepoError> { Self::guard_collection(collection)?; - // SELECT + DELETE in one tx; unlink afterwards (outside the tx). - let mut tx = self.pool.begin().await?; - let row: Option = sqlx::query_as( - "SELECT id, collection, name, content_type, size_bytes, \ - checksum_sha256, created_at, updated_at \ - FROM files WHERE app_id = $1 AND collection = $2 AND id = $3 \ - FOR UPDATE", - ) - .bind(app_id.into_inner()) - .bind(collection) - .bind(id) - .fetch_optional(&mut *tx) - .await?; - - let Some(row) = row else { - tx.rollback().await?; - return Ok(None); - }; - - sqlx::query("DELETE FROM files WHERE app_id = $1 AND collection = $2 AND id = $3") - .bind(app_id.into_inner()) - .bind(collection) - .bind(id) - .execute(&mut *tx) - .await?; - tx.commit().await?; - - // Row is gone; unlink the bytes. A failure here leaves an orphan - // file (reclaimed by a future sweep) — not fatal. - let path = self.final_path(app_id, collection, id); - if let Err(e) = std::fs::remove_file(&path) { - if e.kind() != std::io::ErrorKind::NotFound { - tracing::warn!(path = %path.display(), error = %e, "files: unlink after delete failed (orphan)"); - } + // `DELETE ... RETURNING` is one statement, so the old SELECT-FOR-UPDATE + // + DELETE pair is unnecessary. Unlink only AFTER the row is gone: the + // reverse order would destroy the bytes of a row that a failure keeps. + let meta = delete_meta_on(&self.pool, app_id, collection, id).await?; + if meta.is_some() { + unlink_blob(&self.config.root, &app_owner_dir(app_id), collection, id); } - Ok(Some(row.into_meta())) + Ok(meta) } async fn list( diff --git a/crates/manager-core/src/kv_repo.rs b/crates/manager-core/src/kv_repo.rs index ca16acd..dbfc49b 100644 --- a/crates/manager-core/src/kv_repo.rs +++ b/crates/manager-core/src/kv_repo.rs @@ -99,27 +99,6 @@ const KV_LIST_DEFAULT_LIMIT: u32 = 100; // trigger fan-out on ONE connection inside ONE transaction. // ---------------------------------------------------------------------------- -pub(crate) async fn get_on<'c, E>( - exec: E, - app_id: AppId, - collection: &str, - key: &str, -) -> Result, KvRepoError> -where - E: sqlx::PgExecutor<'c>, -{ - let row: Option<(serde_json::Value,)> = sqlx::query_as( - "SELECT value FROM kv_entries \ - WHERE app_id = $1 AND collection = $2 AND key = $3", - ) - .bind(app_id.into_inner()) - .bind(collection) - .bind(key) - .fetch_optional(exec) - .await?; - Ok(row.map(|(v,)| v)) -} - /// Upsert. Returns the previous value, so the caller knows whether this was an /// `insert` or an `update` for the emitted `ServiceEvent`. pub(crate) async fn set_on<'c, E>( diff --git a/crates/manager-core/tests/atomic_write.rs b/crates/manager-core/tests/atomic_write.rs index 0ff22a0..c98701b 100644 --- a/crates/manager-core/tests/atomic_write.rs +++ b/crates/manager-core/tests/atomic_write.rs @@ -20,6 +20,7 @@ use picloud_manager_core::atomic_write::{ GroupKvTarget, GroupKvWriter, KvWriter, PostgresFilesWriter, PostgresGroupDocsWriter, PostgresGroupFilesWriter, PostgresGroupKvWriter, PostgresKvWriter, }; +use picloud_manager_core::files_repo::{FilesConfig, FilesRepo, FsFilesRepo}; use picloud_manager_core::quota::GroupWriteQuota; use picloud_shared::{ AppId, ExecutionId, FileUpdate, FilesError, GroupDocsError, GroupFilesError, GroupId, @@ -734,6 +735,29 @@ async fn concurrent_uploads_cannot_push_an_app_past_its_disk_ceiling() { stored, "the refused update must not have changed the stored total" ); + // …and, crucially, the file must SURVIVE. `write_atomic_at` renames over the + // final path, so a quota check that ran AFTER the blob write would have + // destroyed the original bytes — the user would permanently lose a file + // merely by exceeding a quota. Read it back through the repo, which verifies + // the stored checksum, so both "bytes gone" and "bytes replaced but metadata + // still has the old checksum" fail here. + let repo = FsFilesRepo::new( + pool.clone(), + FilesConfig { + root: root.clone(), + max_file_size_bytes: 1024 * 1024, + }, + ); + let survivor = repo + .get(AppId::from(f.app), "assets", created[0]) + .await + .expect("the refused update must not corrupt or destroy the existing file") + .expect("the file must still exist"); + assert_eq!( + survivor, + vec![b'x'; 100], + "the original bytes must be intact — a refused update must not touch the blob" + ); cleanup(&pool, f.app).await; let _ = std::fs::remove_dir_all(&root);