Skip to content

Commit

Permalink
wip. tpuclient for transfer_with_client and setup_vote_and_stake_acco…
Browse files Browse the repository at this point in the history 
…unts
  • Loading branch information
@gregcusack
gregcusack committed Apr 10, 2024
1 parent 51f9972 commit 86f190e
Showing 1 changed file with 72 additions and 43 deletions.
115 changes: 72 additions & 43 deletions local-cluster/src/local_cluster.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -449,12 +449,7 @@ impl LocalCluster {
mut voting_keypair: Option<Arc<Keypair>>,
socket_addr_space: SocketAddrSpace,
) -> Pubkey {
let (rpc, tpu) = LegacyContactInfo::try_from(&self.entry_point_info)
.map(|node| {
cluster_tests::get_client_facing_addr(self.connection_cache.protocol(), node)
})
.unwrap();
let client = ThinClient::new(rpc, tpu, self.connection_cache.clone());
let client = self.build_tpu_quic_client().expect("tpu_client");

// Must have enough tokens to fund vote account and set delegate
let should_create_vote_pubkey = voting_keypair.is_none();
Expand All @@ -477,10 +472,12 @@ impl LocalCluster {
&validator_pubkey,
stake * 2 + 2,
);

info!(
"validator {} balance {}",
validator_pubkey, validator_balance
);

Self::setup_vote_and_stake_accounts(
&client,
voting_keypair.as_ref().unwrap(),
Expand Down Expand Up @@ -649,36 +646,39 @@ impl LocalCluster {
}

fn transfer_with_client(
client: &ThinClient,
client: &QuicTpuClient,
source_keypair: &Keypair,
dest_pubkey: &Pubkey,
lamports: u64,
) -> u64 {
trace!("getting leader blockhash");

let (blockhash, _) = client
.rpc_client()
.get_latest_blockhash_with_commitment(CommitmentConfig::processed())
.unwrap();
let mut tx = system_transaction::transfer(source_keypair, dest_pubkey, lamports, blockhash);
let tx = system_transaction::transfer(source_keypair, dest_pubkey, lamports, blockhash);
info!(
"executing transfer of {} from {} to {}",
lamports,
source_keypair.pubkey(),
*dest_pubkey
);
client
.retry_transfer(source_keypair, &mut tx, 10)
.expect("client transfer");
.try_send_transaction(&tx)
.expect("client transfer should succeed");
client
.rpc_client()
.wait_for_balance_with_commitment(
dest_pubkey,
Some(lamports),
CommitmentConfig::processed(),
)
.expect("get balance")
.expect("get balance should succeed")
}

fn setup_vote_and_stake_accounts(
client: &ThinClient,
client: &QuicTpuClient,
vote_account: &Keypair,
from_account: &Arc<Keypair>,
amount: u64,
Expand All @@ -694,6 +694,7 @@ impl LocalCluster {

// Create the vote account if necessary
if client
.rpc_client()
.poll_get_balance_with_commitment(&vote_account_pubkey, CommitmentConfig::processed())
.unwrap_or(0)
== 0
Expand All @@ -703,6 +704,7 @@ impl LocalCluster {
// as the cluster is already running, and using the wrong account size will cause the
// InitializeAccount tx to fail
let use_current_vote_state = client
.rpc_client()
.poll_get_balance_with_commitment(
&feature_set::vote_state_add_vote_latency::id(),
CommitmentConfig::processed(),
Expand All @@ -727,18 +729,20 @@ impl LocalCluster {
},
);
let message = Message::new(&instructions, Some(&from_account.pubkey()));
let mut transaction = Transaction::new(
let transaction = Transaction::new(
&[from_account.as_ref(), vote_account],
message,
client
.rpc_client()
.get_latest_blockhash_with_commitment(CommitmentConfig::processed())
.unwrap()
.0,
);
client
.retry_transfer(from_account, &mut transaction, 10)
.expect("fund vote");
.try_send_transaction(&transaction)
.expect("should fund vote");
client
.rpc_client()
.wait_for_balance_with_commitment(
&vote_account_pubkey,
Some(amount),
Expand All @@ -755,24 +759,21 @@ impl LocalCluster {
amount,
);
let message = Message::new(&instructions, Some(&from_account.pubkey()));
let mut transaction = Transaction::new(
let transaction = Transaction::new(
&[from_account.as_ref(), &stake_account_keypair],
message,
client
.rpc_client()
.get_latest_blockhash_with_commitment(CommitmentConfig::processed())
.unwrap()
.0,
);

client
.send_and_confirm_transaction(
&[from_account.as_ref(), &stake_account_keypair],
&mut transaction,
5,
0,
)
.try_send_transaction(&transaction)
.expect("delegate stake");
client
.rpc_client()
.wait_for_balance_with_commitment(
&stake_account_pubkey,
Some(amount),
Expand All @@ -788,36 +789,64 @@ impl LocalCluster {
info!("Checking for vote account registration of {}", node_pubkey);
match (
client
.rpc_client()
.get_account_with_commitment(&stake_account_pubkey, CommitmentConfig::processed()),
client.get_account_with_commitment(&vote_account_pubkey, CommitmentConfig::processed()),
client
.rpc_client()
.get_account_with_commitment(&vote_account_pubkey, CommitmentConfig::processed()),
) {
(Ok(Some(stake_account)), Ok(Some(vote_account))) => {
match (
stake_state::stake_from(&stake_account),
vote_state::from(&vote_account),
) {
(Some(stake_state), Some(vote_state)) => {
if stake_state.delegation.voter_pubkey != vote_account_pubkey
|| stake_state.delegation.stake != amount
{
Err(Error::new(ErrorKind::Other, "invalid stake account state"))
} else if vote_state.node_pubkey != node_pubkey {
Err(Error::new(ErrorKind::Other, "invalid vote account state"))
} else {
info!("node {} {:?} {:?}", node_pubkey, stake_state, vote_state);

Ok(())
(Ok(stake_account), Ok(vote_account)) => {
match (stake_account.value, vote_account.value) {
(Some(stake_account), Some(vote_account)) => {
match (
stake_state::stake_from(&stake_account),
vote_state::from(&vote_account),
) {
(Some(stake_state), Some(vote_state)) => {
if stake_state.delegation.voter_pubkey != vote_account_pubkey
|| stake_state.delegation.stake != amount
{
return Err(Error::new(
ErrorKind::Other,
"invalid stake account state",
));
} else if vote_state.node_pubkey != node_pubkey {
return Err(Error::new(
ErrorKind::Other,
"invalid vote account state",
));
} else {
info!(
"node {} {:?} {:?}",
node_pubkey, stake_state, vote_state
);

return Ok(());
}
}
(None, _) => {
Err(Error::new(ErrorKind::Other, "invalid stake account data"))
}
(_, None) => {
Err(Error::new(ErrorKind::Other, "invalid vote account data"))
}
}
}
(None, _) => Err(Error::new(ErrorKind::Other, "invalid stake account data")),
(_, None) => Err(Error::new(ErrorKind::Other, "invalid vote account data")),
(None, _) => Err(Error::new(
ErrorKind::Other,
"unable to retrieve stake account data",
)),
(_, None) => Err(Error::new(
ErrorKind::Other,
"unable to retrieve vote account data",
)),
}
}
(Ok(None), _) | (Err(_), _) => Err(Error::new(
(Err(_), _) => Err(Error::new(
ErrorKind::Other,
"unable to retrieve stake account data",
)),
(_, Ok(None)) | (_, Err(_)) => Err(Error::new(
(_, Err(_)) => Err(Error::new(
ErrorKind::Other,
"unable to retrieve vote account data",
)),
Expand Down

0 comments on commit 86f190e

Please sign in to comment.