PiCloud/crates/manager-core/src/quota.rs

//! Storage ceilings — per-app and §11.6 per-group.
//!
//! # Why the two are enforced DIFFERENTLY
//!
//! The group ceilings serialize (a `pg_advisory_xact_lock` across the check and
//! the write) because they must be exact: a group's shared collections are the
//! cross-app blast radius, and with a small ceiling the check-then-write race
//! overshoots wildly (measured: 19 rows stored against a ceiling of 5). Group
//! writes are comparatively rare, so the lock costs little.
//!
//! The per-app KV/docs ceilings deliberately do **not** lock, and check only a
//! ROW count, on INSERT only:
//!
//!   * `kv::set` is the hottest write path in the system. Taking a per-app lock
//!     on every set would serialize an app's entire data plane; a `SUM(...)` scan
//!     on every set would make write cost grow with the app's stored size. Both
//!     are far worse trades than what the lock buys here.
//!   * What it buys is small: without the lock, overshoot is bounded by write
//!     concurrency (`PICLOUD_MAX_CONCURRENT_EXECUTIONS`, 32 by default) against a
//!     ceiling of 100k — roughly 0.03%. These are anti-DoS rails, not billing.
//!   * A byte ceiling is unnecessary: every value is already capped at
//!     `PICLOUD_KV_MAX_VALUE_BYTES`, so `max_rows × max_value_bytes` is ALREADY a
//!     finite bound on an app's stored bytes. Add a second scan for nothing.
//!   * An update adds no row and cannot exceed the per-value cap, so it is exempt
//!     and pays nothing at all — the COUNT runs only when a key/doc is created.
//!
//! Per-app FILES are the exception and DO lock: a single blob may be 100 MB, so
//! 32 racing uploads could overshoot by gigabytes of real disk, and a file write
//! is heavy enough that a lock and a `SUM(size_bytes)` are lost in the noise.
//!
//! §11.6 per-group quotas — global env-var ceilings on a group's shared
//! collections, enforced in the group write path (mirrors the per-value
//! `PICLOUD_*_MAX_*_BYTES` caps). One default applies to every group; per-group
//! configurable limits are deferred.
//!
//! - KV / docs: a per-group ROW-count ceiling AND a per-group TOTAL-BYTES ceiling
//!   (across the group's collections of that kind). Row count exempts an update
//!   of an existing key/doc (net-zero rows); the byte ceiling is checked on the
//!   projected total (old value's bytes subtracted, new value's added) so a
//!   same-or-smaller update near the cap is still allowed.
//! - files: a per-group TOTAL-BYTES ceiling (sum of stored blob sizes).

/// Default per-group shared-KV row ceiling. Override `PICLOUD_GROUP_KV_MAX_ROWS`.
pub const DEFAULT_GROUP_KV_MAX_ROWS: u64 = 100_000;
/// Default per-group shared-docs row ceiling. Override `PICLOUD_GROUP_DOCS_MAX_ROWS`.
pub const DEFAULT_GROUP_DOCS_MAX_ROWS: u64 = 100_000;
/// Default per-group shared-KV total-bytes ceiling (256 MiB). Override
/// `PICLOUD_GROUP_KV_MAX_TOTAL_BYTES`.
pub const DEFAULT_GROUP_KV_MAX_TOTAL_BYTES: u64 = 256 * 1024 * 1024;
/// Default per-group shared-docs total-bytes ceiling (256 MiB). Override
/// `PICLOUD_GROUP_DOCS_MAX_TOTAL_BYTES`.
pub const DEFAULT_GROUP_DOCS_MAX_TOTAL_BYTES: u64 = 256 * 1024 * 1024;
/// Default per-group shared-files total-bytes ceiling (10 GiB). Override
/// `PICLOUD_GROUP_FILES_MAX_TOTAL_BYTES`.
pub const DEFAULT_GROUP_FILES_MAX_TOTAL_BYTES: u64 = 10 * 1024 * 1024 * 1024;

/// Default per-app KV key ceiling. Override `PICLOUD_APP_KV_MAX_ROWS`.
pub const DEFAULT_APP_KV_MAX_ROWS: u64 = 100_000;
/// Default per-app doc ceiling. Override `PICLOUD_APP_DOCS_MAX_ROWS`.
pub const DEFAULT_APP_DOCS_MAX_ROWS: u64 = 100_000;
/// Default per-app stored-blob ceiling (10 GiB). Override
/// `PICLOUD_APP_FILES_MAX_TOTAL_BYTES`.
pub const DEFAULT_APP_FILES_MAX_TOTAL_BYTES: u64 = 10 * 1024 * 1024 * 1024;

fn u64_from_env(var: &str, default: u64) -> u64 {
    if let Ok(v) = std::env::var(var) {
        match v.trim().parse::<u64>() {
            Ok(n) if n > 0 => return n,
            _ => tracing::warn!(value = %v, "ignoring invalid {var} (want a positive integer)"),
        }
    }
    default
}

#[must_use]
pub fn app_kv_max_rows_from_env() -> u64 {
    u64_from_env("PICLOUD_APP_KV_MAX_ROWS", DEFAULT_APP_KV_MAX_ROWS)
}

