Skip to content

Commit

Permalink
Program Runtime: Unify transaction batch program caches (anza-xyz#1399)
Browse files Browse the repository at this point in the history 
* local program cache: add `modified_entries` field

* use `modified_entries` for modified program cache

* invoke context: make `program_cache_for_tx_batch` mutable

* invoke context: unify local program cache instances

* remove `find_program_in_cache` alias
  • Loading branch information
@buffalojoec
buffalojoec authored Jun 17, 2024
1 parent 7ff7fe9 commit 3af8f0c
Show file tree
Hide file tree
Showing 8 changed files with 99 additions and 119 deletions.
27 changes: 7 additions & 20 deletions program-runtime/src/invoke_context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,7 @@ pub struct InvokeContext<'a> {
/// Information about the currently executing transaction.
pub transaction_context: &'a mut TransactionContext,
/// The local program cache for the transaction batch.
pub program_cache_for_tx_batch: &'a ProgramCacheForTxBatch,
pub program_cache_for_tx_batch: &'a mut ProgramCacheForTxBatch,
/// Runtime configurations used to provision the invocation environment.
pub environment_config: EnvironmentConfig<'a>,
/// The compute budget for the current invocation.
Expand All @@ -202,7 +202,6 @@ pub struct InvokeContext<'a> {
/// the designated compute budget during program execution.
compute_meter: RefCell<u64>,
log_collector: Option<Rc<RefCell<LogCollector>>>,
pub programs_modified_by_tx: &'a mut ProgramCacheForTxBatch,
/// Latest measurement not yet accumulated in [ExecuteDetailsTimings::execute_us]
pub execute_time: Option<Measure>,
pub timings: ExecuteDetailsTimings,
Expand All @@ -214,11 +213,10 @@ impl<'a> InvokeContext<'a> {
#[allow(clippy::too_many_arguments)]
pub fn new(
transaction_context: &'a mut TransactionContext,
program_cache_for_tx_batch: &'a ProgramCacheForTxBatch,
program_cache_for_tx_batch: &'a mut ProgramCacheForTxBatch,
environment_config: EnvironmentConfig<'a>,
log_collector: Option<Rc<RefCell<LogCollector>>>,
compute_budget: ComputeBudget,
programs_modified_by_tx: &'a mut ProgramCacheForTxBatch,
) -> Self {
Self {
transaction_context,
Expand All @@ -227,22 +225,13 @@ impl<'a> InvokeContext<'a> {
log_collector,
compute_budget,
compute_meter: RefCell::new(compute_budget.compute_unit_limit),
programs_modified_by_tx,
execute_time: None,
timings: ExecuteDetailsTimings::default(),
syscall_context: Vec::new(),
traces: Vec::new(),
}
}

pub fn find_program_in_cache(&self, pubkey: &Pubkey) -> Option<Arc<ProgramCacheEntry>> {
// First lookup the cache of the programs modified by the current transaction. If not found, lookup
// the cache of the cache of the programs that are loaded for the transaction batch.
self.programs_modified_by_tx
.find(pubkey)
.or_else(|| self.program_cache_for_tx_batch.find(pubkey))
}

pub fn get_environments_for_slot(
&self,
effective_slot: Slot,
Expand Down Expand Up @@ -733,15 +722,13 @@ macro_rules! with_mock_invoke_context {
0,
&sysvar_cache,
);
let program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
let mut programs_modified_by_tx = ProgramCacheForTxBatch::default();
let mut program_cache_for_tx_batch = ProgramCacheForTxBatch::default();
let mut $invoke_context = InvokeContext::new(
&mut $transaction_context,
&program_cache_for_tx_batch,
&mut program_cache_for_tx_batch,
environment_config,
Some(LogCollector::new_ref()),
compute_budget,
&mut programs_modified_by_tx,
);
};
}
Expand Down Expand Up @@ -802,7 +789,7 @@ pub fn mock_process_instruction<F: FnMut(&mut InvokeContext), G: FnMut(&mut Invo
*loader_id,
Arc::new(ProgramCacheEntry::new_builtin(0, 0, builtin_function)),
);
invoke_context.program_cache_for_tx_batch = &program_cache_for_tx_batch;
invoke_context.program_cache_for_tx_batch = &mut program_cache_for_tx_batch;
pre_adjustments(&mut invoke_context);
let result = invoke_context.process_instruction(
instruction_data,
Expand Down Expand Up @@ -1059,7 +1046,7 @@ mod tests {
callee_program_id,
Arc::new(ProgramCacheEntry::new_builtin(0, 1, MockBuiltin::vm)),
);
invoke_context.program_cache_for_tx_batch = &program_cache_for_tx_batch;
invoke_context.program_cache_for_tx_batch = &mut program_cache_for_tx_batch;

// Account modification tests
let cases = vec![
Expand Down Expand Up @@ -1208,7 +1195,7 @@ mod tests {
program_key,
Arc::new(ProgramCacheEntry::new_builtin(0, 0, MockBuiltin::vm)),
);
invoke_context.program_cache_for_tx_batch = &program_cache_for_tx_batch;
invoke_context.program_cache_for_tx_batch = &mut program_cache_for_tx_batch;

// Test: Resize the account to *the same size*, so not consuming any additional size; this must succeed
{
Expand Down
58 changes: 36 additions & 22 deletions program-runtime/src/loaded_programs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -663,6 +663,8 @@ pub struct ProgramCacheForTxBatch {
/// Pubkey is the address of a program.
/// ProgramCacheEntry is the corresponding program entry valid for the slot in which a transaction is being executed.
entries: HashMap<Pubkey, Arc<ProgramCacheEntry>>,
/// Program entries modified during the transaction batch.
modified_entries: HashMap<Pubkey, Arc<ProgramCacheEntry>>,
slot: Slot,
pub environments: ProgramRuntimeEnvironments,
/// Anticipated replacement for `environments` at the next epoch.
Expand All @@ -689,6 +691,7 @@ impl ProgramCacheForTxBatch {
) -> Self {
Self {
entries: HashMap::new(),
modified_entries: HashMap::new(),
slot,
environments,
upcoming_environments,
Expand All @@ -706,6 +709,7 @@ impl ProgramCacheForTxBatch {
) -> Self {
Self {
entries: HashMap::new(),
modified_entries: HashMap::new(),
slot,
environments: cache.get_environments_for_epoch(epoch),
upcoming_environments: cache.get_upcoming_environments_for_epoch(epoch),
Expand All @@ -716,14 +720,6 @@ impl ProgramCacheForTxBatch {
}
}

pub fn entries(&self) -> &HashMap<Pubkey, Arc<ProgramCacheEntry>> {
&self.entries
}

pub fn take_entries(&mut self) -> HashMap<Pubkey, Arc<ProgramCacheEntry>> {
std::mem::take(&mut self.entries)
}

/// Returns the current environments depending on the given epoch
pub fn get_environments_for_epoch(&self, epoch: Epoch) -> &ProgramRuntimeEnvironments {
if epoch != self.latest_root_epoch {
Expand All @@ -747,21 +743,39 @@ impl ProgramCacheForTxBatch {
(self.entries.insert(key, entry.clone()).is_some(), entry)
}

/// Store an entry in `modified_entries` for a program modified during the
/// transaction batch.
pub fn store_modified_entry(&mut self, key: Pubkey, entry: Arc<ProgramCacheEntry>) {
self.modified_entries.insert(key, entry);
}

/// Drain the program cache's modified entries, returning the owned
/// collection.
pub fn drain_modified_entries(&mut self) -> HashMap<Pubkey, Arc<ProgramCacheEntry>> {
std::mem::take(&mut self.modified_entries)
}

pub fn find(&self, key: &Pubkey) -> Option<Arc<ProgramCacheEntry>> {
self.entries.get(key).map(|entry| {
if entry.is_implicit_delay_visibility_tombstone(self.slot) {
// Found a program entry on the current fork, but it's not effective
// yet. It indicates that the program has delayed visibility. Return
// the tombstone to reflect that.
Arc::new(ProgramCacheEntry::new_tombstone(
entry.deployment_slot,
entry.account_owner,
ProgramCacheEntryType::DelayVisibility,
))
} else {
entry.clone()
}
})
// First lookup the cache of the programs modified by the current
// transaction. If not found, lookup the cache of the cache of the
// programs that are loaded for the transaction batch.
self.modified_entries
.get(key)
.or(self.entries.get(key))
.map(|entry| {
if entry.is_implicit_delay_visibility_tombstone(self.slot) {
// Found a program entry on the current fork, but it's not effective
// yet. It indicates that the program has delayed visibility. Return
// the tombstone to reflect that.
Arc::new(ProgramCacheEntry::new_tombstone(
entry.deployment_slot,
entry.account_owner,
ProgramCacheEntryType::DelayVisibility,
))
} else {
entry.clone()
}
})
}

pub fn slot(&self) -> Slot {
Expand Down
39 changes: 21 additions & 18 deletions programs/bpf_loader/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ macro_rules! deploy_program {
environments.program_runtime_v1.clone(),
true,
)?;
if let Some(old_entry) = $invoke_context.find_program_in_cache(&$program_id) {
if let Some(old_entry) = $invoke_context.program_cache_for_tx_batch.find(&$program_id) {
executor.tx_usage_counter.store(
old_entry.tx_usage_counter.load(Ordering::Relaxed),
Ordering::Relaxed
Expand All @@ -165,7 +165,7 @@ macro_rules! deploy_program {
$drop
load_program_metrics.program_id = $program_id.to_string();
load_program_metrics.submit_datapoint(&mut $invoke_context.timings);
$invoke_context.programs_modified_by_tx.replenish($program_id, Arc::new(executor));
$invoke_context.program_cache_for_tx_batch.store_modified_entry($program_id, Arc::new(executor));
}};
}

Expand Down Expand Up @@ -437,7 +437,8 @@ pub fn process_instruction_inner(

let mut get_or_create_executor_time = Measure::start("get_or_create_executor_time");
let executor = invoke_context
.find_program_in_cache(program_account.get_key())
.program_cache_for_tx_batch
.find(program_account.get_key())
.ok_or_else(|| {
ic_logger_msg!(log_collector, "Program is not cached");
InstructionError::InvalidAccountData
Expand Down Expand Up @@ -1109,14 +1110,16 @@ fn process_loader_upgradeable_instruction(
&log_collector,
)?;
let clock = invoke_context.get_sysvar_cache().get_clock()?;
invoke_context.programs_modified_by_tx.replenish(
program_key,
Arc::new(ProgramCacheEntry::new_tombstone(
clock.slot,
ProgramCacheEntryOwner::LoaderV3,
ProgramCacheEntryType::Closed,
)),
);
invoke_context
.program_cache_for_tx_batch
.store_modified_entry(
program_key,
Arc::new(ProgramCacheEntry::new_tombstone(
clock.slot,
ProgramCacheEntryOwner::LoaderV3,
ProgramCacheEntryType::Closed,
)),
);
}
_ => {
ic_logger_msg!(log_collector, "Invalid Program account");
Expand Down Expand Up @@ -1543,11 +1546,11 @@ pub mod test_utils {
false,
) {
invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.set_slot_for_tests(DELAY_VISIBILITY_SLOT_OFFSET);
invoke_context
.programs_modified_by_tx
.replenish(*pubkey, Arc::new(loaded_program));
.program_cache_for_tx_batch
.store_modified_entry(*pubkey, Arc::new(loaded_program));
}
}
}
Expand Down Expand Up @@ -3763,7 +3766,7 @@ mod tests {
latest_access_slot: AtomicU64::new(0),
};
invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.replenish(program_id, Arc::new(program));

assert_matches!(
Expand All @@ -3772,7 +3775,7 @@ mod tests {
);

let updated_program = invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.find(&program_id)
.expect("Didn't find upgraded program in the cache");

Expand Down Expand Up @@ -3807,7 +3810,7 @@ mod tests {
latest_access_slot: AtomicU64::new(0),
};
invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.replenish(program_id, Arc::new(program));

let program_id2 = Pubkey::new_unique();
Expand All @@ -3817,7 +3820,7 @@ mod tests {
);

let program2 = invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.find(&program_id2)
.expect("Didn't find upgraded program in the cache");

Expand Down
24 changes: 15 additions & 9 deletions programs/loader-v4/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -449,7 +449,10 @@ pub fn process_instruction_deploy(
state.slot = current_slot;
state.status = LoaderV4Status::Deployed;

if let Some(old_entry) = invoke_context.find_program_in_cache(program.get_key()) {
if let Some(old_entry) = invoke_context
.program_cache_for_tx_batch
.find(program.get_key())
{
executor.tx_usage_counter.store(
old_entry.tx_usage_counter.load(Ordering::Relaxed),
Ordering::Relaxed,
Expand All @@ -460,8 +463,8 @@ pub fn process_instruction_deploy(
);
}
invoke_context
.programs_modified_by_tx
.replenish(*program.get_key(), Arc::new(executor));
.program_cache_for_tx_batch
.store_modified_entry(*program.get_key(), Arc::new(executor));
Ok(())
}

Expand Down Expand Up @@ -592,7 +595,8 @@ pub fn process_instruction_inner(
}
let mut get_or_create_executor_time = Measure::start("get_or_create_executor_time");
let loaded_program = invoke_context
.find_program_in_cache(program.get_key())
.program_cache_for_tx_batch
.find(program.get_key())
.ok_or_else(|| {
ic_logger_msg!(log_collector, "Program is not cached");
InstructionError::InvalidAccountData
Expand Down Expand Up @@ -661,7 +665,7 @@ mod tests {
if let Ok(loaded_program) = ProgramCacheEntry::new(
&loader_v4::id(),
invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.environments
.program_runtime_v2
.clone(),
Expand All @@ -671,10 +675,12 @@ mod tests {
account.data().len(),
&mut load_program_metrics,
) {
invoke_context.programs_modified_by_tx.set_slot_for_tests(0);
invoke_context
.programs_modified_by_tx
.replenish(*pubkey, Arc::new(loaded_program));
.program_cache_for_tx_batch
.set_slot_for_tests(0);
invoke_context
.program_cache_for_tx_batch
.store_modified_entry(*pubkey, Arc::new(loaded_program));
}
}
}
Expand Down Expand Up @@ -708,7 +714,7 @@ mod tests {
Entrypoint::vm,
|invoke_context| {
invoke_context
.programs_modified_by_tx
.program_cache_for_tx_batch
.environments
.program_runtime_v2 = Arc::new(create_program_runtime_environment_v2(
&ComputeBudget::default(),
Expand Down
Loading

0 comments on commit 3af8f0c

Please sign in to comment.