refactor(collator): cumulative statistics code optimization

Author: drmick

collator/src/collator/do_collate/mod.rs

@@ -1222,73 +1222,27 @@ impl CollatorStdImpl {
         mc_data_stuff: &McDataStuff,
         block_id_short: &BlockIdShort,
     ) -> Result<Option<Vec<QueueShardBoundedRange>>> {
-        let mc_block_id = mc_data_stuff.current.block_id;
-
-        let prov_mc_state = if let Some(ref prev_mc_state) = mc_data_stuff.previous {
-            prev_mc_state
-        } else {
+        let Some(prev_state) = &mc_data_stuff.previous else {
             return Ok(None);
         };
 
-        let prev_mc_block_shards: FastHashMap<ShardIdent, ShardDescriptionShort> =
-            prov_mc_state.shards.iter().cloned().collect();
-
-        let mut ranges = Vec::<QueueShardBoundedRange>::new();
-
-        // Load max_message from masterchain block diff
-        let Some(master_max_msg) = self
-            .get_max_message(&mc_block_id)
-            .await
-            .context("loading diff for mc block")?
-        else {
-            return Ok(None);
-        };
-
-        ranges.push(QueueShardBoundedRange {
-            shard_ident: mc_block_id.shard,
-            from: Bound::Included(master_max_msg),
-            to: Bound::Included(master_max_msg),
-        });
-
-        // Iterate over all updated shard blocks and add their diff ranges
-        for (shard, current_descr) in mc_data_stuff
-            .current
-            .shards
-            .iter()
-            .filter(|(shard, d)| d.top_sc_block_updated && shard != &block_id_short.shard)
-        {
-            let Some(prev_descr) = prev_mc_block_shards.get(shard) else {
-                tracing::warn!(target: tracing_targets::COLLATOR, "prev_mc_block_shards not found: {shard:?}");
-                return Ok(None);
-            };
-
-            if prev_descr.seqno == 0 {
-                return Ok(None);
+        let prev_shards: FastHashMap<_, _> = prev_state.shards.iter().cloned().collect();
+
+        let mut shard_pairs = Vec::new();
+        for (shard, curr) in mc_data_stuff.current.shards.iter() {
+            if curr.top_sc_block_updated && *shard != block_id_short.shard {
+                if let Some(prev) = prev_shards.get(shard) {
+                    shard_pairs.push((
+                        *shard,
+                        prev.get_block_id(*shard),
+                        curr.get_block_id(*shard),
+                    ));
+                }
             }
-
-            let Some(first_diff_msg) = self
-                .get_max_message(&prev_descr.get_block_id(*shard))
-                .await
-                .context("loading first diff msg")?
-            else {
-                return Ok(None);
-            };
-            let Some(last_diff_msg) = self
-                .get_max_message(&current_descr.get_block_id(*shard))
-                .await
-                .context("loading last diff msg")?
-            else {
-                return Ok(None);
-            };
-
-            ranges.push(QueueShardBoundedRange {
-                shard_ident: *shard,
-                from: Bound::Excluded(first_diff_msg),
-                to: Bound::Included(last_diff_msg),
-            });
         }
 
-        Ok(Some(ranges))
+        self.build_ranges(&mc_data_stuff.current.block_id, shard_pairs)
+            .await
     }
 
     /// Collect ranges for loading cumulative statistics if collation block is master
@@ -1299,67 +1253,69 @@ impl CollatorStdImpl {
         mc_data_stuff: &McDataStuff,
         top_shard_blocks_info: Vec<TopBlockDescription>,
     ) -> Result<Option<Vec<QueueShardBoundedRange>>> {
-        let prev_mc_block_id = mc_data_stuff.current.block_id;
-
-        let prev_mc_block_shards: FastHashMap<ShardIdent, ShardDescriptionShort> =
-            mc_data_stuff.current.shards.iter().cloned().collect();
+        let prev_shards: FastHashMap<_, _> = mc_data_stuff.current.shards.iter().cloned().collect();
+
+        let mut shard_pairs = Vec::new();
+        for top in top_shard_blocks_info {
+            if let Some(prev) = prev_shards.get(&top.block_id.shard) {
+                shard_pairs.push((
+                    top.block_id.shard,
+                    prev.get_block_id(top.block_id.shard),
+                    top.block_id,
+                ));
+            }
+        }
 
-        let mut ranges = Vec::<QueueShardBoundedRange>::new();
+        self.build_ranges(&mc_data_stuff.current.block_id, shard_pairs)
+            .await
+    }
 
-        // Load max_message from masterchain block diff
+    async fn build_ranges(
+        &self,
+        master_block_id: &BlockId,
+        shard_pairs: Vec<(ShardIdent, BlockId, BlockId)>,
+    ) -> Result<Option<Vec<QueueShardBoundedRange>>> {
         let Some(master_max_msg) = self
-            .get_max_message(&prev_mc_block_id)
+            .get_diff_max_message(master_block_id)
             .await
             .context("loading diff for mc block")?
         else {
             return Ok(None);
         };
 
-        ranges.push(QueueShardBoundedRange {
-            shard_ident: prev_mc_block_id.shard,
+        // add mc block diff range
+        let mut ranges = vec![QueueShardBoundedRange {
+            shard_ident: master_block_id.shard,
             from: Bound::Included(master_max_msg),
             to: Bound::Included(master_max_msg),
-        });
-
-        // Iterate over all updated shard blocks and add their diff ranges
-        for top_block_description in top_shard_blocks_info {
-            let Some(prev_descr) = prev_mc_block_shards.get(&top_block_description.block_id.shard)
-            else {
-                tracing::warn!(target: tracing_targets::COLLATOR, "prev_mc_block_shards not found: {:?}", top_block_description.block_id.shard);
-                return Ok(None);
-            };
+        }];
 
-            if prev_descr.seqno == 0 {
+        for (shard, prev_id, curr_id) in shard_pairs {
+            if prev_id.seqno == 0 {
                 return Ok(None);
             }
 
-            let Some(first_diff_msg) = self
-                .get_max_message(&prev_descr.get_block_id(top_block_description.block_id.shard))
-                .await
-                .context("loading first diff msg")?
-            else {
-                return Ok(None);
-            };
-            let Some(last_diff_msg) = self
-                .get_max_message(&top_block_description.block_id)
-                .await
-                .context("loading last diff msg")?
-            else {
-                return Ok(None);
-            };
-
-            ranges.push(QueueShardBoundedRange {
-                shard_ident: top_block_description.block_id.shard,
-                from: Bound::Excluded(first_diff_msg),
-                to: Bound::Included(last_diff_msg),
-            });
+            let (first_msg, last_msg) = tokio::try_join!(
+                self.get_diff_max_message(&prev_id),
+                self.get_diff_max_message(&curr_id)
+            )
+            .context("loading shard diffs")?;
+
+            match (first_msg, last_msg) {
+                (Some(first), Some(last)) => ranges.push(QueueShardBoundedRange {
+                    shard_ident: shard,
+                    from: Bound::Excluded(first),
+                    to: Bound::Included(last),
+                }),
+                _ => return Ok(None),
+            }
         }
 
         Ok(Some(ranges))
     }
 
     /// Helper function that retrieves `max_message` from either MQ or state diff
-    async fn get_max_message(&self, block_id: &BlockId) -> Result<Option<QueueKey>> {
+    async fn get_diff_max_message(&self, block_id: &BlockId) -> Result<Option<QueueKey>> {
         if let Some(diff) =
             self.mq_adapter
                 .get_diff_info(&block_id.shard, block_id.seqno, DiffZone::Both)?

collator/src/collator/messages_reader/mod.rs

@@ -226,12 +226,12 @@ impl<V: InternalMessageValue> MessagesReader<V> {
                         );
 
                         // truncate data before range
-                        previous_cumulative_statistics.truncate_before(
-                            &cx.for_shard_id,
-                            cx.prev_state_gen_lt,
-                            cx.mc_state_gen_lt,
-                            &cx.mc_top_shards_end_lts.iter().copied().collect(),
-                        );
+                        // previous_cumulative_statistics.truncate_before(
+                        //     &cx.for_shard_id,
+                        //     cx.prev_state_gen_lt,
+                        //     cx.mc_state_gen_lt,
+                        //     &cx.mc_top_shards_end_lts.iter().copied().collect(),
+                        // );
 
                         // partial load statistics and enrich current value
                         previous_cumulative_statistics.load_partial(
@@ -240,6 +240,25 @@ impl<V: InternalMessageValue> MessagesReader<V> {
                             part_stat_ranges,
                         )?;
 
+                        let mut inner_cumulative_statistics =
+                            CumulativeStatistics::new(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,
+                            &cx.mc_top_shards_end_lts.iter().copied().collect(),
+                        )?;
+
+                        tracing::info!(target: "local_debug", "previous_cumulative_statistics {previous_cumulative_statistics:?}");
+
+                        if !previous_cumulative_statistics
+                            .shards_stats_eq(&inner_cumulative_statistics)
+                        {
+                            panic!("111")
+                        }
+
                         previous_cumulative_statistics
                     }
                     _ => {