feat: wiring for bandwidth scheduler #12234

jancionear · 2024-10-16T13:49:30Z

Add the wiring needed for the bandwidth scheduler algorithm.

Changes:

Add a new ProtocolFeature - BandwidthScheduler, its protocol version is set to nightly
Add a struct that will keep the bandwidth requests generated by the shards
Propagate the bandwidth requests through the blockchain - put the generated bandwidth requests in the shard headers, pass the previous bandwidth requests to the runtime
Add a struct that represents the bandwidth scheduler state, it's stored in the trie and modified on every scheduler invocation.
Mock implementation of bandwidth scheduler - it takes the previous bandwidth requests and the state and mocks the scheduler algorithm. It activates the requests propagation logic and breaks some tests.

Propagation of bandwidth requests

The flow of bandwidth requests looks as follows:

A chunk is applied and generates bandwidth requests. They are put in ApplyResult and ApplyChunkResult
The requests are taken from the apply result and put in ChunkExtra. ChunkExtra is persisted in the database
During chunk production, Client fetches ChunkExtra of the previous chunk and puts the bandwidth requests in chunk header
The produced chunks are included in the block
The new chunks are applied, their ApplyState contains bandwidth requests taken from all the chunk headers in the block that contains the applied chunks.
During the application, bandwidth scheduler looks at the requests created at the previous height and grants banwidth
Receipts are sent out
Then the chunk generates new bandwidth requests
etc

The flow is very similar to the one for congestion info.

Scheduler state

Bandwidth scheduler needs to keep some persistent state. In the future it'll be something like "how much every shard was granted lately", it'll be used to maintain fairness. For now it's just mock data.
Scheduler state should always be the same on all shards. All shards start with the same scheduler state, apply the scheduler at the same heights with the same inputs and always end up with the same scheduler state.
This means that the bandwidth scheduler also needs to be run for missing chunks. Luckily that can be easily achieved thanks to existing apply_old_chunk infrastructure (all missing chunk are applied, it counts as "implicit state transitions").
The state_root will now change at every height, even when there are no receipts to be processed. It breaks some tests which assumed that the state root wouldn't change.

The pull request is meant to be reviewed commit-by-commit, I tried to make the commit history nice.

Add the structs which will be used to represent bandwidth requests generated by a shard. For now the BandwidthRequest doesn't have the requested values, they will be added later.

Chunk headers will contain bandwidth requests generated during the previous chunk application. We will collect bandwidth requests from all the shards and use them in bandwidth scheduler during chunk application.

Bandwidth requests will be generated during chunk application and then they'll be available in the ApplyResult.

Result of chunk application should keep the generated bandwidth requests.

ChunkExtra stores the results of chunk application in a persistent way. Let's put the generated bandwidth requests there and then fetch them when producing the next chunk.

Collect the bandwith requests generated by all shards at the previous height and expose them to the runtime. Runtime needs to have the requests, as they're needed to run the bandwidth scheduler.

Add a struct that keeps the persistent state used by the bandwidth scheduler.

Add a mock implementation of the bandwidth scheduler algorithm. Bandwidth scheduler takes the current state and previous bandwidth requests and generates bandwidth grants. The mock implementation takes the inputs and generates deterministic state changes based on them, but it doesn't generate the bandwidth grants yet. The mock implementation is enough to activate the logic that propagates bandwidth requests throughout the blockchain and break some tests.

This test assumed that the state root doesn't change when there are no receipts, but this is no longer true. Bandwidth scheduler modifies the state on every height, so now the state root changes every time.

state_viewer::apply_chunk has the ability to apply a chunk when the block that contains the chunk isn't available. Initially I passed empty bandwidth requests in the ApplyState, as usually they're taken from the chunk headers in the block that contains the applied chunk, and this block isn't available here. But that breaks test_apply_chunk - a test which applies a chunk in a normal way and compares the result with a chunk that was applied without providing the block. It expects the state roots to be the same, but that's not the case because the bandwidth requests are different and bandwidth scheduler produces different state. To deal with this we can try to fetch the original bandwidth requests from the chunk extra of the previous chunks. It's an opportunistic reconstruction - if the chunk extra is available it adds the requests to the apply_state, if not it leaves them empty. This is enough to fix the state root mismatch.

