all: Track wait stats in connection pool, add LoadManager

Author: lutter

graph/src/components/graphql.rs

@@ -1,6 +1,6 @@
 use futures::prelude::*;
 
-use crate::data::graphql::effort::QueryEffort;
+use crate::data::graphql::effort::LoadManager;
 use crate::data::query::{Query, QueryError, QueryResult};
 use crate::data::subscription::{Subscription, SubscriptionError, SubscriptionResult};
 
@@ -50,5 +50,5 @@ pub trait GraphQlRunner: Send + Sync + 'static {
             })
     }
 
-    fn effort(&self) -> Arc<QueryEffort>;
+    fn load_manager(&self) -> Arc<LoadManager>;
 }

graph/src/components/store.rs

@@ -11,7 +11,7 @@ use std::env;
 use std::fmt;
 use std::str::FromStr;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use std::time::{Duration, Instant};
 use web3::types::{Address, H256};
 
@@ -1214,6 +1214,10 @@ mock! {
     }
 }
 
+// The type that the connection pool uses to track wait times for
+// connection checkouts
+pub type PoolWaitStats = Arc<RwLock<MovingStats>>;
+
 // The store trait must be implemented manually because mockall does not support async_trait, nor borrowing from arguments.
 impl Store for MockStore {
     fn block_ptr(

graph/src/data/graphql/effort.rs

@@ -4,6 +4,7 @@ use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 use std::time::{Duration, Instant};
 
+use crate::components::store::PoolWaitStats;
 use crate::util::stats::{MovingStats, BIN_SIZE, WINDOW_SIZE};
 
 pub struct QueryEffort {
@@ -61,3 +62,21 @@ impl QueryEffortInner {
         self.total.add_at(now, duration);
     }
 }
+
+pub struct LoadManager {
+    effort: QueryEffort,
+    store_wait_stats: PoolWaitStats,
+}
+
+impl LoadManager {
+    pub fn new(store_wait_stats: PoolWaitStats) -> Self {
+        Self {
+            effort: QueryEffort::default(),
+            store_wait_stats,
+        }
+    }
+
+    pub fn add_query(&self, shape_hash: u64, duration: Duration) {
+        self.effort.add(shape_hash, duration);
+    }
+}

graph/src/lib.rs

@@ -101,9 +101,9 @@ pub mod prelude {
         AttributeIndexDefinition, BlockNumber, ChainStore, ChildMultiplicity, EntityCache,
         EntityChange, EntityChangeOperation, EntityCollection, EntityFilter, EntityKey, EntityLink,
         EntityModification, EntityOperation, EntityOrder, EntityQuery, EntityRange, EntityWindow,
-        EthereumCallCache, MetadataOperation, ParentLink, Store, StoreError, StoreEvent,
-        StoreEventStream, StoreEventStreamBox, SubgraphDeploymentStore, TransactionAbortError,
-        WindowAttribute, BLOCK_NUMBER_MAX, SUBSCRIPTION_THROTTLE_INTERVAL,
+        EthereumCallCache, MetadataOperation, ParentLink, PoolWaitStats, Store, StoreError,
+        StoreEvent, StoreEventStream, StoreEventStreamBox, SubgraphDeploymentStore,
+        TransactionAbortError, WindowAttribute, BLOCK_NUMBER_MAX, SUBSCRIPTION_THROTTLE_INTERVAL,
     };
     pub use crate::components::subgraph::{
         BlockState, DataSourceLoader, DataSourceTemplateInfo, HostMetrics, RuntimeHost,

graphql/src/query/mod.rs

@@ -5,7 +5,7 @@ use std::sync::{atomic::AtomicBool, Arc};
 use std::time::Instant;
 use uuid::Uuid;
 
-use graph::data::graphql::effort::QueryEffort;
+use graph::data::graphql::effort::LoadManager;
 
 use crate::execution::*;
 use crate::schema::ast as sast;
@@ -33,7 +33,7 @@ where
     /// Maximum value for the `first` argument.
     pub max_first: u32,
 
-    pub effort: Arc<QueryEffort>,
+    pub load_manager: Arc<LoadManager>,
 }
 
 /// Executes a query and returns a result.
@@ -79,7 +79,7 @@ where
     let start = Instant::now();
     let result = execute_root_selection_set(&ctx, selection_set, query_type, block_ptr);
     let elapsed = start.elapsed();
-    options.effort.add(query.shape_hash, elapsed);
+    options.load_manager.add_query(query.shape_hash, elapsed);
     if *graph::log::LOG_GQL_TIMING {
         info!(
             query_logger,

graphql/src/runner.rs

@@ -11,11 +11,12 @@ use crate::prelude::{
 };
 use crate::query::execute_query;
 use crate::subscription::execute_prepared_subscription;
-use graph::data::graphql::effort::QueryEffort;
+use graph::data::graphql::effort::LoadManager;
 use graph::prelude::{
-    o, shape_hash, EthereumBlockPointer, GraphQlRunner as GraphQlRunnerTrait, Logger, Query,
-    QueryExecutionError, QueryResult, QueryResultFuture, Store, StoreError, SubgraphDeploymentId,
-    SubgraphDeploymentStore, Subscription, SubscriptionError, SubscriptionResultFuture,
+    o, shape_hash, EthereumBlockPointer, GraphQlRunner as GraphQlRunnerTrait, Logger,
+    PoolWaitStats, Query, QueryExecutionError, QueryResult, QueryResultFuture, Store, StoreError,
+    SubgraphDeploymentId, SubgraphDeploymentStore, Subscription, SubscriptionError,
+    SubscriptionResultFuture,
 };
 
 use lazy_static::lazy_static;
@@ -25,7 +26,7 @@ pub struct GraphQlRunner<S> {
     logger: Logger,
     store: Arc<S>,
     expensive: HashMap<u64, Arc<q::Document>>,
-    effort: Arc<QueryEffort>,
+    load_manager: Arc<LoadManager>,
 }
 
 lazy_static! {
@@ -56,7 +57,12 @@ where
     S: Store + SubgraphDeploymentStore,
 {
     /// Creates a new query runner.
-    pub fn new(logger: &Logger, store: Arc<S>, expensive: &Vec<Arc<q::Document>>) -> Self {
+    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), doc.clone()))
@@ -65,7 +71,7 @@ where
             logger: logger.new(o!("component" => "GraphQlRunner")),
             store,
             expensive,
-            effort: Arc::new(QueryEffort::default()),
+            load_manager: Arc::new(LoadManager::new(store_wait_stats)),
         }
     }
 
@@ -135,7 +141,7 @@ where
                     resolver,
                     deadline: GRAPHQL_QUERY_TIMEOUT.map(|t| Instant::now() + t),
                     max_first: max_first.unwrap_or(*GRAPHQL_MAX_FIRST),
-                    effort: self.effort.clone(),
+                    load_manager: self.load_manager.clone(),
                 },
             ) {
                 Err(errs) => errors.extend(errs),
@@ -215,7 +221,7 @@ where
         Box::new(future::result(result))
     }
 
