[resharding] Rewrite state column gc #13598

shreyan-gupta · 2025-05-22T08:04:37Z

Update the state garbage collection code to take into account the new solution for refcount issue in resharding.

I don't completely understand what the old solution was doing, but the new solution divides gc into two parts, one specific to resharding, where we are simply deleting the parent shard_uid state, and second, deleting untracked shard_uids.

The logic to find untracked shard_uids is updated and simplified.

Don't worry if the old code and new code don't make 100% sense. I took several days to figure out what exactly is going on :)

shreyan-gupta · 2025-05-22T08:05:29Z

test-loop-tests/src/tests/resharding_v3.rs

-        1 => {
-            let shards_split_map = vec![vec![ShardId::new(0), ShardId::new(1), ShardId::new(2)]];
-            #[allow(deprecated)]
-            ShardLayout::v1(boundary_accounts, Some(shards_split_map), 3)


ShardLayout::v1 is neither used nor supported. We can remove these tests.

shreyan-gupta · 2025-05-22T08:08:19Z

test-loop-tests/src/utils/resharding.rs

        let (key, _) = kv.unwrap();
        let shard_uid = ShardUId::try_from_slice(&key[0..8]).unwrap();
        shard_uid_prefixes.insert(shard_uid);
    }
-    // TODO(resharding): Handle GC after TrieStateResharder


With the new implementation of refcount fix, we no longer have shard_uid mapping. We can however re-enable the state column GC check that we had commented out earlier. This is the main test addition of this PR.

shreyan-gupta · 2025-05-22T08:31:25Z

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

