Speed up setting lockups (#9849)
* De-dup sending messages * Add --no-wait option for setting lockups * Don't set lockups that are already set * Extend adjacent lockups
This commit is contained in:
parent
a8394317c7
commit
f8ad3aca25
|
@ -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",
|
||||
|
|
|
@ -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" }
|
||||
|
|
|
@ -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<String, Stri
|
|||
lockup_date: unix_timestamp_from_rfc3339_datetime(matches, "lockup_date"),
|
||||
new_custodian: value_t!(matches, "new_custodian", String).ok(),
|
||||
num_accounts: value_t_or_exit!(matches, "num_accounts", usize),
|
||||
no_wait: matches.is_present("no_wait"),
|
||||
unlock_years: value_t!(matches, "unlock_years", f64).ok(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -46,6 +46,8 @@ pub(crate) struct SetLockupArgs<P, K> {
|
|||
pub lockup_date: Option<UnixTimestamp>,
|
||||
pub new_custodian: Option<P>,
|
||||
pub num_accounts: usize,
|
||||
pub no_wait: bool,
|
||||
pub unlock_years: Option<f64>,
|
||||
}
|
||||
|
||||
pub(crate) struct RebaseArgs<P, K> {
|
||||
|
@ -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)
|
||||
}
|
||||
|
|
|
@ -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<Lockup, ClientError> {
|
||||
client
|
||||
.get_account(address)
|
||||
.map(|account| StakeState::lockup_from(&account).unwrap())
|
||||
}
|
||||
|
||||
fn get_lockups(
|
||||
client: &RpcClient,
|
||||
addresses: Vec<Pubkey>,
|
||||
) -> Result<Vec<(Pubkey, Lockup)>, 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<Pubkey, Box<dyn Signer>>,
|
||||
|
@ -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<Pubkey, Box<dyn Signer>>,
|
||||
) -> 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,11 +207,31 @@ fn send_message<S: Signers>(
|
|||
client: &RpcClient,
|
||||
message: Message,
|
||||
signers: &S,
|
||||
no_wait: bool,
|
||||
) -> Result<Signature, ClientError> {
|
||||
let mut transaction = Transaction::new_unsigned(message);
|
||||
client.resign_transaction(&mut transaction, signers)?;
|
||||
if no_wait {
|
||||
client.send_transaction(&transaction)
|
||||
} else {
|
||||
client.send_and_confirm_transaction_with_spinner(&mut transaction, signers)
|
||||
}
|
||||
}
|
||||
|
||||
fn send_messages<S: Signers>(
|
||||
client: &RpcClient,
|
||||
messages: Vec<Message>,
|
||||
signers: &S,
|
||||
no_wait: bool,
|
||||
) -> Result<Vec<Signature>, 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<dyn Error>> {
|
||||
let command_args = parse_args(env::args_os());
|
||||
|
|
|
@ -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::<Vec<_>>()
|
||||
}
|
||||
|
||||
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<f64>,
|
||||
) -> Vec<Message> {
|
||||
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<C: SyncClient>(
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue