Skip to content

Commit 86e9d7b

Browse files
gztensorgztensor
committed
Add Owned map and migration to populate it
1 parent ac74347 commit 86e9d7b

File tree

5 files changed

+113
-17
lines changed

5 files changed

+113
-17
lines changed

‎pallets/subtensor/src/lib.rs‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,9 @@ pub mod pallet {
363363
#[pallet::storage] // --- MAP ( hot ) --> cold | Returns the controlling coldkey for a hotkey.
364364
pub type Owner<T: Config> =
365365
StorageMap<_, Blake2_128Concat, T::AccountId, T::AccountId, ValueQuery, DefaultAccount<T>>;
366+
#[pallet::storage] // --- MAP ( cold ) --> Vec<hot> | Returns the vector of hotkeys controlled by this coldkey.
367+
pub type Owned<T: Config> =
368+
StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
366369
#[pallet::storage] // --- MAP ( hot ) --> take | Returns the hotkey delegation take. And signals that this key is open for delegation.
367370
pub type Delegates<T: Config> =
368371
StorageMap<_, Blake2_128Concat, T::AccountId, u16, ValueQuery, DefaultDefaultTake<T>>;
@@ -1204,6 +1207,14 @@ pub mod pallet {
12041207
// Fill stake information.
12051208
Owner::<T>::insert(hotkey.clone(), coldkey.clone());
12061209

1210+
// Update Owned map
1211+
let mut hotkeys = Owned::<T>::get(&coldkey);
1212+
hotkeys.push(hotkey.clone());
1213+
Owned::<T>::insert(
1214+
&coldkey,
1215+
hotkeys,
1216+
);
1217+
12071218
TotalHotkeyStake::<T>::insert(hotkey.clone(), stake);
12081219
TotalColdkeyStake::<T>::insert(
12091220
coldkey.clone(),
@@ -1325,7 +1336,9 @@ pub mod pallet {
13251336
// Storage version v4 -> v5
13261337
.saturating_add(migration::migrate_delete_subnet_3::<T>())
13271338
// Doesn't check storage version. TODO: Remove after upgrade
1328-
.saturating_add(migration::migration5_total_issuance::<T>(false));
1339+
.saturating_add(migration::migration5_total_issuance::<T>(false))
1340+
// Populate Owned map for coldkey swap. Doesn't update storage vesion.
1341+
.saturating_add(migration::migrate_populate_owned::<T>());
13291342

13301343
weight
13311344
}
@@ -1970,6 +1983,19 @@ pub mod pallet {
19701983
Self::do_swap_hotkey(origin, &hotkey, &new_hotkey)
19711984
}
19721985

1986+
/// The extrinsic for user to change the coldkey
1987+
#[pallet::call_index(71)]
1988+
#[pallet::weight((Weight::from_parts(1_940_000_000, 0)
1989+
.saturating_add(T::DbWeight::get().reads(272))
1990+
.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
1991+
pub fn swap_coldkey(
1992+
origin: OriginFor<T>,
1993+
old_coldkey: T::AccountId,
1994+
new_coldkey: T::AccountId,
1995+
) -> DispatchResultWithPostInfo {
1996+
Self::do_swap_coldkey(origin, &old_coldkey, &new_coldkey)
1997+
}
1998+
19731999
// ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------
19742000

19752001
// ==================================

‎pallets/subtensor/src/migration.rs‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -477,3 +477,44 @@ pub fn migrate_to_v2_fixed_total_stake<T: Config>() -> Weight {
477477
Weight::zero()
478478
}
479479
}
480+
481+
pub fn migrate_populate_owned<T: Config>() -> Weight {
482+
// Setup migration weight
483+
let mut weight = T::DbWeight::get().reads(1);
484+
let migration_name = "Populate Owned map";
485+
486+
// Check if this migration is needed (if Owned map is empty)
487+
let migrate = Owned::<T>::iter().next().is_none();
488+
489+
// Only runs if we haven't already updated version past above new_storage_version.
490+
if migrate {
491+
info!(target: LOG_TARGET_1, ">>> Migration: {}", migration_name);
492+
493+
let mut longest_hotkey_vector = 0;
494+
let mut longest_coldkey: Option<T::AccountId> = None;
495+
Owner::<T>::iter().for_each(|(hotkey, coldkey)| {
496+
let mut hotkeys = Owned::<T>::get(&coldkey);
497+
hotkeys.push(hotkey);
498+
if longest_hotkey_vector < hotkeys.len() {
499+
longest_hotkey_vector = hotkeys.len();
500+
longest_coldkey = Some(coldkey.clone());
501+
}
502+
503+
Owned::<T>::insert(
504+
&coldkey,
505+
hotkeys,
506+
);
507+
508+
weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1));
509+
});
510+
info!(target: LOG_TARGET_1, "Migration {} finished. Longest hotkey vector: {}", migration_name, longest_hotkey_vector);
511+
if let Some(c) = longest_coldkey {
512+
info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {}", c);
513+
}
514+
515+
weight
516+
} else {
517+
info!(target: LOG_TARGET_1, "Migration {} already done!", migration_name);
518+
Weight::zero()
519+
}
520+
}

