working

Author: jayshrivastava

Cargo.lock

@@ -35,6 +35,8 @@ delegate = "0.13.4"
 dashmap = "6.1.0"
 prost = "0.13.5"
 rand = "0.8.5"
+rand_chacha = "0.3.1"
+base64 = "0.22"
 object_store = "0.12.3"
 bytes = "1.10.1"
 log = "0.4"

Cargo.toml

@@ -4,12 +4,13 @@ use datafusion::error::Result;
 use datafusion_distributed::test_utils::{
     fuzz::{FuzzDB, FuzzConfig},
     tpcds::{discover_tpcds_queries, register_available_tpcds_tables, generate_tpcds_data},
+    rand::rng,
 };
+use rand::Rng;
 use datafusion::prelude::SessionContext;
 use log::{debug, info, warn, error};
 use std::process;
 use std::time::Instant;
-use rand::Rng;
 
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
@@ -18,6 +19,10 @@ use rand::Rng;
 struct Cli {
     #[command(subcommand)]
     command: Commands,
+    
+    /// Base64-encoded seed for reproducible randomization
+    #[arg(long)]
+    seed: Option<String>,
 }
 
 /// Trait for defining workloads that can be fuzzed
@@ -124,7 +129,7 @@ async fn main() -> Result<()> {
             info!("   Workload: {}", workload.name());
             info!("   Scale factor: {}", scale_factor);
 
-            if let Err(e) = run_workload(workload).await {
+            if let Err(e) = run_workload(workload, cli.seed).await {
                 error!("❌ Fuzz testing failed: {}", e);
                 process::exit(1);
             }
@@ -146,25 +151,30 @@ fn validate_tpcds_args(scale_factor: &str) -> std::result::Result<(), String> {
     Ok(())
 }
 
-fn randomized_fuzz_config() -> FuzzConfig {
-    let mut rng = rand::thread_rng();
+fn randomized_fuzz_config(seed: Option<String>) -> Result<(FuzzConfig, String)> {
+    let (mut rng, seed_b64) = rng(seed)?;
+    
     let config = FuzzConfig {
         num_workers: rng.gen_range(2..=8),
         files_per_task: rng.gen_range(1..=8),
         cardinality_task_count_factor: rng.gen_range(1.0..=3.0),
     };
-    config
+    
+    Ok((config, seed_b64))
 }
 
 /// Run a workload using the generic workload trait
-async fn run_workload<W>(workload: W) -> Result<()> 
+async fn run_workload<W>(workload: W, seed: Option<String>) -> Result<()> 
 where W: Workload {
     info!("📊 Running {} workload...", workload.name());
 
     // Create FuzzDB with randomized session config
     info!("⚙️  Setting up distributed session with randomized configuration...");
 
-    let fuzz_db = match FuzzDB::new(randomized_fuzz_config(), W::setup).await {
+    let (config, seed_b64) = randomized_fuzz_config(seed)?;
+    info!("🎲 Using seed: {}", seed_b64);
+    
+    let fuzz_db = match FuzzDB::new(config, W::setup).await {
         Ok(db) => db,
         Err(e) => {
             return Err(datafusion::error::DataFusionError::Execution(format!(

src/bin/fuzz.rs

@@ -315,31 +315,6 @@ mod tests {
     use std::sync::Arc;
     use datafusion::logical_expr::LogicalPlanBuilder;
 
-    #[tokio::test]
-    async fn test_session_config_default() {
-        let config = SessionConfig::default();
-        assert_eq!(config.num_workers, 4);
-        assert_eq!(config.tasks_per_file, 4);
-        assert_eq!(config.cardinality_task_count_factor, 4);
-        assert_eq!(config.target_partitions, 8);
-    }
-
-    #[tokio::test]
-    async fn test_randomize_session_config() {
-        let config1 = randomize_session_config();
-        let config2 = randomize_session_config();
-        
-        // Basic bounds checking
-        assert!(config1.num_workers >= 2 && config1.num_workers <= 8);
-        assert!(config1.tasks_per_file >= 1 && config1.tasks_per_file <= 8);
-        assert!(config1.cardinality_task_count_factor >= 1 && config1.cardinality_task_count_factor <= 8);
-        assert!(config1.target_partitions >= 4 && config1.target_partitions <= 16);
-
-        // Configs should be different (with high probability)
-        // Note: There's a small chance they could be equal due to randomness
-        println!("Config 1: {:?}", config1);
-        println!("Config 2: {:?}", config2);
-    }
 
     #[tokio::test]
     async fn test_records_equal_as_sets_empty() {

src/test_utils/fuzz.rs

@@ -6,6 +6,7 @@ pub mod metrics;
 pub mod mock_exec;
 pub mod parquet;
 pub mod plans;
+pub mod rand;
 pub mod session_context;
 pub mod tpcds;
 pub mod tpch;

src/test_utils/mod.rs

@@ -0,0 +1,57 @@
+use rand::{SeedableRng, RngCore};
+use rand_chacha::ChaCha8Rng;
+use base64::{Engine as _, engine::general_purpose::STANDARD};
+use datafusion::error::{DataFusionError, Result};
+
+/// Create a seeded ChaCha8 RNG from a base64 string or generate a new random seed
+pub fn rng(seed: Option<String>) -> Result<(ChaCha8Rng, String)> {
+    let (rng, seed_b64) = if let Some(seed_str) = seed {
+        // Use provided seed
+        let seed_bytes = STANDARD.decode(&seed_str).map_err(|e| {
+            DataFusionError::Execution(format!("Invalid base64 seed: {}", e))
+        })?;
+        
+        if seed_bytes.len() != 32 {
+            return Err(DataFusionError::Execution(
+                "Seed must be 32 bytes (256 bits) when base64-decoded".to_string()
+            ));
+        }
+        
+        let mut seed_array = [0u8; 32];
+        seed_array.copy_from_slice(&seed_bytes);
+        let rng = ChaCha8Rng::from_seed(seed_array);
+        
+        (rng, seed_str)
+    } else {
+        // Generate random seed
+        let mut seed = [0u8; 32];
+        rand::thread_rng().fill_bytes(&mut seed);
+        let seed_b64 = STANDARD.encode(&seed);
+        let rng = ChaCha8Rng::from_seed(seed);
+        
+        (rng, seed_b64)
+    };
+    
+    Ok((rng, seed_b64))
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_rng_with_seed() {
+        // Create a known seed
+        let seed_bytes = [42u8; 32];
+        let seed_b64 = STANDARD.encode(&seed_bytes);
+        
+        let (mut rng1, _) = rng(Some(seed_b64.clone())).unwrap();
+        let (mut rng2, _) = rng(Some(seed_b64)).unwrap();
+        
+        // Both RNGs should generate the same sequence
+        assert_eq!(rng1.next_u64(), rng2.next_u64());
+        assert_eq!(rng1.next_u64(), rng2.next_u64());
+    }
+
+}