diff --git a/crates/arr-db/migrations/0027_subtitle_budgets.sql b/crates/arr-db/migrations/0027_subtitle_budgets.sql new file mode 100644 index 0000000..dc058f0 --- /dev/null +++ b/crates/arr-db/migrations/0027_subtitle_budgets.sql @@ -0,0 +1,19 @@ +-- #197. Token bucket state for §15's provider and translator budgets. The +-- allowance itself is `subtitle_settings.provider_daily_budgets` / +-- `translator_daily_budgets` (#198); this table only tracks what has already +-- been spent today, so a bucket with no allowance configured reads as +-- unlimited rather than zero (0025's own comment on that column). +-- +-- One row per (kind, name, day); a new day is a fresh row rather than a +-- reset column, so the refill is "today has no row yet" and needs no cron. + +CREATE TABLE subtitle_budget_spend ( + kind TEXT NOT NULL CHECK (kind IN ('provider', 'translator')), + name TEXT NOT NULL, + -- UTC calendar day, `YYYY-MM-DD`. + day TEXT NOT NULL, + -- A provider unit is one download; a translator unit is one character of + -- source text sent (DESIGN.md §15: translators bill per character). + spent INTEGER NOT NULL DEFAULT 0 CHECK (spent >= 0), + PRIMARY KEY (kind, name, day) +) STRICT; diff --git a/crates/arr-db/src/lib.rs b/crates/arr-db/src/lib.rs index 4faa1eb..c530100 100644 --- a/crates/arr-db/src/lib.rs +++ b/crates/arr-db/src/lib.rs @@ -7,10 +7,12 @@ use std::path::Path; pub mod blacklist; pub mod policy; +pub mod subtitle_budget; pub mod subtitles; pub use blacklist::Blacklist; pub use policy::{MoviePolicy, PolicyColumns, PolicyError, TitlePolicy}; +pub use subtitle_budget::BudgetKind; pub use subtitles::{ NewSubtitleFile, PendingSubtitle, SubtitleAttempt, SubtitleFile, SubtitleOrigin, SubtitleState, SubtitleSync, diff --git a/crates/arr-db/src/subtitle_budget.rs b/crates/arr-db/src/subtitle_budget.rs new file mode 100644 index 0000000..ea55d57 --- /dev/null +++ b/crates/arr-db/src/subtitle_budget.rs @@ -0,0 +1,222 @@ +//! Token buckets for §15's provider and translator daily budgets (#197). +//! +//! One row per (kind, name, today). The allowance lives in +//! `subtitle_settings` (#198) and is passed in by the caller rather than +//! read here, because the reconcile loop already loads settings once per +//! tick; this module only tracks and gates spend against it. +//! +//! [`try_spend`] is the single gate, and it is the only one: check and +//! increment happen in one statement, so concurrent closes racing the same +//! bucket (§8 runs up to four at once) cannot both slip a spend past the cap +//! between a read and a write. [`spent_today`] is read-only, kept apart for +//! callers that only need to report what a bucket holds. + +use sqlx::SqlitePool; + +/// Which bucket a spend counts against. +#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::Type)] +#[sqlx(rename_all = "lowercase")] +pub enum BudgetKind { + /// One unit per download (DESIGN.md §15's example: a single-digit daily + /// download cap on `OpenSubtitles`' free tier). + Provider, + /// One unit per character of source text sent, not per call (DESIGN.md + /// §15: translators bill per character). + Translator, +} + +/// What has already been spent from `kind`/`name`'s bucket today. +/// +/// # Errors +/// +/// If the query fails. +pub async fn spent_today( + pool: &SqlitePool, + kind: BudgetKind, + name: &str, +) -> Result { + Ok(sqlx::query_scalar!( + r#"SELECT spent AS "spent!: i64" FROM subtitle_budget_spend + WHERE kind = ? AND name = ? AND day = strftime('%Y-%m-%d', 'now')"#, + kind, + name + ) + .fetch_optional(pool) + .await? + .unwrap_or(0)) +} + +/// Try to spend `amount` from `kind`/`name`'s bucket today, atomically. +/// +/// `None` allowance means unlimited (§10): always spends, and nothing is +/// recorded — there is no cap to account spend against. With an allowance, +/// the insert-or-increment and the cap check are one statement, so a spend +/// that would cross `allowance` neither applies nor is recorded. +/// +/// # Errors +/// +/// If the query fails. +pub async fn try_spend( + pool: &SqlitePool, + kind: BudgetKind, + name: &str, + amount: i64, + allowance: Option, +) -> Result { + let Some(allowance) = allowance else { + return Ok(true); + }; + let spent = sqlx::query_scalar!( + r#"INSERT INTO subtitle_budget_spend (kind, name, day, spent) + SELECT ?1, ?2, strftime('%Y-%m-%d', 'now'), ?3 + WHERE ?3 <= ?4 + ON CONFLICT (kind, name, day) DO UPDATE + SET spent = subtitle_budget_spend.spent + excluded.spent + WHERE subtitle_budget_spend.spent + excluded.spent <= ?4 + RETURNING spent AS "spent!: i64""#, + kind, + name, + amount, + allowance + ) + .fetch_optional(pool) + .await?; + Ok(spent.is_some()) +} + +#[cfg(test)] +#[allow(clippy::unwrap_used)] +mod tests { + use super::{spent_today, try_spend, BudgetKind}; + use crate::Db; + + async fn database() -> (Db, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let database = Db::connect(dir.path().join("test.db")).await.unwrap(); + database.migrate().await.unwrap(); + (database, dir) + } + + #[tokio::test] + async fn an_unconfigured_allowance_is_unlimited_and_untracked() { + let (db, _dir) = database().await; + + let spent = try_spend(db.pool(), BudgetKind::Provider, "opensubtitles", 1, None) + .await + .unwrap(); + + assert!(spent); + assert_eq!( + spent_today(db.pool(), BudgetKind::Provider, "opensubtitles") + .await + .unwrap(), + 0, + "nothing to cap means nothing to account" + ); + } + + #[tokio::test] + async fn spend_accumulates_until_the_allowance_is_reached() { + let (db, _dir) = database().await; + let allowance = Some(5); + + for _ in 0..5 { + assert!(try_spend( + db.pool(), + BudgetKind::Provider, + "opensubtitles", + 1, + allowance + ) + .await + .unwrap()); + } + + assert_eq!( + spent_today(db.pool(), BudgetKind::Provider, "opensubtitles") + .await + .unwrap(), + 5 + ); + + let refused = try_spend( + db.pool(), + BudgetKind::Provider, + "opensubtitles", + 1, + allowance, + ) + .await + .unwrap(); + assert!(!refused, "the sixth call is over the cap"); + assert_eq!( + spent_today(db.pool(), BudgetKind::Provider, "opensubtitles") + .await + .unwrap(), + 5, + "a refused spend is not recorded" + ); + } + + #[tokio::test] + async fn a_spend_that_would_cross_the_cap_is_refused_whole() { + let (db, _dir) = database().await; + let allowance = Some(100); + + // A translator bucket counts characters, not calls, so one spend can + // be far larger than what remains. + assert!( + try_spend(db.pool(), BudgetKind::Translator, "deepl", 90, allowance) + .await + .unwrap() + ); + let refused = try_spend(db.pool(), BudgetKind::Translator, "deepl", 20, allowance) + .await + .unwrap(); + assert!(!refused, "90 + 20 would cross 100"); + assert_eq!( + spent_today(db.pool(), BudgetKind::Translator, "deepl") + .await + .unwrap(), + 90, + "the refused spend left the bucket untouched" + ); + + assert!( + try_spend(db.pool(), BudgetKind::Translator, "deepl", 10, allowance) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn buckets_are_independent_per_kind_and_name() { + let (db, _dir) = database().await; + let allowance = Some(1); + + assert!(try_spend( + db.pool(), + BudgetKind::Provider, + "opensubtitles", + 1, + allowance + ) + .await + .unwrap()); + + // Same name, other kind: untouched. + assert_eq!( + spent_today(db.pool(), BudgetKind::Translator, "opensubtitles") + .await + .unwrap(), + 0 + ); + // Same kind, other name: untouched. + assert_eq!( + spent_today(db.pool(), BudgetKind::Provider, "podnapisi") + .await + .unwrap(), + 0 + ); + } +}