feature(collator): remove processed data from reused statistics

Author: drmick

collator/src/collator/do_collate/finalize.rs

@@ -744,7 +744,6 @@ impl Phase<FinalizeState> {
                         .collation_data
                         .mc_shards_processed_to_by_partitions
                         .clone(),
-
                     prev_mc_block_id: Some(self.state.mc_data.block_id),
                 }))
             }

collator/src/collator/do_collate/mod.rs

@@ -105,6 +105,26 @@ impl CollatorStdImpl {
         };
         let created_by = author.to_bytes().into();
 
+        let is_first_block_after_prev_master = is_first_block_after_prev_master(
+            prev_shard_data.blocks_ids()[0], // TODO: consider split/merge
+            &mc_data.shards,
+        );
+        let part_stat_ranges = if is_first_block_after_prev_master && mc_data.block_id.seqno > 0 {
+            if next_block_id_short.is_masterchain() {
+                self.mc_compute_part_stat_ranges(
+                    &mc_data,
+                    &next_block_id_short,
+                    top_shard_blocks_info.clone().unwrap(),
+                )
+                .await?
+            } else {
+                self.compute_part_stat_ranges(&mc_data, &next_block_id_short)
+                    .await?
+            }
+        } else {
+            None
+        };
+
         let collation_data = self.create_collation_data(
             next_block_id_short,
             next_chain_time,
@@ -117,18 +137,6 @@ impl CollatorStdImpl {
         let anchors_cache = std::mem::take(&mut self.anchors_cache);
         let block_serializer_cache = self.block_serializer_cache.clone();
 
-        let is_first_block_after_prev_master = is_first_block_after_prev_master(
-            prev_shard_data.blocks_ids()[0], // TODO: consider split/merge
-            &mc_data.shards,
-        );
-
-        let part_stat_ranges = if is_first_block_after_prev_master && mc_data.block_id.seqno > 0 {
-            self.compute_part_stat_ranges(&mc_data, &collation_data.block_id_short)
-                .await?
-        } else {
-            None
-        };
-
         let state = Box::new(ActualState {
             collation_config,
             collation_data,
@@ -1199,27 +1207,34 @@ impl CollatorStdImpl {
         );
     }
 
+    /// Collect ranges for loading statistics if current block is first after master
+    /// Using previous masterchain block diff for range
+    /// and other shards diffs between previous masterchain block - 1 and previous masterchain block
     async fn compute_part_stat_ranges(
         &self,
         mc_data: &McData,
         block_id_short: &BlockIdShort,
     ) -> Result<Option<Vec<QueueShardBoundedRange>>> {
         let mc_block_id = mc_data.block_id;
+
         let prev_mc_block_id = mc_data
             .prev_mc_block_id
             .context("Prev MC block must be present")?;
 
         if prev_mc_block_id.seqno + 1 != mc_block_id.seqno {
-            bail!(
+            tracing::error!(
+                target: tracing_targets::COLLATOR,
                 "Prev MC block ID has an incorrect sequence. Prev: {prev_mc_block_id:?}. \
              Current: {mc_block_id:?}"
             );
+            return Ok(None);
         }
 
         let prev_mc_state = self
             .state_node_adapter
             .load_state(&prev_mc_block_id)
             .await?;
+
         let prev_mc_block_shards: FastHashMap<ShardIdent, ShardDescriptionShort> = prev_mc_state
             .state_extra()
             .cloned()?
@@ -1257,6 +1272,10 @@ impl CollatorStdImpl {
                 return Ok(None);
             };
 
+            if prev_descr.seqno == 0 {
+                return Ok(None);
+            }
+
             let Some(first_diff_msg) = self
                 .get_max_message(&prev_descr.get_block_id(*shard))
                 .await
@@ -1282,7 +1301,96 @@ impl CollatorStdImpl {
         Ok(Some(ranges))
     }
 
-    // Async helper function that retrieves max_message from either MQ or state diff
+    /// Collect ranges for loading cumulative statistics if collation block is master
+    /// Using range from previous masterchain block diff
+    /// and diffs between previous masterchain top blocks and current top blocks
+    async fn mc_compute_part_stat_ranges(
+        &self,
+        mc_data: &McData,
+        block_id_short: &BlockIdShort,
+        top_shard_blocks_info: Vec<TopBlockDescription>,
+    ) -> Result<Option<Vec<QueueShardBoundedRange>>> {
+        let prev_mc_block_id = mc_data.block_id;
+
+        if prev_mc_block_id.seqno + 1 != block_id_short.seqno {
+            tracing::error!(
+                target: tracing_targets::COLLATOR,
+                "Prev MC block ID has an incorrect sequence. Prev: {prev_mc_block_id:?}.
+             Current: {:?}",
+                block_id_short.seqno
+            );
+            return Ok(None);
+        }
+
+        let prev_mc_state = self
+            .state_node_adapter
+            .load_state(&prev_mc_block_id)
+            .await?;
+
+        let prev_mc_block_shards: FastHashMap<ShardIdent, ShardDescriptionShort> = prev_mc_state
+            .state_extra()
+            .cloned()?
+            .shards
+            .as_vec()?
+            .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(&prev_mc_block_id)
+            .await
+            .context("loading diff for mc block")?
+        else {
+            return Ok(None);
+        };
+
+        ranges.push(QueueShardBoundedRange {
+            shard_ident: prev_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 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 {
+                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),
+            });
+        }
+
+        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>> {
         if let Some(diff) =
             self.mq_adapter

collator/src/collator/messages_reader/mod.rs

@@ -217,16 +217,29 @@ 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
                         previous_cumulative_statistics.update_processed_to_by_partitions(
-                            params.all_shards_processed_to_by_partitions,
+                            params.all_shards_processed_to_by_partitions.clone(),
                         );
+
+                        // 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(),
+                        );
+
+                        // partial load statistics and enrich current value
                         previous_cumulative_statistics.load_partial(
                             mq_adapter.clone(),
                             &partitions,
                             part_stat_ranges,
                         )?;
+
                         previous_cumulative_statistics
                     }
                     _ => {
@@ -240,7 +253,6 @@ impl<V: InternalMessageValue> MessagesReader<V> {
                             cx.mc_state_gen_lt,
                             &cx.mc_top_shards_end_lts.iter().copied().collect(),
                         )?;
-                        cumulative_stats_just_loaded = true;
                         inner_cumulative_statistics
                     }
                 }

collator/src/collator/types.rs

@@ -1,4 +1,5 @@
 use std::collections::{BTreeMap, VecDeque};
+use std::fmt;
 use std::sync::{Arc, OnceLock};
 use std::time::Duration;
 
@@ -1308,6 +1309,7 @@ impl CumulativeStatistics {
         }
     }
 
+    /// Create range and full load statistics and store it
     pub fn load<V: InternalMessageValue>(
         &mut self,
         mq_adapter: Arc<dyn MessageQueueAdapter<V>>,
@@ -1342,6 +1344,7 @@ impl CumulativeStatistics {
         partitions: &FastHashSet<QueuePartitionIdx>,
         ranges: Vec<QueueShardBoundedRange>,
     ) -> Result<()> {
+        self.dirty = true;
         tracing::trace!(
             target: tracing_targets::COLLATOR,
             "cumulative_stats_partial_ranges: {:?}",
@@ -1376,14 +1379,68 @@ impl CumulativeStatistics {
         Ok(())
     }
 
+    /// Create range and remove data before this range
+    pub fn truncate_before(
+        &mut self,
+        current_shard: &ShardIdent,
+        prev_state_gen_lt: Lt,
+        mc_state_gen_lt: Lt,
+        mc_top_shards_end_lts: &FastHashMap<ShardIdent, Lt>,
+    ) {
+        let ranges = Self::compute_cumulative_stats_ranges(
+            current_shard,
+            &self.all_shards_processed_to_by_partitions,
+            prev_state_gen_lt,
+            mc_state_gen_lt,
+            mc_top_shards_end_lts,
+        );
+
+        for r in &ranges {
+            self.truncate_range(r);
+        }
+
+        self.dirty = true;
+    }
+
+    /// Remove data before range
+    fn truncate_range(&mut self, range: &QueueShardBoundedRange) {
+        let cut_key = match &range.from {
+            Bound::Included(k) | Bound::Excluded(k) => *k,
+        };
+
+        if let Some(by_partitions) = self.shards_stats_by_partitions.get_mut(&range.shard_ident) {
+            by_partitions.retain(|_part, diffs| {
+                diffs.retain(|k, _| k > &cut_key);
+                !diffs.is_empty()
+            });
+
+            if by_partitions.is_empty() {
+                self.shards_stats_by_partitions.remove(&range.shard_ident);
+            }
+        }
+    }
+
+    /// Update `all_shards_processed_to_by_partitions` and remove
+    /// processed data
     pub fn update_processed_to_by_partitions(
         &mut self,
-        all_shards_processed_to_by_partitions: FastHashMap<
-            ShardIdent,
-            (bool, ProcessedToByPartitions),
-        >,
+        new_pt: FastHashMap<ShardIdent, (bool, ProcessedToByPartitions)>,
     ) {
-        self.all_shards_processed_to_by_partitions = all_shards_processed_to_by_partitions;
+        for (&dst_shard, (_, ref processed_to)) in &new_pt {
+            let changed = match self.all_shards_processed_to_by_partitions.get(&dst_shard) {
+                Some((_, old_pt)) => old_pt != processed_to,
+                None => true,
+            };
+
+            if changed {
+                self.handle_processed_to_update(dst_shard, processed_to.clone());
+            }
+        }
+
+        // TODO should remove full stat by shard if it doesn't exist in new_pt?
+
+        self.all_shards_processed_to_by_partitions = new_pt;
+        self.dirty = true;
     }
 
     fn compute_cumulative_stats_ranges(
@@ -1444,6 +1501,9 @@ impl CumulativeStatistics {
                 }
             }
         }
+        if diff_partition_stats.is_empty() {
+            return;
+        }
 
         // finally add weeded stats
         self.add_diff_partition_stats(
@@ -1647,3 +1707,32 @@ impl ConcurrentQueueStatistics {
         }
     }
 }
+
+impl fmt::Debug for CumulativeStatistics {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let shards_summary: Vec<String> = self
+            .shards_stats_by_partitions
+            .iter()
+            .map(|(shard_id, by_partitions)| {
+                let parts: Vec<String> = by_partitions
+                    .iter()
+                    .map(|(part, diffs)| format!("p{}:{}", part, diffs.len()))
+                    .collect();
+                format!("{} -> {}", shard_id, parts.join(", "))
+            })
+            .collect();
+
+        f.debug_struct("CumulativeStatistics")
+            .field(
+                "processed_to",
+                &format!(
+                    "{} shards",
+                    self.all_shards_processed_to_by_partitions.len()
+                ),
+            )
+            .field("shards_stats_by_partitions", &shards_summary)
+            .field("result", &format!("{} partitions", self.result.len()))
+            .field("dirty", &self.dirty)
+            .finish()
+    }
+}