//! each repo and each host get a few full refetches per window, because a peer can make us refetch
//! a repo over and over (#sync, chain breaks, newer listings, crawls). an operator's resync is free
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use jacquard_common::types::did::Did;
use miette::{IntoDiagnostic, Result};
use smol_str::SmolStr;
use tracing::debug;
use crate::config::Config;
use crate::db::counts::pds_host;
use crate::db::{Txn, keys, ser_repo_meta};
use crate::state::AppState;
use crate::types::{
Allowance, GaugeState, RepoMetadata, RepoState, RepoStatus, ResyncErrorKind, ResyncState,
};
#[derive(Debug)]
pub(crate) struct RefetchBudget {
repo_limit: Option<u32>,
host_limit: Option<u32>,
window: i64,
/// shared here because a host's repos live on different shards. a spend whose txn never commits
/// still counts until the window ends
hosts: parking_lot::Mutex<HashMap<SmolStr, Allowance>>,
}
impl RefetchBudget {
pub(crate) fn new(config: &Config) -> Self {
Self {
repo_limit: config.repo_refetch_limit,
host_limit: config.host_refetch_limit,
window: config.repo_refetch_window.as_secs().max(1) as i64,
hosts: parking_lot::Mutex::new(HashMap::new()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Charged {
Allowed,
Deferred,
}
/// on `Allowed` only `metadata` changes and the caller writes it. on `Deferred` the repo is
/// parked in `txn`
pub(crate) fn charge(
state: &AppState,
txn: &mut Txn<'_>,
did: &Did,
repo: &mut RepoState<'_>,
metadata: &mut RepoMetadata,
cause: &str,
) -> Result<Charged> {
let budget = &state.refetch_budget;
let now = chrono::Utc::now().timestamp();
let mut allowance = metadata.refetch;
let mut spent = match budget.repo_limit {
Some(limit) => Allowance::take(&mut allowance, now, limit, budget.window),
None => Ok(()),
};
if spent.is_ok() {
spent = take_from_host(state, txn, repo, now)?;
}
if let Err(until) = spent {
debug!(%did, cause, until, "out of refetches, deferring");
state.refetch_denied.fetch_add(1, Ordering::Relaxed);
defer(txn, did, repo, metadata, until)?;
return Ok(Charged::Deferred);
}
metadata.refetch = allowance;
if std::mem::take(&mut metadata.refetch_deferred) {
txn.batch
.remove(&txn.db.indexer.resync, keys::repo_key(did));
}
Ok(Charged::Allowed)
}
/// the inner `Err` is when the host's window ends
fn take_from_host(
state: &AppState,
txn: &mut Txn<'_>,
repo: &RepoState<'_>,
now: i64,
) -> Result<Result<(), i64>> {
let budget = &state.refetch_budget;
let (Some(limit), Some(host)) = (budget.host_limit, repo.pds.as_deref().and_then(pds_host))
else {
return Ok(Ok(()));
};
if state.is_trusted_host(&host) {
return Ok(Ok(()));
}
let key = keys::refetch_allowance_key(&host);
let mut hosts = budget.hosts.lock();
let mut allowance = match hosts.get(&host) {
Some(allowance) => Some(*allowance),
None => txn
.db
.cursors
.get(&key)
.into_diagnostic()?
.map(|bytes| rmp_serde::from_slice::<Allowance>(&bytes).into_diagnostic())
.transpose()?,
};
if let Err(until) = Allowance::take(&mut allowance, now, limit, budget.window) {
return Ok(Err(until));
}
let spent = allowance.expect("take leaves an allowance");
txn.batch.insert(
&txn.db.cursors,
&key,
rmp_serde::to_vec(&spent).into_diagnostic()?,
);
hosts.insert(host, spent);
Ok(Ok(()))
}
fn defer(
txn: &mut Txn<'_>,
did: &Did,
repo: &mut RepoState<'_>,
metadata: &mut RepoMetadata,
until: i64,
) -> Result<()> {
if !repo.active {
return Ok(());
}
metadata.refetch_deferred = true;
let db = txn.db;
txn.batch.insert(
&db.repo_metadata,
keys::repo_metadata_key(did),
ser_repo_meta(metadata)?,
);
let entry = ResyncState::Error {
kind: ResyncErrorKind::Ratelimited,
retry_count: metadata.retry_count,
next_retry: until,
};
txn.batch.insert(
&db.indexer.resync,
keys::repo_key(did),
rmp_serde::to_vec(&entry).into_diagnostic()?,
);
txn.transition_lifecycle(did, GaugeState::Resync(Some(ResyncErrorKind::Ratelimited)))?;
repo.status = RepoStatus::Throttled;
repo.touch();
txn.put_repo(did, repo)?;
Ok(())
}
#[cfg(test)]
pub(crate) mod tests {
use std::time::Duration;
use super::*;
pub(crate) fn state(limit: Option<u32>) -> (tempfile::TempDir, AppState) {
let tmp = tempfile::tempdir().unwrap();
let state = AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
repo_refetch_limit: limit,
repo_refetch_window: Duration::from_secs(3600),
..Default::default()
})
.unwrap();
(tmp, state)
}
pub(crate) fn did() -> Did {
Did::new_static("did:plc:ewvi7nxzyoun6zhxrhs64oiz").unwrap()
}
pub(crate) fn put(state: &AppState, did: &Did, repo: &RepoState<'_>, metadata: &RepoMetadata) {
let mut batch = state.db.inner.batch();
batch.insert(
&state.db.repos,
keys::repo_key(did),
crate::db::ser_repo_state(repo).unwrap(),
);
batch.insert(
&state.db.repo_metadata,
keys::repo_metadata_key(did),
ser_repo_meta(metadata).unwrap(),
);
batch.commit().unwrap();
}
pub(crate) fn put_metadata(state: &AppState, did: &Did, metadata: &RepoMetadata) {
state
.db
.repo_metadata
.insert(
keys::repo_metadata_key(did),
ser_repo_meta(metadata).unwrap(),
)
.unwrap();
}
pub(crate) fn seed(state: &AppState, active: bool, used: Option<(i64, u32)>) -> RepoMetadata {
let did = did();
let mut repo = RepoState::synced();
repo.active = active;
let mut metadata = RepoMetadata::backfilling(7);
metadata.retry_count = 3;
metadata.refetch = used.map(|(since, used)| Allowance { since, used });
put(state, &did, &repo, &metadata);
metadata
}
pub(crate) fn row_of(state: &AppState, did: &Did) -> RepoState<'static> {
use jacquard_common::IntoStatic;
let bytes = state.db.repos.get(keys::repo_key(did)).unwrap().unwrap();
crate::db::deser_repo_state(&bytes).unwrap().into_static()
}
pub(crate) fn row(state: &AppState) -> RepoState<'static> {
row_of(state, &did())
}
pub(crate) fn metadata_of(state: &AppState, did: &Did) -> RepoMetadata {
let bytes = state
.db
.repo_metadata
.get(keys::repo_metadata_key(did))
.unwrap()
.unwrap();
crate::db::deser_repo_meta(&bytes).unwrap()
}
pub(crate) fn put_entry(state: &AppState, did: &Did, entry: &ResyncState) {
state
.db
.indexer
.resync
.insert(keys::repo_key(did), rmp_serde::to_vec(entry).unwrap())
.unwrap();
}
pub(crate) fn entry_of(state: &AppState, did: &Did) -> Option<ResyncState> {
state
.db
.indexer
.resync
.get(keys::repo_key(did))
.unwrap()
.map(|bytes| rmp_serde::from_slice(&bytes).unwrap())
}
fn metadata(state: &AppState) -> RepoMetadata {
metadata_of(state, &did())
}
fn entry(state: &AppState) -> Option<ResyncState> {
entry_of(state, &did())
}
fn charge_once(state: &AppState) -> Charged {
charge_did(state, &did())
}
fn denied_repo() -> (tempfile::TempDir, AppState) {
let (tmp, state) = state(Some(1));
let now = chrono::Utc::now().timestamp();
seed(&state, true, Some((now, 1)));
assert_eq!(charge_once(&state), Charged::Deferred);
assert!(metadata_of(&state, &did()).refetch_deferred);
(tmp, state)
}
#[test]
fn a_status_that_takes_the_resync_entry_away_ends_the_denial_too() {
use crate::ops::transition_repo;
for (status, waits) in [
(RepoStatus::Synced, false),
(RepoStatus::Deleted, false),
(RepoStatus::Deactivated, false),
// the retry worker requeues a desynchronized repo, it keeps waiting
(RepoStatus::Desynchronized, true),
] {
let (_tmp, state) = denied_repo();
let mut txn = Txn::new(&state.db);
transition_repo(
&mut txn,
&mut Vec::new(),
&did(),
row(&state),
status.clone(),
)
.unwrap();
txn.commit().unwrap();
assert_eq!(
metadata_of(&state, &did()).refetch_deferred,
waits,
"{status:?}"
);
}
}
#[test]
fn untracking_or_erasing_a_repo_ends_its_denial() {
use crate::ingest::indexer::message::RepoOp;
let (_tmp, state) = denied_repo();
assert!(
crate::ingest::indexer::repos::repo_op(&state, None, RepoOp::Untrack, &did()).unwrap()
);
assert!(!metadata_of(&state, &did()).refetch_deferred);
assert!(entry_of(&state, &did()).is_none());
let (_tmp, state) = denied_repo();
let mut txn = Txn::new(&state.db);
crate::db::start_erase(&mut txn, &did(), false).unwrap();
txn.commit().unwrap();
assert!(!metadata_of(&state, &did()).refetch_deferred);
assert!(entry_of(&state, &did()).is_none());
}
#[test]
fn the_allowance_counts_a_fixed_window() {
let mut allowance = None;
for used in 1..=2 {
assert_eq!(Allowance::take(&mut allowance, 1000, 2, 60), Ok(()));
assert_eq!(allowance, Some(Allowance { since: 1000, used }));
}
assert_eq!(Allowance::take(&mut allowance, 1059, 2, 60), Err(1060));
assert_eq!(
allowance,
Some(Allowance {
since: 1000,
used: 2
})
);
assert_eq!(Allowance::take(&mut allowance, 1060, 2, 60), Ok(()));
assert_eq!(
allowance,
Some(Allowance {
since: 1060,
used: 1
})
);
}
#[test]
fn an_allowance_charged_by_amount_is_spent_at_its_limit_and_checked_without_taking() {
let mut allowance = None;
assert_eq!(Allowance::check(&allowance, 1000, 10, 60), Ok(()));
assert_eq!(allowance, None);
Allowance::charge(&mut allowance, 1000, 9, 60);
assert_eq!(Allowance::check(&allowance, 1010, 10, 60), Ok(()));
Allowance::charge(&mut allowance, 1010, 4, 60);
assert_eq!(
allowance,
Some(Allowance {
since: 1000,
used: 13
})
);
assert_eq!(Allowance::check(&allowance, 1059, 10, 60), Err(1060));
assert_eq!(Allowance::check(&allowance, 1060, 10, 60), Ok(()));
Allowance::charge(&mut allowance, 1060, u32::MAX, 60);
Allowance::charge(&mut allowance, 1061, u32::MAX, 60);
assert_eq!(
allowance,
Some(Allowance {
since: 1060,
used: u32::MAX
})
);
}
#[test]
fn a_repo_refetches_up_to_its_limit_and_then_waits_for_the_window_to_end() {
let (_tmp, state) = state(Some(2));
seed(&state, true, None);
assert_eq!(charge_once(&state), Charged::Allowed);
assert_eq!(charge_once(&state), Charged::Allowed);
let since = metadata(&state).refetch.unwrap().since;
assert!(entry(&state).is_none());
assert_eq!(charge_once(&state), Charged::Deferred);
let repo = row(&state);
assert!(repo.active, "a denial isn't the account going inactive");
assert_eq!(repo.status, RepoStatus::Throttled);
let metadata = metadata(&state);
assert!(metadata.refetch_deferred);
assert_eq!(metadata.refetch.unwrap().used, 2, "a denial costs nothing");
let Some(ResyncState::Error {
kind,
retry_count,
next_retry,
}) = entry(&state)
else {
panic!("a denied refetch waits as a resync error");
};
assert_eq!(kind, ResyncErrorKind::Ratelimited);
assert_eq!(retry_count, 3, "waiting isn't failing");
assert_eq!(next_retry, since + 3600);
assert_eq!(state.refetch_denied.load(Ordering::Relaxed), 1);
assert_eq!(charge_once(&state), Charged::Deferred);
assert!(row(&state).active);
assert!(
matches!(entry(&state), Some(ResyncState::Error { next_retry: n, .. }) if n == since + 3600)
);
}
#[test]
fn a_new_window_admits_a_denied_repo_and_clears_what_parked_it() {
let (_tmp, state) = state(Some(1));
let long_ago = chrono::Utc::now().timestamp() - 7200;
seed(&state, true, Some((long_ago, 1)));
assert_eq!(charge_once(&state), Charged::Allowed);
let allowance = metadata(&state).refetch.unwrap();
assert!(allowance.since > long_ago && allowance.used == 1);
assert_eq!(charge_once(&state), Charged::Deferred);
assert!(entry(&state).is_some());
let mut stale = metadata(&state);
stale.refetch = Some(Allowance {
since: long_ago,
used: 1,
});
put_metadata(&state, &did(), &stale);
assert_eq!(charge_once(&state), Charged::Allowed);
assert!(!metadata(&state).refetch_deferred);
assert!(
entry(&state).is_none(),
"the refetch it queues replaces the wait"
);
}
#[test]
fn no_limit_charges_nothing() {
let (_tmp, state) = state(None);
seed(&state, true, None);
for _ in 0..20 {
assert_eq!(charge_once(&state), Charged::Allowed);
}
assert_eq!(metadata(&state).refetch, None);
}
#[test]
fn an_inactive_repo_is_denied_without_being_marked() {
let (_tmp, state) = state(Some(1));
seed(&state, false, Some((chrono::Utc::now().timestamp(), 1)));
assert_eq!(charge_once(&state), Charged::Deferred);
let repo = row(&state);
assert!(!repo.active);
assert_eq!(repo.status, RepoStatus::Synced);
assert!(entry(&state).is_none());
assert!(!metadata(&state).refetch_deferred);
}
pub(crate) fn on_a_host(state: &AppState, did: &Did, active: bool) {
let mut repo = RepoState::synced();
repo.active = active;
repo.pds = Some(jacquard_common::CowStr::Borrowed("https://pds.example/"));
put(state, did, &repo, &RepoMetadata::backfilling(7));
}
pub(crate) fn charge_did(state: &AppState, did: &Did) -> Charged {
let mut txn = Txn::new(&state.db);
let mut repo = row_of(state, did);
let mut metadata = metadata_of(state, did);
let charged = charge(state, &mut txn, did, &mut repo, &mut metadata, "test").unwrap();
if charged == Charged::Allowed {
txn.batch.insert(
&state.db.repo_metadata,
keys::repo_metadata_key(did),
ser_repo_meta(&metadata).unwrap(),
);
}
txn.commit().unwrap();
charged
}
pub(crate) fn host_config(
path: &std::path::Path,
repo_limit: u32,
host_limit: Option<u32>,
tier: Option<&str>,
window: Duration,
) -> Config {
Config {
database_path: path.to_path_buf(),
repo_refetch_limit: Some(repo_limit),
host_refetch_limit: host_limit,
repo_refetch_window: window,
tier_policy: crate::pds_meta::TierPolicy {
tiers: Config::default().tier_policy.tiers,
rules: tier
.map(|tier| crate::pds_meta::TierRule {
pattern: "pds.example".into(),
tier_name: tier.into(),
})
.into_iter()
.collect(),
},
..Default::default()
}
}
pub(crate) fn host_state(
host_limit: Option<u32>,
tier: Option<&str>,
) -> (tempfile::TempDir, AppState) {
host_state_in(host_limit, tier, Duration::from_secs(3600))
}
pub(crate) fn host_state_in(
host_limit: Option<u32>,
tier: Option<&str>,
window: Duration,
) -> (tempfile::TempDir, AppState) {
let tmp = tempfile::tempdir().unwrap();
let config = host_config(tmp.path(), 10, host_limit, tier, window);
(tmp, AppState::new(&config).unwrap())
}
pub(crate) fn dids(n: usize) -> Vec<Did> {
(0..n)
.map(|i| Did::new_owned(format!("did:plc:hostrefetchrepo{i:0>4}")).unwrap())
.collect()
}
#[test]
fn the_repos_of_one_host_share_its_refetches() {
let (_tmp, state) = host_state(Some(3), None);
let repos = dids(5);
repos.iter().for_each(|did| on_a_host(&state, did, true));
let charged = repos
.iter()
.map(|did| charge_did(&state, did))
.collect::<Vec<_>>();
assert_eq!(
charged,
[
Charged::Allowed,
Charged::Allowed,
Charged::Allowed,
Charged::Deferred,
Charged::Deferred
]
);
assert_eq!(metadata_of(&state, &repos[3]).refetch, None);
assert_eq!(row_of(&state, &repos[3]).status, RepoStatus::Throttled);
assert!(matches!(
entry_of(&state, &repos[3]),
Some(ResyncState::Error { next_retry, .. })
if next_retry > chrono::Utc::now().timestamp()
));
}
#[test]
fn the_allowances_survive_a_restart() {
let tmp = tempfile::tempdir().unwrap();
let config = host_config(tmp.path(), 1, Some(2), None, Duration::from_secs(3600));
let repos = dids(3);
{
let state = AppState::new(&config).unwrap();
repos.iter().for_each(|did| on_a_host(&state, did, true));
assert_eq!(charge_did(&state, &repos[0]), Charged::Allowed);
assert_eq!(charge_did(&state, &repos[1]), Charged::Allowed);
}
let state = AppState::new(&config).unwrap();
assert_eq!(
charge_did(&state, &repos[2]),
Charged::Deferred,
"the host's"
);
assert_eq!(
charge_did(&state, &repos[0]),
Charged::Deferred,
"the repo's"
);
}
#[test]
fn a_charge_that_never_commits_leaves_nothing_behind_after_a_restart() {
let tmp = tempfile::tempdir().unwrap();
let config = host_config(tmp.path(), 1, Some(1), None, Duration::from_secs(3600));
let repos = dids(2);
{
let state = AppState::new(&config).unwrap();
repos.iter().for_each(|did| on_a_host(&state, did, true));
let mut txn = Txn::new(&state.db);
let mut repo = row_of(&state, &repos[0]);
let mut metadata = metadata_of(&state, &repos[0]);
let charged = charge(
&state,
&mut txn,
&repos[0],
&mut repo,
&mut metadata,
"test",
);
assert_eq!(charged.unwrap(), Charged::Allowed);
}
let state = AppState::new(&config).unwrap();
assert_eq!(charge_did(&state, &repos[1]), Charged::Allowed);
}
#[test]
fn a_denied_charge_writes_no_host_count() {
let (_tmp, state) = host_state(Some(1), None);
let repos = dids(3);
repos.iter().for_each(|did| on_a_host(&state, did, true));
let key = keys::refetch_allowance_key("pds.example");
let count = || {
let bytes = state.db.cursors.get(&key).unwrap().unwrap();
rmp_serde::from_slice::<Allowance>(&bytes).unwrap().used
};
assert_eq!(charge_did(&state, &repos[0]), Charged::Allowed);
assert_eq!(count(), 1);
for did in &repos[1..] {
assert_eq!(charge_did(&state, did), Charged::Deferred);
}
assert_eq!(count(), 1);
}
#[test]
fn a_repo_with_no_known_host_is_charged_only_its_own() {
let (_tmp, state) = host_state(Some(1), None);
let repos = dids(3);
for did in &repos {
put(
&state,
did,
&RepoState::synced(),
&RepoMetadata::backfilling(7),
);
}
for did in &repos {
assert_eq!(charge_did(&state, did), Charged::Allowed);
}
}
#[test]
fn the_trusted_tier_is_not_counted_against_its_host_but_the_default_is() {
// the default tier has trusted limits in an indexer build, only the name tells them apart
for (tier, exempt) in [(Some("trusted"), true), (None, false)] {
let (_tmp, state) = host_state(Some(1), tier);
let repos = dids(3);
repos.iter().for_each(|did| on_a_host(&state, did, true));
let charged = repos
.iter()
.map(|did| charge_did(&state, did) == Charged::Allowed)
.collect::<Vec<_>>();
assert_eq!(charged, [true, exempt, exempt], "{tier:?}");
}
}
#[test]
fn a_trusted_hosts_repos_still_have_their_own_limit() {
let (_tmp, state) = host_state(Some(1), Some("trusted"));
let did = dids(1).remove(0);
on_a_host(&state, &did, true);
for _ in 0..10 {
assert_eq!(charge_did(&state, &did), Charged::Allowed);
}
assert_eq!(charge_did(&state, &did), Charged::Deferred);
}
#[test]
fn the_host_limit_is_off_with_none() {
let (_tmp, state) = host_state(None, None);
for did in dids(20) {
on_a_host(&state, &did, true);
assert_eq!(charge_did(&state, &did), Charged::Allowed);
}
}
}