Skip to content

Commit

Permalink
chore: pr review, make swaps free
Browse files Browse the repository at this point in the history
  • Loading branch information
Samuel Dare committed Jul 4, 2024
1 parent c82f82f commit 0855df7
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 108 deletions.
36 changes: 29 additions & 7 deletions pallets/subtensor/src/migration.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,7 @@ pub fn migrate_to_v2_fixed_total_stake<T: Config>() -> Weight {
}
}

/// Migrate the OwnedHotkeys map to the new storage format
pub fn migrate_populate_owned<T: Config>() -> Weight {
// Setup migration weight
let mut weight = T::DbWeight::get().reads(1);
Expand All @@ -486,27 +487,48 @@ pub fn migrate_populate_owned<T: Config>() -> Weight {
// Check if this migration is needed (if OwnedHotkeys map is empty)
let migrate = OwnedHotkeys::<T>::iter().next().is_none();

// Only runs if we haven't already updated version past above new_storage_version.
// Only runs if the migration is needed
if migrate {
info!(target: LOG_TARGET_1, ">>> Migration: {}", migration_name);
info!(target: LOG_TARGET_1, ">>> Starting Migration: {}", migration_name);

let mut longest_hotkey_vector = 0;
let mut longest_hotkey_vector: usize = 0;
let mut longest_coldkey: Option<T::AccountId> = None;
let mut keys_touched: u64 = 0;
let mut storage_reads: u64 = 0;
let mut storage_writes: u64 = 0;

// Iterate through all Owner entries
Owner::<T>::iter().for_each(|(hotkey, coldkey)| {
storage_reads = storage_reads.saturating_add(1); // Read from Owner storage
let mut hotkeys = OwnedHotkeys::<T>::get(&coldkey);
storage_reads = storage_reads.saturating_add(1); // Read from OwnedHotkeys storage

// Add the hotkey if it's not already in the vector
if !hotkeys.contains(&hotkey) {
hotkeys.push(hotkey);
keys_touched = keys_touched.saturating_add(1);

// Update longest hotkey vector info
if longest_hotkey_vector < hotkeys.len() {
longest_hotkey_vector = hotkeys.len();
longest_coldkey = Some(coldkey.clone());
}
}

OwnedHotkeys::<T>::insert(&coldkey, hotkeys);
// Update the OwnedHotkeys storage
OwnedHotkeys::<T>::insert(&coldkey, hotkeys);
storage_writes = storage_writes.saturating_add(1); // Write to OwnedHotkeys storage
}

weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1));
// Accrue weight for reads and writes
weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
});
info!(target: LOG_TARGET_1, "Migration {} finished. Longest hotkey vector: {}", migration_name, longest_hotkey_vector);

// Log migration results
info!(
target: LOG_TARGET_1,
"Migration {} finished. Keys touched: {}, Longest hotkey vector: {}, Storage reads: {}, Storage writes: {}",
migration_name, keys_touched, longest_hotkey_vector, storage_reads, storage_writes
);
if let Some(c) = longest_coldkey {
info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {:?}", c);
}
Expand Down
52 changes: 4 additions & 48 deletions pallets/subtensor/src/swap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,16 +53,6 @@ impl<T: Config> Pallet<T> {
T::DbWeight::get().reads((TotalNetworks::<T>::get().saturating_add(1u16)) as u64),
);

let swap_cost = Self::get_hotkey_swap_cost();
log::debug!("Swap cost: {:?}", swap_cost);

ensure!(
Self::can_remove_balance_from_coldkey_account(&coldkey, swap_cost),
Error::<T>::NotEnoughBalanceToPaySwapHotKey
);
let actual_burn_amount = Self::remove_balance_from_coldkey_account(&coldkey, swap_cost)?;
Self::burn_tokens(actual_burn_amount);

Self::swap_owner(old_hotkey, new_hotkey, &coldkey, &mut weight);
Self::swap_total_hotkey_stake(old_hotkey, new_hotkey, &mut weight);
Self::swap_delegates(old_hotkey, new_hotkey, &mut weight);
Expand Down Expand Up @@ -125,38 +115,17 @@ impl<T: Config> Pallet<T> {
old_coldkey: &T::AccountId,
new_coldkey: &T::AccountId,
) -> DispatchResultWithPostInfo {
let caller = ensure_signed(origin)?;

// Ensure the caller is the old coldkey
ensure!(caller == *old_coldkey, Error::<T>::NonAssociatedColdKey);
ensure_signed(origin)?;

let mut weight = T::DbWeight::get().reads(2);

// Check if the new coldkey is already associated with any hotkeys
ensure!(
old_coldkey != new_coldkey,
Error::<T>::NewColdKeyIsSameWithOld
!Self::coldkey_has_associated_hotkeys(new_coldkey),
Error::<T>::ColdKeyAlreadyAssociated
);

// // Check if the new coldkey is already associated with any hotkeys
// ensure!(
// !Self::coldkey_has_associated_hotkeys(new_coldkey),
// Error::<T>::ColdKeyAlreadyAssociated
// );

let block: u64 = Self::get_current_block_as_u64();
// ensure!(
// !Self::exceeds_tx_rate_limit(Self::get_last_tx_block(old_coldkey), block),
// Error::<T>::ColdKeySwapTxRateLimitExceeded
// );

// Note: we probably want to make this free
let swap_cost = Self::get_coldkey_swap_cost();
ensure!(
Self::can_remove_balance_from_coldkey_account(old_coldkey, swap_cost),
Error::<T>::NotEnoughBalanceToPaySwapColdKey
);
let actual_burn_amount = Self::remove_balance_from_coldkey_account(old_coldkey, swap_cost)?;
Self::burn_tokens(actual_burn_amount);

// Swap coldkey references in storage maps
Self::swap_total_coldkey_stake(old_coldkey, new_coldkey, &mut weight);
Expand Down Expand Up @@ -671,17 +640,4 @@ impl<T: Config> Pallet<T> {
OwnedHotkeys::<T>::insert(new_coldkey, hotkeys);
weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 2));
}

