all: Use a HashSet to store banned/jailed queries in the LoadManager

Author: lutter

graph/src/data/graphql/effort.rs

@@ -1,7 +1,7 @@
 //! Utilities to keep moving statistics about queries
 
 use rand::{prelude::Rng, thread_rng};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::env;
 use std::str::FromStr;
 use std::sync::{Arc, RwLock};
@@ -142,22 +142,22 @@ pub struct LoadManager {
     store_wait_stats: PoolWaitStats,
     /// List of query shapes that have been statically blocked through
     /// configuration
-    blocked_queries: Vec<u64>,
+    blocked_queries: HashSet<u64>,
     /// List of query shapes that have caused more than `JAIL_THRESHOLD`
     /// proportion of the work while the system was overloaded. Currently,
     /// there is no way for a query to get out of jail other than
     /// restarting the process
-    jailed_queries: RwLock<Vec<u64>>,
+    jailed_queries: RwLock<HashSet<u64>>,
     kill_state: RwLock<KillState>,
 }
 
 impl LoadManager {
-    pub fn new(store_wait_stats: PoolWaitStats, blocked_queries: Vec<u64>) -> Self {
+    pub fn new(store_wait_stats: PoolWaitStats, blocked_queries: HashSet<u64>) -> Self {
         Self {
             effort: QueryEffort::default(),
             store_wait_stats,
             blocked_queries,
-            jailed_queries: RwLock::new(Vec::new()),
+            jailed_queries: RwLock::new(HashSet::new()),
             kill_state: RwLock::new(KillState::new()),
         }
     }
@@ -250,7 +250,7 @@ impl LoadManager {
         if known_query && query_effort / total_effort > *JAIL_THRESHOLD {
             // Any single query that causes at least JAIL_THRESHOLD of the
             // effort in an overload situation gets killed
-            self.jailed_queries.write().unwrap().push(shape_hash);
+            self.jailed_queries.write().unwrap().insert(shape_hash);
             return true;
         }
 

graphql/src/runner.rs

@@ -1,6 +1,6 @@
 use futures01::future;
 use graphql_parser::query as q;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashSet};
 use std::env;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -65,7 +65,7 @@ where
         let expensive = expensive
             .into_iter()
             .map(|doc| shape_hash(&doc))
-            .collect::<Vec<_>>();
+            .collect::<HashSet<_>>();
         GraphQlRunner {
             logger: logger.new(o!("component" => "GraphQlRunner")),
             store,

store/test-store/src/lib.rs

@@ -14,6 +14,7 @@ use graph_store_postgres::{ChainHeadUpdateListener, Store, StoreConfig, Subscrip
 use graphql_parser::query as q;
 use hex_literal::hex;
 use lazy_static::lazy_static;
+use std::collections::HashSet;
 use std::env;
 use std::sync::{Mutex, RwLock};
 use std::time::Instant;
@@ -39,7 +40,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(), vec![]));
+    pub static ref LOAD_MANAGER: Arc<LoadManager> = Arc::new(LoadManager::new(POOL_WAIT_STATS.clone(), HashSet::default()));
 
     // Create Store instance once for use with each of the tests.
     pub static ref STORE: Arc<Store> = {