Skip to content

Commit e7ec1f3

Browse files
bittobybittoby
committed
fix: prevent zero-valued entries in Alpha storage maps and add multi-block on_idle migration to purge existing ones
1 parent 3feb9d2 commit e7ec1f3

14 files changed

Lines changed: 683 additions & 49 deletions

File tree

‎pallets/subtensor/src/coinbase/root.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -587,7 +587,11 @@ impl<T: Config> Pallet<T> {
587587
LastRateLimitedBlock::<T>::get(rate_limit_key)
588588
}
589589
pub fn set_rate_limited_last_block(rate_limit_key: &RateLimitKey<T::AccountId>, block: u64) {
590-
LastRateLimitedBlock::<T>::insert(rate_limit_key, block);
590+
if block == 0 {
591+
LastRateLimitedBlock::<T>::remove(rate_limit_key);
592+
} else {
593+
LastRateLimitedBlock::<T>::insert(rate_limit_key, block);
594+
}
591595
}
592596
pub fn remove_rate_limited_last_block(rate_limit_key: &RateLimitKey<T::AccountId>) {
593597
LastRateLimitedBlock::<T>::remove(rate_limit_key);

‎pallets/subtensor/src/coinbase/run_coinbase.rs‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -598,13 +598,18 @@ impl<T: Config> Pallet<T> {
598598
log::debug!("hotkey: {hotkey:?} alpha_divs: {alpha_divs:?}");
599599
Self::increase_stake_for_hotkey_on_subnet(&hotkey, netuid, tou64!(alpha_divs).into());
600600
// Record dividends for this hotkey.
601-
AlphaDividendsPerSubnet::<T>::mutate(netuid, &hotkey, |divs| {
602-
*divs = divs.saturating_add(tou64!(alpha_divs).into());
603-
});
601+
let alpha_divs_u64: u64 = tou64!(alpha_divs);
602+
if alpha_divs_u64 != 0 {
603+
AlphaDividendsPerSubnet::<T>::mutate(netuid, &hotkey, |divs| {
604+
*divs = divs.saturating_add(alpha_divs_u64.into());
605+
});
606+
}
604607
// Record total hotkey alpha based on which this value of AlphaDividendsPerSubnet
605608
// was calculated
606609
let total_hotkey_alpha = TotalHotkeyAlpha::<T>::get(&hotkey, netuid);
607-
TotalHotkeyAlphaLastEpoch::<T>::insert(hotkey, netuid, total_hotkey_alpha);
610+
if !total_hotkey_alpha.is_zero() {
611+
TotalHotkeyAlphaLastEpoch::<T>::insert(hotkey, netuid, total_hotkey_alpha);
612+
}
608613
}
609614

610615
// Distribute root alpha divs.

‎pallets/subtensor/src/lib.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2441,6 +2441,11 @@ pub mod pallet {
24412441
#[pallet::storage]
24422442
pub type HasMigrationRun<T: Config> = StorageMap<_, Identity, Vec<u8>, bool, ValueQuery>;
24432443

2444+
/// --- Tracks the current phase of the zero-alpha multi-block cleanup.
2445+
/// 0 = inactive/complete, 1-4 = active phases (Alpha, TotalHotkeyShares, etc.)
2446+
#[pallet::storage]
2447+
pub type ZeroAlphaCleanupPhase<T: Config> = StorageValue<_, u8, ValueQuery>;
2448+
24442449
/// Default value for pending childkey cooldown (settable by root).
24452450
/// Uses the same value as DefaultPendingCooldown for consistency.
24462451
#[pallet::type_value]

‎pallets/subtensor/src/macros/hooks.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ mod hooks {
3838
}
3939
}
4040

41+
// ---- Called when the block has leftover weight. Used for multi-block migrations.
42+
fn on_idle(_block_number: BlockNumberFor<T>, remaining_weight: Weight) -> Weight {
43+
migrations::migrate_remove_zero_alpha::on_idle_remove_zero_alpha::<T>(remaining_weight)
44+
}
45+
4146
// ---- Called on the finalization of this pallet. The code weight must be taken into account prior to the execution of this macro.
4247
//
4348
// # Args:
@@ -166,7 +171,9 @@ mod hooks {
166171
// Fix staking hot keys
167172
.saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::<T>())
168173
// Migrate coldkey swap scheduled to announcements
169-
.saturating_add(migrations::migrate_coldkey_swap_scheduled_to_announcements::migrate_coldkey_swap_scheduled_to_announcements::<T>());
174+
.saturating_add(migrations::migrate_coldkey_swap_scheduled_to_announcements::migrate_coldkey_swap_scheduled_to_announcements::<T>())
175+
// Remove zero-valued entries from Alpha and related storage maps
176+
.saturating_add(migrations::migrate_remove_zero_alpha::migrate_remove_zero_alpha::<T>());
170177
weight
171178
}
172179

