all: Measure loading time in EntityCache.as_modifications

This adds a new section "as_modifications_load" that shows us how much time
we spend reading from the database

Author: lutter

core/src/amp_subgraph/runner/data_processing.rs

@@ -139,7 +139,7 @@ async fn process_record_batch_group<AC>(
         entity_lfu_cache,
         evict_stats: _,
     } = entity_cache
-        .as_modifications(block_number.compat())
+        .as_modifications(block_number.compat(), &cx.metrics.stopwatch)
         .await
         .map_err(Error::from)
         .map_err(|e| e.context("failed to extract entity modifications from the state"))?;

core/src/subgraph/runner/mod.rs

@@ -686,7 +686,7 @@ where
             entity_lfu_cache: cache,
             evict_stats,
         } = entity_cache
-            .as_modifications(block_ptr.number)
+            .as_modifications(block_ptr.number, &self.metrics.host.stopwatch)
             .await
             .classify()?;
         section.end();
@@ -1571,7 +1571,7 @@ where
             mods.extend(
                 block_state
                     .entity_cache
-                    .as_modifications(block.number())
+                    .as_modifications(block.number(), &self.metrics.subgraph.stopwatch)
                     .await?
                     .modifications,
             );

graph/src/components/store/entity_cache.rs

@@ -8,7 +8,7 @@ use crate::cheap_clone::CheapClone;
 use crate::components::store::write::EntityModification;
 use crate::components::store::{self as s, Entity, EntityOperation};
 use crate::data::store::{EntityValidationError, Id, IdType, IntoEntityIterator};
-use crate::prelude::{CacheWeight, ENV_VARS};
+use crate::prelude::{CacheWeight, StopwatchMetrics, ENV_VARS};
 use crate::schema::{EntityKey, InputSchema};
 use crate::util::intern::Error as InternError;
 use crate::util::lfu_cache::{EvictStats, LfuCache};
