Add LRU to derived storage

LRU allows to bound the maximum number of *values* that are present in
the table.
This commit is contained in:
Aleksey Kladov 2019-06-06 17:59:54 +03:00
parent f9468e2ac4
commit 3d89c0d817
5 changed files with 143 additions and 1 deletions

View file

@ -17,6 +17,7 @@ indexmap = "1.0.1"
log = "0.4.5"
smallvec = "0.6.5"
salsa-macros = { version = "0.12.1", path = "components/salsa-macros" }
linked-hash-map = "0.5.2"
[dev-dependencies]
diff = "0.1.0"

View file

@ -1,6 +1,7 @@
use crate::debug::TableEntry;
use crate::plumbing::CycleDetected;
use crate::plumbing::DatabaseKey;
use crate::plumbing::LruQueryStorageOps;
use crate::plumbing::QueryFunction;
use crate::plumbing::QueryStorageMassOps;
use crate::plumbing::QueryStorageOps;
@ -11,13 +12,16 @@ use crate::runtime::Runtime;
use crate::runtime::RuntimeId;
use crate::runtime::StampedValue;
use crate::{Database, DiscardIf, DiscardWhat, Event, EventKind, SweepStrategy};
use linked_hash_map::LinkedHashMap;
use log::{debug, info};
use parking_lot::Mutex;
use parking_lot::RwLock;
use rustc_hash::FxHashMap;
use rustc_hash::{FxHashMap, FxHasher};
use smallvec::SmallVec;
use std::hash::BuildHasherDefault;
use std::marker::PhantomData;
use std::ops::Deref;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::Arc;
@ -36,6 +40,8 @@ pub type DependencyStorage<DB, Q> = DerivedStorage<DB, Q, NeverMemoizeValue>;
/// storage requirements.
pub type VolatileStorage<DB, Q> = DerivedStorage<DB, Q, VolatileValue>;
type LinkedHashSet<T> = LinkedHashMap<T, (), BuildHasherDefault<FxHasher>>;
/// Handles storage where the value is 'derived' by executing a
/// function (in contrast to "inputs").
pub struct DerivedStorage<DB, Q, MP>
@ -44,6 +50,10 @@ where
DB: Database,
MP: MemoizationPolicy<DB, Q>,
{
lru_cap: AtomicUsize,
// if `lru_keys` and `map` are locked togeter,
// `lru_keys` is locked first, to prevent deadlocks.
lru_keys: Mutex<LinkedHashSet<Q::Key>>,
map: RwLock<FxHashMap<Q::Key, QueryState<DB, Q>>>,
policy: PhantomData<MP>,
}
@ -237,6 +247,8 @@ where
fn default() -> Self {
DerivedStorage {
map: RwLock::new(FxHashMap::default()),
lru_cap: AtomicUsize::new(0),
lru_keys: Mutex::new(LinkedHashSet::with_hasher(Default::default())),
policy: PhantomData,
}
}
@ -702,6 +714,19 @@ where
) -> Result<Q::Value, CycleDetected> {
let StampedValue { value, changed_at } = self.read(db, key, &database_key)?;
let lru_cap = self.lru_cap.load(Ordering::Relaxed);
if lru_cap > 0 {
let mut lru_keys = self.lru_keys.lock();
lru_keys.insert(key.clone(), ());
if lru_keys.len() > lru_cap {
if let Some((evicted, ())) = lru_keys.pop_front() {
if let Some(QueryState::Memoized(memo)) = self.map.write().get_mut(&evicted) {
memo.value = None;
}
}
}
}
db.salsa_runtime()
.report_query_read(database_key, changed_at);
@ -1011,6 +1036,27 @@ where
}
}
impl<DB, Q, MP> LruQueryStorageOps for DerivedStorage<DB, Q, MP>
where
Q: QueryFunction<DB>,
DB: Database,
MP: MemoizationPolicy<DB, Q>,
{
fn set_lru_capacity(&self, new_capacity: usize) {
let mut lru_keys = self.lru_keys.lock();
let mut map = self.map.write();
self.lru_cap.store(new_capacity, Ordering::SeqCst);
while lru_keys.len() > new_capacity {
let (evicted, ()) = lru_keys.pop_front().unwrap();
if let Some(QueryState::Memoized(memo)) = map.get_mut(&evicted) {
memo.value = None;
}
}
let additional_cap = new_capacity - lru_keys.len();
lru_keys.reserve(additional_cap);
}
}
impl<DB, Q> Memo<DB, Q>
where
Q: QueryFunction<DB>,