‎pallets/subtensor/src/migrations/migrate_remove_zero_alpha.rs‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
use super::*;
2+
use frame_support::{traits::Get, weights::Weight};
3+
use log;
4+
use scale_info::prelude::string::String;
5+
6+
/// The migration name used for the `HasMigrationRun` guard.
7+
const MIGRATION_NAME: &[u8] = b"migrate_remove_zero_alpha_v2";
8+
9+
/// Called from `on_runtime_upgrade`. Schedules the cleanup by setting phase = 1
10+
/// if the migration hasn't run yet. This is O(1) — no iteration.
11+
pub fn migrate_remove_zero_alpha<T: Config>() -> Weight {
12+
let migration_name = MIGRATION_NAME.to_vec();
13+
let mut weight = T::DbWeight::get().reads(1);
14+
15+
if HasMigrationRun::<T>::get(&migration_name) {
16+
log::info!(
17+
"Migration '{}' already completed. Skipping.",
18+
String::from_utf8_lossy(&migration_name)
19+
);
20+
return weight;
21+
}
22+
23+
// Schedule the cleanup to run in on_idle by setting phase to 1
24+
ZeroAlphaCleanupPhase::<T>::put(1u8);
25+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
26+
27+
log::info!(
28+
"Migration '{}' scheduled. Will clean up zero entries via on_idle.",
29+
String::from_utf8_lossy(&migration_name)
30+
);
31+
32+
weight
33+
}
34+
35+
/// Called from `on_idle` each block. Processes one storage map per block,
36+
/// removing all zero-valued entries. Advances to the next phase when done.
37+
///
38+
/// Phases:
39+
/// 0 = inactive/complete
40+
/// 1 = cleaning Alpha
41+
/// 2 = cleaning TotalHotkeyShares
42+
/// 3 = cleaning TotalHotkeyAlphaLastEpoch
43+
/// 4 = cleaning AlphaDividendsPerSubnet
44+
pub fn on_idle_remove_zero_alpha<T: Config>(remaining_weight: Weight) -> Weight {
45+
let phase = ZeroAlphaCleanupPhase::<T>::get();
46+
47+
// Phase 0 means not active or already completed
48+
if phase == 0 {
49+
return Weight::zero();
50+
}
51+
52+
// Minimum weight needed: 1 read (phase) + at least some work
53+
let min_weight = T::DbWeight::get().reads_writes(2, 1);
54+
if remaining_weight.ref_time() < min_weight.ref_time() {
55+
return Weight::zero();
56+
}
57+
58+
let mut weight = T::DbWeight::get().reads(1); // reading phase
59+
60+
match phase {
61+
1 => {
62+
let (consumed, removed) = clean_alpha::<T>();
63+
weight = weight.saturating_add(consumed);
64+
log::info!("Zero-alpha cleanup: Alpha complete. Removed {removed} zero entries.");
65+
ZeroAlphaCleanupPhase::<T>::put(2u8);
66+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
67+
}
68+
2 => {
69+
let (consumed, removed) = clean_total_hotkey_shares::<T>();
70+
weight = weight.saturating_add(consumed);
71+
log::info!(
72+
"Zero-alpha cleanup: TotalHotkeyShares complete. Removed {removed} zero entries."
73+
);
74+
ZeroAlphaCleanupPhase::<T>::put(3u8);
75+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
76+
}
77+
3 => {
78+
let (consumed, removed) = clean_total_hotkey_alpha_last_epoch::<T>();
79+
weight = weight.saturating_add(consumed);
80+
log::info!(
81+
"Zero-alpha cleanup: TotalHotkeyAlphaLastEpoch complete. Removed {removed} zero entries."
82+
);
83+
ZeroAlphaCleanupPhase::<T>::put(4u8);
84+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
85+
}
86+
4 => {
87+
let (consumed, removed) = clean_alpha_dividends_per_subnet::<T>();
88+
weight = weight.saturating_add(consumed);
89+
log::info!(
90+
"Zero-alpha cleanup: AlphaDividendsPerSubnet complete. Removed {removed} zero entries."
91+
);
92+
93+
// All phases complete — mark migration as done
94+
HasMigrationRun::<T>::insert(MIGRATION_NAME.to_vec(), true);
95+
ZeroAlphaCleanupPhase::<T>::put(0u8);
96+
weight = weight.saturating_add(T::DbWeight::get().writes(2));
97+
log::info!("Zero-alpha cleanup: All phases complete. Migration marked as done.");
98+
}
99+
_ => {
100+
// Unknown phase, reset
101+
ZeroAlphaCleanupPhase::<T>::put(0u8);
102+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
103+
}
104+
}
105+
106+
weight
107+
}
108+
109+
/// Remove all zero-valued entries from Alpha.
110+
/// Returns (weight_consumed, entries_removed).
111+
fn clean_alpha<T: Config>() -> (Weight, u64) {
112+
let mut weight = Weight::zero();
113+
let mut removed = 0u64;
114+
115+
let to_remove: Vec<_> = Alpha::<T>::iter()
116+
.filter_map(|((hotkey, coldkey, netuid), value)| {
117+
weight = weight.saturating_add(T::DbWeight::get().reads(1));
118+
if value == 0 {
119+
Some((hotkey, coldkey, netuid))
120+
} else {
121+
None
122+
}
123+
})
124+
.collect();
125+
126+
for (hotkey, coldkey, netuid) in &to_remove {
127+
Alpha::<T>::remove((hotkey, coldkey, *netuid));
128+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
129+
removed = removed.saturating_add(1);
130+
}
131+
132+
(weight, removed)
133+
}
134+
135+
/// Remove all zero-valued entries from TotalHotkeyShares.
136+
fn clean_total_hotkey_shares<T: Config>() -> (Weight, u64) {
137+
let mut weight = Weight::zero();
138+
let mut removed = 0u64;
139+
140+
let to_remove: Vec<_> = TotalHotkeyShares::<T>::iter()
141+
.filter_map(|(hotkey, netuid, value)| {
142+
weight = weight.saturating_add(T::DbWeight::get().reads(1));
143+
if value == 0 {
144+
Some((hotkey, netuid))
145+
} else {
146+
None
147+
}
148+
})
149+
.collect();
150+
151+
for (hotkey, netuid) in &to_remove {
152+
TotalHotkeyShares::<T>::remove(hotkey, *netuid);
153+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
154+
removed = removed.saturating_add(1);
155+
}
156+
157+
(weight, removed)
158+
}
159+
160+
/// Remove all zero-valued entries from TotalHotkeyAlphaLastEpoch.
161+
fn clean_total_hotkey_alpha_last_epoch<T: Config>() -> (Weight, u64) {
162+
let mut weight = Weight::zero();
163+
let mut removed = 0u64;
164+
165+
let to_remove: Vec<_> = TotalHotkeyAlphaLastEpoch::<T>::iter()
166+
.filter_map(|(hotkey, netuid, value)| {
167+
weight = weight.saturating_add(T::DbWeight::get().reads(1));
168+
if value.is_zero() {
169+
Some((hotkey, netuid))
170+
} else {
171+
None
172+
}
173+
})
174+
.collect();
175+
176+
for (hotkey, netuid) in &to_remove {
177+
TotalHotkeyAlphaLastEpoch::<T>::remove(hotkey, *netuid);
178+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
179+
removed = removed.saturating_add(1);
180+
}
181+
182+
(weight, removed)
183+
}
184+
185+
/// Remove all zero-valued entries from AlphaDividendsPerSubnet.
186+
fn clean_alpha_dividends_per_subnet<T: Config>() -> (Weight, u64) {
187+
let mut weight = Weight::zero();
188+
let mut removed = 0u64;
189+
190+
let to_remove: Vec<_> = AlphaDividendsPerSubnet::<T>::iter()
191+
.filter_map(|(netuid, hotkey, value)| {
192+
weight = weight.saturating_add(T::DbWeight::get().reads(1));
193+
if value.is_zero() {
194+
Some((netuid, hotkey))
195+
} else {
196+
None
197+
}
198+
})
199+
.collect();
200+
201+
for (netuid, hotkey) in &to_remove {
202+
AlphaDividendsPerSubnet::<T>::remove(*netuid, hotkey);
203+
weight = weight.saturating_add(T::DbWeight::get().writes(1));
204+
removed = removed.saturating_add(1);
205+
}
206+
207+
(weight, removed)
208+
}