This tests creates a situation where the last few chunks in an epoch are missing. It performs state sync, then takes the state root of the first missing chunk on one node and expects the state roots of all the missing chunks on the other node to be the same as that first state root. This breaks because bandwidth scheduler changes the state at every height - even for missing chunks - so the state root for later missing chunks is not the same as the state root of the first missing chunk. Fix the problem by comparing state roots of missing chunks at the same heights.

This test performs state sync and then does a function call on the synced node to test that the sync worked. The function call at the end of the test started failing with `MissingTrieValue`. I suspect that the function call is done with the wrong state root - it worked previously, when all the state roots were the same, as the chunks don't have any transactions, but broke when bandwidth scheduler started changing the state at every height. The `MissingTrieValue` error stops occuring when the state root is taken from the previous block. My understanding of state sync isn't very good, but I think this theory makes sense.

Add an extra check to ensure that the scheduler state stays the same on all shards.

codecov · 2024-10-16T14:19:24Z

Codecov Report

Attention: Patch coverage is 84.59870% with 71 lines in your changes missing coverage. Please review.

Project coverage is 71.61%. Comparing base (e0f4592) to head (42eae3b).

Files with missing lines	Patch %	Lines
core/primitives/src/views.rs	73.50%	31 Missing ⚠️
chain/chain/src/validate.rs	54.83%	13 Missing and 1 partial ⚠️
runtime/runtime/src/bandwidth_scheduler/mod.rs	90.00%	5 Missing and 2 partials ⚠️
core/primitives/src/bandwidth_scheduler.rs	84.00%	4 Missing ⚠️
tools/state-viewer/src/util.rs	33.33%	4 Missing ⚠️
chain/rosetta-rpc/src/adapters/transactions.rs	0.00%	3 Missing ⚠️
core/primitives/src/types.rs	93.33%	2 Missing ⚠️
tools/state-viewer/src/apply_chunk.rs	84.61%	1 Missing and 1 partial ⚠️
chain/chain-primitives/src/error.rs	0.00%	1 Missing ⚠️
..._validation/chunk_validator/orphan_witness_pool.rs	50.00%	1 Missing ⚠️
... and 2 more

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #12234      +/-   ##
==========================================
+ Coverage   71.57%   71.61%   +0.04%     
==========================================
  Files         836      838       +2     
  Lines      168073   168463     +390     
  Branches   168073   168463     +390     
==========================================
+ Hits       120294   120644     +350     
- Misses      42524    42561      +37     
- Partials     5255     5258       +3

