mirror of
https://github.com/salsa-rs/salsa.git
synced 2024-12-25 05:29:43 +00:00
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:
parent
f9468e2ac4
commit
3d89c0d817
5 changed files with 143 additions and 1 deletions
|
@ -17,6 +17,7 @@ indexmap = "1.0.1"
|
||||||
log = "0.4.5"
|
log = "0.4.5"
|
||||||
smallvec = "0.6.5"
|
smallvec = "0.6.5"
|
||||||
salsa-macros = { version = "0.12.1", path = "components/salsa-macros" }
|
salsa-macros = { version = "0.12.1", path = "components/salsa-macros" }
|
||||||
|
linked-hash-map = "0.5.2"
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
diff = "0.1.0"
|
diff = "0.1.0"
|
||||||
|
|
|
@ -1,6 +1,7 @@
|
||||||
use crate::debug::TableEntry;
|
use crate::debug::TableEntry;
|
||||||
use crate::plumbing::CycleDetected;
|
use crate::plumbing::CycleDetected;
|
||||||
use crate::plumbing::DatabaseKey;
|
use crate::plumbing::DatabaseKey;
|
||||||
|
use crate::plumbing::LruQueryStorageOps;
|
||||||
use crate::plumbing::QueryFunction;
|
use crate::plumbing::QueryFunction;
|
||||||
use crate::plumbing::QueryStorageMassOps;
|
use crate::plumbing::QueryStorageMassOps;
|
||||||
use crate::plumbing::QueryStorageOps;
|
use crate::plumbing::QueryStorageOps;
|
||||||
|
@ -11,13 +12,16 @@ use crate::runtime::Runtime;
|
||||||
use crate::runtime::RuntimeId;
|
use crate::runtime::RuntimeId;
|
||||||
use crate::runtime::StampedValue;
|
use crate::runtime::StampedValue;
|
||||||
use crate::{Database, DiscardIf, DiscardWhat, Event, EventKind, SweepStrategy};
|
use crate::{Database, DiscardIf, DiscardWhat, Event, EventKind, SweepStrategy};
|
||||||
|
use linked_hash_map::LinkedHashMap;
|
||||||
use log::{debug, info};
|
use log::{debug, info};
|
||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
use parking_lot::RwLock;
|
use parking_lot::RwLock;
|
||||||
use rustc_hash::FxHashMap;
|
use rustc_hash::{FxHashMap, FxHasher};
|
||||||
use smallvec::SmallVec;
|
use smallvec::SmallVec;
|
||||||
|
use std::hash::BuildHasherDefault;
|
||||||
use std::marker::PhantomData;
|
use std::marker::PhantomData;
|
||||||
use std::ops::Deref;
|
use std::ops::Deref;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::sync::mpsc::{self, Receiver, Sender};
|
use std::sync::mpsc::{self, Receiver, Sender};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
@ -36,6 +40,8 @@ pub type DependencyStorage<DB, Q> = DerivedStorage<DB, Q, NeverMemoizeValue>;
|
||||||
/// storage requirements.
|
/// storage requirements.
|
||||||
pub type VolatileStorage<DB, Q> = DerivedStorage<DB, Q, VolatileValue>;
|
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
|
/// Handles storage where the value is 'derived' by executing a
|
||||||
/// function (in contrast to "inputs").
|
/// function (in contrast to "inputs").
|
||||||
pub struct DerivedStorage<DB, Q, MP>
|
pub struct DerivedStorage<DB, Q, MP>
|
||||||
|
@ -44,6 +50,10 @@ where
|
||||||
DB: Database,
|
DB: Database,
|
||||||
MP: MemoizationPolicy<DB, Q>,
|
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>>>,
|
map: RwLock<FxHashMap<Q::Key, QueryState<DB, Q>>>,
|
||||||
policy: PhantomData<MP>,
|
policy: PhantomData<MP>,
|
||||||
}
|
}
|
||||||
|
@ -237,6 +247,8 @@ where
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
DerivedStorage {
|
DerivedStorage {
|
||||||
map: RwLock::new(FxHashMap::default()),
|
map: RwLock::new(FxHashMap::default()),
|
||||||
|
lru_cap: AtomicUsize::new(0),
|
||||||
|
lru_keys: Mutex::new(LinkedHashSet::with_hasher(Default::default())),
|
||||||
policy: PhantomData,
|
policy: PhantomData,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -702,6 +714,19 @@ where
|
||||||
) -> Result<Q::Value, CycleDetected> {
|
) -> Result<Q::Value, CycleDetected> {
|
||||||
let StampedValue { value, changed_at } = self.read(db, key, &database_key)?;
|
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()
|
db.salsa_runtime()
|
||||||
.report_query_read(database_key, changed_at);
|
.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>
|
impl<DB, Q> Memo<DB, Q>
|
||||||
where
|
where
|
||||||
Q: QueryFunction<DB>,
|
Q: QueryFunction<DB>,
|
||||||
|
|
16
src/lib.rs
16
src/lib.rs
|
@ -24,6 +24,7 @@ use crate::plumbing::CycleDetected;
|
||||||
use crate::plumbing::InputQueryStorageOps;
|
use crate::plumbing::InputQueryStorageOps;
|
||||||
use crate::plumbing::QueryStorageMassOps;
|
use crate::plumbing::QueryStorageMassOps;
|
||||||
use crate::plumbing::QueryStorageOps;
|
use crate::plumbing::QueryStorageOps;
|
||||||
|
use crate::plumbing::LruQueryStorageOps;
|
||||||
use derive_new::new;
|
use derive_new::new;
|
||||||
use std::fmt::{self, Debug};
|
use std::fmt::{self, Debug};
|
||||||
use std::hash::Hash;
|
use std::hash::Hash;
|
||||||
|
@ -534,6 +535,21 @@ where
|
||||||
self.storage
|
self.storage
|
||||||
.set_constant(self.db, &key, &self.database_key(&key), value);
|
.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.
|
// Re-export the procedural macros.
|
||||||
|
|
|
@ -203,3 +203,14 @@ where
|
||||||
new_value: Q::Value,
|
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
68
tests/lru.rs
Normal 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);
|
||||||
|
}
|
Loading…
Reference in a new issue