Add cache clearing to stop the memory leak (#1106)

* Add cache clearing

* Add TODO comment

---------

Co-authored-by: Keavon Chambers <keavon@keavon.com>
This commit is contained in:
0HyperCube 2023-04-09 23:30:57 +01:00 committed by Keavon Chambers
parent 9c4164291c
commit a58d51d685
4 changed files with 51 additions and 18 deletions

View file

@ -34,7 +34,12 @@ where
Box::new(self.node.eval(*input))
}
}
fn reset(self: std::pin::Pin<&mut Self>) {
let wrapped_node = unsafe { self.map_unchecked_mut(|e| &mut e.node) };
Node::reset(wrapped_node);
}
}
impl<_I, _O, S0> DynAnyRefNode<_I, _O, S0> {
pub const fn new(node: S0) -> Self {
Self { node, _i: core::marker::PhantomData }

View file

@ -2,6 +2,8 @@ use graphene_core::Node;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use xxhash_rust::xxh3::Xxh3;
/// Caches the output of a given Node and acts as a proxy
@ -9,7 +11,7 @@ use xxhash_rust::xxh3::Xxh3;
pub struct CacheNode<T, CachedNode> {
// We have to use an append only data structure to make sure the references
// to the cache entries are always valid
cache: boxcar::Vec<(u64, T)>,
cache: boxcar::Vec<(u64, T, AtomicBool)>,
node: CachedNode,
}
impl<'i, T: 'i, I: 'i + Hash, CachedNode: 'i> Node<'i, I> for CacheNode<T, CachedNode>
@ -22,17 +24,25 @@ where
input.hash(&mut hasher);
let hash = hasher.finish();
if let Some((_, cached_value)) = self.cache.iter().find(|(h, _)| *h == hash) {
if let Some((_, cached_value, keep)) = self.cache.iter().find(|(h, _, _)| *h == hash) {
keep.store(true, std::sync::atomic::Ordering::Relaxed);
return cached_value;
} else {
trace!("Cache miss");
let output = self.node.eval(input);
let index = self.cache.push((hash, output));
let index = self.cache.push((hash, output, AtomicBool::new(true)));
return &self.cache[index].1;
}
}
fn reset(mut self: Pin<&mut Self>) {
let old_cache = std::mem::take(&mut self.cache);
self.cache = old_cache.into_iter().filter(|(_, _, keep)| keep.swap(false, std::sync::atomic::Ordering::Relaxed)).collect();
}
}
impl<T, CachedNode> std::marker::Unpin for CacheNode<T, CachedNode> {}
impl<T, CachedNode> CacheNode<T, CachedNode> {
pub fn new(node: CachedNode) -> CacheNode<T, CachedNode> {
CacheNode { cache: boxcar::Vec::new(), node }
@ -72,8 +82,16 @@ impl<'i, T: 'i + Hash> Node<'i, Option<T>> for LetNode<T> {
None => &self.cache.iter().last().expect("Let node was not initialized").1,
}
}
fn reset(mut self: Pin<&mut Self>) {
if let Some(last) = std::mem::take(&mut self.cache).into_iter().last() {
self.cache = boxcar::vec![last];
}
}
}
impl<T> std::marker::Unpin for LetNode<T> {}
impl<T> LetNode<T> {
pub fn new() -> LetNode<T> {
LetNode { cache: boxcar::Vec::new() }