core: thread EntityLfuCache instead of EntityCache across blocks

The AMP runner previously created an EntityCache with a fake
SeqGenerator(0) and shared it across block boundaries, even though an
EntityCache should always be tied to a single block. Now only the
LfuCache is threaded through, and EntityCache is created locally in
process_record_batch_group with the correct block's SeqGenerator.

Author: lutter

core/src/amp_subgraph/runner/data_processing.rs

@@ -11,22 +11,22 @@ use graph::{
     },
     blockchain::block_stream::FirehoseCursor,
     cheap_clone::CheapClone,
-    components::store::{EntityCache, ModificationsAndCache, SeqGenerator},
+    components::store::{EntityCache, EntityLfuCache, ModificationsAndCache, SeqGenerator},
 };
 use slog::{debug, trace};
 
 use super::{data_stream::TablePtr, Compat, Context, Error};
 
 pub(super) async fn process_record_batch_groups<AC>(
     cx: &mut Context<AC>,
-    mut entity_cache: EntityCache,
+    mut entity_lfu_cache: EntityLfuCache,
     record_batch_groups: RecordBatchGroups,
     stream_table_ptr: Arc<[TablePtr]>,
     latest_block: BlockNumber,
-) -> Result<EntityCache, Error> {
+) -> Result<EntityLfuCache, Error> {
     if record_batch_groups.is_empty() {
         debug!(cx.logger, "Received no record batch groups");
-        return Ok(entity_cache);
+        return Ok(entity_lfu_cache);
     }
 
     let from_block = record_batch_groups
@@ -50,9 +50,9 @@ pub(super) async fn process_record_batch_groups<AC>(
             "record_batches_count" => record_batch_group.record_batches.len()
         );
 
-        entity_cache = process_record_batch_group(
+        entity_lfu_cache = process_record_batch_group(
             cx,
-            entity_cache,
+            entity_lfu_cache,
             block_number,
             block_hash,
             record_batch_group,
@@ -79,32 +79,36 @@ pub(super) async fn process_record_batch_groups<AC>(
         "to_block" => to_block
     );
 
-    Ok(entity_cache)
+    Ok(entity_lfu_cache)
 }
 
 async fn process_record_batch_group<AC>(
     cx: &mut Context<AC>,
-    mut entity_cache: EntityCache,
+    entity_lfu_cache: EntityLfuCache,
     block_number: BlockNumber,
     block_hash: BlockHash,
     record_batch_group: RecordBatchGroup,
     stream_table_ptr: &[TablePtr],
     latest_block: BlockNumber,
-) -> Result<EntityCache, Error> {
+) -> Result<EntityLfuCache, Error> {
     let _section = cx
         .metrics
         .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() {
         debug!(cx.logger, "Record batch group is empty");
-        return Ok(entity_cache);
+        return Ok(entity_lfu_cache);
     }
 
+    let mut entity_cache = EntityCache::with_current(
+        cx.store.cheap_clone(),
+        entity_lfu_cache,
+        SeqGenerator::new(block_number.compat()),
+    );
+
     let block_timestamp = if cx.manifest.schema.has_aggregations() {
         decode_block_timestamp(&record_batches)
             .map_err(|e| e.context("failed to decode block timestamp"))?
@@ -135,7 +139,6 @@ 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,
@@ -171,11 +174,7 @@ async fn process_record_batch_group<AC>(
         cx.metrics.deployment_synced.record(true);
     }
 
-    Ok(EntityCache::with_current(
-        cx.store.cheap_clone(),
-        entity_lfu_cache,
-        vid_gen,
-    ))
+    Ok(entity_lfu_cache)
 }
 
 async fn process_record_batch<AC>(

core/src/amp_subgraph/runner/mod.rs

@@ -11,10 +11,8 @@ use std::time::{Duration, Instant};
 use anyhow::Result;
 use futures::StreamExt;
 use graph::{
-    amp::Client,
-    cheap_clone::CheapClone,
-    components::store::{EntityCache, SeqGenerator},
-    data::subgraph::schema::SubgraphError,
+    amp::Client, cheap_clone::CheapClone, components::store::EntityLfuCache,
+    data::subgraph::schema::SubgraphError, util::lfu_cache::LfuCache,
 };
 use slog::{debug, error, warn};
 use tokio_util::sync::CancellationToken;
@@ -106,18 +104,15 @@ where
             .update(latest_block.min(cx.end_block()));
 
         let mut deployment_is_failed = cx.store.health().await?.is_failed();
-        // 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 entity_lfu_cache: EntityLfuCache = LfuCache::new();
         let mut stream = new_data_stream(cx, latest_block);
 
         while let Some(result) = stream.next().await {
             let (record_batch_groups, stream_table_ptr) = result?;
 
-            entity_cache = process_record_batch_groups(
+            entity_lfu_cache = process_record_batch_groups(
                 cx,
-                entity_cache,
+                entity_lfu_cache,
                 record_batch_groups,
                 stream_table_ptr,
                 latest_block,