//! the one-off `listRepos` walk an indexer owes a PDS it starts following, and the
//! rewalk a followed PDS can ask for through `requestCrawl`.
//!
//! a new PDS is tailed from the relay's seq, so repos that already lived on it
//! never reach us as events and nothing else lists them. only hosts that are new
//! to the network get the walk: a seed relay's hosts count as bootstrap until its
//! `listHosts` has been read to the end once, otherwise a fresh node would walk
//! every PDS out there. a relay has nothing to backfill, so there it's a no-op.
//!
//! a rewalk lists the host again to pick up repos we missed and resync the ones of
//! its repos listed at a newer rev than ours, or stuck desynchronized or failing.
//! a host's rewalks may read [`crate::config::Config::rewalk_mb`] of `listRepos` responses per
//! rewalk window. a walk that spends them stops, keeps its cursor and goes on once the window
//! ends, and asking while one is queued, running or waiting does nothing.
#[cfg(feature = "indexer")]
mod enabled {
use miette::{IntoDiagnostic, Result};
use smol_str::SmolStr;
use url::Url;
use crate::db::{Db, keys};
use crate::state::AppState;
use crate::types::Allowance;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Rewalk {
Queued,
AlreadyWalking,
OutOfAllowance,
}
/// the `listRepos` bytes a host's rewalks may read per `window` seconds, counted from the
/// first page one reads. the walk that reads the page that overshoots `limit` finishes it.
pub(crate) struct Budget {
key: Vec<u8>,
limit: u32,
window: i64,
allowance: Option<Allowance>,
}
impl Budget {
pub(crate) fn load(db: &Db, host: &str, limit: u32, window: i64) -> Result<Self> {
let key = keys::rewalk_allowance_key(host);
let allowance = db
.cursors
.get(&key)
.into_diagnostic()?
.map(|bytes| rmp_serde::from_slice::<Allowance>(&bytes).into_diagnostic())
.transpose()?;
Ok(Self {
key,
limit,
window,
allowance,
})
}
/// the unix time the window ends, once it has spent its bytes
pub(crate) fn spent_until(&self, now: i64) -> Option<i64> {
Allowance::check(&self.allowance, now, self.limit, self.window).err()
}
/// the key and value that persist the charge
pub(crate) fn charge(&mut self, now: i64, bytes: usize) -> Result<(Vec<u8>, Vec<u8>)> {
let bytes = u32::try_from(bytes).unwrap_or(u32::MAX);
Allowance::charge(&mut self.allowance, now, bytes, self.window);
let value = rmp_serde::to_vec(&self.allowance).into_diagnostic()?;
Ok((self.key.clone(), value))
}
}
/// queue a rewalk of `host` unless it's walking already or its window has spent `limit` bytes.
/// call [`wake`] once it's queued.
pub(crate) fn request_rewalk(
db: &Db,
host: &str,
now: i64,
window: i64,
limit: u32,
) -> Result<Rewalk> {
let walk_key = keys::rewalk_key(host);
if db.cursors.contains_key(&walk_key).into_diagnostic()?
|| db
.cursors
.contains_key(keys::new_pds_crawl_key(host))
.into_diagnostic()?
{
return Ok(Rewalk::AlreadyWalking);
}
if Budget::load(db, host, limit, window)?
.spent_until(now)
.is_some()
{
return Ok(Rewalk::OutOfAllowance);
}
db.cursors
.insert(walk_key, encode_cursor(None)?)
.into_diagnostic()?;
Ok(Rewalk::Queued)
}
/// stage the walk in the batch that admits the host, so a restart can't lose it.
pub(crate) fn queue(batch: &mut fjall::OwnedWriteBatch, db: &Db, host: &str) -> Result<()> {
batch.insert(
&db.cursors,
keys::new_pds_crawl_key(host),
encode_cursor(None)?,
);
Ok(())
}
/// call once the batch from [`queue`] has committed.
pub(crate) fn wake(state: &AppState) {
state.new_pds_crawl_notify.notify_one();
}
pub(crate) fn seed_listed(db: &Db, seed: &Url) -> Result<bool> {
db.cursors
.contains_key(keys::seed_listed_key(seed.as_str()))
.into_diagnostic()
}
pub(crate) fn mark_seed_listed(db: &Db, seed: &Url) -> Result<()> {
db.cursors
.insert(keys::seed_listed_key(seed.as_str()), b"")
.into_diagnostic()
}
pub(crate) fn encode_cursor(cursor: Option<&SmolStr>) -> Result<Vec<u8>> {
rmp_serde::to_vec(&cursor).into_diagnostic()
}
pub(crate) fn decode_cursor(bytes: &[u8]) -> Result<Option<SmolStr>> {
rmp_serde::from_slice(bytes).into_diagnostic()
}
}
#[cfg(feature = "indexer")]
pub(crate) use enabled::*;
#[cfg(not(feature = "indexer"))]
mod noop {
use miette::Result;
use url::Url;
use crate::db::Db;
use crate::state::AppState;
pub(crate) fn queue(_batch: &mut fjall::OwnedWriteBatch, _db: &Db, _host: &str) -> Result<()> {
Ok(())
}
pub(crate) fn wake(_state: &AppState) {}
pub(crate) fn seed_listed(_db: &Db, _seed: &Url) -> Result<bool> {
Ok(false)
}
pub(crate) fn mark_seed_listed(_db: &Db, _seed: &Url) -> Result<()> {
Ok(())
}
}
#[cfg(not(feature = "indexer"))]
pub(crate) use noop::*;
#[cfg(all(test, feature = "indexer"))]
mod tests {
use super::*;
use crate::config::Config;
use crate::db::keys;
use crate::state::AppState;
use miette::{IntoDiagnostic, Result};
const WINDOW: i64 = 1800;
const LIMIT: u32 = 1000;
fn finish_walk(state: &AppState, host: &str) -> Result<()> {
state
.db
.cursors
.remove(keys::rewalk_key(host))
.into_diagnostic()
}
fn read_bytes(state: &AppState, host: &str, now: i64, bytes: usize) -> Result<()> {
let mut budget = Budget::load(&state.db, host, LIMIT, WINDOW)?;
let (key, value) = budget.charge(now, bytes)?;
state.db.cursors.insert(key, value).into_diagnostic()
}
#[test]
fn the_allowance_comes_back_once_its_window_is_over() -> Result<()> {
let dir = tempfile::tempdir().into_diagnostic()?;
let state = AppState::new(&Config {
database_path: dir.path().to_path_buf(),
..Default::default()
})?;
for at in [0, 10, 20] {
assert_eq!(
request_rewalk(&state.db, "pds.example", at, WINDOW, LIMIT)?,
Rewalk::Queued,
"walks don't use the allowance, their bytes do"
);
finish_walk(&state, "pds.example")?;
}
read_bytes(&state, "pds.example", 30, LIMIT as usize - 1)?;
assert_eq!(
request_rewalk(&state.db, "pds.example", WINDOW - 1, WINDOW, LIMIT)?,
Rewalk::Queued,
"one byte is left"
);
finish_walk(&state, "pds.example")?;
read_bytes(&state, "pds.example", 40, 1)?;
assert_eq!(
request_rewalk(&state.db, "pds.example", 29 + WINDOW, WINDOW, LIMIT)?,
Rewalk::OutOfAllowance
);
assert_eq!(
request_rewalk(&state.db, "other.example", WINDOW - 1, WINDOW, LIMIT)?,
Rewalk::Queued,
"another host's allowance is its own"
);
assert_eq!(
request_rewalk(&state.db, "pds.example", 30 + WINDOW, WINDOW, LIMIT)?,
Rewalk::Queued
);
Ok(())
}
#[test]
fn the_spent_bytes_and_a_queued_rewalk_survive_a_restart() -> Result<()> {
let dir = tempfile::tempdir().into_diagnostic()?;
let config = Config {
database_path: dir.path().to_path_buf(),
..Default::default()
};
let state = AppState::new(&config)?;
assert_eq!(
request_rewalk(&state.db, "pds.example", 0, WINDOW, LIMIT)?,
Rewalk::Queued
);
read_bytes(&state, "pds.example", 0, LIMIT as usize)?;
state.db.persist()?;
drop(state);
let state = AppState::new(&config)?;
assert_eq!(
request_rewalk(&state.db, "pds.example", 1, WINDOW, LIMIT)?,
Rewalk::AlreadyWalking
);
assert!(
state
.db
.cursors
.contains_key(keys::rewalk_key("pds.example"))
.into_diagnostic()?
);
finish_walk(&state, "pds.example")?;
assert_eq!(
request_rewalk(&state.db, "pds.example", 1, WINDOW, LIMIT)?,
Rewalk::OutOfAllowance
);
let budget = Budget::load(&state.db, "pds.example", LIMIT, WINDOW)?;
assert_eq!(budget.spent_until(1), Some(WINDOW));
assert_eq!(budget.spent_until(WINDOW), None);
Ok(())
}
}