Skip to content

Commit

Permalink
Feature/multi-iterator-scanner-read-locks (#28862)
Browse files Browse the repository at this point in the history
  • Loading branch information
apfitzge authored Nov 28, 2022
1 parent 9327658 commit bdd1624
Show file tree
Hide file tree
Showing 5 changed files with 376 additions and 42 deletions.
2 changes: 1 addition & 1 deletion core/src/banking_stage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -651,7 +651,7 @@ impl BankingStage {

// Clear payload for next iteration
payload.sanitized_transactions.clear();
payload.write_accounts.clear();
payload.account_locks.clear();

let ProcessTransactionsSummary {
reached_max_poh_height,
Expand Down
1 change: 1 addition & 0 deletions core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ pub mod poh_timing_report_service;
pub mod poh_timing_reporter;
pub mod progress_map;
pub mod qos_service;
pub mod read_write_account_set;
pub mod repair_generic_traversal;
pub mod repair_response;
pub mod repair_service;
Expand Down
354 changes: 354 additions & 0 deletions core/src/read_write_account_set.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,354 @@
use {
solana_sdk::{
message::{SanitizedMessage, VersionedMessage},
pubkey::Pubkey,
},
std::collections::HashSet,
};

/// Wrapper struct to check account locks for a batch of transactions.
#[derive(Debug, Default)]
pub struct ReadWriteAccountSet {
/// Set of accounts that are locked for read
read_set: HashSet<Pubkey>,
/// Set of accounts that are locked for write
write_set: HashSet<Pubkey>,
}

impl ReadWriteAccountSet {
/// Check static account locks for a transaction message.
pub fn check_static_account_locks(&self, message: &VersionedMessage) -> bool {
!message
.static_account_keys()
.iter()
.enumerate()
.any(|(index, pubkey)| {
if message.is_maybe_writable(index) {
!self.can_write(pubkey)
} else {
!self.can_read(pubkey)
}
})
}

/// Check all account locks and if they are available, lock them.
/// Returns true if all account locks are available and false otherwise.
pub fn try_locking(&mut self, message: &SanitizedMessage) -> bool {
if self.check_sanitized_message_account_locks(message) {
self.add_sanitized_message_account_locks(message);
true
} else {
false
}
}

/// Clears the read and write sets
pub fn clear(&mut self) {
self.read_set.clear();
self.write_set.clear();
}

/// Check if a sanitized message's account locks are available.
fn check_sanitized_message_account_locks(&self, message: &SanitizedMessage) -> bool {
!message
.account_keys()
.iter()
.enumerate()
.any(|(index, pubkey)| {
if message.is_writable(index) {
!self.can_write(pubkey)
} else {
!self.can_read(pubkey)
}
})
}

/// Insert the read and write locks for a sanitized message.
fn add_sanitized_message_account_locks(&mut self, message: &SanitizedMessage) {
message
.account_keys()
.iter()
.enumerate()
.for_each(|(index, pubkey)| {
if message.is_writable(index) {
self.add_write(pubkey);
} else {
self.add_read(pubkey);
}
});
}

/// Check if an account can be read-locked
fn can_read(&self, pubkey: &Pubkey) -> bool {
!self.write_set.contains(pubkey)
}

/// Check if an account can be write-locked
fn can_write(&self, pubkey: &Pubkey) -> bool {
!self.write_set.contains(pubkey) && !self.read_set.contains(pubkey)
}

/// Add an account to the read-set.
/// Should only be called after `can_read()` returns true
fn add_read(&mut self, pubkey: &Pubkey) {
self.read_set.insert(*pubkey);
}

/// Add an account to the write-set.
/// Should only be called after `can_write()` returns true
fn add_write(&mut self, pubkey: &Pubkey) {
self.write_set.insert(*pubkey);
}
}

#[cfg(test)]
mod tests {
use {
super::ReadWriteAccountSet,
solana_address_lookup_table_program::state::{AddressLookupTable, LookupTableMeta},
solana_ledger::genesis_utils::GenesisConfigInfo,
solana_runtime::{bank::Bank, genesis_utils::create_genesis_config},
solana_sdk::{
account::AccountSharedData,
hash::Hash,
message::{
v0::{self, MessageAddressTableLookup},
MessageHeader, VersionedMessage,
},
pubkey::Pubkey,
signature::Keypair,
signer::Signer,
transaction::{MessageHash, SanitizedTransaction, VersionedTransaction},
},
std::{borrow::Cow, sync::Arc},
};

fn create_test_versioned_message(
write_keys: &[Pubkey],
read_keys: &[Pubkey],
address_table_lookups: Vec<MessageAddressTableLookup>,
) -> VersionedMessage {
VersionedMessage::V0(v0::Message {
header: MessageHeader {
num_required_signatures: write_keys.len() as u8,
num_readonly_signed_accounts: 0,
num_readonly_unsigned_accounts: read_keys.len() as u8,
},
recent_blockhash: Hash::default(),
account_keys: write_keys.iter().chain(read_keys.iter()).copied().collect(),
address_table_lookups,
instructions: vec![],
})
}

fn create_test_sanitized_transaction(
write_keypair: &Keypair,
read_keys: &[Pubkey],
address_table_lookups: Vec<MessageAddressTableLookup>,
bank: &Bank,
) -> SanitizedTransaction {
let message = create_test_versioned_message(
&[write_keypair.pubkey()],
read_keys,
address_table_lookups,
);
SanitizedTransaction::try_create(
VersionedTransaction::try_new(message, &[write_keypair]).unwrap(),
MessageHash::Compute,
Some(false),
bank,
true, // require_static_program_ids
)
.unwrap()
}

fn create_test_address_lookup_table(
bank: Arc<Bank>,
num_addresses: usize,
) -> (Arc<Bank>, Pubkey) {
let mut addresses = Vec::with_capacity(num_addresses);
addresses.resize_with(num_addresses, Pubkey::new_unique);
let address_lookup_table = AddressLookupTable {
meta: LookupTableMeta {
authority: None,
..LookupTableMeta::default()
},
addresses: Cow::Owned(addresses),
};

let address_table_key = Pubkey::new_unique();
let data = address_lookup_table.serialize_for_tests().unwrap();
let mut account =
AccountSharedData::new(1, data.len(), &solana_address_lookup_table_program::id());
account.set_data(data);
bank.store_account(&address_table_key, &account);

(
Arc::new(Bank::new_from_parent(
&bank,
&Pubkey::new_unique(),
bank.slot() + 1,
)),
address_table_key,
)
}

fn create_test_bank() -> Arc<Bank> {
let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);
Arc::new(Bank::new_no_wallclock_throttle_for_tests(&genesis_config))
}

// Helper function (could potentially use test_case in future).
// conflict_index = 0 means write lock conflict
// conflict_index = 1 means read lock conflict
fn test_check_static_account_locks(conflict_index: usize, add_write: bool, expectation: bool) {
let message =
create_test_versioned_message(&[Pubkey::new_unique()], &[Pubkey::new_unique()], vec![]);

let mut account_locks = ReadWriteAccountSet::default();
assert!(account_locks.check_static_account_locks(&message));

let conflict_key = message.static_account_keys().get(conflict_index).unwrap();
if add_write {
account_locks.add_write(conflict_key);
} else {
account_locks.add_read(conflict_key);
}
assert_eq!(
expectation,
account_locks.check_static_account_locks(&message)
);
}

#[test]
fn test_check_static_account_locks_write_write_conflict() {
test_check_static_account_locks(0, true, false);
}

#[test]
fn test_check_static_account_locks_read_write_conflict() {
test_check_static_account_locks(0, false, false);
}

#[test]
fn test_check_static_account_locks_write_read_conflict() {
test_check_static_account_locks(1, true, false);
}

#[test]
fn test_check_static_account_locks_read_read_non_conflict() {
test_check_static_account_locks(1, false, true);
}

// Helper function (could potentially use test_case in future).
// conflict_index = 0 means write lock conflict with static key
// conflict_index = 1 means read lock conflict with static key
// conflict_index = 2 means write lock conflict with address table key
// conflict_index = 3 means read lock conflict with address table key
fn test_check_sanitized_message_account_locks(
conflict_index: usize,
add_write: bool,
expectation: bool,
) {
let bank = create_test_bank();
let (bank, table_address) = create_test_address_lookup_table(bank, 2);
let tx = create_test_sanitized_transaction(
&Keypair::new(),
&[Pubkey::new_unique()],
vec![MessageAddressTableLookup {
account_key: table_address,
writable_indexes: vec![0],
readonly_indexes: vec![1],
}],
&bank,
);
let message = tx.message();

let mut account_locks = ReadWriteAccountSet::default();
assert!(account_locks.check_sanitized_message_account_locks(message));

let conflict_key = message.account_keys().get(conflict_index).unwrap();
if add_write {
account_locks.add_write(conflict_key);
} else {
account_locks.add_read(conflict_key);
}
assert_eq!(
expectation,
account_locks.check_sanitized_message_account_locks(message)
);
}

#[test]
fn test_check_sanitized_message_account_locks_write_write_conflict() {
test_check_sanitized_message_account_locks(0, true, false); // static key conflict
test_check_sanitized_message_account_locks(2, true, false); // lookup key conflict
}

#[test]
fn test_check_sanitized_message_account_locks_read_write_conflict() {
test_check_sanitized_message_account_locks(0, false, false); // static key conflict
test_check_sanitized_message_account_locks(2, false, false); // lookup key conflict
}

#[test]
fn test_check_sanitized_message_account_locks_write_read_conflict() {
test_check_sanitized_message_account_locks(1, true, false); // static key conflict
test_check_sanitized_message_account_locks(3, true, false); // lookup key conflict
}

#[test]
fn test_check_sanitized_message_account_locks_read_read_non_conflict() {
test_check_sanitized_message_account_locks(1, false, true); // static key conflict
test_check_sanitized_message_account_locks(3, false, true); // lookup key conflict
}

#[test]
pub fn test_write_write_conflict() {
let mut account_locks = ReadWriteAccountSet::default();
let account = Pubkey::new_unique();
assert!(account_locks.can_write(&account));
account_locks.add_write(&account);
assert!(!account_locks.can_write(&account));
}

#[test]
pub fn test_read_write_conflict() {
let mut account_locks = ReadWriteAccountSet::default();
let account = Pubkey::new_unique();
assert!(account_locks.can_read(&account));
account_locks.add_read(&account);
assert!(!account_locks.can_write(&account));
assert!(account_locks.can_read(&account));
}

#[test]
pub fn test_write_read_conflict() {
let mut account_locks = ReadWriteAccountSet::default();
let account = Pubkey::new_unique();
assert!(account_locks.can_write(&account));
account_locks.add_write(&account);
assert!(!account_locks.can_write(&account));
assert!(!account_locks.can_read(&account));
}

#[test]
pub fn test_read_read_non_conflict() {
let mut account_locks = ReadWriteAccountSet::default();
let account = Pubkey::new_unique();
assert!(account_locks.can_read(&account));
account_locks.add_read(&account);
assert!(account_locks.can_read(&account));
}

#[test]
pub fn test_write_write_different_keys() {
let mut account_locks = ReadWriteAccountSet::default();
let account1 = Pubkey::new_unique();
let account2 = Pubkey::new_unique();
assert!(account_locks.can_write(&account1));
account_locks.add_write(&account1);
assert!(account_locks.can_write(&account2));
assert!(account_locks.can_read(&account2));
}
}
Loading

0 comments on commit bdd1624

Please sign in to comment.