feature(collator): remove cumulative stats recalculation

Author: drmick

collator/src/collator/messages_reader/internals_reader.rs

@@ -356,7 +356,6 @@ impl<V: InternalMessageValue> InternalsPartitionReader<V> {
 
             ranges.push(QueueShardBoundedRange {
                 shard_ident: *shard_id,
-                // TODO: may be we should pass `from` here because we should load stats for the full range
                 from: Bound::Excluded(shard_reader_state.current_position),
                 to: Bound::Included(shard_range_to),
             });

collator/src/collator/messages_reader/mod.rs

@@ -217,7 +217,6 @@ impl<V: InternalMessageValue> MessagesReader<V> {
                     .copied()
                     .collect();
 
-                cumulative_stats_just_loaded = true;
                 match (previous_cumulative_statistics, cx.part_stat_ranges) {
                     (Some(mut previous_cumulative_statistics), Some(part_stat_ranges)) => {
                         // update all_shards_processed_to_by_partitions and truncate processed data
@@ -235,11 +234,13 @@ impl<V: InternalMessageValue> MessagesReader<V> {
                         previous_cumulative_statistics
                     }
                     _ => {
-                        let mut inner_cumulative_statistics =
-                            CumulativeStatistics::new(params.all_shards_processed_to_by_partitions);
+                        cumulative_stats_just_loaded = true;
+                        let mut inner_cumulative_statistics = CumulativeStatistics::new(
+                            cx.for_shard_id,
+                            params.all_shards_processed_to_by_partitions,
+                        );
                         inner_cumulative_statistics.load(
                             mq_adapter.clone(),
-                            &cx.for_shard_id,
                             &partitions,
                             cx.prev_state_gen_lt,
                             cx.mc_state_gen_lt,

collator/src/collator/types.rs

@@ -24,7 +24,6 @@ use tycho_block_util::state::{RefMcStateHandle, ShardStateStuff};
 use tycho_core::global_config::MempoolGlobalConfig;
 use tycho_executor::{AccountMeta, PublicLibraryChange, TransactionMeta};
 use tycho_network::PeerId;
-use tycho_util::metrics::HistogramGuard;
 use tycho_util::{DashMapEntry, FastDashMap, FastHashMap, FastHashSet};
 
 use super::do_collate::work_units::PrepareMsgGroupsWu;
@@ -1281,6 +1280,9 @@ pub struct QueueStatisticsWithRemaning {
 }
 
 pub struct CumulativeStatistics {
+    /// Cumulative statistics created for this shard. When this shard reads statistics, it decrements `remaining messages`
+    /// Another shard can decrements the statistics remaining only by using `update_processed_to_by_partitions`
+    for_shard: ShardIdent,
     /// Actual processed to info for master and all shards
     all_shards_processed_to_by_partitions: FastHashMap<ShardIdent, (bool, ProcessedToByPartitions)>,
 
@@ -1290,38 +1292,35 @@ pub struct CumulativeStatistics {
 
     /// The final aggregated statistics (across all shards) by partitions.
     result: FastHashMap<QueuePartitionIdx, QueueStatisticsWithRemaning>,
-
-    /// A flag indicating that data has changed, and we need to recalculate before returning `result`.
-    dirty: bool,
 }
 
 impl CumulativeStatistics {
     pub fn new(
+        for_shard: ShardIdent,
         all_shards_processed_to_by_partitions: FastHashMap<
             ShardIdent,
             (bool, ProcessedToByPartitions),
         >,
     ) -> Self {
         Self {
+            for_shard,
             all_shards_processed_to_by_partitions,
             shards_stats_by_partitions: Default::default(),
             result: Default::default(),
-            dirty: false,
         }
     }
 
     /// Create range and full load statistics and store it
     pub fn load<V: InternalMessageValue>(
         &mut self,
         mq_adapter: Arc<dyn MessageQueueAdapter<V>>,
-        current_shard: &ShardIdent,
         partitions: &FastHashSet<QueuePartitionIdx>,
         prev_state_gen_lt: Lt,
         mc_state_gen_lt: Lt,
         mc_top_shards_end_lts: &FastHashMap<ShardIdent, Lt>,
     ) -> Result<()> {
         let ranges = Self::compute_cumulative_stats_ranges(
-            current_shard,
+            &self.for_shard,
             &self.all_shards_processed_to_by_partitions,
             prev_state_gen_lt,
             mc_state_gen_lt,
@@ -1345,7 +1344,6 @@ impl CumulativeStatistics {
         partitions: &FastHashSet<QueuePartitionIdx>,
         ranges: Vec<QueueShardBoundedRange>,
     ) -> Result<()> {
-        self.dirty = true;
         tracing::trace!(
             target: tracing_targets::COLLATOR,
             "cumulative_stats_partial_ranges: {:?}",
@@ -1407,7 +1405,6 @@ impl CumulativeStatistics {
             });
 
         self.all_shards_processed_to_by_partitions = new_pt;
-        self.dirty = true;
     }
 
     fn compute_cumulative_stats_ranges(
@@ -1469,15 +1466,23 @@ impl CumulativeStatistics {
             }
         }
 
+        let entry = self
+            .result
+            .entry(partition)
+            .or_insert_with(|| QueueStatisticsWithRemaning {
+                initial_stats: QueueStatistics::default(),
+                remaning_stats: ConcurrentQueueStatistics::default(),
+            });
+        entry.initial_stats.append(&diff_partition_stats);
+        entry.remaning_stats.append(&diff_partition_stats);
+
         // finally add weeded stats
         self.add_diff_partition_stats(
             partition,
             diff_shard,
             diff_max_message,
             diff_partition_stats,
         );
-
-        self.dirty = true;
     }
 
     fn add_diff_partition_stats(
@@ -1541,6 +1546,12 @@ impl CumulativeStatistics {
                                         cumulative_stats
                                             .initial_stats
                                             .decrement_for_account(dst_acc.clone(), *count);
+
+                                        if self.for_shard != dst_shard {
+                                            cumulative_stats
+                                                .remaning_stats
+                                                .decrement_for_account(dst_acc.clone(), *count);
+                                        }
                                         false
                                     } else {
                                         true
@@ -1571,15 +1582,12 @@ impl CumulativeStatistics {
     /// Returns  a reference to the aggregated stats by partitions.
     /// If the data is marked as dirty, it triggers a lazy recalculation first.
     pub fn result(&mut self) -> &FastHashMap<QueuePartitionIdx, QueueStatisticsWithRemaning> {
-        self.ensure_finalized();
         &self.result
     }
 
     /// Calc aggregated stats among all partitions.
     /// If the data is marked as dirty, it triggers a lazy recalculation first.
     pub fn get_aggregated_result(&mut self) -> QueueStatistics {
-        self.ensure_finalized();
-
         let mut res: Option<QueueStatistics> = None;
         for stats in self.result.values() {
             if let Some(aggregated) = res.as_mut() {
@@ -1590,42 +1598,6 @@ impl CumulativeStatistics {
         }
         res.unwrap_or_default()
     }
-
-    /// A helper function to trigger a recalculation if `dirty` is set.
-    fn ensure_finalized(&mut self) {
-        if self.dirty {
-            self.recalculate();
-        }
-    }
-
-    /// Clears the existing result and aggregates all data from `shards_statistics`.
-    fn recalculate(&mut self) {
-        let _histogram = HistogramGuard::begin("tycho_do_collate_recalculate_statistics_time");
-        self.result.clear();
-
-        for shard_stats_by_partitions in self.shards_stats_by_partitions.values() {
-            for (&partition, diffs) in shard_stats_by_partitions {
-                let mut partition_stats = AccountStatistics::new();
-                for diff_stats in diffs.values() {
-                    for (account, &count) in diff_stats {
-                        *partition_stats.entry(account.clone()).or_default() += count;
-                    }
-                }
-                self.result
-                    .entry(partition)
-                    .and_modify(|stats| {
-                        stats.initial_stats.append(&partition_stats);
-                        stats.remaning_stats.append(&partition_stats);
-                    })
-                    .or_insert(QueueStatisticsWithRemaning {
-                        initial_stats: QueueStatistics::with_statistics(partition_stats.clone()),
-                        remaning_stats: ConcurrentQueueStatistics::with_statistics(partition_stats),
-                    });
-            }
-        }
-
-        self.dirty = false;
-    }
 }
 
 #[derive(Debug, Default, Clone)]
@@ -1634,12 +1606,6 @@ pub struct ConcurrentQueueStatistics {
 }
 
 impl ConcurrentQueueStatistics {
-    pub fn with_statistics(statistics: AccountStatistics) -> Self {
-        Self {
-            statistics: Arc::new(statistics.into_iter().collect()),
-        }
-    }
-
     pub fn statistics(&self) -> &FastDashMap<IntAddr, u64> {
         &self.statistics
     }
@@ -1694,7 +1660,6 @@ impl fmt::Debug for CumulativeStatistics {
             )
             .field("shards_stats_by_partitions", &shards_summary)
             .field("result", &format!("{} partitions", self.result.len()))
-            .field("dirty", &self.dirty)
             .finish()
     }
 }