/// Returns the cost of swapping a coldkey.
///
/// # Returns
///
/// * `u64` - The cost of swapping a coldkey in Rao.
///
/// # Note
///
/// This function returns a hardcoded value. In a production environment, this should be configurable or determined dynamically.
pub fn get_coldkey_swap_cost() -> u64 {
1_000_000 // Example cost in Rao
}
}
57 changes: 5 additions & 52 deletions pallets/subtensor/tests/swap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1057,16 +1057,12 @@ fn test_do_swap_coldkey_success() {
let hotkey = U256::from(3);
let netuid = 1u16;
let stake_amount = 1000u64;
let swap_cost = SubtensorModule::get_coldkey_swap_cost();
let free_balance = 12345;

// Setup initial state
add_network(netuid, 13, 0);
register_ok_neuron(netuid, hotkey, old_coldkey, 0);
SubtensorModule::add_balance_to_coldkey_account(
&old_coldkey,
stake_amount + swap_cost + free_balance,
);
SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + free_balance);

// Add stake to the neuron
assert_ok!(SubtensorModule::add_stake(
Expand All @@ -1093,7 +1089,7 @@ fn test_do_swap_coldkey_success() {

// Get coldkey free balance before swap
let balance = SubtensorModule::get_coldkey_balance(&old_coldkey);
assert_eq!(balance, free_balance + swap_cost);
assert_eq!(balance, free_balance);

// Perform the swap
assert_ok!(SubtensorModule::do_swap_coldkey(
Expand All @@ -1112,10 +1108,7 @@ fn test_do_swap_coldkey_success() {
assert!(!OwnedHotkeys::<Test>::contains_key(old_coldkey));

// Verify balance transfer
assert_eq!(
SubtensorModule::get_coldkey_balance(&new_coldkey),
balance - swap_cost
);
assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance);
assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0);

// Verify event emission
Expand All @@ -1129,45 +1122,6 @@ fn test_do_swap_coldkey_success() {
});
}

#[test]
fn test_do_swap_coldkey_not_enough_balance() {
new_test_ext(1).execute_with(|| {
let old_coldkey = U256::from(1);
let new_coldkey = U256::from(2);
let swap_cost = SubtensorModule::get_coldkey_swap_cost();

// Setup initial state with insufficient balance
SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, swap_cost - 1);

// Attempt the swap
assert_err!(
SubtensorModule::do_swap_coldkey(
<<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
&old_coldkey,
&new_coldkey
),
Error::<Test>::NotEnoughBalanceToPaySwapColdKey
);
});
}

#[test]
fn test_do_swap_coldkey_same_keys() {
new_test_ext(1).execute_with(|| {
let coldkey = U256::from(1);

// Attempt the swap with same old and new coldkeys
assert_err!(
SubtensorModule::do_swap_coldkey(
<<Test as Config>::RuntimeOrigin>::signed(coldkey),
&coldkey,
&coldkey
),
Error::<Test>::NewColdKeyIsSameWithOld
);
});
}

#[test]
fn test_swap_total_coldkey_stake() {
new_test_ext(1).execute_with(|| {
Expand Down Expand Up @@ -1339,8 +1293,7 @@ fn test_do_swap_coldkey_with_subnet_ownership() {
let new_coldkey = U256::from(2);
let hotkey = U256::from(3);
let netuid = 1u16;
let stake_amount = 1000u64;
let swap_cost = SubtensorModule::get_coldkey_swap_cost();
let stake_amount: u64 = 1000u64;

// Setup initial state
add_network(netuid, 13, 0);
Expand All @@ -1349,7 +1302,7 @@ fn test_do_swap_coldkey_with_subnet_ownership() {
// Set TotalNetworks because swap relies on it
pallet_subtensor::TotalNetworks::<Test>::set(1);

SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + swap_cost);
SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount);
SubnetOwner::<Test>::insert(netuid, old_coldkey);

// Populate OwnedHotkeys map
Expand Down
2 changes: 1 addition & 1 deletion runtime/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
// `spec_version`, and `authoring_version` are the same between Wasm and native.
// This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
// the compatible custom types.
spec_version: 154,
spec_version: 156,
impl_version: 1,
apis: RUNTIME_API_VERSIONS,
transaction_version: 1,
Expand Down

0 comments on commit 0855df7

Please sign in to comment.