-    fn effort(&self) -> Arc<QueryEffort> {
-        self.effort.clone()
+    fn load_manager(&self) -> Arc<LoadManager> {
+        self.load_manager.clone()
     }
 }

graphql/tests/introspection.rs

@@ -4,16 +4,15 @@ extern crate pretty_assertions;
 use graphql_parser::{query as q, schema as s};
 use std::collections::HashMap;
 use std::sync::Arc;
-use std::time::Duration;
 
-use graph::data::graphql::effort::QueryEffort;
 use graph::prelude::{
     o, slog, Logger, Query, QueryExecutionError, QueryResult, Schema, SubgraphDeploymentId,
 };
 use graph_graphql::prelude::{
     api_schema, execute_query, object, object_value, ExecutionContext, ObjectOrInterface,
     Query as PreparedQuery, QueryExecutionOptions, Resolver,
 };
+use test_store::LOAD_MANAGER;
 
 /// Mock resolver used in tests that don't need a resolver.
 #[derive(Clone)]
@@ -560,10 +559,7 @@ fn introspection_query(schema: Schema, query: &str) -> QueryResult {
         resolver: MockResolver,
         deadline: None,
         max_first: std::u32::MAX,
-        effort: Arc::new(QueryEffort::new(
-            Duration::from_millis(0),
-            Duration::from_millis(0),
-        )),
+        load_manager: LOAD_MANAGER.clone(),
     };
 
     let result = PreparedQuery::new(query, None, 100)

