yolo/arr
3
0
Fork 0
Code Issues 12 Pull Requests Actions Packages Projects Releases Wiki Activity

feat(arr): token bucket table for subtitle budgets

Browse Source 
One row per provider/translator per day; try_spend is a single
atomic upsert so concurrent reconcile closes can't both slip a spend
past the daily cap. No allowance configured reads as unlimited.
This commit is contained in:
Miguel Palhas
2026-08-25 02:40:06 +01:00
parent f49507b7d1
commit 4d84625332
3 changed files with 243 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/arr-db/migrations/0027_subtitle_budgets.sql
+19
View File
@@ -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;
crates/arr-db/src/lib.rs
+2
View File
@@ -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,
crates/arr-db/src/subtitle_budget.rs
+222
View File
@@ -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<i64, sqlx::Error> {
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<i64>,
) -> Result<bool, sqlx::Error> {
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
);
}
}