@@ -474,6 +474,7 @@ impl EntityCache {
     pub async fn as_modifications(
         mut self,
         block: BlockNumber,
+        stopwatch: &StopwatchMetrics,
     ) -> Result<ModificationsAndCache, StoreError> {
         assert!(!self.in_handler);
 
@@ -491,10 +492,14 @@ impl EntityCache {
         // is wrong and the store already has a version of the entity from a
         // previous block, the attempt to insert will trigger a constraint
         // violation in the database, ensuring correctness
-        let missing = missing.filter(|key| !key.entity_type.is_immutable());
+        {
+            let _section = stopwatch.start_section("as_modifications_load");
 
-        for (entity_key, entity) in self.store.get_many(missing.cloned().collect()).await? {
-            self.current.insert(entity_key, Some(Arc::new(entity)));
+            let missing = missing.filter(|key| !key.entity_type.is_immutable());
+
+            for (entity_key, entity) in self.store.get_many(missing.cloned().collect()).await? {
+                self.current.insert(entity_key, Some(Arc::new(entity)));
+            }
         }
 
         let mut mods = Vec::new();

runtime/test/src/test.rs

@@ -20,7 +20,7 @@ use graph_runtime_wasm::{
 use semver::Version;
 use std::collections::{BTreeMap, HashMap};
 use std::str::FromStr;
-use test_store::{LOGGER, STORE};
+use test_store::{LOGGER, STOPWATCH, STORE, SUBGRAPH_STORE};
 use wasmtime::{AsContext, AsContextMut};
 
 use crate::common::{mock_context, mock_data_source};
@@ -44,20 +44,20 @@ fn subgraph_id_with_api_version(subgraph_id: &str, api_version: Version) -> Stri
     )
 }
 
-async fn test_valid_module_and_store(
+async fn test_module_and_deployment(
     subgraph_id: &str,
     data_source: DataSource,
     api_version: Version,
-) -> (WasmInstance, Arc<impl SubgraphStore>, DeploymentLocator) {
-    test_valid_module_and_store_with_timeout(subgraph_id, data_source, api_version, None).await
+) -> (WasmInstance, DeploymentLocator) {
+    test_module_and_deployment_with_timeout(subgraph_id, data_source, api_version, None).await
 }
 
-async fn test_valid_module_and_store_with_timeout(
+async fn test_module_and_deployment_with_timeout(
     subgraph_id: &str,
     data_source: DataSource,
     api_version: Version,
     timeout: Option<Duration>,
-) -> (WasmInstance, Arc<impl SubgraphStore>, DeploymentLocator) {
+) -> (WasmInstance, DeploymentLocator) {
     let logger = Logger::root(slog::Discard, o!());
     let subgraph_id_with_api_version =
         subgraph_id_with_api_version(subgraph_id, api_version.clone());
@@ -80,7 +80,7 @@ async fn test_valid_module_and_store_with_timeout(
         }",
     )
     .await;
-    let stopwatch_metrics = StopwatchMetrics::new(
+    let stopwatch = StopwatchMetrics::new(
         logger.clone(),
         deployment_id.clone(),
         "test",
@@ -93,7 +93,7 @@ async fn test_valid_module_and_store_with_timeout(
     let host_metrics = Arc::new(HostMetrics::new(
         metrics_registry,
         deployment_id.as_str(),
-        stopwatch_metrics,
+        stopwatch.cheap_clone(),
         gas_metrics,
     ));
 
@@ -115,15 +115,15 @@ async fn test_valid_module_and_store_with_timeout(
     .await
     .unwrap();
 
-    (module, store.subgraph_store(), deployment)
+    (module, deployment)
 }
 
 pub async fn test_module(
     subgraph_id: &str,
     data_source: DataSource,
     api_version: Version,
 ) -> WasmInstance {
-    test_valid_module_and_store(subgraph_id, data_source, api_version)
+    test_module_and_deployment(subgraph_id, data_source, api_version)
         .await
         .0
 }
@@ -135,9 +135,7 @@ pub async fn test_module_latest(subgraph_id: &str, wasm_file: &str) -> WasmInsta
         &wasm_file_path(wasm_file, API_VERSION_0_0_5),
         version.clone(),
     );
-    test_valid_module_and_store(subgraph_id, ds, version)
-        .await
-        .0
+    test_module_and_deployment(subgraph_id, ds, version).await.0
 }
 
 pub trait SyncWasmTy: wasmtime::WasmTy + Sync {}
@@ -529,7 +527,7 @@ async fn run_ipfs_map(
             .to_owned()
     };
 
-    let (mut instance, _, _) = test_valid_module_and_store(
+    let mut instance = test_module(
         subgraph_id,
         mock_data_source(
             &wasm_file_path("ipfs_map.wasm", api_version.clone()),
@@ -557,7 +555,7 @@ async fn run_ipfs_map(
         .take_ctx()
         .take_state()
         .entity_cache
-        .as_modifications(0)
+        .as_modifications(0, &STOPWATCH)
         .await?
         .modifications;
 
@@ -1008,7 +1006,8 @@ async fn ens_name_by_hash_v0_0_5() {
 }
 
 async fn test_entity_store(api_version: Version) {
-    let (mut instance, store, deployment) = test_valid_module_and_store(
+    let store = SUBGRAPH_STORE.clone();
+    let (mut instance, deployment) = test_module_and_deployment(
         "entityStore",
         mock_data_source(
             &wasm_file_path("store.wasm", api_version.clone()),
@@ -1073,7 +1072,11 @@ async fn test_entity_store(api_version: Version) {
         &mut ctx.ctx.state.entity_cache,
         EntityCache::new(Arc::new(writable.clone())),
     );
-    let mut mods = cache.as_modifications(0).await.unwrap().modifications;
+    let mut mods = cache
+        .as_modifications(0, &STOPWATCH)
+        .await
+        .unwrap()
+        .modifications;
     assert_eq!(1, mods.len());
     match mods.pop().unwrap() {
         EntityModification::Overwrite { data, .. } => {
@@ -1093,7 +1096,7 @@ async fn test_entity_store(api_version: Version) {
         .take_ctx()
         .take_state()
         .entity_cache
-        .as_modifications(0)
+        .as_modifications(0, &STOPWATCH)
         .await
         .unwrap()
         .modifications;
@@ -1626,7 +1629,7 @@ async fn generate_id() {
 
     let entity_cache = host.ctx.state.entity_cache;
     let mods = entity_cache
-        .as_modifications(12)
+        .as_modifications(12, &STOPWATCH)
         .await
         .unwrap()
         .modifications;

runtime/test/src/test/abi.rs

@@ -7,7 +7,7 @@ use super::*;
 
 async fn test_unbounded_loop(api_version: Version) {
     // Set handler timeout to 3 seconds.
-    let mut instance = test_valid_module_and_store_with_timeout(
+    let mut instance = test_module_and_deployment_with_timeout(
         "unboundedLoop",
         mock_data_source(
             &wasm_file_path("non_terminating.wasm", api_version.clone()),

store/test-store/src/store.rs

@@ -68,6 +68,13 @@ lazy_static! {
     pub static ref NODE_ID: NodeId = NodeId::new("test").unwrap();
     pub static ref SUBGRAPH_STORE: Arc<DieselSubgraphStore> = STORE.subgraph_store();
     static ref BLOCK_STORE: DieselBlockStore = STORE.block_store();
+    pub static ref STOPWATCH: StopwatchMetrics = StopwatchMetrics::new(
+        Logger::root(slog::Discard, o!()),
+        DeploymentHash::new("test").unwrap(),
+        "test",
+        METRICS_REGISTRY.clone(),
+        "dummy".to_string(),
+    );
     pub static ref GENESIS_PTR: BlockPtr = (
         B256::from(hex!(
             "bd34884280958002c51d3f7b5f853e6febeba33de0f40d15b0363006533c924f"
@@ -362,13 +369,6 @@ pub async fn transact_entities_and_dynamic_data_sources(
         Arc::new(manifest_idx_and_name),
     ))?;
 
-    let mut entity_cache = EntityCache::new(Arc::new(store.clone()));
-    entity_cache.append(ops);
-    let mods = entity_cache
-        .as_modifications(block_ptr_to.number)
-        .await
-        .expect("failed to convert to modifications")
-        .modifications;
     let metrics_registry = Arc::new(MetricsRegistry::mock());
     let stopwatch_metrics = StopwatchMetrics::new(
         Logger::root(slog::Discard, o!()),
@@ -377,6 +377,14 @@ pub async fn transact_entities_and_dynamic_data_sources(
         metrics_registry.clone(),
         store.shard().to_string(),
     );
+
+    let mut entity_cache = EntityCache::new(Arc::new(store.clone()));
+    entity_cache.append(ops);
+    let mods = entity_cache
+        .as_modifications(block_ptr_to.number, &stopwatch_metrics)
+        .await
+        .expect("failed to convert to modifications")
+        .modifications;
     let block_time = BlockTime::for_test(&block_ptr_to);
     store
         .transact_block_operations(

store/test-store/tests/graph/entity_cache.rs

@@ -195,7 +195,7 @@ fn sort_by_entity_key(mut mods: Vec<EntityModification>) -> Vec<EntityModificati
 async fn empty_cache_modifications() {
     let store = Arc::new(MockStore::new(BTreeMap::new()));
     let cache = EntityCache::new(store);
-    let result = cache.as_modifications(0).await;
+    let result = cache.as_modifications(0, &STOPWATCH).await;
     assert_eq!(result.unwrap().modifications, vec![]);
 }
 
@@ -225,7 +225,7 @@ async fn insert_modifications() {
     mogwai_data.set_vid(100).unwrap();
     sigurros_data.set_vid(101).unwrap();
 
-    let result = cache.as_modifications(0).await;
+    let result = cache.as_modifications(0, &STOPWATCH).await;
     assert_eq!(
         sort_by_entity_key(result.unwrap().modifications),
         sort_by_entity_key(vec![
@@ -276,7 +276,7 @@ async fn overwrite_modifications() {
     mogwai_data.set_vid(100).unwrap();
     sigurros_data.set_vid(101).unwrap();
 
-    let result = cache.as_modifications(0).await;
+    let result = cache.as_modifications(0, &STOPWATCH).await;
     assert_eq!(
         sort_by_entity_key(result.unwrap().modifications),
         sort_by_entity_key(vec![
@@ -316,7 +316,7 @@ async fn consecutive_modifications() {
 
     // We expect a single overwrite modification for the above that leaves "id"
     // and "name" untouched, sets "founded" and removes the "label" field.
-    let result = cache.as_modifications(0).await;
+    let result = cache.as_modifications(0, &STOPWATCH).await;
     assert_eq!(
         sort_by_entity_key(result.unwrap().modifications),
         sort_by_entity_key(vec![EntityModification::overwrite(
@@ -344,7 +344,7 @@ async fn check_vid_sequence() {
             .unwrap();
     }
 
-    let result = cache.as_modifications(0).await;
+    let result = cache.as_modifications(0, &STOPWATCH).await;
     let mods = result.unwrap().modifications;
     for m in mods {
         match m {
@@ -386,7 +386,7 @@ async fn offchain_trigger_vid_collision_without_fix() {
         .set(band1_key.clone(), band1_data, block, None)
         .await
         .unwrap();
-    let result1 = cache1.as_modifications(block).await.unwrap();
+    let result1 = cache1.as_modifications(block, &STOPWATCH).await.unwrap();
 
     // Simulate second offchain trigger: another fresh EntityCache (vid_seq ALSO starts at 100)
     let store2 = Arc::new(MockStore::new(BTreeMap::new()));
@@ -397,7 +397,7 @@ async fn offchain_trigger_vid_collision_without_fix() {
         .set(band2_key.clone(), band2_data, block, None)
         .await
         .unwrap();
-    let result2 = cache2.as_modifications(block).await.unwrap();
+    let result2 = cache2.as_modifications(block, &STOPWATCH).await.unwrap();
 
     // Extract VIDs from both modifications
     let vid1 = match &result1.modifications[0] {
@@ -445,7 +445,7 @@ async fn offchain_trigger_vid_no_collision_with_fix() {
 
     // THE FIX: capture vid_seq BEFORE as_modifications consumes the cache
     let next_vid_seq = cache1.vid_seq;
-    let result1 = cache1.as_modifications(block).await.unwrap();
+    let result1 = cache1.as_modifications(block, &STOPWATCH).await.unwrap();
 
     // Second offchain trigger: set vid_seq to where first trigger left off
     let store2 = Arc::new(MockStore::new(BTreeMap::new()));
@@ -457,7 +457,7 @@ async fn offchain_trigger_vid_no_collision_with_fix() {
         .set(band2_key.clone(), band2_data, block, None)
         .await
         .unwrap();
-    let result2 = cache2.as_modifications(block).await.unwrap();
+    let result2 = cache2.as_modifications(block, &STOPWATCH).await.unwrap();
 
     let vid1 = match &result1.modifications[0] {
         EntityModification::Insert { data, .. } => data.vid(),

store/test-store/tests/postgres/aggregation.rs

@@ -26,7 +26,8 @@ use graph::{
 };
 use graph_store_postgres::Store as DieselStore;
 use test_store::{
-    create_test_subgraph, remove_subgraphs, run_test_sequentially, BLOCKS, LOGGER, METRICS_REGISTRY,
+    create_test_subgraph, remove_subgraphs, run_test_sequentially, BLOCKS, LOGGER,
+    METRICS_REGISTRY, STOPWATCH,
 };
 
 const SCHEMA: &str = r#"
@@ -110,7 +111,7 @@ pub async fn insert_entities(
     let mut entity_cache = EntityCache::new(Arc::new(store.clone()));
     entity_cache.append(ops);
     let mods = entity_cache
-        .as_modifications(block_ptr_to.number)
+        .as_modifications(block_ptr_to.number, &STOPWATCH)
         .await
         .expect("failed to convert to modifications")
         .modifications;