query cache: Mem limit, turn on by default

Author: leoyvens

core/src/subgraph/instance_manager.rs

@@ -21,10 +21,10 @@ use super::SubgraphInstance;
 lazy_static! {
     /// Size limit of the entity LFU cache, in bytes.
     // Multiplied by 1000 because the env var is in KB.
-    pub static ref ENTITY_CACHE_SIZE: u64 = 1000
+    pub static ref ENTITY_CACHE_SIZE: usize = 1000
         * std::env::var("GRAPH_ENTITY_CACHE_SIZE")
             .unwrap_or("10000".into())
-            .parse::<u64>()
+            .parse::<usize>()
             .expect("invalid GRAPH_ENTITY_CACHE_SIZE");
 }
 

graph/src/data/store/mod.rs

@@ -1,6 +1,5 @@
 use crate::data::subgraph::SubgraphDeploymentId;
 use crate::prelude::{format_err, EntityKey, QueryExecutionError};
-use crate::util::lfu_cache::CacheWeight;
 use failure::Error;
 use graphql_parser::query;
 use graphql_parser::schema;
@@ -596,32 +595,6 @@ pub trait ToEntityKey {
     fn to_entity_key(&self, subgraph: SubgraphDeploymentId) -> EntityKey;
 }
 
-impl CacheWeight for Value {
-    fn weight(&self) -> u64 {
-        use std::mem::size_of_val;
-        size_of_val(self) as u64
-            + match self {
-                Value::String(s) => s.len() as u64,
-                Value::BigDecimal(d) => (d.digits() as f32).log2() as u64,
-                Value::List(values) => values.iter().map(|value| value.weight()).sum(),
-                Value::Bytes(bytes) => bytes.as_slice().len() as u64,
-                Value::BigInt(n) => n.bits() / 8 as u64,
-                Value::Int(_) | Value::Bool(_) | Value::Null => 0,
-            }
-    }
-}
-
-impl CacheWeight for Entity {
-    /// The weight of an entity in the cache is the approximate amount of bytes occupied by it.
-    fn weight(&self) -> u64 {
-        use std::mem::size_of_val;
-        self.0
-            .iter()
-            .map(|(key, value)| size_of_val(key) as u64 + key.len() as u64 + value.weight())
-            .sum()
-    }
-}
-
 #[test]
 fn value_bytes() {
     let graphql_value = query::Value::String("0x8f494c66afc1d3f8ac1b45df21f02a46".to_owned());

graph/src/data/store/scalar.rs

@@ -320,8 +320,8 @@ impl BigInt {
         BigInt(self.0.pow(&exponent))
     }
 
-    pub fn bits(&self) -> u64 {
-        self.0.bits() as u64
+    pub fn bits(&self) -> usize {
+        self.0.bits()
     }
 }
 

graph/src/lib.rs

@@ -149,6 +149,7 @@ pub mod prelude {
         ComponentLoggerConfig, ElasticComponentLoggerConfig, LoggerFactory,
     };
     pub use crate::log::split::split_logger;
+    pub use crate::util::cache_weight::CacheWeight;
     pub use crate::util::error::CompatErr;
     pub use crate::util::futures::{retry, TimeoutError};
     pub use crate::util::stats::MovingStats;

graph/src/util/cache_weight.rs

@@ -0,0 +1,95 @@
+use crate::prelude::{BigDecimal, BigInt, Entity, Value};
+
+/// Estimate of how much memory a value consumes.
+/// Useful for measuring the size of caches.
+pub trait CacheWeight {
+    /// Total weight of the value.
+    fn weight(&self) -> usize {
+        std::mem::size_of_val(&self) + self.indirect_weight()
+    }
+
+    /// The weight of values pointed to by this value but logically owned by it, which is not
+    /// accounted for by `size_of`.
+    fn indirect_weight(&self) -> usize;
+}
+
+impl<T: CacheWeight> CacheWeight for Option<T> {
+    fn indirect_weight(&self) -> usize {
+        match self {
+            Some(x) => x.indirect_weight(),
+            None => 0,
+        }
+    }
+}
+
+impl<T: CacheWeight> CacheWeight for Vec<T> {
+    fn indirect_weight(&self) -> usize {
+        self.iter().map(CacheWeight::indirect_weight).sum()
+    }
+}
+
+impl<T: CacheWeight, U: CacheWeight> CacheWeight for std::collections::BTreeMap<T, U> {
+    fn indirect_weight(&self) -> usize {
+        self.iter()
+            .map(|(key, value)| key.weight() + value.weight())
+            .sum()
+    }
+}
+
+impl CacheWeight for &'_ [u8] {
+    fn indirect_weight(&self) -> usize {
+        self.len()
+    }
+}
+
+impl CacheWeight for String {
+    fn indirect_weight(&self) -> usize {
+        self.len()
+    }
+}
+
+impl CacheWeight for BigDecimal {
+    fn indirect_weight(&self) -> usize {
+        (self.digits() as f32).log2() as usize
+    }
+}
+
+impl CacheWeight for BigInt {
+    fn indirect_weight(&self) -> usize {
+        self.bits() / 8
+    }
+}
+
+impl CacheWeight for Value {
+    fn indirect_weight(&self) -> usize {
+        match self {
+            Value::String(s) => s.indirect_weight(),
+            Value::BigDecimal(d) => d.indirect_weight(),
+            Value::List(values) => values.indirect_weight(),
+            Value::Bytes(bytes) => bytes.as_slice().indirect_weight(),
+            Value::BigInt(n) => n.indirect_weight(),
+            Value::Int(_) | Value::Bool(_) | Value::Null => 0,
+        }
+    }
+}
+
+impl CacheWeight for Entity {
+    fn indirect_weight(&self) -> usize {
+        self.iter()
+            .map(|(key, value)| key.weight() + value.weight())
+            .sum()
+    }
+}
+
+impl CacheWeight for graphql_parser::query::Value {
+    fn indirect_weight(&self) -> usize {
+        use graphql_parser::query as q;
+
+        match self {
+            q::Value::Boolean(_) | q::Value::Int(_) | q::Value::Null | q::Value::Float(_) => 0,
+            q::Value::Enum(s) | q::Value::String(s) | q::Value::Variable(s) => s.indirect_weight(),
+            q::Value::List(l) => l.indirect_weight(),
+            q::Value::Object(o) => o.indirect_weight(),
+        }
+    }
+}