‎pallets/subtensor/src/migrations/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ pub mod migrate_remove_tao_dividends;
4141
pub mod migrate_remove_total_hotkey_coldkey_stakes_this_interval;
4242
pub mod migrate_remove_unknown_neuron_axon_cert_prom;
4343
pub mod migrate_remove_unused_maps_and_values;
44+
pub mod migrate_remove_zero_alpha;
4445
pub mod migrate_remove_zero_total_hotkey_alpha;
4546
pub mod migrate_reset_bonds_moving_average;
4647
pub mod migrate_reset_max_burn;

‎pallets/subtensor/src/staking/claim_root.rs‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,8 +263,12 @@ impl<T: Config> Pallet<T> {
263263
.saturating_to_num(),
264264
);
265265

266-
// Set the new root claimed value.
267-
RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
266+
// Set the new root claimed value, or remove if zero to avoid storage bloat.
267+
if new_root_claimed != 0 {
268+
RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
269+
} else {
270+
RootClaimed::<T>::remove((netuid, hotkey, coldkey));
271+
}
268272
}
269273
}
270274

@@ -290,8 +294,12 @@ impl<T: Config> Pallet<T> {
290294
.saturating_to_num(),
291295
);
292296

293-
// Set the new root_claimed value.
294-
RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
297+
// Set the new root_claimed value, removing if zero to avoid storage bloat.
298+
if new_root_claimed != 0 {
299+
RootClaimed::<T>::insert((netuid, hotkey, coldkey), new_root_claimed);
300+
} else {
301+
RootClaimed::<T>::remove((netuid, hotkey, coldkey));
302+
}
295303
}
296304
}
297305

