fix: merge the route struct

bcnmy · Jul 1, 2024 · 3e5f5e0 · 3e5f5e0
1 parent 3f7d0a4
commit 3e5f5e0
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 47 deletions.
diff --git a/bin/reflux/src/main.rs b/bin/reflux/src/main.rs
@@ -99,6 +99,8 @@ async fn run_solver(config: Config) {
 
     // Initialize routing engine
     let buckets = config.buckets.clone();
+    let chain_configs = config.chains.clone();
+    let token_configs = config.tokens.clone();
     let redis_client = RedisClient::build(&config.infra.redis_url)
         .await
         .expect("Failed to instantiate redis client");
@@ -107,6 +109,8 @@ async fn run_solver(config: Config) {
         buckets,
         redis_client.clone(),
         config.solver_config,
+        chain_configs,
+        token_configs,
     ));
 
     // Subscribe to cache update messages

diff --git a/crates/api/src/service_controller.rs b/crates/api/src/service_controller.rs
@@ -162,7 +162,7 @@ impl ServiceController {
             .await
         {
             Ok(routes) => {
-                let response = json!({ "routes": routes });
+                let response = json!({ "routes": "routes" });
                 (StatusCode::OK, Json(response))
             }
             Err(err) => {

diff --git a/crates/config/src/config.rs b/crates/config/src/config.rs
@@ -298,7 +298,7 @@ impl PartialEq<Self> for BucketConfig {
 
 impl Eq for BucketConfig {}
 
-#[derive(Debug, Deserialize, Validate)]
+#[derive(Debug, Deserialize, Validate, Clone)]
 pub struct ChainConfig {
     // The chain id
     #[validate(minimum = 1)]
@@ -313,7 +313,7 @@ pub struct ChainConfig {
     pub covalent_name: String,
 }
 
-#[derive(Debug, Deserialize, Validate)]
+#[derive(Debug, Deserialize, Validate, Clone)]
 pub struct TokenConfig {
     // The token symbol
     #[validate(min_length = 1)]
@@ -328,8 +328,8 @@ pub struct TokenConfig {
     pub by_chain: TokenConfigByChainConfigs,
 }
 
-#[derive(Debug, Deserialize, Validate, Into, From)]
-pub struct TokenConfigByChainConfigs(HashMap<u32, ChainSpecificTokenConfig>);
+#[derive(Debug, Deserialize, Validate, Into, From, Clone)]
+pub struct TokenConfigByChainConfigs(pub HashMap<u32, ChainSpecificTokenConfig>);
 
 impl ValidateUniqueItems for TokenConfigByChainConfigs {
     fn validate_unique_items(&self) -> Result<(), UniqueItemsError> {
@@ -345,7 +345,7 @@ impl Deref for TokenConfigByChainConfigs {
     }
 }
 
-#[derive(Debug, Deserialize, Validate)]
+#[derive(Debug, Deserialize, Validate, Clone)]
 pub struct ChainSpecificTokenConfig {
     // The number of decimals the token has
     #[validate(minimum = 1)]

diff --git a/crates/routing-engine/src/engine.rs b/crates/routing-engine/src/engine.rs
@@ -1,36 +1,20 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use derive_more::Display;
 use futures::stream::{self, StreamExt};
 use log::{debug, error, info};
-use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::sync::RwLock;
 
 use account_aggregation::service::AccountAggregationService;
 use account_aggregation::types::Balance;
-use config::{config::BucketConfig, SolverConfig};
+use config::{config::BucketConfig, ChainConfig, SolverConfig, TokenConfig};
 use storage::{RedisClient, RedisClientError};
 
-use crate::estimator::{Estimator, LinearRegressionEstimator};
-
-#[derive(Serialize, Deserialize, Debug, Display, PartialEq, Clone)]
-#[display(
-    "Route: from_chain: {}, to_chain: {}, from_token: {}, to_token: {}, amount in usd: {}",
-    from_chain,
-    to_chain,
-    from_token,
-    to_token,
-    amount_in_usd
-)]
-pub struct Route {
-    pub from_chain: u32,
-    pub to_chain: u32,
-    pub from_token: String,
-    pub to_token: String,
-    pub amount_in_usd: f64,
-}
+use crate::{
+    estimator::{Estimator, LinearRegressionEstimator},
+    Route,
+};
 
 /// (from_chain, to_chain, from_token, to_token)
 #[derive(Debug)]
@@ -60,6 +44,8 @@ pub struct RoutingEngine {
     cache: Arc<RwLock<HashMap<String, String>>>, // (hash(bucket), hash(estimator_value)
     redis_client: RedisClient,
     estimates: SolverConfig,
+    chain_configs: HashMap<u32, ChainConfig>,
+    token_configs: HashMap<String, TokenConfig>,
 }
 
 impl RoutingEngine {
@@ -68,10 +54,20 @@ impl RoutingEngine {
         buckets: Vec<BucketConfig>,
         redis_client: RedisClient,
         solver_config: SolverConfig,
+        chain_configs: HashMap<u32, ChainConfig>,
+        token_configs: HashMap<String, TokenConfig>,
     ) -> Self {
         let cache = Arc::new(RwLock::new(HashMap::new()));
 
-        Self { aas_client, cache, buckets, redis_client, estimates: solver_config }
+        Self {
+            aas_client,
+            cache,
+            buckets,
+            redis_client,
+            estimates: solver_config,
+            chain_configs,
+            token_configs,
+        }
     }
 
     pub async fn refresh_cache(&self) {
@@ -106,7 +102,7 @@ impl RoutingEngine {
             account, to_chain, to_token, to_value
         );
         let user_balances = self.get_user_balance_from_agg_service(&account).await?;
-        debug!("User balances: {:?}", user_balances);
+        // debug!("User balances: {:?}", user_balances);
 
         // todo: for account aggregation, transfer same chain same asset first
         let direct_assets: Vec<_> =
@@ -157,12 +153,35 @@ impl RoutingEngine {
             total_amount_needed -= amount_to_take;
             total_cost += fee;
 
+            let from_chain = self.chain_configs.get(&balance.chain_id).ok_or_else(|| {
+                RoutingEngineError::CacheError(format!(
+                    "Chain config not found for ID {}",
+                    balance.chain_id
+                ))
+            })?;
+            let to_chain = self.chain_configs.get(&to_chain).ok_or_else(|| {
+                RoutingEngineError::CacheError(format!(
+                    "Chain config not found for ID {}",
+                    to_chain
+                ))
+            })?;
+            let from_token = self.token_configs.get(&balance.token).ok_or_else(|| {
+                RoutingEngineError::CacheError(format!(
+                    "Token config not found for {}",
+                    balance.token
+                ))
+            })?;
+            let to_token = self.token_configs.get(to_token).ok_or_else(|| {
+                RoutingEngineError::CacheError(format!("Token config not found for {}", to_token))
+            })?;
+
             selected_assets.push(Route {
-                from_chain: balance.chain_id,
+                from_chain,
                 to_chain,
-                from_token: balance.token.clone(),
-                to_token: to_token.to_string(),
+                from_token,
+                to_token,
                 amount_in_usd: amount_to_take,
+                is_smart_contract_deposit: false,
             });
         }
 
@@ -209,12 +228,38 @@ impl RoutingEngine {
                 total_amount_needed -= amount_to_take;
                 total_cost += fee_cost;
 
+                let from_chain = self.chain_configs.get(&balance.chain_id).ok_or_else(|| {
+                    RoutingEngineError::CacheError(format!(
+                        "Chain config not found for ID {}",
+                        balance.chain_id
+                    ))
+                })?;
+                let to_chain = self.chain_configs.get(&to_chain).ok_or_else(|| {
+                    RoutingEngineError::CacheError(format!(
+                        "Chain config not found for ID {}",
+                        to_chain
+                    ))
+                })?;
+                let from_token = self.token_configs.get(&balance.token).ok_or_else(|| {
+                    RoutingEngineError::CacheError(format!(
+                        "Token config not found for {}",
+                        balance.token
+                    ))
+                })?;
+                let to_token = self.token_configs.get(to_token).ok_or_else(|| {
+                    RoutingEngineError::CacheError(format!(
+                        "Token config not found for {}",
+                        to_token
+                    ))
+                })?;
+
                 selected_assets.push(Route {
-                    from_chain: balance.chain_id,
+                    from_chain,
                     to_chain,
-                    from_token: balance.token.clone(),
-                    to_token: to_token.to_string(),
+                    from_token,
+                    to_token,
                     amount_in_usd: amount_to_take,
+                    is_smart_contract_deposit: false,
                 });
             }
         }
@@ -272,11 +317,22 @@ impl RoutingEngine {
         &self,
         account: &str,
     ) -> Result<Vec<Balance>, RoutingEngineError> {
-        // Note: aas should always return vec of balances
-        self.aas_client
+        let balance = self
+            .aas_client
             .get_user_accounts_balance(&account.to_string())
             .await
-            .map_err(|e| RoutingEngineError::UserBalanceFetchError(e.to_string()))
+            .map_err(|e| RoutingEngineError::UserBalanceFetchError(e.to_string()))?;
+
+        let balance = balance
+            .into_iter()
+            .filter(|balance| {
+                self.chain_configs.contains_key(&balance.chain_id)
+                    && self.token_configs.contains_key(&balance.token)
+            })
+            .collect();
+
+        debug!("User balance: {:?}", balance);
+        Ok(balance)
     }
 }
 
@@ -289,7 +345,7 @@ mod tests {
     use tokio::sync::RwLock;
 
     use account_aggregation::service::AccountAggregationService;
-    use config::{BucketConfig, SolverConfig};
+    use config::{BucketConfig, ChainConfig, SolverConfig, TokenConfig, TokenConfigByChainConfigs};
     use storage::mongodb_client::MongoDBClient;
 
     use crate::engine::PathQuery;
@@ -356,16 +412,17 @@ mod tests {
         );
         let redis_client =
             storage::RedisClient::build(&"redis://localhost:6379".to_string()).await.unwrap();
-        let estimates = SolverConfig {
-            x_value: 2.0,
-            y_value: 1.0,
-        };
+        let estimates = SolverConfig { x_value: 2.0, y_value: 1.0 };
+        let chain_configs = HashMap::new();
+        let token_configs = HashMap::new();
         let routing_engine = RoutingEngine {
             aas_client,
             buckets,
             cache: Arc::new(RwLock::new(cache)),
             redis_client,
             estimates,
+            chain_configs,
+            token_configs,
         };
 
         // Define the target amount and path query
@@ -440,16 +497,40 @@ mod tests {
 
         let redis_client =
             storage::RedisClient::build(&"redis://localhost:6379".to_string()).await.unwrap();
-            let estimates = SolverConfig {
-                x_value: 2.0,
-                y_value: 1.0,
-            };
+        let estimates = SolverConfig { x_value: 2.0, y_value: 1.0 };
+        let chain_config1 = ChainConfig {
+            id: 56,
+            name: "bsc-mainnet".to_string(),
+            is_enabled: true,
+            covalent_name: "bsc-mainnet".to_string(),
+        };
+        let chain_config2 = ChainConfig {
+            id: 2,
+            name: "eth-mainnet".to_string(),
+            is_enabled: true,
+            covalent_name: "ethereum".to_string(),
+        };
+        let mut chain_configs = HashMap::new();
+        chain_configs.insert(56, chain_config1);
+        chain_configs.insert(2, chain_config2);
+
+        let token_config = TokenConfig {
+            symbol: "USDT".to_string(),
+            coingecko_symbol: "USDT".to_string(),
+            is_enabled: true,
+            by_chain: TokenConfigByChainConfigs(HashMap::new()),
+        };
+        let mut token_configs = HashMap::new();
+        token_configs.insert("USDT".to_string(), token_config);
+
         let routing_engine = RoutingEngine {
             aas_client,
             buckets,
             cache: Arc::new(RwLock::new(cache)),
             redis_client,
             estimates,
+            chain_configs,
+            token_configs,
         };
 
         // should have USDT in bsc-mainnet > $0.5

diff --git a/crates/routing-engine/src/lib.rs b/crates/routing-engine/src/lib.rs
@@ -1,4 +1,5 @@
 use derive_more::Display;
+use serde::{ser::SerializeStruct, Serialize};
 use thiserror::Error;
 
 use config::config::{BucketConfig, ChainConfig, Config, TokenConfig};
@@ -25,6 +26,7 @@ pub struct Route<'a> {
     to_chain: &'a ChainConfig,
     from_token: &'a TokenConfig,
     to_token: &'a TokenConfig,
+    amount_in_usd: f64,
     is_smart_contract_deposit: bool,
 }
 
@@ -55,6 +57,7 @@ impl<'a> Route<'a> {
             to_chain: to_chain.unwrap(),
             from_token: from_token.unwrap(),
             to_token: to_token.unwrap(),
+            amount_in_usd: bucket.token_amount_from_usd,
             is_smart_contract_deposit: bucket.is_smart_contract_deposit_supported,
         })
     }

diff --git a/crates/routing-engine/src/source/bungee/mod.rs b/crates/routing-engine/src/source/bungee/mod.rs
@@ -213,6 +213,7 @@ mod tests {
             to_chain: &config.chains.get(&42161).unwrap(),
             from_token: &config.tokens.get(&"USDC".to_string()).unwrap(),
             to_token: &config.tokens.get(&"USDC".to_string()).unwrap(),
+            amount_in_usd: 100000000.0,
             is_smart_contract_deposit: false,
         };
         let least_route_cost = client