graph, runtime, core: introduce VidGenerator for shared VID sequences

Replace the per-EntityCache vid_seq and seq fields with a shared
VidGenerator backed by Arc<AtomicU32>. Created once per block and
cheaply cloned across all EntityCache instances, this eliminates VID
collisions in ipfs.map() callbacks and offchain triggers without
manual sequence threading.

Author: lutter

core/src/amp_subgraph/runner/data_processing.rs

@@ -11,7 +11,7 @@ use graph::{
     },
     blockchain::block_stream::FirehoseCursor,
     cheap_clone::CheapClone,
-    components::store::{EntityCache, ModificationsAndCache},
+    components::store::{EntityCache, ModificationsAndCache, SeqGenerator},
 };
 use slog::{debug, trace};
 
@@ -96,6 +96,8 @@ async fn process_record_batch_group<AC>(
         .stopwatch
         .start_section("process_record_batch_group");
 
+    entity_cache.seq_gen = SeqGenerator::new(block_number.compat());
+
     let RecordBatchGroup { record_batches } = record_batch_group;
 
     if record_batches.is_empty() {
@@ -121,7 +123,6 @@ async fn process_record_batch_group<AC>(
         process_record_batch(
             cx,
             &mut entity_cache,
-            block_number,
             record_batch,
             stream_table_ptr[stream_index],
         )
@@ -134,6 +135,7 @@ async fn process_record_batch_group<AC>(
     }
 
     let section = cx.metrics.stopwatch.start_section("as_modifications");
+    let vid_gen = entity_cache.vid_gen();
     let ModificationsAndCache {
         modifications,
         entity_lfu_cache,
@@ -172,13 +174,13 @@ async fn process_record_batch_group<AC>(
     Ok(EntityCache::with_current(
         cx.store.cheap_clone(),
         entity_lfu_cache,
+        vid_gen,
     ))
 }
 
 async fn process_record_batch<AC>(
     cx: &mut Context<AC>,
     entity_cache: &mut EntityCache,
-    block_number: BlockNumber,
     record_batch: RecordBatch,
     (i, j): TablePtr,
 ) -> Result<(), Error> {
@@ -209,13 +211,11 @@ async fn process_record_batch<AC>(
         let key = match key {
             Some(key) => key,
             None => {
-                let entity_id = entity_cache
-                    .generate_id(id_type, block_number.compat())
-                    .map_err(|e| {
-                        Error::Deterministic(e.context(format!(
-                            "failed to generate a new id for an entity of type '{entity_name}'"
-                        )))
-                    })?;
+                let entity_id = entity_cache.seq_gen.id(id_type).map_err(|e| {
+                    Error::Deterministic(e.context(format!(
+                        "failed to generate a new id for an entity of type '{entity_name}'"
+                    )))
+                })?;
 
                 entity_data.push(("id".into(), entity_id.clone().into()));
                 entity_type.key(entity_id)
@@ -229,14 +229,11 @@ async fn process_record_batch<AC>(
             )))
         })?;
 
-        entity_cache
-            .set(key, entity, block_number.compat(), None)
-            .await
-            .map_err(|e| {
-                Error::Deterministic(e.context(format!(
-                    "failed to store a new entity of type '{entity_name}' with id '{entity_id}'"
-                )))
-            })?;
+        entity_cache.set(key, entity, None).await.map_err(|e| {
+            Error::Deterministic(e.context(format!(
+                "failed to store a new entity of type '{entity_name}' with id '{entity_id}'"
+            )))
+        })?;
     }
 
     Ok(())

core/src/amp_subgraph/runner/mod.rs

@@ -11,7 +11,9 @@ use std::time::{Duration, Instant};
 use anyhow::Result;
 use futures::StreamExt;
 use graph::{
-    amp::Client, cheap_clone::CheapClone, components::store::EntityCache,
+    amp::Client,
+    cheap_clone::CheapClone,
+    components::store::{EntityCache, SeqGenerator},
     data::subgraph::schema::SubgraphError,
 };
 use slog::{debug, error, warn};
@@ -104,7 +106,10 @@ where
             .update(latest_block.min(cx.end_block()));
 
         let mut deployment_is_failed = cx.store.health().await?.is_failed();
-        let mut entity_cache = EntityCache::new(cx.store.cheap_clone());
+        // The SeqGenerator gets replaced with one for the correct block
+        // number in `process_record_batch_groups`, so the initial value
+        // doesn't matter much.
+        let mut entity_cache = EntityCache::new(cx.store.cheap_clone(), SeqGenerator::new(0));
         let mut stream = new_data_stream(cx, latest_block);
 
         while let Some(result) = stream.next().await {

core/src/subgraph/runner/mod.rs

@@ -17,7 +17,9 @@ use graph::blockchain::{
     Block, BlockTime, Blockchain, DataSource as _, SubgraphFilter, Trigger, TriggerFilter as _,
     TriggerFilterWrapper,
 };
-use graph::components::store::{EmptyStore, GetScope, ReadStore, StoredDynamicDataSource};
+use graph::components::store::{
+    EmptyStore, GetScope, ReadStore, SeqGenerator, StoredDynamicDataSource,
+};
 use graph::components::subgraph::InstanceDSTemplate;
 use graph::components::trigger_processor::RunnableTriggers;
 use graph::components::{
@@ -1019,9 +1021,9 @@ where
             .ready_offchain_events()
             .non_deterministic()?;
 
-        let onchain_vid_seq = block_state.entity_cache.vid_seq;
+        let vid_gen = block_state.seq_gen();
         let (offchain_mods, processed_offchain_data_sources, persisted_off_chain_data_sources) =
-            self.handle_offchain_triggers(offchain_events, block, onchain_vid_seq)
+            self.handle_offchain_triggers(offchain_events, block, vid_gen)
                 .await
                 .non_deterministic()?;
 
@@ -1066,9 +1068,11 @@ where
         // Causality region for onchain triggers.
         let causality_region = PoICausalityRegion::from_network(&self.inputs.network);
 
+        let vid_gen = SeqGenerator::new(block_ptr.number);
         let mut block_state = BlockState::new(
             self.inputs.store.clone(),
             std::mem::take(&mut self.state.entity_lfu_cache),
+            vid_gen,
         );
 
         let _section = self
@@ -1471,7 +1475,7 @@ where
         &mut self,
         triggers: Vec<offchain::TriggerData>,
         block: &Arc<C::Block>,
-        mut next_vid_seq: u32,
+        vid_gen: SeqGenerator,
     ) -> Result<
         (
             Vec<EntityModification>,
@@ -1488,12 +1492,11 @@ where
             // Using an `EmptyStore` and clearing the cache for each trigger is a makeshift way to
             // get causality region isolation.
             let schema = ReadStore::input_schema(&self.inputs.store);
-            let mut block_state = BlockState::new(EmptyStore::new(schema), LfuCache::new());
-
-            // Continue the vid sequence from the previous trigger (or from
-            // onchain processing) so that each offchain trigger does not
-            // reset to RESERVED_VIDS and produce duplicate VIDs.
-            block_state.entity_cache.vid_seq = next_vid_seq;
+            let mut block_state = BlockState::new(
+                EmptyStore::new(schema),
+                LfuCache::new(),
+                vid_gen.cheap_clone(),
+            );
 
             // PoI ignores offchain events.
             // See also: poi-ignores-offchain
@@ -1560,10 +1563,6 @@ where
                 return Err(anyhow!("{}", err));
             }
 
-            // Carry forward the vid sequence so the next iteration doesn't
-            // reset to RESERVED_VIDS and produce duplicate VIDs.
-            next_vid_seq = block_state.entity_cache.vid_seq;
-
             mods.extend(
                 block_state
                     .entity_cache
@@ -1685,7 +1684,6 @@ async fn update_proof_of_indexing(
         key: EntityKey,
         digest: Bytes,
         block_time: BlockTime,
-        block: BlockNumber,
     ) -> Result<(), Error> {
         let digest_name = entity_cache.schema.poi_digest();
         let mut data = vec![
@@ -1700,12 +1698,11 @@ async fn update_proof_of_indexing(
             data.push((entity_cache.schema.poi_block_time(), block_time));
         }
         let poi = entity_cache.make_entity(data)?;
-        entity_cache.set(key, poi, block, None).await
+        entity_cache.set(key, poi, None).await
     }
 
     let _section_guard = stopwatch.start_section("update_proof_of_indexing");
 
-    let block_number = proof_of_indexing.get_block();
     let mut proof_of_indexing = proof_of_indexing.take();
 
     for (causality_region, stream) in proof_of_indexing.drain() {
@@ -1742,7 +1739,6 @@ async fn update_proof_of_indexing(
             entity_key,
             updated_proof_of_indexing,
             block_time,
-            block_number,
         )
         .await?;
     }

graph/src/components/store/entity_cache.rs

@@ -2,6 +2,7 @@ use anyhow::{anyhow, bail};
 use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::fmt::{self, Debug};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::Arc;
 
 use crate::cheap_clone::CheapClone;
@@ -21,6 +22,52 @@ pub type EntityLfuCache = LfuCache<EntityKey, Option<Arc<Entity>>>;
 // Currently none is used, but lets reserve a few more.
 const RESERVED_VIDS: u32 = 100;
 
+/// Shared generator for VID and entity ID sequences within a block.
+///
+/// Created once per block and shared (via `Arc`) across all `EntityCache`
+/// instances that operate on the same block. This prevents VID collisions
+/// when multiple isolated caches (e.g. ipfs.map callbacks, offchain
+/// triggers) write entities in the same block.
+#[derive(Clone, Debug)]
+pub struct SeqGenerator {
+    block: BlockNumber,
+    vid_seq: Arc<AtomicU32>,
+    id_seq: Arc<AtomicU32>,
+}
+
+impl CheapClone for SeqGenerator {
+    fn cheap_clone(&self) -> Self {
+        self.clone()
+    }
+}
+
+impl SeqGenerator {
+    pub fn new(block: BlockNumber) -> Self {
+        SeqGenerator {
+            block,
+            vid_seq: Arc::new(AtomicU32::new(RESERVED_VIDS)),
+            id_seq: Arc::new(AtomicU32::new(0)),
+        }
+    }
+
+    /// Return the next VID. The VID encodes the block number in the upper
+    /// 32 bits and a monotonically increasing sequence in the lower 32
+    /// bits.
+    pub fn vid(&self) -> i64 {
+        let seq = self.vid_seq.fetch_add(1, Ordering::Relaxed);
+        ((self.block as i64) << 32) + seq as i64
+    }
+
+    /// Generate the next entity ID for the given ID type. The ID encodes
+    /// the block number in the upper 32 bits and a monotonically
+    /// increasing sequence in the lower 32 bits.
+    pub fn id(&self, id_type: IdType) -> anyhow::Result<Id> {
+        let seq = self.id_seq.fetch_add(1, Ordering::Relaxed);
+
+        id_type.generate_id(self.block, seq)
+    }
+}
+
 /// The scope in which the `EntityCache` should perform a `get` operation
 pub enum GetScope {
     /// Get from all previously stored entities in the store
@@ -103,16 +150,8 @@ pub struct EntityCache {
 
     pub schema: InputSchema,
 
-    /// A sequence number for generating entity IDs. We use one number for
-    /// all id's as the id's are scoped by block and a u32 has plenty of
-    /// room for all changes in one block. To ensure reproducability of
-    /// generated IDs, the `EntityCache` needs to be newly instantiated for
-    /// each block
-    seq: u32,
-
-    // Sequence number of the next VID value for this block. The value written
-    // in the database consist of a block number and this SEQ number.
-    pub vid_seq: u32,
+    /// Shared sequence generator for VIDs and entity IDs within a block.
+    pub seq_gen: SeqGenerator,
 }
 
 impl Debug for EntityCache {
@@ -131,16 +170,15 @@ pub struct ModificationsAndCache {
 }
 
 impl EntityCache {
-    pub fn new(store: Arc<dyn s::ReadStore>) -> Self {
+    pub fn new(store: Arc<dyn s::ReadStore>, seq_gen: SeqGenerator) -> Self {
         Self {
             current: LfuCache::new(),
             updates: HashMap::new(),
             handler_updates: HashMap::new(),
             in_handler: false,
             schema: store.input_schema(),
             store,
-            seq: 0,
-            vid_seq: RESERVED_VIDS,
+            seq_gen,
         }
     }
 
@@ -152,19 +190,26 @@ impl EntityCache {
         self.schema.make_entity(iter)
     }
 
-    pub fn with_current(store: Arc<dyn s::ReadStore>, current: EntityLfuCache) -> EntityCache {
+    pub fn with_current(
+        store: Arc<dyn s::ReadStore>,
+        current: EntityLfuCache,
+        seq_gen: SeqGenerator,
+    ) -> EntityCache {
         EntityCache {
             current,
             updates: HashMap::new(),
             handler_updates: HashMap::new(),
             in_handler: false,
             schema: store.input_schema(),
             store,
-            seq: 0,
-            vid_seq: RESERVED_VIDS,
+            seq_gen,
         }
     }
 
+    pub fn seq_gen(&self) -> SeqGenerator {
+        self.seq_gen.cheap_clone()
+    }
+
     pub(crate) fn enter_handler(&mut self) {
         assert!(!self.in_handler);
         self.in_handler = true;
@@ -368,7 +413,6 @@ impl EntityCache {
         &mut self,
         key: EntityKey,
         entity: Entity,
-        block: BlockNumber,
         write_capacity_remaining: Option<&mut usize>,
     ) -> Result<(), anyhow::Error> {
         // check the validate for derived fields
@@ -386,9 +430,7 @@ impl EntityCache {
             *write_capacity_remaining -= weight;
         }
 
-        // The next VID is based on a block number and a sequence within the block
-        let vid = ((block as i64) << 32) + self.vid_seq as i64;
-        self.vid_seq += 1;
+        let vid = self.seq_gen.vid();
         let mut entity = entity;
         let old_vid = entity.set_vid(vid).expect("the vid should be set");
         // Make sure that there was no VID previously set for this entity.
@@ -457,16 +499,6 @@ impl EntityCache {
         for (key, op) in other.updates {
             self.entity_op(key, op);
         }
-        // Carry forward vid_seq to prevent VID collisions when the caller
-        // continues writing entities after merging.
-        self.vid_seq = self.vid_seq.max(other.vid_seq);
-    }
-
-    /// Generate an id.
-    pub fn generate_id(&mut self, id_type: IdType, block: BlockNumber) -> anyhow::Result<Id> {
-        let id = id_type.generate_id(block, self.seq)?;
-        self.seq += 1;
-        Ok(id)
     }
 
     /// Return the changes that have been made via `set` and `remove` as

graph/src/components/store/mod.rs

@@ -9,7 +9,9 @@ use diesel::pg::Pg;
 use diesel::serialize::{Output, ToSql};
 use diesel::sql_types::Integer;
 use diesel_derives::{AsExpression, FromSqlRow};
-pub use entity_cache::{EntityCache, EntityLfuCache, GetScope, ModificationsAndCache};
+pub use entity_cache::{
+    EntityCache, EntityLfuCache, GetScope, ModificationsAndCache, SeqGenerator,
+};
 use slog::Logger;
 use tokio_stream::wrappers::ReceiverStream;