fix: [resharding] fix gas and balance redistribution during resharding (

#9449)

During resharding we need to redistribute the gas and balance to the children shards. 

The old implementation relied on the fact that the split shards and indexed from 0. In the future reshardings this may not be the case, for example when splitting shard 3, the split shard ids will be 3 & 4. 

The principle idea of the distribution remains the same. Both quantities are distributed evenly across the children shards. If there is any remainder := quantity % num_split_shards, then we assign 1 extra unit to each of the first `remainder` children shards. 

e.g. splitting shard 3 -> 3' & 4'
there is 5 gas to redistribute to children
shard 3' will get 3 units of gas
shard 4' will get 2 units of gas

the old implementation relied on a math trick to do that, the new one takes a more engineering leaning approach

chain/chain/src/chain.rs

@@ -5067,8 +5067,16 @@ impl<'a> ChainUpdate<'a> {
     ) -> Result<(), Error> {
         let block_hash = block.hash();
         let prev_hash = block.header().prev_hash();
+        let height = block.header().height();
         match apply_results_or_state_changes {
-            ApplySplitStateResultOrStateChanges::ApplySplitStateResults(results) => {
+            ApplySplitStateResultOrStateChanges::ApplySplitStateResults(mut results) => {
+                tracing::debug!(target: "resharding", height, ?shard_uid, "process_split_state apply");
+
+                // Sort the results so that the gas reassignment is deterministic.
+                results.sort_unstable_by_key(|r| r.shard_uid);
+                // Drop the mutability as we no longer need it.
+                let results = results;
+
                 // Split validator_proposals, gas_burnt, balance_burnt to each split shard
                 // and store the chunk extra for split shards
                 // Note that here we do not split outcomes by the new shard layout, we simply store
@@ -5085,12 +5093,12 @@ impl<'a> ChainUpdate<'a> {
 
                 let mut validator_proposals_by_shard: HashMap<_, Vec<_>> = HashMap::new();
                 for validator_proposal in chunk_extra.validator_proposals() {
-                    let shard_id = account_id_to_shard_uid(
+                    let shard_uid = account_id_to_shard_uid(
                         validator_proposal.account_id(),
                         &next_epoch_shard_layout,
                     );
                     validator_proposals_by_shard
-                        .entry(shard_id)
+                        .entry(shard_uid)
                         .or_default()
                         .push(validator_proposal);
                 }
@@ -5099,21 +5107,49 @@ impl<'a> ChainUpdate<'a> {
                     .get_split_shard_uids(shard_uid.shard_id())
                     .unwrap_or_else(|| panic!("invalid shard layout {:?}", next_epoch_shard_layout))
                     .len() as NumShards;
+
                 let total_gas_used = chunk_extra.gas_used();
                 let total_balance_burnt = chunk_extra.balance_burnt();
-                let gas_res = total_gas_used % num_split_shards;
+
+                // The gas remainder, the split shards will be reassigned one
+                // unit each until its depleted.
+                let mut gas_res = total_gas_used % num_split_shards;
+                // The gas quotient, the split shards will be reassigned the
+                // full value each.
                 let gas_split = total_gas_used / num_split_shards;
-                let balance_res = (total_balance_burnt % num_split_shards as u128) as NumShards;
+
+                // The balance remainder, the split shards will be reassigned one
+                // unit each until its depleted.
+                let mut balance_res = (total_balance_burnt % num_split_shards as u128) as NumShards;
+                // The balance quotient, the split shards will be reassigned the
+                // full value each.
                 let balance_split = total_balance_burnt / (num_split_shards as u128);
+
                 let gas_limit = chunk_extra.gas_limit();
                 let outcome_root = *chunk_extra.outcome_root();
 
                 let mut sum_gas_used = 0;
                 let mut sum_balance_burnt = 0;
+
+                // The gas and balance distribution assumes that we have a result for every split shard.
+                // TODO(resharding) make sure that is the case.
+                assert_eq!(num_split_shards, results.len() as u64);
+
                 for result in results {
-                    let shard_id = result.shard_uid.shard_id();
-                    let gas_burnt = gas_split + if shard_id < gas_res { 1 } else { 0 };
-                    let balance_burnt = balance_split + if shard_id < balance_res { 1 } else { 0 };
+                    let gas_burnt = if gas_res > 0 {
+                        gas_res -= 1;
+                        gas_split + 1
+                    } else {
+                        gas_split
+                    };
+
+                    let balance_burnt = if balance_res > 0 {
+                        balance_res -= 1;
+                        balance_split + 1
+                    } else {
+                        balance_split
+                    };
+
                     let new_chunk_extra = ChunkExtra::new(
                         &result.new_root,
                         outcome_root,
@@ -5150,6 +5186,7 @@ impl<'a> ChainUpdate<'a> {
                 assert_eq!(sum_balance_burnt, total_balance_burnt);
             }
             ApplySplitStateResultOrStateChanges::StateChangesForSplitStates(state_changes) => {
+                tracing::debug!(target: "resharding", height, ?shard_uid, "process_split_state store");
                 self.chain_store_update.add_state_changes_for_split_states(
                     *block_hash,
                     shard_uid.shard_id(),

chain/chain/src/types.rs

@@ -89,7 +89,11 @@ pub struct ApplySplitStateResult {
 // otherwise, it simply returns the state changes needed to be applied to split states
 #[derive(Debug)]
 pub enum ApplySplitStateResultOrStateChanges {
+    /// Immediately apply the split state result.
+    /// Happens during IsCaughtUp and CatchingUp
     ApplySplitStateResults(Vec<ApplySplitStateResult>),
+    /// Store the split state results so that they can be applied later.
+    /// Happens during NotCaughtUp.
     StateChangesForSplitStates(StateChangesForSplitStates),
 }