use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock, Weak};
use std::time::Duration;
use super::group;
use super::track::{self, TrackState};
pub(crate) const ENTRY_OVERHEAD: u64 = group::CACHE_OVERHEAD + track::CACHE_OVERHEAD;
const WRITE_BOOST: u64 = 1;
const READ_BOOST: u64 = 2;
const ACCESS_SHIFT: u32 = 2;
const TICK_MS: u64 = 100;
pub const DEFAULT_EXPIRY: Duration = Duration::from_secs(30);
#[derive(Clone, Debug, Default)]
pub struct Config {
capacity: Option<u64>,
expiry: Option<Duration>,
}
impl Config {
pub fn with_capacity(mut self, capacity: impl Into<Option<u64>>) -> Self {
self.capacity = capacity.into();
self
}
pub fn with_expiry(mut self, expiry: impl Into<Option<Duration>>) -> Self {
self.expiry = expiry.into();
self
}
}
#[derive(Clone)]
pub struct Pool {
inner: Arc<Inner>,
}
impl Default for Pool {
fn default() -> Self {
Self::unbounded()
}
}
struct Inner {
used: AtomicU64,
capacity: AtomicU64,
expiry: u64,
clock: Mutex<Option<Clock>>,
tick: AtomicU64,
access_sum: AtomicU64,
access_count: AtomicU64,
tracks: kio::Lock<slab::Slab<Weak<Track>>>,
}
struct Clock {
epoch: crate::time::Instant,
now: crate::time::Instant,
sweep: Option<crate::time::Instant>,
}
impl Pool {
pub fn new(config: Config) -> Self {
let expiry = config.expiry.map_or(u64::MAX, |expiry| {
let ms = u64::try_from(expiry.as_millis()).unwrap_or(u64::MAX);
if ms == u64::MAX {
return u64::MAX;
}
ms.max(1).div_ceil(TICK_MS).saturating_mul(TICK_MS)
});
let pool = Self {
inner: Arc::new(Inner {
used: AtomicU64::new(0),
capacity: AtomicU64::new(config.capacity.unwrap_or(u64::MAX)),
expiry,
clock: Mutex::new(None),
tick: AtomicU64::new(0),
access_sum: AtomicU64::new(0),
access_count: AtomicU64::new(0),
tracks: kio::Lock::new(slab::Slab::new()),
}),
};
#[cfg(test)]
crate::model::clock::register(&pool);
pool
}
pub fn unbounded() -> Self {
Self::new(Config::default())
}
pub fn capacity(&self) -> Option<u64> {
match self.inner.capacity.load(Ordering::Relaxed) {
u64::MAX => None,
capacity => Some(capacity),
}
}
pub fn used(&self) -> u64 {
self.inner.used.load(Ordering::Relaxed)
}
pub fn resize(&self, capacity: impl Into<Option<u64>>) {
let capacity = capacity.into().unwrap_or(u64::MAX);
self.inner.capacity.store(capacity, Ordering::Relaxed);
}
pub fn expiry(&self) -> Option<Duration> {
match self.inner.expiry {
u64::MAX => None,
ms => Some(Duration::from_millis(ms)),
}
}
pub(crate) fn expiry_ticks(&self) -> u64 {
match self.inner.expiry {
u64::MAX => u64::MAX,
ms => ms / TICK_MS,
}
}
pub(crate) fn sweep_interval(&self) -> Option<Duration> {
self.expiry()
.or_else(|| self.capacity().map(|_| DEFAULT_EXPIRY))
.map(|window| window / 2)
}
pub(crate) fn sweep(&self) {
let tracks: Vec<_> = self
.inner
.tracks
.lock()
.iter()
.filter_map(|(_, track)| track.upgrade())
.collect();
for track in tracks {
track.sweep();
}
}
pub fn gc(&self, now: crate::time::Instant) -> Option<crate::time::Instant> {
self.advance(now, true)
}
fn advance(&self, now: crate::time::Instant, sweep: bool) -> Option<crate::time::Instant> {
let mut clock = self.inner.clock.lock().unwrap();
let clock = clock.get_or_insert(Clock {
epoch: now,
now,
sweep: None,
});
let now = now.max(clock.now);
let tick = u64::try_from(now.duration_since(clock.epoch).as_millis() / u128::from(TICK_MS))
.expect("cache clock overflow");
self.inner.tick.store(tick, Ordering::Relaxed);
clock.now = now;
if self.sweep_interval().is_none() {
clock.sweep = None;
} else if sweep && clock.sweep.is_none_or(|at| at <= now) {
self.sweep();
clock.sweep = self.sweep_interval().and_then(|interval| now.checked_add(interval));
}
clock.sweep
}
#[cfg(test)]
pub(crate) fn advance_test(&self, now: crate::time::Instant) {
self.advance(now, false);
let tracks: Vec<_> = self
.inner
.tracks
.lock()
.iter()
.filter_map(|(_, track)| track.upgrade())
.collect();
for track in tracks {
if let Some(state) = track.state.upgrade() {
state.read().date_cache_accesses(self.now());
}
}
}
fn register(&self, track: &Arc<Track>) -> Option<usize> {
self.expiry()?;
Some(self.inner.tracks.lock().insert(Arc::downgrade(track)))
}
fn unregister(&self, key: usize) {
self.inner.tracks.lock().remove(key);
}
#[cfg(test)]
pub(crate) fn same_pool(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.inner, &other.inner)
}
pub fn downgrade(&self) -> PoolWeak {
PoolWeak {
inner: Arc::downgrade(&self.inner),
}
}
pub(crate) fn add(&self, n: u64) {
self.inner.used.fetch_add(n, Ordering::Relaxed);
}
pub(crate) fn sub(&self, n: u64) {
self.inner.used.fetch_sub(n, Ordering::Relaxed);
}
pub(crate) fn now(&self) -> u64 {
self.inner.tick.load(Ordering::Relaxed)
}
fn stamp(&self, boost: u64) -> u64 {
self.now().saturating_mul(1 << ACCESS_SHIFT).saturating_add(boost)
}
pub(crate) fn average(&self) -> Option<u64> {
let count = self.inner.access_count.load(Ordering::Relaxed);
if count == 0 {
return None;
}
Some(self.inner.access_sum.load(Ordering::Relaxed) / count)
}
pub(crate) fn access_insert(&self, ts: u64) {
self.inner.access_sum.fetch_add(ts, Ordering::Relaxed);
self.inner.access_count.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn access_remove(&self, ts: u64) {
self.inner.access_sum.fetch_sub(ts, Ordering::Relaxed);
self.inner.access_count.fetch_sub(1, Ordering::Relaxed);
}
pub(crate) fn access_refresh(&self, old: u64, new: u64) {
self.inner
.access_sum
.fetch_add(new.wrapping_sub(old), Ordering::Relaxed);
}
pub(crate) fn accrue(&self, written: u64) -> Option<u64> {
let used = self.inner.used.load(Ordering::Relaxed);
let capacity = self.inner.capacity.load(Ordering::Relaxed);
if used <= capacity {
return None;
}
let debt = written as u128 * used as u128 / capacity.max(1) as u128;
Some(u64::try_from(debt).unwrap_or(u64::MAX))
}
}
impl std::fmt::Debug for Pool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Pool")
.field("used", &self.used())
.field("capacity", &self.capacity())
.field("expiry", &self.expiry())
.finish()
}
}
#[derive(Clone)]
pub struct PoolWeak {
inner: std::sync::Weak<Inner>,
}
impl PoolWeak {
pub fn upgrade(&self) -> Option<Pool> {
self.inner.upgrade().map(|inner| Pool { inner })
}
}
impl std::fmt::Debug for PoolWeak {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.upgrade() {
Some(pool) => pool.fmt(f),
None => f.debug_struct("PoolWeak").finish_non_exhaustive(),
}
}
}
const WRITE_CHARGE_THRESHOLD: u64 = 256 * 1024;
const EXPIRY_SCAN_TICKS: u64 = 1000 / TICK_MS;
#[derive(Default)]
pub(crate) struct Track {
pool: Pool,
written: AtomicU64,
next_expiry: AtomicU64,
expiry_cursor: AtomicUsize,
state: kio::Weak<TrackState>,
sweep: OnceLock<usize>,
}
impl Track {
pub(crate) fn new(pool: Pool, state: kio::Weak<TrackState>) -> Arc<Self> {
let track = Arc::new(Self {
pool,
written: AtomicU64::new(0),
next_expiry: AtomicU64::new(0),
expiry_cursor: AtomicUsize::new(0),
state,
sweep: OnceLock::new(),
});
if let Some(key) = track.pool.register(&track) {
let _ = track.sweep.set(key);
}
track
}
pub(crate) fn pool(&self) -> &Pool {
&self.pool
}
pub(crate) fn charge(self: &Arc<Self>) -> Charge {
self.pool.add(ENTRY_OVERHEAD);
self.written.fetch_add(ENTRY_OVERHEAD, Ordering::Relaxed);
let access = Arc::new(Access::new(self.pool.stamp(0)));
Charge {
track: Some(self.clone()),
bytes: ENTRY_OVERHEAD,
access,
counted: false,
}
}
pub(crate) fn take_written(&self) -> u64 {
self.written.swap(0, Ordering::Relaxed)
}
pub(crate) fn settle(&self, now: Option<u64>) {
self.settle_inner(now, false);
}
pub(crate) fn sweep(&self) {
self.settle_inner(None, true);
}
fn settle_inner(&self, now: Option<u64>, full: bool) {
let settle_debt = self.written.load(Ordering::Relaxed) >= WRITE_CHARGE_THRESHOLD;
let scan_expiry = if full {
self.pool.expiry().is_some()
} else {
self.expiry_due(now)
};
if !settle_debt && !scan_expiry {
return;
}
let Some(state) = self.state.upgrade() else { return };
let expiry = if scan_expiry {
let state = state.read();
let scan = if full {
state.expiry_scan_drain()
} else {
state.expiry_scan()
};
state.expiry_mutation_due(scan).then_some(scan)
} else {
None
};
if !settle_debt && expiry.is_none() {
return;
}
if let Ok(mut state) = state.write() {
if settle_debt {
state.charge_debt();
}
if let Some(scan) = expiry {
state.evict_expired_scan(scan);
}
}
}
pub(crate) fn next_expiry_scan(&self, width: usize) -> usize {
self.expiry_cursor.fetch_add(width, Ordering::Relaxed)
}
fn expiry_due(&self, now: Option<u64>) -> bool {
let expiry = self.pool.expiry_ticks();
if expiry == u64::MAX {
return false;
}
let now = now.unwrap_or_else(|| self.pool.now());
let interval = expiry.clamp(1, EXPIRY_SCAN_TICKS);
let deadline = now.saturating_add(interval);
let next = self.next_expiry.load(Ordering::Relaxed);
if now < next {
return false;
}
self.next_expiry
.compare_exchange(next, deadline, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
}
}
impl Drop for Track {
fn drop(&mut self) {
if let Some(key) = self.sweep.get() {
self.pool.unregister(*key);
}
}
}
#[derive(Default)]
pub(crate) struct Charge {
track: Option<Arc<Track>>,
bytes: u64,
access: Arc<Access>,
counted: bool,
}
#[derive(Default)]
pub(crate) struct Access {
stamp: AtomicU64,
expires: AtomicU64,
}
impl Access {
fn new(stamp: u64) -> Self {
Self {
stamp: AtomicU64::new(stamp),
expires: AtomicU64::new(u64::MAX),
}
}
pub(crate) fn get(&self) -> u64 {
self.stamp.load(Ordering::Relaxed)
}
pub(crate) fn touch(&self) {
self.expires.store(u64::MAX, Ordering::Relaxed);
}
pub(crate) fn tick(&self, now: Option<u64>) -> Option<u64> {
let tick = match now {
Some(now) => match self
.expires
.compare_exchange(u64::MAX, now, Ordering::Relaxed, Ordering::Relaxed)
{
Ok(_) => now,
Err(tick) => tick,
},
None => self.expires.load(Ordering::Relaxed),
};
(tick != u64::MAX).then_some(tick)
}
fn bump(&self, target: u64) -> u64 {
self.stamp.fetch_max(target, Ordering::Relaxed)
}
}
impl Charge {
pub(crate) fn add(&mut self, n: u64) -> Option<u64> {
if let Some(track) = &self.track {
track.pool.add(n);
track.written.fetch_add(n, Ordering::Relaxed);
self.bytes += n;
}
self.touch(WRITE_BOOST)
}
pub(crate) fn size(&self) -> u64 {
self.bytes
}
pub(crate) fn access(&self) -> Arc<Access> {
self.access.clone()
}
pub(crate) fn accessed(&self) -> u64 {
self.access.get()
}
pub(crate) fn demote(&mut self) {
if let Some(track) = &self.track
&& !self.counted
{
track.pool.access_insert(self.accessed());
self.counted = true;
}
}
pub(crate) fn refresh(&self) {
self.touch(READ_BOOST);
}
pub(crate) fn record_write(&mut self) -> Option<u64> {
self.touch(WRITE_BOOST)
}
fn touch(&self, boost: u64) -> Option<u64> {
let track = self.track.as_ref()?;
self.access.touch();
let target = track.pool.stamp(boost);
let prev = self.access.bump(target);
if target > prev && self.counted {
track.pool.access_refresh(prev, target);
}
Some(target >> ACCESS_SHIFT)
}
pub(crate) fn clear(&mut self) {
if let Some(track) = &self.track {
track.pool.sub(self.bytes);
self.bytes = 0;
if self.counted {
track.pool.access_remove(self.accessed());
self.counted = false;
}
}
}
}
impl Drop for Charge {
fn drop(&mut self) {
self.clear();
}
}
#[cfg(test)]
mod test {
use super::*;
fn charge(pool: &Pool) -> Charge {
Track::new(pool.clone(), kio::Weak::new()).charge()
}
fn bounded(capacity: u64) -> Pool {
let config = Config::default().with_capacity(capacity).with_expiry(DEFAULT_EXPIRY);
Pool::new(config)
}
#[test]
fn unbounded_never_accrues() {
let pool = Pool::unbounded();
let mut charge = charge(&pool);
charge.add(1 << 40);
assert_eq!(pool.accrue(1 << 30), None);
assert_eq!(pool.used(), (1 << 40) + ENTRY_OVERHEAD);
drop(charge);
assert_eq!(pool.used(), 0);
}
#[test]
fn config_applies_capacity_and_expiry() {
let pool = bounded(1000);
assert_eq!(pool.capacity(), Some(1000));
assert_eq!(pool.expiry(), Some(DEFAULT_EXPIRY));
}
#[test]
fn weak_follows_the_last_strong_handle() {
let pool = bounded(1000);
let clone = pool.clone();
let weak = pool.downgrade();
drop(pool);
let upgraded = weak.upgrade().expect("a strong handle remains");
assert!(upgraded.same_pool(&clone));
drop(upgraded);
drop(clone);
assert!(weak.upgrade().is_none());
}
#[test]
fn accrue_none_under_capacity() {
let pool = bounded(ENTRY_OVERHEAD + 1000);
let mut charge = charge(&pool);
charge.add(500);
assert_eq!(pool.accrue(100), None);
}
#[test]
fn accrue_proportional_over_capacity() {
let pool = bounded(1000);
let mut charge = charge(&pool);
charge.add(2000 - ENTRY_OVERHEAD);
assert_eq!(pool.accrue(100), Some(200));
assert_eq!(pool.accrue(0), Some(0));
}
#[test]
fn average_tracks_evictable_population() {
let pool = bounded(1000);
assert_eq!(pool.average(), None);
pool.access_insert(10);
pool.access_insert(20);
assert_eq!(pool.average(), Some(15));
pool.access_refresh(10, 40);
assert_eq!(pool.average(), Some(30));
pool.access_remove(40);
assert_eq!(pool.average(), Some(20));
pool.access_remove(20);
assert_eq!(pool.average(), None);
}
#[test]
fn charge_raii() {
let pool = bounded(1000);
let mut charge = charge(&pool);
assert_eq!(pool.used(), ENTRY_OVERHEAD);
charge.add(100);
assert_eq!(pool.used(), ENTRY_OVERHEAD + 100);
charge.clear();
assert_eq!(pool.used(), 0);
charge.clear();
drop(charge);
assert_eq!(pool.used(), 0);
}
#[test]
fn detached_charge_is_noop() {
let mut charge = Charge::default();
charge.add(123);
charge.clear();
}
#[test]
fn accrue_saturates() {
let pool = bounded(1);
let mut c = charge(&pool);
c.add(1 << 40);
assert_eq!(pool.accrue(1 << 40), Some(u64::MAX));
}
#[test]
fn charge_counts_gross_writes() {
let track = Track::new(bounded(1000), kio::Weak::new());
let mut c = track.charge();
c.add(100);
assert_eq!(track.take_written(), ENTRY_OVERHEAD + 100);
assert_eq!(track.take_written(), 0, "taking it drains the counter");
}
#[test]
fn charge_owns_access_sample() {
let pool = bounded(1000);
let mut c = charge(&pool);
assert_eq!(pool.average(), None, "not evictable until demoted");
c.demote();
c.demote(); assert!(pool.average().is_some());
c.clear();
assert_eq!(pool.average(), None, "aborted groups leave no ghost sample");
drop(c);
assert_eq!(pool.average(), None);
}
#[test]
fn refresh_updates_a_counted_sample() {
let pool = bounded(1000);
let mut c = charge(&pool);
c.demote();
c.refresh();
assert_eq!(pool.average(), Some(c.accessed()));
c.clear();
assert_eq!(pool.average(), None);
}
#[test]
fn refresh_protects_within_a_tick() {
let pool = bounded(1000);
let mut c = charge(&pool);
c.demote();
let average = pool.average().unwrap();
c.refresh();
assert!(c.accessed() > average);
assert_eq!(c.access().tick(None), None, "undated access is protected until cleanup");
assert_eq!(
c.access().tick(Some(pool.now())),
Some(pool.now()),
"cleanup dates the access"
);
let stamped = c.accessed();
c.refresh();
assert_eq!(c.accessed(), stamped);
}
#[test]
fn expiry_config() {
let pool = Pool::unbounded();
assert_eq!(pool.expiry(), None);
let pool = Pool::new(Config::default().with_expiry(Duration::from_secs(1)));
assert_eq!(pool.expiry(), Some(Duration::from_secs(1)));
assert_eq!(pool.expiry_ticks(), 10);
let pool = Pool::new(Config::default().with_expiry(Duration::from_millis(1)));
assert_eq!(pool.expiry(), Some(Duration::from_millis(TICK_MS)));
assert_eq!(pool.expiry_ticks(), 1);
let pool = Pool::new(Config::default());
assert_eq!(pool.expiry(), None);
assert_eq!(pool.expiry_ticks(), u64::MAX);
}
#[test]
fn expiry_gate_reuses_a_supplied_tick() {
let pool = Pool::new(Config::default().with_expiry(Duration::from_secs(1)));
let track = Track::new(pool, kio::Weak::new());
assert!(track.expiry_due(Some(0)));
assert!(!track.expiry_due(Some(9)));
assert!(track.expiry_due(Some(10)));
assert!(!track.expiry_due(None));
}
#[test]
fn sweep_interval_is_half_the_window() {
assert_eq!(Pool::unbounded().sweep_interval(), None);
let pool = Pool::new(Config::default().with_expiry(Duration::from_secs(4)));
assert_eq!(pool.sweep_interval(), Some(Duration::from_secs(2)));
}
#[test]
fn the_sweep_registry_follows_account_lifetime() {
let pool = bounded(1000);
assert!(pool.inner.tracks.lock().is_empty());
let track = Track::new(pool.clone(), kio::Weak::new());
assert_eq!(pool.inner.tracks.lock().len(), 1);
pool.sweep();
drop(track);
assert!(pool.inner.tracks.lock().is_empty(), "a dropped account leaves no entry");
}
#[test]
fn an_inert_pool_registers_nothing() {
let pool = Pool::unbounded();
let track = Track::new(pool.clone(), kio::Weak::new());
assert!(pool.inner.tracks.lock().is_empty());
pool.sweep();
drop(track);
}
#[test]
fn expiry_gate_stays_closed_without_a_window() {
let track = Track::new(Pool::unbounded(), kio::Weak::new());
assert!(!track.expiry_due(None));
assert!(!track.expiry_due(Some(u64::MAX)));
}
#[test]
fn standalone_origin_enables_default_expiry() {
assert_eq!(crate::origin::Config::default().pool.expiry(), Some(DEFAULT_EXPIRY));
}
#[test]
fn collecting_before_the_deadline_does_not_postpone_it() {
let pool = Pool::new(Config::default().with_expiry(Duration::from_secs(2)));
let now = crate::model::clock::now();
let deadline = pool.gc(now);
assert_eq!(pool.gc(now + Duration::from_millis(500)), deadline);
}
#[test]
fn bounded_pools_sample_recency_without_expiration() {
let pool = Pool::unbounded();
let now = crate::model::clock::now();
assert_eq!(pool.gc(now), None);
pool.resize(1024);
assert_eq!(pool.gc(now), Some(now + DEFAULT_EXPIRY / 2));
pool.gc(now + DEFAULT_EXPIRY);
assert!(pool.now() > 0);
pool.resize(None);
assert_eq!(pool.gc(now + DEFAULT_EXPIRY), None);
}
#[test]
fn resize() {
let pool = Pool::unbounded();
assert_eq!(pool.capacity(), None);
let mut charge = charge(&pool);
charge.add(1000);
pool.resize(100);
assert_eq!(pool.capacity(), Some(100));
assert!(pool.used() > 100);
assert!(pool.accrue(50).unwrap() > 50);
pool.resize(None);
assert_eq!(pool.capacity(), None);
assert_eq!(pool.accrue(50), None);
}
}