Merge pull request #61 from heliaxdev/yuji/remove-rng

Remove `rng` from `State`

Author: yito88

check/src/checks/status.rs

@@ -9,8 +9,9 @@ pub struct StatusCheck;
 
 impl DoCheck for StatusCheck {
     async fn check(&self, sdk: &Sdk, _state: &mut crate::state::State) -> Result<(), String> {
-        let client = sdk.namada.clone_client();
-        let status = client
+        let status = sdk
+            .namada
+            .client
             .status()
             .await
             .map_err(|e| format!("Failed to query status: {e}"))?;

workload/src/config.rs

@@ -12,8 +12,6 @@ pub struct AppConfig {
     #[arg(required = true)]
     pub chain_id: String,
     #[clap(long, env)]
-    pub seed: Option<u64>,
-    #[clap(long, env)]
     pub id: u64,
     #[clap(long, env)]
     #[arg(required = true)]

workload/src/executor.rs

@@ -135,8 +135,8 @@ impl WorkloadExecutor {
         step_type.is_valid(&self.sdk, &self.state).await
     }
 
-    pub async fn build(&mut self, step_type: &StepType) -> Result<Vec<Task>, StepError> {
-        step_type.build_task(&self.sdk, &mut self.state).await
+    pub async fn build(&self, step_type: &StepType) -> Result<Vec<Task>, StepError> {
+        step_type.build_task(&self.sdk, &self.state).await
     }
 
     pub async fn build_check(&self, tasks: &[Task]) -> Result<Vec<Check>, StepError> {

workload/src/main.rs

@@ -61,7 +61,7 @@ async fn inner_main() -> Code {
         Ok(result) => result,
         Err(StateError::EmptyFile) => {
             tracing::warn!("State file is empty, creating new one");
-            match State::create_new(config.id, config.seed) {
+            match State::create_new(config.id) {
                 Ok(result) => result,
                 Err(e) => return Code::StateFatal(e),
             }
@@ -70,7 +70,6 @@ async fn inner_main() -> Code {
     };
 
     tracing::info!("Using base dir: {}", state.base_dir.as_path().display());
-    tracing::info!("Using seed: {}", state.seed);
 
     let url = Url::from_str(&config.rpc).expect("invalid RPC address");
     tracing::debug!("Opening connection to {url}");

workload/src/sdk/namada.rs

@@ -42,7 +42,7 @@ impl Sdk {
             .insert_keypair("faucet".to_string(), true, sk, None, Some(address), None)
             .unwrap();
 
-        let native_token = rpc::query_native_token(&namada.clone_client())
+        let native_token = rpc::query_native_token(&namada.client)
             .await
             .map_err(|e| e.to_string())?;
         namada_wallet

workload/src/state.rs

@@ -1,11 +1,10 @@
-use std::{
-    collections::{BTreeSet, HashMap, HashSet},
-    env, fs,
-    path::PathBuf,
-};
+use std::collections::{BTreeSet, HashMap, HashSet};
+use std::path::PathBuf;
+use std::{env, fs};
 
+use antithesis_sdk::random::AntithesisRng;
 use fs2::FileExt;
-use rand::{seq::IteratorRandom, Rng, RngCore};
+use rand::seq::IteratorRandom;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
@@ -82,15 +81,12 @@ pub struct State {
     pub deactivated_validators: HashMap<Alias, Account>,
     pub proposals: HashMap<u64, (u64, u64)>,
     pub id: u64,
-    pub seed: u64,
-    pub rng: AntithesisRng,
     pub base_dir: PathBuf,
     pub stats: HashMap<String, u64>,
 }
 
 impl State {
-    pub fn new(id: u64, seed: Option<u64>) -> Self {
-        let seed = seed.unwrap_or(rand::thread_rng().gen_range(0..u64::MAX));
+    pub fn new(id: u64) -> Self {
         Self {
             accounts: HashMap::default(),
             masp_accounts: HashMap::default(),
@@ -103,8 +99,6 @@ impl State {
             deactivated_validators: HashMap::default(),
             proposals: HashMap::default(),
             id,
-            seed,
-            rng: AntithesisRng::default(),
             base_dir: env::current_dir().unwrap().join("base"),
             stats: HashMap::default(),
         }
@@ -130,11 +124,11 @@ impl State {
         Ok(())
     }
 
-    pub fn create_new(id: u64, seed: Option<u64>) -> Result<(Self, fs::File), StateError> {
+    pub fn create_new(id: u64) -> Result<(Self, fs::File), StateError> {
         // Lock the state file before writing the new
         let file = Self::lock_state_file(id)?;
 
-        let state = Self::new(id, seed);
+        let state = Self::new(id);
         state.save(None)?;
 
         Ok((state, file))
@@ -269,16 +263,16 @@ impl State {
 
     /// GET
 
-    pub fn random_account(&mut self, blacklist: Vec<Alias>) -> Option<Account> {
+    pub fn random_account(&self, blacklist: Vec<Alias>) -> Option<Account> {
         self.accounts
             .iter()
             .filter(|(alias, _)| !blacklist.contains(alias))
-            .choose(&mut self.rng)
+            .choose(&mut AntithesisRng)
             .map(|(_, account)| account.clone())
     }
 
     pub fn random_masp_account_with_min_balance(
-        &mut self,
+        &self,
         blacklist: Vec<Alias>,
         min_value: u64,
     ) -> Option<MaspAccount> {
@@ -294,74 +288,74 @@ impl State {
                     None
                 }
             })
-            .choose(&mut self.rng)
+            .choose(&mut AntithesisRng)
     }
 
-    pub fn random_payment_address(&mut self, blacklist: Vec<Alias>) -> Option<MaspAccount> {
+    pub fn random_payment_address(&self, blacklist: Vec<Alias>) -> Option<MaspAccount> {
         self.masp_accounts
             .iter()
             .filter(|(alias, _)| !blacklist.contains(alias))
-            .choose(&mut self.rng)
+            .choose(&mut AntithesisRng)
             .map(|(_, account)| account.clone())
     }
 
     pub fn random_implicit_accounts(
-        &mut self,
+        &self,
         blacklist: Vec<Alias>,
         sample_size: usize,
     ) -> Vec<Account> {
         self.accounts
             .iter()
             .filter(|(alias, _)| !blacklist.contains(alias))
             .filter(|(_, account)| account.is_implicit())
-            .choose_multiple(&mut self.rng, sample_size)
+            .choose_multiple(&mut AntithesisRng, sample_size)
             .into_iter()
             .map(|(_, account)| account.clone())
             .collect()
     }
 
     pub fn random_enstablished_account(
-        &mut self,
+        &self,
         blacklist: Vec<Alias>,
         sample_size: usize,
     ) -> Vec<Account> {
         self.accounts
             .iter()
             .filter(|(alias, _)| !blacklist.contains(alias))
             .filter(|(_, account)| account.is_enstablished())
-            .choose_multiple(&mut self.rng, sample_size)
+            .choose_multiple(&mut AntithesisRng, sample_size)
             .into_iter()
             .map(|(_, account)| account.clone())
             .collect()
     }
 
-    pub fn random_validator(&mut self, blacklist: Vec<Alias>, sample_size: usize) -> Vec<Account> {
+    pub fn random_validator(&self, blacklist: Vec<Alias>, sample_size: usize) -> Vec<Account> {
         self.validators
             .iter()
             .filter(|(alias, _)| !blacklist.contains(alias))
             .filter(|(_, account)| account.is_enstablished())
-            .choose_multiple(&mut self.rng, sample_size)
+            .choose_multiple(&mut AntithesisRng, sample_size)
             .into_iter()
             .map(|(_, account)| account.clone())
             .collect()
     }
 
     pub fn random_deactivated_validator(
-        &mut self,
+        &self,
         blacklist: Vec<Alias>,
         sample_size: usize,
     ) -> Vec<Account> {
         self.deactivated_validators
             .iter()
             .filter(|(alias, _)| !blacklist.contains(alias))
             .filter(|(_, account)| account.is_enstablished())
-            .choose_multiple(&mut self.rng, sample_size)
+            .choose_multiple(&mut AntithesisRng, sample_size)
             .into_iter()
             .map(|(_, account)| account.clone())
             .collect()
     }
 
-    pub fn random_bond(&mut self) -> Bond {
+    pub fn random_bond(&self) -> Bond {
         self.bonds
             .iter()
             .flat_map(|(source, bonds)| {
@@ -377,12 +371,12 @@ impl State {
                     }
                 })
             })
-            .choose(&mut self.rng)
+            .choose(&mut AntithesisRng)
             .unwrap()
     }
 
     pub fn random_account_with_min_balance(
-        &mut self,
+        &self,
         blacklist: Vec<Alias>,
         min_balance: u64,
     ) -> Option<Account> {
@@ -398,7 +392,7 @@ impl State {
                     None
                 }
             })
-            .choose(&mut self.rng)
+            .choose(&mut AntithesisRng)
     }
 
     pub fn get_account_by_alias(&self, alias: &Alias) -> Account {
@@ -423,14 +417,14 @@ impl State {
             .unwrap_or_default()
     }
 
-    pub fn get_redelegations_targets_for(&mut self, alias: &Alias) -> HashSet<String> {
+    pub fn get_redelegations_targets_for(&self, alias: &Alias) -> HashSet<String> {
         self.redelegations
             .get(alias)
             .map(|data| data.keys().cloned().collect::<HashSet<String>>())
             .unwrap_or_default()
     }
 
-    pub fn random_votable_proposal(&mut self, current_epoch: u64) -> u64 {
+    pub fn random_votable_proposal(&self, current_epoch: u64) -> u64 {
         self.proposals
             .iter()
             .filter_map(|(proposal_id, (start_epoch, end_epoch))| {
@@ -440,7 +434,7 @@ impl State {
                     None
                 }
             })
-            .choose(&mut self.rng)
+            .choose(&mut AntithesisRng)
             .unwrap()
     }
 
@@ -626,35 +620,3 @@ impl State {
             .insert(latest_proposal_id, (start_epoch, end_epoch));
     }
 }
-
-#[derive(Clone, Debug, Default, Serialize, Deserialize)]
-pub struct AntithesisRng {}
-
-impl RngCore for AntithesisRng {
-    fn next_u32(&mut self) -> u32 {
-        (antithesis_sdk::random::get_random() & 0xFFFF_FFFF) as u32
-    }
-
-    fn next_u64(&mut self) -> u64 {
-        antithesis_sdk::random::get_random()
-    }
-
-    fn fill_bytes(&mut self, dest: &mut [u8]) {
-        let mut i = 0;
-        while i + 8 <= dest.len() {
-            let random = antithesis_sdk::random::get_random();
-            dest[i..i + 8].copy_from_slice(&random.to_le_bytes());
-            i += 8;
-        }
-        if i < dest.len() {
-            let random = antithesis_sdk::random::get_random();
-            let dest_len = dest.len();
-            dest[i..].copy_from_slice(&random.to_le_bytes()[..dest_len - i]);
-        }
-    }
-
-    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
-        self.fill_bytes(dest);
-        Ok(())
-    }
-}

workload/src/step.rs

@@ -104,7 +104,7 @@ pub trait StepContext {
     async fn is_valid(&self, sdk: &Sdk, state: &State) -> Result<bool, StepError>;
 
     #[allow(async_fn_in_trait)]
-    async fn build_task(&self, sdk: &Sdk, state: &mut State) -> Result<Vec<Task>, StepError>;
+    async fn build_task(&self, sdk: &Sdk, state: &State) -> Result<Vec<Task>, StepError>;
 
     fn assert(&self, code: &Code);
 }

workload/src/step/batch.rs

@@ -1,3 +1,4 @@
+use antithesis_sdk::random::AntithesisRng;
 use rand::seq::SliceRandom;
 use serde_json::json;
 
@@ -23,7 +24,7 @@ impl StepContext for BatchBond {
         Ok(state.min_n_account_with_min_balance(3, MIN_TRANSFER_BALANCE))
     }
 
-    async fn build_task(&self, sdk: &Sdk, state: &mut State) -> Result<Vec<Task>, StepError> {
+    async fn build_task(&self, sdk: &Sdk, state: &State) -> Result<Vec<Task>, StepError> {
         Box::pin(build_batch(
             sdk,
             vec![StepType::Bond(Default::default())],
@@ -67,7 +68,7 @@ impl StepContext for BatchRandom {
         Ok(state.min_n_account_with_min_balance(3, MIN_TRANSFER_BALANCE) && state.min_bonds(3))
     }
 
-    async fn build_task(&self, sdk: &Sdk, state: &mut State) -> Result<Vec<Task>, StepError> {
+    async fn build_task(&self, sdk: &Sdk, state: &State) -> Result<Vec<Task>, StepError> {
         Box::pin(build_batch(
             sdk,
             vec![
@@ -111,16 +112,16 @@ async fn build_batch(
     sdk: &Sdk,
     possibilities: Vec<StepType>,
     max_size: u64,
-    state: &mut State,
+    state: &State,
 ) -> Result<Vec<Task>, StepError> {
     let mut tmp_state = state.clone();
 
     let mut batch_tasks = vec![];
     for _ in 0..max_size {
         let step = possibilities
-            .choose(&mut state.rng)
+            .choose(&mut AntithesisRng)
             .expect("at least one StepType should exist");
-        let tasks = step.build_task(sdk, &mut tmp_state).await?;
+        let tasks = step.build_task(sdk, &tmp_state).await?;
         for task in &tasks {
             task.update_state(&mut tmp_state, false);
         }

workload/src/step/become_validator.rs

@@ -24,14 +24,14 @@ impl StepContext for BecomeValidator {
         Ok(state.min_n_enstablished_accounts(1))
     }
 
-    async fn build_task(&self, _sdk: &Sdk, state: &mut State) -> Result<Vec<Task>, StepError> {
-        let commission_rate = utils::random_between::<u64>(state, 0, 100);
+    async fn build_task(&self, _sdk: &Sdk, state: &State) -> Result<Vec<Task>, StepError> {
+        let commission_rate = utils::random_between::<u64>(0, 100);
         let commission_rate = Dec::new(commission_rate as i128, 2).unwrap();
 
-        let commission_rate_change = utils::random_between::<u64>(state, 0, 100);
+        let commission_rate_change = utils::random_between::<u64>(0, 100);
         let commission_rate_change = Dec::new(commission_rate_change as i128, 2).unwrap();
 
-        let random_alias = utils::random_alias(state);
+        let random_alias = utils::random_alias();
         let consensus_key_alias = format!("{}-consensus", random_alias.name);
         let eth_cold_key_alias = format!("{}-eth-cold", random_alias.name);
         let eth_hot_key_alias = format!("{}-eth-hot", random_alias.name);