use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use super::track::TrackState;
const ENTRY_OVERHEAD: u64 = 256;
const WRITE_BOOST: u64 = 1;
const READ_BOOST: u64 = 2;
const TICK_MS: u64 = 100;
#[derive(Clone, Default)]
pub struct Pool {
inner: Arc<Inner>,
}
struct Inner {
used: AtomicU64,
capacity: AtomicU64,
epoch: web_async::time::Instant,
access_sum: AtomicU64,
access_count: AtomicU64,
}
impl Default for Inner {
fn default() -> Self {
Self {
used: AtomicU64::new(0),
capacity: AtomicU64::new(u64::MAX),
epoch: web_async::time::Instant::now(),
access_sum: AtomicU64::new(0),
access_count: AtomicU64::new(0),
}
}
}
impl Pool {
pub fn new(capacity: u64) -> Self {
let pool = Self::default();
pool.inner.capacity.store(capacity, Ordering::Relaxed);
pool
}
pub fn unbounded() -> Self {
Self::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 same_pool(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.inner, &other.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.epoch.elapsed().as_millis() as u64 / TICK_MS
}
pub(crate) fn ticks(duration: Duration) -> u64 {
u64::try_from(duration.as_millis() / TICK_MS as u128).unwrap_or(u64::MAX)
}
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())
.finish()
}
}
const WRITE_CHARGE_THRESHOLD: u64 = 256 * 1024;
#[derive(Default)]
pub(crate) struct Track {
pool: Pool,
written: AtomicU64,
state: kio::Weak<TrackState>,
}
impl Track {
pub(crate) fn new(pool: Pool, state: kio::Weak<TrackState>) -> Arc<Self> {
Arc::new(Self {
pool,
written: AtomicU64::new(0),
state,
})
}
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 last = self.pool.now();
Charge {
track: Some(self.clone()),
bytes: ENTRY_OVERHEAD,
last,
counted: false,
}
}
pub(crate) fn take_written(&self) -> u64 {
self.written.swap(0, Ordering::Relaxed)
}
pub(crate) fn settle(&self) {
if self.written.load(Ordering::Relaxed) < WRITE_CHARGE_THRESHOLD {
return;
}
let Some(state) = self.state.upgrade() else { return };
if let Ok(mut state) = state.write() {
state.charge_debt();
}
}
}
#[derive(Default)]
pub(crate) struct Charge {
track: Option<Arc<Track>>,
bytes: u64,
last: u64,
counted: bool,
}
impl Charge {
pub(crate) fn add(&mut self, n: 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 sub(&mut self, n: u64) {
if let Some(track) = &self.track {
track.pool.sub(n);
self.bytes = self.bytes.saturating_sub(n);
}
}
pub(crate) fn size(&self) -> u64 {
self.bytes
}
pub(crate) fn accessed(&self) -> u64 {
self.last
}
pub(crate) fn demote(&mut self) {
if let Some(track) = &self.track
&& !self.counted
{
track.pool.access_insert(self.last);
self.counted = true;
}
}
pub(crate) fn refresh(&mut self) {
self.touch(READ_BOOST);
}
fn touch(&mut self, boost: u64) {
let Some(track) = &self.track else { return };
let target = track.pool.now().saturating_add(boost);
if target <= self.last {
return;
}
if self.counted {
track.pool.access_refresh(self.last, target);
}
self.last = target;
}
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.last);
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()
}
#[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 accrue_none_under_capacity() {
let pool = Pool::new(1000);
let mut charge = charge(&pool);
charge.add(500);
assert_eq!(pool.accrue(100), None);
}
#[test]
fn accrue_proportional_over_capacity() {
let pool = Pool::new(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 = Pool::new(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 = Pool::new(1000);
let mut charge = charge(&pool);
assert_eq!(pool.used(), ENTRY_OVERHEAD);
charge.add(100);
assert_eq!(pool.used(), ENTRY_OVERHEAD + 100);
charge.sub(40);
assert_eq!(pool.used(), ENTRY_OVERHEAD + 60);
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.sub(23);
charge.clear();
}
#[test]
fn accrue_saturates() {
let pool = Pool::new(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(Pool::new(1000), kio::Weak::new());
let mut c = track.charge();
c.add(100);
c.sub(40); 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 = Pool::new(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_protects_within_a_tick() {
let pool = Pool::new(1000);
let mut c = charge(&pool);
c.demote();
let average = pool.average().unwrap();
c.refresh();
assert!(c.accessed() > average);
let stamped = c.accessed();
c.refresh();
assert_eq!(c.accessed(), stamped);
}
#[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);
}
}