core, graph, store: Adjust vids of cached entities to reflect db changes

When entities are inserted or updated, their vid changes. That change needs
to be propagated to the entities that stay in the EntityCache across blocks
so that a future update to them modifies the right version (vid)

Author: lutter

core/src/subgraph/instance_manager.rs

@@ -992,7 +992,7 @@ async fn process_block<T: RuntimeHostBuilder<C>, C: Blockchain>(
         data_sources,
         deterministic_errors,
     ) {
-        Ok(_) => {
+        Ok(vid_map) => {
             // For subgraphs with `nonFatalErrors` feature disabled, we consider
             // any error as fatal.
             //
@@ -1020,6 +1020,10 @@ async fn process_block<T: RuntimeHostBuilder<C>, C: Blockchain>(
                 return Err(BlockProcessingError::Canceled);
             }
 
+            // Adjust the vids of cached entities because inserts and
+            // updates will have changed them
+            ctx.state.entity_lfu_cache.update_vids(vid_map);
+
             Ok(needs_restart)
         }
 

graph/src/components/store.rs

@@ -1048,7 +1048,7 @@ pub trait WritableStore: Send + Sync + 'static {
         stopwatch: StopwatchMetrics,
         data_sources: Vec<StoredDynamicDataSource>,
         deterministic_errors: Vec<SubgraphError>,
-    ) -> Result<(), StoreError>;
+    ) -> Result<Vec<(EntityKey, Vid)>, StoreError>;
 
     /// Look up multiple entities as of the latest block. Returns a map of
     /// entities by type.
@@ -1234,7 +1234,7 @@ impl WritableStore for MockStore {
         _: StopwatchMetrics,
         _: Vec<StoredDynamicDataSource>,
         _: Vec<SubgraphError>,
-    ) -> Result<(), StoreError> {
+    ) -> Result<Vec<(EntityKey, Vid)>, StoreError> {
         unimplemented!()
     }
 
@@ -1801,6 +1801,19 @@ impl LfuCache<EntityKey, Option<EntityVersion>> {
             Some(data) => Ok(data.to_owned()),
         }
     }
+
+    /// Update the `vid` of cached entities to reflect changes made in the
+    /// database. When entities stay cached across insert/update operations,
+    /// their vid in the database changes as a result of these operations
+    /// and needs to be updated
+    pub fn update_vids(&mut self, vid_map: Vec<(EntityKey, Vid)>) {
+        for (key, vid) in vid_map {
+            assert!(vid.is_some());
+            if let Some(Some(x)) = self.get_mut(key) {
+                x.vid = vid;
+            }
+        }
+    }
 }
 
 /// Determines which columns should be selected in a table.

graph/src/util/lfu_cache.rs

