all: Make LoadManager and argument to GraphQlRunner::new

Author: lutter

graph/src/data/graphql/effort.rs

@@ -1,16 +1,17 @@
 //! Utilities to keep moving statistics about queries
 
+use graphql_parser::query as q;
+use lazy_static::lazy_static;
 use rand::{prelude::Rng, thread_rng};
 use std::collections::{HashMap, HashSet};
 use std::env;
 use std::str::FromStr;
 use std::sync::{Arc, RwLock};
 use std::time::{Duration, Instant};
 
-use lazy_static::lazy_static;
-
 use crate::components::store::PoolWaitStats;
-use crate::prelude::{info, warn, Logger};
+use crate::data::graphql::shape_hash::shape_hash;
+use crate::prelude::{info, o, warn, Logger};
 use crate::util::stats::{MovingStats, BIN_SIZE, WINDOW_SIZE};
 
 lazy_static! {
@@ -198,10 +199,15 @@ pub struct LoadManager {
 
 impl LoadManager {
     pub fn new(
-        logger: Logger,
+        logger: &Logger,
         store_wait_stats: PoolWaitStats,
-        blocked_queries: HashSet<u64>,
+        blocked_queries: Vec<Arc<q::Document>>,
     ) -> Self {
+        let logger = logger.new(o!("component" => "LoadManager"));
+        let blocked_queries = blocked_queries
+            .into_iter()
+            .map(|doc| shape_hash(&doc))
+            .collect::<HashSet<_>>();
         info!(
             logger,
             "Creating LoadManager with disabled={}", *LOAD_MANAGEMENT_DISABLED,

graphql/src/runner.rs

@@ -1,6 +1,6 @@
 use futures01::future;
 use graphql_parser::query as q;
-use std::collections::{BTreeMap, HashSet};
+use std::collections::BTreeMap;
 use std::env;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -13,10 +13,9 @@ use crate::query::execute_query;
 use crate::subscription::execute_prepared_subscription;
 use graph::data::graphql::effort::LoadManager;
 use graph::prelude::{
-    o, shape_hash, EthereumBlockPointer, GraphQlRunner as GraphQlRunnerTrait, Logger,
-    PoolWaitStats, Query, QueryExecutionError, QueryResult, QueryResultFuture, Store, StoreError,
-    SubgraphDeploymentId, SubgraphDeploymentStore, Subscription, SubscriptionError,
-    SubscriptionResultFuture,
+    o, EthereumBlockPointer, GraphQlRunner as GraphQlRunnerTrait, Logger, Query,
+    QueryExecutionError, QueryResult, QueryResultFuture, Store, StoreError, SubgraphDeploymentId,
+    SubgraphDeploymentStore, Subscription, SubscriptionError, SubscriptionResultFuture,
 };
 
 use lazy_static::lazy_static;
@@ -56,22 +55,12 @@ where
     S: Store + SubgraphDeploymentStore,
 {
     /// Creates a new query runner.
-    pub fn new(
-        logger: &Logger,
-        store: Arc<S>,
-        expensive: &Vec<Arc<q::Document>>,
-        store_wait_stats: PoolWaitStats,
-    ) -> Self {
-        let expensive = expensive
-            .into_iter()
-            .map(|doc| shape_hash(&doc))
-            .collect::<HashSet<_>>();
+    pub fn new(logger: &Logger, store: Arc<S>, load_manager: Arc<LoadManager>) -> Self {
         let logger = logger.new(o!("component" => "GraphQlRunner"));
-        let load_logger = logger.new(o!("component" => "LoadManager"));
         GraphQlRunner {
             logger,
             store,
-            load_manager: Arc::new(LoadManager::new(load_logger, store_wait_stats, expensive)),
+            load_manager,
         }
     }
 

graphql/tests/query.rs

@@ -8,6 +8,7 @@ use std::iter::FromIterator;
 use std::sync::{Arc, RwLock};
 use std::time::{Duration, Instant};
 
+use graph::data::graphql::effort::LoadManager;
 use graph::prelude::{
     futures03::stream::StreamExt, futures03::FutureExt, futures03::TryFutureExt, o, slog, tokio,
     Entity, EntityKey, EntityOperation, EthereumBlockPointer, FutureExtension, Logger, MovingStats,
@@ -231,7 +232,8 @@ fn execute_query_document_with_variables(
     variables: Option<QueryVariables>,
 ) -> QueryResult {
     let stats = Arc::new(RwLock::new(MovingStats::default()));
-    let runner = GraphQlRunner::new(&*LOGGER, STORE.clone(), &vec![], stats);
+    let load_manager = Arc::new(LoadManager::new(&*LOGGER, stats, vec![]));
+    let runner = GraphQlRunner::new(&*LOGGER, STORE.clone(), load_manager);
     let query = Query::new(Arc::new(api_test_schema()), query, variables);
 
     return_err!(runner.execute(query, None, None, None))

node/src/main.rs

@@ -13,6 +13,7 @@ use std::time::Duration;
 use tokio::sync::mpsc;
 
 use graph::components::forward;
+use graph::data::graphql::effort::LoadManager;
 use graph::log::logger;
 use graph::prelude::{
     EthereumAdapter as EthereumAdapterTrait, IndexNodeServer as _, JsonRpcServer as _, *,
@@ -588,11 +589,15 @@ async fn main() {
         .and_then(move |stores| {
             let generic_store = stores.values().next().expect("error creating stores");
 
+            let load_manager = Arc::new(LoadManager::new(
+                &logger,
+                wait_stats.clone(),
+                expensive_queries,
+            ));
             let graphql_runner = Arc::new(GraphQlRunner::new(
                 &logger,
                 generic_store.clone(),
-                &expensive_queries,
-                wait_stats.clone(),
+                load_manager,
             ));
             let mut graphql_server = GraphQLQueryServer::new(
                 &logger_factory,

store/test-store/src/lib.rs

@@ -14,7 +14,6 @@ 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;
@@ -40,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(LOGGER.clone(), POOL_WAIT_STATS.clone(), HashSet::default()));
+    pub static ref LOAD_MANAGER: Arc<LoadManager> = Arc::new(LoadManager::new(&*LOGGER, POOL_WAIT_STATS.clone(), Vec::new()));
 
     // Create Store instance once for use with each of the tests.
     pub static ref STORE: Arc<Store> = {