refactor(torii): load token_id into cache on startup

dojoengine · lambda-0x · Oct 2, 2024 · Oct 1, 2024 · Oct 2, 2024 · Oct 2, 2024
commit a7e55cea35040d0625a3265ed6fabf4b6b52b995
diff --git a/bin/torii/torii.toml b/bin/torii/torii.toml
@@ -3,13 +3,4 @@
 #     { type = "WORLD", address = "<WORLD_CONTRACT_ADDRESS>" },
 #     { type = "ERC20", address = "<ERC20_CONTRACT_ADDRESS>" },
 #     { type = "ERC721", address = "<ERC721_CONTRACT_ADDRESS>" },
-# ]
-
-contracts = [
-    # num-sepolia
-    { type = "WORLD", address = "0x50c46c6f69ea9f12b24fcdd1569a584ae8893db0c70ea7884b605384da78ebc" },
-    # LORDS
-    { type = "ERC20", address = "0x044e6bcc627e6201ce09f781d1aae44ea4c21c2fdef299e34fce55bef2d02210" },
-    # LSVR
-    { type = "ERC721", address = "0x0046506cccd3ffb110bc04445d25d94f1d26d645a30e6ac8a17c24e084bed08b" },
-]
+# ]
diff --git a/crates/torii/core/src/sql/cache.rs b/crates/torii/core/src/sql/cache.rs
@@ -2,13 +2,14 @@ use std::collections::{HashMap, HashSet};
 
 use dojo_types::schema::Ty;
 use dojo_world::contracts::abi::model::Layout;
-use sqlx::SqlitePool;
+use sqlx::{Pool, Sqlite, SqlitePool};
 use starknet_crypto::Felt;
 use tokio::sync::RwLock;
 
 use crate::error::{Error, ParseError, QueryError};
 use crate::model::{parse_sql_model_members, SqlModelMember};
 use crate::sql::utils::I256;
