feat(torii): token balances subscription

crates/torii/core/src/executor/erc.rs

@@ -15,8 +15,9 @@
 use super::{ApplyBalanceDiffQuery, Executor};
 use crate::constants::{IPFS_CLIENT_MAX_RETRY, SQL_FELT_DELIMITER, TOKEN_BALANCE_TABLE};
 use crate::executor::LOG_TARGET;
+use crate::simple_broker::SimpleBroker;
 use crate::sql::utils::{felt_to_sql_string, sql_string_to_u256, u256_to_sql_string, I256};
-use crate::types::ContractType;
+use crate::types::{ContractType, TokenBalance};
 use crate::utils::fetch_content_from_ipfs;
 
 #[derive(Debug, Clone)]
@@ -159,18 +160,21 @@
         }
 
         // write the new balance to the database
-        sqlx::query(&format!(
+        let token_balance: TokenBalance = sqlx::query_as(&format!(
             "INSERT OR REPLACE INTO {TOKEN_BALANCE_TABLE} (id, contract_address, account_address, \
-             token_id, balance) VALUES (?, ?, ?, ?, ?)",
+             token_id, balance) VALUES (?, ?, ?, ?, ?) RETURNING *",
         ))
         .bind(id)
         .bind(contract_address)
         .bind(account_address)
         .bind(token_id)
         .bind(u256_to_sql_string(&balance))
-        .execute(&mut **tx)
+        .fetch_one(&mut **tx)
         .await?;
 
+        debug!(target: LOG_TARGET, token_balance = ?token_balance, "Applied balance diff");
+        SimpleBroker::publish(token_balance);
+
         Ok(())
     }
 

crates/torii/grpc/proto/world.proto