-/// Run one more epoch.
-/// "Restart from the snapshot" is to ensure that we can continue producing blocks without relying on caches.
-#[test]
-fn test_long_chain_with_restart_from_snapshot() {


Don't know why this test is failing but it has been superseded by testloop restarts tests.

codecov · 2025-05-22T08:39:34Z

Codecov Report

Attention: Patch coverage is 80.00000% with 16 lines in your changes missing coverage. Please review.

Project coverage is 69.48%. Comparing base (6c7c2d4) to head (d6fc327).
Report is 4 commits behind head on master.

Files with missing lines	Patch %	Lines
chain/chain/src/garbage_collection.rs	77.27%	0 Missing and 15 partials ⚠️
core/primitives/src/shard_layout.rs	88.88%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #13598   +/-   ##
=======================================
  Coverage   69.48%   69.48%           
=======================================
  Files         856      856           
  Lines      168182   168114   -68     
  Branches   168182   168114   -68     
=======================================
- Hits       116857   116816   -41     
+ Misses      46542    46520   -22     
+ Partials     4783     4778    -5

Flag	Coverage Δ
pytests	`1.52% <0.00%> (+<0.01%)`	⬆️
pytests-nightly	`1.62% <0.00%> (+<0.01%)`	⬆️
unittests	`69.10% <80.00%> (-0.01%)`	⬇️
unittests-nightly	`69.01% <80.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.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

wacban

LGTM

wacban · 2025-05-22T10:11:16Z

chain/chain/src/garbage_collection.rs

@@ -1134,133 +1128,104 @@ impl<'a> ChainStoreUpdate<'a> {
    }
 }

-/// Returns shards that we tracked in an epoch, given a hash of the last block in the epoch.
-/// The block has to be available, so this function has to be called before gc is run for the block.
+/// Cleans up the state of the parent shard once we stop tracking the epoch where the resharding happened.


Can you clarify a bit the "tracking the epoch" part?

Changed the language a bit, but high level I'm trying to convey, we need to cleanup parent if this was the last epoch in which the parent existed, i.e. was resharded.

wacban · 2025-05-22T10:13:31Z

chain/chain/src/garbage_collection.rs

+    let mut trie_store_update = store.trie_store().store_update();
+    for parent_shard_uid in shard_layout.get_split_parent_shard_uids()? {
+        // Delete the state of the parent shard
+        tracing::info!(target: "garbage_collection", ?parent_shard_uid, "resharding state cleanup");


Does it deserve to be info?

This ideally happens just once for the parent_shard_uid, and is quite instructive, so might as well be an info instead of debug. If you'd like me to change to debug, I can do that.

Changed to debug

wacban · 2025-05-22T10:14:13Z

chain/chain/src/garbage_collection.rs

+    for parent_shard_uid in shard_layout.get_split_parent_shard_uids()? {
+        // Delete the state of the parent shard
+        tracing::info!(target: "garbage_collection", ?parent_shard_uid, "resharding state cleanup");
+        trie_store_update.delete_shard_uid_prefixed_state(parent_shard_uid);


sanity check - this should not be called on the archival nodes

Yes, from my understanding, GC code only runs on normal nodes.

Doesn't GC run differently on archival nodes?
It's like a normal node GC with the addition that there's a copy loop to cold DB. Plus GC is stalling to proceed if cold DB lag is high.

I might be confusing the high level vision with the impl details though..

Hmm, I might lack the understanding of cold storage loop. I think it's better to take this discussion offline.

wacban · 2025-05-22T10:15:14Z

chain/chain/src/garbage_collection.rs

-        {
-            tracked_shards.push(shard_uid);
-        }
+        assert_eq!(get_shard_uid_mapping(&store, shard_uid), shard_uid, "Incomplete Resharding");


maybe do that before actually deleting the data :)

it wouldn't be deleted yet if I understand? seems it is just added to the batch so far.

This is mostly just a sanity check that resharding should have definitely been completed by now. Any node that continues to have this mapping would likely fail.

wacban · 2025-05-22T10:16:15Z

chain/chain/src/garbage_collection.rs

-/// where each `ShardUId` is potentially mapped to its ancestor to get the database key prefix.
-/// We only remove a shard State if all its descendants are ready to be cleaned up,
-/// in which case, we also remove the mapping from `StateShardUIdMapping`.
+/// block_hash is the last block of the epoch we are cleaning up.


technically it can be any block, it's just this method with be a no-op for any block other than the last of the epoch

wacban · 2025-05-22T10:21:20Z

chain/chain/src/garbage_collection.rs

+    let latest_block_hash = chain_store_update.head()?.last_block_hash;
+    let mut block_info = epoch_manager.get_block_info(&latest_block_hash)?;


The naming is very confusing since the block_info is not the block of the block_hash. Can you rename those to something like target_block_x and current_block_x?

wacban · 2025-05-22T10:26:13Z

chain/chain/src/garbage_collection.rs

+    let mut shards_to_cleanup = epoch_manager
+        .get_shard_layout(epoch_manager.get_block_info(block_hash)?.epoch_id())?
+        .shard_uids()


nit1: There is a lot crammed into a single line, can you split it down?
nit2: Can you use retain for consistency with the other logic below?

let target_block_info = let target_epoch_id = let target_shard_layout = .. let mut shards_to_cleanup = shard_layout.shard_uids(); shards_to_cleanup.retain(...)

wacban · 2025-05-22T10:29:20Z

chain/chain/src/garbage_collection.rs

+        .shard_uids()
+        .filter(|shard_uid| {
+            // Remove shards that we are tracking right now
+            !shard_tracker.cares_about_shard_this_or_next_epoch(


I'm guessing you're using different logic for the current epoch because it's not guaranteed that TrieChanges will already be present?

Yes precisely. This was leading to scenarios where we continued to track state of child shards somehow, but I'm not sure how exactly was that happening.

wacban · 2025-05-22T10:30:16Z

chain/chain/src/garbage_collection.rs

+
+    // reverse iterate over the epochs
+    let store = chain_store_update.store();
+    while !shards_to_cleanup.is_empty() && block_info.hash() != block_hash {


I'd be tempted to remove the empty condition, it's seems like a tiny and irrelevant optimization. Correct me if I'm wrong ;)

Yeah, mostly as an optimization... Removed it

wacban · 2025-05-22T10:31:49Z

chain/chain/src/garbage_collection.rs

    }

    // Delete State of `shards_to_cleanup` and associated ShardUId mapping.
-    tracing::info!(target: "garbage_collection", ?shards_to_cleanup, ?shard_uid_mappings_to_remove, "state_cleanup");
+    tracing::info!(target: "garbage_collection", ?shards_to_cleanup, "state_cleanup");


again I'm not sure if this deserves info

Changed to debug

staffik · 2025-05-22T10:11:18Z

chain/chain/src/garbage_collection.rs

-        return Ok(None);
+) -> Result<(), Error> {
+    // If we are GC'ing the resharding block, i.e. the last block of the epoch, clear out state for the parent shard
+    if !epoch_manager.will_shard_layout_change(block_hash)? {


If I understand correctly, that returns true for each block in the epoch before resharding. Was not the intention to run gc_resharding for the last block only?

Good catch, updated to include condition is_last_block_in_finished_epoch

staffik · 2025-05-22T10:27:25Z

chain/chain/src/garbage_collection.rs

+        .get_shard_layout(epoch_manager.get_block_info(block_hash)?.epoch_id())?
+        .shard_uids()
+        .filter(|shard_uid| {
+            // Remove shards that we are tracking right now


"not tracking"

Updated comment to say // Remove shards that we are tracking from shards_to_cleanup

staffik · 2025-05-22T10:48:02Z

chain/chain/src/garbage_collection.rs

+            // Remove shards that we are tracking right now
+            !shard_tracker.cares_about_shard_this_or_next_epoch(
+                me,
+                block_info.prev_hash(),


nit: maybe rename it to last_block_info?

staffik · 2025-05-22T10:48:05Z

chain/chain/src/garbage_collection.rs

+
+    // reverse iterate over the epochs
+    let store = chain_store_update.store();
+    while !shards_to_cleanup.is_empty() && block_info.hash() != block_hash {


What about renaming block_hash to last_block_hash_in_gced_epoch.
block_info.hash() != block_hash might be confusing

staffik · 2025-05-22T10:50:06Z

chain/chain/src/garbage_collection.rs

+    let mut block_info = epoch_manager.get_block_info(&latest_block_hash)?;
+
+    // Get all the shards that belong to the epoch we are cleaning up
+    let mut shards_to_cleanup = epoch_manager


nit: maybe rename it to potential_shards_to_cleanup, as it is superset of shards that will be cleaned up

Umm.. I feel that makes the name of the variable a bit too long without adding too much extra information about it. I did try it but would prefer to stick with shards_to_cleanup

darioush · 2025-05-22T13:20:23Z

test-loop-tests/src/tests/resharding_v3.rs

-        .build();
+fn slow_test_resharding_v3_load_memtrie() {
+    let params =
+        TestReshardingParametersBuilder::default().load_memtries_for_tracked_shards(false).build();


kind of confusing why the test for _load_memtrie would set load_memtries_for_tracked_shards(false). should we add a comment?

Honestly I don't understand this test too much either to add any sensible comment :(

The reason to set load_memtries_for_tracked_shards(false) is that...
🥁
... we want to check if memtries are loaded (requirement for resharding) even when the node doesn't automatically load memtries for tracked shard.

Basically memtrie of parent must be loaded regardless of config

darioush · 2025-05-22T15:55:41Z

chain/chain/src/garbage_collection.rs

-        {
-            tracked_shards.push(shard_uid);
-        }
+        assert_eq!(get_shard_uid_mapping(&store, shard_uid), shard_uid, "Incomplete Resharding");


it wouldn't be deleted yet if I understand? seems it is just added to the batch so far.

[resharding] Rewrite state column gc

51009d0

shreyan-gupta requested review from wacban, darioush, Trisfald and staffik May 22, 2025 08:04

shreyan-gupta requested a review from a team as a code owner May 22, 2025 08:04

shreyan-gupta commented May 22, 2025

View reviewed changes

add missing comment

eb0491c

shreyan-gupta commented May 22, 2025

View reviewed changes

shreyan-gupta added 2 commits May 22, 2025 01:15

build

75b990b

remove failing test

6f20531

shreyan-gupta commented May 22, 2025

View reviewed changes

wacban approved these changes May 22, 2025

View reviewed changes

staffik reviewed May 22, 2025

View reviewed changes

This comment was marked as off-topic.

Sign in to view

darioush reviewed May 22, 2025

View reviewed changes

shreyan-gupta added 2 commits May 22, 2025 12:13

resolve comments

ca7305f

fix

d6fc327

shreyan-gupta added this pull request to the merge queue May 22, 2025

Merged via the queue into master with commit 89a06ac May 22, 2025
28 checks passed

shreyan-gupta deleted the shreyan/resharding/gc_new branch May 22, 2025 21:20

		let latest_block_hash = chain_store_update.head()?.last_block_hash;
		let mut block_info = epoch_manager.get_block_info(&latest_block_hash)?;

[resharding] Rewrite state column gc #13598

[resharding] Rewrite state column gc #13598

Conversation

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!