blockworks-foundation
/
solana-with-rpc-optimizations
mirror of https://github.com/blockworks-foundation/solana-with-rpc-optimizations.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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:
Greg Fitzgerald 2020-05-04 13:05:04 -06:00 committed by GitHub
parent a8394317c7
commit f8ad3aca25
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 191 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@ -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",

1
stake-accounts/Cargo.toml
View File

@ -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" }

16
stake-accounts/src/arg_parser.rs
View File

@ -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(),
}
}

4
stake-accounts/src/args.rs
View File

@ -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)
}

98
stake-accounts/src/main.rs
View File

@ -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,10 +207,30 @@ 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)?;
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<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>> {

110
stake-accounts/src/stake_accounts.rs
View File

@ -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);
}
}