+use crate::types::ContractType;
 
 #[derive(Debug, Clone)]
 pub struct Model {
@@ -126,12 +127,24 @@ impl ModelCache {
 
 #[derive(Debug)]
 pub struct LocalCache {
-    pub erc_cache: HashMap<String, I256>,
+    pub erc_cache: HashMap<(ContractType, String), I256>,
     pub token_id_registry: HashSet<String>,
 }
 
 impl LocalCache {
-    pub fn new() -> Self {
+    pub async fn new(pool: Pool<Sqlite>) -> Self {
+        // read existing token_id's from balances table and cache them
+        let token_id_registry: Vec<(String,)> = sqlx::query_as("SELECT token_id FROM balances")
+            .fetch_all(&pool)
+            .await
+            .expect("Should be able to read token_id's from blances table");
+
+        let token_id_registry = token_id_registry.into_iter().map(|token_id| token_id.0).collect();
+
+        Self { erc_cache: HashMap::new(), token_id_registry }
+    }
+
+    pub fn empty() -> Self {
         Self { erc_cache: HashMap::new(), token_id_registry: HashSet::new() }
     }
 

diff --git a/crates/torii/core/src/sql/erc.rs b/crates/torii/core/src/sql/erc.rs
@@ -9,6 +9,7 @@ use super::query_queue::{Argument, QueryType};
 use super::utils::{sql_string_to_u256, u256_to_sql_string, I256};
 use super::{Sql, FELT_DELIMITER};
 use crate::sql::utils::{felt_and_u256_to_sql_string, felt_to_sql_string, felts_to_sql_string};
+use crate::types::ContractType;
 use crate::utils::utc_dt_string_from_timestamp;
 
 impl Sql {
@@ -40,18 +41,21 @@ impl Sql {
             block_timestamp,
         );
 
-        self.query_queue.execute_all().await?;
-
         if from_address != Felt::ZERO {
             // from_address/contract_address/
             let from_balance_id = felts_to_sql_string(&[from_address, contract_address]);
-            let from_balance = self.local_cache.erc_cache.entry(from_balance_id).or_default();
+            let from_balance = self
+                .local_cache
+                .erc_cache
+                .entry((ContractType::ERC20, from_balance_id))
+                .or_default();
             *from_balance -= I256::from(amount);
         }
 
         if to_address != Felt::ZERO {
             let to_balance_id = felts_to_sql_string(&[to_address, contract_address]);
-            let to_balance = self.local_cache.erc_cache.entry(to_balance_id).or_default();
+            let to_balance =
+                self.local_cache.erc_cache.entry((ContractType::ERC20, to_balance_id)).or_default();
             *to_balance += I256::from(amount);
         }
 
@@ -89,20 +93,26 @@ impl Sql {
             block_timestamp,
         );
 
-        self.query_queue.execute_all().await?;
-
         // from_address/contract_address:id
         if from_address != Felt::ZERO {
             let from_balance_id =
                 format!("{}{FELT_DELIMITER}{}", felt_to_sql_string(&from_address), &token_id);
-            let from_balance = self.local_cache.erc_cache.entry(from_balance_id).or_default();
+            let from_balance = self
+                .local_cache
+                .erc_cache
+                .entry((ContractType::ERC721, from_balance_id))
+                .or_default();
             *from_balance -= I256::from(1u8);
         }
 
         if to_address != Felt::ZERO {
             let to_balance_id =
                 format!("{}{FELT_DELIMITER}{}", felt_to_sql_string(&to_address), &token_id);
-            let to_balance = self.local_cache.erc_cache.entry(to_balance_id).or_default();
+            let to_balance = self
+                .local_cache
+                .erc_cache
+                .entry((ContractType::ERC721, to_balance_id))
+                .or_default();
             *to_balance += I256::from(1u8);
         }
 
@@ -319,11 +329,16 @@ impl Sql {
     }
 
     pub async fn apply_cache_diff(&mut self) -> Result<()> {
-        for (id_str, balance) in self.local_cache.erc_cache.iter() {
+        // execute the query queue to apply any queries registering token metadata
+        self.query_queue.execute_all().await?;
+
+        for ((contract_type, id_str), balance) in self.local_cache.erc_cache.iter() {
             let id = id_str.split(FELT_DELIMITER).collect::<Vec<&str>>();
-            match id.len() {
-                // account_address/contract_address:id => ERC721
-                2 => {
+            match contract_type {
+                ContractType::WORLD => unreachable!(),
+                ContractType::ERC721 => {
+                    // account_address/contract_address:id => ERC721
+                    assert!(id.len() == 2);
                     let account_address = id[0];
                     let token_id = id[1];
                     let mid = token_id.split(":").collect::<Vec<&str>>();
@@ -338,8 +353,9 @@ impl Sql {
                     )
                     .await?;
                 }
-                // account_address/contract_address/ => ERC20
-                3 => {
+                ContractType::ERC20 => {
+                    // account_address/contract_address/ => ERC20
+                    assert!(id.len() == 3);
                     let account_address = id[0];
                     let contract_address = id[1];
                     let token_id = id[1];
@@ -353,7 +369,6 @@ impl Sql {
                     )
                     .await?;
                 }
-                _ => unreachable!(),
             }
         }
 

diff --git a/crates/torii/core/src/sql/mod.rs b/crates/torii/core/src/sql/mod.rs
@@ -60,7 +60,7 @@ impl Clone for Sql {
             pool: self.pool.clone(),
             query_queue: QueryQueue::new(self.pool.clone()),
             model_cache: self.model_cache.clone(),
-            local_cache: LocalCache::new(),
+            local_cache: LocalCache::empty(),
         }
     }
 }
@@ -88,12 +88,14 @@ impl Sql {
 
         query_queue.execute_all().await?;
 
+        let local_cache = LocalCache::new(pool.clone()).await;
+
         Ok(Self {
             pool: pool.clone(),
             world_address,
             query_queue,
             model_cache: Arc::new(ModelCache::new(pool)),
-            local_cache: LocalCache::new(),
+            local_cache,
         })
     }
 
@@ -804,11 +806,7 @@ impl Sql {
             Ty::Enum(e) => {
                 if e.options.iter().all(
                     |o| {
-                        if let Ty::Tuple(t) = &o.ty {
-                            t.is_empty()
-                        } else {
-                            false
-                        }
+                        if let Ty::Tuple(t) = &o.ty { t.is_empty() } else { false }
                     },
                 ) {
                     return;

diff --git a/scripts/verify_db_balances/src/main.rs b/scripts/verify_db_balances/src/main.rs
@@ -115,7 +115,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
             if db_balance != starknet_balance {
                 error!(
-                    "Mismatch for account {} and contract {}: DB balance = {}, Starknet balance = {}",
+                    "Mismatch for account {} and contract {}: DB balance = {}, Starknet balance = \
+                     {}",
                     account_address, contract_address, db_balance, starknet_balance
                 );
             } else {