View file

@ -24,6 +24,7 @@ use crate::plumbing::CycleDetected;
use crate::plumbing::InputQueryStorageOps;
use crate::plumbing::QueryStorageMassOps;
use crate::plumbing::QueryStorageOps;
use crate::plumbing::LruQueryStorageOps;
use derive_new::new;
use std::fmt::{self, Debug};
use std::hash::Hash;
@ -534,6 +535,21 @@ where
self.storage
.set_constant(self.db, &key, &self.database_key(&key), value);
}
/// Sets the size of LRU cache of values for this query table.
///
/// That is, at most `cap` values will be preset in the table at the same
/// time. This helps with keeping maximum memory usage under control, at the
/// cost of potential extra recalculations of evicted values.
///
/// If `cap` is zero, all values are preserved, this is the default.
pub fn set_lru_capacity(&self, cap: usize)
where
Q::Storage: plumbing::LruQueryStorageOps,
{
self.storage
.set_lru_capacity(cap);
}
}
// Re-export the procedural macros.

View file

@ -203,3 +203,14 @@ where
new_value: Q::Value,
);
}
/// An optional trait that is implemented for "user mutable" storage:
/// that is, storage whose value is not derived from other storage but
/// is set independently.
pub trait LruQueryStorageOps: Default
{
fn set_lru_capacity(
&self,
new_capacity: usize,
);
}

68
tests/lru.rs Normal file
View file

@ -0,0 +1,68 @@
//! Test setting LRU actually limits the number of things in the database;
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use salsa::Database as _;
#[derive(Debug, PartialEq, Eq)]
struct HotPotato(u32);
static N_POTATOES: AtomicUsize = AtomicUsize::new(0);
impl HotPotato {
fn new(id: u32) -> HotPotato {
N_POTATOES.fetch_add(1, Ordering::SeqCst);
HotPotato(id)
}
}
impl Drop for HotPotato {
fn drop(&mut self) {
N_POTATOES.fetch_sub(1, Ordering::SeqCst);
}
}
#[salsa::query_group(QueryGroupStorage)]
trait QueryGroup {
fn get(&self, x: u32) -> Arc<HotPotato>;
}
fn get(_db: &impl QueryGroup, x: u32) -> Arc<HotPotato> {
Arc::new(HotPotato::new(x))
}
#[salsa::database(QueryGroupStorage)]
#[derive(Default)]
struct Database {
runtime: salsa::Runtime<Database>,
}
impl salsa::Database for Database {
fn salsa_runtime(&self) -> &salsa::Runtime<Database> {
&self.runtime
}
}
#[test]
fn lru_works() {
let mut db = Database::default();
let cap = 32;
db.query_mut(GetQuery).set_lru_capacity(32);
assert_eq!(N_POTATOES.load(Ordering::SeqCst), 0);
for i in 0..128u32 {
let p = db.get(i);
assert_eq!(p.0, i)
}
assert_eq!(N_POTATOES.load(Ordering::SeqCst), cap);
for i in 0..128u32 {
let p = db.get(i);
assert_eq!(p.0, i)
}
assert_eq!(N_POTATOES.load(Ordering::SeqCst), cap);
drop(db);
assert_eq!(N_POTATOES.load(Ordering::SeqCst), 0);
}