graph: ensure removal paths go through remove_entity

Author: MoonBoi9001

graph/src/components/store/write.rs

@@ -498,31 +498,29 @@ impl RowGroup {
     }
 
     pub fn track_metrics(&self, metrics: &mut BlockStateMetrics) {
+        // Track entity count changes
         let changes: Vec<i32> = self
             .rows
             .iter()
             .map(|row| row.entity_count_change())
             .collect();
+        metrics.track_entity_count_change_batch(&self.entity_type, &changes);
 
-        let (writes, removals): (Vec<_>, Vec<_>) = self
+        // Track writes only
+        let writes: Vec<Entity> = self
             .rows
             .iter()
             .filter_map(|row| match row {
-                EntityModification::Insert { data, .. }
-                | EntityModification::Overwrite { data, .. } => Some((data, false)),
+                EntityModification::Insert { data, .. } 
+                | EntityModification::Overwrite { data, .. } => Some(data.as_ref().clone()),
                 EntityModification::Remove { .. } => None,
             })
-            .unzip();
-
-        metrics.track_entity_count_change_batch(&self.entity_type, &changes);
+            .collect();
 
         if !writes.is_empty() {
             metrics.track_entity_write_batch(&self.entity_type, &writes);
             metrics.track_storage_size_change_batch(&self.entity_type, &writes, false);
         }
-        if !removals.is_empty() {
-            metrics.track_storage_size_change_batch(&self.entity_type, &removals, true);
-        }
     }
 }
 

runtime/wasm/src/host_exports.rs

@@ -1249,6 +1249,24 @@ impl HostExports {
             .map(|mut tokens| tokens.pop().unwrap())
             .context("Failed to decode")
     }
+
+    fn remove_entity(&mut self, key: &EntityKey) -> Result<(), HostExportError> {
+        let entity_type = key.entity_type.clone();
+        
+        if let Some(entity) = state.entity_cache.get(&key, GetScope::Store)? {
+            // Track removal metrics while we still have the entity data
+            state
+                .metrics
+                .track_storage_size_change(&entity_type, &entity, true);
+            state.metrics.track_entity_count_change(&entity_type, -1);
+        }
+
+        state.modifications.push(EntityModification::remove(
+            key.clone(),
+            state.block.number,
+        ));
+        Ok(())
+    }
 }
 
 fn string_to_h160(string: &str) -> Result<H160, DeterministicHostError> {