txn generator workflow

crates/transaction-emitter-lib/src/emitter/mod.rs

@@ -620,6 +620,7 @@ impl TxnEmitter {
         };
         let (txn_generator_creator, _, _) = create_txn_generator_creator(
             &req.transaction_mix_per_phase,
+            root_account,
             &mut all_accounts,
             vec![],
             &txn_executor,

crates/transaction-generator-lib/src/account_generator.rs

@@ -1,22 +1,21 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
-use crate::{create_account_transaction, TransactionGenerator, TransactionGeneratorCreator};
-use aptos_infallible::RwLock;
-use aptos_logger::{info, sample, sample::SampleRate};
+use crate::{
+    create_account_transaction, ObjectPool, TransactionGenerator, TransactionGeneratorCreator,
+};
 use aptos_sdk::{
     move_types::account_address::AccountAddress,
     transaction_builder::TransactionFactory,
     types::{transaction::SignedTransaction, LocalAccount},
 };
-use rand::{rngs::StdRng, Rng, SeedableRng};
-use std::{sync::Arc, time::Duration};
+use rand::{rngs::StdRng, SeedableRng};
+use std::sync::Arc;
 
 pub struct AccountGenerator {
     rng: StdRng,
     txn_factory: TransactionFactory,
-    addresses_pool: Arc<RwLock<Vec<AccountAddress>>>,
-    accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
-    add_created_accounts_to_pool: bool,
+    addresses_pool: Option<Arc<ObjectPool<AccountAddress>>>,
+    accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
     max_working_set: usize,
     creation_balance: u64,
 }
@@ -25,9 +24,8 @@ impl AccountGenerator {
     pub fn new(
         rng: StdRng,
         txn_factory: TransactionFactory,
-        addresses_pool: Arc<RwLock<Vec<AccountAddress>>>,
-        accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
-        add_created_accounts_to_pool: bool,
+        addresses_pool: Option<Arc<ObjectPool<AccountAddress>>>,
+        accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
         max_working_set: usize,
         creation_balance: u64,
     ) -> Self {
@@ -36,41 +34,12 @@ impl AccountGenerator {
             txn_factory,
             addresses_pool,
             accounts_pool,
-            add_created_accounts_to_pool,
             max_working_set,
             creation_balance,
         }
     }
 }
 
-fn add_to_sized_pool<T>(
-    pool: &RwLock<Vec<T>>,
-    mut addition: Vec<T>,
-    max_working_set: usize,
-    rng: &mut StdRng,
-) {
-    let mut current = pool.write();
-    if current.len() < max_working_set {
-        current.append(&mut addition);
-        sample!(
-            SampleRate::Duration(Duration::from_secs(120)),
-            info!("Accounts working set increased to {}", current.len())
-        );
-    } else {
-        let start = rng.gen_range(0, current.len() - addition.len());
-        current[start..start + addition.len()].swap_with_slice(&mut addition);
-
-        sample!(
-            SampleRate::Duration(Duration::from_secs(120)),
-            info!(
-                "Already at limit {} > {}, so exchanged accounts in working set",
-                current.len(),
-                max_working_set
-            )
-        );
-    }
-}
-
 impl TransactionGenerator for AccountGenerator {
     fn generate_transactions(
         &mut self,
@@ -94,52 +63,41 @@ impl TransactionGenerator for AccountGenerator {
             new_account_addresses.push(receiver_address);
         }
 
-        if self.add_created_accounts_to_pool {
-            add_to_sized_pool(
-                self.accounts_pool.as_ref(),
-                new_accounts,
-                self.max_working_set,
-                &mut self.rng,
-            );
-            add_to_sized_pool(
-                self.addresses_pool.as_ref(),
+        if let Some(addresses_pool) = &self.addresses_pool {
+            addresses_pool.add_to_pool_bounded(
                 new_account_addresses,
                 self.max_working_set,
                 &mut self.rng,
             );
         }
+        if let Some(accounts_pool) = &self.accounts_pool {
+            accounts_pool.add_to_pool_bounded(new_accounts, self.max_working_set, &mut self.rng);
+        }
+
         requests
     }
 }
 
 pub struct AccountGeneratorCreator {
     txn_factory: TransactionFactory,
-    addresses_pool: Arc<RwLock<Vec<AccountAddress>>>,
-    accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
-    add_created_accounts_to_pool: bool,
+    addresses_pool: Option<Arc<ObjectPool<AccountAddress>>>,
+    accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
     max_working_set: usize,
     creation_balance: u64,
 }
 
 impl AccountGeneratorCreator {
     pub fn new(
         txn_factory: TransactionFactory,
-        addresses_pool: Arc<RwLock<Vec<AccountAddress>>>,
-        accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
-        add_created_accounts_to_pool: bool,
+        addresses_pool: Option<Arc<ObjectPool<AccountAddress>>>,
+        accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
         max_working_set: usize,
         creation_balance: u64,
     ) -> Self {
-        if add_created_accounts_to_pool {
-            addresses_pool.write().reserve(max_working_set);
-            accounts_pool.write().reserve(max_working_set);
-        }
-
         Self {
             txn_factory,
             addresses_pool,
             accounts_pool,
-            add_created_accounts_to_pool,
             max_working_set,
             creation_balance,
         }
@@ -153,7 +111,6 @@ impl TransactionGeneratorCreator for AccountGeneratorCreator {
             self.txn_factory.clone(),
             self.addresses_pool.clone(),
             self.accounts_pool.clone(),
-            self.add_created_accounts_to_pool,
             self.max_working_set,
             self.creation_balance,
         ))

crates/transaction-generator-lib/src/accounts_pool_wrapper.rs

@@ -1,9 +1,9 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{get_account_to_burn_from_pool, TransactionGenerator, TransactionGeneratorCreator};
-use aptos_infallible::RwLock;
+use crate::{ObjectPool, TransactionGenerator, TransactionGeneratorCreator};
 use aptos_sdk::types::{transaction::SignedTransaction, LocalAccount};
+use rand::{rngs::StdRng, SeedableRng};
 use std::sync::Arc;
 
 /// Wrapper that allows inner transaction generator to have unique accounts
@@ -12,18 +12,24 @@ use std::sync::Arc;
 /// and burning (removing accounts from the pool) them - basically using them only once.
 /// (we cannot use more as sequence number is not updated on failure)
 pub struct AccountsPoolWrapperGenerator {
-    creator: Box<dyn TransactionGenerator>,
-    accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
+    rng: StdRng,
+    generator: Box<dyn TransactionGenerator>,
+    source_accounts_pool: Arc<ObjectPool<LocalAccount>>,
+    destination_accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
 }
 
 impl AccountsPoolWrapperGenerator {
     pub fn new(
-        creator: Box<dyn TransactionGenerator>,
-        accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
+        rng: StdRng,
+        generator: Box<dyn TransactionGenerator>,
+        source_accounts_pool: Arc<ObjectPool<LocalAccount>>,
+        destination_accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
     ) -> Self {
         Self {
-            creator,
-            accounts_pool,
+            rng,
+            generator,
+            source_accounts_pool,
+            destination_accounts_pool,
         }
     }
 }
@@ -34,40 +40,51 @@ impl TransactionGenerator for AccountsPoolWrapperGenerator {
         _account: &LocalAccount,
         num_to_create: usize,
     ) -> Vec<SignedTransaction> {
-        let mut accounts_to_burn =
-            get_account_to_burn_from_pool(&self.accounts_pool, num_to_create);
-        if accounts_to_burn.is_empty() {
+        let mut accounts_to_use =
+            self.source_accounts_pool
+                .take_from_pool(num_to_create, true, &mut self.rng);
+        if accounts_to_use.is_empty() {
             return Vec::new();
         }
-        accounts_to_burn
+        let txns = accounts_to_use
             .iter_mut()
-            .flat_map(|account| self.creator.generate_transactions(account, 1))
-            .collect()
+            .flat_map(|account| self.generator.generate_transactions(account, 1))
+            .collect();
+
+        if let Some(destination_accounts_pool) = &self.destination_accounts_pool {
+            destination_accounts_pool.add_to_pool(accounts_to_use);
+        }
+        txns
     }
 }
 
 pub struct AccountsPoolWrapperCreator {
     creator: Box<dyn TransactionGeneratorCreator>,
-    accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
+    source_accounts_pool: Arc<ObjectPool<LocalAccount>>,
+    destination_accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
 }
 
 impl AccountsPoolWrapperCreator {
     pub fn new(
         creator: Box<dyn TransactionGeneratorCreator>,
-        accounts_pool: Arc<RwLock<Vec<LocalAccount>>>,
+        source_accounts_pool: Arc<ObjectPool<LocalAccount>>,
+        destination_accounts_pool: Option<Arc<ObjectPool<LocalAccount>>>,
     ) -> Self {
         Self {
             creator,
-            accounts_pool,
+            source_accounts_pool,
+            destination_accounts_pool,
         }
     }
 }
 
 impl TransactionGeneratorCreator for AccountsPoolWrapperCreator {
     fn create_transaction_generator(&self) -> Box<dyn TransactionGenerator> {
         Box::new(AccountsPoolWrapperGenerator::new(
+            StdRng::from_entropy(),
             self.creator.create_transaction_generator(),
-            self.accounts_pool.clone(),
+            self.source_accounts_pool.clone(),
+            self.destination_accounts_pool.clone(),
         ))
     }
 }

crates/transaction-generator-lib/src/batch_transfer.rs

@@ -1,21 +1,20 @@
 // Copyright © Aptos Foundation
 
-use crate::{TransactionGenerator, TransactionGeneratorCreator};
-use aptos_infallible::RwLock;
+use crate::{ObjectPool, TransactionGenerator, TransactionGeneratorCreator};
 use aptos_sdk::{
     move_types::account_address::AccountAddress,
     transaction_builder::{aptos_stdlib, TransactionFactory},
     types::{transaction::SignedTransaction, LocalAccount},
 };
-use rand::{rngs::StdRng, seq::SliceRandom, SeedableRng};
+use rand::{rngs::StdRng, SeedableRng};
 use std::sync::Arc;
 
 pub struct BatchTransferTransactionGenerator {
     rng: StdRng,
     batch_size: usize,
     send_amount: u64,
     txn_factory: TransactionFactory,
-    all_addresses: Arc<RwLock<Vec<AccountAddress>>>,
+    all_addresses: Arc<ObjectPool<AccountAddress>>,
 }
 
 impl BatchTransferTransactionGenerator {
@@ -24,7 +23,7 @@ impl BatchTransferTransactionGenerator {
         batch_size: usize,
         send_amount: u64,
         txn_factory: TransactionFactory,
-        all_addresses: Arc<RwLock<Vec<AccountAddress>>>,
+        all_addresses: Arc<ObjectPool<AccountAddress>>,
     ) -> Self {
         Self {
             rng,
@@ -46,10 +45,7 @@ impl TransactionGenerator for BatchTransferTransactionGenerator {
         for _ in 0..num_to_create {
             let receivers = self
                 .all_addresses
-                .read()
-                .choose_multiple(&mut self.rng, self.batch_size)
-                .cloned()
-                .collect::<Vec<_>>();
+                .clone_from_pool(self.batch_size, &mut self.rng);
             requests.push(
                 account.sign_with_transaction_builder(self.txn_factory.payload(
                     aptos_stdlib::aptos_account_batch_transfer(receivers, vec![
@@ -67,15 +63,15 @@ impl TransactionGenerator for BatchTransferTransactionGenerator {
 pub struct BatchTransferTransactionGeneratorCreator {
     txn_factory: TransactionFactory,
     amount: u64,
-    all_addresses: Arc<RwLock<Vec<AccountAddress>>>,
+    all_addresses: Arc<ObjectPool<AccountAddress>>,
     batch_size: usize,
 }
 
 impl BatchTransferTransactionGeneratorCreator {
     pub fn new(
         txn_factory: TransactionFactory,
         amount: u64,
-        all_addresses: Arc<RwLock<Vec<AccountAddress>>>,
+        all_addresses: Arc<ObjectPool<AccountAddress>>,
         batch_size: usize,
     ) -> Self {
         Self {