refactor: use typed builder for state and args

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

crates/tap-agent/src/agent.rs

@@ -76,17 +76,16 @@ pub async fn start_agent(
 
     let config = SenderAccountConfig::from_config(&config);
 
-    let args = SenderAccountsManagerArgs {
-        config,
-        domain_separator,
-        pgpool,
-        indexer_allocations,
-        escrow_accounts,
-        escrow_subgraph,
-        network_subgraph,
-        sender_aggregator_endpoints: sender_aggregator_endpoints.clone(),
-        prefix: None,
-    };
+    let args = SenderAccountsManagerArgs::builder()
+        .config(config)
+        .domain_separator(domain_separator)
+        .pgpool(pgpool)
+        .indexer_allocations(indexer_allocations)
+        .escrow_accounts(escrow_accounts)
+        .network_subgraph(network_subgraph)
+        .escrow_subgraph(escrow_subgraph)
+        .sender_aggregator_endpoints(sender_aggregator_endpoints.clone())
+        .build();
 
     SenderAccountsManager::spawn(None, SenderAccountsManager, args)
         .await

crates/tap-agent/src/agent/sender_account.rs

@@ -24,6 +24,7 @@ use std::{
 use tap_core::rav::SignedRAV;
 use tokio::sync::watch::Receiver;
 use tracing::{error, warn, Level};
+use typed_builder::TypedBuilder;
 
 use crate::{
     adaptative_concurrency::AdaptiveLimiter,
@@ -35,8 +36,7 @@ use crate::{
         sender_allocation::SenderAllocationMessage,
         unaggregated_receipts::UnaggregatedReceipts,
     },
-    backoff::BackoffInfo,
-    tracker::{SenderFeeTracker, SimpleFeeTracker},
+    tracker::SenderFeeTracker,
 };
 
 mod state;
@@ -72,21 +72,22 @@ pub enum SenderAccountMessage {
     IsSchedulerEnabled(ractor::RpcReplyPort<bool>),
 }
 
+#[derive(TypedBuilder)]
 pub struct SenderAccountArgs {
-    pub config: &'static SenderAccountConfig,
-
-    pub pgpool: PgPool,
-    pub sender_id: Address,
-    pub escrow_accounts: Receiver<EscrowAccounts>,
-    pub indexer_allocations: Receiver<HashSet<Address>>,
-    pub escrow_subgraph: &'static SubgraphClient,
-    pub network_subgraph: &'static SubgraphClient,
-    pub domain_separator: Eip712Domain,
-    pub sender_aggregator_endpoint: Url,
-    pub allocation_ids: HashSet<Address>,
-    pub prefix: Option<String>,
-
-    pub retry_interval: Duration,
+    config: &'static SenderAccountConfig,
+
+    pgpool: PgPool,
+    sender_id: Address,
+    escrow_accounts: Receiver<EscrowAccounts>,
+    indexer_allocations: Receiver<HashSet<Address>>,
+    escrow_subgraph: &'static SubgraphClient,
+    network_subgraph: &'static SubgraphClient,
+    domain_separator: Eip712Domain,
+    sender_aggregator_endpoint: Url,
+    allocation_ids: HashSet<Address>,
+    prefix: Option<String>,
+
+    retry_interval: Duration,
 }
 
 pub struct SenderAccountConfig {
@@ -290,29 +291,25 @@ impl Actor for SenderAccount {
             .request_timeout(config.rav_request_timeout)
             .build(&sender_aggregator_endpoint)?;
 
-        let state = State {
-            prefix,
-            sender_fee_tracker: SenderFeeTracker::new(config.rav_request_buffer),
-            rav_tracker: SimpleFeeTracker::default(),
-            invalid_receipts_tracker: SimpleFeeTracker::default(),
-            allocation_ids: allocation_ids.clone(),
-            _indexer_allocations_handle,
-            _escrow_account_monitor,
-            scheduled_rav_request: None,
-            sender: sender_id,
-            denied,
-            sender_balance,
-            retry_interval,
-            adaptive_limiter: AdaptiveLimiter::new(INITIAL_RAV_REQUEST_CONCURRENT, 1..50),
-            escrow_accounts,
-            escrow_subgraph,
-            network_subgraph,
-            domain_separator,
-            pgpool,
-            sender_aggregator,
-            backoff_info: BackoffInfo::default(),
-            config,
-        };
+        let state = State::builder()
+            .prefix(prefix)
+            .sender_fee_tracker(SenderFeeTracker::new(config.rav_request_buffer))
+            .allocation_ids(allocation_ids.clone())
+            ._indexer_allocations_handle(_indexer_allocations_handle)
+            ._escrow_account_monitor(_escrow_account_monitor)
+            .sender(sender_id)
+            .denied(denied)
+            .sender_balance(sender_balance)
+            .retry_interval(retry_interval)
+            .adaptive_limiter(AdaptiveLimiter::new(INITIAL_RAV_REQUEST_CONCURRENT, 1..50))
+            .escrow_accounts(escrow_accounts)
+            .escrow_subgraph(escrow_subgraph)
+            .network_subgraph(network_subgraph)
+            .domain_separator(domain_separator)
+            .pgpool(pgpool)
+            .sender_aggregator(sender_aggregator)
+            .config(config)
+            .build();
 
         stream::iter(allocation_ids)
             // Create a sender allocation for each allocation

crates/tap-agent/src/agent/sender_account/state.rs

@@ -16,6 +16,7 @@ use std::{collections::HashSet, str::FromStr, time::Duration};
 use tap_core::rav::SignedRAV;
 use tokio::{sync::watch::Receiver, task::JoinHandle};
 use tracing::error;
+use typed_builder::TypedBuilder;
 
 use crate::{
     adaptative_concurrency::AdaptiveLimiter,
@@ -34,14 +35,18 @@ use super::{
     SenderAccountConfig, SenderAccountMessage, PENDING_RAV, SENDER_FEE_TRACKER, UNAGGREGATED_FEES,
 };
 
+#[derive(TypedBuilder)]
 pub struct State {
-    pub(super) prefix: Option<String>,
+    prefix: Option<String>,
     pub(super) sender_fee_tracker: SenderFeeTracker,
+    #[builder(default)]
     pub(super) rav_tracker: SimpleFeeTracker,
+    #[builder(default)]
     pub(super) invalid_receipts_tracker: SimpleFeeTracker,
     pub(super) allocation_ids: HashSet<Address>,
-    pub(super) _indexer_allocations_handle: JoinHandle<()>,
-    pub(super) _escrow_account_monitor: JoinHandle<()>,
+    _indexer_allocations_handle: JoinHandle<()>,
+    _escrow_account_monitor: JoinHandle<()>,
+    #[builder(default)]
     pub(super) scheduled_rav_request:
         Option<JoinHandle<Result<(), MessagingErr<SenderAccountMessage>>>>,
 
@@ -56,16 +61,17 @@ pub struct State {
     pub(super) adaptive_limiter: AdaptiveLimiter,
 
     // Receivers
-    pub(super) escrow_accounts: Receiver<EscrowAccounts>,
+    escrow_accounts: Receiver<EscrowAccounts>,
 
-    pub(super) escrow_subgraph: &'static SubgraphClient,
-    pub(super) network_subgraph: &'static SubgraphClient,
+    escrow_subgraph: &'static SubgraphClient,
+    network_subgraph: &'static SubgraphClient,
 
-    pub(super) domain_separator: Eip712Domain,
+    domain_separator: Eip712Domain,
     pub(super) pgpool: PgPool,
-    pub(super) sender_aggregator: jsonrpsee::http_client::HttpClient,
+    sender_aggregator: jsonrpsee::http_client::HttpClient,
 
     // Backoff info
+    #[builder(default)]
     pub(super) backoff_info: BackoffInfo,
 
     // Config
@@ -83,17 +89,17 @@ impl State {
             %allocation_id,
             "SenderAccount is creating allocation."
         );
-        let args = SenderAllocationArgs {
-            pgpool: self.pgpool.clone(),
-            allocation_id,
-            sender: self.sender,
-            escrow_accounts: self.escrow_accounts.clone(),
-            escrow_subgraph: self.escrow_subgraph,
-            domain_separator: self.domain_separator.clone(),
-            sender_account_ref: sender_account_ref.clone(),
-            sender_aggregator: self.sender_aggregator.clone(),
-            config: AllocationConfig::from_sender_config(self.config),
-        };
+        let args = SenderAllocationArgs::builder()
+            .pgpool(self.pgpool.clone())
+            .allocation_id(allocation_id)
+            .sender(self.sender)
+            .escrow_accounts(self.escrow_accounts.clone())
+            .escrow_subgraph(self.escrow_subgraph)
+            .domain_separator(self.domain_separator.clone())
+            .sender_account_ref(sender_account_ref.clone())
+            .sender_aggregator(self.sender_aggregator.clone())
+            .config(AllocationConfig::from_sender_config(self.config))
+            .build();
 
         SenderAllocation::spawn_linked(
             Some(self.format_sender_allocation(&allocation_id)),

crates/tap-agent/src/agent/sender_accounts_manager.rs

@@ -20,6 +20,7 @@ use tokio::{
     sync::watch::{self, Receiver},
 };
 use tracing::error;
+use typed_builder::TypedBuilder;
 
 mod receipt_watcher;
 mod state;
@@ -42,19 +43,22 @@ pub enum SenderAccountsManagerMessage {
     UpdateSenderAccounts(HashSet<Address>),
 }
 
+#[derive(TypedBuilder)]
 pub struct SenderAccountsManagerArgs {
-    pub config: &'static SenderAccountConfig,
-    pub domain_separator: Eip712Domain,
-
-    pub pgpool: PgPool,
-    pub indexer_allocations: Receiver<HashMap<Address, Allocation>>,
-    pub escrow_accounts: Receiver<EscrowAccounts>,
-    pub escrow_subgraph: &'static SubgraphClient,
-    pub network_subgraph: &'static SubgraphClient,
-    pub sender_aggregator_endpoints: HashMap<Address, Url>,
-
-    pub prefix: Option<String>,
+    config: &'static SenderAccountConfig,
+    domain_separator: Eip712Domain,
+
+    pgpool: PgPool,
+    indexer_allocations: Receiver<HashMap<Address, Allocation>>,
+    escrow_accounts: Receiver<EscrowAccounts>,
+    escrow_subgraph: &'static SubgraphClient,
+    network_subgraph: &'static SubgraphClient,
+    sender_aggregator_endpoints: HashMap<Address, Url>,
+
+    #[builder(default)]
+    prefix: Option<String>,
 }
+
 #[async_trait::async_trait]
 impl Actor for SenderAccountsManager {
     type Msg = SenderAccountsManagerMessage;
@@ -105,20 +109,18 @@ impl Actor for SenderAccountsManager {
                 async {}
             });
 
-        let mut state = State {
-            config,
-            domain_separator,
-            sender_ids: HashSet::new(),
-            new_receipts_watcher_handle: None,
-            _eligible_allocations_senders_handle,
-            pgpool,
-            indexer_allocations: allocations_rx,
-            escrow_accounts: escrow_accounts.clone(),
-            escrow_subgraph,
-            network_subgraph,
-            sender_aggregator_endpoints,
-            prefix: prefix.clone(),
-        };
+        let mut state = State::builder()
+            .config(config)
+            .domain_separator(domain_separator)
+            .sender_aggregator_endpoints(sender_aggregator_endpoints)
+            .indexer_allocations(allocations_rx)
+            .pgpool(pgpool)
+            ._eligible_allocations_senders_handle(_eligible_allocations_senders_handle)
+            .escrow_subgraph(escrow_subgraph)
+            .network_subgraph(network_subgraph)
+            .escrow_accounts(escrow_accounts.clone())
+            .prefix(prefix.clone())
+            .build();
         let sender_allocation = select! {
             sender_allocation = state.get_pending_sender_allocation_id() => sender_allocation,
             _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => {

crates/tap-agent/src/agent/sender_accounts_manager/state.rs

@@ -16,21 +16,27 @@ use reqwest::Url;
 use sqlx::PgPool;
 use tokio::sync::watch::Receiver;
 use tracing::{error, warn};
+use typed_builder::TypedBuilder;
 
+#[derive(TypedBuilder)]
 pub struct State {
-    pub(super) sender_ids: HashSet<Address>,
-    pub(super) new_receipts_watcher_handle: Option<tokio::task::JoinHandle<()>>,
-    pub(super) _eligible_allocations_senders_handle: JoinHandle<()>,
-
-    pub(super) config: &'static SenderAccountConfig,
-    pub(super) domain_separator: Eip712Domain,
-    pub(super) pgpool: PgPool,
-    pub(super) indexer_allocations: Receiver<HashSet<Address>>,
-    pub(super) escrow_accounts: Receiver<EscrowAccounts>,
-    pub(super) escrow_subgraph: &'static SubgraphClient,
-    pub(super) network_subgraph: &'static SubgraphClient,
-    pub(super) sender_aggregator_endpoints: HashMap<Address, Url>,
-    pub(super) prefix: Option<String>,
+    #[builder(default)]
+    pub sender_ids: HashSet<Address>,
+
+    #[builder(default)]
+    pub new_receipts_watcher_handle: Option<tokio::task::JoinHandle<()>>,
+    _eligible_allocations_senders_handle: JoinHandle<()>,
+
+    config: &'static SenderAccountConfig,
+    domain_separator: Eip712Domain,
+    pgpool: PgPool,
+    indexer_allocations: Receiver<HashSet<Address>>,
+    escrow_accounts: Receiver<EscrowAccounts>,
+    escrow_subgraph: &'static SubgraphClient,
+    network_subgraph: &'static SubgraphClient,
+    sender_aggregator_endpoints: HashMap<Address, Url>,
+    #[builder(default)]
+    prefix: Option<String>,
 }
 
 impl State {
@@ -201,26 +207,27 @@ impl State {
         sender_id: &Address,
         allocation_ids: HashSet<Address>,
     ) -> Result<SenderAccountArgs> {
-        Ok(SenderAccountArgs {
-            config: self.config,
-            pgpool: self.pgpool.clone(),
-            sender_id: *sender_id,
-            escrow_accounts: self.escrow_accounts.clone(),
-            indexer_allocations: self.indexer_allocations.clone(),
-            escrow_subgraph: self.escrow_subgraph,
-            network_subgraph: self.network_subgraph,
-            domain_separator: self.domain_separator.clone(),
-            sender_aggregator_endpoint: self
-                .sender_aggregator_endpoints
-                .get(sender_id)
-                .ok_or(anyhow!(
-                    "No sender_aggregator_endpoints found for sender {}",
-                    sender_id
-                ))?
-                .clone(),
-            allocation_ids,
-            prefix: self.prefix.clone(),
-            retry_interval: Duration::from_secs(30),
-        })
+        Ok(SenderAccountArgs::builder()
+            .config(self.config)
+            .pgpool(self.pgpool.clone())
+            .sender_id(*sender_id)
+            .escrow_accounts(self.escrow_accounts.clone())
+            .indexer_allocations(self.indexer_allocations.clone())
+            .escrow_subgraph(self.escrow_subgraph)
+            .network_subgraph(self.network_subgraph)
+            .domain_separator(self.domain_separator.clone())
+            .sender_aggregator_endpoint(
+                self.sender_aggregator_endpoints
+                    .get(sender_id)
+                    .ok_or(anyhow!(
+                        "No sender_aggregator_endpoints found for sender {}",
+                        sender_id
+                    ))?
+                    .clone(),
+            )
+            .allocation_ids(allocation_ids)
+            .prefix(self.prefix.clone())
+            .retry_interval(Duration::from_secs(30))
+            .build())
     }
 }