Get token balances by querying accountBalance (#5326)

backend/canisters/market_maker/CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Switch ICDex over to using `Pool` mode ([#5325](https://github.com/open-chat-labs/open-chat/pull/5325))
+- Get token balances by querying `accountBalance` ([#5326](https://github.com/open-chat-labs/open-chat/pull/5326))
 
 ## [[2.0.1044](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1044-market_maker)] - 2024-02-03
 

backend/canisters/market_maker/impl/src/exchanges/icdex.rs

@@ -43,6 +43,10 @@ impl<M: Fn(MakeOrderRequest) + Send + Sync, C: Fn(CancelOrderRequest) + Send + S
         }
         Ok(())
     }
+
+    async fn account_balances(&self) -> CallResult<Vec<(CanisterId, u128)>> {
+        self.account_balances().await
+    }
 }
 
 pub async fn deposit_funds() {

backend/canisters/market_maker/impl/src/exchanges/mod.rs

@@ -1,6 +1,6 @@
 use async_trait::async_trait;
 use ic_cdk::api::call::CallResult;
-use types::{AggregatedOrders, CancelOrderRequest, MakeOrderRequest, MarketState, Order, TokenInfo};
+use types::{AggregatedOrders, CancelOrderRequest, CanisterId, MakeOrderRequest, MarketState, Order, TokenInfo};
 
 pub mod icdex;
 
@@ -13,6 +13,7 @@ pub trait Exchange: Send + Sync {
     async fn orderbook(&self) -> CallResult<AggregatedOrders>;
     async fn make_orders(&self, orders: Vec<MakeOrderRequest>) -> CallResult<()>;
     async fn cancel_orders(&self, orders: Vec<CancelOrderRequest>) -> CallResult<()>;
+    async fn account_balances(&self) -> CallResult<Vec<(CanisterId, u128)>>;
     async fn market_state(&self) -> CallResult<MarketState> {
         let (latest_price, my_open_orders, orderbook) =
             futures::future::try_join3(self.latest_price(), self.my_open_orders(), self.orderbook()).await?;

backend/canisters/market_maker/impl/src/jobs/calculate_balances.rs

@@ -1,10 +1,9 @@
 use crate::exchanges::Exchange;
 use crate::{mutate_state, read_state, CanisterBalances, RuntimeState};
 use ic_cdk::api::call::CallResult;
-use itertools::Itertools;
 use std::collections::BTreeMap;
 use std::time::Duration;
-use types::{CanisterId, Milliseconds, OrderType, TimestampMillis};
+use types::{CanisterId, Milliseconds, TimestampMillis};
 use utils::canister_timers::run_now_then_interval;
 use utils::time::HOUR_IN_MS;
 
@@ -15,19 +14,14 @@ pub fn start_job() {
 }
 
 fn run() {
-    let PrepareResult {
-        exchange_clients,
-        this_canister_id,
-        now,
-    } = read_state(prepare);
+    let PrepareResult { exchange_clients, now } = read_state(prepare);
     if !exchange_clients.is_empty() {
-        ic_cdk::spawn(run_async(exchange_clients, this_canister_id, now));
+        ic_cdk::spawn(run_async(exchange_clients, now));
     }
 }
 
 struct PrepareResult {
     exchange_clients: Vec<Box<dyn Exchange>>,
-    this_canister_id: CanisterId,
     now: TimestampMillis,
 }
 
@@ -42,32 +36,16 @@ fn prepare(state: &RuntimeState) -> PrepareResult {
 
     PrepareResult {
         exchange_clients,
-        this_canister_id: state.env.canister_id(),
         now: state.env.now(),
     }
 }
 
-async fn run_async(exchange_clients: Vec<Box<dyn Exchange>>, this_canister_id: CanisterId, now: TimestampMillis) {
-    let ledger_canister_ids: Vec<_> = exchange_clients
-        .iter()
-        .flat_map(|e| [e.quote_token().ledger, e.base_token().ledger])
-        .unique()
-        .collect();
-
-    let ledger_balance_futures = futures::future::try_join_all(
-        ledger_canister_ids
-            .into_iter()
-            .map(|l| get_ledger_balance(l, this_canister_id)),
-    );
-
-    let exchange_balance_futures = futures::future::try_join_all(exchange_clients.into_iter().map(get_exchange_balances));
-
-    if let Ok((ledger_balances, exchange_balances)) =
-        futures::future::try_join(ledger_balance_futures, exchange_balance_futures).await
+async fn run_async(exchange_clients: Vec<Box<dyn Exchange>>, now: TimestampMillis) {
+    if let Ok(exchange_balances) = futures::future::try_join_all(exchange_clients.into_iter().map(get_exchange_balances)).await
     {
         let mut balances = BTreeMap::new();
 
-        for (ledger_canister_id, balance) in ledger_balances.into_iter().chain(exchange_balances.into_iter().flatten()) {
+        for (ledger_canister_id, balance) in exchange_balances.into_iter().flatten() {
             *balances.entry(ledger_canister_id).or_default() += balance;
         }
 
@@ -84,31 +62,6 @@ async fn run_async(exchange_clients: Vec<Box<dyn Exchange>>, this_canister_id: C
     }
 }
 
-async fn get_ledger_balance(ledger_canister_id: CanisterId, this_canister_id: CanisterId) -> CallResult<(CanisterId, u128)> {
-    let balance = icrc_ledger_canister_c2c_client::icrc1_balance_of(ledger_canister_id, &this_canister_id.into()).await?;
-
-    Ok((ledger_canister_id, balance.0.try_into().unwrap()))
-}
-
 async fn get_exchange_balances(exchange_client: Box<dyn Exchange>) -> CallResult<Vec<(CanisterId, u128)>> {
-    let my_open_orders = exchange_client.my_open_orders().await?;
-
-    let mut quote_token_balance = 0;
-    let mut base_token_balance = 0;
-
-    let quote_token = exchange_client.quote_token();
-
-    for order in my_open_orders {
-        match order.order_type {
-            OrderType::Bid => {
-                quote_token_balance += (order.amount * order.price / 10u64.pow(quote_token.decimals as u32)) as u128
-            }
-            OrderType::Ask => base_token_balance += order.amount as u128,
-        }
-    }
-
-    Ok(vec![
-        (exchange_client.quote_token().ledger, quote_token_balance),
-        (exchange_client.base_token().ledger, base_token_balance),
-    ])
+    exchange_client.account_balances().await
 }

backend/external_canisters/icdex/api/src/queries/accountBalance.rs

@@ -0,0 +1,16 @@
+use candid::{CandidType, Nat};
+use serde::{Deserialize, Serialize};
+
+pub type Args = String;
+
+#[derive(CandidType, Serialize, Deserialize)]
+pub struct Response {
+    pub token0: Balance,
+    pub token1: Balance,
+}
+
+#[derive(CandidType, Serialize, Deserialize)]
+pub struct Balance {
+    pub locked: Nat,
+    pub available: Nat,
+}

backend/external_canisters/icdex/api/src/queries/mod.rs

@@ -1,4 +1,6 @@
 #[allow(non_snake_case)]
+pub mod accountBalance;
+#[allow(non_snake_case)]
 pub mod getTxAccount;
 pub mod level10;
 pub mod pending;

backend/external_canisters/icdex/c2c_client/src/lib.rs

@@ -1,8 +1,9 @@
 #![allow(non_snake_case)]
-use canister_client::generate_candid_c2c_call_tuple_args;
+use canister_client::{generate_candid_c2c_call, generate_candid_c2c_call_tuple_args};
 use icdex_canister::*;
 
 // Queries
+generate_candid_c2c_call!(accountBalance);
 generate_candid_c2c_call_tuple_args!(getTxAccount);
 generate_candid_c2c_call_tuple_args!(level10);
 generate_candid_c2c_call_tuple_args!(pending);

backend/libraries/icdex_client/src/lib.rs

@@ -116,6 +116,22 @@ impl<M: Fn(MakeOrderRequest), C: Fn(CancelOrderRequest)> ICDexClient<M, C> {
         Ok(())
     }
 
+    pub async fn account_balances(&self) -> CallResult<Vec<(CanisterId, u128)>> {
+        let response =
+            icdex_canister_c2c_client::accountBalance(self.dex_canister_id, &self.this_canister_id.to_string()).await?;
+
+        Ok(vec![
+            (
+                self.base_token.ledger,
+                u128::try_from((response.token0.available + response.token0.locked).0).unwrap(),
+            ),
+            (
+                self.quote_token.ledger,
+                u128::try_from((response.token1.available + response.token1.locked).0).unwrap(),
+            ),
+        ])
+    }
+
     fn convert_order(&self, order: TradingOrder) -> Order {
         let (order_type, amount) = match order.remaining.quantity {
             OrderQuantity::Buy(n, _) => (OrderType::Bid, n),