fix(bandwidth_scheduler) - use u16 for shard id in bandwidth requests (…

…near#12511) Bandwidth requests used `u8` to keep the shard id. It was done to save space, but it's not very future proof. Let's change it to `u16`. I also removed the manual conversion from `ShardId` to `u8/u16`. It now uses the `Into<u16>` implementation of `ShardId`.
saketh-are · Nov 25, 2024 · 5f36b1f · 5f36b1f
1 parent 2b35548
commit 5f36b1f
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 62 deletions.
diff --git a/core/primitives-core/src/types.rs b/core/primitives-core/src/types.rs
@@ -183,14 +183,6 @@ impl Into<usize> for ShardId {
     }
 }
 
-impl TryInto<u8> for ShardId {
-    type Error = std::num::TryFromIntError;
-
-    fn try_into(self) -> Result<u8, Self::Error> {
-        self.0.try_into()
-    }
-}
-
 impl<T> Add<T> for ShardId
 where
     T: Add<u64, Output = u64>,

diff --git a/core/primitives/src/bandwidth_scheduler.rs b/core/primitives/src/bandwidth_scheduler.rs
@@ -78,7 +78,7 @@ pub struct BandwidthRequestsV1 {
 )]
 pub struct BandwidthRequest {
     /// Requesting bandwidth to this shard.
-    pub to_shard: u8,
+    pub to_shard: u16,
     /// Bitmap which describes what values of bandwidth are requested.
     pub requested_values_bitmap: BandwidthRequestBitmap,
 }
