diff --git a/Cargo.lock b/Cargo.lock index 41b85130f..1d1bbb20b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4845,6 +4845,7 @@ name = "solana-stake-accounts" version = "1.2.0" dependencies = [ "clap", + "itertools 0.9.0", "solana-clap-utils", "solana-cli-config", "solana-client", diff --git a/stake-accounts/Cargo.toml b/stake-accounts/Cargo.toml index 085d1ed26..7bec4d18b 100644 --- a/stake-accounts/Cargo.toml +++ b/stake-accounts/Cargo.toml @@ -10,6 +10,7 @@ homepage = "https://solana.com/" [dependencies] clap = "2.33.0" +itertools = "0.9.0" solana-clap-utils = { path = "../clap-utils", version = "1.2.0" } solana-cli-config = { path = "../cli-config", version = "1.2.0" } solana-client = { path = "../client", version = "1.2.0" } diff --git a/stake-accounts/src/arg_parser.rs b/stake-accounts/src/arg_parser.rs index a6bbe281c..812c72900 100644 --- a/stake-accounts/src/arg_parser.rs +++ b/stake-accounts/src/arg_parser.rs @@ -244,7 +244,19 @@ where .arg(lockup_epoch_arg()) .arg(lockup_date_arg()) .arg(new_custodian_arg()) - .arg(num_accounts_arg()), + .arg(num_accounts_arg()) + .arg( + Arg::with_name("no_wait") + .long("no-wait") + .help("Send transactions without waiting for confirmation"), + ) + .arg( + Arg::with_name("unlock_years") + .long("unlock-years") + .takes_value(true) + .value_name("NUMBER") + .help("Years to unlock after the cliff"), + ), ) .subcommand( SubCommand::with_name("rebase") @@ -316,6 +328,8 @@ fn parse_set_lockup_args(matches: &ArgMatches<'_>) -> SetLockupArgs { pub lockup_date: Option, pub new_custodian: Option

, pub num_accounts: usize, + pub no_wait: bool, + pub unlock_years: Option, } pub(crate) struct RebaseArgs { @@ -191,6 +193,8 @@ fn resolve_set_lockup_args( lockup_date: args.lockup_date, new_custodian: resolve_new_custodian(wallet_manager, &args.new_custodian)?, num_accounts: args.num_accounts, + no_wait: args.no_wait, + unlock_years: args.unlock_years, }; Ok(resolved_args) } diff --git a/stake-accounts/src/main.rs b/stake-accounts/src/main.rs index 7546b0df2..2e83ac7b9 100644 --- a/stake-accounts/src/main.rs +++ b/stake-accounts/src/main.rs @@ -6,6 +6,7 @@ use crate::arg_parser::parse_args; use crate::args::{ resolve_command, AuthorizeArgs, Command, MoveArgs, NewArgs, RebaseArgs, SetLockupArgs, }; +use itertools::Itertools; use solana_cli_config::Config; use solana_client::client_error::ClientError; use solana_client::rpc_client::RpcClient; @@ -17,7 +18,10 @@ use solana_sdk::{ signers::Signers, transaction::Transaction, }; -use solana_stake_program::stake_instruction::LockupArgs; +use solana_stake_program::{ + stake_instruction::LockupArgs, + stake_state::{Lockup, StakeState}, +}; use std::env; use std::error::Error; @@ -45,6 +49,26 @@ fn get_balances( .collect() } +fn get_lockup(client: &RpcClient, address: &Pubkey) -> Result { + client + .get_account(address) + .map(|account| StakeState::lockup_from(&account).unwrap()) +} + +fn get_lockups( + client: &RpcClient, + addresses: Vec, +) -> Result, ClientError> { + addresses + .into_iter() + .map(|pubkey| get_lockup(client, &pubkey).map(|bal| (pubkey, bal))) + .collect() +} + +fn unique_signers(signers: Vec<&dyn Signer>) -> Vec<&dyn Signer> { + signers.into_iter().unique_by(|s| s.pubkey()).collect_vec() +} + fn process_new_stake_account( client: &RpcClient, args: &NewArgs>, @@ -59,12 +83,12 @@ fn process_new_stake_account( &Pubkey::default(), args.index, ); - let signers = vec![ + let signers = unique_signers(vec![ &*args.fee_payer, &*args.funding_keypair, &*args.base_keypair, - ]; - let signature = send_message(client, message, &signers)?; + ]); + let signature = send_message(client, message, &signers, false)?; Ok(signature) } @@ -81,15 +105,12 @@ fn process_authorize_stake_accounts( &args.new_withdraw_authority, args.num_accounts, ); - let signers = vec![ + let signers = unique_signers(vec![ &*args.fee_payer, &*args.stake_authority, &*args.withdraw_authority, - ]; - for message in messages { - let signature = send_message(client, message, &signers)?; - println!("{}", signature); - } + ]); + send_messages(client, messages, &signers, false)?; Ok(()) } @@ -97,6 +118,10 @@ fn process_lockup_stake_accounts( client: &RpcClient, args: &SetLockupArgs>, ) -> Result<(), ClientError> { + let addresses = + stake_accounts::derive_stake_account_addresses(&args.base_pubkey, args.num_accounts); + let existing_lockups = get_lockups(&client, addresses)?; + let lockup = LockupArgs { epoch: args.lockup_epoch, unix_timestamp: args.lockup_date, @@ -104,16 +129,17 @@ fn process_lockup_stake_accounts( }; let messages = stake_accounts::lockup_stake_accounts( &args.fee_payer.pubkey(), - &args.base_pubkey, &args.custodian.pubkey(), &lockup, - args.num_accounts, + &existing_lockups, + args.unlock_years, ); - let signers = vec![&*args.fee_payer, &*args.custodian]; - for message in messages { - let signature = send_message(client, message, &signers)?; - println!("{}", signature); + if messages.is_empty() { + eprintln!("No work to do"); + return Ok(()); } + let signers = unique_signers(vec![&*args.fee_payer, &*args.custodian]); + send_messages(client, messages, &signers, args.no_wait)?; Ok(()) } @@ -135,15 +161,12 @@ fn process_rebase_stake_accounts( eprintln!("No accounts found"); return Ok(()); } - let signers = vec![ + let signers = unique_signers(vec![ &*args.fee_payer, &*args.new_base_keypair, &*args.stake_authority, - ]; - for message in messages { - let signature = send_message(client, message, &signers)?; - println!("{}", signature); - } + ]); + send_messages(client, messages, &signers, false)?; Ok(()) } @@ -170,16 +193,13 @@ fn process_move_stake_accounts( eprintln!("No accounts found"); return Ok(()); } - let signers = vec![ + let signers = unique_signers(vec![ &*args.fee_payer, &*args.new_base_keypair, &*args.stake_authority, &*authorize_args.withdraw_authority, - ]; - for message in messages { - let signature = send_message(client, message, &signers)?; - println!("{}", signature); - } + ]); + send_messages(client, messages, &signers, false)?; Ok(()) } @@ -187,10 +207,30 @@ fn send_message( client: &RpcClient, message: Message, signers: &S, + no_wait: bool, ) -> Result { let mut transaction = Transaction::new_unsigned(message); client.resign_transaction(&mut transaction, signers)?; - client.send_and_confirm_transaction_with_spinner(&mut transaction, signers) + if no_wait { + client.send_transaction(&transaction) + } else { + client.send_and_confirm_transaction_with_spinner(&mut transaction, signers) + } +} + +fn send_messages( + client: &RpcClient, + messages: Vec, + signers: &S, + no_wait: bool, +) -> Result, ClientError> { + let mut signatures = vec![]; + for message in messages { + let signature = send_message(client, message, signers, no_wait)?; + signatures.push(signature); + println!("{}", signature); + } + Ok(signatures) } fn main() -> Result<(), Box> { diff --git a/stake-accounts/src/stake_accounts.rs b/stake-accounts/src/stake_accounts.rs index f2fc7502b..bafafd81f 100644 --- a/stake-accounts/src/stake_accounts.rs +++ b/stake-accounts/src/stake_accounts.rs @@ -1,9 +1,14 @@ -use solana_sdk::{instruction::Instruction, message::Message, pubkey::Pubkey}; +use solana_sdk::{ + clock::SECONDS_PER_DAY, instruction::Instruction, message::Message, pubkey::Pubkey, +}; use solana_stake_program::{ stake_instruction::{self, LockupArgs}, stake_state::{Authorized, Lockup, StakeAuthorize}, }; +const DAYS_PER_YEAR: f64 = 365.25; +const SECONDS_PER_YEAR: i64 = (SECONDS_PER_DAY as f64 * DAYS_PER_YEAR) as i64; + pub(crate) fn derive_stake_account_address(base_pubkey: &Pubkey, i: usize) -> Pubkey { Pubkey::create_with_seed(base_pubkey, &i.to_string(), &solana_stake_program::id()).unwrap() } @@ -157,19 +162,63 @@ pub(crate) fn authorize_stake_accounts( .collect::>() } +fn extend_lockup(lockup: &LockupArgs, years: f64) -> LockupArgs { + let offset = (SECONDS_PER_YEAR as f64 * years) as i64; + let unix_timestamp = lockup.unix_timestamp.map(|x| x + offset); + let epoch = lockup.epoch.map(|_| todo!()); + LockupArgs { + unix_timestamp, + epoch, + custodian: lockup.custodian, + } +} + +fn apply_lockup_changes(lockup: &LockupArgs, existing_lockup: &Lockup) -> LockupArgs { + let custodian = match lockup.custodian { + Some(x) if x == existing_lockup.custodian => None, + x => x, + }; + let epoch = match lockup.epoch { + Some(x) if x == existing_lockup.epoch => None, + x => x, + }; + let unix_timestamp = match lockup.unix_timestamp { + Some(x) if x == existing_lockup.unix_timestamp => None, + x => x, + }; + LockupArgs { + custodian, + epoch, + unix_timestamp, + } +} + pub(crate) fn lockup_stake_accounts( fee_payer_pubkey: &Pubkey, - base_pubkey: &Pubkey, custodian_pubkey: &Pubkey, lockup: &LockupArgs, - num_accounts: usize, + existing_lockups: &[(Pubkey, Lockup)], + unlock_years: Option, ) -> Vec { - let stake_account_addresses = derive_stake_account_addresses(base_pubkey, num_accounts); - stake_account_addresses + let default_lockup = LockupArgs::default(); + existing_lockups .iter() - .map(|address| { + .enumerate() + .filter_map(|(index, (address, existing_lockup))| { + let lockup = if let Some(unlock_years) = unlock_years { + let unlocks = existing_lockups.len() - 1; + let years = (unlock_years / unlocks as f64) * index as f64; + extend_lockup(lockup, years) + } else { + *lockup + }; + let lockup = apply_lockup_changes(&lockup, existing_lockup); + if lockup == default_lockup { + return None; + } let instruction = stake_instruction::set_lockup(address, &lockup, custodian_pubkey); - Message::new_with_payer(&[instruction], Some(&fee_payer_pubkey)) + let message = Message::new_with_payer(&[instruction], Some(&fee_payer_pubkey)); + Some(message) }) .collect() } @@ -274,6 +323,21 @@ mod tests { .collect() } + fn get_lockups( + client: &C, + base_pubkey: &Pubkey, + num_accounts: usize, + ) -> Vec<(Pubkey, Lockup)> { + (0..num_accounts) + .into_iter() + .map(|i| { + let address = derive_stake_account_address(&base_pubkey, i); + let account = client.get_account(&address).unwrap().unwrap(); + (address, StakeState::lockup_from(&account).unwrap()) + }) + .collect() + } + #[test] fn test_new_derived_stake_account() { let (bank, funding_keypair, rent) = create_bank(10_000_000); @@ -396,15 +460,16 @@ mod tests { let signers = [&funding_keypair, &fee_payer_keypair, &base_keypair]; bank_client.send_message(&signers, message).unwrap(); + let lockups = get_lockups(&bank_client, &base_pubkey, 1); let messages = lockup_stake_accounts( &fee_payer_pubkey, - &base_pubkey, &custodian_pubkey, &LockupArgs { unix_timestamp: Some(1), ..LockupArgs::default() }, - 1, + &lockups, + None, ); let signers = [&fee_payer_keypair, &custodian_keypair]; @@ -416,6 +481,20 @@ mod tests { let lockup = StakeState::lockup_from(&account).unwrap(); assert_eq!(lockup.unix_timestamp, 1); assert_eq!(lockup.epoch, 0); + + // Assert no work left to do + let lockups = get_lockups(&bank_client, &base_pubkey, 1); + let messages = lockup_stake_accounts( + &fee_payer_pubkey, + &custodian_pubkey, + &LockupArgs { + unix_timestamp: Some(1), + ..LockupArgs::default() + }, + &lockups, + None, + ); + assert_eq!(messages.len(), 0); } #[test] @@ -557,4 +636,17 @@ mod tests { assert_eq!(authorized.staker, new_stake_authority_pubkey); assert_eq!(authorized.withdrawer, new_withdraw_authority_pubkey); } + + #[test] + fn test_extend_lockup() { + let lockup = LockupArgs { + unix_timestamp: Some(1), + ..LockupArgs::default() + }; + let expected_lockup = LockupArgs { + unix_timestamp: Some(1 + SECONDS_PER_YEAR), + ..LockupArgs::default() + }; + assert_eq!(extend_lockup(&lockup, 1.0), expected_lockup); + } }