Skip to main content

rialo_s_program_runtime/
loaded_programs.rs

1// Copyright (c) Subzero Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3// This file is either (a) original to Subzero Labs, Inc. or (b) derived from the Anza codebase and modified by Subzero Labs, Inc.
4
5use std::{
6    collections::{hash_map::Entry, HashMap},
7    fmt::{Debug, Formatter},
8};
9
10use log::{debug, error, log_enabled, trace};
11use percentage::PercentageInteger;
12use polkavm::{ArcBytes, ProgramBlob, ProgramParts};
13use rialo_s_clock::{Epoch, Slot};
14use rialo_s_pubkey::Pubkey;
15use rialo_s_sdk_ids::{
16    bpf_loader, bpf_loader_deprecated, bpf_loader_upgradeable, loader_v4, native_loader,
17};
18use rialo_s_type_overrides::{
19    rand::{thread_rng, Rng},
20    sync::{
21        atomic::{AtomicU64, Ordering},
22        Arc, Condvar, Mutex,
23    },
24    thread,
25};
26use solana_sbpf::{
27    elf::Executable, program::BuiltinProgram, verifier::RequisiteVerifier, vm::Config,
28};
29#[cfg(feature = "metrics")]
30use {rialo_s_measure::measure::Measure, rialo_s_timings::ExecuteDetailsTimings};
31
32use crate::invoke_context::{BuiltinFunctionWithContext, InvokeContext};
33
34pub type ProgramRuntimeEnvironment = Arc<BuiltinProgram<InvokeContext<'static>>>;
35pub const MAX_LOADED_ENTRY_COUNT: usize = 512;
36pub const DELAY_VISIBILITY_SLOT_OFFSET: Slot = 1;
37
38/// Relationship between two fork IDs
39#[derive(Copy, Clone, Debug, PartialEq)]
40pub enum BlockRelation {
41    /// The slot is on the same fork and is an ancestor of the other slot
42    Ancestor,
43    /// The two slots are equal and are on the same fork
44    Equal,
45    /// The slot is on the same fork and is a descendant of the other slot
46    Descendant,
47    /// The slots are on two different forks and may have had a common ancestor at some point
48    Unrelated,
49    /// Either one or both of the slots are either older than the latest root, or are in future
50    Unknown,
51}
52
53/// Maps relationship between two slots.
54pub trait ForkGraph: Copy + Clone + Send + Sync {
55    /// Returns the BlockRelation of A to B
56    fn relationship(&self, a: Slot, b: Slot) -> BlockRelation;
57}
58
59/// The owner of a programs accounts, thus the loader of a program
60#[derive(Default, Clone, Copy, PartialEq, Eq, Debug)]
61pub enum ProgramCacheEntryOwner {
62    #[default]
63    NativeLoader,
64    LoaderV1,
65    LoaderV2,
66    LoaderV3,
67    LoaderV4,
68    RiscVLoader,
69}
70
71impl TryFrom<&Pubkey> for ProgramCacheEntryOwner {
72    type Error = ();
73    fn try_from(loader_key: &Pubkey) -> Result<Self, ()> {
74        if native_loader::check_id(loader_key) {
75            Ok(ProgramCacheEntryOwner::NativeLoader)
76        } else if bpf_loader_deprecated::check_id(loader_key) {
77            Ok(ProgramCacheEntryOwner::LoaderV1)
78        } else if bpf_loader::check_id(loader_key) {
79            Ok(ProgramCacheEntryOwner::LoaderV2)
80        } else if bpf_loader_upgradeable::check_id(loader_key) {
81            Ok(ProgramCacheEntryOwner::LoaderV3)
82        } else if loader_v4::check_id(loader_key) {
83            Ok(ProgramCacheEntryOwner::LoaderV4)
84        } else if rialo_risc_v_loader_interface::check_id(loader_key) {
85            Ok(ProgramCacheEntryOwner::RiscVLoader)
86        } else {
87            Err(())
88        }
89    }
90}
91
92impl From<ProgramCacheEntryOwner> for Pubkey {
93    fn from(program_cache_entry_owner: ProgramCacheEntryOwner) -> Self {
94        match program_cache_entry_owner {
95            ProgramCacheEntryOwner::NativeLoader => native_loader::id(),
96            ProgramCacheEntryOwner::LoaderV1 => bpf_loader_deprecated::id(),
97            ProgramCacheEntryOwner::LoaderV2 => bpf_loader::id(),
98            ProgramCacheEntryOwner::LoaderV3 => bpf_loader_upgradeable::id(),
99            ProgramCacheEntryOwner::LoaderV4 => loader_v4::id(),
100            ProgramCacheEntryOwner::RiscVLoader => rialo_risc_v_loader_interface::id(),
101        }
102    }
103}
104
105/*
106    The possible ProgramCacheEntryType transitions:
107
108    DelayVisibility is special in that it is never stored in the cache.
109    It is only returned by ProgramCacheForTxBatch::find() when a Loaded entry
110    is encountered which is not effective yet.
111
112    Builtin re/deployment:
113    - Empty => Builtin in TransactionBatchProcessor::add_builtin
114    - Builtin => Builtin in TransactionBatchProcessor::add_builtin
115
116    Un/re/deployment (with delay and cooldown):
117    - Empty / Closed => Loaded in UpgradeableLoaderInstruction::DeployWithMaxDataLen
118    - Loaded / FailedVerification => Loaded in UpgradeableLoaderInstruction::Upgrade
119    - Loaded / FailedVerification => Closed in UpgradeableLoaderInstruction::Close
120
121    Eviction and unloading (in the same slot):
122    - Unloaded => Loaded in ProgramCache::assign_program
123    - Loaded => Unloaded in ProgramCache::unload_program_entry
124
125    At epoch boundary (when feature set and environment changes):
126    - Loaded => FailedVerification in Bank::_new_from_parent
127    - FailedVerification => Loaded in Bank::_new_from_parent
128
129    Through pruning (when on orphan fork or overshadowed on the rooted fork):
130    - Closed / Unloaded / Loaded / Builtin => Empty in ProgramCache::prune
131*/
132
133/// Executable kind.
134pub enum ExecutableKind {
135    /// eBPF.
136    Ebpf(Executable<InvokeContext<'static>>),
137
138    /// RISC-V.
139    RiscV(Vec<u8>),
140}
141
142/// Actual payload of [ProgramCacheEntry].
143#[derive(Default)]
144pub enum ProgramCacheEntryType {
145    /// Tombstone for programs which currently do not pass the verifier but could if the feature set changed.
146    FailedVerification(ProgramRuntimeEnvironment),
147    /// Tombstone for programs that were either explicitly closed or never deployed.
148    ///
149    /// It's also used for accounts belonging to program loaders, that don't actually contain program code (e.g. buffer accounts for LoaderV3 programs).
150    #[default]
151    Closed,
152    /// Tombstone for programs which have recently been modified but the new version is not visible yet.
153    DelayVisibility,
154    /// Successfully verified but not currently compiled.
155    ///
156    /// It continues to track usage statistics even when the compiled executable of the program is evicted from memory.
157    Unloaded(ProgramRuntimeEnvironment),
158    /// Verified and compiled program
159    Loaded(ExecutableKind),
160    /// A built-in program which is not stored on-chain but backed into and distributed with the validator
161    Builtin(BuiltinProgram<InvokeContext<'static>>),
162}
163
164impl Debug for ProgramCacheEntryType {
165    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
166        match self {
167            ProgramCacheEntryType::FailedVerification(_) => {
168                write!(f, "ProgramCacheEntryType::FailedVerification")
169            }
170            ProgramCacheEntryType::Closed => write!(f, "ProgramCacheEntryType::Closed"),
171            ProgramCacheEntryType::DelayVisibility => {
172                write!(f, "ProgramCacheEntryType::DelayVisibility")
173            }
174            ProgramCacheEntryType::Unloaded(_) => write!(f, "ProgramCacheEntryType::Unloaded"),
175            ProgramCacheEntryType::Loaded(_) => write!(f, "ProgramCacheEntryType::Loaded"),
176            ProgramCacheEntryType::Builtin(_) => write!(f, "ProgramCacheEntryType::Builtin"),
177        }
178    }
179}
180
181impl ProgramCacheEntryType {
182    /// Returns a reference to its environment if it has one
183    pub fn get_environment(&self) -> Option<&ProgramRuntimeEnvironment> {
184        match self {
185            ProgramCacheEntryType::Loaded(ExecutableKind::Ebpf(program)) => {
186                Some(program.get_loader())
187            }
188            ProgramCacheEntryType::FailedVerification(env)
189            | ProgramCacheEntryType::Unloaded(env) => Some(env),
190            _ => None,
191        }
192    }
193}
194
195/// Holds a program version at a specific address and on a specific slot / fork.
196///
197/// It contains the actual program in [ProgramCacheEntryType] and a bunch of meta-data.
198#[derive(Debug, Default)]
199pub struct ProgramCacheEntry {
200    /// The program of this entry
201    pub program: ProgramCacheEntryType,
202    /// The loader of this entry
203    pub account_owner: ProgramCacheEntryOwner,
204    /// Size of account that stores the program and program data
205    pub account_size: usize,
206    /// Slot in which the program was (re)deployed
207    pub deployment_slot: Slot,
208    /// Slot in which this entry will become active (can be in the future)
209    pub effective_slot: Slot,
210    /// How often this entry was used by a transaction
211    pub tx_usage_counter: AtomicU64,
212    /// How often this entry was used by an instruction
213    pub ix_usage_counter: AtomicU64,
214    /// Latest slot in which the entry was used
215    pub latest_access_slot: AtomicU64,
216}
217
218/// Global cache statistics for [ProgramCache].
219#[derive(Debug, Default)]
220pub struct ProgramCacheStats {
221    /// a program was already in the cache
222    pub hits: AtomicU64,
223    /// a program was not found and loaded instead
224    pub misses: AtomicU64,
225    /// a compiled executable was unloaded
226    pub evictions: HashMap<Pubkey, u64>,
227    /// an unloaded program was loaded again (opposite of eviction)
228    pub reloads: AtomicU64,
229    /// a program was loaded or un/re/deployed
230    pub insertions: AtomicU64,
231    /// a program was loaded but can not be extracted on its own fork anymore
232    pub lost_insertions: AtomicU64,
233    /// a program which was already in the cache was reloaded by mistake
234    pub replacements: AtomicU64,
235    /// a program was only used once before being unloaded
236    pub one_hit_wonders: AtomicU64,
237    /// a program became unreachable in the fork graph because of rerooting
238    pub prunes_orphan: AtomicU64,
239    /// a program got pruned because it was not recompiled for the next epoch
240    pub prunes_environment: AtomicU64,
241    /// a program had no entries because all slot versions got pruned
242    pub empty_entries: AtomicU64,
243    /// water level of loaded entries currently cached
244    pub water_level: AtomicU64,
245}
246
247impl ProgramCacheStats {
248    pub fn reset(&mut self) {
249        *self = ProgramCacheStats::default();
250    }
251    pub fn log(&self) {
252        let hits = self.hits.load(Ordering::Relaxed);
253        let misses = self.misses.load(Ordering::Relaxed);
254        let evictions: u64 = self.evictions.values().sum();
255        let reloads = self.reloads.load(Ordering::Relaxed);
256        let insertions = self.insertions.load(Ordering::Relaxed);
257        let lost_insertions = self.lost_insertions.load(Ordering::Relaxed);
258        let replacements = self.replacements.load(Ordering::Relaxed);
259        let one_hit_wonders = self.one_hit_wonders.load(Ordering::Relaxed);
260        let prunes_orphan = self.prunes_orphan.load(Ordering::Relaxed);
261        let prunes_environment = self.prunes_environment.load(Ordering::Relaxed);
262        let empty_entries = self.empty_entries.load(Ordering::Relaxed);
263        let water_level = self.water_level.load(Ordering::Relaxed);
264        debug!(
265            "Loaded Programs Cache Stats -- Hits: {}, Misses: {}, Evictions: {}, Reloads: {}, Insertions: {}, Lost-Insertions: {}, Replacements: {}, One-Hit-Wonders: {}, Prunes-Orphan: {}, Prunes-Environment: {}, Empty: {}, Water-Level: {}",
266            hits, misses, evictions, reloads, insertions, lost_insertions, replacements, one_hit_wonders, prunes_orphan, prunes_environment, empty_entries, water_level
267        );
268        if log_enabled!(log::Level::Trace) && !self.evictions.is_empty() {
269            let mut evictions = self.evictions.iter().collect::<Vec<_>>();
270            evictions.sort_by_key(|e| e.1);
271            let evictions = evictions
272                .into_iter()
273                .rev()
274                .map(|(program_id, evictions)| {
275                    format!("  {:<44}  {}", program_id.to_string(), evictions)
276                })
277                .collect::<Vec<_>>();
278            let evictions = evictions.join("\n");
279            trace!(
280                "Eviction Details:\n  {:<44}  {}\n{}",
281                "Program",
282                "Count",
283                evictions
284            );
285        }
286    }
287}
288
289#[cfg(feature = "metrics")]
290/// Time measurements for loading a single [ProgramCacheEntry].
291#[derive(Debug, Default)]
292pub struct LoadProgramMetrics {
293    /// Program address, but as text
294    pub program_id: String,
295    /// Microseconds it took to `create_program_runtime_environment`
296    pub register_syscalls_us: u64,
297    /// Microseconds it took to `Executable::<InvokeContext>::load`
298    pub load_elf_us: u64,
299    /// Microseconds it took to `executable.verify::<RequisiteVerifier>`
300    pub verify_code_us: u64,
301    /// Microseconds it took to `executable.jit_compile`
302    pub jit_compile_us: u64,
303}
304
305#[cfg(feature = "metrics")]
306impl LoadProgramMetrics {
307    pub fn submit_datapoint(&self, timings: &mut ExecuteDetailsTimings) {
308        timings.create_executor_register_syscalls_us += self.register_syscalls_us;
309        timings.create_executor_load_elf_us += self.load_elf_us;
310        timings.create_executor_verify_code_us += self.verify_code_us;
311        timings.create_executor_jit_compile_us += self.jit_compile_us;
312        datapoint_trace!(
313            "create_executor_trace",
314            ("program_id", self.program_id, String),
315            ("register_syscalls_us", self.register_syscalls_us, i64),
316            ("load_elf_us", self.load_elf_us, i64),
317            ("verify_code_us", self.verify_code_us, i64),
318            ("jit_compile_us", self.jit_compile_us, i64),
319        );
320    }
321}
322
323impl PartialEq for ProgramCacheEntry {
324    fn eq(&self, other: &Self) -> bool {
325        self.effective_slot == other.effective_slot
326            && self.deployment_slot == other.deployment_slot
327            && self.is_tombstone() == other.is_tombstone()
328    }
329}
330
331impl ProgramCacheEntry {
332    /// Creates an eBPF new user program
333    pub fn new(
334        loader_key: &Pubkey,
335        program_runtime_environment: ProgramRuntimeEnvironment,
336        deployment_slot: Slot,
337        effective_slot: Slot,
338        elf_bytes: &[u8],
339        account_size: usize,
340        #[cfg(feature = "metrics")] metrics: &mut LoadProgramMetrics,
341    ) -> Result<Self, Box<dyn std::error::Error>> {
342        Self::new_internal(
343            loader_key,
344            program_runtime_environment,
345            deployment_slot,
346            effective_slot,
347            elf_bytes,
348            account_size,
349            #[cfg(feature = "metrics")]
350            metrics,
351            false, /* reloading */
352        )
353    }
354
355    /// Creates a RISC-V new user program
356    pub fn new_risc_v(
357        loader_key: &Pubkey,
358        deployment_slot: Slot,
359        effective_slot: Slot,
360        elf_bytes: &[u8],
361        account_size: usize,
362    ) -> Result<Self, Box<dyn std::error::Error>> {
363        Self::new_internal_risc_v(
364            loader_key,
365            deployment_slot,
366            effective_slot,
367            elf_bytes,
368            account_size,
369            false, /* reloading */
370        )
371    }
372
373    /// Reloads an eBPF user program, *without* running the verifier.
374    ///
375    /// # Safety
376    ///
377    /// This method is unsafe since it assumes that the program has already been verified. Should
378    /// only be called when the program was previously verified and loaded in the cache, but was
379    /// unloaded due to inactivity. It should also be checked that the `program_runtime_environment`
380    /// hasn't changed since it was unloaded.
381    pub unsafe fn reload(
382        loader_key: &Pubkey,
383        program_runtime_environment: Arc<BuiltinProgram<InvokeContext<'static>>>,
384        deployment_slot: Slot,
385        effective_slot: Slot,
386        elf_bytes: &[u8],
387        account_size: usize,
388        #[cfg(feature = "metrics")] metrics: &mut LoadProgramMetrics,
389    ) -> Result<Self, Box<dyn std::error::Error>> {
390        Self::new_internal(
391            loader_key,
392            program_runtime_environment,
393            deployment_slot,
394            effective_slot,
395            elf_bytes,
396            account_size,
397            #[cfg(feature = "metrics")]
398            metrics,
399            true, /* reloading */
400        )
401    }
402
403    /// Reloads a RISC-V user program, *without* running the verifier.
404    ///
405    /// # Safety
406    ///
407    /// This method is unsafe since it assumes that the program has already been verified. Should
408    /// only be called when the program was previously verified and loaded in the cache, but was
409    /// unloaded due to inactivity. It should also be checked that the `program_runtime_environment`
410    /// hasn't changed since it was unloaded.
411    pub unsafe fn reload_risc_v(
412        loader_key: &Pubkey,
413        deployment_slot: Slot,
414        effective_slot: Slot,
415        elf_bytes: &[u8],
416        account_size: usize,
417    ) -> Result<Self, Box<dyn std::error::Error>> {
418        Self::new_internal_risc_v(
419            loader_key,
420            deployment_slot,
421            effective_slot,
422            elf_bytes,
423            account_size,
424            true, /* reloading */
425        )
426    }
427
428    fn new_internal(
429        loader_key: &Pubkey,
430        program_runtime_environment: Arc<BuiltinProgram<InvokeContext<'static>>>,
431        deployment_slot: Slot,
432        effective_slot: Slot,
433        elf_bytes: &[u8],
434        account_size: usize,
435        #[cfg(feature = "metrics")] metrics: &mut LoadProgramMetrics,
436        reloading: bool,
437    ) -> Result<Self, Box<dyn std::error::Error>> {
438        #[cfg(feature = "metrics")]
439        let load_elf_time = Measure::start("load_elf_time");
440        // The following unused_mut exception is needed for architectures that do not
441        // support JIT compilation.
442        #[allow(unused_mut)]
443        let mut executable = Executable::load(elf_bytes, program_runtime_environment.clone())?;
444        #[cfg(feature = "metrics")]
445        {
446            metrics.load_elf_us = load_elf_time.end_as_us();
447        }
448
449        if !reloading {
450            #[cfg(feature = "metrics")]
451            let verify_code_time = Measure::start("verify_code_time");
452            executable.verify::<RequisiteVerifier>()?;
453            #[cfg(feature = "metrics")]
454            {
455                metrics.verify_code_us = verify_code_time.end_as_us();
456            }
457        }
458
459        #[cfg(all(not(target_os = "windows"), target_arch = "x86_64"))]
460        {
461            #[cfg(feature = "metrics")]
462            let jit_compile_time = Measure::start("jit_compile_time");
463            executable.jit_compile()?;
464            #[cfg(feature = "metrics")]
465            {
466                metrics.jit_compile_us = jit_compile_time.end_as_us();
467            }
468        }
469
470        Ok(Self {
471            deployment_slot,
472            account_owner: ProgramCacheEntryOwner::try_from(loader_key).unwrap(),
473            account_size,
474            effective_slot,
475            tx_usage_counter: AtomicU64::new(0),
476            program: ProgramCacheEntryType::Loaded(ExecutableKind::Ebpf(executable)),
477            ix_usage_counter: AtomicU64::new(0),
478            latest_access_slot: AtomicU64::new(0),
479        })
480    }
481
482    fn new_internal_risc_v(
483        loader_key: &Pubkey,
484        deployment_slot: Slot,
485        effective_slot: Slot,
486        elf_bytes: &[u8],
487        account_size: usize,
488        reloading: bool,
489    ) -> Result<Self, Box<dyn std::error::Error>> {
490        // verify that `elf_bytes` is a valid polkavm program
491        //
492        // the program is considered valid if it:
493        //  * parses into a valid polkavm program blob
494        //  * is a 64-bit program
495        //  * has stack size <= 0x40000
496        //  * exports an entrypoint
497        if !reloading {
498            const MAX_STACK_SIZE: u32 = 262144;
499
500            let program = Into::<ArcBytes>::into(elf_bytes);
501
502            let program_parts = ProgramParts::from_bytes(program.clone())
503                .map_err(|error| format!("failed to parse program: {error:?}"))?;
504
505            if !program_parts.is_64_bit {
506                return Err(String::from("program is not a 64-bit program").into());
507            }
508
509            if program_parts.stack_size > MAX_STACK_SIZE {
510                return Err(String::from("program stack is too large").into());
511            }
512
513            let blob = ProgramBlob::parse(program)
514                .map_err(|error| format!("failed to parse program into program blob: {error:?}"))?;
515
516            if !blob
517                .exports()
518                .any(|entry| entry.symbol().as_bytes() == b"entrypoint")
519            {
520                return Err(String::from("program doesn't have an entrypoint").into());
521            }
522        }
523
524        Ok(Self {
525            deployment_slot,
526            account_owner: ProgramCacheEntryOwner::try_from(loader_key).unwrap(),
527            account_size,
528            effective_slot,
529            tx_usage_counter: AtomicU64::new(0),
530            program: ProgramCacheEntryType::Loaded(ExecutableKind::RiscV(elf_bytes.to_vec())),
531            ix_usage_counter: AtomicU64::new(0),
532            latest_access_slot: AtomicU64::new(0),
533        })
534    }
535
536    pub fn to_unloaded(&self) -> Option<Self> {
537        match &self.program {
538            ProgramCacheEntryType::Loaded(_) => {}
539            ProgramCacheEntryType::FailedVerification(_)
540            | ProgramCacheEntryType::Closed
541            | ProgramCacheEntryType::DelayVisibility
542            | ProgramCacheEntryType::Unloaded(_)
543            | ProgramCacheEntryType::Builtin(_) => {
544                return None;
545            }
546        }
547        Some(Self {
548            program: ProgramCacheEntryType::Unloaded(self.program.get_environment()?.clone()),
549            account_owner: self.account_owner,
550            account_size: self.account_size,
551            deployment_slot: self.deployment_slot,
552            effective_slot: self.effective_slot,
553            tx_usage_counter: AtomicU64::new(self.tx_usage_counter.load(Ordering::Relaxed)),
554            ix_usage_counter: AtomicU64::new(self.ix_usage_counter.load(Ordering::Relaxed)),
555            latest_access_slot: AtomicU64::new(self.latest_access_slot.load(Ordering::Relaxed)),
556        })
557    }
558
559    /// Creates a new built-in program
560    pub fn new_builtin(
561        deployment_slot: Slot,
562        account_size: usize,
563        builtin_function: BuiltinFunctionWithContext,
564    ) -> Self {
565        let mut program = BuiltinProgram::new_builtin();
566        program
567            .register_function("entrypoint", builtin_function)
568            .unwrap();
569        Self {
570            deployment_slot,
571            account_owner: ProgramCacheEntryOwner::NativeLoader,
572            account_size,
573            effective_slot: deployment_slot,
574            tx_usage_counter: AtomicU64::new(0),
575            program: ProgramCacheEntryType::Builtin(program),
576            ix_usage_counter: AtomicU64::new(0),
577            latest_access_slot: AtomicU64::new(0),
578        }
579    }
580
581    pub fn new_tombstone(
582        slot: Slot,
583        account_owner: ProgramCacheEntryOwner,
584        reason: ProgramCacheEntryType,
585    ) -> Self {
586        let tombstone = Self {
587            program: reason,
588            account_owner,
589            account_size: 0,
590            deployment_slot: slot,
591            effective_slot: slot,
592            tx_usage_counter: AtomicU64::default(),
593            ix_usage_counter: AtomicU64::default(),
594            latest_access_slot: AtomicU64::new(0),
595        };
596        debug_assert!(tombstone.is_tombstone());
597        tombstone
598    }
599
600    pub fn is_tombstone(&self) -> bool {
601        matches!(
602            self.program,
603            ProgramCacheEntryType::FailedVerification(_)
604                | ProgramCacheEntryType::Closed
605                | ProgramCacheEntryType::DelayVisibility
606        )
607    }
608
609    fn is_implicit_delay_visibility_tombstone(&self, slot: Slot) -> bool {
610        !matches!(self.program, ProgramCacheEntryType::Builtin(_))
611            && self.effective_slot.saturating_sub(self.deployment_slot)
612                == DELAY_VISIBILITY_SLOT_OFFSET
613            && slot >= self.deployment_slot
614            && slot < self.effective_slot
615    }
616
617    pub fn update_access_slot(&self, slot: Slot) {
618        let _ = self.latest_access_slot.fetch_max(slot, Ordering::Relaxed);
619    }
620
621    pub fn decayed_usage_counter(&self, now: Slot) -> u64 {
622        let last_access = self.latest_access_slot.load(Ordering::Relaxed);
623        // Shifting the u64 value for more than 63 will cause an overflow.
624        let decaying_for = std::cmp::min(63, now.saturating_sub(last_access));
625        self.tx_usage_counter.load(Ordering::Relaxed) >> decaying_for
626    }
627
628    pub fn account_owner(&self) -> Pubkey {
629        self.account_owner.into()
630    }
631}
632
633/// Globally shared RBPF config and syscall registry
634///
635/// This is only valid in an epoch range as long as no feature affecting RBPF is activated.
636#[derive(Clone, Debug)]
637pub struct ProgramRuntimeEnvironments {
638    /// For program runtime V1
639    pub program_runtime_v1: ProgramRuntimeEnvironment,
640    /// For program runtime V2
641    pub program_runtime_v2: ProgramRuntimeEnvironment,
642}
643
644impl Default for ProgramRuntimeEnvironments {
645    fn default() -> Self {
646        let empty_loader = Arc::new(BuiltinProgram::new_loader(Config::default()));
647        Self {
648            program_runtime_v1: empty_loader.clone(),
649            program_runtime_v2: empty_loader,
650        }
651    }
652}
653
654#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
655pub struct LoadingTaskCookie(u64);
656
657impl LoadingTaskCookie {
658    fn new() -> Self {
659        Self(0)
660    }
661
662    fn update(&mut self) {
663        let LoadingTaskCookie(cookie) = self;
664        *cookie = cookie.wrapping_add(1);
665    }
666}
667
668/// Suspends the thread in case no cooprative loading task was assigned
669#[derive(Debug, Default)]
670pub struct LoadingTaskWaiter {
671    cookie: Mutex<LoadingTaskCookie>,
672    cond: Condvar,
673}
674
675impl LoadingTaskWaiter {
676    pub fn new() -> Self {
677        Self {
678            cookie: Mutex::new(LoadingTaskCookie::new()),
679            cond: Condvar::new(),
680        }
681    }
682
683    pub fn cookie(&self) -> LoadingTaskCookie {
684        *self.cookie.lock().unwrap()
685    }
686
687    pub fn notify(&self) {
688        let mut cookie = self.cookie.lock().unwrap();
689        cookie.update();
690        self.cond.notify_all();
691    }
692
693    pub fn wait(&self, cookie: LoadingTaskCookie) -> LoadingTaskCookie {
694        let cookie_guard = self.cookie.lock().unwrap();
695        *self
696            .cond
697            .wait_while(cookie_guard, |current_cookie| *current_cookie == cookie)
698            .unwrap()
699    }
700}
701
702#[derive(Debug)]
703enum IndexImplementation {
704    /// Fork-graph aware index implementation
705    V1 {
706        /// A two level index:
707        ///
708        /// - the first level is for the address at which programs are deployed
709        /// - the second level for the slot (and thus also fork), sorted by slot number.
710        entries: HashMap<Pubkey, Vec<Arc<ProgramCacheEntry>>>,
711        /// The entries that are getting loaded and have not yet finished loading.
712        ///
713        /// The key is the program address, the value is a tuple of the slot in which the program is
714        /// being loaded and the thread ID doing the load.
715        ///
716        /// It is possible that multiple TX batches from different slots need different versions of a
717        /// program. The deployment slot of a program is only known after load tho,
718        /// so all loads for a given program key are serialized.
719        loading_entries: Mutex<HashMap<Pubkey, (Slot, thread::ThreadId)>>,
720    },
721}
722
723/// This structure is the global cache of loaded, verified and compiled programs.
724///
725/// It ...
726/// - is validator global and fork graph aware, so it can optimize the commonalities across banks.
727/// - handles the visibility rules of un/re/deployments.
728/// - stores the usage statistics and verification status of each program.
729/// - is elastic and uses a probabilistic eviction stragety based on the usage statistics.
730/// - also keeps the compiled executables around, but only for the most used programs.
731/// - supports various kinds of tombstones to avoid loading programs which can not be loaded.
732/// - cleans up entries on orphan branches when the block store is rerooted.
733/// - supports the cache preparation phase before feature activations which can change cached programs.
734/// - manages the environments of the programs and upcoming environments for the next epoch.
735/// - allows for cooperative loading of TX batches which hit the same missing programs simultaneously.
736/// - enforces that all programs used in a batch are eagerly loaded ahead of execution.
737/// - is not persisted to disk or a snapshot, so it needs to cold start and warm up first.
738pub struct ProgramCache {
739    /// Index of the cached entries and cooperative loading tasks
740    index: IndexImplementation,
741    /// The slot of the last rerooting
742    pub latest_root_slot: Slot,
743    /// The epoch of the last rerooting
744    pub latest_root_epoch: Epoch,
745    /// Environments of the current epoch
746    pub environments: ProgramRuntimeEnvironments,
747    /// Anticipated replacement for `environments` at the next epoch
748    ///
749    /// This is `None` during most of an epoch, and only `Some` around the boundaries (at the end and beginning of an epoch).
750    /// More precisely, it starts with the cache preparation phase a few hundred slots before the epoch boundary,
751    /// and it ends with the first rerooting after the epoch boundary.
752    pub upcoming_environments: Option<ProgramRuntimeEnvironments>,
753    /// List of loaded programs which should be recompiled before the next epoch (but don't have to).
754    pub programs_to_recompile: Vec<(Pubkey, Arc<ProgramCacheEntry>)>,
755    /// Statistics counters
756    pub stats: ProgramCacheStats,
757    /// Coordinates TX batches waiting for others to complete their task during cooperative loading
758    pub loading_task_waiter: Arc<LoadingTaskWaiter>,
759}
760
761impl Debug for ProgramCache {
762    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
763        f.debug_struct("ProgramCache")
764            .field("root slot", &self.latest_root_slot)
765            .field("root epoch", &self.latest_root_epoch)
766            .field("stats", &self.stats)
767            .field("index", &self.index)
768            .finish()
769    }
770}
771
772/// Local view into [ProgramCache] which was extracted for a specific TX batch.
773///
774/// This isolation enables the global [ProgramCache] to continue to evolve (e.g. evictions),
775/// while the TX batch is guaranteed it will continue to find all the programs it requires.
776/// For program management instructions this also buffers them before they are merged back into the global [ProgramCache].
777#[derive(Clone, Debug, Default)]
778pub struct ProgramCacheForTxBatch {
779    /// Pubkey is the address of a program.
780    /// ProgramCacheEntry is the corresponding program entry valid for the slot in which a transaction is being executed.
781    entries: HashMap<Pubkey, Arc<ProgramCacheEntry>>,
782    /// Program entries modified during the transaction batch.
783    modified_entries: HashMap<Pubkey, Arc<ProgramCacheEntry>>,
784    slot: Slot,
785    pub environments: ProgramRuntimeEnvironments,
786    /// Anticipated replacement for `environments` at the next epoch.
787    ///
788    /// This is `None` during most of an epoch, and only `Some` around the boundaries (at the end and beginning of an epoch).
789    /// More precisely, it starts with the cache preparation phase a few hundred slots before the epoch boundary,
790    /// and it ends with the first rerooting after the epoch boundary.
791    /// Needed when a program is deployed at the last slot of an epoch, becomes effective in the next epoch.
792    /// So needs to be compiled with the environment for the next epoch.
793    pub upcoming_environments: Option<ProgramRuntimeEnvironments>,
794    /// The epoch of the last rerooting
795    pub latest_root_epoch: Epoch,
796    pub hit_max_limit: bool,
797    pub loaded_missing: bool,
798    pub merged_modified: bool,
799}
800
801impl ProgramCacheForTxBatch {
802    pub fn new(
803        slot: Slot,
804        environments: ProgramRuntimeEnvironments,
805        upcoming_environments: Option<ProgramRuntimeEnvironments>,
806        latest_root_epoch: Epoch,
807    ) -> Self {
808        Self {
809            entries: HashMap::new(),
810            modified_entries: HashMap::new(),
811            slot,
812            environments,
813            upcoming_environments,
814            latest_root_epoch,
815            hit_max_limit: false,
816            loaded_missing: false,
817            merged_modified: false,
818        }
819    }
820
821    pub fn new_from_cache(slot: Slot, epoch: Epoch, cache: &ProgramCache) -> Self {
822        Self {
823            entries: HashMap::new(),
824            modified_entries: HashMap::new(),
825            slot,
826            environments: cache.get_environments_for_epoch(epoch),
827            upcoming_environments: cache.get_upcoming_environments_for_epoch(epoch),
828            latest_root_epoch: cache.latest_root_epoch,
829            hit_max_limit: false,
830            loaded_missing: false,
831            merged_modified: false,
832        }
833    }
834
835    /// Returns the current environments depending on the given epoch
836    pub fn get_environments_for_epoch(&self, epoch: Epoch) -> &ProgramRuntimeEnvironments {
837        if epoch != self.latest_root_epoch {
838            if let Some(upcoming_environments) = self.upcoming_environments.as_ref() {
839                return upcoming_environments;
840            }
841        }
842        &self.environments
843    }
844
845    /// Refill the cache with a single entry. It's typically called during transaction loading, and
846    /// transaction processing (for program management instructions).
847    /// It replaces the existing entry (if any) with the provided entry. The return value contains
848    /// `true` if an entry existed.
849    /// The function also returns the newly inserted value.
850    pub fn replenish(
851        &mut self,
852        key: Pubkey,
853        entry: Arc<ProgramCacheEntry>,
854    ) -> (bool, Arc<ProgramCacheEntry>) {
855        (self.entries.insert(key, entry.clone()).is_some(), entry)
856    }
857
858    /// Store an entry in `modified_entries` for a program modified during the
859    /// transaction batch.
860    pub fn store_modified_entry(&mut self, key: Pubkey, entry: Arc<ProgramCacheEntry>) {
861        self.modified_entries.insert(key, entry);
862    }
863
864    /// Drain the program cache's modified entries, returning the owned
865    /// collection.
866    pub fn drain_modified_entries(&mut self) -> HashMap<Pubkey, Arc<ProgramCacheEntry>> {
867        std::mem::take(&mut self.modified_entries)
868    }
869
870    pub fn find(&self, key: &Pubkey) -> Option<Arc<ProgramCacheEntry>> {
871        // First lookup the cache of the programs modified by the current
872        // transaction. If not found, lookup the cache of the cache of the
873        // programs that are loaded for the transaction batch.
874        self.modified_entries
875            .get(key)
876            .or(self.entries.get(key))
877            .map(|entry| {
878                if entry.is_implicit_delay_visibility_tombstone(self.slot) {
879                    // Found a program entry on the current fork, but it's not effective
880                    // yet. It indicates that the program has delayed visibility. Return
881                    // the tombstone to reflect that.
882                    Arc::new(ProgramCacheEntry::new_tombstone(
883                        entry.deployment_slot,
884                        entry.account_owner,
885                        ProgramCacheEntryType::DelayVisibility,
886                    ))
887                } else {
888                    entry.clone()
889                }
890            })
891    }
892
893    pub fn slot(&self) -> Slot {
894        self.slot
895    }
896
897    #[cfg(any(test, feature = "dev-context-only-utils"))]
898    pub fn set_slot_for_tests(&mut self, slot: Slot) {
899        self.slot = slot;
900    }
901
902    pub fn merge(&mut self, modified_entries: &HashMap<Pubkey, Arc<ProgramCacheEntry>>) {
903        modified_entries.iter().for_each(|(key, entry)| {
904            self.merged_modified = true;
905            self.replenish(*key, entry.clone());
906        })
907    }
908
909    pub fn is_empty(&self) -> bool {
910        self.entries.is_empty()
911    }
912}
913
914pub enum ProgramCacheMatchCriteria {
915    DeployedOnOrAfterSlot(Slot),
916    Tombstone,
917    NoCriteria,
918}
919
920impl ProgramCache {
921    pub fn new(root_slot: Slot, root_epoch: Epoch) -> Self {
922        Self {
923            index: IndexImplementation::V1 {
924                entries: HashMap::new(),
925                loading_entries: Mutex::new(HashMap::new()),
926            },
927            latest_root_slot: root_slot,
928            latest_root_epoch: root_epoch,
929            environments: ProgramRuntimeEnvironments::default(),
930            upcoming_environments: None,
931            programs_to_recompile: Vec::default(),
932            stats: ProgramCacheStats::default(),
933            loading_task_waiter: Arc::new(LoadingTaskWaiter::default()),
934        }
935    }
936
937    /// Returns the current environments depending on the given epoch
938    pub fn get_environments_for_epoch(&self, epoch: Epoch) -> ProgramRuntimeEnvironments {
939        if epoch != self.latest_root_epoch {
940            if let Some(upcoming_environments) = self.upcoming_environments.as_ref() {
941                return upcoming_environments.clone();
942            }
943        }
944        self.environments.clone()
945    }
946
947    /// Returns the upcoming environments depending on the given epoch
948    pub fn get_upcoming_environments_for_epoch(
949        &self,
950        epoch: Epoch,
951    ) -> Option<ProgramRuntimeEnvironments> {
952        if epoch == self.latest_root_epoch {
953            return self.upcoming_environments.clone();
954        }
955        None
956    }
957
958    /// Insert a single entry. It's typically called during transaction loading,
959    /// when the cache doesn't contain the entry corresponding to program `key`.
960    pub fn assign_program(&mut self, key: Pubkey, entry: Arc<ProgramCacheEntry>) -> bool {
961        debug_assert!(!matches!(
962            &entry.program,
963            ProgramCacheEntryType::DelayVisibility
964        ));
965        // This function always returns `true` during normal operation.
966        // Only during the cache preparation phase this can return `false`
967        // for entries with `upcoming_environments`.
968        fn is_current_env(
969            environments: &ProgramRuntimeEnvironments,
970            env_opt: Option<&ProgramRuntimeEnvironment>,
971        ) -> bool {
972            env_opt
973                .map(|env| {
974                    Arc::ptr_eq(env, &environments.program_runtime_v1)
975                        || Arc::ptr_eq(env, &environments.program_runtime_v2)
976                })
977                .unwrap_or(true)
978        }
979        match &mut self.index {
980            IndexImplementation::V1 { entries, .. } => {
981                let slot_versions = &mut entries.entry(key).or_default();
982                match slot_versions.binary_search_by(|at| {
983                    at.effective_slot
984                        .cmp(&entry.effective_slot)
985                        .then(at.deployment_slot.cmp(&entry.deployment_slot))
986                        .then(
987                            // This `.then()` has no effect during normal operation.
988                            // Only during the cache preparation phase this does allow entries
989                            // which only differ in their environment to be interleaved in `slot_versions`.
990                            is_current_env(&self.environments, at.program.get_environment()).cmp(
991                                &is_current_env(
992                                    &self.environments,
993                                    entry.program.get_environment(),
994                                ),
995                            ),
996                        )
997                }) {
998                    Ok(index) => {
999                        let existing = slot_versions.get_mut(index).unwrap();
1000                        match (&existing.program, &entry.program) {
1001                            (
1002                                ProgramCacheEntryType::Builtin(_),
1003                                ProgramCacheEntryType::Builtin(_),
1004                            )
1005                            | (
1006                                ProgramCacheEntryType::Unloaded(_),
1007                                ProgramCacheEntryType::Loaded(_),
1008                            ) => {}
1009                            _ => {
1010                                // Something is wrong, I can feel it ...
1011                                error!("ProgramCache::assign_program() failed key={:?} existing={:?} entry={:?}", key, slot_versions, entry);
1012                                debug_assert!(false, "Unexpected replacement of an entry");
1013                                self.stats.replacements.fetch_add(1, Ordering::Relaxed);
1014                                return true;
1015                            }
1016                        }
1017                        // Copy over the usage counter to the new entry
1018                        entry.tx_usage_counter.fetch_add(
1019                            existing.tx_usage_counter.load(Ordering::Relaxed),
1020                            Ordering::Relaxed,
1021                        );
1022                        entry.ix_usage_counter.fetch_add(
1023                            existing.ix_usage_counter.load(Ordering::Relaxed),
1024                            Ordering::Relaxed,
1025                        );
1026                        *existing = Arc::clone(&entry);
1027                        self.stats.reloads.fetch_add(1, Ordering::Relaxed);
1028                    }
1029                    Err(index) => {
1030                        self.stats.insertions.fetch_add(1, Ordering::Relaxed);
1031                        slot_versions.insert(index, Arc::clone(&entry));
1032                    }
1033                }
1034            }
1035        }
1036        false
1037    }
1038
1039    pub fn prune_by_deployment_slot(&mut self, slot: Slot) {
1040        match &mut self.index {
1041            IndexImplementation::V1 { entries, .. } => {
1042                for second_level in entries.values_mut() {
1043                    second_level.retain(|entry| entry.deployment_slot != slot);
1044                }
1045                self.remove_programs_with_no_entries();
1046            }
1047        }
1048    }
1049
1050    /// Before rerooting the blockstore this removes all superfluous entries
1051    pub fn prune(&mut self, new_root_slot: Slot, new_root_epoch: Epoch) {
1052        let mut preparation_phase_ends = false;
1053        if self.latest_root_epoch != new_root_epoch {
1054            self.latest_root_epoch = new_root_epoch;
1055            if let Some(upcoming_environments) = self.upcoming_environments.take() {
1056                preparation_phase_ends = true;
1057                self.environments = upcoming_environments;
1058                self.programs_to_recompile.clear();
1059            }
1060        }
1061        match &mut self.index {
1062            IndexImplementation::V1 { entries, .. } => {
1063                for second_level in entries.values_mut() {
1064                    // Remove entries un/re/deployed on orphan forks
1065                    let mut first_ancestor_found = false;
1066                    let mut first_ancestor_env = None;
1067                    *second_level = second_level
1068                        .iter()
1069                        .rev()
1070                        .filter(|entry| {
1071                            if entry.deployment_slot >= new_root_slot {
1072                                true
1073                            } else if entry.deployment_slot < new_root_slot
1074                                || entry.deployment_slot <= self.latest_root_slot
1075                            {
1076                                if !first_ancestor_found {
1077                                    first_ancestor_found = true;
1078                                    first_ancestor_env = entry.program.get_environment();
1079                                    return true;
1080                                }
1081                                // Do not prune the entry if the runtime environment of the entry is different
1082                                // than the entry that was previously found (stored in first_ancestor_env).
1083                                // Different environment indicates that this entry might belong to an older
1084                                // epoch that had a different environment (e.g. different feature set).
1085                                // Once the root moves to the new/current epoch, the entry will get pruned.
1086                                // But, until then the entry might still be getting used by an older slot.
1087                                if let Some(entry_env) = entry.program.get_environment() {
1088                                    if let Some(env) = first_ancestor_env {
1089                                        if !Arc::ptr_eq(entry_env, env) {
1090                                            return true;
1091                                        }
1092                                    }
1093                                }
1094                                self.stats.prunes_orphan.fetch_add(1, Ordering::Relaxed);
1095                                false
1096                            } else {
1097                                self.stats.prunes_orphan.fetch_add(1, Ordering::Relaxed);
1098                                false
1099                            }
1100                        })
1101                        .filter(|entry| {
1102                            // Remove outdated environment of previous feature set
1103                            if preparation_phase_ends
1104                                && !Self::matches_environment(entry, &self.environments)
1105                            {
1106                                self.stats
1107                                    .prunes_environment
1108                                    .fetch_add(1, Ordering::Relaxed);
1109                                return false;
1110                            }
1111                            true
1112                        })
1113                        .cloned()
1114                        .collect();
1115                    second_level.reverse();
1116                }
1117            }
1118        }
1119        self.remove_programs_with_no_entries();
1120        debug_assert!(self.latest_root_slot <= new_root_slot);
1121        self.latest_root_slot = new_root_slot;
1122    }
1123
1124    fn matches_environment(
1125        entry: &Arc<ProgramCacheEntry>,
1126        environments: &ProgramRuntimeEnvironments,
1127    ) -> bool {
1128        let Some(environment) = entry.program.get_environment() else {
1129            return true;
1130        };
1131        Arc::ptr_eq(environment, &environments.program_runtime_v1)
1132            || Arc::ptr_eq(environment, &environments.program_runtime_v2)
1133    }
1134
1135    fn matches_criteria(
1136        program: &Arc<ProgramCacheEntry>,
1137        criteria: &ProgramCacheMatchCriteria,
1138    ) -> bool {
1139        match criteria {
1140            ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(slot) => {
1141                program.deployment_slot >= *slot
1142            }
1143            ProgramCacheMatchCriteria::Tombstone => program.is_tombstone(),
1144            ProgramCacheMatchCriteria::NoCriteria => true,
1145        }
1146    }
1147
1148    /// Extracts a subset of the programs relevant to a transaction batch
1149    /// and returns which program accounts the accounts DB needs to load.
1150    pub fn extract(
1151        &self,
1152        search_for: &mut Vec<(Pubkey, (ProgramCacheMatchCriteria, u64))>,
1153        loaded_programs_for_tx_batch: &mut ProgramCacheForTxBatch,
1154        is_first_round: bool,
1155    ) -> Option<(Pubkey, u64)> {
1156        let mut cooperative_loading_task = None;
1157        match &self.index {
1158            IndexImplementation::V1 {
1159                entries,
1160                loading_entries,
1161            } => {
1162                search_for.retain(|(key, (match_criteria, usage_count))| {
1163                    if let Some(second_level) = entries.get(key) {
1164                        for entry in second_level.iter().rev() {
1165                            if entry.deployment_slot <= self.latest_root_slot
1166                                || entry.deployment_slot <= loaded_programs_for_tx_batch.slot
1167                            {
1168                                let entry_to_return = if loaded_programs_for_tx_batch.slot
1169                                    >= entry.effective_slot
1170                                    && Self::matches_environment(
1171                                        entry,
1172                                        &loaded_programs_for_tx_batch.environments,
1173                                    ) {
1174                                    if !Self::matches_criteria(entry, match_criteria) {
1175                                        break;
1176                                    }
1177                                    if let ProgramCacheEntryType::Unloaded(_environment) =
1178                                        &entry.program
1179                                    {
1180                                        break;
1181                                    }
1182                                    entry.clone()
1183                                } else if entry.is_implicit_delay_visibility_tombstone(
1184                                    loaded_programs_for_tx_batch.slot,
1185                                ) {
1186                                    // Found a program entry on the current fork, but it's not effective
1187                                    // yet. It indicates that the program has delayed visibility. Return
1188                                    // the tombstone to reflect that.
1189                                    Arc::new(ProgramCacheEntry::new_tombstone(
1190                                        entry.deployment_slot,
1191                                        entry.account_owner,
1192                                        ProgramCacheEntryType::DelayVisibility,
1193                                    ))
1194                                } else {
1195                                    continue;
1196                                };
1197                                entry_to_return
1198                                    .update_access_slot(loaded_programs_for_tx_batch.slot);
1199                                entry_to_return
1200                                    .tx_usage_counter
1201                                    .fetch_add(*usage_count, Ordering::Relaxed);
1202                                loaded_programs_for_tx_batch
1203                                    .entries
1204                                    .insert(*key, entry_to_return);
1205                                return false;
1206                            }
1207                        }
1208                    }
1209                    if cooperative_loading_task.is_none() {
1210                        let mut loading_entries = loading_entries.lock().unwrap();
1211                        let entry = loading_entries.entry(*key);
1212                        if let Entry::Vacant(entry) = entry {
1213                            entry.insert((
1214                                loaded_programs_for_tx_batch.slot,
1215                                thread::current().id(),
1216                            ));
1217                            cooperative_loading_task = Some((*key, *usage_count));
1218                        }
1219                    }
1220                    true
1221                });
1222            }
1223        }
1224        if is_first_round {
1225            self.stats
1226                .misses
1227                .fetch_add(search_for.len() as u64, Ordering::Relaxed);
1228            self.stats.hits.fetch_add(
1229                loaded_programs_for_tx_batch.entries.len() as u64,
1230                Ordering::Relaxed,
1231            );
1232        }
1233        cooperative_loading_task
1234    }
1235
1236    /// Called by Bank::replenish_program_cache() for each program that is done loading.
1237    pub fn finish_cooperative_loading_task(
1238        &mut self,
1239        slot: Slot,
1240        key: Pubkey,
1241        loaded_program: Arc<ProgramCacheEntry>,
1242    ) -> bool {
1243        match &mut self.index {
1244            IndexImplementation::V1 {
1245                loading_entries, ..
1246            } => {
1247                let loading_thread = loading_entries.get_mut().unwrap().remove(&key);
1248                debug_assert_eq!(loading_thread, Some((slot, thread::current().id())));
1249                // Check that it will be visible to our own fork once inserted
1250                if loaded_program.deployment_slot > self.latest_root_slot
1251                    && loaded_program.deployment_slot > slot
1252                {
1253                    self.stats.lost_insertions.fetch_add(1, Ordering::Relaxed);
1254                }
1255                let was_occupied = self.assign_program(key, loaded_program);
1256                self.loading_task_waiter.notify();
1257                was_occupied
1258            }
1259        }
1260    }
1261
1262    pub fn merge(&mut self, modified_entries: &HashMap<Pubkey, Arc<ProgramCacheEntry>>) {
1263        modified_entries.iter().for_each(|(key, entry)| {
1264            self.assign_program(*key, entry.clone());
1265        })
1266    }
1267
1268    /// Returns the list of entries which are verified and compiled.
1269    pub fn get_flattened_entries(
1270        &self,
1271        include_program_runtime_v1: bool,
1272        _include_program_runtime_v2: bool,
1273    ) -> Vec<(Pubkey, Arc<ProgramCacheEntry>)> {
1274        match &self.index {
1275            IndexImplementation::V1 { entries, .. } => entries
1276                .iter()
1277                .flat_map(|(id, second_level)| {
1278                    second_level
1279                        .iter()
1280                        .filter_map(move |program| match program.program {
1281                            ProgramCacheEntryType::Loaded(_) => {
1282                                if include_program_runtime_v1 {
1283                                    Some((*id, program.clone()))
1284                                } else {
1285                                    None
1286                                }
1287                            }
1288                            _ => None,
1289                        })
1290                })
1291                .collect(),
1292        }
1293    }
1294
1295    /// Returns the list of all entries in the cache.
1296    #[cfg(any(test, feature = "dev-context-only-utils"))]
1297    pub fn get_flattened_entries_for_tests(&self) -> Vec<(Pubkey, Arc<ProgramCacheEntry>)> {
1298        match &self.index {
1299            IndexImplementation::V1 { entries, .. } => entries
1300                .iter()
1301                .flat_map(|(id, second_level)| {
1302                    second_level.iter().map(|program| (*id, program.clone()))
1303                })
1304                .collect(),
1305        }
1306    }
1307
1308    /// Returns the slot versions for the given program id.
1309    #[cfg(any(test, feature = "dev-context-only-utils"))]
1310    pub fn get_slot_versions_for_tests(&self, key: &Pubkey) -> &[Arc<ProgramCacheEntry>] {
1311        match &self.index {
1312            IndexImplementation::V1 { entries, .. } => entries
1313                .get(key)
1314                .map(|second_level| second_level.as_ref())
1315                .unwrap_or(&[]),
1316        }
1317    }
1318
1319    /// Unloads programs which were used infrequently
1320    pub fn sort_and_unload(&mut self, shrink_to: PercentageInteger) {
1321        let mut sorted_candidates = self.get_flattened_entries(true, true);
1322        sorted_candidates
1323            .sort_by_cached_key(|(_id, program)| program.tx_usage_counter.load(Ordering::Relaxed));
1324        let num_to_unload = sorted_candidates
1325            .len()
1326            .saturating_sub(shrink_to.apply_to(MAX_LOADED_ENTRY_COUNT));
1327        self.unload_program_entries(sorted_candidates.iter().take(num_to_unload));
1328    }
1329
1330    /// Evicts programs using 2's random selection, choosing the least used program out of the two entries.
1331    /// The eviction is performed enough number of times to reduce the cache usage to the given percentage.
1332    pub fn evict_using_2s_random_selection(&mut self, shrink_to: PercentageInteger, now: Slot) {
1333        let mut candidates = self.get_flattened_entries(true, true);
1334        self.stats
1335            .water_level
1336            .store(candidates.len() as u64, Ordering::Relaxed);
1337        let num_to_unload = candidates
1338            .len()
1339            .saturating_sub(shrink_to.apply_to(MAX_LOADED_ENTRY_COUNT));
1340        fn random_index_and_usage_counter(
1341            candidates: &[(Pubkey, Arc<ProgramCacheEntry>)],
1342            now: Slot,
1343        ) -> (usize, u64) {
1344            let mut rng = thread_rng();
1345            let index = rng.gen_range(0..candidates.len());
1346            let usage_counter = candidates
1347                .get(index)
1348                .expect("Failed to get cached entry")
1349                .1
1350                .decayed_usage_counter(now);
1351            (index, usage_counter)
1352        }
1353
1354        for _ in 0..num_to_unload {
1355            let (index1, usage_counter1) = random_index_and_usage_counter(&candidates, now);
1356            let (index2, usage_counter2) = random_index_and_usage_counter(&candidates, now);
1357
1358            let (program, entry) = if usage_counter1 < usage_counter2 {
1359                candidates.swap_remove(index1)
1360            } else {
1361                candidates.swap_remove(index2)
1362            };
1363            self.unload_program_entry(&program, &entry);
1364        }
1365    }
1366
1367    /// Removes all the entries at the given keys, if they exist
1368    pub fn remove_programs(&mut self, keys: impl Iterator<Item = Pubkey>) {
1369        match &mut self.index {
1370            IndexImplementation::V1 { entries, .. } => {
1371                for k in keys {
1372                    entries.remove(&k);
1373                }
1374            }
1375        }
1376    }
1377
1378    /// This function removes the given entry for the given program from the cache.
1379    /// The function expects that the program and entry exists in the cache. Otherwise it'll panic.
1380    fn unload_program_entry(&mut self, program: &Pubkey, remove_entry: &Arc<ProgramCacheEntry>) {
1381        match &mut self.index {
1382            IndexImplementation::V1 { entries, .. } => {
1383                let second_level = entries.get_mut(program).expect("Cache lookup failed");
1384                let candidate = second_level
1385                    .iter_mut()
1386                    .find(|entry| entry == &remove_entry)
1387                    .expect("Program entry not found");
1388
1389                // Certain entry types cannot be unloaded, such as tombstones, or already unloaded entries.
1390                // For such entries, `to_unloaded()` will return None.
1391                // These entry types do not occupy much memory.
1392                if let Some(unloaded) = candidate.to_unloaded() {
1393                    if candidate.tx_usage_counter.load(Ordering::Relaxed) == 1 {
1394                        self.stats.one_hit_wonders.fetch_add(1, Ordering::Relaxed);
1395                    }
1396                    self.stats
1397                        .evictions
1398                        .entry(*program)
1399                        .and_modify(|c| *c = c.saturating_add(1))
1400                        .or_insert(1);
1401                    *candidate = Arc::new(unloaded);
1402                }
1403            }
1404        }
1405    }
1406
1407    fn unload_program_entries<'a>(
1408        &mut self,
1409        remove: impl Iterator<Item = &'a (Pubkey, Arc<ProgramCacheEntry>)>,
1410    ) {
1411        for (program, entry) in remove {
1412            self.unload_program_entry(program, entry);
1413        }
1414    }
1415
1416    fn remove_programs_with_no_entries(&mut self) {
1417        match &mut self.index {
1418            IndexImplementation::V1 { entries, .. } => {
1419                let num_programs_before_removal = entries.len();
1420                entries.retain(|_key, second_level| !second_level.is_empty());
1421                if entries.len() < num_programs_before_removal {
1422                    self.stats.empty_entries.fetch_add(
1423                        num_programs_before_removal.saturating_sub(entries.len()) as u64,
1424                        Ordering::Relaxed,
1425                    );
1426                }
1427            }
1428        }
1429    }
1430}
1431
1432#[cfg(feature = "frozen-abi")]
1433impl rialo_frozen_abi::abi_example::AbiExample for ProgramCacheEntry {
1434    fn example() -> Self {
1435        // ProgramCacheEntry isn't serializable by definition.
1436        Self::default()
1437    }
1438}
1439
1440#[cfg(feature = "frozen-abi")]
1441impl rialo_frozen_abi::abi_example::AbiExample for ProgramCache {
1442    fn example() -> Self {
1443        // ProgramCache isn't serializable by definition.
1444        Self::new(Slot::default(), Epoch::default())
1445    }
1446}
1447
1448#[cfg(false)] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
1449#[cfg(test)]
1450mod tests {
1451    use std::{
1452        fs::File,
1453        io::Read,
1454        sync::{
1455            atomic::{AtomicU64, Ordering},
1456            Arc,
1457        },
1458    };
1459
1460    use assert_matches::assert_matches;
1461    use percentage::Percentage;
1462    use rialo_s_clock::Slot;
1463    use rialo_s_pubkey::Pubkey;
1464    use solana_sbpf::{elf::Executable, program::BuiltinProgram};
1465    use test_case::{test_case, test_matrix};
1466
1467    use crate::loaded_programs::{
1468        ProgramCache, ProgramCacheEntry, ProgramCacheEntryOwner, ProgramCacheEntryType,
1469        ProgramCacheForTxBatch, ProgramCacheMatchCriteria, ProgramRuntimeEnvironment,
1470        ProgramRuntimeEnvironments, DELAY_VISIBILITY_SLOT_OFFSET,
1471    };
1472
1473    static MOCK_ENVIRONMENT: std::sync::OnceLock<ProgramRuntimeEnvironment> =
1474        std::sync::OnceLock::<ProgramRuntimeEnvironment>::new();
1475
1476    fn get_mock_env() -> ProgramRuntimeEnvironment {
1477        MOCK_ENVIRONMENT
1478            .get_or_init(|| Arc::new(BuiltinProgram::new_mock()))
1479            .clone()
1480    }
1481
1482    fn new_mock_cache() -> ProgramCache {
1483        let mut cache = ProgramCache::new(0, 0);
1484        cache.environments.program_runtime_v1 = get_mock_env();
1485        cache
1486    }
1487
1488    fn new_test_entry(deployment_slot: Slot, effective_slot: Slot) -> Arc<ProgramCacheEntry> {
1489        new_test_entry_with_usage(deployment_slot, effective_slot, AtomicU64::default())
1490    }
1491
1492    fn new_loaded_entry(env: ProgramRuntimeEnvironment) -> ProgramCacheEntryType {
1493        let mut elf = Vec::new();
1494        File::open("../programs/bpf_loader/test_elfs/out/noop_aligned.so")
1495            .unwrap()
1496            .read_to_end(&mut elf)
1497            .unwrap();
1498        let executable = Executable::load(&elf, env).unwrap();
1499        ProgramCacheEntryType::Loaded(super::ExecutableKind::Ebpf(executable))
1500    }
1501
1502    fn new_test_entry_with_usage(
1503        deployment_slot: Slot,
1504        effective_slot: Slot,
1505        usage_counter: AtomicU64,
1506    ) -> Arc<ProgramCacheEntry> {
1507        Arc::new(ProgramCacheEntry {
1508            program: new_loaded_entry(get_mock_env()),
1509            account_owner: ProgramCacheEntryOwner::LoaderV2,
1510            account_size: 0,
1511            deployment_slot,
1512            effective_slot,
1513            tx_usage_counter: usage_counter,
1514            ix_usage_counter: AtomicU64::default(),
1515            latest_access_slot: AtomicU64::new(deployment_slot),
1516        })
1517    }
1518
1519    fn new_test_builtin_entry(
1520        deployment_slot: Slot,
1521        effective_slot: Slot,
1522    ) -> Arc<ProgramCacheEntry> {
1523        Arc::new(ProgramCacheEntry {
1524            program: ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
1525            account_owner: ProgramCacheEntryOwner::NativeLoader,
1526            account_size: 0,
1527            deployment_slot,
1528            effective_slot,
1529            tx_usage_counter: AtomicU64::default(),
1530            ix_usage_counter: AtomicU64::default(),
1531            latest_access_slot: AtomicU64::default(),
1532        })
1533    }
1534
1535    fn set_tombstone(
1536        cache: &mut ProgramCache,
1537        key: Pubkey,
1538        slot: Slot,
1539        reason: ProgramCacheEntryType,
1540    ) -> Arc<ProgramCacheEntry> {
1541        let program = Arc::new(ProgramCacheEntry::new_tombstone(
1542            slot,
1543            ProgramCacheEntryOwner::LoaderV2,
1544            reason,
1545        ));
1546        cache.assign_program(key, program.clone());
1547        program
1548    }
1549
1550    fn insert_unloaded_entry(
1551        cache: &mut ProgramCache,
1552        key: Pubkey,
1553        slot: Slot,
1554    ) -> Arc<ProgramCacheEntry> {
1555        let loaded = new_test_entry_with_usage(slot, slot.saturating_add(1), AtomicU64::default());
1556        let unloaded = Arc::new(loaded.to_unloaded().expect("Failed to unload the program"));
1557        cache.assign_program(key, unloaded.clone());
1558        unloaded
1559    }
1560
1561    fn num_matching_entries<P>(cache: &ProgramCache, predicate: P) -> usize
1562    where
1563        P: Fn(&ProgramCacheEntryType) -> bool,
1564    {
1565        cache
1566            .get_flattened_entries_for_tests()
1567            .iter()
1568            .filter(|(_key, program)| predicate(&program.program))
1569            .count()
1570    }
1571
1572    #[test]
1573    fn test_usage_counter_decay() {
1574        let _cache = new_mock_cache();
1575        let program = new_test_entry_with_usage(10, 11, AtomicU64::new(32));
1576        program.update_access_slot(15);
1577        assert_eq!(program.decayed_usage_counter(15), 32);
1578        assert_eq!(program.decayed_usage_counter(16), 16);
1579        assert_eq!(program.decayed_usage_counter(17), 8);
1580        assert_eq!(program.decayed_usage_counter(18), 4);
1581        assert_eq!(program.decayed_usage_counter(19), 2);
1582        assert_eq!(program.decayed_usage_counter(20), 1);
1583        assert_eq!(program.decayed_usage_counter(21), 0);
1584        assert_eq!(program.decayed_usage_counter(15), 32);
1585        assert_eq!(program.decayed_usage_counter(14), 32);
1586
1587        program.update_access_slot(18);
1588        assert_eq!(program.decayed_usage_counter(15), 32);
1589        assert_eq!(program.decayed_usage_counter(16), 32);
1590        assert_eq!(program.decayed_usage_counter(17), 32);
1591        assert_eq!(program.decayed_usage_counter(18), 32);
1592        assert_eq!(program.decayed_usage_counter(19), 16);
1593        assert_eq!(program.decayed_usage_counter(20), 8);
1594        assert_eq!(program.decayed_usage_counter(21), 4);
1595
1596        // Decay for 63 or more slots
1597        assert_eq!(program.decayed_usage_counter(18 + 63), 0);
1598        assert_eq!(program.decayed_usage_counter(100), 0);
1599    }
1600
1601    fn program_deploy_test_helper(
1602        cache: &mut ProgramCache,
1603        program: Pubkey,
1604        deployment_slots: Vec<Slot>,
1605        usage_counters: Vec<u64>,
1606        programs: &mut Vec<(Pubkey, Slot, u64)>,
1607    ) {
1608        // Add multiple entries for program
1609        deployment_slots
1610            .iter()
1611            .enumerate()
1612            .for_each(|(i, deployment_slot)| {
1613                let usage_counter = *usage_counters.get(i).unwrap_or(&0);
1614                cache.assign_program(
1615                    program,
1616                    new_test_entry_with_usage(
1617                        *deployment_slot,
1618                        (*deployment_slot).saturating_add(2),
1619                        AtomicU64::new(usage_counter),
1620                    ),
1621                );
1622                programs.push((program, *deployment_slot, usage_counter));
1623            });
1624
1625        // Add tombstones entries for program
1626        let env = Arc::new(BuiltinProgram::new_mock());
1627        for slot in 21..31 {
1628            set_tombstone(
1629                cache,
1630                program,
1631                slot,
1632                ProgramCacheEntryType::FailedVerification(env.clone()),
1633            );
1634        }
1635
1636        // Add unloaded entries for program
1637        for slot in 31..41 {
1638            insert_unloaded_entry(cache, program, slot);
1639        }
1640    }
1641
1642    #[test]
1643    fn test_random_eviction() {
1644        let mut programs = vec![];
1645
1646        let mut cache = new_mock_cache();
1647
1648        // This test adds different kind of entries to the cache.
1649        // Tombstones and unloaded entries are expected to not be evicted.
1650        // It also adds multiple entries for three programs as it tries to create a typical cache instance.
1651
1652        // Program 1
1653        program_deploy_test_helper(
1654            &mut cache,
1655            Pubkey::new_unique(),
1656            vec![0, 10, 20],
1657            vec![4, 5, 25],
1658            &mut programs,
1659        );
1660
1661        // Program 2
1662        program_deploy_test_helper(
1663            &mut cache,
1664            Pubkey::new_unique(),
1665            vec![5, 11],
1666            vec![0, 2],
1667            &mut programs,
1668        );
1669
1670        // Program 3
1671        program_deploy_test_helper(
1672            &mut cache,
1673            Pubkey::new_unique(),
1674            vec![0, 5, 15],
1675            vec![100, 3, 20],
1676            &mut programs,
1677        );
1678
1679        // 1 for each deployment slot
1680        let num_loaded_expected = 8;
1681        // 10 for each program
1682        let num_unloaded_expected = 30;
1683        // 10 for each program
1684        let num_tombstones_expected = 30;
1685
1686        // Count the number of loaded, unloaded and tombstone entries.
1687        programs.sort_by_key(|(_id, _slot, usage_count)| *usage_count);
1688        let num_loaded = num_matching_entries(&cache, |program_type| {
1689            matches!(program_type, ProgramCacheEntryType::Loaded(_))
1690        });
1691        let num_unloaded = num_matching_entries(&cache, |program_type| {
1692            matches!(program_type, ProgramCacheEntryType::Unloaded(_))
1693        });
1694        let num_tombstones = num_matching_entries(&cache, |program_type| {
1695            matches!(
1696                program_type,
1697                ProgramCacheEntryType::DelayVisibility
1698                    | ProgramCacheEntryType::FailedVerification(_)
1699                    | ProgramCacheEntryType::Closed
1700            )
1701        });
1702
1703        // Test that the cache is constructed with the expected number of entries.
1704        assert_eq!(num_loaded, num_loaded_expected);
1705        assert_eq!(num_unloaded, num_unloaded_expected);
1706        assert_eq!(num_tombstones, num_tombstones_expected);
1707
1708        // Evict entries from the cache
1709        let eviction_pct = 1;
1710
1711        let num_loaded_expected =
1712            Percentage::from(eviction_pct).apply_to(crate::loaded_programs::MAX_LOADED_ENTRY_COUNT);
1713        let num_unloaded_expected = num_unloaded_expected + num_loaded - num_loaded_expected;
1714        cache.evict_using_2s_random_selection(Percentage::from(eviction_pct), 21);
1715
1716        // Count the number of loaded, unloaded and tombstone entries.
1717        let num_loaded = num_matching_entries(&cache, |program_type| {
1718            matches!(program_type, ProgramCacheEntryType::Loaded(_))
1719        });
1720        let num_unloaded = num_matching_entries(&cache, |program_type| {
1721            matches!(program_type, ProgramCacheEntryType::Unloaded(_))
1722        });
1723        let num_tombstones = num_matching_entries(&cache, |program_type| {
1724            matches!(program_type, ProgramCacheEntryType::FailedVerification(_))
1725        });
1726
1727        // However many entries are left after the shrink
1728        assert_eq!(num_loaded, num_loaded_expected);
1729        // The original unloaded entries + the evicted loaded entries
1730        assert_eq!(num_unloaded, num_unloaded_expected);
1731        // The original tombstones are not evicted
1732        assert_eq!(num_tombstones, num_tombstones_expected);
1733    }
1734
1735    #[test]
1736    fn test_eviction() {
1737        let mut programs = vec![];
1738        let mut cache = new_mock_cache();
1739
1740        // Program 1
1741        program_deploy_test_helper(
1742            &mut cache,
1743            Pubkey::new_unique(),
1744            vec![0, 10, 20],
1745            vec![4, 5, 25],
1746            &mut programs,
1747        );
1748
1749        // Program 2
1750        program_deploy_test_helper(
1751            &mut cache,
1752            Pubkey::new_unique(),
1753            vec![5, 11],
1754            vec![0, 2],
1755            &mut programs,
1756        );
1757
1758        // Program 3
1759        program_deploy_test_helper(
1760            &mut cache,
1761            Pubkey::new_unique(),
1762            vec![0, 5, 15],
1763            vec![100, 3, 20],
1764            &mut programs,
1765        );
1766
1767        // 1 for each deployment slot
1768        let num_loaded_expected = 8;
1769        // 10 for each program
1770        let num_unloaded_expected = 30;
1771        // 10 for each program
1772        let num_tombstones_expected = 30;
1773
1774        // Count the number of loaded, unloaded and tombstone entries.
1775        programs.sort_by_key(|(_id, _slot, usage_count)| *usage_count);
1776        let num_loaded = num_matching_entries(&cache, |program_type| {
1777            matches!(program_type, ProgramCacheEntryType::Loaded(_))
1778        });
1779        let num_unloaded = num_matching_entries(&cache, |program_type| {
1780            matches!(program_type, ProgramCacheEntryType::Unloaded(_))
1781        });
1782        let num_tombstones = num_matching_entries(&cache, |program_type| {
1783            matches!(program_type, ProgramCacheEntryType::FailedVerification(_))
1784        });
1785
1786        // Test that the cache is constructed with the expected number of entries.
1787        assert_eq!(num_loaded, num_loaded_expected);
1788        assert_eq!(num_unloaded, num_unloaded_expected);
1789        assert_eq!(num_tombstones, num_tombstones_expected);
1790
1791        // Evict entries from the cache
1792        let eviction_pct = 1;
1793
1794        let num_loaded_expected =
1795            Percentage::from(eviction_pct).apply_to(crate::loaded_programs::MAX_LOADED_ENTRY_COUNT);
1796        let num_unloaded_expected = num_unloaded_expected + num_loaded - num_loaded_expected;
1797
1798        cache.sort_and_unload(Percentage::from(eviction_pct));
1799
1800        // Check that every program is still in the cache.
1801        let entries = cache.get_flattened_entries_for_tests();
1802        programs.iter().for_each(|entry| {
1803            assert!(entries.iter().any(|(key, _entry)| key == &entry.0));
1804        });
1805
1806        let unloaded = entries
1807            .iter()
1808            .filter_map(|(key, program)| {
1809                matches!(program.program, ProgramCacheEntryType::Unloaded(_))
1810                    .then_some((*key, program.tx_usage_counter.load(Ordering::Relaxed)))
1811            })
1812            .collect::<Vec<(Pubkey, u64)>>();
1813
1814        for index in 0..3 {
1815            let expected = programs.get(index).expect("Missing program");
1816            assert!(unloaded.contains(&(expected.0, expected.2)));
1817        }
1818
1819        // Count the number of loaded, unloaded and tombstone entries.
1820        let num_loaded = num_matching_entries(&cache, |program_type| {
1821            matches!(program_type, ProgramCacheEntryType::Loaded(_))
1822        });
1823        let num_unloaded = num_matching_entries(&cache, |program_type| {
1824            matches!(program_type, ProgramCacheEntryType::Unloaded(_))
1825        });
1826        let num_tombstones = num_matching_entries(&cache, |program_type| {
1827            matches!(
1828                program_type,
1829                ProgramCacheEntryType::DelayVisibility
1830                    | ProgramCacheEntryType::FailedVerification(_)
1831                    | ProgramCacheEntryType::Closed
1832            )
1833        });
1834
1835        // However many entries are left after the shrink
1836        assert_eq!(num_loaded, num_loaded_expected);
1837        // The original unloaded entries + the evicted loaded entries
1838        assert_eq!(num_unloaded, num_unloaded_expected);
1839        // The original tombstones are not evicted
1840        assert_eq!(num_tombstones, num_tombstones_expected);
1841    }
1842
1843    #[test]
1844    fn test_usage_count_of_unloaded_program() {
1845        let mut cache = new_mock_cache();
1846
1847        let program = Pubkey::new_unique();
1848        let evict_to_pct = 2;
1849        let cache_capacity_after_shrink =
1850            Percentage::from(evict_to_pct).apply_to(crate::loaded_programs::MAX_LOADED_ENTRY_COUNT);
1851        // Add enough programs to the cache to trigger 1 eviction after shrinking.
1852        let num_total_programs = (cache_capacity_after_shrink + 1) as u64;
1853        (0..num_total_programs).for_each(|i| {
1854            cache.assign_program(
1855                program,
1856                new_test_entry_with_usage(i, i + 2, AtomicU64::new(i + 10)),
1857            );
1858        });
1859
1860        cache.sort_and_unload(Percentage::from(evict_to_pct));
1861
1862        let num_unloaded = num_matching_entries(&cache, |program_type| {
1863            matches!(program_type, ProgramCacheEntryType::Unloaded(_))
1864        });
1865        assert_eq!(num_unloaded, 1);
1866
1867        cache
1868            .get_flattened_entries_for_tests()
1869            .iter()
1870            .for_each(|(_key, program)| {
1871                if matches!(program.program, ProgramCacheEntryType::Unloaded(_)) {
1872                    // Test that the usage counter is retained for the unloaded program
1873                    assert_eq!(program.tx_usage_counter.load(Ordering::Relaxed), 10);
1874                    assert_eq!(program.deployment_slot, 0);
1875                    assert_eq!(program.effective_slot, 2);
1876                }
1877            });
1878
1879        // Replenish the program that was just unloaded. Use 0 as the usage counter. This should be
1880        // updated with the usage counter from the unloaded program.
1881        cache.assign_program(program, new_test_entry_with_usage(0, 2, AtomicU64::new(0)));
1882
1883        cache
1884            .get_flattened_entries_for_tests()
1885            .iter()
1886            .for_each(|(_key, program)| {
1887                if matches!(program.program, ProgramCacheEntryType::Unloaded(_))
1888                    && program.deployment_slot == 0
1889                    && program.effective_slot == 2
1890                {
1891                    // Test that the usage counter was correctly updated.
1892                    assert_eq!(program.tx_usage_counter.load(Ordering::Relaxed), 10);
1893                }
1894            });
1895    }
1896
1897    #[test]
1898    fn test_fuzz_assign_program_order() {
1899        use rand::prelude::SliceRandom;
1900        const EXPECTED_ENTRIES: [(u64, u64); 7] =
1901            [(1, 2), (5, 5), (5, 6), (5, 10), (9, 10), (10, 10), (3, 12)];
1902        let mut rng = rand::thread_rng();
1903        let program_id = Pubkey::new_unique();
1904        for _ in 0..1000 {
1905            let mut entries = EXPECTED_ENTRIES.to_vec();
1906            entries.shuffle(&mut rng);
1907            let mut cache = new_mock_cache();
1908            for (deployment_slot, effective_slot) in entries {
1909                assert!(!cache
1910                    .assign_program(program_id, new_test_entry(deployment_slot, effective_slot)));
1911            }
1912            for ((deployment_slot, effective_slot), entry) in EXPECTED_ENTRIES
1913                .iter()
1914                .zip(cache.get_slot_versions_for_tests(&program_id).iter())
1915            {
1916                assert_eq!(entry.deployment_slot, *deployment_slot);
1917                assert_eq!(entry.effective_slot, *effective_slot);
1918            }
1919        }
1920    }
1921
1922    #[test_matrix(
1923        (
1924            ProgramCacheEntryType::Closed,
1925            ProgramCacheEntryType::FailedVerification(get_mock_env()),
1926            new_loaded_entry(get_mock_env()),
1927        ),
1928        (
1929            ProgramCacheEntryType::FailedVerification(get_mock_env()),
1930            ProgramCacheEntryType::Closed,
1931            ProgramCacheEntryType::Unloaded(get_mock_env()),
1932            new_loaded_entry(get_mock_env()),
1933            ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
1934        )
1935    )]
1936    #[test_matrix(
1937        (
1938            ProgramCacheEntryType::Unloaded(get_mock_env()),
1939        ),
1940        (
1941            ProgramCacheEntryType::FailedVerification(get_mock_env()),
1942            ProgramCacheEntryType::Closed,
1943            ProgramCacheEntryType::Unloaded(get_mock_env()),
1944            ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
1945        )
1946    )]
1947    #[test_matrix(
1948        (ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),),
1949        (
1950            ProgramCacheEntryType::FailedVerification(get_mock_env()),
1951            ProgramCacheEntryType::Closed,
1952            ProgramCacheEntryType::Unloaded(get_mock_env()),
1953            new_loaded_entry(get_mock_env()),
1954        )
1955    )]
1956    #[should_panic(expected = "Unexpected replacement of an entry")]
1957    fn test_assign_program_failure(old: ProgramCacheEntryType, new: ProgramCacheEntryType) {
1958        let mut cache = new_mock_cache();
1959        let program_id = Pubkey::new_unique();
1960        assert!(!cache.assign_program(
1961            program_id,
1962            Arc::new(ProgramCacheEntry {
1963                program: old,
1964                account_owner: ProgramCacheEntryOwner::LoaderV2,
1965                account_size: 0,
1966                deployment_slot: 10,
1967                effective_slot: 11,
1968                tx_usage_counter: AtomicU64::default(),
1969                ix_usage_counter: AtomicU64::default(),
1970                latest_access_slot: AtomicU64::default(),
1971            }),
1972        ));
1973        cache.assign_program(
1974            program_id,
1975            Arc::new(ProgramCacheEntry {
1976                program: new,
1977                account_owner: ProgramCacheEntryOwner::LoaderV2,
1978                account_size: 0,
1979                deployment_slot: 10,
1980                effective_slot: 11,
1981                tx_usage_counter: AtomicU64::default(),
1982                ix_usage_counter: AtomicU64::default(),
1983                latest_access_slot: AtomicU64::default(),
1984            }),
1985        );
1986    }
1987
1988    #[test_case(
1989        ProgramCacheEntryType::Unloaded(Arc::new(BuiltinProgram::new_mock())),
1990        new_loaded_entry(get_mock_env())
1991    )]
1992    #[test_case(
1993        ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
1994        ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock())
1995    )]
1996    fn test_assign_program_success(old: ProgramCacheEntryType, new: ProgramCacheEntryType) {
1997        let mut cache = new_mock_cache();
1998        let program_id = Pubkey::new_unique();
1999        assert!(!cache.assign_program(
2000            program_id,
2001            Arc::new(ProgramCacheEntry {
2002                program: old,
2003                account_owner: ProgramCacheEntryOwner::LoaderV2,
2004                account_size: 0,
2005                deployment_slot: 10,
2006                effective_slot: 11,
2007                tx_usage_counter: AtomicU64::default(),
2008                ix_usage_counter: AtomicU64::default(),
2009                latest_access_slot: AtomicU64::default(),
2010            }),
2011        ));
2012        assert!(!cache.assign_program(
2013            program_id,
2014            Arc::new(ProgramCacheEntry {
2015                program: new,
2016                account_owner: ProgramCacheEntryOwner::LoaderV2,
2017                account_size: 0,
2018                deployment_slot: 10,
2019                effective_slot: 11,
2020                tx_usage_counter: AtomicU64::default(),
2021                ix_usage_counter: AtomicU64::default(),
2022                latest_access_slot: AtomicU64::default(),
2023            }),
2024        ));
2025    }
2026
2027    #[test]
2028    fn test_tombstone() {
2029        let env = Arc::new(BuiltinProgram::new_mock());
2030        let tombstone = ProgramCacheEntry::new_tombstone(
2031            0,
2032            ProgramCacheEntryOwner::LoaderV2,
2033            ProgramCacheEntryType::FailedVerification(env.clone()),
2034        );
2035        assert_matches!(
2036            tombstone.program,
2037            ProgramCacheEntryType::FailedVerification(_)
2038        );
2039        assert!(tombstone.is_tombstone());
2040        assert_eq!(tombstone.deployment_slot, 0);
2041        assert_eq!(tombstone.effective_slot, 0);
2042
2043        let tombstone = ProgramCacheEntry::new_tombstone(
2044            100,
2045            ProgramCacheEntryOwner::LoaderV2,
2046            ProgramCacheEntryType::Closed,
2047        );
2048        assert_matches!(tombstone.program, ProgramCacheEntryType::Closed);
2049        assert!(tombstone.is_tombstone());
2050        assert_eq!(tombstone.deployment_slot, 100);
2051        assert_eq!(tombstone.effective_slot, 100);
2052
2053        let mut cache = new_mock_cache();
2054        let program1 = Pubkey::new_unique();
2055        let tombstone = set_tombstone(
2056            &mut cache,
2057            program1,
2058            10,
2059            ProgramCacheEntryType::FailedVerification(env.clone()),
2060        );
2061        let slot_versions = cache.get_slot_versions_for_tests(&program1);
2062        assert_eq!(slot_versions.len(), 1);
2063        assert!(slot_versions.first().unwrap().is_tombstone());
2064        assert_eq!(tombstone.deployment_slot, 10);
2065        assert_eq!(tombstone.effective_slot, 10);
2066
2067        // Add a program at slot 50, and a tombstone for the program at slot 60
2068        let program2 = Pubkey::new_unique();
2069        cache.assign_program(program2, new_test_builtin_entry(50, 51));
2070        let slot_versions = cache.get_slot_versions_for_tests(&program2);
2071        assert_eq!(slot_versions.len(), 1);
2072        assert!(!slot_versions.first().unwrap().is_tombstone());
2073
2074        let tombstone = set_tombstone(
2075            &mut cache,
2076            program2,
2077            60,
2078            ProgramCacheEntryType::FailedVerification(env),
2079        );
2080        let slot_versions = cache.get_slot_versions_for_tests(&program2);
2081        assert_eq!(slot_versions.len(), 2);
2082        assert!(!slot_versions.first().unwrap().is_tombstone());
2083        assert!(slot_versions.get(1).unwrap().is_tombstone());
2084        assert!(tombstone.is_tombstone());
2085        assert_eq!(tombstone.deployment_slot, 60);
2086        assert_eq!(tombstone.effective_slot, 60);
2087    }
2088
2089    #[test]
2090    fn test_prune_empty() {
2091        let mut cache = new_mock_cache();
2092
2093        cache.prune(0, 0);
2094        assert!(cache.get_flattened_entries_for_tests().is_empty());
2095
2096        cache.prune(10, 0);
2097        assert!(cache.get_flattened_entries_for_tests().is_empty());
2098
2099        let mut cache = new_mock_cache();
2100
2101        cache.prune(0, 0);
2102        assert!(cache.get_flattened_entries_for_tests().is_empty());
2103
2104        cache.prune(10, 0);
2105        assert!(cache.get_flattened_entries_for_tests().is_empty());
2106
2107        let mut cache = new_mock_cache();
2108
2109        cache.prune(0, 0);
2110        assert!(cache.get_flattened_entries_for_tests().is_empty());
2111
2112        cache.prune(10, 0);
2113        assert!(cache.get_flattened_entries_for_tests().is_empty());
2114
2115        let mut cache = new_mock_cache();
2116
2117        cache.prune(0, 0);
2118        assert!(cache.get_flattened_entries_for_tests().is_empty());
2119
2120        cache.prune(10, 0);
2121        assert!(cache.get_flattened_entries_for_tests().is_empty());
2122    }
2123
2124    #[test]
2125    fn test_prune_different_env() {
2126        let mut cache = new_mock_cache();
2127
2128        let program1 = Pubkey::new_unique();
2129        cache.assign_program(program1, new_test_entry(10, 10));
2130
2131        let new_env = Arc::new(BuiltinProgram::new_mock());
2132        cache.upcoming_environments = Some(ProgramRuntimeEnvironments {
2133            program_runtime_v1: new_env.clone(),
2134            program_runtime_v2: new_env.clone(),
2135        });
2136        let updated_program = Arc::new(ProgramCacheEntry {
2137            program: new_loaded_entry(new_env.clone()),
2138            account_owner: ProgramCacheEntryOwner::LoaderV2,
2139            account_size: 0,
2140            deployment_slot: 20,
2141            effective_slot: 20,
2142            tx_usage_counter: AtomicU64::default(),
2143            ix_usage_counter: AtomicU64::default(),
2144            latest_access_slot: AtomicU64::default(),
2145        });
2146        cache.assign_program(program1, updated_program.clone());
2147
2148        // Test that there are 2 entries for the program
2149        assert_eq!(cache.get_slot_versions_for_tests(&program1).len(), 2);
2150
2151        cache.prune(21, cache.latest_root_epoch);
2152
2153        // Test that prune didn't remove the entry, since environments are different.
2154        assert_eq!(cache.get_slot_versions_for_tests(&program1).len(), 2);
2155
2156        cache.prune(22, cache.latest_root_epoch.saturating_add(1));
2157
2158        // Test that prune removed 1 entry, since epoch changed
2159        assert_eq!(cache.get_slot_versions_for_tests(&program1).len(), 1);
2160
2161        let entry = cache
2162            .get_slot_versions_for_tests(&program1)
2163            .first()
2164            .expect("Failed to get the program")
2165            .clone();
2166        // Test that the correct entry remains in the cache
2167        assert_eq!(entry, updated_program);
2168    }
2169
2170    fn get_entries_to_load(
2171        cache: &ProgramCache,
2172        loading_slot: Slot,
2173        keys: &[Pubkey],
2174    ) -> Vec<(Pubkey, (ProgramCacheMatchCriteria, u64))> {
2175        let entries = cache.get_flattened_entries_for_tests();
2176        keys.iter()
2177            .filter_map(|key| {
2178                entries
2179                    .iter()
2180                    .rev()
2181                    .find(|(program_id, entry)| {
2182                        program_id == key && entry.deployment_slot <= loading_slot
2183                    })
2184                    .map(|(program_id, _entry)| {
2185                        (*program_id, (ProgramCacheMatchCriteria::NoCriteria, 1))
2186                    })
2187            })
2188            .collect()
2189    }
2190
2191    fn match_slot(
2192        extracted: &ProgramCacheForTxBatch,
2193        program: &Pubkey,
2194        deployment_slot: Slot,
2195        working_slot: Slot,
2196    ) -> bool {
2197        assert_eq!(extracted.slot, working_slot);
2198        extracted
2199            .entries
2200            .get(program)
2201            .map(|entry| entry.deployment_slot == deployment_slot)
2202            .unwrap_or(false)
2203    }
2204
2205    fn match_missing(
2206        missing: &[(Pubkey, (ProgramCacheMatchCriteria, u64))],
2207        program: &Pubkey,
2208        expected_result: bool,
2209    ) -> bool {
2210        missing.iter().any(|(key, _)| key == program) == expected_result
2211    }
2212
2213    #[test]
2214    fn test_fork_extract_and_prune() {
2215        let mut cache = new_mock_cache();
2216
2217        // Fork graph created for the test
2218        //                   0
2219        //                 /   \
2220        //                10    5
2221        //                |     |
2222        //                20    11
2223        //                |     | \
2224        //                22   15  25
2225        //                      |   |
2226        //                     16  27
2227        //                      |
2228        //                     19
2229        //                      |
2230        //                     23
2231        let program1 = Pubkey::new_unique();
2232        cache.assign_program(program1, new_test_entry(0, 1));
2233        cache.assign_program(program1, new_test_entry(10, 11));
2234        cache.assign_program(program1, new_test_entry(20, 21));
2235
2236        let program2 = Pubkey::new_unique();
2237        cache.assign_program(program2, new_test_entry(5, 6));
2238        cache.assign_program(
2239            program2,
2240            new_test_entry(11, 11 + DELAY_VISIBILITY_SLOT_OFFSET),
2241        );
2242
2243        let program3 = Pubkey::new_unique();
2244        cache.assign_program(program3, new_test_entry(25, 26));
2245
2246        let program4 = Pubkey::new_unique();
2247        cache.assign_program(program4, new_test_entry(0, 1));
2248        cache.assign_program(program4, new_test_entry(5, 6));
2249        // The following is a special case, where effective slot is 3 slots in the future
2250        cache.assign_program(
2251            program4,
2252            new_test_entry(15, 15 + DELAY_VISIBILITY_SLOT_OFFSET),
2253        );
2254
2255        // Current fork graph
2256        //                   0
2257        //                 /   \
2258        //                10    5
2259        //                |     |
2260        //                20    11
2261        //                |     | \
2262        //                22   15  25
2263        //                      |   |
2264        //                     16  27
2265        //                      |
2266        //                     19
2267        //                      |
2268        //                     23
2269
2270        // Testing fork 0 - 10 - 12 - 22 with current slot at 22
2271        let mut missing =
2272            get_entries_to_load(&cache, 22, &[program1, program2, program3, program4]);
2273        assert!(match_missing(&missing, &program2, false));
2274        assert!(match_missing(&missing, &program3, false));
2275        let mut extracted = ProgramCacheForTxBatch::new(22, cache.environments.clone(), None, 0);
2276        cache.extract(&mut missing, &mut extracted, true);
2277        assert!(match_slot(&extracted, &program1, 20, 22));
2278        assert!(match_slot(&extracted, &program4, 0, 22));
2279
2280        // Testing fork 0 - 5 - 11 - 15 - 16 with current slot at 15
2281        let mut missing =
2282            get_entries_to_load(&cache, 15, &[program1, program2, program3, program4]);
2283        assert!(match_missing(&missing, &program3, false));
2284        let mut extracted = ProgramCacheForTxBatch::new(15, cache.environments.clone(), None, 0);
2285        cache.extract(&mut missing, &mut extracted, true);
2286        assert!(match_slot(&extracted, &program1, 0, 15));
2287        assert!(match_slot(&extracted, &program2, 11, 15));
2288        // The effective slot of program4 deployed in slot 15 is 19. So it should not be usable in slot 16.
2289        // A delay visibility tombstone should be returned here.
2290        let tombstone = extracted
2291            .find(&program4)
2292            .expect("Failed to find the tombstone");
2293        assert_matches!(tombstone.program, ProgramCacheEntryType::DelayVisibility);
2294        assert_eq!(tombstone.deployment_slot, 15);
2295
2296        // Testing the same fork above, but current slot is now 18 (equal to effective slot of program4).
2297        let mut missing =
2298            get_entries_to_load(&cache, 18, &[program1, program2, program3, program4]);
2299        assert!(match_missing(&missing, &program3, false));
2300        let mut extracted = ProgramCacheForTxBatch::new(18, cache.environments.clone(), None, 0);
2301        cache.extract(&mut missing, &mut extracted, true);
2302        assert!(match_slot(&extracted, &program1, 0, 18));
2303        assert!(match_slot(&extracted, &program2, 11, 18));
2304        // The effective slot of program4 deployed in slot 15 is 18. So it should be usable in slot 18.
2305        assert!(match_slot(&extracted, &program4, 15, 18));
2306
2307        // Testing the same fork above, but current slot is now 23 (future slot than effective slot of program4).
2308        let mut missing =
2309            get_entries_to_load(&cache, 23, &[program1, program2, program3, program4]);
2310        assert!(match_missing(&missing, &program3, false));
2311        let mut extracted = ProgramCacheForTxBatch::new(23, cache.environments.clone(), None, 0);
2312        cache.extract(&mut missing, &mut extracted, true);
2313        assert!(match_slot(&extracted, &program1, 0, 23));
2314        assert!(match_slot(&extracted, &program2, 11, 23));
2315        // The effective slot of program4 deployed in slot 15 is 19. So it should be usable in slot 23.
2316        assert!(match_slot(&extracted, &program4, 15, 23));
2317
2318        // Testing fork 0 - 5 - 11 - 15 - 16 with current slot at 11
2319        let mut missing =
2320            get_entries_to_load(&cache, 11, &[program1, program2, program3, program4]);
2321        assert!(match_missing(&missing, &program3, false));
2322        let mut extracted = ProgramCacheForTxBatch::new(11, cache.environments.clone(), None, 0);
2323        cache.extract(&mut missing, &mut extracted, true);
2324        assert!(match_slot(&extracted, &program1, 0, 11));
2325        // program2 was updated at slot 11, but is not effective till slot 12. The result should contain a tombstone.
2326        let tombstone = extracted
2327            .find(&program2)
2328            .expect("Failed to find the tombstone");
2329        assert_matches!(tombstone.program, ProgramCacheEntryType::DelayVisibility);
2330        assert_eq!(tombstone.deployment_slot, 11);
2331        assert!(match_slot(&extracted, &program4, 5, 11));
2332
2333        cache.prune(5, 0);
2334
2335        // Fork graph after pruning
2336        //                   0
2337        //                   |
2338        //                   5
2339        //                   |
2340        //                   11
2341        //                   | \
2342        //                  15  25
2343        //                   |   |
2344        //                  16  27
2345        //                   |
2346        //                  19
2347        //                   |
2348        //                  23
2349
2350        // Testing fork 11 - 15 - 16- 19 - 22 with root at 5 and current slot at 22
2351        let mut missing =
2352            get_entries_to_load(&cache, 21, &[program1, program2, program3, program4]);
2353        assert!(match_missing(&missing, &program3, false));
2354        let mut extracted = ProgramCacheForTxBatch::new(21, cache.environments.clone(), None, 0);
2355        cache.extract(&mut missing, &mut extracted, true);
2356        // Since the fork was pruned, we should not find the entry deployed at slot 20.
2357        assert!(match_slot(&extracted, &program1, 0, 21));
2358        assert!(match_slot(&extracted, &program2, 11, 21));
2359        assert!(match_slot(&extracted, &program4, 15, 21));
2360
2361        // Testing fork 0 - 5 - 11 - 25 - 27 with current slot at 27
2362        let mut missing =
2363            get_entries_to_load(&cache, 27, &[program1, program2, program3, program4]);
2364        let mut extracted = ProgramCacheForTxBatch::new(27, cache.environments.clone(), None, 0);
2365        cache.extract(&mut missing, &mut extracted, true);
2366        assert!(match_slot(&extracted, &program1, 0, 27));
2367        assert!(match_slot(&extracted, &program2, 11, 27));
2368        assert!(match_slot(&extracted, &program3, 25, 27));
2369        assert!(match_slot(&extracted, &program4, 5, 27));
2370
2371        cache.prune(15, 0);
2372
2373        // Fork graph after pruning
2374        //                  0
2375        //                  |
2376        //                  5
2377        //                  |
2378        //                  11
2379        //                  |
2380        //                  15
2381        //                  |
2382        //                  16
2383        //                  |
2384        //                  19
2385        //                  |
2386        //                  23
2387
2388        // Testing fork 16, 19, 23, with root at 15, current slot at 23
2389        let mut missing =
2390            get_entries_to_load(&cache, 23, &[program1, program2, program3, program4]);
2391        assert!(match_missing(&missing, &program3, false));
2392        let mut extracted = ProgramCacheForTxBatch::new(23, cache.environments.clone(), None, 0);
2393        cache.extract(&mut missing, &mut extracted, true);
2394        assert!(match_slot(&extracted, &program1, 0, 23));
2395        assert!(match_slot(&extracted, &program2, 11, 23));
2396        assert!(match_slot(&extracted, &program4, 15, 23));
2397    }
2398
2399    #[test]
2400    fn test_extract_using_deployment_slot() {
2401        let mut cache = new_mock_cache();
2402
2403        // Fork graph created for the test
2404        //                   0
2405        //                 /   \
2406        //                10    5
2407        //                |     |
2408        //                20    11
2409        //                |     | \
2410        //                22   15  25
2411        //                      |   |
2412        //                     16  27
2413        //                      |
2414        //                     19
2415        //                      |
2416        //                     23
2417        let program1 = Pubkey::new_unique();
2418        cache.assign_program(program1, new_test_entry(0, 1));
2419        cache.assign_program(program1, new_test_entry(20, 21));
2420
2421        let program2 = Pubkey::new_unique();
2422        cache.assign_program(program2, new_test_entry(5, 6));
2423        cache.assign_program(program2, new_test_entry(11, 12));
2424
2425        let program3 = Pubkey::new_unique();
2426        cache.assign_program(program3, new_test_entry(25, 26));
2427
2428        // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 19
2429        let mut missing = get_entries_to_load(&cache, 12, &[program1, program2, program3]);
2430        assert!(match_missing(&missing, &program3, false));
2431        let mut extracted = ProgramCacheForTxBatch::new(12, cache.environments.clone(), None, 0);
2432        cache.extract(&mut missing, &mut extracted, true);
2433        assert!(match_slot(&extracted, &program1, 0, 12));
2434        assert!(match_slot(&extracted, &program2, 11, 12));
2435
2436        // Test the same fork, but request the program modified at a later slot than what's in the cache.
2437        let mut missing = get_entries_to_load(&cache, 12, &[program1, program2, program3]);
2438        missing.get_mut(0).unwrap().1 .0 = ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(5);
2439        missing.get_mut(1).unwrap().1 .0 = ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(5);
2440        assert!(match_missing(&missing, &program3, false));
2441        let mut extracted = ProgramCacheForTxBatch::new(12, cache.environments.clone(), None, 0);
2442        cache.extract(&mut missing, &mut extracted, true);
2443        assert!(match_missing(&missing, &program1, true));
2444        assert!(match_slot(&extracted, &program2, 11, 12));
2445    }
2446
2447    #[test]
2448    fn test_extract_unloaded() {
2449        let mut cache = new_mock_cache();
2450
2451        // Fork graph created for the test
2452        //                   0
2453        //                 /   \
2454        //                10    5
2455        //                |     |
2456        //                20    11
2457        //                |     | \
2458        //                22   15  25
2459        //                      |   |
2460        //                     16  27
2461        //                      |
2462        //                     19
2463        //                      |
2464        //                     23
2465
2466        let program1 = Pubkey::new_unique();
2467        cache.assign_program(program1, new_test_entry(0, 1));
2468        cache.assign_program(program1, new_test_entry(20, 21));
2469
2470        let program2 = Pubkey::new_unique();
2471        cache.assign_program(program2, new_test_entry(5, 6));
2472        cache.assign_program(program2, new_test_entry(11, 12));
2473
2474        let program3 = Pubkey::new_unique();
2475        // Insert an unloaded program with correct/cache's environment at slot 25
2476        let _ = insert_unloaded_entry(&mut cache, program3, 25);
2477
2478        // Insert another unloaded program with a different environment at slot 20
2479        // Since this entry's environment won't match cache's environment, looking up this
2480        // entry should return missing instead of unloaded entry.
2481        cache.assign_program(
2482            program3,
2483            Arc::new(
2484                new_test_entry(20, 21)
2485                    .to_unloaded()
2486                    .expect("Failed to create unloaded program"),
2487            ),
2488        );
2489
2490        // Testing fork 0 - 5 - 11 - 15 - 16 - 19 - 21 - 23 with current slot at 19
2491        let mut missing = get_entries_to_load(&cache, 19, &[program1, program2, program3]);
2492        assert!(match_missing(&missing, &program3, false));
2493        let mut extracted = ProgramCacheForTxBatch::new(19, cache.environments.clone(), None, 0);
2494        cache.extract(&mut missing, &mut extracted, true);
2495        assert!(match_slot(&extracted, &program1, 0, 19));
2496        assert!(match_slot(&extracted, &program2, 11, 19));
2497
2498        // Testing fork 0 - 5 - 11 - 25 - 27 with current slot at 27
2499        let mut missing = get_entries_to_load(&cache, 27, &[program1, program2, program3]);
2500        let mut extracted = ProgramCacheForTxBatch::new(27, cache.environments.clone(), None, 0);
2501        cache.extract(&mut missing, &mut extracted, true);
2502        assert!(match_slot(&extracted, &program1, 0, 27));
2503        assert!(match_slot(&extracted, &program2, 11, 27));
2504        assert!(match_missing(&missing, &program3, true));
2505
2506        // Testing fork 0 - 10 - 20 - 22 with current slot at 22
2507        let mut missing = get_entries_to_load(&cache, 22, &[program1, program2, program3]);
2508        assert!(match_missing(&missing, &program2, false));
2509        let mut extracted = ProgramCacheForTxBatch::new(22, cache.environments.clone(), None, 0);
2510        cache.extract(&mut missing, &mut extracted, true);
2511        assert!(match_slot(&extracted, &program1, 20, 22));
2512        assert!(match_missing(&missing, &program3, true));
2513    }
2514
2515    #[test]
2516    fn test_extract_nonexistent() {
2517        let cache = new_mock_cache();
2518
2519        let program1 = Pubkey::new_unique();
2520        let mut missing = vec![(program1, (ProgramCacheMatchCriteria::NoCriteria, 1))];
2521        let mut extracted = ProgramCacheForTxBatch::new(0, cache.environments.clone(), None, 0);
2522        cache.extract(&mut missing, &mut extracted, true);
2523        assert!(match_missing(&missing, &program1, true));
2524    }
2525
2526    #[test]
2527    fn test_unloaded() {
2528        let mut cache = new_mock_cache();
2529        for program_cache_entry_type in [
2530            ProgramCacheEntryType::FailedVerification(
2531                cache.environments.program_runtime_v1.clone(),
2532            ),
2533            ProgramCacheEntryType::Closed,
2534            ProgramCacheEntryType::Unloaded(cache.environments.program_runtime_v1.clone()),
2535            ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
2536        ] {
2537            let entry = Arc::new(ProgramCacheEntry {
2538                program: program_cache_entry_type,
2539                account_owner: ProgramCacheEntryOwner::LoaderV2,
2540                account_size: 0,
2541                deployment_slot: 0,
2542                effective_slot: 0,
2543                tx_usage_counter: AtomicU64::default(),
2544                ix_usage_counter: AtomicU64::default(),
2545                latest_access_slot: AtomicU64::default(),
2546            });
2547            assert!(entry.to_unloaded().is_none());
2548
2549            // Check that unload_program_entry() does nothing for this entry
2550            let program_id = Pubkey::new_unique();
2551            cache.assign_program(program_id, entry.clone());
2552            cache.unload_program_entry(&program_id, &entry);
2553            assert_eq!(cache.get_slot_versions_for_tests(&program_id).len(), 1);
2554            assert!(cache.stats.evictions.is_empty());
2555        }
2556
2557        let entry = new_test_entry_with_usage(1, 2, AtomicU64::new(3));
2558        let unloaded_entry = entry.to_unloaded().unwrap();
2559        assert_eq!(unloaded_entry.deployment_slot, 1);
2560        assert_eq!(unloaded_entry.effective_slot, 2);
2561        assert_eq!(unloaded_entry.latest_access_slot.load(Ordering::Relaxed), 1);
2562        assert_eq!(unloaded_entry.tx_usage_counter.load(Ordering::Relaxed), 3);
2563
2564        // Check that unload_program_entry() does its work
2565        let program_id = Pubkey::new_unique();
2566        cache.assign_program(program_id, entry.clone());
2567        cache.unload_program_entry(&program_id, &entry);
2568        assert!(cache.stats.evictions.contains_key(&program_id));
2569    }
2570
2571    #[test]
2572    fn test_fork_prune_find_first_ancestor() {
2573        let mut cache = new_mock_cache();
2574
2575        // Fork graph created for the test
2576        //                   0
2577        //                 /   \
2578        //                10    5
2579        //                |
2580        //                20
2581
2582        // Deploy program on slot 0, and slot 5.
2583        // Prune the fork that has slot 5. The cache should still have the program
2584        // deployed at slot 0.
2585        let program1 = Pubkey::new_unique();
2586        cache.assign_program(program1, new_test_entry(0, 1));
2587        cache.assign_program(program1, new_test_entry(5, 6));
2588
2589        cache.prune(10, 0);
2590
2591        let mut missing = get_entries_to_load(&cache, 20, &[program1]);
2592        let mut extracted = ProgramCacheForTxBatch::new(20, cache.environments.clone(), None, 0);
2593        cache.extract(&mut missing, &mut extracted, true);
2594
2595        // The cache should have the program deployed at slot 0
2596        assert_eq!(
2597            extracted
2598                .find(&program1)
2599                .expect("Did not find the program")
2600                .deployment_slot,
2601            0
2602        );
2603    }
2604
2605    #[test]
2606    fn test_prune_by_deployment_slot() {
2607        let mut cache = new_mock_cache();
2608
2609        // Fork graph created for the test
2610        //                   0
2611        //                 /   \
2612        //                10    5
2613        //                |
2614        //                20
2615
2616        // Deploy program on slot 0, and slot 5.
2617        // Prune the fork that has slot 5. The cache should still have the program
2618        // deployed at slot 0.
2619        let program1 = Pubkey::new_unique();
2620        cache.assign_program(program1, new_test_entry(0, 1));
2621        cache.assign_program(program1, new_test_entry(5, 6));
2622
2623        let program2 = Pubkey::new_unique();
2624        cache.assign_program(program2, new_test_entry(10, 11));
2625
2626        let mut missing = get_entries_to_load(&cache, 20, &[program1, program2]);
2627        let mut extracted = ProgramCacheForTxBatch::new(20, cache.environments.clone(), None, 0);
2628        cache.extract(&mut missing, &mut extracted, true);
2629        assert!(match_slot(&extracted, &program1, 0, 20));
2630        assert!(match_slot(&extracted, &program2, 10, 20));
2631
2632        let mut missing = get_entries_to_load(&cache, 6, &[program1, program2]);
2633        assert!(match_missing(&missing, &program2, false));
2634        let mut extracted = ProgramCacheForTxBatch::new(6, cache.environments.clone(), None, 0);
2635        cache.extract(&mut missing, &mut extracted, true);
2636        assert!(match_slot(&extracted, &program1, 5, 6));
2637
2638        // Pruning slot 5 will remove program1 entry deployed at slot 5.
2639        // On fork chaining from slot 5, the entry deployed at slot 0 will become visible.
2640        cache.prune_by_deployment_slot(5);
2641
2642        let mut missing = get_entries_to_load(&cache, 20, &[program1, program2]);
2643        let mut extracted = ProgramCacheForTxBatch::new(20, cache.environments.clone(), None, 0);
2644        cache.extract(&mut missing, &mut extracted, true);
2645        assert!(match_slot(&extracted, &program1, 0, 20));
2646        assert!(match_slot(&extracted, &program2, 10, 20));
2647
2648        let mut missing = get_entries_to_load(&cache, 6, &[program1, program2]);
2649        assert!(match_missing(&missing, &program2, false));
2650        let mut extracted = ProgramCacheForTxBatch::new(6, cache.environments.clone(), None, 0);
2651        cache.extract(&mut missing, &mut extracted, true);
2652        assert!(match_slot(&extracted, &program1, 0, 6));
2653
2654        // Pruning slot 10 will remove program2 entry deployed at slot 10.
2655        // As there is no other entry for program2, extract() will return it as missing.
2656        cache.prune_by_deployment_slot(10);
2657
2658        let mut missing = get_entries_to_load(&cache, 20, &[program1, program2]);
2659        assert!(match_missing(&missing, &program2, false));
2660        let mut extracted = ProgramCacheForTxBatch::new(20, cache.environments.clone(), None, 0);
2661        cache.extract(&mut missing, &mut extracted, true);
2662        assert!(match_slot(&extracted, &program1, 0, 20));
2663    }
2664
2665    #[test]
2666    fn test_usable_entries_for_slot() {
2667        new_mock_cache();
2668        let tombstone = Arc::new(ProgramCacheEntry::new_tombstone(
2669            0,
2670            ProgramCacheEntryOwner::LoaderV2,
2671            ProgramCacheEntryType::Closed,
2672        ));
2673
2674        assert!(ProgramCache::matches_criteria(
2675            &tombstone,
2676            &ProgramCacheMatchCriteria::NoCriteria
2677        ));
2678
2679        assert!(ProgramCache::matches_criteria(
2680            &tombstone,
2681            &ProgramCacheMatchCriteria::Tombstone
2682        ));
2683
2684        assert!(ProgramCache::matches_criteria(
2685            &tombstone,
2686            &ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(0)
2687        ));
2688
2689        assert!(!ProgramCache::matches_criteria(
2690            &tombstone,
2691            &ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(1)
2692        ));
2693
2694        let program = new_test_entry(0, 1);
2695
2696        assert!(ProgramCache::matches_criteria(
2697            &program,
2698            &ProgramCacheMatchCriteria::NoCriteria
2699        ));
2700
2701        assert!(!ProgramCache::matches_criteria(
2702            &program,
2703            &ProgramCacheMatchCriteria::Tombstone
2704        ));
2705
2706        assert!(ProgramCache::matches_criteria(
2707            &program,
2708            &ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(0)
2709        ));
2710
2711        assert!(!ProgramCache::matches_criteria(
2712            &program,
2713            &ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(1)
2714        ));
2715
2716        let program = Arc::new(new_test_entry_with_usage(0, 1, AtomicU64::default()));
2717
2718        assert!(ProgramCache::matches_criteria(
2719            &program,
2720            &ProgramCacheMatchCriteria::NoCriteria
2721        ));
2722
2723        assert!(!ProgramCache::matches_criteria(
2724            &program,
2725            &ProgramCacheMatchCriteria::Tombstone
2726        ));
2727
2728        assert!(ProgramCache::matches_criteria(
2729            &program,
2730            &ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(0)
2731        ));
2732
2733        assert!(!ProgramCache::matches_criteria(
2734            &program,
2735            &ProgramCacheMatchCriteria::DeployedOnOrAfterSlot(1)
2736        ));
2737    }
2738}