@@ -87,7 +87,7 @@ impl BandwidthRequest {
     /// Creates a bandwidth request based on the sizes of receipts in the outgoing buffer.
     /// Returns None when a request is not needed (receipt size below base bandwidth).
     pub fn make_from_receipt_sizes<E>(
-        to_shard: u8,
+        to_shard: ShardId,
         receipt_sizes: impl Iterator<Item = Result<u64, E>>,
         params: &BandwidthSchedulerParams,
     ) -> Result<Option<BandwidthRequest>, E> {
@@ -128,7 +128,7 @@ impl BandwidthRequest {
             return Ok(None);
         }
 
-        Ok(Some(BandwidthRequest { to_shard, requested_values_bitmap: bitmap }))
+        Ok(Some(BandwidthRequest { to_shard: to_shard.into(), requested_values_bitmap: bitmap }))
     }
 
     /// Create a basic bandwidth request when receipt sizes are not available.
@@ -138,7 +138,7 @@ impl BandwidthRequest {
     /// It is used only during the protocol upgrade while the outgoing buffer metadata
     /// is not built for receipts that were buffered before the upgrade.
     pub fn make_max_receipt_size_request(
-        to_shard: u8,
+        to_shard: ShardId,
         params: &BandwidthSchedulerParams,
     ) -> BandwidthRequest {
         let mut bitmap = BandwidthRequestBitmap::new();
@@ -150,7 +150,7 @@ impl BandwidthRequest {
             .expect("max_receipt_size should be in the values list");
         bitmap.set_bit(max_receipt_size_value_pos, true);
 
-        BandwidthRequest { to_shard, requested_values_bitmap: bitmap }
+        BandwidthRequest { to_shard: to_shard.into(), requested_values_bitmap: bitmap }
     }
 }
 
@@ -376,6 +376,7 @@ mod tests {
     use rand::{Rng, SeedableRng};
 
     use crate::bandwidth_scheduler::{interpolate, BANDWIDTH_REQUEST_VALUES_NUM};
+    use crate::shard_layout::ShardUId;
 
     use super::{
         BandwidthRequest, BandwidthRequestBitmap, BandwidthRequestValues, BandwidthSchedulerParams,
@@ -499,7 +500,7 @@ mod tests {
     // Make a bandwidth request to shard 0 with a bitmap which has ones at the specified indices.
     fn make_request_with_ones(ones_indexes: &[usize]) -> BandwidthRequest {
         let mut req = BandwidthRequest {
-            to_shard: 0,
+            to_shard: ShardUId::single_shard().shard_id().into(),
             requested_values_bitmap: BandwidthRequestBitmap::new(),
         };
         for i in ones_indexes {
@@ -524,8 +525,12 @@ mod tests {
         let values = BandwidthRequestValues::new(&params).values;
 
         let get_request = |receipt_sizes: &[u64]| -> Option<BandwidthRequest> {
-            BandwidthRequest::make_from_receipt_sizes(0, make_sizes_iter(receipt_sizes), &params)
-                .unwrap()
+            BandwidthRequest::make_from_receipt_sizes(
+                ShardUId::single_shard().shard_id(),
+                make_sizes_iter(receipt_sizes),
+                &params,
+            )
+            .unwrap()
         };
 
         // No receipts - no bandwidth request.
@@ -614,7 +619,7 @@ mod tests {
                 .collect();
 
             let request = BandwidthRequest::make_from_receipt_sizes(
-                0,
+                ShardUId::single_shard().shard_id(),
                 make_sizes_iter(&receipt_sizes),
                 &params,
             )
@@ -633,7 +638,7 @@ mod tests {
         params: &BandwidthSchedulerParams,
     ) -> Option<BandwidthRequest> {
         let mut request = BandwidthRequest {
-            to_shard: 0,
+            to_shard: ShardUId::single_shard().shard_id().into(),
             requested_values_bitmap: BandwidthRequestBitmap::new(),
         };
         let values = BandwidthRequestValues::new(params).values;

diff --git a/core/store/src/trie/outgoing_metadata.rs b/core/store/src/trie/outgoing_metadata.rs
@@ -427,7 +427,7 @@ mod tests {
     use near_primitives::bandwidth_scheduler::{
         BandwidthRequest, BandwidthRequestValues, BandwidthSchedulerParams,
     };
-    use near_primitives::shard_layout::ShardLayout;
+    use near_primitives::shard_layout::{ShardLayout, ShardUId};
     use near_primitives::types::{Gas, ShardId};
     use rand::{Rng, SeedableRng};
     use rand_chacha::ChaCha20Rng;
@@ -731,14 +731,14 @@ mod tests {
 
                 // Ideal bandwidth request produced from individual receipt sizes.
                 let ideal_bandwidth_request = BandwidthRequest::make_from_receipt_sizes(
-                    0,
+                    ShardUId::single_shard().shard_id(),
                     buffered_receipts.iter().map(|s| Ok::<u64, Infallible>(s.as_u64())),
                     &scheduler_params,
                 )
                 .unwrap();
                 // Bandwidth request produced from receipt groups.
                 let groups_bandwidth_request = BandwidthRequest::make_from_receipt_sizes(
-                    0,
+                    ShardUId::single_shard().shard_id(),
                     test_queue.groups.iter().map(|g| Ok::<u64, Infallible>(g.total_size as u64)),
                     &scheduler_params,
                 )

diff --git a/pytest/lib/messages/block.py b/pytest/lib/messages/block.py
@@ -1100,7 +1100,7 @@ class ChunkProductionKey:
         BandwidthRequest, {
             'kind':
                 'struct',
-            'fields': [['to_shard', 'u8'],
+            'fields': [['to_shard', 'u16'],
                        ['requested_values_bitmap', BandwidthRequestBitmap]]
         }
     ],

diff --git a/runtime/runtime/src/congestion_control.rs b/runtime/runtime/src/congestion_control.rs
@@ -458,8 +458,6 @@ impl ReceiptSinkV2 {
             return Ok(None);
         };
 
-        let shard_u8: u8 = to_shard.try_into().expect("ShardId doesn't fit into u8");
-
         // To make a proper bandwidth request we need the metadata for the outgoing buffer to be fully initialized
         // (i.e. contain data about all of the receipts in the outgoing buffer). There is a moment right after the
         // protocol upgrade where the outgoing buffer contains receipts which were buffered in the previous protocol
@@ -473,17 +471,17 @@ impl ReceiptSinkV2 {
         // proper bandwidth requests.
         let Some(metadata) = self.outgoing_metadatas.get_metadata_for_shard(&to_shard) else {
             // Metadata not ready, make a basic request.
-            return Ok(Some(BandwidthRequest::make_max_receipt_size_request(shard_u8, params)));
+            return Ok(Some(BandwidthRequest::make_max_receipt_size_request(to_shard, params)));
         };
 
         if metadata.total_receipts_num() != outgoing_receipts_buffer_len {
             // Metadata not ready, contains less receipts than the outgoing buffer. Make a basic request.
-            return Ok(Some(BandwidthRequest::make_max_receipt_size_request(shard_u8, params)));
+            return Ok(Some(BandwidthRequest::make_max_receipt_size_request(to_shard, params)));
         }
 
         // Metadata is fully initialized, make a proper bandwidth request using it.
         let receipt_sizes_iter = metadata.iter_receipt_group_sizes(trie, side_effects);
-        BandwidthRequest::make_from_receipt_sizes(shard_u8, receipt_sizes_iter, params)
+        BandwidthRequest::make_from_receipt_sizes(to_shard, receipt_sizes_iter, params)
     }
 }
 

diff --git a/tools/protocol-schema-check/res/protocol_schema.toml b/tools/protocol-schema-check/res/protocol_schema.toml
@@ -16,16 +16,16 @@ Approval = 593918844
 ApprovalInner = 3210929495
 ApprovalMessage = 1343934820
 BalanceMismatchError = 2525009456
-BandwidthRequest = 210685111
+BandwidthRequest = 234831851
 BandwidthRequestBitmap = 2138002689
-BandwidthRequests = 2480804143
-BandwidthRequestsV1 = 956890102
+BandwidthRequests = 984876287
+BandwidthRequestsV1 = 3810915065
 BandwidthSchedulerState = 2671146930
 BitArray = 3709965115
-Block = 2469526753
-BlockBody = 3091196344
-BlockBodyV1 = 2211289526
-BlockBodyV2 = 334862132
+Block = 3541579558
+BlockBody = 206872245
+BlockBodyV1 = 2404380656
+BlockBodyV2 = 4061842967
 BlockChunkValidatorStats = 2108136564
 BlockDoubleSign = 3280983623
 BlockExtra = 1007391376
@@ -46,14 +46,14 @@ BlockInfoV1 = 735547821
 BlockInfoV2 = 1224525771
 BlockInfoV3 = 3120095857
 BlockV1 = 4242169352
-BlockV2 = 3136491312
-BlockV3 = 624321527
-BlockV4 = 4031156318
+BlockV2 = 3555701995
+BlockV3 = 1912022225
+BlockV4 = 619721361
 BlockWithChangesInfo = 887507517
 BufferedReceiptIndices = 2030010377
 CachedParts = 1180507252
-Challenge = 395035922
-ChallengeBody = 3963829050
+Challenge = 1892506304
+ChallengeBody = 652213846
 ChunkContractAccesses = 266426785
 ChunkContractAccessesInner = 2811580521
 ChunkContractAccessesV1 = 3680796018
@@ -67,10 +67,10 @@ ChunkExtraV1 = 774877102
 ChunkHash = 1471814478
 ChunkHashHeight = 825215623
 ChunkProductionKey = 2508733236
-ChunkProofs = 455145743
-ChunkState = 1097416384
+ChunkProofs = 368992087
+ChunkState = 1462720897
 ChunkStateTransition = 307448170
-ChunkStateWitness = 432744472
+ChunkStateWitness = 2738191235
 ChunkStateWitnessAck = 177881908
 ChunkStats = 4176245277
 CodeBytes = 2940589161
@@ -106,10 +106,10 @@ EdgeInner = 3441978250
 EdgeRepr = 733491460
 EdgeState = 2217555236
 EncodedChunkStateWitness = 329848903
-EncodedShardChunk = 4054694684
+EncodedShardChunk = 3513967731
 EncodedShardChunkBody = 2481614037
 EncodedShardChunkV1 = 23151435
-EncodedShardChunkV2 = 885978559
+EncodedShardChunkV2 = 2872968441
 EpochId = 1173955846
 EpochInfo = 820386104
 EpochInfoAggregator = 2599467180
@@ -155,7 +155,7 @@ LatestKnown = 2945167085
 LatestWitnessesInfo = 2488443612
 LegacyAccount = 1291371319
 MainTransitionKey = 3721480128
-MaybeEncodedShardChunk = 3950487412
+MaybeEncodedShardChunk = 2688433530
 MerklePathItem = 2615629611
 MessageDiscriminant = 3240833245
 MethodResolveError = 1206790835
@@ -164,13 +164,13 @@ NextEpochValidatorInfo = 3660299258
 NonDelegateAction = 3255205790
 ParentSplitParameters = 1570407998
 PartialEdgeInfo = 1350359189
-PartialEncodedChunk = 2722484497
+PartialEncodedChunk = 3453254449
 PartialEncodedChunkForwardMsg = 68012243
 PartialEncodedChunkPart = 194051090
 PartialEncodedChunkRequestMsg = 1470767646
 PartialEncodedChunkResponseMsg = 2957212759
 PartialEncodedChunkV1 = 3642706173
-PartialEncodedChunkV2 = 2811493841
+PartialEncodedChunkV2 = 889431033
 PartialEncodedContractDeploys = 3216562245
 PartialEncodedContractDeploysInner = 2549441552
 PartialEncodedContractDeploysPart = 1672852427
@@ -182,7 +182,7 @@ PeerChainInfoV2 = 1260985250
 PeerId = 2447445523
 PeerIdOrHash = 4080492546
 PeerInfo = 3831734408
-PeerMessage = 72369337
+PeerMessage = 2449118209
 Ping = 2783493472
 Pong = 3159638327
 PrepareError = 4009037507
@@ -211,35 +211,35 @@ ReceiptV1 = 2994842769
 ReceiptValidationError = 551721215
 ReceivedData = 3601438283
 RootProof = 3135729669
-RoutedMessage = 1383052756
-RoutedMessageBody = 4271987761
+RoutedMessage = 4094023612
+RoutedMessageBody = 2434166470
 RoutingTableUpdate = 2987752645
 Secp256K1PublicKey = 4117078281
 Secp256K1Signature = 3687154735
 ServerError = 2338793369
-ShardChunk = 3411174761
-ShardChunkHeader = 837506379
-ShardChunkHeaderInner = 3017297308
+ShardChunk = 350912142
+ShardChunkHeader = 2471921769
+ShardChunkHeaderInner = 4085026561
 ShardChunkHeaderInnerV1 = 1271245459
 ShardChunkHeaderInnerV2 = 2664186997
 ShardChunkHeaderInnerV3 = 2843221286
-ShardChunkHeaderInnerV4 = 1349627561
+ShardChunkHeaderInnerV4 = 3066669719
 ShardChunkHeaderV1 = 47891389
 ShardChunkHeaderV2 = 226996174
-ShardChunkHeaderV3 = 1154712329
+ShardChunkHeaderV3 = 3315420662
 ShardChunkV1 = 1956351688
-ShardChunkV2 = 387353750
+ShardChunkV2 = 61923749
 ShardLayout = 1639977238
 ShardLayoutV0 = 3139625127
 ShardLayoutV1 = 2054829142
 ShardLayoutV2 = 997571636
 ShardProof = 1787648268
-ShardStateSyncResponse = 587140030
+ShardStateSyncResponse = 85505566
 ShardStateSyncResponseHeaderV1 = 1491041593
-ShardStateSyncResponseHeaderV2 = 2480490611
+ShardStateSyncResponseHeaderV2 = 3976991370
 ShardStateSyncResponseV1 = 1376844594
-ShardStateSyncResponseV2 = 153327553
-ShardStateSyncResponseV3 = 4173646451
+ShardStateSyncResponseV2 = 2436600135
+ShardStateSyncResponseV3 = 3559507986
 ShardUId = 2410086023
 Signature = 3997391707
 SignedDelegateAction = 2482265228
@@ -252,9 +252,9 @@ StateChangeCause = 3890585134
 StateHeaderKey = 1666317019
 StatePartKey = 1083277414
 StatePartRequest = 1911936050
-StateResponseInfo = 3989607602
+StateResponseInfo = 3497809959
 StateResponseInfoV1 = 226548439
-StateResponseInfoV2 = 1445527353
+StateResponseInfoV2 = 3314983982
 StateRootNode = 1865105129
 StateStoredReceipt = 311659268
 StateStoredReceiptMetadata = 2895538362