Flag	Coverage Δ
backward-compatibility	`0.16% <0.00%> (-0.01%)`	⬇️
db-migration	`0.16% <0.00%> (-0.01%)`	⬇️
genesis-check	`1.24% <0.00%> (-0.01%)`	⬇️
integration-tests	`38.94% <76.57%> (+0.13%)`	⬆️
linux	`71.14% <47.28%> (-0.08%)`	⬇️
linux-nightly	`71.18% <78.52%> (+0.01%)`	⬆️
macos	`54.19% <44.37%> (-0.07%)`	⬇️
pytests	`1.55% <0.00%> (-0.01%)`	⬇️
sanity-checks	`1.36% <0.00%> (-0.01%)`	⬇️
unittests	`65.38% <82.60%> (+0.04%)`	⬆️
upgradability	`0.21% <0.00%> (-0.01%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

… is enabled All chunks produced in the protocol version where bandwidth scheduler is enabled should use ShardChunkHeaderInner::V4, I missed this in the previous commit.

The test iterates over all items in the trie and creates a StateRecord for each of them. The problem is that some types of trie entries don't have a corresponding StateRecord variant. For example outgoing buffers, yield resume data, and bandwidth scheduler state can't be made into a StateRecord. The test started failing because it tries to unwrap result of `StateRecord::from_raw_key_value` for a trie entry that represents BandwidthSchedulerState. The function returns None and the unwrap panics. Fix the problem by removing the unwrap and instead looking for `Some` value. The test only looks for one type of StateRecord, it doesn't matter if it skips over the scheduler state.

…ate missing chunks

jancionear · 2024-10-17T22:19:48Z

For the testloop test I need another fix to be merged: #12253
For now I cherry-picked commits from the fix (a0ea12d, ca1ebd6) to get the test running. I will revert them when the fix gets merged.

This reverts commit e6bd44e.

This reverts commit ca1ebd6.

…to simulate missing chunks" This reverts commit a0ea12d.

jancionear · 2024-10-18T12:28:41Z

I will revert them when the fix gets merged.

The fix got merged. Reverted commits with the old version of the test, merged master and added a test using the functionality on master.

wacban

I went through half of it, so far looks good to me. The only piece that is a bit concerning is the validate - please see my comment

@tayfunelmas Do you remember what was that trouble we had with congestion control initialization in genesis? Would we need to handle it similarly here?

wacban · 2024-10-23T11:38:43Z

chain/chain/src/types.rs

+    /// Requests for bandwidth to send receipts to other shards.
+    /// Will be None for protocol versions that don't have the BandwidthScheduler feature enabled.
+    pub bandwidth_requests: Option<BandwidthRequests>,
+    /// Used only for a sanity check.
+    pub bandwidth_scheduler_state_hash: CryptoHash,


mini nit: Maybe put it next to congestion control? They are very similar.

wacban · 2024-10-23T11:41:23Z

chain/chain/src/validate.rs

+    if extra_bandwidth_requests.is_none()
+        && header_bandwidth_requests == Some(&BandwidthRequests::empty())
+    {
+        // This corner case happens for the first chunk that has the BandwidthScheduler feature enabled.
+        // The previous chunk was applied with a protocol version which doesn't have bandwidth scheduler
+        // enabled and because of that the bandwidth requests in ChunkExtra are None.
+        // The header was produced in the new protocol version, and the newer version of header always
+        // has some bandwidth requests, it's not an `Option`. Because of that the header requests are `Some(BandwidthRequests::empty())`.
+        return Ok(());
+    }


Can you do it the same way as it worked for congestion info? Basically they always have to match. I don't remember the exact reasons but it was way cleaner way to handle protocol upgrade this way.

Hmm I'm not sure if it's better, the way congestion control does it causes the protocol to accept two different chunk header versions for one protocol version :/
I don't know, I can try it in a follow up PR. I'll merge this one as is, there's enough mess here already. Let's discuss it in a different PR.

wacban · 2024-10-23T11:46:16Z

chain/client/src/stateless_validation/chunk_validator/orphan_witness_pool.rs

+                ShardChunkHeaderInner::V1(_) => unimplemented!(),
                ShardChunkHeaderInner::V2(inner) => inner.encoded_length = encoded_length,
                ShardChunkHeaderInner::V3(inner) => inner.encoded_length = encoded_length,
-                _ => unimplemented!(),
+                ShardChunkHeaderInner::V4(inner) => inner.encoded_length = encoded_length,


nice, always better to be explicit and exhaust all variants, ❤️

wacban · 2024-10-23T11:46:57Z

chain/client/src/test_utils/client.rs

@@ -227,6 +227,7 @@ pub fn create_chunk(
            decoded_chunk.prev_outgoing_receipts(),
            header.prev_outgoing_receipts_root(),
            header.congestion_info(),
+            header.bandwidth_requests().cloned(),


Just out of curiosity why do you need cloned here for requests and not for congestion info? Is congestion info Copy or something?

Yup CongestionInfo is Copy. Bandwidth requests are not because they contain a vec.

wacban · 2024-10-23T11:48:08Z

core/primitives-core/src/version.rs

            // TODO(#11099): Move this feature to Nightly.
-            ProtocolFeature::ExcludeContractCodeFromStateWitness => 146,
+            ProtocolFeature::ExcludeContractCodeFromStateWitness => 147,


lol, there is nightly and then there is after nightly xD

wacban · 2024-10-23T12:12:35Z

integration-tests/src/test_loop/tests/bandwidth_scheduler_protocol_upgrade.rs

+        HashMap::from_iter(
+            [
+                (ShardId::new(0), 0..0),
+                (ShardId::new(1), -2..0),
+                (ShardId::new(2), 0..2),
+                (ShardId::new(3), -2..2),
+            ]
+            .into_iter(),
+        ),


that's sweet

wacban · 2024-10-23T12:13:22Z

integration-tests/src/test_loop/tests/resharding_v3.rs

+            let state_record = StateRecord::from_raw_key_value(key, value);
+            if let Some(StateRecord::Account { account_id, .. }) = state_record {


This is an actual change in semantics, is that intentional?

Intentional, explained in the commit description

wacban · 2024-10-23T12:14:43Z

integration-tests/src/tests/client/process_blocks.rs

+        let chunk_extra =
+            env.clients[0].chain.get_chunk_extra(block.header().prev_hash(), &shard_uid).unwrap();


That's strange, what happened here?

Check out commit description

wacban · 2024-10-23T12:15:51Z

integration-tests/src/tests/client/sync_state_nodes.rs

+                    let expected_height = if ProtocolFeature::BandwidthScheduler
+                        .enabled(genesis.config.protocol_version)
+                    {
+                        height
+                    } else {
+                        last_chunk_height
+                    };


Do I want to know?

I wrote about it in the commit description

wacban · 2024-10-23T12:17:07Z

runtime/runtime/src/bandwidth_scheduler/mod.rs

TODO(wacban) - continue reviewing from here

Ah crap, maybe I should've mentioned that the PR is meant to be reviewed commit-by-commit. I made an effort to make the commit history nice, and added some additional comments about the changes in commit descriptions to make reviewing easier.

jancionear · 2024-10-23T13:06:03Z

@tayfunelmas Do you remember what was that trouble we had with congestion control initialization in genesis? Would we need to handle it similarly here?

The good thing about bandwidth requests is that they can always be empty, it doesn't matter what the state is. AFAIK congestion info needs to be initialized based on the shard state, but for requests this is not a problem. Maybe there's some issue I'm missing, but IMO it should be fine.

Clippy started complaining that these structs are too large (>1024 bytes) and recommended to put them in Box. I did as clippy suggested. Putting structs in boxes shouldn't change the serialization format.

cargo fmt went wild again 0_o

wacban

LGTM

wacban · 2024-10-23T14:21:24Z

runtime/runtime/src/lib.rs

+            bandwidth_requests,
            contract_accesses,
+            bandwidth_scheduler_state_hash: bandwidth_scheduler_output
+                .as_ref()
+                .map(|o| o.scheduler_state_hash)
+                .unwrap_or_default(),


nit: fix the ordering

I think it's fixed now, the master merge fiddled with it.

wacban · 2024-10-23T15:33:01Z

core/primitives/src/trie_key.rs

@@ -56,6 +56,7 @@ pub mod col {
    /// backpressure on the receiving shard.
    /// (`primitives::receipt::Receipt`).
    pub const BUFFERED_RECEIPT: u8 = 14;
+    pub const BANDWIDTH_SCHEDULER_STATE: u8 = 15;


Ah sorry, I mixed that up with nibbles. You're right, we should be just fine. That's awesome :)

jancionear · 2024-10-23T17:43:45Z

core/primitives/src/sharding.rs

+            ShardChunkHeader::V3(header) => header.inner.bandwidth_requests(),
+        }
+    }
+
    /// Returns whether the header is valid for given `ProtocolVersion`.
    pub fn valid_for(&self, version: ProtocolVersion) -> bool {


Reminder to myself - make sure that we check this in witness validation. Didn't matter before, but might matter now.

jancionear added 17 commits October 16, 2024 11:20

Add ProtocolFeature::BandwidthScheduler

e782412

Add BandwidthRequest and BandwidthRequests structs

3e22208

Add the structs which will be used to represent bandwidth requests generated by a shard. For now the BandwidthRequest doesn't have the requested values, they will be added later.

Add bandwidth requests to chunk header

5614b06

Chunk headers will contain bandwidth requests generated during the previous chunk application. We will collect bandwidth requests from all the shards and use them in bandwidth scheduler during chunk application.

Add bandwidth requests to ApplyResult

a217022

Bandwidth requests will be generated during chunk application and then they'll be available in the ApplyResult.

Add banwidth requests to ApplyChunkResult

dc156ac

Result of chunk application should keep the generated bandwidth requests.

Add bandwidth requests to ChunkExtra

24dd104

ChunkExtra stores the results of chunk application in a persistent way. Let's put the generated bandwidth requests there and then fetch them when producing the next chunk.

Add bandwidth requests to ApplyState

c9edcb9

Collect the bandwith requests generated by all shards at the previous height and expose them to the runtime. Runtime needs to have the requests, as they're needed to run the bandwidth scheduler.

Add bandwidth scheduler state

66103e5

Add a struct that keeps the persistent state used by the bandwidth scheduler.

Propagate previous bandwidth requests for missing chunks

40db96e

Fix test_empty_apply

db96527

This test assumed that the state root doesn't change when there are no receipts, but this is no longer true. Bandwidth scheduler modifies the state on every height, so now the state root changes every time.

Fix test_flat_storage_iter

dc1cd5f

Fix test_archival_save_trie_changes

d7bfc1e

Add a sanity check for scheduler state

3e63448

Add an extra check to ensure that the scheduler state stays the same on all shards.

jancionear requested a review from wacban October 16, 2024 13:49

jancionear requested a review from a team as a code owner October 16, 2024 13:49

Fix python formatting

caed220

jancionear added 9 commits October 16, 2024 15:43

Merge branch 'master' into bandsim-wires

1a27050

Don't allow ShardChunkHeaderInner::V3 when BandwidthScheduler feature…

54ff731

… is enabled All chunks produced in the protocol version where bandwidth scheduler is enabled should use ShardChunkHeaderInner::V4, I missed this in the previous commit.

Validate that bandwidth requests in chunk extra match the chunk header

dfb11b6

Merge branch 'master' into bandsim-wires

272d486

Merge branch 'master' into bandsim-wires

ccec51a

fix(test_loop) - drop endorsements instead of partial chunks to simul…

a0ea12d

…ate missing chunks

Add a testloop protocol upgrade test

ca1ebd6

Test protocol upgrade to version with bandwidth scheduler

e6bd44e

jancionear added 5 commits October 18, 2024 12:17

Revert "Test protocol upgrade to version with bandwidth scheduler"

bb3f370

This reverts commit e6bd44e.

Revert "Add a testloop protocol upgrade test"

2af440c

This reverts commit ca1ebd6.

Revert "fix(test_loop) - drop endorsements instead of partial chunks …

aeea61f

…to simulate missing chunks" This reverts commit a0ea12d.

Merge branch 'master' into bandsim-wires

b4fca01

Test protocol upgrade to version with bandwidth scheduler

a539c3a

jancionear mentioned this pull request Oct 21, 2024

[Project tracking] Cross shard bandwidth limits #11920

Open

Merge branch 'master' into bandsim-wires

8ad224b

wacban reviewed Oct 23, 2024

View reviewed changes

jancionear added 3 commits October 23, 2024 14:29

Box structs that are too large

c4fde1a

Clippy started complaining that these structs are too large (>1024 bytes) and recommended to put them in Box. I did as clippy suggested. Putting structs in boxes shouldn't change the serialization format.

Merge branch 'master' into bandsim-wires

d8068ee

Add a commit about BandwidthSchedulerStateUpdate

42eae3b

cargo fmt went wild again 0_o

wacban approved these changes Oct 23, 2024

View reviewed changes

jancionear added this pull request to the merge queue Oct 23, 2024

jancionear commented Oct 23, 2024

View reviewed changes

github-merge-queue bot removed this pull request from the merge queue due to failed status checks Oct 23, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: wiring for bandwidth scheduler #12234

feat: wiring for bandwidth scheduler #12234

jancionear commented Oct 16, 2024 •

edited

Loading

codecov bot commented Oct 16, 2024 •

edited

Loading

jancionear commented Oct 17, 2024

jancionear commented Oct 18, 2024

wacban left a comment

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

wacban Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

jancionear commented Oct 23, 2024

wacban left a comment

wacban Oct 23, 2024

jancionear Oct 23, 2024

wacban Oct 23, 2024

jancionear Oct 23, 2024

		let state_record = StateRecord::from_raw_key_value(key, value);
		if let Some(StateRecord::Account { account_id, .. }) = state_record {

		let chunk_extra =
		env.clients[0].chain.get_chunk_extra(block.header().prev_hash(), &shard_uid).unwrap();

feat: wiring for bandwidth scheduler #12234

Are you sure you want to change the base?

feat: wiring for bandwidth scheduler #12234

Conversation

jancionear commented Oct 16, 2024 • edited Loading

Propagation of bandwidth requests

Scheduler state

codecov bot commented Oct 16, 2024 • edited Loading

Codecov Report

jancionear commented Oct 17, 2024

jancionear commented Oct 18, 2024

wacban left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jancionear commented Oct 23, 2024

wacban left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jancionear commented Oct 16, 2024 •

edited

Loading

codecov bot commented Oct 16, 2024 •

edited

Loading