feat: support sweeping rebalancer transactions (#827)

## 📝 Summary

<!--- A general summary of your changes -->

Support sweeping rebalancer transactions. If the "sweep" rule is
specified and after applying all "fund" rules the source balance exceeds
the specified max amount, transfer the funds to the provided
destination.

## ✅ I have completed the following steps:

* [x] Run `make lint`
* [x] Run `make test`
* [ ] Added tests (if applicable)

Author: efbig

crates/rbuilder-rebalancer/README.md

@@ -28,9 +28,18 @@ secret = "<SECRET>" # Private key or env variable name containing the private ke
 min_balance = "<U256>"
 
 [[rule]]
-description = "" # information description of the rebalancing rule  
+description = "" # informational description of the rebalancing rule  
 source_id = "<ACCOUNT ID>" # rebalancing account ID referencing an account entry in `accounts`
 destination = "<ADDRESS>" # destination target address
-destination_min_balance = "2500000000000000000" # minimum balance. after going below it, the account will be topped up
+type = "fund" # funding type rule 
+destination_min_balance = "2500000000000000000" # minimum destination balance. after going below it, the destination account will be topped up
 destination_target_balance = "5000000000000000000" # the target balance to top up to
-```
+
+[[rule]]
+description = "" # informational description of the rebalancing rule  
+source_id = "<ACCOUNT ID>" # rebalancing account ID referencing an account entry in `accounts`
+destination = "<ADDRESS>" # destination address
+type = "sweep" # sweeping type rule 
+source_target_balance = "2500000000000000000" # the target balance that should be remained after sweep
+source_max_balance = "5000000000000000000" # maximum source balance. after exceeding it, the source account balance will be swept
+```

crates/rbuilder-rebalancer/src/bin/rbuilder-rebalancer.rs

@@ -1,7 +1,10 @@
 use alloy_provider::ProviderBuilder;
 use alloy_signer_local::PrivateKeySigner;
 use clap::Parser;
-use rbuilder_rebalancer::{config::RebalancerConfig, rebalancer::Rebalancer};
+use rbuilder_rebalancer::{
+    config::{RebalancerConfig, RebalancerRuleType},
+    rebalancer::Rebalancer,
+};
 use std::{path::PathBuf, str::FromStr, time::Duration};
 use tracing::*;
 
@@ -30,13 +33,37 @@ impl Cli {
         }
 
         for rule in &config.rules {
-            if rule.destination_min_balance >= rule.destination_target_balance {
-                eyre::bail!("Invalid configuration for rule `{}`: minimum balance must be lower than the target", rule.description);
-            }
-
             if !config.accounts.iter().any(|acc| acc.id == rule.source_id) {
                 eyre::bail!("Invalid configuration for rule `{}`: account entry is missing for source account {}", rule.description, rule.source_id);
             }
+
+            match rule.ty {
+                RebalancerRuleType::Fund {
+                    destination_target_balance,
+                    destination_min_balance,
+                } => {
+                    if destination_min_balance >= destination_target_balance {
+                        eyre::bail!("Invalid configuration for rule `{}`: minimum balance must be lower than the target", rule.description);
+                    }
+                }
+                RebalancerRuleType::Sweep {
+                    source_target_balance,
+                    source_max_balance,
+                } => {
+                    if source_target_balance >= source_max_balance {
+                        eyre::bail!("Invalid configuration for rule `{}`: target balance must be lower than the max", rule.description);
+                    }
+
+                    let source = config
+                        .accounts
+                        .iter()
+                        .find(|acc| acc.id == rule.source_id)
+                        .expect("exists");
+                    if source_target_balance < source.min_balance {
+                        eyre::bail!("Invalid configuration for rule `{}`: target balance must be higher than the min source", rule.description);
+                    }
+                }
+            };
         }
 
         let rpc_provider = ProviderBuilder::new().connect(&config.rpc_url).await?;

crates/rbuilder-rebalancer/src/config.rs

@@ -67,10 +67,38 @@ pub struct RebalancerRule {
     pub source_id: String,
     /// Destination address.
     pub destination: Address,
-    /// Destination target balance.
-    pub destination_target_balance: U256,
-    /// Destination minimum threshold balance after which rebalancing will be triggered.
-    pub destination_min_balance: U256,
+    /// Rebalancer rule type and details.
+    #[serde(flatten)]
+    pub ty: RebalancerRuleType,
+}
+
+#[derive(PartialEq, Eq, Debug, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum RebalancerRuleType {
+    Fund {
+        /// Destination target balance.
+        destination_target_balance: U256,
+        /// Destination minimum threshold balance after which funding will be triggered.
+        destination_min_balance: U256,
+    },
+    Sweep {
+        /// Source target balance.
+        source_target_balance: U256,
+        /// Source maximum threshold balance after which sweeping will be triggered.
+        source_max_balance: U256,
+    },
+}
+
+impl RebalancerRuleType {
+    /// Returns `true` if the rule is of type [`RebalancerRuleType::Fund`].
+    pub fn is_fund(&self) -> bool {
+        matches!(self, Self::Fund { .. })
+    }
+
+    /// Returns `true` if the rule is of type [`RebalancerRuleType::Sweep`].
+    pub fn is_sweep(&self) -> bool {
+        matches!(self, Self::Sweep { .. })
+    }
 }
 
 #[cfg(test)]
