//! how much memory whole-car backfills may hold at once. a whole car keeps every block until its
//! last step is in, so its peak follows the car's size, and a few big repos fetched together can
//! take far more than the concurrency limit suggests.
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};
use tokio::sync::Notify;
/// peak memory above baseline per byte of car, the larger of the bench repos measured with the
/// perf-backfill probe: jay.bsky.team (15 MB) peaked at about 6.3x its car and pfrazee.com
/// (83 MB) at about 4.5x, so small repos carry more fixed cost per byte
const PEAK_PER_CAR_BYTE: f64 = 6.3;
/// what a whole-car backfill of `car_bytes` is expected to hold at its peak
pub(crate) fn peak_for(car_bytes: usize) -> usize {
(car_bytes as f64 * PEAK_PER_CAR_BYTE) as usize
}
/// memory per block a spilled car still holds: its 24 byte index entry, and its share of the
/// mst walk's visited set, which keeps a cid for each node (about a quarter of the blocks)
const SPILLED_PER_BLOCK: usize = 64;
/// the bench repos average 260 bytes a block
const BYTES_PER_BLOCK: usize = 256;
/// one step in flight and the page of stored records the merge holds, with room to spare
const SPILLED_STEP: usize = 16 * 1024 * 1024;
/// what a backfill whose car of `car_bytes` spilled to disk holds at its peak: the index and
/// one step, not the car
pub(crate) fn spilled_peak_for(car_bytes: usize) -> usize {
car_bytes / BYTES_PER_BLOCK * SPILLED_PER_BLOCK + SPILLED_STEP
}
#[derive(Clone)]
pub(crate) struct BackfillBudget {
inner: Arc<Inner>,
}
struct Inner {
limit: usize,
state: Mutex<State>,
freed: Notify,
}
#[derive(Default)]
struct State {
used: usize,
next_ticket: u64,
/// bytes held by each reservation, by when it was made
holders: BTreeMap<u64, usize>,
}
impl BackfillBudget {
pub(crate) fn new(limit: usize) -> Self {
Self {
inner: Arc::new(Inner {
limit,
state: Mutex::default(),
freed: Notify::new(),
}),
}
}
/// an empty reservation, placed behind every one made before it
pub(crate) fn reservation(&self) -> Reservation {
let mut state = self.inner.state.lock().unwrap();
let ticket = state.next_ticket;
state.next_ticket += 1;
state.holders.insert(ticket, 0);
Reservation {
inner: self.inner.clone(),
ticket,
bytes: 0,
}
}
#[cfg(test)]
pub(crate) fn used(&self) -> usize {
self.inner.state.lock().unwrap().used
}
}
/// a share of the budget, given back when dropped, so a backfill that fails, times out or is
/// cancelled can't leak it
pub(crate) struct Reservation {
inner: Arc<Inner>,
ticket: u64,
bytes: usize,
}
impl Reservation {
/// waits until this reservation holds at least `bytes`. the oldest reservation never
/// waits and may go over the limit: reservations grow as their cars arrive, so two big
/// ones could otherwise each hold part of the budget and wait on each other forever
pub(crate) async fn grow_to(&mut self, bytes: usize) {
if bytes <= self.bytes {
return;
}
let more = bytes - self.bytes;
loop {
let freed = self.inner.freed.notified();
tokio::pin!(freed);
// registered before checking, so a release between the check and the wait isn't
// missed
freed.as_mut().enable();
{
let mut state = self.inner.state.lock().unwrap();
let oldest = state.holders.keys().next() == Some(&self.ticket);
if oldest || state.used + more <= self.inner.limit {
state.used += more;
*state.holders.get_mut(&self.ticket).unwrap() += more;
self.bytes = bytes;
return;
}
}
freed.await;
}
}
/// holds `bytes` from now on without waiting, which may go over the limit. a spilled car
/// grows by its index alone, and a backfill holding a spill slot must never wait on the
/// budget, since one waiting for that slot can hold the budget it'd need
pub(crate) fn resize_now(&mut self, bytes: usize) {
let mut state = self.inner.state.lock().unwrap();
state.used = state.used - self.bytes + bytes;
*state.holders.get_mut(&self.ticket).unwrap() = bytes;
let shrank = bytes < self.bytes;
self.bytes = bytes;
drop(state);
if shrank {
self.inner.freed.notify_waiters();
}
}
}
impl std::fmt::Debug for Reservation {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Reservation")
.field("bytes", &self.bytes)
.finish()
}
}
impl Drop for Reservation {
fn drop(&mut self) {
let mut state = self.inner.state.lock().unwrap();
state.holders.remove(&self.ticket);
state.used -= self.bytes;
drop(state);
self.inner.freed.notify_waiters();
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::FutureExt;
use std::time::Duration;
#[tokio::test]
async fn small_backfills_never_wait_on_the_default_budget() {
let budget = BackfillBudget::new(crate::config::Config::default().backfill_memory_bytes);
let mut held = Vec::new();
for _ in 0..64 {
let mut reservation = budget.reservation();
// a 1 MB car, bigger than most repos
reservation
.grow_to(peak_for(1024 * 1024))
.now_or_never()
.expect("a small backfill waited for the budget");
held.push(reservation);
}
drop(held);
assert_eq!(budget.used(), 0);
}
#[tokio::test]
async fn two_growing_backfills_bigger_than_the_budget_both_finish() {
let budget = BackfillBudget::new(100);
let grow = |budget: BackfillBudget| async move {
let mut reservation = budget.reservation();
for step in 1..=10 {
reservation.grow_to(step * 30).await;
tokio::task::yield_now().await;
}
reservation.bytes
};
let both = futures::future::join(grow(budget.clone()), grow(budget.clone()));
let done = tokio::time::timeout(Duration::from_secs(5), both)
.await
.expect("two growing reservations waited on each other");
assert_eq!(done, (300, 300));
assert_eq!(budget.used(), 0);
}
#[tokio::test]
async fn a_newer_reservation_waits_until_an_older_one_lets_go() {
let budget = BackfillBudget::new(100);
let mut older = budget.reservation();
older.grow_to(150).await;
let mut newer = budget.reservation();
assert!(newer.grow_to(10).now_or_never().is_none());
drop(older);
newer
.grow_to(10)
.now_or_never()
.expect("freed budget was handed on");
assert_eq!(budget.used(), 10);
}
#[tokio::test]
async fn a_cancelled_wait_takes_nothing() {
let budget = BackfillBudget::new(100);
let mut older = budget.reservation();
older.grow_to(100).await;
let mut newer = budget.reservation();
assert!(newer.grow_to(50).now_or_never().is_none());
drop(newer);
drop(older);
assert_eq!(budget.used(), 0);
}
}