‎pallets/subtensor/src/staking/set_children.rs‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -412,15 +412,11 @@ impl<T: Config> Pallet<T> {
412412
for (parent, _) in relations.parents().iter() {
413413
let mut ck = ChildKeys::<T>::get(parent.clone(), netuid);
414414
PCRelations::<T>::remove_edge(&mut ck, old_hotkey);
415-
ChildKeys::<T>::insert(parent.clone(), netuid, ck);
415+
Self::set_childkeys(parent.clone(), netuid, ck);
416416
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
417417
}
418418
// 2c) Clear direct maps of old_hotkey
419-
ChildKeys::<T>::insert(
420-
old_hotkey.clone(),
421-
netuid,
422-
Vec::<(u64, T::AccountId)>::new(),
423-
);
419+
ChildKeys::<T>::remove(old_hotkey.clone(), netuid);
424420
Self::set_parentkeys(
425421
old_hotkey.clone(),
426422
netuid,

‎pallets/subtensor/src/swap/swap_coldkey.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,9 @@ impl<T: Config> Pallet<T> {
160160
}
161161

162162
StakingHotkeys::<T>::remove(old_coldkey);
163-
StakingHotkeys::<T>::insert(new_coldkey, new_staking_hotkeys);
163+
if !new_staking_hotkeys.is_empty() {
164+
StakingHotkeys::<T>::insert(new_coldkey, new_staking_hotkeys);
165+
}
164166
}
165167

166168
/// Transfer the ownership of the hotkeys owned by the old coldkey to the new coldkey.
@@ -178,6 +180,8 @@ impl<T: Config> Pallet<T> {
178180
}
179181
}
180182
OwnedHotkeys::<T>::remove(old_coldkey);
181-
OwnedHotkeys::<T>::insert(new_coldkey, new_owned_hotkeys);
183+
if !new_owned_hotkeys.is_empty() {
184+
OwnedHotkeys::<T>::insert(new_coldkey, new_owned_hotkeys);
185+
}
182186
}
183187
}

0 commit comments

Comments
 (0)