fix(congestion_control) - proper protocol upgrade handling for conges…

…tion seed (#12594)

Fixing unhandled protocol upgrade introduced in #12555.

runtime/runtime/src/lib.rs

@@ -2083,18 +2083,27 @@ impl Runtime {
         let mut own_congestion_info = receipt_sink.own_congestion_info();
         if let Some(congestion_info) = &mut own_congestion_info {
             delayed_receipts.apply_congestion_changes(congestion_info)?;
-            let shard_layout = epoch_info_provider.shard_layout(&apply_state.epoch_id)?;
-            let shard_ids = shard_layout.shard_ids().collect_vec();
-            let shard_index = shard_layout
-                .get_shard_index(apply_state.shard_id)
-                .map_err(Into::<EpochError>::into)?
-                .try_into()
-                .expect("Shard Index must fit within u64");
-
-            let congestion_seed = apply_state.block_height.wrapping_add(shard_index);
+            let protocol_version = apply_state.current_protocol_version;
+
+            let (all_shards, shard_seed) =
+                if ProtocolFeature::SimpleNightshadeV4.enabled(protocol_version) {
+                    let shard_layout = epoch_info_provider.shard_layout(&apply_state.epoch_id)?;
+                    let shard_ids = shard_layout.shard_ids().collect_vec();
+                    let shard_index = shard_layout
+                        .get_shard_index(apply_state.shard_id)
+                        .map_err(Into::<EpochError>::into)?
+                        .try_into()
+                        .expect("Shard Index must fit within u64");
+
+                    (shard_ids, shard_index)
+                } else {
+                    (apply_state.congestion_info.all_shards(), apply_state.shard_id.into())
+                };
+
+            let congestion_seed = apply_state.block_height.wrapping_add(shard_seed);
             congestion_info.finalize_allowed_shard(
                 apply_state.shard_id,
-                shard_ids.as_slice(),
+                &all_shards,
                 congestion_seed,
             );
         }