From 8b4ab1a99be22f1d02f86cd4227133c536a664e4 Mon Sep 17 00:00:00 2001 From: guibescos <59208140+guibescos@users.noreply.github.com> Date: Tue, 12 Nov 2024 15:25:43 +0000 Subject: [PATCH] feat: add script to snapshot all staking accounts (#542) * checkpoint * checkpoint * checkpoint * go * imports * clippy * some clean up * add custody authority --- staking/Cargo.lock | 2 + staking/cli/Cargo.toml | 2 + staking/cli/src/cli.rs | 1 + staking/cli/src/instructions.rs | 281 +++++++++++++++++++++++++++- staking/cli/src/main.rs | 4 + staking/programs/staking/src/lib.rs | 2 +- 6 files changed, 289 insertions(+), 3 deletions(-) diff --git a/staking/Cargo.lock b/staking/Cargo.lock index 7e1c4523..97cef11a 100644 --- a/staking/Cargo.lock +++ b/staking/Cargo.lock @@ -5279,6 +5279,7 @@ dependencies = [ "anchor-spl", "base64 0.22.1", "byteorder", + "chrono", "clap 3.2.25", "integration-tests", "integrity-pool", @@ -5289,6 +5290,7 @@ dependencies = [ "serde_json", "serde_wormhole", "shellexpand", + "solana-account-decoder", "solana-client", "solana-remote-wallet", "solana-sdk", diff --git a/staking/cli/Cargo.toml b/staking/cli/Cargo.toml index 47302437..9a3e3cf8 100644 --- a/staking/cli/Cargo.toml +++ b/staking/cli/Cargo.toml @@ -25,3 +25,5 @@ reqwest = "0.11" serde_json = "1.0.128" uriparse = "0.6.4" solana-remote-wallet = "1.18.16" +solana-account-decoder = "1.18.16" +chrono = "0.4.38" diff --git a/staking/cli/src/cli.rs b/staking/cli/src/cli.rs index 10fdb75a..174572bd 100644 --- a/staking/cli/src/cli.rs +++ b/staking/cli/src/cli.rs @@ -109,6 +109,7 @@ pub enum Action { #[clap(long, help = "Publisher caps")] publisher_caps: Pubkey, }, + SaveStakeAccountsSnapshot {}, } pub enum SignerSource { diff --git a/staking/cli/src/instructions.rs b/staking/cli/src/instructions.rs index 277ab620..83a575e7 100644 --- a/staking/cli/src/instructions.rs +++ b/staking/cli/src/instructions.rs @@ -1,12 +1,16 @@ use { anchor_lang::{ AccountDeserialize, + Discriminator, InstructionData, ToAccountMetas, }, anchor_spl::{ associated_token::spl_associated_token_account, - token::spl_token, + token::{ + spl_token, + TokenAccount, + }, }, base64::Engine, integration_tests::{ @@ -42,9 +46,19 @@ use { }, reqwest::blocking::Client, serde_wormhole::RawMessage, + solana_account_decoder::UiAccountEncoding, solana_client::{ rpc_client::RpcClient, - rpc_config::RpcSendTransactionConfig, + rpc_config::{ + RpcAccountInfoConfig, + RpcProgramAccountsConfig, + RpcSendTransactionConfig, + }, + rpc_filter::{ + Memcmp, + MemcmpEncodedBytes, + RpcFilterType, + }, }, solana_sdk::{ commitment_config::CommitmentConfig, @@ -67,9 +81,28 @@ use { TransactionError, }, }, + staking::{ + state::{ + global_config::GlobalConfig, + max_voter_weight_record::MAX_VOTER_WEIGHT, + positions::{ + DynamicPositionArrayAccount, + PositionData, + PositionState, + }, + stake_account::StakeAccountMetadataV2, + }, + utils::voter_weight::compute_voter_weight, + }, std::{ cmp::min, + collections::HashMap, convert::TryInto, + fs::File, + io::{ + BufWriter, + Write, + }, mem::size_of, }, wormhole_core_bridge_solana::sdk::{ @@ -498,6 +531,11 @@ pub fn initialize_pool( .unwrap(); } +pub fn get_current_time(rpc_client: &RpcClient) -> i64 { + let slot = rpc_client.get_slot().unwrap(); + rpc_client.get_block_time(slot).unwrap() +} + pub fn get_current_epoch(rpc_client: &RpcClient) -> u64 { let slot = rpc_client.get_slot().unwrap(); let blocktime = rpc_client.get_block_time(slot).unwrap(); @@ -825,3 +863,242 @@ pub fn update_y(rpc_client: &RpcClient, signer: &dyn Signer, y: u64) { process_transaction(rpc_client, &[instruction], &[signer]).unwrap(); } + +pub fn save_stake_accounts_snapshot(rpc_client: &RpcClient) { + let data: Vec<(Pubkey, DynamicPositionArrayAccount, Pubkey, Pubkey, Pubkey)> = rpc_client + .get_program_accounts_with_config( + &staking::ID, + RpcProgramAccountsConfig { + filters: Some(vec![RpcFilterType::Memcmp(Memcmp::new( + 0, + MemcmpEncodedBytes::Bytes(PositionData::discriminator().to_vec()), + ))]), + account_config: RpcAccountInfoConfig { + encoding: Some(UiAccountEncoding::Base64Zstd), + data_slice: None, + commitment: None, + min_context_slot: None, + }, + with_context: None, + }, + ) + .unwrap() + .into_iter() + .map(|(pubkey, account)| { + ( + pubkey, + DynamicPositionArrayAccount { + key: pubkey, + lamports: account.lamports, + data: account.data.clone(), + }, + get_stake_account_metadata_address(pubkey), + get_stake_account_custody_address(pubkey), + get_stake_account_custody_authority_address(pubkey), + ) + }) + .collect::>(); + + let metadata_accounts_data = rpc_client + .get_program_accounts_with_config( + &staking::ID, + RpcProgramAccountsConfig { + filters: Some(vec![RpcFilterType::Memcmp(Memcmp::new( + 0, + MemcmpEncodedBytes::Bytes(StakeAccountMetadataV2::discriminator().to_vec()), + ))]), + account_config: RpcAccountInfoConfig { + encoding: Some(UiAccountEncoding::Base64Zstd), + data_slice: None, + commitment: None, + min_context_slot: None, + }, + with_context: None, + }, + ) + .unwrap() + .into_iter() + .map(|(pubkey, account)| { + ( + pubkey, + StakeAccountMetadataV2::try_deserialize(&mut account.data.as_slice()).unwrap(), + ) + }) + .collect::>(); + + let config = GlobalConfig::try_deserialize( + &mut rpc_client + .get_account_data(&get_config_address()) + .unwrap() + .as_slice(), + ) + .unwrap(); + let current_time = get_current_time(rpc_client); + + let metadata_account_data_locked: HashMap = metadata_accounts_data + .iter() + .map(|(pubkey, metadata)| { + ( + *pubkey, + metadata + .lock + .get_unvested_balance(current_time, config.pyth_token_list_time) + .unwrap(), + ) + }) + .collect::>(); + + let data = data + .into_iter() + .map( + |(pubkey, account, metadata_pubkey, custody_pubkey, custody_authority_pubkey)| { + ( + pubkey, + account, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + *metadata_account_data_locked.get(&metadata_pubkey).unwrap(), + ) + }, + ) + .collect::>(); + + // We need to check the actual tokens accounts, since you can initialize a stake account with an + // arbitrary vesting schedule but 0 tokens + let locked_token_accounts_pubkeys = data + .iter() + .filter(|(_, _, _, _, _, locked_amount)| *locked_amount > 0u64) + .map(|(_, _, _, token_account_pubkey, _, _)| *token_account_pubkey) + .collect::>(); + + let mut locked_token_accounts_actual_amounts: HashMap = HashMap::new(); + for chunk in locked_token_accounts_pubkeys.chunks(100) { + rpc_client + .get_multiple_accounts(chunk) + .unwrap() + .into_iter() + .enumerate() + .for_each(|(index, account)| { + locked_token_accounts_actual_amounts.insert( + chunk[index], + TokenAccount::try_deserialize(&mut account.unwrap().data.as_slice()) + .unwrap() + .amount, + ); + }); + } + + let data = data + .into_iter() + .map( + |( + pubkey, + account, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + locked_amount, + )| { + ( + pubkey, + account, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + min( + locked_amount, + *locked_token_accounts_actual_amounts + .get(&custody_pubkey) + .unwrap_or(&0u64), + ), + ) + }, + ) + .collect::>(); + + let current_epoch = get_current_epoch(rpc_client); + let data = data + .into_iter() + .map( + |( + pubkey, + mut account, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + locked_amount, + )| { + let dynamic_position_array = account.to_dynamic_position_array(); + let owner = dynamic_position_array.owner().unwrap(); + let staked_in_governance = compute_voter_weight( + &dynamic_position_array, + current_epoch, + MAX_VOTER_WEIGHT, + MAX_VOTER_WEIGHT, + ) + .unwrap(); + + let staked_in_ois = { + let mut amount = 0u64; + for i in 0..dynamic_position_array.get_position_capacity() { + if let Some(position) = dynamic_position_array.read_position(i).unwrap() { + match position.get_current_position(current_epoch).unwrap() { + PositionState::LOCKED | PositionState::PREUNLOCKING => { + if !position.is_voting() { + amount += position.amount; + } + } + _ => {} + } + } + } + amount + }; + ( + pubkey, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + owner, + locked_amount, + staked_in_governance, + staked_in_ois, + ) + }, + ) + .collect::>(); + + let timestamp = chrono::Utc::now().format("%Y-%m-%d_%H:%M:%S").to_string(); + let file = File::create(format!("snapshots/snapshot-{}.csv", timestamp)).unwrap(); + let mut writer = BufWriter::new(file); + + // Write the header + writeln!(writer, "positions_pubkey,metadata_pubkey,custody_pubkey,custody_authority_pubkey,owner,locked_amount,staked_in_governance,staked_in_ois").unwrap(); + // Write the data + for ( + pubkey, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + owner, + locked_amount, + staked_in_governance, + staked_in_ois, + ) in data + { + writeln!( + writer, + "{},{},{},{},{},{},{},{}", + pubkey, + metadata_pubkey, + custody_pubkey, + custody_authority_pubkey, + owner, + locked_amount, + staked_in_governance, + staked_in_ois + ) + .unwrap(); + } +} diff --git a/staking/cli/src/main.rs b/staking/cli/src/main.rs index 5c6a0cc8..a06c0b9f 100644 --- a/staking/cli/src/main.rs +++ b/staking/cli/src/main.rs @@ -13,6 +13,7 @@ use { fetch_publisher_caps_and_advance, initialize_pool, initialize_reward_custody, + save_stake_accounts_snapshot, set_publisher_stake_account, slash, update_delegation_fee, @@ -93,5 +94,8 @@ fn main() { Action::ClosePublisherCaps { publisher_caps } => { close_publisher_caps(&rpc_client, keypair.as_ref(), publisher_caps) } + Action::SaveStakeAccountsSnapshot {} => { + save_stake_accounts_snapshot(&rpc_client); + } } } diff --git a/staking/programs/staking/src/lib.rs b/staking/programs/staking/src/lib.rs index fc3226a4..c2412cea 100644 --- a/staking/programs/staking/src/lib.rs +++ b/staking/programs/staking/src/lib.rs @@ -41,7 +41,7 @@ use { pub mod context; pub mod error; pub mod state; -mod utils; +pub mod utils; #[cfg(feature = "wasm")] pub mod wasm;