From 8102f398d4483d46615da67b5c6a03748b919bef Mon Sep 17 00:00:00 2001 From: Dylan Wolff Date: Mon, 6 Oct 2025 19:34:29 +0000 Subject: [PATCH] Adding global config for Shuttle --- shuttle/Cargo.toml | 1 + shuttle/shuttle_config_example.toml | 62 ++++++++ shuttle/src/lib.rs | 208 +++++++++++++++++++++++++- shuttle/src/scheduler/dfs.rs | 7 + shuttle/src/scheduler/mod.rs | 2 + shuttle/src/scheduler/pct.rs | 11 ++ shuttle/src/scheduler/random.rs | 10 ++ shuttle/src/scheduler/registry.rs | 98 ++++++++++++ shuttle/src/scheduler/replay.rs | 12 ++ shuttle/src/scheduler/round_robin.rs | 6 + shuttle/src/scheduler/urw.rs | 10 ++ shuttle/tests/basic/execution.rs | 8 +- shuttle/tests/basic/rwlock.rs | 8 +- shuttle/tests/config/mod.rs | 204 +++++++++++++++++++++++++ shuttle/tests/config/shuttle_dfs.toml | 5 + shuttle/tests/mod.rs | 1 + 16 files changed, 641 insertions(+), 12 deletions(-) create mode 100644 shuttle/shuttle_config_example.toml create mode 100644 shuttle/src/scheduler/registry.rs create mode 100644 shuttle/tests/config/mod.rs create mode 100644 shuttle/tests/config/shuttle_dfs.toml diff --git a/shuttle/Cargo.toml b/shuttle/Cargo.toml index e01acb4a..b90a0fbd 100644 --- a/shuttle/Cargo.toml +++ b/shuttle/Cargo.toml @@ -27,6 +27,7 @@ regex = { version = "1.10.6", optional = true } serde = { version = "1.0", features = ["derive"], optional = true } serde_json = { version = "1.0", optional = true } const-siphasher = "1.0.2" +config = { version = "0.15.18", default-features = false, features = ["toml"] } [dev-dependencies] criterion = { version = "0.4.0", features = ["html_reports"] } diff --git a/shuttle/shuttle_config_example.toml b/shuttle/shuttle_config_example.toml new file mode 100644 index 00000000..dc8b213f --- /dev/null +++ b/shuttle/shuttle_config_example.toml @@ -0,0 +1,62 @@ +# Shuttle Configuration Example +# This file demonstrates available configuration options for Shuttle + +# Stack size for each thread (in bytes) +stack_size = 32768 + +# How to persist failing schedules: "none", "print", or "file" +failure_persistence = "print" + +# Optional: Directory path for file persistence (only used when failure_persistence = "file") +# failure_persistence_path = "./shuttle_failures" + +# Maximum number of steps before taking action (0 means no limit when max_steps_behavior = "none") +max_steps = 1000000 + +# What to do when max_steps is reached: "fail", "continue", or "none" +max_steps_behavior = "fail" + +# Optional: Maximum time in seconds for a single test iteration +# max_time_secs = 60 + +# Suppress warning messages +silence_warnings = false + +# Record execution steps in tracing spans +record_steps_in_span = false + +# Return immediately when a panic occurs (vs continuing to explore other schedules) +immediately_return_on_panic = false + +# Enable metrics collection +enable_metrics = false + +# Check for uncontrolled nondeterminism +check_uncontrolled_nondeterminism = false + +# Scheduler configuration +[scheduler] +# Scheduler type: "random", "pct", "dfs", "replay", "round_robin", or "urw" +# Additional schedulers can be registered using the scheduler registry +type = "random" + +# Number of iterations to run +iterations = 100 + +# Optional: Random seed for reproducible runs +seed = 42 + +# PCT-specific: Number of priority change points +depth = 3 + +# DFS-specific: Maximum iterations before stopping +max_iterations = 10000 + +# DFS-specific: Allow random data generation +allow_random_data = false + +# Replay-specific: Path to schedule file +schedule_file = "./schedule.json" + +# Replay-specific: Inline schedule string +schedule = "..." diff --git a/shuttle/src/lib.rs b/shuttle/src/lib.rs index 2a932428..450607d1 100644 --- a/shuttle/src/lib.rs +++ b/shuttle/src/lib.rs @@ -193,10 +193,14 @@ pub mod scheduler; mod runtime; +use std::path::Path; +use std::path::PathBuf; + pub use runtime::runner::{PortfolioRunner, Runner}; +pub use scheduler::registry::{register_scheduler, SchedulerFactory}; /// Configuration parameters for Shuttle -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Eq, PartialEq)] #[non_exhaustive] pub struct Config { /// Stack size allocated for each thread @@ -265,6 +269,43 @@ impl Config { immediately_return_on_panic: false, } } + + /// Create Config for a single Shuttle test from a global config::Config + pub fn from_global_config(settings: &config::Config) -> Self { + Self::try_from_global_config(settings).expect("Failed to load configuration") + } + + fn try_from_global_config(global_config: &config::Config) -> Result { + let stack_size = global_config.get_int("stack_size")? as usize; + let failure_persistence = FailurePersistence::variant_from_string( + &global_config.get_string("failure_persistence")?, + global_config + .get_string("failure_persistence_path") + .ok() + .map(PathBuf::from), + ); + let max_steps = MaxSteps::variant_from_string( + &global_config.get_string("max_steps_behavior")?, + global_config.get_int("max_steps")? as usize, + ); + let max_time = global_config + .get_int("max_time_secs") + .ok() + .map(|s| std::time::Duration::from_secs(s as u64)); + let silence_warnings = global_config.get_bool("silence_warnings")?; + let record_steps_in_span = global_config.get_bool("record_steps_in_span")?; + let immediately_return_on_panic = global_config.get_bool("immediately_return_on_panic")?; + + Ok(Self { + stack_size, + failure_persistence, + max_steps, + max_time, + silence_warnings, + record_steps_in_span, + immediately_return_on_panic, + }) + } } impl Default for Config { @@ -290,6 +331,25 @@ pub enum FailurePersistence { File(Option), } +impl FailurePersistence { + fn variant_to_string(&self) -> &str { + match self { + FailurePersistence::None => "none", + FailurePersistence::Print => "print", + FailurePersistence::File(_) => "file", + } + } + + fn variant_from_string(s: &str, path: Option) -> Self { + match s { + "none" => FailurePersistence::None, + "file" => FailurePersistence::File(path), + "print" => FailurePersistence::Print, + _ => panic!("Unexpected failure_persistence: {s}"), + } + } +} + /// Specifies an upper bound on the number of steps a single iteration of a Shuttle test can take, /// and how to react when the bound is reached. /// @@ -319,19 +379,155 @@ pub enum MaxSteps { ContinueAfter(usize), } -/// Run the given function once under a round-robin concurrency scheduler. -// TODO consider removing this -- round robin scheduling is never what you want. -#[doc(hidden)] +impl MaxSteps { + fn variant_to_string(&self) -> &str { + match self { + MaxSteps::None => "none", + MaxSteps::FailAfter(_) => "fail", + MaxSteps::ContinueAfter(_) => "continue", + } + } + + fn variant_from_string(s: &str, value: usize) -> Self { + match s { + "none" => MaxSteps::None, + "fail" => MaxSteps::FailAfter(value), + "continue" => MaxSteps::ContinueAfter(value), + _ => panic!("Unexpected max_steps_behavior: {s}"), + } + } +} + +fn should_check_uncontrolled_nondeterminism(global_config: &config::Config) -> bool { + global_config.get_bool("check_uncontrolled_nondeterminism").unwrap() +} + +fn should_enable_metrics(global_config: &config::Config) -> bool { + global_config.get_bool("enable_metrics").unwrap() +} + +/// Run the given function once under the globally configured Shuttle scheduler. +/// +/// # Configuration Sources +/// +/// Configuration is loaded in the following order (later sources override earlier ones): +/// 1. Default values from [`Config::default`] +/// 2. TOML file (path determined by `SHUTTLE_CONFIG_FILE` env var or `shuttle.toml` by default) +/// 3. Environment variables with the `SHUTTLE.` prefix +/// +/// # Environment Variables +/// +/// Any configuration option can be overridden via environment variables using the pattern +/// `SHUTTLE.`. Nested fields (such as scheduler-specific configuration) are also delineated +/// by `.`s. For example: +/// - `SHUTTLE.STACK_SIZE=65536` sets the stack size +/// - `SHUTTLE.SCHEDULER.TYPE=random` sets the scheduler type +/// - `SHUTTLE.SCHEDULER.ITERATIONS=200` sets the number of iterations +/// +/// # Custom Schedulers +/// +/// Custom scheduler implementations can be registered with [`register_scheduler`] to make them +/// available via configuration. This allows custom schedulers to be used with `check` without +/// modifying test code: +/// +/// ```no_run +/// use shuttle::scheduler::{Schedule, Scheduler, Task, TaskId}; +/// use shuttle::register_scheduler; +/// +/// struct MyScheduler; +/// impl Scheduler for MyScheduler { +/// fn new_execution(&mut self) -> Option { Some(Schedule::new(0)) } +/// fn next_task(&mut self, runnable: &[&Task], _: Option, _: bool) -> Option { +/// runnable.first().map(|t| t.id()) +/// } +/// fn next_u64(&mut self) -> u64 { 0 } +/// } +/// +/// register_scheduler("my_scheduler", |_config| Box::new(MyScheduler)); +/// ``` +/// +/// # Example +/// +/// ```no_run +/// use shuttle::sync::{Arc, Mutex}; +/// use shuttle::thread; +/// +/// // Set the config file path via environment variable +/// std::env::set_var("SHUTTLE_CONFIG_FILE", "shuttle_config_example.toml"); +/// +/// shuttle::check(|| { +/// let lock = Arc::new(Mutex::new(0u64)); +/// let lock2 = lock.clone(); +/// +/// thread::spawn(move || { +/// *lock.lock().unwrap() = 1; +/// }); +/// +/// let _ = *lock2.lock().unwrap(); +/// }); +/// ``` pub fn check(f: F) where F: Fn() + Send + Sync + 'static, { - use crate::scheduler::RoundRobinScheduler; + let global_config = load_global_config(); + let config = Config::from_global_config(&global_config); + let scheduler_type = global_config + .get_string("scheduler.type") + .expect("No scheduler type found in global config!"); + let mut scheduler = scheduler::registry::create_scheduler(&scheduler_type, &global_config); + if should_enable_metrics(&global_config) { + scheduler = Box::new(scheduler::metrics::MetricsScheduler::new(scheduler)); + } + if should_check_uncontrolled_nondeterminism(&global_config) { + scheduler = Box::new(scheduler::UncontrolledNondeterminismCheckScheduler::new(scheduler)); + } - let runner = Runner::new(RoundRobinScheduler::new(1), Default::default()); + let runner = Runner::new(scheduler, config); runner.run(f); } +/// Load configuration from TOML files and environment variables +#[doc(hidden)] +pub fn load_global_config() -> config::Config { + load_global_config_from(None::<&str>) +} + +/// Load configuration from TOML files and environment variables with an optional file path +#[doc(hidden)] +pub fn load_global_config_from>(config_path: Option

) -> config::Config { + try_load_global_config_from(config_path).expect("Failed to load configuration") +} + +fn try_load_global_config_from>(config_path: Option

) -> Result { + let defaults = Config::default(); + + let max_steps_value = match defaults.max_steps { + MaxSteps::None => 0_i128, + MaxSteps::FailAfter(n) | MaxSteps::ContinueAfter(n) => n as i128, + }; + let max_steps_behavior = defaults.max_steps.variant_to_string(); + + let config_path = config_path + .map(|p| p.as_ref().to_path_buf()) + .or_else(|| std::env::var("SHUTTLE_CONFIG_FILE").ok().map(PathBuf::from)) + .unwrap_or_else(|| PathBuf::from("shuttle")); + + config::Config::builder() + .set_default("stack_size", defaults.stack_size as i128)? + .set_default("failure_persistence", defaults.failure_persistence.variant_to_string())? + .set_default("max_steps", max_steps_value)? + .set_default("max_steps_behavior", max_steps_behavior)? + .set_default("silence_warnings", defaults.silence_warnings)? + .set_default("record_steps_in_span", defaults.record_steps_in_span)? + .set_default("immediately_return_on_panic", defaults.immediately_return_on_panic)? + .set_default("enable_metrics", false)? + .set_default("check_uncontrolled_nondeterminism", false)? + .add_source(config::File::from(config_path.as_path()).required(false)) + .add_source(config::Environment::with_prefix("SHUTTLE").separator(".")) + .build() +} + /// Run the given function under a *uniformly* random scheduler for some number of iterations. /// Each iteration will run a (potentially) different randomized schedule. pub fn check_urw(f: F, iterations: usize) diff --git a/shuttle/src/scheduler/dfs.rs b/shuttle/src/scheduler/dfs.rs index c07d1946..9733b208 100644 --- a/shuttle/src/scheduler/dfs.rs +++ b/shuttle/src/scheduler/dfs.rs @@ -38,6 +38,13 @@ impl DfsScheduler { } } + /// Construct a new DfsScheduler from configuration. + pub fn from_config(config: &config::Config) -> Self { + let max_iterations = config.get_int("scheduler.max_iterations").ok().map(|i| i as usize); + let allow_random_data = config.get_bool("scheduler.allow_random_data").unwrap_or(false); + Self::new(max_iterations, allow_random_data) + } + /// Check if there are any scheduling points at or below the `index`th level that have remaining /// schedulable tasks to explore. // TODO probably should memoize this -- at each iteration, just need to know the largest i diff --git a/shuttle/src/scheduler/mod.rs b/shuttle/src/scheduler/mod.rs index 581a659e..c3adc854 100644 --- a/shuttle/src/scheduler/mod.rs +++ b/shuttle/src/scheduler/mod.rs @@ -6,6 +6,8 @@ mod data; mod dfs; mod pct; mod random; +/// Global registry for custom schedulers +pub mod registry; mod replay; mod round_robin; mod uncontrolled_nondeterminism; diff --git a/shuttle/src/scheduler/pct.rs b/shuttle/src/scheduler/pct.rs index 8135c4e1..ad93bcff 100644 --- a/shuttle/src/scheduler/pct.rs +++ b/shuttle/src/scheduler/pct.rs @@ -38,6 +38,17 @@ impl PctScheduler { Self::new_from_seed(OsRng.next_u64(), max_depth, max_iterations) } + /// Construct a new PctScheduler from configuration. + pub fn from_config(config: &config::Config) -> Self { + let depth = config.get_int("scheduler.depth").unwrap_or(3) as usize; + let iterations = config.get_int("scheduler.iterations").unwrap_or(100) as usize; + if let Ok(seed) = config.get_int("scheduler.seed") { + Self::new_from_seed(seed as u64, depth, iterations) + } else { + Self::new(depth, iterations) + } + } + /// Construct a new PCTScheduler with a given seed. /// /// If the `SHUTTLE_RANDOM_SEED` environment variable is set, then that seed will be used instead. diff --git a/shuttle/src/scheduler/random.rs b/shuttle/src/scheduler/random.rs index ea45ec7f..560a6800 100644 --- a/shuttle/src/scheduler/random.rs +++ b/shuttle/src/scheduler/random.rs @@ -52,6 +52,16 @@ impl RandomScheduler { Self::new_from_seed(OsRng.next_u64(), max_iterations) } + /// Construct a new RandomScheduler from configuration. + pub fn from_config(config: &config::Config) -> Self { + let iterations = config.get_int("scheduler.iterations").unwrap_or(100) as usize; + if let Ok(seed) = config.get_int("scheduler.seed") { + Self::new_from_seed(seed as u64, iterations) + } else { + Self::new(iterations) + } + } + /// Construct a new RandomScheduler with a given seed. /// /// Two RandomSchedulers initialized with the same seed will make the same scheduling decisions diff --git a/shuttle/src/scheduler/registry.rs b/shuttle/src/scheduler/registry.rs new file mode 100644 index 00000000..1e7a9dcc --- /dev/null +++ b/shuttle/src/scheduler/registry.rs @@ -0,0 +1,98 @@ +use crate::scheduler::{ + DfsScheduler, PctScheduler, RandomScheduler, ReplayScheduler, RoundRobinScheduler, Scheduler, UrwRandomScheduler, +}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; + +/// Factory function type for creating schedulers from configuration +pub type SchedulerFactory = Arc Box + Send + Sync>; + +static REGISTRY: OnceLock>> = OnceLock::new(); + +fn get_registry() -> &'static Mutex> { + REGISTRY.get_or_init(|| { + let mut map = HashMap::new(); + + // Register built-in schedulers + map.insert( + "random".to_string(), + Arc::new(|config: &config::Config| { + Box::new(RandomScheduler::from_config(config)) as Box + }) as SchedulerFactory, + ); + + map.insert( + "dfs".to_string(), + Arc::new(|config: &config::Config| Box::new(DfsScheduler::from_config(config)) as Box) + as SchedulerFactory, + ); + + map.insert( + "roundrobin".to_string(), + Arc::new(|config: &config::Config| { + Box::new(RoundRobinScheduler::from_config(config)) as Box + }) as SchedulerFactory, + ); + + map.insert( + "urw".to_string(), + Arc::new(|config: &config::Config| { + Box::new(UrwRandomScheduler::from_config(config)) as Box + }) as SchedulerFactory, + ); + + map.insert( + "pct".to_string(), + Arc::new(|config: &config::Config| Box::new(PctScheduler::from_config(config)) as Box) + as SchedulerFactory, + ); + + map.insert( + "replay".to_string(), + Arc::new(|config: &config::Config| { + Box::new(ReplayScheduler::from_config(config)) as Box + }) as SchedulerFactory, + ); + + Mutex::new(map) + }) +} + +/// Register a custom scheduler factory with the global registry +/// +/// # Example +/// +/// ```no_run +/// use shuttle::{register_scheduler, scheduler::Scheduler}; +/// use config::Config; +/// +/// struct MyScheduler; +/// impl Scheduler for MyScheduler { +/// // ... implementation +/// # fn new_execution(&mut self) -> Option { None } +/// # fn next_task(&mut self, _: &[&shuttle::scheduler::Task], _: Option, _: bool) -> Option { None } +/// # fn next_u64(&mut self) -> u64 { 0 } +/// } +/// +/// register_scheduler("my_scheduler", |config| { +/// let iterations = config.get_int("scheduler.iterations").unwrap_or(100) as usize; +/// Box::new(MyScheduler) +/// }); +/// ``` +pub fn register_scheduler(name: &str, factory: F) +where + F: Fn(&config::Config) -> Box + Send + Sync + 'static, +{ + let registry = get_registry(); + let mut map = registry.lock().unwrap(); + map.insert(name.to_string(), Arc::new(factory)); +} + +/// Create a scheduler from the registry +pub(crate) fn create_scheduler(name: &str, config: &config::Config) -> Box { + let registry = get_registry(); + let map = registry.lock().unwrap(); + map.get(name) + .map(|factory| factory(config)) + .expect("Could not create scheduler") +} diff --git a/shuttle/src/scheduler/replay.rs b/shuttle/src/scheduler/replay.rs index f97e91a7..2b2202bb 100644 --- a/shuttle/src/scheduler/replay.rs +++ b/shuttle/src/scheduler/replay.rs @@ -21,6 +21,18 @@ pub struct ReplayScheduler { } impl ReplayScheduler { + /// Construct a new ReplayScheduler from configuration. + pub fn from_config(config: &config::Config) -> Self { + if let Ok(path) = config.get_string("scheduler.schedule_file") { + Self::new_from_file(path).expect("failed to load schedule from file") + } else { + let schedule = config + .get_string("scheduler.schedule") + .expect("replay scheduler requires 'scheduler.schedule' or 'scheduler.schedule_file' config"); + Self::new_from_encoded(&schedule) + } + } + /// Given an encoded schedule, construct a new [`ReplayScheduler`] that will execute threads in /// the order specified in the schedule. pub fn new_from_encoded(encoded_schedule: &str) -> Self { diff --git a/shuttle/src/scheduler/round_robin.rs b/shuttle/src/scheduler/round_robin.rs index 86f5635c..45df3656 100644 --- a/shuttle/src/scheduler/round_robin.rs +++ b/shuttle/src/scheduler/round_robin.rs @@ -21,6 +21,12 @@ impl RoundRobinScheduler { data_source: RandomDataSource::initialize(0), } } + + /// Construct a new RoundRobinScheduler from configuration. + pub fn from_config(config: &config::Config) -> Self { + let iterations = config.get_int("scheduler.iterations").unwrap_or(1) as usize; + Self::new(iterations) + } } impl Scheduler for RoundRobinScheduler { diff --git a/shuttle/src/scheduler/urw.rs b/shuttle/src/scheduler/urw.rs index fa8453d9..e2392f31 100644 --- a/shuttle/src/scheduler/urw.rs +++ b/shuttle/src/scheduler/urw.rs @@ -60,6 +60,16 @@ impl UrwRandomScheduler { Self::new_from_seed(OsRng.next_u64(), max_iterations) } + /// Construct a new UrwRandomScheduler from configuration. + pub fn from_config(config: &config::Config) -> Self { + let iterations = config.get_int("scheduler.iterations").unwrap_or(100) as usize; + if let Ok(seed) = config.get_int("scheduler.seed") { + Self::new_from_seed(seed as u64, iterations) + } else { + Self::new(iterations) + } + } + /// Construct a UniformRandomScheduler with a given seed. /// Two UniformRandomSchedulers initialized with the same seed will make the same scheduling decisions when executing the same workloads. /// If the `SHUTTLE_RANDOM_SEED` environment variable is set, then that seed will be used instead. diff --git a/shuttle/tests/basic/execution.rs b/shuttle/tests/basic/execution.rs index 1720a1db..c65432e5 100644 --- a/shuttle/tests/basic/execution.rs +++ b/shuttle/tests/basic/execution.rs @@ -1,6 +1,6 @@ use shuttle::{ - check, check_dfs, current, - scheduler::{DfsScheduler, RandomScheduler}, + check_dfs, current, + scheduler::{DfsScheduler, RandomScheduler, RoundRobinScheduler}, thread, Config, MaxSteps, Runner, }; use std::panic::{catch_unwind, AssertUnwindSafe}; @@ -15,7 +15,9 @@ fn basic_scheduler_test() { let counter = Arc::new(AtomicUsize::new(0)); let counter_clone = Arc::clone(&counter); - check(move || { + let runner = Runner::new(RoundRobinScheduler::new(1), Default::default()); + + runner.run(move || { counter.fetch_add(1, Ordering::SeqCst); let counter_clone = Arc::clone(&counter); thread::spawn(move || { diff --git a/shuttle/tests/basic/rwlock.rs b/shuttle/tests/basic/rwlock.rs index c90da6aa..123625ed 100644 --- a/shuttle/tests/basic/rwlock.rs +++ b/shuttle/tests/basic/rwlock.rs @@ -1,6 +1,6 @@ -use shuttle::scheduler::PctScheduler; +use shuttle::scheduler::{PctScheduler, RoundRobinScheduler}; use shuttle::sync::{mpsc::channel, RwLock}; -use shuttle::{check, check_dfs, check_random, thread, Runner}; +use shuttle::{check_dfs, check_random, thread, Runner}; use std::collections::HashSet; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, TryLockError}; @@ -77,8 +77,10 @@ fn deadlock() { #[test] #[should_panic(expected = "deadlock")] fn deadlock_default() { + let runner = Runner::new(RoundRobinScheduler::new(1), Default::default()); + // Round-robin should always fail this deadlock test - check(deadlock); + runner.run(deadlock); } #[test] diff --git a/shuttle/tests/config/mod.rs b/shuttle/tests/config/mod.rs new file mode 100644 index 00000000..52abd24d --- /dev/null +++ b/shuttle/tests/config/mod.rs @@ -0,0 +1,204 @@ +use shuttle::scheduler::{DfsScheduler, PctScheduler, RandomScheduler, UrwRandomScheduler}; +use shuttle::{load_global_config, load_global_config_from}; +use std::env; +use std::sync::{Arc, RwLock}; + +static ENV_LOCK: RwLock<()> = RwLock::new(()); + +#[test] +fn test_config_from_env() { + let settings = { + let _guard = ENV_LOCK.write().unwrap(); + env::set_var("SHUTTLE.STACK_SIZE", "0"); + let settings = load_global_config(); + env::remove_var("SHUTTLE.STACK_SIZE"); + settings + }; + + let config = shuttle::Config::from_global_config(&settings); + + assert_eq!(config.stack_size, 0); +} + +#[test] +fn test_config_from_file() { + let settings = { + let _guard = ENV_LOCK.read().unwrap(); + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/config/shuttle_dfs.toml"); + load_global_config_from(Some(path)) + }; + assert_eq!(settings.get_string("scheduler.type").unwrap(), "dfs"); +} + +#[test] +fn test_custom_scheduler_custom_config_field() { + use shuttle::register_scheduler; + use shuttle::scheduler::{Schedule, Scheduler, Task, TaskId}; + + struct FooScheduler { + _foo: String, + } + + impl Scheduler for FooScheduler { + fn new_execution(&mut self) -> Option { + Some(Schedule::new(0)) + } + + fn next_task(&mut self, runnable: &[&Task], _: Option, _: bool) -> Option { + runnable.first().map(|t| t.id()) + } + + fn next_u64(&mut self) -> u64 { + 0 + } + } + + register_scheduler("foo_scheduler", |config| { + let foo = config.get_string("scheduler.foo").unwrap(); + Box::new(FooScheduler { _foo: foo }) + }); + let settings = { + let _guard = ENV_LOCK.write().unwrap(); + env::set_var("SHUTTLE.SCHEDULER.TYPE", "foo_scheduler"); + env::set_var("SHUTTLE.SCHEDULER.FOO", "bar"); + + let settings = load_global_config(); + env::remove_var("SHUTTLE.SCHEDULER.TYPE"); + env::remove_var("SHUTTLE.SCHEDULER.FOO"); + settings + }; + assert_eq!(settings.get_string("scheduler.type").unwrap(), "foo_scheduler"); + assert_eq!(settings.get_string("scheduler.foo").unwrap(), "bar"); +} + +mod main_task_only_scheduler { + use shuttle::scheduler::{Schedule, Scheduler, Task, TaskId}; + use std::sync::Once; + + static REGISTER: Once = Once::new(); + + pub struct MainTaskOnlyScheduler { + iterations: usize, + current: usize, + } + + impl Scheduler for MainTaskOnlyScheduler { + fn new_execution(&mut self) -> Option { + if self.current < self.iterations { + self.current += 1; + Some(Schedule::new(0)) + } else { + None + } + } + + fn next_task(&mut self, runnable: &[&Task], current: Option, _: bool) -> Option { + if current.is_none() { + Some(runnable[0].id()) + } else { + None + } + } + + fn next_u64(&mut self) -> u64 { + 0 + } + } + + pub fn register() { + REGISTER.call_once(|| { + shuttle::register_scheduler("main_task_only", |config| { + let iterations = config.get_int("scheduler.iterations").unwrap_or(1) as usize; + Box::new(MainTaskOnlyScheduler { iterations, current: 0 }) + }); + }); + } +} + +#[test] +fn test_run_custom_scheduler() { + main_task_only_scheduler::register(); + + let thread1_ran_outer = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let thread1_ran_inner = thread1_ran_outer.clone(); + { + let _guard = ENV_LOCK.write().unwrap(); + env::set_var("SHUTTLE.SCHEDULER.TYPE", "main_task_only"); + env::set_var("SHUTTLE.SCHEDULER.ITERATIONS", "1"); + + shuttle::check(move || { + use shuttle::sync::atomic::{AtomicBool, Ordering}; + use shuttle::sync::Arc; + + let thread1_ran = Arc::new(AtomicBool::new(false)); + let t1 = thread1_ran.clone(); + + shuttle::thread::spawn(move || { + t1.store(true, Ordering::SeqCst); + }); + + shuttle::thread::yield_now(); + + _ = thread1_ran_inner.fetch_or(thread1_ran.load(Ordering::SeqCst), Ordering::SeqCst); + }); + + env::remove_var("SHUTTLE.SCHEDULER.TYPE"); + env::remove_var("SHUTTLE.SCHEDULER.ITERATIONS"); + } + assert!(!thread1_ran_outer.load(std::sync::atomic::Ordering::SeqCst)); +} + +#[test] +fn test_config_example_file() { + let settings = { + let _guard = ENV_LOCK.read().unwrap(); + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("shuttle_config_example.toml"); + load_global_config_from(Some(path)) + }; + + let expected = crate::Config::default(); + let actual = crate::Config::from_global_config(&settings); + + assert_eq!(expected, actual); + + RandomScheduler::from_config(&settings); +} + +#[test] +fn test_config_example_scheduler_dfs() { + let settings = { + let _guard = ENV_LOCK.write().unwrap(); + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("shuttle_config_example.toml"); + env::set_var("SHUTTLE.SCHEDULER.TYPE", "dfs"); + let settings = load_global_config_from(Some(path)); + env::remove_var("SHUTTLE.SCHEDULER.TYPE"); + settings + }; + DfsScheduler::from_config(&settings); +} + +#[test] +fn test_config_example_scheduler_urw() { + let settings = { + let _guard = ENV_LOCK.write().unwrap(); + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("shuttle_config_example.toml"); + env::set_var("SHUTTLE.SCHEDULER.TYPE", "urw"); + let settings = load_global_config_from(Some(path)); + env::remove_var("SHUTTLE.SCHEDULER.TYPE"); + settings + }; + UrwRandomScheduler::from_config(&settings); +} + +#[test] +fn test_config_example_scheduler_pct() { + let settings = { + let _guard = ENV_LOCK.write().unwrap(); + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("shuttle_config_example.toml"); + env::set_var("SHUTTLE.SCHEDULER.TYPE", "pct"); + let settings = load_global_config_from(Some(path)); + env::remove_var("SHUTTLE.SCHEDULER.TYPE"); + settings + }; + PctScheduler::from_config(&settings); +} diff --git a/shuttle/tests/config/shuttle_dfs.toml b/shuttle/tests/config/shuttle_dfs.toml new file mode 100644 index 00000000..5b76e304 --- /dev/null +++ b/shuttle/tests/config/shuttle_dfs.toml @@ -0,0 +1,5 @@ +# DFS scheduler configuration +[scheduler] +type = "dfs" +max_iterations = 100 +allow_random_data = false diff --git a/shuttle/tests/mod.rs b/shuttle/tests/mod.rs index 4fd7010e..46fa7aa8 100644 --- a/shuttle/tests/mod.rs +++ b/shuttle/tests/mod.rs @@ -2,6 +2,7 @@ mod advanced; mod basic; +mod config; mod data; mod demo; mod future;