use std::collections::HashSet;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use redis::aio::ConnectionManager;
use redis::{AsyncCommands, Script, ScriptInvocation};
use super::{
Admission, BudgetError, BudgetKey, BudgetStore, Denial, ExceededScope, Reservation,
SharedSettings, Uncertain,
};
use crate::backends::health::BackendHealth;
use crate::policy::PolicyHold;
use crate::redis_support::RedisHealth;
use crate::telemetry::metrics;
const BACKEND: &str = "redis";
const POLICY_STORE: &str = crate::policy::ungoverned::BUDGET_REDIS;
const LAYOUT_V2: &str = "v2";
const LAYOUT_MIGRATING: &str = "v2-migrating";
const RESERVE: &str = r#"
local now = tonumber(ARGV[1])
local ttl_ms = tonumber(ARGV[2])
local limit = tonumber(ARGV[3])
local amount = tonumber(ARGV[4])
local id = ARGV[5]
local held = 0
local reservations = redis.call('HGETALL', KEYS[2])
for i = 1, #reservations, 2 do
local separator = string.find(reservations[i + 1], ':')
local value = tonumber(string.sub(reservations[i + 1], 1, separator - 1))
local expires_at = tonumber(string.sub(reservations[i + 1], separator + 1))
if expires_at <= now then
redis.call('HDEL', KEYS[2], reservations[i])
else
held = held + value
end
end
local spent = tonumber(redis.call('GET', KEYS[1]) or '0')
if spent + held + amount > limit then
return 0
end
redis.call('HSET', KEYS[2], id, amount .. ':' .. (now + ttl_ms))
redis.call('PEXPIRE', KEYS[2], ttl_ms * 2)
return 1
"#;
const SETTLE: &str = r#"
redis.call('HDEL', KEYS[2], ARGV[1])
local actual = tonumber(ARGV[2])
if actual > 0 then
redis.call('INCRBY', KEYS[1], actual)
end
return 1
"#;
const RESERVE_V2: &str = r#"
local now = tonumber(ARGV[1])
local ttl_ms = tonumber(ARGV[2])
local subject_limit = tonumber(ARGV[3])
local namespace_limit = tonumber(ARGV[4])
local amount = tonumber(ARGV[5])
local id = ARGV[6]
local function held(hash)
local total = 0
local reservations = redis.call('HGETALL', hash)
for i = 1, #reservations, 2 do
local separator = string.find(reservations[i + 1], ':')
local value = tonumber(string.sub(reservations[i + 1], 1, separator - 1))
local expires_at = tonumber(string.sub(reservations[i + 1], separator + 1))
if expires_at <= now then
redis.call('HDEL', hash, reservations[i])
else
total = total + value
end
end
return total
end
local subject_held = held(KEYS[2])
local namespace_held = held(KEYS[4])
local subject_spent = tonumber(redis.call('GET', KEYS[1]) or '0')
local namespace_spent = tonumber(redis.call('GET', KEYS[3]) or '0')
if subject_spent + subject_held + amount > subject_limit then
return 0
end
if namespace_spent + namespace_held + amount > namespace_limit then
return 2
end
local hold = amount .. ':' .. (now + ttl_ms)
redis.call('HSET', KEYS[2], id, hold)
redis.call('PEXPIRE', KEYS[2], ttl_ms * 2)
redis.call('HSET', KEYS[4], id, hold)
redis.call('PEXPIRE', KEYS[4], ttl_ms * 2)
return 1
"#;
const SETTLE_V2: &str = r#"
redis.call('HDEL', KEYS[2], ARGV[1])
redis.call('HDEL', KEYS[4], ARGV[1])
local actual = tonumber(ARGV[2])
if actual > 0 then
redis.call('INCRBY', KEYS[1], actual)
redis.call('INCRBY', KEYS[3], actual)
end
return 1
"#;
const DRAIN_V1: &str = r#"
local pending = redis.call('GET', KEYS[2])
if pending then
return pending
end
local spent = redis.call('GET', KEYS[1])
if not spent then
return false
end
local seq = redis.call('INCR', KEYS[3])
redis.call('DEL', KEYS[1])
local claim = seq .. ':' .. spent
redis.call('SET', KEYS[2], claim)
return claim
"#;
const APPLY_CLAIM: &str = r#"
local amount = tonumber(ARGV[2])
if redis.call('HSETNX', KEYS[3], ARGV[1], amount) == 0 then
return 0
end
if amount > 0 then
redis.call('INCRBY', KEYS[1], amount)
redis.call('INCRBY', KEYS[2], amount)
end
return amount
"#;
pub struct RedisBudget {
settings: SharedSettings,
key_prefix: String,
connection: ConnectionManager,
reserve: Script,
settle: Script,
health: Arc<RedisHealth>,
}
const PROBE_BOUND: Duration = Duration::from_secs(5);
impl RedisBudget {
pub async fn connect(
url: &str,
key_prefix: String,
settings: SharedSettings,
) -> Result<Self, BudgetError> {
let client = ::redis::Client::open(url)
.map_err(|e| BudgetError::invalid(BACKEND, format!("unusable URL: {e}")))?;
let mut connection = ConnectionManager::new(client).await?;
::redis::cmd("PING")
.query_async::<String>(&mut connection)
.await?;
let (reserve, settle) = if settings.enforces_namespace_cap() {
require_migrated_layout(&mut connection, &key_prefix).await?;
(Script::new(RESERVE_V2), Script::new(SETTLE_V2))
} else {
require_unmigrated_layout(&mut connection, &key_prefix).await?;
(Script::new(RESERVE), Script::new(SETTLE))
};
let health = {
let probed = connection.clone();
Arc::new(RedisHealth::new(PROBE_BOUND, move || probed.clone()))
};
Ok(Self {
settings,
key_prefix,
connection,
reserve,
settle,
health,
})
}
fn keys(&self, key: &BudgetKey) -> (String, String) {
let scope = format!("{}:{{{}|{}}}", self.key_prefix, key.namespace, key.subject);
(format!("{scope}:spent"), format!("{scope}:reservations"))
}
fn script<'a>(&self, script: &'a Script, key: &BudgetKey) -> ScriptInvocation<'a> {
let mut invocation = script.prepare_invoke();
if self.settings.enforces_namespace_cap() {
let scopes = v2_keys(&self.key_prefix, key);
invocation
.key(scopes.subject_spent)
.key(scopes.subject_reservations)
.key(scopes.namespace_spent)
.key(scopes.namespace_reservations);
} else {
let (spent, reservations) = self.keys(key);
invocation.key(spent).key(reservations);
}
invocation
}
}
struct V2Keys {
subject_spent: String,
subject_reservations: String,
namespace_spent: String,
namespace_reservations: String,
}
fn v2_keys(key_prefix: &str, key: &BudgetKey) -> V2Keys {
let namespace = namespace_scope(key_prefix, &key.namespace);
let subject = format!("{namespace}:subject:{}", escaped(&key.subject));
V2Keys {
subject_spent: format!("{subject}:spent"),
subject_reservations: format!("{subject}:reservations"),
namespace_spent: format!("{namespace}:namespace:spent"),
namespace_reservations: format!("{namespace}:namespace:reservations"),
}
}
fn namespace_scope(key_prefix: &str, namespace: &str) -> String {
format!("{key_prefix}:v2:{{{}}}", escaped(namespace))
}
fn escaped(part: &str) -> String {
let mut out = String::with_capacity(part.len());
for character in part.chars() {
match character {
'%' => out.push_str("%25"),
'{' => out.push_str("%7B"),
'}' => out.push_str("%7D"),
other => out.push(other),
}
}
out
}
fn layout_key(key_prefix: &str) -> String {
format!("{key_prefix}:layout")
}
fn legacy_patterns(key_prefix: &str) -> [String; 2] {
[
format!("{key_prefix}:{{*}}:spent"),
format!("{key_prefix}:{{*}}:reservations"),
]
}
async fn require_migrated_layout(
connection: &mut ConnectionManager,
key_prefix: &str,
) -> Result<(), BudgetError> {
let marker: Option<String> = connection.get(layout_key(key_prefix)).await?;
if marker.as_deref() == Some(LAYOUT_MIGRATING) {
return Err(unfinished_migration(key_prefix));
}
if marker.as_deref() != Some(LAYOUT_V2) {
return Err(BudgetError::invalid(
BACKEND,
format!(
"`namespace_limit_microdollars` needs the v2 key layout, but `{}` is not marked \
migrated. Stop every replica and run `axond budget migrate-redis`, which carries \
existing spend forward rather than restarting it from zero.",
layout_key(key_prefix)
),
));
}
require_no_pending_claims(connection, key_prefix).await?;
let legacy = count_legacy_keys(connection, key_prefix).await?;
if legacy > 0 {
return Err(BudgetError::invalid(
BACKEND,
format!(
"{legacy} v1 budget key(s) exist under `{key_prefix}` after the migration to the \
v2 layout: a gateway binary without namespace-cap support is still writing them, \
and the two layouts would each enforce half the traffic. Stop the v1 replicas and \
re-run `axond budget migrate-redis`."
),
));
}
Ok(())
}
async fn require_unmigrated_layout(
connection: &mut ConnectionManager,
key_prefix: &str,
) -> Result<(), BudgetError> {
let marker: Option<String> = connection.get(layout_key(key_prefix)).await?;
if marker.as_deref() == Some(LAYOUT_V2) {
return Err(BudgetError::invalid(
BACKEND,
format!(
"`{}` is marked migrated to the v2 key layout, so `namespace_limit_microdollars` \
must stay set: the v1 keys this configuration would use no longer hold the \
accumulated spend.",
layout_key(key_prefix)
),
));
}
if marker.as_deref() == Some(LAYOUT_MIGRATING) {
return Err(unfinished_migration(key_prefix));
}
Ok(())
}
fn unfinished_migration(key_prefix: &str) -> BudgetError {
BudgetError::invalid(
BACKEND,
format!(
"`{}` is marked `{LAYOUT_MIGRATING}`: a `axond budget migrate-redis` run did not \
finish, so some subjects hold their spend in the v2 layout and the rest in v1 and \
neither configuration can enforce a whole ledger. Re-run `axond budget migrate-redis`, \
which resumes where it stopped and carries each subject exactly once.",
layout_key(key_prefix)
),
)
}
async fn require_no_pending_claims(
connection: &mut ConnectionManager,
key_prefix: &str,
) -> Result<(), BudgetError> {
let pending = tally(connection, &format!("{key_prefix}:{{*}}{PENDING}")).await?;
let Some(example) = pending.example else {
return Ok(());
};
Err(BudgetError::invalid(
BACKEND,
format!(
"{} interrupted migration claim(s) are outstanding under `{key_prefix}` (for example \
`{example}`): a `axond budget migrate-redis` run was cut short after taking spend off \
the v1 key and before adding it to the v2 counters, so that spend is in neither. \
Re-run `axond budget migrate-redis`, which finishes them exactly once.",
pending.count
),
))
}
async fn count_legacy_keys(
connection: &mut ConnectionManager,
key_prefix: &str,
) -> Result<usize, BudgetError> {
let mut total = 0;
for pattern in legacy_patterns(key_prefix) {
total += tally(connection, &pattern).await?.count;
}
Ok(total)
}
struct Tally {
count: usize,
example: Option<String>,
}
async fn tally(connection: &mut ConnectionManager, pattern: &str) -> Result<Tally, BudgetError> {
let mut tally = Tally {
count: 0,
example: None,
};
scan(connection, pattern, |key| {
tally.count += 1;
tally.example.get_or_insert(key);
})
.await?;
Ok(tally)
}
async fn scan(
connection: &mut ConnectionManager,
pattern: &str,
mut each: impl FnMut(String),
) -> Result<(), BudgetError> {
let mut cursor = 0u64;
loop {
let (next, keys): (u64, Vec<String>) = ::redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(500)
.query_async(connection)
.await?;
for key in keys {
each(key);
}
cursor = next;
if cursor == 0 {
return Ok(());
}
}
}
async fn scanned(
connection: &mut ConnectionManager,
pattern: &str,
) -> Result<Vec<String>, BudgetError> {
let mut found = Vec::new();
scan(connection, pattern, |key| found.push(key)).await?;
Ok(found)
}
#[derive(Debug, Default, PartialEq, Eq)]
pub struct MigrationReport {
pub subjects: usize,
pub reservation_hashes: usize,
pub namespaces: usize,
pub carried_microdollars: u64,
}
pub async fn migrate_v1_to_v2(
url: &str,
key_prefix: &str,
namespaces: &[String],
) -> Result<MigrationReport, BudgetError> {
let client = ::redis::Client::open(url)
.map_err(|e| BudgetError::invalid(BACKEND, format!("unusable URL: {e}")))?;
let mut connection = ConnectionManager::new(client).await?;
let drain = Script::new(DRAIN_V1);
let apply = Script::new(APPLY_CLAIM);
let mut report = MigrationReport::default();
let spend_keys = scanned(&mut connection, &legacy_patterns(key_prefix)[0]).await?;
let reservation_keys = scanned(&mut connection, &legacy_patterns(key_prefix)[1]).await?;
let claimed_keys = scanned(&mut connection, &format!("{key_prefix}:{{*}}{PENDING}")).await?;
let mut scopes = Vec::with_capacity(spend_keys.len() + claimed_keys.len());
for legacy in &spend_keys {
scopes.push(resolve_legacy_scope(
key_prefix, legacy, ":spent", namespaces,
)?);
}
for claimed in &claimed_keys {
scopes.push(resolve_legacy_scope(
key_prefix, claimed, PENDING, namespaces,
)?);
}
for legacy in &reservation_keys {
resolve_legacy_scope(key_prefix, legacy, ":reservations", namespaces)?;
}
if !scopes.is_empty() || !reservation_keys.is_empty() {
let _: () = connection
.set(layout_key(key_prefix), LAYOUT_MIGRATING)
.await?;
}
let mut seen = HashSet::new();
let mut touched = HashSet::new();
for key in scopes {
if !seen.insert((key.namespace.clone(), key.subject.clone())) {
continue;
}
touched.insert(key.namespace.clone());
let carried = carry_forward(&mut connection, &drain, &apply, key_prefix, &key).await?;
report.subjects += 1;
report.carried_microdollars = report.carried_microdollars.saturating_add(carried);
}
for legacy in &reservation_keys {
let _: i64 = connection.del(legacy).await?;
report.reservation_hashes += 1;
}
report.namespaces = touched.len();
for namespace in &touched {
let _: i64 = connection
.del(format!(
"{}:migration:applied",
namespace_scope(key_prefix, namespace)
))
.await?;
}
for key in &seen {
let _: i64 = connection
.del(format!("{key_prefix}:{{{}|{}}}{SEQ}", key.0, key.1))
.await?;
}
let _: () = connection.set(layout_key(key_prefix), LAYOUT_V2).await?;
Ok(report)
}
const PENDING: &str = ":migration_pending";
const SEQ: &str = ":migration_seq";
async fn carry_forward(
connection: &mut ConnectionManager,
drain: &Script,
apply: &Script,
key_prefix: &str,
key: &BudgetKey,
) -> Result<u64, BudgetError> {
let scope = format!("{key_prefix}:{{{}|{}}}", key.namespace, key.subject);
let pending = format!("{scope}{PENDING}");
let v2 = v2_keys(key_prefix, key);
let applied = format!(
"{}:migration:applied",
namespace_scope(key_prefix, &key.namespace)
);
let mut carried: u64 = 0;
loop {
let claim: Option<String> = drain
.prepare_invoke()
.key(format!("{scope}:spent"))
.key(&pending)
.key(format!("{scope}{SEQ}"))
.invoke_async(connection)
.await?;
let Some(claim) = claim else {
return Ok(carried);
};
let (sequence, amount) = claim.split_once(':').ok_or_else(|| {
BudgetError::invalid(
BACKEND,
format!("the migration claim in `{pending}` is malformed: `{claim}`"),
)
})?;
let amount: i64 = amount.parse().unwrap_or_default();
let added: i64 = apply
.prepare_invoke()
.key(&v2.subject_spent)
.key(&v2.namespace_spent)
.key(&applied)
.arg(format!("{}#{sequence}", escaped(&key.subject)))
.arg(amount.max(0))
.invoke_async(connection)
.await?;
let _: i64 = connection.del(&pending).await?;
carried = carried.saturating_add(added.max(0) as u64);
}
}
fn resolve_legacy_scope(
key_prefix: &str,
key: &str,
suffix: &str,
namespaces: &[String],
) -> Result<BudgetKey, BudgetError> {
let ambiguous = |detail: &str| {
BudgetError::invalid(
BACKEND,
format!(
"the v1 budget key `{key}` cannot be attributed: {detail}. The v1 \
`{{namespace|subject}}` tag is unescaped, so it is resolved against the \
configured `[[namespace]]` ids rather than split; nothing has been migrated or \
deleted. Configure the namespace this key belongs to and re-run the migration \
— or, if the tag cannot be read one way whatever is configured, carry this key \
into the v2 layout by hand (or delete it, accepting the reset of that ledger)."
),
)
};
let opening = format!("{key_prefix}:{{");
let closing = format!("}}{suffix}");
let tag = key
.strip_prefix(&opening)
.and_then(|rest| rest.strip_suffix(&closing))
.ok_or_else(|| ambiguous("it is not a `<prefix>:{namespace|subject}` key"))?;
let mut matched = namespaces
.iter()
.filter_map(|namespace| {
tag.strip_prefix(namespace.as_str())
.and_then(|rest| rest.strip_prefix('|'))
.map(|subject| BudgetKey {
namespace: namespace.clone(),
subject: subject.to_owned(),
})
})
.collect::<Vec<_>>();
matched.sort_by(|a, b| (&a.namespace, &a.subject).cmp(&(&b.namespace, &b.subject)));
matched.dedup();
match matched.len() {
1 if !matched[0].subject.contains('|') => Ok(matched.remove(0)),
1 => Err(ambiguous(&format!(
"`{}` is a configured namespace id, but the tag's remainder (`{}`) holds a `|` too, so \
the tag reads equally as a longer namespace this config does not declare",
matched[0].namespace, matched[0].subject
))),
0 => Err(ambiguous(
"no configured namespace id is a prefix of its tag",
)),
found => Err(ambiguous(&format!(
"{found} configured namespace ids are prefixes of its tag, so the split between \
namespace and subject is ambiguous"
))),
}
}
#[cfg(test)]
fn hash_tag(key: &str) -> Option<&str> {
let start = key.find('{')? + 1;
let end = key[start..].find('}')? + start;
Some(&key[start..end])
}
#[async_trait]
impl BudgetStore for RedisBudget {
fn name(&self) -> &'static str {
"redis"
}
fn health(&self) -> Option<Arc<dyn BackendHealth>> {
Some(Arc::clone(&self.health) as Arc<dyn BackendHealth>)
}
async fn reserve(&self, key: &BudgetKey, estimated_microdollars: u64) -> Admission {
let Some(governing) = self.settings.caps(POLICY_STORE, &key.namespace) else {
return Admission::Denied(Denial::StoreUnavailable);
};
let caps = governing.caps;
let reservation = Reservation {
id: Reservation::next_id(),
estimate_microdollars: estimated_microdollars,
generation: governing.generation,
};
let hold = PolicyHold::take(&self.settings.ceilings, reservation.generation);
let ttl_ms = caps.reservation_ttl.as_millis() as u64;
let mut invocation = self.script(&self.reserve, key);
invocation.arg(now_ms()).arg(ttl_ms);
invocation.arg(caps.subject_microdollars);
if let Some(namespace_limit) = caps.namespace_microdollars {
invocation.arg(namespace_limit);
}
let admitted: Result<i64, ::redis::RedisError> = invocation
.arg(estimated_microdollars)
.arg(&reservation.id)
.invoke_async(&mut self.connection.clone())
.await;
match admitted {
Ok(1) => {
hold.kept();
Admission::Allowed(reservation)
}
Ok(2) => exceeded(key, ExceededScope::Namespace),
Ok(_) => exceeded(key, ExceededScope::Subject),
Err(e) => self.settings.unavailable.admission(
BACKEND,
&e,
Some(Uncertain {
hold,
reservation_ttl: caps.reservation_ttl,
}),
),
}
}
async fn settle(&self, key: &BudgetKey, reservation: &Reservation, actual_microdollars: u64) {
if reservation.id.is_empty() {
return;
}
self.settings.ceilings.exit(reservation.generation);
let settled: Result<i64, ::redis::RedisError> = self
.script(&self.settle, key)
.arg(&reservation.id)
.arg(actual_microdollars)
.invoke_async(&mut self.connection.clone())
.await;
if let Err(e) = settled {
tracing::error!(
error = %e,
namespace = %key.namespace,
actual_microdollars,
"budget settlement was lost; the reservation expires on its own deadline"
);
}
}
}
fn exceeded(key: &BudgetKey, scope: ExceededScope) -> Admission {
if scope == ExceededScope::Namespace {
metrics::record_budget_namespace_denial();
tracing::info!(
namespace = %key.namespace,
"namespace spend cap is exhausted; denying"
);
}
Admission::Denied(Denial::Exceeded)
}
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or(Duration::ZERO)
.as_millis() as u64
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use tokio::sync::Barrier;
use super::super::UnavailablePolicy;
use super::super::tests::key;
use super::*;
use crate::policy::BudgetCaps;
fn settings(limit: u64) -> SharedSettings {
expiring_settings(limit, None, Duration::from_secs(300))
}
fn namespace_settings(limit: u64, namespace_limit: u64) -> SharedSettings {
expiring_settings(limit, Some(namespace_limit), Duration::from_secs(300))
}
fn expiring_settings(
limit: u64,
namespace_limit: Option<u64>,
reservation_ttl: Duration,
) -> SharedSettings {
SharedSettings::fixed(
BudgetCaps {
subject_microdollars: limit,
namespace_microdollars: namespace_limit,
reservation_ttl,
},
UnavailablePolicy::Deny,
)
}
fn prefix() -> String {
format!("axond:test:{}", Reservation::next_id())
}
fn namespaces() -> Vec<String> {
vec!["acme".to_owned()]
}
fn tag(key: &str) -> &str {
hash_tag(key).expect("hash tag")
}
#[test]
fn a_budgets_keys_share_one_hash_slot() {
let scope = format!("axond:budget:{{{}|{}}}", "acme", "subject");
let spent = format!("{scope}:spent");
let reservations = format!("{scope}:reservations");
assert_eq!(tag(&spent), tag(&reservations));
assert_eq!(tag(&spent), "acme|subject");
}
#[test]
fn every_v2_key_in_a_namespace_shares_the_namespace_hash_tag() {
let keys = v2_keys(
"axond:budget",
&BudgetKey {
namespace: "acme".into(),
subject: "subject-a".into(),
},
);
let other = v2_keys(
"axond:budget",
&BudgetKey {
namespace: "acme".into(),
subject: "subject-b".into(),
},
);
for k in [
&keys.subject_spent,
&keys.subject_reservations,
&keys.namespace_spent,
&keys.namespace_reservations,
&other.subject_spent,
] {
assert_eq!(tag(k), "acme", "{k}");
}
assert_ne!(keys.subject_spent, other.subject_spent);
assert_eq!(keys.namespace_spent, other.namespace_spent);
}
#[test]
fn braces_in_identifiers_cannot_move_the_hash_tag() {
let keys = v2_keys(
"axond:budget",
&BudgetKey {
namespace: "ac}me".into(),
subject: "sub{ject}".into(),
},
);
assert_eq!(tag(&keys.subject_spent), "ac%7Dme");
assert_eq!(tag(&keys.subject_spent), tag(&keys.namespace_spent));
assert_ne!(escaped("a%7B"), escaped("a{"));
}
#[test]
fn a_v2_key_is_not_mistaken_for_legacy_state() {
let keys = v2_keys("axond:budget", &key());
let [spent_pattern, _] = legacy_patterns("axond:budget");
assert!(spent_pattern.starts_with("axond:budget:{"));
assert!(keys.subject_spent.starts_with("axond:budget:v2:{"));
}
#[test]
fn a_legacy_key_is_attributed_to_a_configured_namespace() {
let resolve = |key: &str, namespaces: &[&str]| {
let namespaces: Vec<String> = namespaces.iter().map(|n| (*n).to_owned()).collect();
resolve_legacy_scope("axond:budget", key, ":spent", &namespaces)
};
let unattributable = resolve("axond:budget:{acme|sub|ject}:spent", &["acme"])
.expect_err("a `|` left in the subject is a second reading of the tag");
assert!(
format!("{unattributable}").contains("holds a `|` too"),
"{unattributable}"
);
let parsed =
resolve("axond:budget:{team|west|sub}:spent", &["team|west"]).expect("resolved");
assert_eq!(parsed.namespace, "team|west");
assert_eq!(parsed.subject, "sub");
let parsed = resolve("axond:budget:{ac}me|sub{ject}:spent", &["ac}me"]).expect("resolved");
assert_eq!(parsed.namespace, "ac}me");
assert_eq!(parsed.subject, "sub{ject");
let ambiguous = resolve("axond:budget:{team|west|sub}:spent", &["team", "team|west"])
.expect_err("two candidate namespaces are ambiguous");
assert!(format!("{ambiguous}").contains("ambiguous"), "{ambiguous}");
let parsed = resolve("axond:budget:{acme|sub}:spent", &["acme", "acme"])
.expect("a duplicate namespace id resolves to the one key it describes");
assert_eq!(parsed.namespace, "acme");
assert_eq!(parsed.subject, "sub");
for (key, namespaces) in [
("axond:budget:{gone|sub}:spent", &["acme"][..]),
("axond:budget:layout", &["acme"][..]),
("axond:budget:{acme}:spent", &["acme"][..]),
] {
resolve(key, namespaces).expect_err(key);
}
}
#[test]
fn the_reserve_script_reclaims_expired_holds_before_deciding() {
assert!(RESERVE.contains("HDEL"));
assert!(RESERVE.contains("expires_at <= now"));
assert!(RESERVE.contains("spent + held + amount > limit"));
}
#[test]
fn the_composite_script_decides_both_caps_before_it_holds_either() {
let decisions = RESERVE_V2.find("subject_limit then").expect("subject cap");
let namespace = RESERVE_V2
.find("namespace_limit then")
.expect("namespace cap");
let first_write = RESERVE_V2.find("HSET").expect("the hold");
assert!(decisions < first_write);
assert!(namespace < first_write);
assert_eq!(SETTLE_V2.matches("INCRBY").count(), 2);
assert_eq!(SETTLE_V2.matches("HDEL").count(), 2);
}
#[tokio::test]
async fn two_stores_sharing_one_redis_enforce_a_single_cap() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let replica_a = RedisBudget::connect(&url, prefix.clone(), settings(1_000))
.await
.expect("connect");
let replica_b = RedisBudget::connect(&url, prefix, settings(1_000))
.await
.expect("connect");
let k = key();
let held = replica_a.reserve(&k, 700).await;
assert_eq!(
replica_b.reserve(&k, 700).await,
Admission::Denied(Denial::Exceeded)
);
let Admission::Allowed(reservation) = held else {
panic!("the first reservation must be admitted");
};
replica_a.settle(&k, &reservation, 100).await;
let second = replica_b.reserve(&k, 700).await;
assert!(matches!(second, Admission::Allowed(_)));
let Admission::Allowed(reservation) = second else {
unreachable!("just asserted")
};
replica_b.settle(&k, &reservation, 700).await;
assert_eq!(
replica_a.reserve(&k, 300).await,
Admission::Denied(Denial::Exceeded)
);
}
#[tokio::test]
async fn an_expired_reservation_stops_counting_against_the_cap() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let expiring = expiring_settings(1_000, None, Duration::from_millis(50));
let store = RedisBudget::connect(&url, prefix(), expiring)
.await
.expect("connect");
let k = key();
assert!(matches!(
store.reserve(&k, 900).await,
Admission::Allowed(_)
));
assert_eq!(
store.reserve(&k, 900).await,
Admission::Denied(Denial::Exceeded)
);
tokio::time::sleep(Duration::from_millis(80)).await;
assert!(matches!(
store.reserve(&k, 900).await,
Admission::Allowed(_)
));
}
#[tokio::test]
async fn an_unreachable_server_denies_by_default() {
let err = RedisBudget::connect(
"redis://127.0.0.1:1/",
"axond:budget".to_owned(),
settings(1),
)
.await
.err()
.expect("an unreachable server must fail at boot");
assert!(matches!(err, BudgetError::Redis(_)), "{err:?}");
}
async fn namespace_store(
url: &str,
prefix: &str,
subject_limit: u64,
namespace_limit: u64,
) -> RedisBudget {
migrate_v1_to_v2(url, prefix, &namespaces())
.await
.expect("migrate");
RedisBudget::connect(
url,
prefix.to_owned(),
namespace_settings(subject_limit, namespace_limit),
)
.await
.expect("connect")
}
#[tokio::test]
async fn two_subjects_cannot_collectively_exceed_the_namespace_cap() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let store = namespace_store(&url, &prefix, 1_000, 1_200).await;
let first = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let second = BudgetKey {
namespace: "acme".into(),
subject: "second".into(),
};
let Admission::Allowed(held) = store.reserve(&first, 800).await else {
panic!("the first subject fits both caps");
};
assert_eq!(
store.reserve(&second, 800).await,
Admission::Denied(Denial::Exceeded)
);
assert!(matches!(
store.reserve(&second, 300).await,
Admission::Allowed(_)
));
store.settle(&first, &held, 800).await;
assert_eq!(
store.reserve(&second, 200).await,
Admission::Denied(Denial::Exceeded)
);
}
#[tokio::test]
async fn a_subject_cap_still_binds_under_a_generous_namespace_cap() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let store = namespace_store(&url, &prefix, 500, 1_000_000).await;
let k = key();
let Admission::Allowed(held) = store.reserve(&k, 500).await else {
panic!("the subject cap has room");
};
assert_eq!(
store.reserve(&k, 1).await,
Admission::Denied(Denial::Exceeded)
);
store.settle(&k, &held, 0).await;
assert!(matches!(
store.reserve(&k, 500).await,
Admission::Allowed(_)
));
}
#[tokio::test]
async fn namespaces_do_not_share_a_cap() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let store = namespace_store(&url, &prefix, 1_000, 1_000).await;
let acme = BudgetKey {
namespace: "acme".into(),
subject: "s".into(),
};
let other = BudgetKey {
namespace: "other".into(),
subject: "s".into(),
};
let Admission::Allowed(held) = store.reserve(&acme, 1_000).await else {
panic!("acme fits its own cap");
};
store.settle(&acme, &held, 1_000).await;
assert_eq!(
store.reserve(&acme, 1).await,
Admission::Denied(Denial::Exceeded)
);
assert!(matches!(
store.reserve(&other, 1_000).await,
Admission::Allowed(_)
));
}
#[tokio::test]
async fn releasing_frees_the_estimate_in_both_scopes() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let store = namespace_store(&url, &prefix, 1_000, 1_000).await;
let k = key();
let Admission::Allowed(held) = store.reserve(&k, 1_000).await else {
panic!("an empty namespace admits");
};
store.release(&k, &held).await;
let other = BudgetKey {
namespace: k.namespace.clone(),
subject: "another".into(),
};
assert!(matches!(
store.reserve(&other, 1_000).await,
Admission::Allowed(_)
));
}
#[tokio::test]
async fn a_partial_settlement_charges_both_scopes_once() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let store = namespace_store(&url, &prefix, 1_000, 1_000).await;
let k = key();
let Admission::Allowed(held) = store.reserve(&k, 900).await else {
panic!("an empty namespace admits");
};
store.settle(&k, &held, 100).await;
let other = BudgetKey {
namespace: k.namespace.clone(),
subject: "another".into(),
};
assert!(matches!(
store.reserve(&other, 900).await,
Admission::Allowed(_)
));
assert_eq!(
store.reserve(&other, 1).await,
Admission::Denied(Denial::Exceeded)
);
}
#[tokio::test]
async fn an_expired_hold_frees_the_namespace_cap_too() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate");
let expiring = expiring_settings(1_000, Some(1_000), Duration::from_secs(600));
let store = RedisBudget::connect(&url, prefix.clone(), expiring)
.await
.expect("connect");
let first = BudgetKey {
namespace: "acme".into(),
subject: "died".into(),
};
let second = BudgetKey {
namespace: "acme".into(),
subject: "alive".into(),
};
let Admission::Allowed(held) = store.reserve(&first, 900).await else {
panic!("the first reservation must be admitted");
};
assert_eq!(
store.reserve(&second, 900).await,
Admission::Denied(Denial::Exceeded)
);
{
let keys = v2_keys(&prefix, &first);
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let stale = format!("900:{}", now_ms() - 1);
for key in [&keys.subject_reservations, &keys.namespace_reservations] {
let rewritten: i64 = connection
.hset(key, &held.id, &stale)
.await
.expect("backdate the hold");
assert_eq!(rewritten, 0, "the hold is rewritten, not added");
}
}
assert!(matches!(
store.reserve(&second, 900).await,
Admission::Allowed(_)
));
}
#[tokio::test]
async fn two_replicas_enforce_one_namespace_cap_under_contention() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let replica_a = namespace_store(&url, &prefix, 1_000_000, 1_000).await;
let replica_b =
RedisBudget::connect(&url, prefix.clone(), namespace_settings(1_000_000, 1_000))
.await
.expect("connect");
let replica_a = Arc::new(replica_a);
let replica_b = Arc::new(replica_b);
let contenders = 40;
let start = Arc::new(Barrier::new(contenders));
let mut tasks = Vec::with_capacity(contenders);
for index in 0..contenders {
let key = BudgetKey {
namespace: "acme".into(),
subject: format!("subject-{index}"),
};
let store = if index % 2 == 0 {
Arc::clone(&replica_a)
} else {
Arc::clone(&replica_b)
};
let start = Arc::clone(&start);
tasks.push(tokio::spawn(async move {
start.wait().await;
match store.reserve(&key, 100).await {
Admission::Allowed(held) => {
store.settle(&key, &held, 100).await;
true
}
Admission::Denied(_) => false,
}
}));
}
let mut admitted = 0;
for task in tasks {
if task.await.expect("no task panicked") {
admitted += 1;
}
}
assert_eq!(admitted, 10);
}
#[tokio::test]
async fn the_migration_carries_v1_spend_into_both_scopes() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let v1 = RedisBudget::connect(&url, prefix.clone(), settings(1_000))
.await
.expect("connect");
for subject in ["first", "second"] {
let k = BudgetKey {
namespace: "acme".into(),
subject: subject.into(),
};
let Admission::Allowed(held) = v1.reserve(&k, 400).await else {
panic!("each subject has its own v1 cap");
};
v1.settle(&k, &held, 400).await;
}
let report = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate");
assert_eq!(report.subjects, 2);
assert_eq!(report.namespaces, 1);
assert_eq!(report.carried_microdollars, 800);
let again = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("re-migrate");
assert_eq!(again.carried_microdollars, 0);
assert_eq!(again.subjects, 0);
let store = RedisBudget::connect(&url, prefix, namespace_settings(1_000, 1_000))
.await
.expect("connect");
assert_eq!(
store
.reserve(
&BudgetKey {
namespace: "acme".into(),
subject: "third".into(),
},
201,
)
.await,
Admission::Denied(Denial::Exceeded)
);
assert_eq!(
store
.reserve(
&BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
},
601,
)
.await,
Admission::Denied(Denial::Exceeded)
);
}
#[tokio::test]
async fn a_v1_write_after_the_migration_is_added_not_dropped() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let k = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let v1 = RedisBudget::connect(&url, prefix.clone(), settings(10_000))
.await
.expect("connect");
let Admission::Allowed(held) = v1.reserve(&k, 400).await else {
panic!("admitted");
};
v1.settle(&k, &held, 400).await;
let first = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate");
assert_eq!(first.carried_microdollars, 400);
let Admission::Allowed(held) = v1.reserve(&k, 150).await else {
panic!("admitted");
};
v1.settle(&k, &held, 150).await;
let recovery = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("re-migrate");
assert_eq!(
recovery.carried_microdollars, 150,
"the stray write is added to the 400 already carried"
);
assert_eq!(spent(&url, &prefix, &k).await, 550);
assert_eq!(namespace_spent(&url, &prefix, "acme").await, 550);
let again = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("re-migrate");
assert_eq!(again.carried_microdollars, 0);
assert_eq!(spent(&url, &prefix, &k).await, 550);
}
#[tokio::test]
async fn recovering_a_stray_v1_write_keeps_the_v2_spend_beside_it() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let carried_subject = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let v2_only_subject = BudgetKey {
namespace: "acme".into(),
subject: "second".into(),
};
let v1 = RedisBudget::connect(&url, prefix.clone(), settings(10_000))
.await
.expect("connect");
let Admission::Allowed(held) = v1.reserve(&carried_subject, 400).await else {
panic!("admitted");
};
v1.settle(&carried_subject, &held, 400).await;
migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate");
let v2 = RedisBudget::connect(&url, prefix.clone(), namespace_settings(10_000, 10_000))
.await
.expect("connect");
let Admission::Allowed(held) = v2.reserve(&v2_only_subject, 250).await else {
panic!("admitted");
};
v2.settle(&v2_only_subject, &held, 250).await;
let _: () = ::redis::Client::open(url.as_str())
.expect("client")
.get_multiplexed_async_connection()
.await
.expect("connect")
.set(format!("{prefix}:{{acme|first}}:spent"), 150)
.await
.expect("stray v1 spend");
let recovery = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("re-migrate");
assert_eq!(recovery.carried_microdollars, 150);
assert_eq!(spent(&url, &prefix, &carried_subject).await, 550);
assert_eq!(
spent(&url, &prefix, &v2_only_subject).await,
250,
"the v2-only subject is untouched by the carry"
);
assert_eq!(
namespace_spent(&url, &prefix, "acme").await,
800,
"400 carried, 250 settled under v2, 150 recovered: the recovery added \
its delta instead of overwriting the total"
);
}
#[tokio::test]
async fn a_configured_namespace_that_is_only_a_prefix_does_not_claim_the_key() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let orphan = format!("{prefix}:{{team|west|abc}}:spent");
let _: () = connection.set(&orphan, 70).await.expect("orphan write");
let err = migrate_v1_to_v2(&url, &prefix, &["team".to_owned()])
.await
.expect_err("a prefix match over a `|`-bearing remainder must abort the migration");
assert!(format!("{err}").contains(&orphan), "{err}");
let left: Option<i64> = connection.get(&orphan).await.expect("read");
assert_eq!(left, Some(70), "nothing is moved or deleted on an abort");
assert_eq!(namespace_spent(&url, &prefix, "team").await, 0);
let marker: Option<String> = connection
.get(layout_key(&prefix))
.await
.expect("read marker");
assert_eq!(marker, None, "and nothing is stamped");
let report = migrate_v1_to_v2(&url, &prefix, &["team|west".to_owned()])
.await
.expect("the owning namespace resolves it");
assert_eq!(report.carried_microdollars, 70);
assert_eq!(namespace_spent(&url, &prefix, "team|west").await, 70);
assert_eq!(namespace_spent(&url, &prefix, "team").await, 0);
}
#[tokio::test]
async fn a_pipe_bearing_subject_is_unambiguous_under_the_v2_layout() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate an empty prefix");
let store = RedisBudget::connect(&url, prefix.clone(), namespace_settings(1_000, 1_000))
.await
.expect("connect");
let piped = BudgetKey {
namespace: "acme".into(),
subject: "sub|ject".into(),
};
let plain = BudgetKey {
namespace: "acme".into(),
subject: "sub".into(),
};
let Admission::Allowed(held) = store.reserve(&piped, 300).await else {
panic!("within both caps");
};
store.settle(&piped, &held, 300).await;
assert_eq!(spent(&url, &prefix, &piped).await, 300);
assert_eq!(spent(&url, &prefix, &plain).await, 0, "no key collision");
assert_eq!(namespace_spent(&url, &prefix, "acme").await, 300);
}
#[tokio::test]
async fn a_finished_migration_leaves_no_bookkeeping_behind() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let k = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let _: () = connection
.set(format!("{prefix}:{{acme|first}}:spent"), 250)
.await
.expect("v1 spend");
migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate");
assert_eq!(spent(&url, &prefix, &k).await, 250);
for key in [
format!("{prefix}:{{acme|first}}{SEQ}"),
format!("{}:migration:applied", namespace_scope(&prefix, "acme")),
] {
let left: bool = connection.exists(&key).await.expect("read");
assert!(!left, "`{key}` outlived the migration");
}
let _: () = connection
.set(format!("{prefix}:{{acme|first}}:spent"), 100)
.await
.expect("stray v1 spend");
let recovery = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("recover");
assert_eq!(recovery.carried_microdollars, 100);
assert_eq!(spent(&url, &prefix, &k).await, 350);
assert_eq!(namespace_spent(&url, &prefix, "acme").await, 350);
}
#[tokio::test]
async fn a_migration_that_stopped_part_way_serves_neither_layout() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let carried = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let waiting = BudgetKey {
namespace: "acme".into(),
subject: "second".into(),
};
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
for (key, spend) in [(&carried, 300), (&waiting, 200)] {
let _: () = connection
.set(
format!("{prefix}:{{{}|{}}}:spent", key.namespace, key.subject),
spend,
)
.await
.expect("v1 spend");
}
let _: () = connection
.set(layout_key(&prefix), LAYOUT_MIGRATING)
.await
.expect("mark mid-migration");
let scope = format!("{prefix}:{{acme|first}}");
let claim: Option<String> = Script::new(DRAIN_V1)
.prepare_invoke()
.key(format!("{scope}:spent"))
.key(format!("{scope}{PENDING}"))
.key(format!("{scope}{SEQ}"))
.invoke_async(&mut connection)
.await
.expect("claim");
let (sequence, amount) = claim.as_deref().expect("claimed").split_once(':').unwrap();
let _: i64 = Script::new(APPLY_CLAIM)
.prepare_invoke()
.key(v2_keys(&prefix, &carried).subject_spent)
.key(v2_keys(&prefix, &carried).namespace_spent)
.key(format!(
"{}:migration:applied",
namespace_scope(&prefix, "acme")
))
.arg(format!("first#{sequence}"))
.arg(amount.parse::<i64>().unwrap())
.invoke_async(&mut connection)
.await
.expect("apply");
let _: i64 = connection
.del(format!("{scope}{PENDING}"))
.await
.expect("finish the carry");
for shared in [namespace_settings(1_000, 1_000), settings(1_000)] {
let refused = RedisBudget::connect(&url, prefix.clone(), shared)
.await
.err()
.expect("a part-way migration must fail at boot");
assert!(
format!("{refused}").contains("did not finish"),
"the error must name the unfinished migration: {refused}"
);
}
let resumed = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("resume");
assert_eq!(
resumed.carried_microdollars, 200,
"the already-carried subject is not charged again"
);
assert_eq!(spent(&url, &prefix, &carried).await, 300);
assert_eq!(spent(&url, &prefix, &waiting).await, 200);
assert_eq!(namespace_spent(&url, &prefix, "acme").await, 500);
RedisBudget::connect(&url, prefix.clone(), namespace_settings(1_000, 1_000))
.await
.expect("a finished migration boots");
}
#[tokio::test]
async fn an_outstanding_migration_claim_stops_either_configuration_booting() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let k = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let store = namespace_store(&url, &prefix, 1_000, 1_000).await;
let Admission::Allowed(held) = store.reserve(&k, 400).await else {
panic!("migrated state serves as usual");
};
store.settle(&k, &held, 400).await;
drop(store);
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let scope = format!("{prefix}:{{acme|first}}");
let _: () = connection
.set(format!("{scope}:spent"), 150)
.await
.expect("stray v1 spend");
let claim: Option<String> = Script::new(DRAIN_V1)
.prepare_invoke()
.key(format!("{scope}:spent"))
.key(format!("{scope}{PENDING}"))
.key(format!("{scope}{SEQ}"))
.invoke_async(&mut connection)
.await
.expect("claim");
assert!(claim.is_some(), "the claim took the v1 counter");
let refused = RedisBudget::connect(&url, prefix.clone(), namespace_settings(1_000, 1_000))
.await
.err()
.expect("an outstanding claim must fail at boot");
assert!(
format!("{refused}").contains("migrate-redis"),
"the error must point at the migration: {refused}"
);
let refused = RedisBudget::connect(&url, prefix.clone(), settings(1_000))
.await
.err()
.expect("migrated state must fail at boot without the cap");
assert!(
format!("{refused}").contains("must stay set"),
"the marker is what refuses the old layout: {refused}"
);
let recovery = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("resume");
assert_eq!(recovery.carried_microdollars, 150);
assert_eq!(spent(&url, &prefix, &k).await, 550);
assert_eq!(namespace_spent(&url, &prefix, "acme").await, 550);
RedisBudget::connect(&url, prefix.clone(), namespace_settings(1_000, 1_000))
.await
.expect("a finished migration boots");
}
#[tokio::test]
async fn an_interrupted_carry_is_finished_exactly_once() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let k = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let scope = format!("{prefix}:{{acme|first}}");
let drain = Script::new(DRAIN_V1);
let apply = Script::new(APPLY_CLAIM);
let _: () = connection
.set(format!("{scope}:spent"), 300)
.await
.expect("v1 spend");
let claim: Option<String> = drain
.prepare_invoke()
.key(format!("{scope}:spent"))
.key(format!("{scope}{PENDING}"))
.key(format!("{scope}{SEQ}"))
.invoke_async(&mut connection)
.await
.expect("claim");
assert_eq!(claim.as_deref(), Some("1:300"));
let v1_gone: Option<i64> = connection
.get(format!("{scope}:spent"))
.await
.expect("read");
assert_eq!(v1_gone, None, "a claimed counter is drained atomically");
let resumed = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("resume");
assert_eq!(
resumed.carried_microdollars, 300,
"the orphaned claim is finished, not lost"
);
assert_eq!(spent(&url, &prefix, &k).await, 300);
let claim = "2:75";
let _: () = connection
.set(format!("{scope}{PENDING}"), claim)
.await
.expect("orphan claim");
let _: () = connection
.set(format!("{scope}{SEQ}"), 2)
.await
.expect("sequence");
for expected in [75, 0] {
let added: i64 = apply
.prepare_invoke()
.key(v2_keys(&prefix, &k).subject_spent)
.key(v2_keys(&prefix, &k).namespace_spent)
.key(format!(
"{}:migration:applied",
namespace_scope(&prefix, "acme")
))
.arg("first#2")
.arg(75)
.invoke_async(&mut connection)
.await
.expect("apply");
assert_eq!(added, expected, "a claim is applied at most once");
}
let recovered = migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("resume");
assert_eq!(
recovered.carried_microdollars, 0,
"the already-applied claim is not charged again"
);
assert_eq!(spent(&url, &prefix, &k).await, 375);
assert_eq!(namespace_spent(&url, &prefix, "acme").await, 375);
}
async fn read(url: &str, key: String) -> u64 {
let client = ::redis::Client::open(url).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let value: Option<i64> = connection.get(key).await.expect("read");
value.unwrap_or_default().max(0) as u64
}
async fn spent(url: &str, prefix: &str, key: &BudgetKey) -> u64 {
read(url, v2_keys(prefix, key).subject_spent).await
}
async fn namespace_spent(url: &str, prefix: &str, namespace: &str) -> u64 {
read(
url,
format!("{}:namespace:spent", namespace_scope(prefix, namespace)),
)
.await
}
#[tokio::test]
async fn an_unattributable_v1_key_aborts_the_migration_intact() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
let v1 = RedisBudget::connect(&url, prefix.clone(), settings(1_000))
.await
.expect("connect");
let k = BudgetKey {
namespace: "acme".into(),
subject: "first".into(),
};
let Admission::Allowed(held) = v1.reserve(&k, 400).await else {
panic!("the v1 cap admits it");
};
v1.settle(&k, &held, 400).await;
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let orphan = format!("{prefix}:{{retired|sub}}:spent");
let _: () = connection.set(&orphan, 25).await.expect("orphan write");
for namespaces in [namespaces(), vec!["acme".to_owned(), "team".to_owned()]] {
let err = migrate_v1_to_v2(&url, &prefix, &namespaces)
.await
.expect_err("an unattributable key must abort the migration");
assert!(format!("{err}").contains(&orphan), "{err}");
}
let v1_spend: Option<i64> = connection
.get(format!("{prefix}:{{acme|first}}:spent"))
.await
.expect("read");
assert_eq!(v1_spend, Some(400), "the v1 ledger is untouched");
let orphan_spend: Option<i64> = connection.get(&orphan).await.expect("read");
assert_eq!(orphan_spend, Some(25));
let marker: Option<String> = connection
.get(layout_key(&prefix))
.await
.expect("read marker");
assert_eq!(marker, None, "an aborted migration marks nothing");
let report = migrate_v1_to_v2(&url, &prefix, &["acme".to_owned(), "retired".to_owned()])
.await
.expect("migrate");
assert_eq!(report.carried_microdollars, 425);
}
#[tokio::test]
async fn the_namespace_cap_refuses_to_boot_against_unmigrated_state() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let err = RedisBudget::connect(&url, prefix(), namespace_settings(1_000, 1_000))
.await
.err()
.expect("un-migrated state must fail at boot");
assert!(
format!("{err}").contains("migrate-redis"),
"the error must name the migration: {err}"
);
}
#[tokio::test]
async fn a_v1_key_written_after_the_migration_is_rejected() {
let Some(url) = crate::test_services::redis_url() else {
return;
};
let prefix = prefix();
migrate_v1_to_v2(&url, &prefix, &namespaces())
.await
.expect("migrate");
let v1 = RedisBudget::connect(&url, prefix.clone(), settings(1_000))
.await
.err()
.expect("a v1 configuration must not boot against migrated state");
assert!(
format!("{v1}").contains("namespace_limit_microdollars"),
"{v1}"
);
let client = ::redis::Client::open(url.as_str()).expect("client");
let mut connection = ConnectionManager::new(client).await.expect("connect");
let _: () = connection
.set(format!("{prefix}:{{acme|stale}}:spent"), 10)
.await
.expect("legacy write");
let err = RedisBudget::connect(&url, prefix, namespace_settings(1_000, 1_000))
.await
.err()
.expect("mixed binaries must fail at boot");
assert!(format!("{err}").contains("v1 budget key"), "{err}");
}
}