graph, graphql, node: Record GraphQL query result sizes in metrics

Author: lutter

graph/src/components/metrics/mod.rs

@@ -213,6 +213,18 @@ pub trait MetricsRegistry: Send + Sync + 'static {
         Ok(histogram)
     }
 
+    fn new_histogram(
+        &self,
+        name: &str,
+        help: &str,
+        buckets: Vec<f64>,
+    ) -> Result<Box<Histogram>, PrometheusError> {
+        let opts = HistogramOpts::new(name, help).buckets(buckets);
+        let histogram = Box::new(Histogram::with_opts(opts)?);
+        self.register(name, histogram.clone());
+        Ok(histogram)
+    }
+
     fn new_histogram_vec(
         &self,
         name: &str,

graphql/src/lib.rs

@@ -40,5 +40,6 @@ pub mod prelude {
 
 #[cfg(debug_assertions)]
 pub mod test_support {
+    pub use super::runner::ResultSizeMetrics;
     pub use super::runner::INITIAL_DEPLOYMENT_STATE_FOR_TESTS;
 }

graphql/src/runner.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::env;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -6,6 +7,8 @@ use std::time::{Duration, Instant};
 use crate::prelude::{QueryExecutionOptions, StoreResolver, SubscriptionExecutionOptions};
 use crate::query::execute_query;
 use crate::subscription::execute_prepared_subscription;
+use graph::prelude::MetricsRegistry;
+use graph::prometheus::{Gauge, Histogram};
 use graph::{
     components::store::SubscriptionManager,
     prelude::{
@@ -21,12 +24,59 @@ use graph::{
 
 use lazy_static::lazy_static;
 
+pub struct ResultSizeMetrics {
+    histogram: Box<Histogram>,
+    max_gauge: Box<Gauge>,
+}
+
+impl ResultSizeMetrics {
+    fn new(registry: Arc<impl MetricsRegistry>) -> Self {
+        // Divide the Histogram into exponentially sized buckets between 1k and 32G
+        let bins = (10..26).map(|n| 2u64.pow(n) as f64).collect::<Vec<_>>();
+        let histogram = registry
+            .new_histogram(
+                "query_result_size",
+                "the size of the result of successful GraphQL queries (in CacheWeight)",
+                bins,
+            )
+            .unwrap();
+
+        let max_gauge = registry
+            .new_gauge(
+                "query_result_max",
+                "the maximum size of a query result (in CacheWeight)",
+                HashMap::new(),
+            )
+            .unwrap();
+
+        Self {
+            histogram,
+            max_gauge,
+        }
+    }
+
+    // Tests need to construct one of these, but normal code doesn't
+    #[cfg(debug_assertions)]
+    pub fn make(registry: Arc<impl MetricsRegistry>) -> Self {
+        Self::new(registry)
+    }
+
+    pub fn observe(&self, size: usize) {
+        let size = size as f64;
+        self.histogram.observe(size);
+        if self.max_gauge.get() < size {
+            self.max_gauge.set(size);
+        }
+    }
+}
+
 /// GraphQL runner implementation for The Graph.
 pub struct GraphQlRunner<S, SM> {
     logger: Logger,
     store: Arc<S>,
     subscription_manager: Arc<SM>,
     load_manager: Arc<LoadManager>,
+    result_size: Arc<ResultSizeMetrics>,
 }
 
 lazy_static! {
@@ -81,13 +131,16 @@ where
         store: Arc<S>,
         subscription_manager: Arc<SM>,
         load_manager: Arc<LoadManager>,
+        registry: Arc<impl MetricsRegistry>,
     ) -> Self {
         let logger = logger.new(o!("component" => "GraphQlRunner"));
+        let result_size = Arc::new(ResultSizeMetrics::new(registry));
         GraphQlRunner {
             logger,
             store,
             subscription_manager,
             load_manager,
+            result_size,
         }
     }
 
@@ -129,6 +182,7 @@ where
         max_depth: Option<u8>,
         max_first: Option<u32>,
         max_skip: Option<u32>,
+        result_size: Arc<ResultSizeMetrics>,
     ) -> Result<QueryResults, QueryResults> {
         // We need to use the same `QueryStore` for the entire query to ensure
         // we have a consistent view if the world, even when replicas, which
@@ -181,6 +235,7 @@ where
                 bc,
                 error_policy,
                 query.schema.id().clone(),
+                result_size.cheap_clone(),
             )
             .await?;
             max_block = max_block.max(resolver.block_number());
@@ -242,6 +297,7 @@ where
             max_depth,
             max_first,
             max_skip,
+            self.result_size.cheap_clone(),
         )
         .await
         .unwrap_or_else(|e| e)
@@ -288,6 +344,7 @@ where
                 max_depth: *GRAPHQL_MAX_DEPTH,
                 max_first: *GRAPHQL_MAX_FIRST,
                 max_skip: *GRAPHQL_MAX_SKIP,
+                result_size: self.result_size.clone(),
             },
         )
         .await

graphql/src/store/prefetch.rs

@@ -24,6 +24,7 @@ use graph::{
 
 use crate::execution::{ExecutionContext, Resolver};
 use crate::query::ast as qast;
+use crate::runner::ResultSizeMetrics;
 use crate::schema::ast as sast;
 use crate::store::{build_query, StoreResolver};
 
@@ -497,8 +498,10 @@ pub fn run(
     resolver: &StoreResolver,
     ctx: &ExecutionContext<impl Resolver>,
     selection_set: &q::SelectionSet,
+    result_size: &ResultSizeMetrics,
 ) -> Result<q::Value, Vec<QueryExecutionError>> {
     execute_root_selection_set(resolver, ctx, selection_set).map(|nodes| {
+        result_size.observe(nodes.weight());
         let map = BTreeMap::default();
         q::Value::Object(nodes.into_iter().fold(map, |mut map, node| {
             // For root nodes, we only care about the children

graphql/src/store/resolver.rs

@@ -10,6 +10,7 @@ use graph::prelude::*;
 use graph::{components::store::*, data::schema::BLOCK_FIELD_TYPE};
 
 use crate::query::ext::BlockConstraint;
+use crate::runner::ResultSizeMetrics;
 use crate::schema::ast as sast;
 use crate::{prelude::*, schema::api::ErrorPolicy};
 
@@ -25,6 +26,7 @@ pub struct StoreResolver {
     deployment: DeploymentHash,
     has_non_fatal_errors: bool,
     error_policy: ErrorPolicy,
+    result_size: Arc<ResultSizeMetrics>,
 }
 
 impl CheapClone for StoreResolver {}
@@ -39,6 +41,7 @@ impl StoreResolver {
         deployment: DeploymentHash,
         store: Arc<dyn QueryStore>,
         subscription_manager: Arc<dyn SubscriptionManager>,
+        result_size: Arc<ResultSizeMetrics>,
     ) -> Self {
         StoreResolver {
             logger: logger.new(o!("component" => "StoreResolver")),
@@ -50,6 +53,7 @@ impl StoreResolver {
             // Checking for non-fatal errors does not work with subscriptions.
             has_non_fatal_errors: false,
             error_policy: ErrorPolicy::Deny,
+            result_size,
         }
     }
 
@@ -65,6 +69,7 @@ impl StoreResolver {
         bc: BlockConstraint,
         error_policy: ErrorPolicy,
         deployment: DeploymentHash,
+        result_size: Arc<ResultSizeMetrics>,
     ) -> Result<Self, QueryExecutionError> {
         let store_clone = store.cheap_clone();
         let deployment2 = deployment.clone();
@@ -87,6 +92,7 @@ impl StoreResolver {
             deployment,
             has_non_fatal_errors,
             error_policy,
+            result_size,
         };
         Ok(resolver)
     }
@@ -221,7 +227,7 @@ impl Resolver for StoreResolver {
         ctx: &ExecutionContext<Self>,
         selection_set: &q::SelectionSet,
     ) -> Result<Option<q::Value>, Vec<QueryExecutionError>> {
-        super::prefetch::run(self, ctx, selection_set).map(Some)
+        super::prefetch::run(self, ctx, selection_set, &self.result_size).map(Some)
     }
 
     fn resolve_objects(

graphql/src/subscription/mod.rs

@@ -5,6 +5,7 @@ use std::time::{Duration, Instant};
 
 use graph::{components::store::SubscriptionManager, prelude::*};
 
+use crate::runner::ResultSizeMetrics;
 use crate::{
     execution::*,
     prelude::{BlockConstraint, StoreResolver},
@@ -35,6 +36,8 @@ pub struct SubscriptionExecutionOptions {
 
     /// Maximum value for the `skip` argument.
     pub max_skip: u32,
+
+    pub result_size: Arc<ResultSizeMetrics>,
 }
 
 pub async fn execute_subscription(
@@ -83,6 +86,7 @@ async fn create_source_event_stream(
         query.schema.id().clone(),
         options.store.clone(),
         options.subscription_manager.cheap_clone(),
+        options.result_size.cheap_clone(),
     );
     let ctx = ExecutionContext {
         logger: options.logger.cheap_clone(),
@@ -158,6 +162,7 @@ fn map_source_to_response_stream(
         max_depth: _,
         max_first,
         max_skip,
+        result_size,
     } = options;
 
     Box::new(
@@ -177,6 +182,7 @@ fn map_source_to_response_stream(
                     timeout,
                     max_first,
                     max_skip,
+                    result_size.cheap_clone(),
                 )
                 .boxed(),
             }),
@@ -192,6 +198,7 @@ async fn execute_subscription_event(
     timeout: Option<Duration>,
     max_first: u32,
     max_skip: u32,
+    result_size: Arc<ResultSizeMetrics>,
 ) -> Arc<QueryResult> {
     debug!(logger, "Execute subscription event"; "event" => format!("{:?}", event));
 
@@ -202,6 +209,7 @@ async fn execute_subscription_event(
         BlockConstraint::Latest,
         ErrorPolicy::Deny,
         query.schema.id().clone(),
+        result_size,
     )
     .await
     {

graphql/tests/query.rs

@@ -30,8 +30,9 @@ use graph::{
 use graph_graphql::{prelude::*, subscription::execute_subscription};
 use test_store::{
     deployment_state, execute_subgraph_query_with_complexity, execute_subgraph_query_with_deadline,
-    revert_block, run_test_sequentially, transact_entity_operations, transact_errors, Store,
-    BLOCK_ONE, GENESIS_PTR, LOAD_MANAGER, LOGGER, STORE, SUBSCRIPTION_MANAGER,
+    result_size_metrics, revert_block, run_test_sequentially, transact_entity_operations,
+    transact_errors, Store, BLOCK_ONE, GENESIS_PTR, LOAD_MANAGER, LOGGER, METRICS_REGISTRY, STORE,
+    SUBSCRIPTION_MANAGER,
 };
 
 const NETWORK_NAME: &str = "fake_network";
@@ -258,6 +259,7 @@ async fn execute_query_document_with_variables(
         STORE.clone(),
         SUBSCRIPTION_MANAGER.clone(),
         LOAD_MANAGER.clone(),
+        METRICS_REGISTRY.clone(),
     ));
     let target = QueryTarget::Deployment(id.clone());
     let query = Query::new(query, variables);
@@ -885,6 +887,7 @@ fn query_complexity_subscriptions() {
             max_depth: 100,
             max_first: std::u32::MAX,
             max_skip: std::u32::MAX,
+            result_size: result_size_metrics(),
         };
         let schema = STORE.subgraph_store().api_schema(&deployment.hash).unwrap();
 
@@ -933,6 +936,7 @@ fn query_complexity_subscriptions() {
             max_depth: 100,
             max_first: std::u32::MAX,
             max_skip: std::u32::MAX,
+            result_size: result_size_metrics(),
         };
 
         // The extra introspection causes the complexity to go over.
@@ -1311,6 +1315,7 @@ fn subscription_gets_result_even_without_events() {
             max_depth: 100,
             max_first: std::u32::MAX,
             max_skip: std::u32::MAX,
+            result_size: result_size_metrics(),
         };
         // Execute the subscription and expect at least one result to be
         // available in the result stream

node/src/bin/manager.rs

@@ -432,13 +432,14 @@ impl Context {
         let store = self.store();
 
         let subscription_manager = Arc::new(PanicSubscriptionManager);
-        let load_manager = Arc::new(LoadManager::new(&logger, vec![], registry));
+        let load_manager = Arc::new(LoadManager::new(&logger, vec![], registry.clone()));
 
         Arc::new(GraphQlRunner::new(
             &logger,
             store,
             subscription_manager,
             load_manager,
+            registry,
         ))
     }
 }

node/src/main.rs

@@ -231,6 +231,7 @@ async fn main() {
             network_store.clone(),
             subscription_manager.clone(),
             load_manager,
+            metrics_registry.clone(),
         ));
         let mut graphql_server = GraphQLQueryServer::new(
             &logger_factory,

store/test-store/src/store.rs

@@ -14,6 +14,7 @@ use graph::{
 use graph_graphql::prelude::{
     execute_query, Query as PreparedQuery, QueryExecutionOptions, StoreResolver,
 };
+use graph_graphql::test_support::ResultSizeMetrics;
 use graph_mock::MockMetricsRegistry;
 use graph_node::config::{Config, Opt};
 use graph_node::store_builder::StoreBuilder;
@@ -45,10 +46,12 @@ lazy_static! {
     static ref SEQ_LOCK: Mutex<()> = Mutex::new(());
     pub static ref STORE_RUNTIME: Runtime =
         Builder::new_multi_thread().enable_all().build().unwrap();
+    pub static ref METRICS_REGISTRY: Arc<MockMetricsRegistry> =
+        Arc::new(MockMetricsRegistry::new());
     pub static ref LOAD_MANAGER: Arc<LoadManager> = Arc::new(LoadManager::new(
         &*LOGGER,
         Vec::new(),
-        Arc::new(MockMetricsRegistry::new()),
+        METRICS_REGISTRY.clone(),
     ));
     static ref STORE_POOL_CONFIG: (Arc<Store>, ConnectionPool, Config, Arc<SubscriptionManager>) =
         build_store();
@@ -359,6 +362,10 @@ macro_rules! return_err {
     };
 }
 
+pub fn result_size_metrics() -> Arc<ResultSizeMetrics> {
+    Arc::new(ResultSizeMetrics::make(METRICS_REGISTRY.clone()))
+}
+
 fn execute_subgraph_query_internal(
     query: Query,
     target: QueryTarget,
@@ -403,7 +410,8 @@ fn execute_subgraph_query_internal(
             SUBSCRIPTION_MANAGER.clone(),
             bc,
             error_policy,
-            query.schema.id().clone()
+            query.schema.id().clone(),
+            result_size_metrics()
         )));
         result.append(rt.block_on(execute_query(
             query.clone(),