@@ -84,16 +112,25 @@ mod tests {
             rpc_url = ""
             builder_url = ""
             transfer_max_priority_fee_per_gas = "123"
-            
+
             env_filter = "info"
             log_color = true
 
             [[rule]]
             description = ""
             source_id = ""
             destination = "0x0000000000000000000000000000000000000000"
+            type = "fund"
             destination_target_balance = "1"
             destination_min_balance = "1"
+
+            [[rule]]
+            description = ""
+            source_id = ""
+            destination = "0x0000000000000000000000000000000000000000"
+            type = "sweep"
+            source_target_balance = "1"
+            source_max_balance = "1"
         "#,
         )
         .unwrap();
@@ -105,13 +142,26 @@ mod tests {
                 transfer_max_priority_fee_per_gas: U256::from(123),
                 logger: LoggerConfig::dev(),
                 accounts: Vec::new(),
-                rules: Vec::from([RebalancerRule {
-                    description: String::new(),
-                    source_id: String::new(),
-                    destination: Address::ZERO,
-                    destination_min_balance: U256::from(1),
-                    destination_target_balance: U256::from(1),
-                }])
+                rules: Vec::from([
+                    RebalancerRule {
+                        description: String::new(),
+                        source_id: String::new(),
+                        destination: Address::ZERO,
+                        ty: RebalancerRuleType::Fund {
+                            destination_target_balance: U256::from(1),
+                            destination_min_balance: U256::from(1),
+                        }
+                    },
+                    RebalancerRule {
+                        description: String::new(),
+                        source_id: String::new(),
+                        destination: Address::ZERO,
+                        ty: RebalancerRuleType::Sweep {
+                            source_target_balance: U256::from(1),
+                            source_max_balance: U256::from(1),
+                        }
+                    }
+                ])
             }
         );
     }

crates/rbuilder-rebalancer/src/rebalancer.rs

@@ -18,7 +18,7 @@ use std::{cell::OnceCell, time::Duration};
 use tokio::time;
 use tracing::*;
 
-use crate::config::{RebalancerAccount, RebalancerRule};
+use crate::config::{RebalancerAccount, RebalancerRule, RebalancerRuleType};
 
 pub struct Rebalancer<P> {
     provider: P,
@@ -109,19 +109,28 @@ impl<P: Provider> Rebalancer<P> {
         let accounts = self.fetch_accounts(header.hash).await?;
         debug!(target: "rebalancer", number = header.number, hash = %header.hash, accounts = accounts.len(), "Updated account infos for tracked accounts");
         let mut transfers_by_source = HashMap::<String, Vec<Transfer>>::default();
+
+        // First evaluate all destination funding rules.
         for rule in &self.rules {
+            let RebalancerRuleType::Fund {
+                destination_target_balance,
+                destination_min_balance,
+            } = rule.ty
+            else {
+                continue;
+            };
+
             let destination_balance = accounts
                 .get(&rule.destination)
-                .ok_or(eyre::eyre!("missing account for {}", rule.destination))?
+                .ok_or(eyre::eyre!("missing account {}", rule.destination))?
                 .balance;
 
-            if destination_balance > rule.destination_min_balance {
-                trace!(target: "rebalancer", number = header.number, hash = %header.hash, %rule.description, %rule.destination, %rule.destination_min_balance, %destination_balance, "Rebalancing destination balance above minimum");
+            if destination_balance > destination_min_balance {
+                trace!(target: "rebalancer", number = header.number, hash = %header.hash, %rule.description, %rule.destination, %destination_min_balance, %destination_balance, "Rebalancing destination balance above minimum");
                 continue;
             }
 
-            let destination_target_delta = rule
-                .destination_target_balance
+            let destination_target_delta = destination_target_balance
                 .checked_sub(destination_balance)
                 .expect("misconfiguration");
 
@@ -136,6 +145,56 @@ impl<P: Provider> Rebalancer<P> {
                 .push(transfer);
         }
 
+        // Then evaluate all source sweeping rules.
+        for rule in &self.rules {
+            let RebalancerRuleType::Sweep {
+                source_target_balance,
+                source_max_balance,
+            } = rule.ty
+            else {
+                continue;
+            };
+
+            let source = self
+                .accounts
+                .get(&rule.source_id)
+                .ok_or(eyre::eyre!("missing source {}", rule.source_id))?;
+            let source_address = source.secret.address();
+            let source_balance = accounts
+                .get(&source_address)
+                .ok_or(eyre::eyre!("missing account {source_address}"))?
+                .balance;
+
+            let total_amount_out = transfers_by_source
+                .get(&rule.source_id)
+                .map_or(U256::ZERO, |transfers| {
+                    transfers.iter().map(|t| t.amount).sum()
+                });
+
+            if source_balance
+                .checked_sub(total_amount_out)
+                .is_none_or(|balance| balance <= source_max_balance)
+            {
+                trace!(target: "rebalancer", number = header.number, hash = %header.hash, %rule.description, %rule.destination, %source_max_balance, %source_balance, %total_amount_out, "Rebalancing source balance below maximum");
+                continue;
+            }
+
+            let source_target_delta = source_balance
+                .checked_sub(total_amount_out)
+                .unwrap() // checked above
+                .checked_sub(source_target_balance)
+                .expect("misconfiguration");
+            let transfer = Transfer {
+                destination: rule.destination,
+                amount: source_target_delta,
+                description: rule.description.clone(),
+            };
+            transfers_by_source
+                .entry(rule.source_id.clone())
+                .or_default()
+                .push(transfer);
+        }
+
         for (source_id, transfers) in transfers_by_source {
             let total_amount_out = transfers.iter().map(|t| t.amount).sum();