#[must_use]
pub fn app_docs_max_rows_from_env() -> u64 {
    u64_from_env("PICLOUD_APP_DOCS_MAX_ROWS", DEFAULT_APP_DOCS_MAX_ROWS)
}

#[must_use]
pub fn app_files_max_total_bytes_from_env() -> u64 {
    u64_from_env(
        "PICLOUD_APP_FILES_MAX_TOTAL_BYTES",
        DEFAULT_APP_FILES_MAX_TOTAL_BYTES,
    )
}

#[must_use]
pub fn group_kv_max_rows_from_env() -> u64 {
    u64_from_env("PICLOUD_GROUP_KV_MAX_ROWS", DEFAULT_GROUP_KV_MAX_ROWS)
}

#[must_use]
pub fn group_docs_max_rows_from_env() -> u64 {
    u64_from_env("PICLOUD_GROUP_DOCS_MAX_ROWS", DEFAULT_GROUP_DOCS_MAX_ROWS)
}

#[must_use]
pub fn group_kv_max_total_bytes_from_env() -> u64 {
    u64_from_env(
        "PICLOUD_GROUP_KV_MAX_TOTAL_BYTES",
        DEFAULT_GROUP_KV_MAX_TOTAL_BYTES,
    )
}

#[must_use]
pub fn group_docs_max_total_bytes_from_env() -> u64 {
    u64_from_env(
        "PICLOUD_GROUP_DOCS_MAX_TOTAL_BYTES",
        DEFAULT_GROUP_DOCS_MAX_TOTAL_BYTES,
    )
}

#[must_use]
pub fn group_files_max_total_bytes_from_env() -> u64 {
    u64_from_env(
        "PICLOUD_GROUP_FILES_MAX_TOTAL_BYTES",
        DEFAULT_GROUP_FILES_MAX_TOTAL_BYTES,
    )
}

// ----------------------------------------------------------------------------
// The row/byte decision, shared by every backend
// ----------------------------------------------------------------------------

/// The ceilings one shared-collection write must fit under.
#[derive(Debug, Clone, Copy)]
pub struct GroupWriteQuota {
    pub max_rows: u64,
    pub max_total_bytes: u64,
    /// The per-value cap. Not a ceiling of its own here — it is what makes the
    /// fast path below sound.
    pub max_value_bytes: usize,
}

/// The two usage reads a quota decision needs. Implemented over the writing
/// transaction's connection (the real path — see `crate::atomic_write`) and over
/// the in-memory test store, so the POLICY below has one home regardless of
/// backend.
///
/// **The reads must come from the same transaction as the write they gate.**
/// Reading usage on one connection and writing on another is the check-then-write
/// race that lets N concurrent writers each observe the pre-write total, each
/// pass, and together blow the ceiling.
#[async_trait::async_trait]
pub trait GroupUsage: Send {
    /// Current row count for the group, across its collections of this kind.
    async fn count_rows(&mut self) -> Result<u64, String>;

    /// Total stored bytes for the group AFTER this write — the current total,
    /// minus the bytes of the row being replaced (if any), plus the new value's.
    async fn projected_total_bytes(&mut self) -> Result<u64, String>;
}

/// What the ceilings say about a pending write.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuotaOutcome {
    Allowed,
    RowsExceeded { limit: u64, actual: u64 },
    BytesExceeded { limit: u64, actual: u64 },
}

/// Decide whether a shared-collection write fits. `adds_row` = the write will
/// ADD a row rather than replace one (a new key, or a doc create) — a
/// compare-and-swap against an absent key with a `Some` precondition adds
/// nothing, because it cannot swap at all.
///
/// # Errors
/// Propagates a backend failure from either usage read, as a string.
pub async fn check_group_write(
    q: GroupWriteQuota,
    adds_row: bool,
    usage: &mut dyn GroupUsage,
) -> Result<QuotaOutcome, String> {
    // One cheap COUNT backs BOTH the row ceiling and the byte fast path below.
    let count = usage.count_rows().await?;

    // Row ceiling: only a write that adds a row consumes one. An update is
    // net-zero and exempt.
    if adds_row && count >= q.max_rows {
        return Ok(QuotaOutcome::RowsExceeded {
            limit: q.max_rows,
            actual: count,
        });
    }

    // Byte ceiling, on the PROJECTED total — so an update that is the same size
    // or smaller still goes through when the group is sitting near its cap.
    //
    // Fast path: every stored value is capped at `max_value_bytes`, so the group
    // can hold at most `rows_after * max_value_bytes`. When even that upper bound
    // fits under the ceiling, no write can breach it — skip the
    // O(collection-size) `SUM(octet_length(...))` scan entirely. The scan runs
    // only for a group actually near its byte cap.
    let rows_after = if adds_row { count + 1 } else { count };
    let upper_bound = rows_after.saturating_mul(q.max_value_bytes as u64);
    if upper_bound > q.max_total_bytes {
        let projected = usage.projected_total_bytes().await?;
        if projected > q.max_total_bytes {
            return Ok(QuotaOutcome::BytesExceeded {
                limit: q.max_total_bytes,
                actual: projected,
            });
        }
    }
    Ok(QuotaOutcome::Allowed)
}