graphql/tests/query.rs

@@ -5,13 +5,13 @@ use graphql_parser::{query as q, Pos};
 use lazy_static::lazy_static;
 use std::collections::HashMap;
 use std::iter::FromIterator;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use std::time::{Duration, Instant};
 
 use graph::prelude::{
     futures03::stream::StreamExt, futures03::FutureExt, futures03::TryFutureExt, o, slog, tokio,
-    Entity, EntityKey, EntityOperation, EthereumBlockPointer, FutureExtension, Logger, Query,
-    QueryError, QueryExecutionError, QueryResult, QueryVariables, Schema, Store,
+    Entity, EntityKey, EntityOperation, EthereumBlockPointer, FutureExtension, Logger, MovingStats,
+    Query, QueryError, QueryExecutionError, QueryResult, QueryVariables, Schema, Store,
     SubgraphDeploymentEntity, SubgraphDeploymentId, SubgraphDeploymentStore, SubgraphManifest,
     Subscription, SubscriptionError, Value,
 };
@@ -230,7 +230,8 @@ fn execute_query_document_with_variables(
     query: q::Document,
     variables: Option<QueryVariables>,
 ) -> QueryResult {
-    let runner = GraphQlRunner::new(&*LOGGER, STORE.clone(), &vec![]);
+    let stats = Arc::new(RwLock::new(MovingStats::default()));
+    let runner = GraphQlRunner::new(&*LOGGER, STORE.clone(), &vec![], stats);
     let query = Query::new(Arc::new(api_test_schema()), query, variables);
 
     return_err!(runner.execute(query, None, None, None))

node/src/main.rs

@@ -8,6 +8,7 @@ use std::env;
 use std::io::{BufRead, BufReader};
 use std::path::Path;
 use std::str::FromStr;
+use std::sync::RwLock;
 use std::time::Duration;
 use tokio::sync::mpsc;
 
@@ -518,12 +519,14 @@ async fn main() {
     let stores_error_logger = logger.clone();
     let stores_eth_adapters = eth_adapters.clone();
     let contention_logger = logger.clone();
+    let wait_stats = Arc::new(RwLock::new(MovingStats::default()));
 
     let postgres_conn_pool = create_connection_pool(
         postgres_url.clone(),
         store_conn_pool_size,
         &logger,
         connection_pool_registry,
+        wait_stats.clone(),
     );
 
     let chain_head_update_listener = Arc::new(PostgresChainHeadUpdateListener::new(
@@ -589,6 +592,7 @@ async fn main() {
                 &logger,
                 generic_store.clone(),
                 &expensive_queries,
+                wait_stats.clone(),
             ));
             let mut graphql_server = GraphQLQueryServer::new(
                 &logger_factory,

server/http/src/service.rs

@@ -475,7 +475,7 @@ mod tests {
     use hyper::{Body, Method, Request};
     use std::collections::BTreeMap;
 
-    use graph::data::graphql::effort::QueryEffort;
+    use graph::data::graphql::effort::LoadManager;
     use graph::prelude::*;
     use graph_mock::{mock_store_with_users_subgraph, MockMetricsRegistry};
     use graphql_parser::query as q;
@@ -515,7 +515,7 @@ mod tests {
             unreachable!();
         }
 
-        fn effort(&self) -> Arc<QueryEffort> {
+        fn load_manager(&self) -> Arc<LoadManager> {
             unimplemented!()
         }
     }