@@ -87,7 +87,7 @@ impl<K: Clone + Ord + Eq + Hash + Debug + CacheWeight, V: CacheWeight + Default>
     /// Updates and bumps freceny if already present.
     pub fn insert(&mut self, key: K, value: V) {
         let weight = CacheEntry::weight(&key, &value);
-        match self.get_mut(key.clone()) {
+        match self.get_mut_entry(key.clone()) {
             None => {
                 self.total_weight += weight;
                 self.queue.push(
@@ -119,7 +119,7 @@ impl<K: Clone + Ord + Eq + Hash + Debug + CacheWeight, V: CacheWeight + Default>
             .unwrap_or(0)
     }
 
-    fn get_mut(&mut self, key: K) -> Option<&mut CacheEntry<K, V>> {
+    fn get_mut_entry(&mut self, key: K) -> Option<&mut CacheEntry<K, V>> {
         // Increment the frequency by 1
         let key_entry = CacheEntry::cache_key(key);
         self.queue
@@ -130,8 +130,12 @@ impl<K: Clone + Ord + Eq + Hash + Debug + CacheWeight, V: CacheWeight + Default>
         })
     }
 
+    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
+        self.get_mut_entry(key).map(|x| &mut x.value)
+    }
+
     pub fn get(&mut self, key: &K) -> Option<&V> {
-        self.get_mut(key.clone()).map(|x| &x.value)
+        self.get_mut_entry(key.clone()).map(|x| &x.value)
     }
 
     pub fn remove(&mut self, key: &K) -> Option<V> {

store/postgres/src/deployment_store.rs

@@ -311,7 +311,7 @@ impl DeploymentStore {
         mods: &[EntityModification],
         ptr: &BlockPtr,
         stopwatch: StopwatchMetrics,
-    ) -> Result<i32, StoreError> {
+    ) -> Result<(i32, Vec<(EntityKey, Vid)>), StoreError> {
         use EntityModification::*;
         let mut count = 0;
 
@@ -349,25 +349,38 @@ impl DeploymentStore {
         }
 
         // Apply modification groups.
+        let mut vid_map = Vec::new();
+
         // Inserts:
         for (entity_type, mut entities) in inserts.into_iter() {
-            count +=
+            for (id, vid) in
                 self.insert_entities(&entity_type, &mut entities, conn, layout, ptr, &stopwatch)?
-                    as i32
+            {
+                count += 1;
+                vid_map.push((
+                    EntityKey::data(layout.site.deployment.clone(), entity_type.to_string(), id),
+                    vid,
+                ))
+            }
         }
 
         // Overwrites:
         for (entity_type, (mut entities, vids)) in overwrites.into_iter() {
             // we do not update the count since the number of entities remains the same
-            self.overwrite_entities(
+            for (id, vid) in self.overwrite_entities(
                 &entity_type,
                 &mut entities,
                 vids.as_slice(),
                 conn,
                 layout,
                 ptr,
                 &stopwatch,
-            )?;
+            )? {
+                vid_map.push((
+                    EntityKey::data(layout.site.deployment.clone(), entity_type.to_string(), id),
+                    vid,
+                ))
+            }
         }
 
         // Removals
@@ -376,7 +389,7 @@ impl DeploymentStore {
                 self.remove_entities(&entity_type, vids.as_slice(), conn, layout, ptr, &stopwatch)?
                     as i32;
         }
-        Ok(count)
+        Ok((count, vid_map))
     }
 
     fn insert_entities<'a>(
@@ -387,7 +400,7 @@ impl DeploymentStore {
         layout: &'a Layout,
         ptr: &BlockPtr,
         stopwatch: &StopwatchMetrics,
-    ) -> Result<usize, StoreError> {
+    ) -> Result<Vec<(String, Vid)>, StoreError> {
         let section = stopwatch.start_section("check_interface_entity_uniqueness");
         for (key, _) in data.iter() {
             // WARNING: This will potentially execute 2 queries for each entity key.
@@ -408,7 +421,7 @@ impl DeploymentStore {
         layout: &'a Layout,
         ptr: &BlockPtr,
         stopwatch: &StopwatchMetrics,
-    ) -> Result<usize, StoreError> {
+    ) -> Result<Vec<(String, Vid)>, StoreError> {
         let section = stopwatch.start_section("check_interface_entity_uniqueness");
         for (key, _) in data.iter() {
             // WARNING: This will potentially execute 2 queries for each entity key.
@@ -863,7 +876,7 @@ impl DeploymentStore {
         stopwatch: StopwatchMetrics,
         data_sources: &[StoredDynamicDataSource],
         deterministic_errors: &[SubgraphError],
-    ) -> Result<StoreEvent, StoreError> {
+    ) -> Result<(StoreEvent, Vec<(EntityKey, Vid)>), StoreError> {
         // All operations should apply only to data or metadata for this subgraph
         if mods
             .iter()
@@ -881,7 +894,7 @@ impl DeploymentStore {
             self.get_conn()?
         };
 
-        let event = conn.transaction(|| -> Result<_, StoreError> {
+        conn.transaction(|| -> Result<_, StoreError> {
             // Emit a store event for the changes we are about to make. We
             // wait with sending it until we have done all our other work
             // so that we do not hold a lock on the notification queue
@@ -891,7 +904,7 @@ impl DeploymentStore {
             // Make the changes
             let layout = self.layout(&conn, site.clone())?;
             let section = stopwatch.start_section("apply_entity_modifications");
-            let count = self.apply_entity_modifications(
+            let (count, vid_map) = self.apply_entity_modifications(
                 &conn,
                 layout.as_ref(),
                 mods,
@@ -925,10 +938,8 @@ impl DeploymentStore {
                 }
             }
 
-            Ok(event)
-        })?;
-
-        Ok(event)
+            Ok((event, vid_map))
+        })
     }
 
     fn rewind_with_conn(

store/postgres/src/relational.rs

@@ -18,6 +18,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::convert::{From, TryFrom};
 use std::env;
 use std::fmt::{self, Write};
+use std::num::NonZeroU64;
 use std::str::FromStr;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
@@ -568,21 +569,21 @@ impl Layout {
         entities: &'a mut [(&'a EntityKey, Cow<'a, Entity>)],
         block: BlockNumber,
         stopwatch: &StopwatchMetrics,
-    ) -> Result<usize, StoreError> {
+    ) -> Result<Vec<(String, Vid)>, StoreError> {
         let table = self.table_for_entity(entity_type)?;
         let _section = stopwatch.start_section("insert_modification_insert_query");
-        let mut count = 0;
+        let mut vid_map = Vec::with_capacity(entities.len());
         // Each operation must respect the maximum number of bindings allowed in PostgreSQL queries,
         // so we need to act in chunks whose size is defined by the number of entities times the
         // number of attributes each entity type has.
         // We add 1 to account for the `block_range` bind parameter
         let chunk_size = POSTGRES_MAX_PARAMETERS / (table.columns.len() + 1);
         for chunk in entities.chunks_mut(chunk_size) {
-            count += InsertQuery::new(table, chunk, block)?
-                .get_results(conn)
-                .map(|ids| ids.len())?
+            for red in InsertQuery::new(table, chunk, block)?.get_results(conn)? {
+                vid_map.push((red.id, NonZeroU64::new(red.vid as u64)))
+            }
         }
-        Ok(count)
+        Ok(vid_map)
     }
 
     pub fn conflicting_entity(
@@ -684,25 +685,27 @@ impl Layout {
         vids: &'a [Vid],
         block: BlockNumber,
         stopwatch: &StopwatchMetrics,
-    ) -> Result<usize, StoreError> {
+    ) -> Result<Vec<(String, Vid)>, StoreError> {
         let table = self.table_for_entity(&entity_type)?;
 
         let section = stopwatch.start_section("update_modification_clamp_range_query");
         ClampRangeQuery::new(table, &vids, block)?.execute(conn)?;
         section.end();
 
         let _section = stopwatch.start_section("update_modification_insert_query");
-        let mut count = 0;
 
         // Each operation must respect the maximum number of bindings allowed in PostgreSQL queries,
         // so we need to act in chunks whose size is defined by the number of entities times the
         // number of attributes each entity type has.
         // We add 1 to account for the `block_range` bind parameter
         let chunk_size = POSTGRES_MAX_PARAMETERS / (table.columns.len() + 1);
+        let mut vid_map = Vec::with_capacity(entities.len());
         for chunk in entities.chunks_mut(chunk_size) {
-            count += InsertQuery::new(table, chunk, block)?.execute(conn)?;
+            for red in InsertQuery::new(table, chunk, block)?.get_results(conn)? {
+                vid_map.push((red.id, NonZeroU64::new(red.vid as u64)))
+            }
         }
-        Ok(count)
+        Ok(vid_map)
     }
 
     pub fn delete(

store/postgres/src/relational_queries.rs

@@ -1384,7 +1384,7 @@ impl<'a> QueryFragment<Pg> for InsertQuery<'a> {
         }
         out.push_sql("\nreturning ");
         out.push_sql(PRIMARY_KEY_COLUMN);
-        out.push_sql("::text");
+        out.push_sql("::text, vid");
 
         Ok(())
     }
@@ -2647,6 +2647,8 @@ impl<'a, Conn> RunQueryDsl<Conn> for ClampRangeQuery<'a> {}
 pub struct ReturnedEntityData {
     #[sql_type = "Text"]
     pub id: String,
+    #[sql_type = "BigInt"]
+    pub vid: i64,
 }
 
 impl ReturnedEntityData {
@@ -2689,7 +2691,7 @@ impl<'a> QueryFragment<Pg> for RevertRemoveQuery<'a> {
         out.push_bind_param::<Integer, _>(&self.block)?;
         out.push_sql("\nreturning ");
         out.push_sql(PRIMARY_KEY_COLUMN);
-        out.push_sql("::text");
+        out.push_sql("::text, vid");
         Ok(())
     }
 }
@@ -2760,7 +2762,7 @@ impl<'a> QueryFragment<Pg> for RevertClampQuery<'a> {
         out.push_sql("), 2147483647) < 2147483647");
         out.push_sql("\nreturning ");
         out.push_sql(PRIMARY_KEY_COLUMN);
-        out.push_sql("::text");
+        out.push_sql("::text, vid");
         Ok(())
     }
 }

store/postgres/src/subgraph_store.rs

@@ -17,7 +17,10 @@ use graph::{
     constraint_violation,
     data::query::QueryTarget,
     data::subgraph::schema::SubgraphError,
-    data::{store::EntityVersion, subgraph::status},
+    data::{
+        store::{EntityVersion, Vid},
+        subgraph::status,
+    },
     prelude::StoreEvent,
     prelude::SubgraphDeploymentEntity,
     prelude::{
@@ -1240,13 +1243,13 @@ impl WritableStoreTrait for WritableStore {
         stopwatch: StopwatchMetrics,
         data_sources: Vec<StoredDynamicDataSource>,
         deterministic_errors: Vec<SubgraphError>,
-    ) -> Result<(), StoreError> {
+    ) -> Result<Vec<(EntityKey, Vid)>, StoreError> {
         assert!(
             same_subgraph(&mods, &self.site.deployment),
             "can only transact operations within one shard"
         );
         self.retry("transact_block_operations", move || {
-            let event = self.writable.transact_block_operations(
+            let (event, vid_map) = self.writable.transact_block_operations(
                 self.site.clone(),
                 &block_ptr_to,
                 firehose_cursor.as_deref(),
@@ -1257,7 +1260,8 @@ impl WritableStoreTrait for WritableStore {
             )?;
 
             let _section = stopwatch.start_section("send_store_event");
-            self.try_send_store_event(event)
+            self.try_send_store_event(event)?;
+            Ok(vid_map)
         })
     }
 

store/postgres/tests/relational.rs

@@ -235,7 +235,8 @@ fn insert_entity(
             0,
             &MOCK_STOPWATCH,
         )
-        .expect(&errmsg);
+        .expect(&errmsg)
+        .len();
     assert_eq!(inserted, entities_with_keys_owned.len());
     let ids_for_type = {
         let ids: Vec<_> = entities_with_keys_owned
@@ -288,7 +289,8 @@ fn update_entity(
             0,
             &MOCK_STOPWATCH,
         )
-        .expect(&errmsg);
+        .expect(&errmsg)
+        .len();
     assert_eq!(updated, 1);
 }