Add Shuttle multithreading test infrastructure (#1634)

anza-xyz · Jun 17, 2024 · ebd063e · ebd063e
1 parent 2e54ed2
commit ebd063e
Show file tree

Hide file tree

Showing 22 changed files with 320 additions and 63 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -116,6 +116,7 @@ members = [
     "transaction-metrics-tracker",
     "transaction-status",
     "turbine",
+    "type-overrides",
     "udp-client",
     "unified-scheduler-logic",
     "unified-scheduler-pool",
@@ -309,6 +310,7 @@ serde_yaml = "0.9.34"
 serial_test = "2.0.0"
 sha2 = "0.10.8"
 sha3 = "0.10.8"
+shuttle = "0.7.1"
 signal-hook = "0.3.17"
 siphasher = "0.3.11"
 smallvec = "1.13.2"
@@ -394,6 +396,7 @@ solana-tpu-client = { path = "tpu-client", version = "=2.0.0", default-features
 solana-transaction-status = { path = "transaction-status", version = "=2.0.0" }
 solana-transaction-metrics-tracker = { path = "transaction-metrics-tracker", version = "=2.0.0" }
 solana-turbine = { path = "turbine", version = "=2.0.0" }
+solana-type-overrides = { path = "type-overrides", version = "=2.0.0" }
 solana-udp-client = { path = "udp-client", version = "=2.0.0" }
 solana-version = { path = "version", version = "=2.0.0" }
 solana-vote = { path = "vote", version = "=2.0.0" }

diff --git a/ledger-tool/Cargo.toml b/ledger-tool/Cargo.toml
@@ -48,6 +48,7 @@ solana-stake-program = { workspace = true }
 solana-storage-bigtable = { workspace = true }
 solana-streamer = { workspace = true }
 solana-transaction-status = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana-unified-scheduler-pool = { workspace = true }
 solana-version = { workspace = true }
 solana-vote-program = { workspace = true }

diff --git a/program-runtime/Cargo.toml b/program-runtime/Cargo.toml
@@ -28,6 +28,7 @@ solana-frozen-abi-macro = { workspace = true, optional = true }
 solana-measure = { workspace = true }
 solana-metrics = { workspace = true }
 solana-sdk = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana-vote = { workspace = true }
 solana_rbpf = { workspace = true }
 thiserror = { workspace = true }
@@ -55,3 +56,4 @@ frozen-abi = [
     "solana-compute-budget/frozen-abi",
     "solana-sdk/frozen-abi",
 ]
+shuttle-test = ["solana-type-overrides/shuttle-test"]
diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs
@@ -36,13 +36,13 @@ use {
             IndexOfAccount, InstructionAccount, TransactionAccount, TransactionContext,
         },
     },
+    solana_type_overrides::sync::{atomic::Ordering, Arc},
     solana_vote::vote_account::VoteAccountsHashMap,
     std::{
         alloc::Layout,
         cell::RefCell,
         fmt::{self, Debug},
         rc::Rc,
-        sync::{atomic::Ordering, Arc},
     },
 };
 
@@ -681,7 +681,7 @@ macro_rules! with_mock_invoke_context {
                 account::ReadableAccount, feature_set::FeatureSet, hash::Hash, sysvar::rent::Rent,
                 transaction_context::TransactionContext,
             },
-            std::sync::Arc,
+            solana_type_overrides::sync::Arc,
             $crate::{
                 invoke_context::{EnvironmentConfig, InvokeContext},
                 loaded_programs::ProgramCacheForTxBatch,

diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
@@ -5,7 +5,6 @@ use {
     },
     log::{debug, error, log_enabled, trace},
     percentage::PercentageInteger,
-    rand::{thread_rng, Rng},
     solana_measure::measure::Measure,
     solana_rbpf::{
         elf::Executable,
@@ -20,13 +19,17 @@ use {
         pubkey::Pubkey,
         saturating_add_assign,
     },
-    std::{
-        collections::{hash_map::Entry, HashMap},
-        fmt::{Debug, Formatter},
+    solana_type_overrides::{
+        rand::{thread_rng, Rng},
         sync::{
             atomic::{AtomicU64, Ordering},
             Arc, Condvar, Mutex, RwLock,
         },
+        thread,
+    },
+    std::{
+        collections::{hash_map::Entry, HashMap},
+        fmt::{Debug, Formatter},
     },
 };
 
@@ -598,7 +601,7 @@ enum IndexImplementation {
         /// It is possible that multiple TX batches from different slots need different versions of a
         /// program. The deployment slot of a program is only known after load tho,
         /// so all loads for a given program key are serialized.
-        loading_entries: Mutex<HashMap<Pubkey, (Slot, std::thread::ThreadId)>>,
+        loading_entries: Mutex<HashMap<Pubkey, (Slot, thread::ThreadId)>>,
     },
 }
 
@@ -1122,7 +1125,7 @@ impl<FG: ForkGraph> ProgramCache<FG> {
                         if let Entry::Vacant(entry) = entry {
                             entry.insert((
                                 loaded_programs_for_tx_batch.slot,
-                                std::thread::current().id(),
+                                thread::current().id(),
                             ));
                             cooperative_loading_task = Some((*key, *usage_count));
                         }
@@ -1156,7 +1159,7 @@ impl<FG: ForkGraph> ProgramCache<FG> {
                 loading_entries, ..
             } => {
                 let loading_thread = loading_entries.get_mut().unwrap().remove(&key);
-                debug_assert_eq!(loading_thread, Some((slot, std::thread::current().id())));
+                debug_assert_eq!(loading_thread, Some((slot, thread::current().id())));
                 // Check that it will be visible to our own fork once inserted
                 if loaded_program.deployment_slot > self.latest_root_slot
                     && !matches!(

diff --git a/program-runtime/src/sysvar_cache.rs b/program-runtime/src/sysvar_cache.rs
@@ -13,7 +13,7 @@ use {
         },
         transaction_context::{IndexOfAccount, InstructionContext, TransactionContext},
     },
-    std::sync::Arc,
+    solana_type_overrides::sync::Arc,
 };
 
 #[cfg(all(RUSTC_WITH_SPECIALIZATION, feature = "frozen-abi"))]

diff --git a/programs/bpf_loader/Cargo.toml b/programs/bpf_loader/Cargo.toml
@@ -20,6 +20,7 @@ solana-measure = { workspace = true }
 solana-poseidon = { workspace = true }
 solana-program-runtime = { workspace = true }
 solana-sdk = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana-zk-token-sdk = { workspace = true }
 solana_rbpf = { workspace = true }
 thiserror = { workspace = true }
@@ -38,3 +39,6 @@ name = "solana_bpf_loader_program"
 
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]
+
+[features]
+shuttle-test = ["solana-type-overrides/shuttle-test", "solana-program-runtime/shuttle-test"]
diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
@@ -46,12 +46,8 @@ use {
         system_instruction::{self, MAX_PERMITTED_DATA_LENGTH},
         transaction_context::{IndexOfAccount, InstructionContext, TransactionContext},
     },
-    std::{
-        cell::RefCell,
-        mem,
-        rc::Rc,
-        sync::{atomic::Ordering, Arc},
-    },
+    solana_type_overrides::sync::{atomic::Ordering, Arc},
+    std::{cell::RefCell, mem, rc::Rc},
     syscalls::{create_program_runtime_environment_v1, morph_into_deployment_environment_v1},
 };
 
@@ -324,7 +320,7 @@ macro_rules! create_vm {
 #[macro_export]
 macro_rules! mock_create_vm {
     ($vm:ident, $additional_regions:expr, $accounts_metadata:expr, $invoke_context:expr $(,)?) => {
-        let loader = std::sync::Arc::new(BuiltinProgram::new_mock());
+        let loader = solana_type_overrides::sync::Arc::new(BuiltinProgram::new_mock());
         let function_registry = solana_rbpf::program::FunctionRegistry::default();
         let executable = solana_rbpf::elf::Executable::<InvokeContext>::from_text_bytes(
             &[0x95, 0, 0, 0, 0, 0, 0, 0],

diff --git a/programs/bpf_loader/src/syscalls/mod.rs b/programs/bpf_loader/src/syscalls/mod.rs
@@ -58,12 +58,12 @@ use {
         sysvar::{Sysvar, SysvarId},
         transaction_context::{IndexOfAccount, InstructionAccount},
     },
+    solana_type_overrides::sync::Arc,
     std::{
         alloc::Layout,
         mem::{align_of, size_of},
         slice::from_raw_parts_mut,
         str::{from_utf8, Utf8Error},
-        sync::Arc,
     },
     thiserror::Error as ThisError,
 };

diff --git a/programs/loader-v4/Cargo.toml b/programs/loader-v4/Cargo.toml
@@ -14,6 +14,7 @@ solana-compute-budget = { workspace = true }
 solana-measure = { workspace = true }
 solana-program-runtime = { workspace = true }
 solana-sdk = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana_rbpf = { workspace = true }
 
 [dev-dependencies]
@@ -25,3 +26,6 @@ name = "solana_loader_v4_program"
 
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]
+
+[features]
+shuttle-test = ["solana-type-overrides/shuttle-test", "solana-program-runtime/shuttle-test"]
diff --git a/programs/loader-v4/src/lib.rs b/programs/loader-v4/src/lib.rs
@@ -30,11 +30,8 @@ use {
         saturating_add_assign,
         transaction_context::{BorrowedAccount, InstructionContext},
     },
-    std::{
-        cell::RefCell,
-        rc::Rc,
-        sync::{atomic::Ordering, Arc},
-    },
+    solana_type_overrides::sync::{atomic::Ordering, Arc},
+    std::{cell::RefCell, rc::Rc},
 };
 
 pub const DEFAULT_COMPUTE_UNITS: u64 = 2_000;

diff --git a/programs/sbf/Cargo.lock b/programs/sbf/Cargo.lock
diff --git a/programs/sbf/Cargo.toml b/programs/sbf/Cargo.toml
@@ -55,6 +55,7 @@ solana-sbf-rust-sysvar = { path = "rust/sysvar", version = "=2.0.0" }
 solana-sdk = { path = "../../sdk", version = "=2.0.0" }
 solana-svm = { path = "../../svm", version = "=2.0.0" }
 solana-transaction-status = { path = "../../transaction-status", version = "=2.0.0" }
+solana-type-overrides = { path = "../../type-overrides", version = "=2.0.0" }
 agave-validator = { path = "../../validator", version = "=2.0.0" }
 solana-zk-token-sdk = { path = "../../zk-token-sdk", version = "=2.0.0" }
 solana_rbpf = "=0.8.1"
@@ -113,6 +114,7 @@ solana-sbf-rust-sysvar = { workspace = true }
 solana-sdk = { workspace = true, features = ["dev-context-only-utils"] }
 solana-svm = { workspace = true }
 solana-transaction-status = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana_rbpf = { workspace = true }
 
 [[bench]]

diff --git a/programs/stake/Cargo.toml b/programs/stake/Cargo.toml
@@ -15,6 +15,7 @@ log = { workspace = true }
 solana-config-program = { workspace = true }
 solana-program-runtime = { workspace = true }
 solana-sdk = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana-vote-program = { workspace = true }
 
 [dev-dependencies]

diff --git a/programs/system/Cargo.toml b/programs/system/Cargo.toml
@@ -16,6 +16,7 @@ serde = { workspace = true }
 serde_derive = { workspace = true }
 solana-program-runtime = { workspace = true }
 solana-sdk = { workspace = true }
+solana-type-overrides = { workspace = true }
 
 [dev-dependencies]
 assert_matches = { workspace = true }

diff --git a/svm/Cargo.toml b/svm/Cargo.toml
@@ -26,6 +26,7 @@ solana-metrics = { workspace = true }
 solana-program-runtime = { workspace = true }
 solana-sdk = { workspace = true }
 solana-system-program = { workspace = true }
+solana-type-overrides = { workspace = true }
 solana-vote = { workspace = true }
 
 [lib]
@@ -62,3 +63,9 @@ frozen-abi = [
     "solana-program-runtime/frozen-abi",
     "solana-sdk/frozen-abi",
 ]
+shuttle-test = [
+    "solana-type-overrides/shuttle-test",
+    "solana-program-runtime/shuttle-test",
+    "solana-bpf-loader-program/shuttle-test",
+    "solana-loader-v4-program/shuttle-test",
+]
diff --git a/svm/src/message_processor.rs b/svm/src/message_processor.rs
@@ -657,7 +657,7 @@ mod tests {
         // copies the `random` implementation at:
         // https://docs.rs/libsecp256k1/latest/src/libsecp256k1/lib.rs.html#430
         let secret_key = {
-            use rand::RngCore;
+            use solana_type_overrides::rand::RngCore;
             let mut rng = rand::thread_rng();
             loop {
                 let mut ret = [0u8; libsecp256k1::util::SECRET_KEY_SIZE];

diff --git a/svm/src/program_loader.rs b/svm/src/program_loader.rs
@@ -17,7 +17,7 @@ use {
         loader_v4::{self, LoaderV4State, LoaderV4Status},
         pubkey::Pubkey,
     },
-    std::sync::Arc,
+    solana_type_overrides::sync::Arc,
 };
 
 #[derive(Debug)]

diff --git a/svm/src/transaction_processor.rs b/svm/src/transaction_processor.rs
@@ -53,16 +53,13 @@ use {
         transaction::{self, SanitizedTransaction, TransactionError},
         transaction_context::{ExecutionRecord, TransactionContext},
     },
+    solana_type_overrides::sync::{atomic::Ordering, Arc, RwLock},
     solana_vote::vote_account::VoteAccountsHashMap,
     std::{
         cell::RefCell,
         collections::{hash_map::Entry, HashMap, HashSet},
         fmt::{Debug, Formatter},
         rc::Rc,
-        sync::{
-            atomic::Ordering::{self, Relaxed},
-            Arc, RwLock,
-        },
     },
 };
 
@@ -649,12 +646,18 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
                     self.slot,
                     false,
                 ) {
-                    recompiled
-                        .tx_usage_counter
-                        .fetch_add(program_to_recompile.tx_usage_counter.load(Relaxed), Relaxed);
-                    recompiled
-                        .ix_usage_counter
-                        .fetch_add(program_to_recompile.ix_usage_counter.load(Relaxed), Relaxed);
+                    recompiled.tx_usage_counter.fetch_add(
+                        program_to_recompile
+                            .tx_usage_counter
+                            .load(Ordering::Relaxed),
+                        Ordering::Relaxed,
+                    );
+                    recompiled.ix_usage_counter.fetch_add(
+                        program_to_recompile
+                            .ix_usage_counter
+                            .load(Ordering::Relaxed),
+                        Ordering::Relaxed,
+                    );
                     let mut program_cache = self.program_cache.write().unwrap();
                     program_cache.assign_program(key, recompiled);
                 }