@@ -34,6 +34,12 @@ service World {
     // Update entity subscription
     rpc UpdateEventMessagesSubscription (UpdateEventMessagesSubscriptionRequest) returns (google.protobuf.Empty);
 
+    // Subscribe to token balance updates.
+    rpc SubscribeTokenBalances (RetrieveTokenBalancesRequest) returns (stream SubscribeTokenBalancesResponse);
+
+    // Update token balance subscription
+    rpc UpdateTokenBalancesSubscription (UpdateTokenBalancesSubscriptionRequest) returns (google.protobuf.Empty);
+
     // Retrieve entities
     rpc RetrieveEventMessages (RetrieveEventMessagesRequest) returns (RetrieveEntitiesResponse);
 
@@ -50,6 +56,24 @@ service World {
     rpc RetrieveTokenBalances (RetrieveTokenBalancesRequest) returns (RetrieveTokenBalancesResponse);
 }
 
+// A request to update a token balance subscription
+message UpdateTokenBalancesSubscriptionRequest {
+    // The subscription ID
+    uint64 subscription_id = 1;
+    // The list of contract addresses to subscribe to
+    repeated bytes contract_addresses = 2;
+    // The list of account addresses to subscribe to
+    repeated bytes account_addresses = 3;
+}
+
+// A response containing token balances
+message SubscribeTokenBalancesResponse {
+    // The subscription ID
+    uint64 subscription_id = 1;
+    // The token balance
+    types.TokenBalance balance = 2;
+}
+
 // A request to retrieve tokens
 message RetrieveTokensRequest {
     // The list of contract addresses to retrieve tokens for

crates/torii/grpc/src/server/mod.rs

@@ -34,6 +34,7 @@
 use starknet::providers::JsonRpcClient;
 use subscriptions::event::EventManager;
 use subscriptions::indexer::IndexerManager;
+use subscriptions::token_balance::TokenBalanceManager;
 use tokio::net::TcpListener;
 use tokio::sync::mpsc::{channel, Receiver};
 use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
@@ -55,11 +56,7 @@
 use crate::proto::types::LogicalOperator;
 use crate::proto::world::world_server::WorldServer;
 use crate::proto::world::{
-    RetrieveEntitiesStreamingResponse, RetrieveEventMessagesRequest, RetrieveTokenBalancesRequest,
-    RetrieveTokenBalancesResponse, RetrieveTokensRequest, RetrieveTokensResponse,
-    SubscribeEntitiesRequest, SubscribeEntityResponse, SubscribeEventMessagesRequest,
-    SubscribeEventsResponse, SubscribeIndexerRequest, SubscribeIndexerResponse,
-    UpdateEventMessagesSubscriptionRequest, WorldMetadataRequest, WorldMetadataResponse,
+    RetrieveEntitiesStreamingResponse, RetrieveEventMessagesRequest, RetrieveTokenBalancesRequest, RetrieveTokenBalancesResponse, RetrieveTokensRequest, RetrieveTokensResponse, SubscribeEntitiesRequest, SubscribeEntityResponse, SubscribeEventMessagesRequest, SubscribeEventsResponse, SubscribeIndexerRequest, SubscribeIndexerResponse, SubscribeTokenBalancesResponse, UpdateEventMessagesSubscriptionRequest, UpdateTokenBalancesSubscriptionRequest, WorldMetadataRequest, WorldMetadataResponse
 };
 use crate::proto::{self};
 use crate::types::schema::SchemaError;
@@ -123,6 +120,7 @@
     event_manager: Arc<EventManager>,
     state_diff_manager: Arc<StateDiffManager>,
     indexer_manager: Arc<IndexerManager>,
+    token_balance_manager: Arc<TokenBalanceManager>,
 }
 
 impl DojoWorld {
@@ -138,6 +136,7 @@
         let event_manager = Arc::new(EventManager::default());
         let state_diff_manager = Arc::new(StateDiffManager::default());
         let indexer_manager = Arc::new(IndexerManager::default());
+        let token_balance_manager = Arc::new(TokenBalanceManager::default());
 
         tokio::task::spawn(subscriptions::model_diff::Service::new_with_block_rcv(
             block_rx,
@@ -156,6 +155,8 @@
 
         tokio::task::spawn(subscriptions::indexer::Service::new(Arc::clone(&indexer_manager)));
 
+        tokio::task::spawn(subscriptions::token_balance::Service::new(Arc::clone(&token_balance_manager)));
+
         Self {
             pool,
             world_address,
@@ -165,6 +166,7 @@
             event_manager,
             state_diff_manager,
             indexer_manager,
+            token_balance_manager,
         }
     }
 }
@@ -1056,6 +1058,14 @@
         Ok(RetrieveTokenBalancesResponse { balances })
     }
 
+    async fn subscribe_token_balances(
+        &self,
+        contract_addresses: Vec<Felt>,
+        account_addresses: Vec<Felt>,
+    ) -> Result<Receiver<Result<proto::world::SubscribeTokenBalancesResponse, tonic::Status>>, Error> {
+        self.token_balance_manager.add_subscriber(contract_addresses, account_addresses).await
+    }
+
     async fn subscribe_indexer(
         &self,
         contract_address: Felt,
@@ -1508,6 +1518,8 @@
     Pin<Box<dyn Stream<Item = Result<SubscribeIndexerResponse, Status>> + Send>>;
 type RetrieveEntitiesStreamingResponseStream =
     Pin<Box<dyn Stream<Item = Result<RetrieveEntitiesStreamingResponse, Status>> + Send>>;
+type SubscribeTokenBalancesResponseStream =
+    Pin<Box<dyn Stream<Item = Result<SubscribeTokenBalancesResponse, Status>> + Send>>;
 
 #[tonic::async_trait]
 impl proto::world::world_server::World for DojoWorld {
@@ -1517,6 +1529,7 @@
     type SubscribeEventsStream = SubscribeEventsResponseStream;
     type SubscribeIndexerStream = SubscribeIndexerResponseStream;
     type RetrieveEntitiesStreamingStream = RetrieveEntitiesStreamingResponseStream;
+    type SubscribeTokenBalancesStream = SubscribeTokenBalancesResponseStream;
 
     async fn world_metadata(
         &self,
@@ -1619,6 +1632,42 @@
         Ok(Response::new(()))
     }
 
+    async fn subscribe_token_balances(
+        &self,
+        request: Request<RetrieveTokenBalancesRequest>,
+    ) -> ServiceResult<Self::SubscribeTokenBalancesStream> {
+        let RetrieveTokenBalancesRequest { contract_addresses, account_addresses } = request.into_inner();
+        let contract_addresses = contract_addresses
+            .iter()
+            .map(|address| Felt::from_bytes_be_slice(address))
+            .collect::<Vec<_>>();
+        let account_addresses = account_addresses
+            .iter()
+            .map(|address| Felt::from_bytes_be_slice(address))
+            .collect::<Vec<_>>();
+
+        let rx = self.subscribe_token_balances(contract_addresses, account_addresses).await.map_err(|e| Status::internal(e.to_string()))?;
+        Ok(Response::new(Box::pin(ReceiverStream::new(rx)) as Self::SubscribeTokenBalancesStream))
+    }
+
+    async fn update_token_balances_subscription(
+        &self,
+        request: Request<UpdateTokenBalancesSubscriptionRequest>,
+    ) -> ServiceResult<()> {
+        let UpdateTokenBalancesSubscriptionRequest { subscription_id, contract_addresses, account_addresses } = request.into_inner();
+        let contract_addresses = contract_addresses
+            .iter()
+            .map(|address| Felt::from_bytes_be_slice(address))
+            .collect::<Vec<_>>();
+        let account_addresses = account_addresses
+            .iter()
+            .map(|address| Felt::from_bytes_be_slice(address))
+            .collect::<Vec<_>>();
+
+        self.token_balance_manager.update_subscriber(subscription_id, contract_addresses, account_addresses).await;
+        Ok(Response::new(()))
+    }
+
     async fn retrieve_entities(
         &self,
         request: Request<RetrieveEntitiesRequest>,

crates/torii/grpc/src/server/subscriptions/mod.rs

@@ -9,6 +9,7 @@ pub mod event;
 pub mod event_message;
 pub mod indexer;
 pub mod model_diff;
+pub mod token_balance;
 
 pub(crate) fn match_entity_keys(
     id: Felt,