Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactor nonce rollback to remove overly complex code #1428

Merged
merged 3 commits into from
Jun 6, 2024
Merged

Refactor nonce rollback to remove overly complex code #1428

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
24f02c1
Refactor nonce rollback to remove overly complex code
jstarry May 17, 2024
54a7104
fix tests
jstarry May 18, 2024
6785293
feedback
jstarry Jun 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
235 changes: 67 additions & 168 deletions accounts-db/src/accounts.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -719,11 +719,14 @@ impl Accounts {
TransactionExecutionResult::NotExecuted(_) => continue,
};

let maybe_nonce = match (execution_status, &loaded_transaction.nonce) {
(Ok(_), _) => None, // Success, don't do any additional nonce processing
(Err(_), Some(nonce)) => {
Some((nonce, true /* rollback */))
}
Comment on lines -722 to -726
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before, maybe_nonce would only be Some(..) if the tx execution status was failed and the transaction used a nonce. Also, rollback is always true in this case, it's never false. This means that prepare_if_nonce_account doesn't need to check rollback and doesn't need to check execution status if maybe_nonce is Some(..)

enum TransactionStatus<'a> {
Successful,
FailedWithNonce { nonce: &'a NonceFull },
}
Copy link
Member

@ryoqun ryoqun Jun 3, 2024 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TransactionStatus is a bit too general.

how about these?

enum StoredAccountProcessingStatus {
    Normal,
    FailedWithNonce(...)
}

// or

enum AccountCollectionMode {
    Normal,
    FailedWithNonce(...)
}

// or 

enum NonceProcessing {
    Skip,
    UpdateAfterFailure(...)
}


let tx_status = match (execution_status, &loaded_transaction.nonce) {
(Ok(_), _) => TransactionStatus::Successful, // Success, don't do any additional nonce processing
(Err(_), Some(nonce)) => TransactionStatus::FailedWithNonce { nonce },
(Err(_), None) => {
// Fees for failed transactions which don't use durable nonces are
// deducted in Bank::filter_program_errors_and_collect_fee
Expand Down Expand Up @@ -753,18 +756,24 @@ impl Accounts {
.filter(|(i, _)| is_storable_account(message, *i))
{
if message.is_writable(i) {
let is_fee_payer = i == 0;
let is_nonce_account = prepare_if_nonce_account(
address,
account,
execution_status,
is_fee_payer,
maybe_nonce,
durable_nonce,
lamports_per_signature,
);

if execution_status.is_ok() || is_nonce_account || is_fee_payer {
let should_store_account = match tx_status {
TransactionStatus::Successful => true,
TransactionStatus::FailedWithNonce { nonce } => {
let is_fee_payer = i == 0;
Copy link
Member

@ryoqun ryoqun Jun 3, 2024 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: since you're touching very sensible piece of code, i'd take the chance to further simplify the code by factoring out the address == nonce.address() into caller by defining let is_nonce_account = ... just below let is_fee_payer = ... here.

Then, make prepare_if_nonce_account return (), which in turn make the nested if is_fee_payer { ... } go way. Also, stop passing address.

Finally, after that, prepare_if_nonce_account()'s job clearer: it does its jobs only if is_nonce_account || is_fee_payer.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also, rename to post_process_failed_nonce?

let is_nonce_account = prepare_if_nonce_account(
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since prepare_if_nonce_account always exited early if maybe_nonce was None, we can refactor to only call it for failed nonce transactions where the nonce is available.

address,
account,
is_fee_payer,
nonce,
durable_nonce,
lamports_per_signature,
);

is_fee_payer || is_nonce_account
}
};

if should_store_account {
// Add to the accounts to store
accounts.push((&*address, &*account));
transactions.push(Some(tx));
Expand All @@ -779,52 +788,42 @@ impl Accounts {
fn prepare_if_nonce_account(
address: &Pubkey,
account: &mut AccountSharedData,
execution_result: &Result<()>,
is_fee_payer: bool,
maybe_nonce: Option<(&NonceFull, bool)>,
nonce: &NonceFull,
&durable_nonce: &DurableNonce,
lamports_per_signature: u64,
) -> bool {
if let Some((nonce, rollback)) = maybe_nonce {
if address == nonce.address() {
if rollback {
// The transaction failed which would normally drop the account
// processing changes, since this account is now being included
// in the accounts written back to the db, roll it back to
// pre-processing state.
*account = nonce.account().clone();
}

// Advance the stored blockhash to prevent fee theft by someone
// replaying nonce transactions that have failed with an
// `InstructionError`.
//
// Since we know we are dealing with a valid nonce account,
// unwrap is safe here
let nonce_versions = StateMut::<NonceVersions>::state(nonce.account()).unwrap();
if let NonceState::Initialized(ref data) = nonce_versions.state() {
let nonce_state = NonceState::new_initialized(
&data.authority,
durable_nonce,
lamports_per_signature,
);
let nonce_versions = NonceVersions::new(nonce_state);
account.set_state(&nonce_versions).unwrap();
}
true
} else {
if execution_result.is_err() && is_fee_payer {
if let Some(fee_payer_account) = nonce.fee_payer_account() {
// Instruction error and fee-payer for this nonce tx is not
// the nonce account itself, rollback the fee payer to the
// fee-paid original state.
*account = fee_payer_account.clone();
}
}

false
if address == nonce.address() {
// The transaction failed which would normally drop the account
// processing changes, since this account is now being included
// in the accounts written back to the db, roll it back to
// pre-processing state.
*account = nonce.account().clone();

// Advance the stored blockhash to prevent fee theft by someone
// replaying nonce transactions that have failed with an
// `InstructionError`.
//
// Since we know we are dealing with a valid nonce account,
// unwrap is safe here
let nonce_versions = StateMut::<NonceVersions>::state(nonce.account()).unwrap();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: nonce.account() => account for slight increase of readability?

if let NonceState::Initialized(ref data) = nonce_versions.state() {
let nonce_state =
NonceState::new_initialized(&data.authority, durable_nonce, lamports_per_signature);
let nonce_versions = NonceVersions::new(nonce_state);
account.set_state(&nonce_versions).unwrap();
}
true
} else {
if is_fee_payer {
if let Some(fee_payer_account) = nonce.fee_payer_account() {
// Instruction error and fee-payer for this nonce tx is not
// the nonce account itself, rollback the fee payer to the
// fee-paid original state.
*account = fee_payer_account.clone();
}
}

false
}
}
Expand Down Expand Up @@ -1685,29 +1684,26 @@ mod tests {
fn run_prepare_if_nonce_account_test(
account_address: &Pubkey,
account: &mut AccountSharedData,
tx_result: &Result<()>,
is_fee_payer: bool,
maybe_nonce: Option<(&NonceFull, bool)>,
nonce: &NonceFull,
durable_nonce: &DurableNonce,
lamports_per_signature: u64,
expect_account: &AccountSharedData,
) -> bool {
// Verify expect_account's relationship
if !is_fee_payer {
match maybe_nonce {
Some((nonce, _)) if nonce.address() == account_address => {
assert_ne!(expect_account, nonce.account())
}
_ => assert_eq!(expect_account, account),
if nonce.address() == account_address {
assert_ne!(expect_account, nonce.account());
} else {
assert_eq!(expect_account, account);
}
}

prepare_if_nonce_account(
account_address,
account,
tx_result,
is_fee_payer,
maybe_nonce,
nonce,
durable_nonce,
lamports_per_signature,
);
Expand Down Expand Up @@ -1742,35 +1738,8 @@ mod tests {
assert!(run_prepare_if_nonce_account_test(
&post_account_address,
&mut post_account,
&Ok(()),
false,
Some((&nonce, true)),
&blockhash,
lamports_per_signature,
&expect_account,
));
}

#[test]
fn test_prepare_if_nonce_account_not_nonce_tx() {
let (
pre_account_address,
_pre_account,
_post_account,
blockhash,
lamports_per_signature,
_maybe_fee_payer_account,
) = create_accounts_prepare_if_nonce_account();
let post_account_address = pre_account_address;

let mut post_account = AccountSharedData::default();
let expect_account = post_account.clone();
assert!(run_prepare_if_nonce_account_test(
&post_account_address,
&mut post_account,
&Ok(()),
false,
None,
&nonce,
&blockhash,
lamports_per_signature,
&expect_account,
Expand All @@ -1795,45 +1764,8 @@ mod tests {
assert!(run_prepare_if_nonce_account_test(
&Pubkey::from([1u8; 32]),
&mut post_account,
&Ok(()),
false,
Some((&nonce, true)),
&blockhash,
lamports_per_signature,
&expect_account,
));
}

#[test]
fn test_prepare_if_nonce_account_tx_error() {
let (
pre_account_address,
pre_account,
mut post_account,
blockhash,
lamports_per_signature,
maybe_fee_payer_account,
) = create_accounts_prepare_if_nonce_account();
let post_account_address = pre_account_address;
let mut expect_account = pre_account.clone();

let nonce = NonceFull::new(pre_account_address, pre_account, maybe_fee_payer_account);

expect_account
.set_state(&NonceVersions::new(NonceState::Initialized(
nonce::state::Data::new(Pubkey::default(), blockhash, lamports_per_signature),
)))
.unwrap();

assert!(run_prepare_if_nonce_account_test(
&post_account_address,
&mut post_account,
&Err(TransactionError::InstructionError(
0,
InstructionError::InvalidArgument,
)),
false,
Some((&nonce, true)),
&nonce,
&blockhash,
lamports_per_signature,
&expect_account,
Expand All @@ -1845,7 +1777,7 @@ mod tests {
let nonce_account = AccountSharedData::new_data(1, &(), &system_program::id()).unwrap();
let pre_fee_payer_account =
AccountSharedData::new_data(42, &(), &system_program::id()).unwrap();
let mut post_fee_payer_account =
let post_fee_payer_account =
AccountSharedData::new_data(84, &[1, 2, 3, 4], &system_program::id()).unwrap();
let nonce = NonceFull::new(
Pubkey::new_unique(),
Expand All @@ -1856,12 +1788,8 @@ mod tests {
assert!(run_prepare_if_nonce_account_test(
&Pubkey::new_unique(),
&mut post_fee_payer_account.clone(),
&Err(TransactionError::InstructionError(
0,
InstructionError::InvalidArgument,
)),
false,
Some((&nonce, true)),
&nonce,
&DurableNonce::default(),
1,
&post_fee_payer_account,
Expand All @@ -1870,37 +1798,8 @@ mod tests {
assert!(run_prepare_if_nonce_account_test(
&Pubkey::new_unique(),
&mut post_fee_payer_account.clone(),
&Ok(()),
true,
Some((&nonce, true)),
&DurableNonce::default(),
1,
&post_fee_payer_account,
));

assert!(run_prepare_if_nonce_account_test(
&Pubkey::new_unique(),
&mut post_fee_payer_account.clone(),
&Err(TransactionError::InstructionError(
0,
InstructionError::InvalidArgument,
)),
true,
None,
&DurableNonce::default(),
1,
&post_fee_payer_account,
));

assert!(run_prepare_if_nonce_account_test(
&Pubkey::new_unique(),
&mut post_fee_payer_account,
&Err(TransactionError::InstructionError(
0,
InstructionError::InvalidArgument,
)),
true,
Some((&nonce, true)),
&nonce,
&DurableNonce::default(),
1,
&pre_fee_payer_account,
Expand Down