all: Move blocking of expensive queries into the LoadManager

Author: lutter

graph/src/data/graphql/effort.rs

@@ -66,17 +66,25 @@ impl QueryEffortInner {
 pub struct LoadManager {
     effort: QueryEffort,
     store_wait_stats: PoolWaitStats,
+    blocked_queries: Vec<u64>,
 }
 
 impl LoadManager {
-    pub fn new(store_wait_stats: PoolWaitStats) -> Self {
+    pub fn new(store_wait_stats: PoolWaitStats, blocked_queries: Vec<u64>) -> Self {
         Self {
             effort: QueryEffort::default(),
             store_wait_stats,
+            blocked_queries,
         }
     }
 
     pub fn add_query(&self, shape_hash: u64, duration: Duration) {
         self.effort.add(shape_hash, duration);
     }
+
+    /// Return `true` if we should decline running the query with this
+    /// `ShapeHash`
+    pub fn decline(&self, shape_hash: u64) -> bool {
+        self.blocked_queries.contains(&shape_hash)
+    }
 }

graphql/src/runner.rs

@@ -1,6 +1,6 @@
 use futures01::future;
 use graphql_parser::query as q;
-use std::collections::{BTreeMap, HashMap};
+use std::collections::BTreeMap;
 use std::env;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -25,7 +25,6 @@ use lazy_static::lazy_static;
 pub struct GraphQlRunner<S> {
     logger: Logger,
     store: Arc<S>,
-    expensive: HashMap<u64, Arc<q::Document>>,
     load_manager: Arc<LoadManager>,
 }
 
@@ -65,13 +64,12 @@ where
     ) -> Self {
         let expensive = expensive
             .into_iter()
-            .map(|doc| (shape_hash(&doc), doc.clone()))
-            .collect::<HashMap<_, _>>();
+            .map(|doc| shape_hash(&doc))
+            .collect::<Vec<_>>();
         GraphQlRunner {
             logger: logger.new(o!("component" => "GraphQlRunner")),
             store,
-            expensive,
-            load_manager: Arc::new(LoadManager::new(store_wait_stats)),
+            load_manager: Arc::new(LoadManager::new(store_wait_stats, expensive)),
         }
     }
 
@@ -156,7 +154,7 @@ where
     }
 
     pub fn check_too_expensive(&self, query: &Query) -> Result<(), Vec<QueryExecutionError>> {
-        if self.expensive.contains_key(&query.shape_hash) {
+        if self.load_manager.decline(query.shape_hash) {
             Err(vec![QueryExecutionError::TooExpensive])
         } else {
             Ok(())

store/test-store/src/lib.rs

@@ -39,7 +39,7 @@ lazy_static! {
 
     pub static ref POOL_WAIT_STATS: PoolWaitStats = Arc::new(RwLock::new(MovingStats::default()));
 
-    pub static ref LOAD_MANAGER: Arc<LoadManager> = Arc::new(LoadManager::new(POOL_WAIT_STATS.clone()));
+    pub static ref LOAD_MANAGER: Arc<LoadManager> = Arc::new(LoadManager::new(POOL_WAIT_STATS.clone(), vec![]));
 
     // Create Store instance once for use with each of the tests.
     pub static ref STORE: Arc<Store> = {