‎pallets/subtensor/src/staking.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -560,6 +560,14 @@ impl<T: Config> Pallet<T> {
560560
if !Self::hotkey_account_exists(hotkey) {
561561
Stake::<T>::insert(hotkey, coldkey, 0);
562562
Owner::<T>::insert(hotkey, coldkey);
563+
564+
// Update Owned map
565+
let mut hotkeys = Owned::<T>::get(&coldkey);
566+
hotkeys.push(hotkey.clone());
567+
Owned::<T>::insert(
568+
&coldkey,
569+
hotkeys,
570+
);
563571
}
564572
}
565573

‎pallets/subtensor/src/swap.rs‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,15 @@ impl<T: Config> Pallet<T> {
223223
) {
224224
Owner::<T>::remove(old_hotkey);
225225
Owner::<T>::insert(new_hotkey, coldkey.clone());
226+
227+
// Update Owned map
228+
let mut hotkeys = Owned::<T>::get(&coldkey);
229+
hotkeys.push(new_hotkey.clone());
230+
Owned::<T>::insert(
231+
&coldkey,
232+
hotkeys,
233+
);
234+
226235
weight.saturating_accrue(T::DbWeight::get().writes(2));
227236
}
228237

@@ -535,11 +544,15 @@ impl<T: Config> Pallet<T> {
535544
new_coldkey: &T::AccountId,
536545
weight: &mut Weight,
537546
) {
538-
for (hotkey, stake) in Stake::<T>::iter_prefix(old_coldkey) {
539-
Stake::<T>::remove(old_coldkey, &hotkey);
540-
Stake::<T>::insert(new_coldkey, &hotkey, stake);
547+
// Find all hotkeys for this coldkey
548+
let hotkeys = Owned::<T>::get(old_coldkey);
549+
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0));
550+
for hotkey in hotkeys.iter() {
551+
let stake = Stake::<T>::get(&hotkey, old_coldkey);
552+
Stake::<T>::remove(&hotkey, old_coldkey);
553+
Stake::<T>::insert(&hotkey, new_coldkey, stake);
554+
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
541555
}
542-
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
543556
}
544557

545558
/// Swaps the owner of all hotkeys from the old coldkey to the new coldkey.
@@ -559,12 +572,17 @@ impl<T: Config> Pallet<T> {
559572
new_coldkey: &T::AccountId,
560573
weight: &mut Weight,
561574
) {
562-
for (hotkey, _) in Owner::<T>::iter() {
563-
if Owner::<T>::get(&hotkey) == *old_coldkey {
564-
Owner::<T>::insert(&hotkey, new_coldkey);
565-
}
575+
let hotkeys = Owned::<T>::get(old_coldkey);
576+
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 0));
577+
for hotkey in hotkeys.iter() {
578+
Owner::<T>::insert(&hotkey, new_coldkey);
579+
weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1));
566580
}
567-
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
581+
582+
// Update Owned map with new coldkey
583+
Owned::<T>::remove(old_coldkey);
584+
Owned::<T>::insert(new_coldkey, hotkeys);
585+
weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2));
568586
}
569587

570588
/// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey.

‎pallets/subtensor/tests/swap.rs‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1210,20 +1210,23 @@ fn test_swap_stake_for_coldkey() {
12101210
let mut weight = Weight::zero();
12111211

12121212
// Initialize Stake for old_coldkey
1213-
Stake::<Test>::insert(old_coldkey, hotkey1, stake_amount1);
1214-
Stake::<Test>::insert(old_coldkey, hotkey2, stake_amount2);
1213+
Stake::<Test>::insert(hotkey1, old_coldkey, stake_amount1);
1214+
Stake::<Test>::insert(hotkey2, old_coldkey, stake_amount2);
1215+
1216+
// Populate Owned map
1217+
Owned::<Test>::insert(old_coldkey, vec![hotkey1, hotkey2]);
12151218

12161219
// Perform the swap
12171220
SubtensorModule::swap_stake_for_coldkey(&old_coldkey, &new_coldkey, &mut weight);
12181221

12191222
// Verify the swap
1220-
assert_eq!(Stake::<Test>::get(new_coldkey, hotkey1), stake_amount1);
1221-
assert_eq!(Stake::<Test>::get(new_coldkey, hotkey2), stake_amount2);
1222-
assert!(!Stake::<Test>::contains_key(old_coldkey, hotkey1));
1223-
assert!(!Stake::<Test>::contains_key(old_coldkey, hotkey2));
1223+
assert_eq!(Stake::<Test>::get(hotkey1, new_coldkey), stake_amount1);
1224+
assert_eq!(Stake::<Test>::get(hotkey2, new_coldkey), stake_amount2);
1225+
assert!(!Stake::<Test>::contains_key(hotkey1, old_coldkey));
1226+
assert!(!Stake::<Test>::contains_key(hotkey2, old_coldkey));
12241227

12251228
// Verify weight update
1226-
let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(1, 4);
1229+
let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(3, 4);
12271230
assert_eq!(weight, expected_weight);
12281231
});
12291232
}

0 commit comments

Comments
 (0)