fix token supplies

Author: Fraccaman

chain/src/main.rs

@@ -17,7 +17,6 @@ use chain::services::{
 use chrono::{NaiveDateTime, Utc};
 use clap::Parser;
 use deadpool_diesel::postgres::Object;
-use futures::future::FutureExt;
 use futures::stream::StreamExt;
 use namada_sdk::time::DateTimeUtc;
 use orm::migrations::CustomMigrationSource;
@@ -287,13 +286,14 @@ async fn crawling_fn(
             native_token.clone(),
         ));
     let addresses = block.addresses_with_balance_change(&native_token);
+    let all_changed_tokens_supply = addresses
+        .iter()
+        .map(|bc| bc.token.clone())
+        .collect::<HashSet<_>>();
 
-    let token_supplies = first_block_in_epoch
-        .eq(&block_height)
-        .then(|| query_token_supplies(&client, &conn, &native_token, epoch))
-        .future()
-        .await
-        .transpose()?;
+    let token_supplies =
+        query_token_supplies(&client, &all_changed_tokens_supply, epoch)
+            .await?;
 
     let validators_addresses = if first_block_in_epoch.eq(&block_height) {
         let previous_epoch = epoch.saturating_sub(1);
@@ -492,7 +492,7 @@ async fn crawling_fn(
 
                 repository::balance::insert_token_supplies(
                     transaction_conn,
-                    token_supplies.into_iter().flatten(),
+                    token_supplies,
                 )?;
 
                 repository::balance::insert_ibc_rate_limits(
@@ -641,10 +641,12 @@ async fn try_initial_query(
         .into_rpc_error()
     };
     let token_supplies_fut = async {
-        let native_token = namada_service::get_native_token(client)
-            .await
-            .into_rpc_error()?;
-        query_token_supplies(client, conn, &native_token, epoch).await
+        query_token_supplies(
+            client,
+            &tokens.iter().cloned().collect::<HashSet<_>>(),
+            epoch,
+        )
+        .await
     };
 
     let (rate_limits, token_supplies) =
@@ -874,22 +876,31 @@ async fn get_block(
     Ok((block, tm_block_response, epoch))
 }
 
-async fn query_non_native_supplies(
+async fn query_token_supplies(
     client: &HttpClient,
-    conn: &Object,
+    tokens: &HashSet<Token>,
     epoch: u32,
 ) -> Result<Vec<TokenSupply>, MainError> {
-    let token_addresses = db_service::get_non_native_tokens(conn)
-        .await
-        .into_db_error()?;
-
-    let mut buffer = Vec::with_capacity(1);
-
-    let mut stream = futures::stream::iter(token_addresses)
-        .map(|address| async move {
-            namada_service::get_token_supply(client, address, epoch)
+    let mut buffer = Vec::with_capacity(tokens.len());
+
+    let mut stream = futures::stream::iter(tokens)
+        .map(|token| async move {
+            match token {
+                Token::Ibc(ibc_token) => namada_service::get_token_supply(
+                    client,
+                    ibc_token.address.to_string(),
+                    epoch,
+                )
                 .await
-                .into_rpc_error()
+                .into_rpc_error(),
+                Token::Native(address) => {
+                    namada_service::get_native_token_supply(
+                        client, address, epoch,
+                    )
+                    .await
+                    .into_rpc_error()
+                }
+            }
         })
         .buffer_unordered(32);
 
@@ -900,23 +911,3 @@ async fn query_non_native_supplies(
 
     Ok(buffer)
 }
-
-async fn query_token_supplies(
-    client: &HttpClient,
-    conn: &Object,
-    native_token: &Id,
-    epoch: u32,
-) -> Result<Vec<TokenSupply>, MainError> {
-    let native_fut =
-        namada_service::get_native_token_supply(client, native_token, epoch)
-            .map(|result| result.into_rpc_error());
-
-    let non_native_fut = query_non_native_supplies(client, conn, epoch);
-
-    let (native, non_native) = futures::try_join!(native_fut, non_native_fut)?;
-
-    let mut supplies = non_native;
-    supplies.push(native);
-
-    Ok(supplies)
-}

chain/src/repository/balance.rs

@@ -283,9 +283,9 @@ mod tests {
                 address: Id::Account(
                     "tnam1q9rhgyv3ydq0zu3whnftvllqnvhvhm270qxay5tn".to_string(),
                 ),
-                trace: Id::Account(
+                trace: Some(Id::Account(
                     "tnam1q9rhgyv3ydq0zu3whnftvllqnvhvhm270qxay5tn".to_string(),
-                ),
+                )),
             });
 
             let new_amount = Amount::from(NamadaAmount::from_u64(200));

chain/src/services/namada.rs

@@ -218,7 +218,7 @@ async fn query_ibc_tokens(
 
                 let token = IbcToken {
                     address: Id::from(ibc_token_addr),
-                    trace: Id::IbcTrace(ibc_trace.clone()),
+                    trace: Some(Id::IbcTrace(ibc_trace.clone())),
                 };
 
                 tokens.insert(token);

orm/src/token.rs

@@ -51,7 +51,7 @@ impl IbcTokenDb {
         match token {
             Token::Ibc(token) => Some(IbcTokenDb {
                 address: token.address.to_string(),
-                ibc_trace: token.trace.to_string(),
+                ibc_trace: token.clone().trace.unwrap_or_default().to_string(),
             }),
             Token::Native(_) => None,
         }

shared/src/block.rs

@@ -746,7 +746,11 @@ impl Block {
                         transfer_changes.0.keys().map(|account| {
                             BalanceChange::new(
                                 Id::Account(account.owner()),
-                                Token::Native(Id::Account(account.token())),
+                                Token::new(
+                                    &account.token(),
+                                    None,
+                                    &native_token.to_string(),
+                                ),
                             )
                         })
                     })

shared/src/token.rs

@@ -8,7 +8,7 @@ use crate::id::Id;
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
 pub struct IbcToken {
     pub address: Id,
-    pub trace: Id,
+    pub trace: Option<Id>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
@@ -26,6 +26,23 @@ impl Display for Token {
     }
 }
 
+impl Token {
+    pub fn new(
+        token: &str,
+        ibc_trace: Option<String>,
+        native_address: &str,
+    ) -> Self {
+        if !token.eq(native_address) {
+            Token::Ibc(IbcToken {
+                address: Id::Account(token.to_string()),
+                trace: ibc_trace.map(Id::IbcTrace),
+            })
+        } else {
+            Token::Native(Id::Account(token.to_string()))
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct IbcRateLimit {
     /// Address of the token in Namada

shared/src/utils.rs

@@ -440,7 +440,7 @@ fn get_ibc_token(
             token_address.clone(),
             crate::token::Token::Ibc(crate::token::IbcToken {
                 address: token_address.into(),
-                trace: Id::IbcTrace(ibc_trace),
+                trace: Some(Id::IbcTrace(ibc_trace)),
             }),
         )
     } else {

webserver/src/response/chain.rs

@@ -100,7 +100,7 @@ impl From<SharedToken> for Token {
             }),
             SharedToken::Ibc(token) => Token::Ibc(IbcToken {
                 address: token.address.to_string(),
-                trace: token.trace.to_string(),
+                trace: token.trace.unwrap_or_default().to_string(),
             }),
         }
     }

webserver/src/service/chain.rs

@@ -59,7 +59,7 @@ impl ChainService {
             .map(|(token, ibc_token)| match ibc_token {
                 Some(ibc_token) => Token::Ibc(IbcToken {
                     address: Id::Account(ibc_token.address),
-                    trace: Id::IbcTrace(ibc_token.ibc_trace),
+                    trace: Some(Id::IbcTrace(ibc_token.ibc_trace)),
                 }),
                 None => Token::Native(Id::Account(token.address)),
             })