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

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

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

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

stake-accounts/src/args.rs

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

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<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 {
+    send_messages(client, messages, &signers, false)?;
-        let signature = send_message(client, message, &signers)?;
-        println!("{}", signature);
-    }
     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];
+    if messages.is_empty() {
-    for message in messages {
+        eprintln!("No work to do");
-        let signature = send_message(client, message, &signers)?;
+        return Ok(());
-        println!("{}", signature);
     }
+    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 {
+    send_messages(client, messages, &signers, false)?;
-        let signature = send_message(client, message, &signers)?;
-        println!("{}", signature);
-    }
     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 {
+    send_messages(client, messages, &signers, false)?;
-        let signature = send_message(client, message, &signers)?;
-        println!("{}", signature);
-    }
     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>> {

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::<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);
+    let default_lockup = LockupArgs::default();
-    stake_account_addresses
+    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);
+    }
 }