graph/src/util/lfu_cache.rs

@@ -1,3 +1,4 @@
+use crate::prelude::CacheWeight;
 use priority_queue::PriorityQueue;
 use std::cmp::Reverse;
 use std::fmt::Debug;
@@ -6,23 +7,10 @@ use std::hash::{Hash, Hasher};
 // The number of `evict` calls without access after which an entry is considered stale.
 const STALE_PERIOD: u64 = 100;
 
-pub trait CacheWeight {
-    fn weight(&self) -> u64;
-}
-
-impl<T: CacheWeight> CacheWeight for Option<T> {
-    fn weight(&self) -> u64 {
-        match self {
-            Some(x) => x.weight(),
-            None => 0,
-        }
-    }
-}
-
 /// `PartialEq` and `Hash` are delegated to the `key`.
 #[derive(Clone, Debug)]
 pub struct CacheEntry<K, V> {
-    weight: u64,
+    weight: usize,
     key: K,
     value: V,
     will_stale: bool,
@@ -66,7 +54,7 @@ type Priority = (bool, Reverse<u64>);
 #[derive(Clone, Debug)]
 pub struct LfuCache<K: Eq + Hash, V> {
     queue: PriorityQueue<CacheEntry<K, V>, Priority>,
-    total_weight: u64,
+    total_weight: usize,
     stale_counter: u64,
 }
 
@@ -157,7 +145,7 @@ impl<K: Clone + Ord + Eq + Hash + Debug, V: CacheWeight + Default> LfuCache<K, V
         self.queue.len()
     }
 
-    pub fn evict(&mut self, max_weight: u64) {
+    pub fn evict(&mut self, max_weight: usize) {
         if self.total_weight <= max_weight {
             return;
         }
@@ -204,18 +192,25 @@ impl<K: Ord + Eq + Hash, V> Extend<(CacheEntry<K, V>, Priority)> for LfuCache<K,
 
 #[test]
 fn entity_lru_cache() {
-    impl CacheWeight for u64 {
-        fn weight(&self) -> u64 {
-            *self
+    #[derive(Default, Debug, PartialEq, Eq)]
+    struct Weight(usize);
+
+    impl CacheWeight for Weight {
+        fn weight(&self) -> usize {
+            self.0
+        }
+
+        fn indirect_weight(&self) -> usize {
+            0
         }
     }
 
-    let mut cache: LfuCache<&'static str, u64> = LfuCache::new();
-    cache.insert("panda", 2);
-    cache.insert("cow", 1);
+    let mut cache: LfuCache<&'static str, Weight> = LfuCache::new();
+    cache.insert("panda", Weight(2));
+    cache.insert("cow", Weight(1));
 
-    assert_eq!(cache.get(&"cow"), Some(&1));
-    assert_eq!(cache.get(&"panda"), Some(&2));
+    assert_eq!(cache.get(&"cow"), Some(&Weight(1)));
+    assert_eq!(cache.get(&"panda"), Some(&Weight(2)));
 
     // Total weight is 3, nothing is evicted.
     cache.evict(3);
@@ -226,28 +221,28 @@ fn entity_lru_cache() {
     cache.evict(2);
     assert!(cache.get(&"panda").is_none());
 
-    cache.insert("alligator", 2);
+    cache.insert("alligator", Weight(2));
 
     // Give "cow" and "alligator" a high frequency.
     for _ in 0..1000 {
         cache.get(&"cow");
         cache.get(&"alligator");
     }
 
-    cache.insert("lion", 3);
+    cache.insert("lion", Weight(3));
 
     // Make "cow" and "alligator" stale and remove them.
     for _ in 0..(2 * STALE_PERIOD) {
         cache.get(&"lion");
 
         // The "whale" is something to evict so the stale counter moves.
-        cache.insert("whale", 100);
+        cache.insert("whale", Weight(100));
         cache.evict(10);
     }
 
     // Either "cow" and "alligator" fit in the cache, or just "lion".
     // "lion" will be kept, it had lower frequency but was not stale.
     assert!(cache.get(&"cow").is_none());
     assert!(cache.get(&"alligator").is_none());
-    assert_eq!(cache.get(&"lion"), Some(&3));
+    assert_eq!(cache.get(&"lion"), Some(&Weight(3)));
 }

graph/src/util/mod.rs

@@ -12,3 +12,5 @@ pub mod lfu_cache;
 pub mod error;
 
 pub mod stats;
+
+pub mod cache_weight;