chore: implement get_balance via grpc (#2935)

* chore: implement get_balance via grpc

Author: wbbradley

.github/workflows/code.yml

@@ -113,7 +113,7 @@ jobs:
     runs-on: ubuntu-ghcloud
     strategy:
       matrix:
-        grpc_migration_level: [0, 1]
+        grpc_migration_level: [0, 100]
     env:
       RUST_LOG: walrus=info,error,sui_=off,consensus_core=off
       OPENSSL_STATIC: "1"
@@ -161,11 +161,11 @@ jobs:
       - name: Run tests (seed=2)
         run: MSIM_TEST_SEED=2 cargo simtest simtest --profile simtest
       - name: Run tests (seed=3, grpc-client-enabled)
-        run: MSIM_TEST_SEED=3 WALRUS_GRPC_MIGRATION_LEVEL=1 cargo simtest simtest --profile simtest
+        run: MSIM_TEST_SEED=3 WALRUS_GRPC_MIGRATION_LEVEL=100 cargo simtest simtest --profile simtest
       - name: Run tests (seed=4, grpc-client-enabled)
-        run: MSIM_TEST_SEED=4 WALRUS_GRPC_MIGRATION_LEVEL=1 cargo simtest simtest --profile simtest
+        run: MSIM_TEST_SEED=4 WALRUS_GRPC_MIGRATION_LEVEL=100 cargo simtest simtest --profile simtest
       - name: Run tests (seed=5, grpc-client-enabled)
-        run: MSIM_TEST_SEED=5 WALRUS_GRPC_MIGRATION_LEVEL=1 cargo simtest simtest --profile simtest
+        run: MSIM_TEST_SEED=5 WALRUS_GRPC_MIGRATION_LEVEL=100 cargo simtest simtest --profile simtest
 
   simtests-build:
     name: Build all simtests

crates/walrus-e2e-tests/tests/test_client.rs

@@ -98,6 +98,7 @@ use walrus_sui::{
         SuiContractClient,
         retry_client::{RetriableSuiClient, retriable_sui_client::LazySuiClientBuilder},
     },
+    coin::Coin,
     config::WalletConfig,
     test_utils::{self, fund_addresses, wallet_for_testing},
     types::{
@@ -2527,19 +2528,19 @@ pub async fn test_select_coins_max_objects() -> TestResult {
         ExponentialBackoffConfig::default(),
     )?;
 
-    let balance = retry_client.get_balance(address, None).await?;
-    assert_eq!(balance.total_balance, u128::from(sui(4)));
+    let balance = retry_client.get_total_balance(address, Coin::SUI).await?;
+    assert_eq!(balance, sui(4));
 
     // The maximum number of coins that can be selected to reach the amount.
     let max_num_coins = 2;
 
     let result = retry_client
-        .select_coins(address, None, sui(1).into(), vec![], max_num_coins)
+        .select_coins(address, Coin::SUI, sui(1).into(), vec![], max_num_coins)
         .await;
     assert!(result.is_ok(), "1 SUI can be constructed with <= 2 coins");
 
     let result = retry_client
-        .select_coins(address, None, sui(3).into(), vec![], max_num_coins)
+        .select_coins(address, Coin::SUI, sui(3).into(), vec![], max_num_coins)
         .await;
     if let Err(error) = result {
         assert!(
@@ -2551,7 +2552,7 @@ pub async fn test_select_coins_max_objects() -> TestResult {
     }
 
     let result = retry_client
-        .select_coins(address, None, sui(5).into(), vec![], max_num_coins)
+        .select_coins(address, Coin::SUI, sui(5).into(), vec![], max_num_coins)
         .await;
     if let Err(error) = result {
         assert!(
@@ -2796,10 +2797,9 @@ async fn test_store_with_upload_relay_with_tip() {
 
     // Get initial balance of relay wallet to verify tip payment
     let initial_relay_balance = retry_client
-        .get_balance(relay_address, None)
+        .get_total_balance(relay_address, Coin::SUI)
         .await
-        .expect("get balance")
-        .total_balance;
+        .expect("get balance");
 
     const BLOB_SIZE: usize = 40000;
     match basic_store_and_read(
@@ -2819,10 +2819,9 @@ async fn test_store_with_upload_relay_with_tip() {
 
     // Verify that the relay wallet received a tip
     let final_relay_balance = retry_client
-        .get_balance(relay_address, None)
+        .get_total_balance(relay_address, Coin::SUI)
         .await
-        .expect("get balance")
-        .total_balance;
+        .expect("get balance");
 
     tracing::info!(
         "Relay address balance - Initial: {initial_relay_balance}, Final: {final_relay_balance}",
@@ -2839,8 +2838,7 @@ async fn test_store_with_upload_relay_with_tip() {
         encoded_blob_length_for_n_shards(n_shards, BLOB_SIZE as u64, EncodingType::RS2)
             .expect("encoded blob size should be valid");
 
-    let expected_tip_lower_bound =
-        u128::from(TIP_BASE + encoded_blob_size.div_ceil(1024) * TIP_MULTIPLIER);
+    let expected_tip_lower_bound = TIP_BASE + encoded_blob_size.div_ceil(1024) * TIP_MULTIPLIER;
     let actual_tip = final_relay_balance - initial_relay_balance;
 
     tracing::info!(

crates/walrus-sdk/src/client/upload_relay_client.rs

@@ -42,12 +42,7 @@ use crate::{
         metadata::{BlobMetadataApi, VerifiedBlobMetadataWithId},
     },
     sui::{
-        client::{
-            BlobPersistence,
-            CoinType,
-            SuiContractClient,
-            transaction_builder::WalrusPtbBuilder,
-        },
+        client::{BlobPersistence, SuiContractClient, transaction_builder::WalrusPtbBuilder},
         config::WalletConfig,
         types::{BlobEvent, BlobRegistered},
         wallet::Wallet,
@@ -221,6 +216,7 @@ impl UploadRelayClient {
             self.gas_budget,
             0, // No additional gas budget.
             tip_amount,
+            None,
         )
         .await?;
 

crates/walrus-service/src/client/multiplexer.rs

@@ -43,6 +43,7 @@ use walrus_sui::{
         SuiReadClient,
         retry_client::RetriableSuiClient,
     },
+    coin::Coin,
     config::load_wallet_context_from_path,
     types::move_structs::BlobWithAttribute,
     utils::create_wallet,
@@ -496,13 +497,13 @@ impl<'a> SubClientLoader<'a> {
             self.config.communication_config.sui_client_request_timeout,
         )?;
 
-        if should_refill(&sui_client, address, None, min_balance).await {
+        if should_refill(&sui_client, address, Coin::SUI, min_balance).await {
             self.refiller.send_gas_request(address).await?;
         } else {
             tracing::debug!(%address, "sub-wallet has enough SUI, skipping refill");
         }
 
-        if should_refill(&sui_client, address, Some(wal_coin_type), min_balance).await {
+        if should_refill(&sui_client, address, wal_coin_type, min_balance).await {
             self.refiller.send_wal_request(address).await?;
         } else {
             tracing::debug!(%address, "sub-wallet has enough WAL, skipping refill");

crates/walrus-service/src/client/refill.rs

@@ -11,7 +11,10 @@ use sui_sdk::types::base_types::SuiAddress;
 use tokio::{task::JoinHandle, time::MissedTickBehavior};
 use tracing::Instrument;
 use walrus_sdk::client::metrics::ClientMetrics;
-use walrus_sui::client::{SuiContractClient, retry_client::RetriableSuiClient};
+use walrus_sui::{
+    client::{SuiContractClient, retry_client::RetriableSuiClient},
+    coin::Coin,
+};
 
 /// Refills gas and WAL for the clients.
 #[derive(Debug, Clone)]
@@ -77,7 +80,7 @@ impl Refiller {
             addresses,
             period,
             sui_client,
-            None, // Use SUI
+            Coin::SUI,
             self.min_balance,
             move |refiller, address| {
                 let metrics = metrics.clone();
@@ -103,7 +106,7 @@ impl Refiller {
             addresses,
             period,
             sui_client,
-            Some(self.wal_coin_type()),
+            self.wal_coin_type(),
             self.min_balance,
             move |refiller, address| {
                 let metrics = metrics.clone();
@@ -122,7 +125,7 @@ impl Refiller {
         addresses: Vec<SuiAddress>,
         period: Duration,
         sui_client: RetriableSuiClient,
-        coin_type: Option<String>,
+        coin_type: &str,
         min_coin_value: u64,
         inner_action: F,
     ) -> JoinHandle<anyhow::Result<()>>
@@ -134,6 +137,7 @@ impl Refiller {
         interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
 
         let refiller = self.contract_client.clone();
+        let coin_type = coin_type.to_owned();
         tokio::spawn(async move {
             loop {
                 let span =
@@ -143,17 +147,10 @@ impl Refiller {
                     interval.tick().await;
                     let sui_client = &sui_client;
                     let _ = try_join_all(addresses.iter().cloned().map(|address| {
-                        let coin_type_inner = coin_type.clone();
                         let inner_fut = inner_action(refiller.clone(), address);
-
+                        let coin_type = coin_type.clone();
                         async move {
-                            if should_refill(
-                                sui_client,
-                                address,
-                                coin_type_inner.clone(),
-                                min_coin_value,
-                            )
-                            .await
+                            if should_refill(sui_client, address, &coin_type, min_coin_value).await
                             {
                                 inner_fut.await
                             } else {
@@ -182,11 +179,8 @@ impl Refiller {
     }
 
     /// The WAL coin type.
-    pub fn wal_coin_type(&self) -> String {
-        self.contract_client
-            .read_client()
-            .wal_coin_type()
-            .to_owned()
+    pub fn wal_coin_type(&self) -> &str {
+        self.contract_client.read_client().wal_coin_type()
     }
 
     /// Sends SUI to the specified address.
@@ -223,13 +217,13 @@ pub struct RefillHandles {
 pub async fn should_refill(
     sui_client: &RetriableSuiClient,
     address: SuiAddress,
-    coin_type: Option<String>,
+    coin_type: &str,
     min_balance: u64,
 ) -> bool {
     sui_client
-        .get_balance(address, coin_type)
+        .get_total_balance(address, coin_type)
         .await
-        .map(|balance| balance.total_balance < u128::from(min_balance))
+        .map(|balance| balance < min_balance)
         .inspect_err(|error| {
             tracing::debug!(
                 ?error,

crates/walrus-service/src/node/contract_service.rs

@@ -21,14 +21,14 @@ use walrus_core::{Epoch, PublicKey, messages::InvalidBlobCertificate};
 use walrus_sui::{
     client::{
         BlobObjectMetadata,
-        CoinType,
         FixedSystemParameters,
         ReadClient,
         SuiClientError,
         SuiClientMetricSet,
         SuiContractClient,
         SuiReadClient,
     },
+    coin::CoinType,
     types::{
         StorageNodeCap,
         UpdatePublicKeyParams,
@@ -346,8 +346,8 @@ async fn monitor_sui_balance(
         async {
             tracing::trace!("querying wallet SUI balance");
 
-            let balance_mist = match client.balance(CoinType::Sui).await {
-                Ok(balance_mist) => balance_mist,
+            let balance_mist = match client.total_balance(CoinType::Sui).await {
+                Ok(balance) => balance,
                 Err(error) => {
                     tracing::warn!(?error, "failed to get SUI balance, skipping interval");
                     return; // Exit the async closure, does not exit the outer function.

crates/walrus-stress/src/main.rs

@@ -28,7 +28,8 @@ use walrus_stress::single_client_workload::{
     single_client_workload_arg::SingleClientWorkloadArgs,
 };
 use walrus_sui::{
-    client::{CoinType, MIN_STAKING_THRESHOLD, ReadClient, SuiContractClient},
+    client::{MIN_STAKING_THRESHOLD, ReadClient, SuiContractClient},
+    coin::CoinType,
     config::WalletConfig,
     types::StorageNode,
     utils::SuiNetwork,
@@ -181,8 +182,8 @@ async fn run_stress(
     .context("Failed to load wallet context")?;
     let contract_client = client_config.new_contract_client(wallet, None).await?;
 
-    let wal_balance = contract_client.balance(CoinType::Wal).await?;
-    let sui_balance = contract_client.balance(CoinType::Sui).await?;
+    let wal_balance = contract_client.total_balance(CoinType::Wal).await?;
+    let sui_balance = contract_client.total_balance(CoinType::Sui).await?;
     tracing::info!(wal_balance, sui_balance, "initial balances");
 
     let refiller = Refiller::new(
@@ -240,8 +241,8 @@ async fn run_staking(config: ClientConfig, _metrics: Arc<ClientMetrics>) -> anyh
         tokio::time::sleep(restaking_period).await;
         let mut committee = contract_client.read_client().current_committee().await?;
         let current_epoch = committee.epoch;
-        let wal_balance = contract_client.balance(CoinType::Wal).await?;
-        let sui_balance = contract_client.balance(CoinType::Sui).await?;
+        let wal_balance = contract_client.total_balance(CoinType::Wal).await?;
+        let sui_balance = contract_client.total_balance(CoinType::Sui).await?;
         tracing::info!(current_epoch, ?mode, wal_balance, sui_balance, "woke up");
 
         match mode {

crates/walrus-sui/src/balance.rs

@@ -0,0 +1,81 @@
+// Copyright (c) Walrus Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+//! A module defining the Balance struct.
+
+use anyhow::Context;
+
+use crate::coin::Coin;
+
+/// An error type for balance retrieval errors.
+#[derive(Debug, thiserror::Error)]
+#[error("error retrieving balance information: {0}")]
+pub struct BalanceRetrievalError(anyhow::Error);
+
+impl From<anyhow::Error> for BalanceRetrievalError {
+    fn from(err: anyhow::Error) -> Self {
+        BalanceRetrievalError(err)
+    }
+}
+
+/// A struct representing the balance of a specific coin type.
+#[derive(Debug, Clone)]
+pub struct Balance {
+    /// The total balance across all coin objects.
+    total_balance: u64,
+    /// The number of coin objects or the actual coin objects.
+    coin_balance: CoinBalance,
+}
+
+/// An enum representing either the count of coin objects or the actual coin objects.
+#[derive(Debug, Clone)]
+pub enum CoinBalance {
+    /// The count of coin objects.
+    Count(usize),
+    /// The actual coin objects.
+    Coins(Vec<Coin>),
+}
+
+impl Balance {
+    /// Creates a new Balance from a coin type and a list of coins.
+    pub fn try_from_coins(coins: Vec<Coin>) -> Result<Self, BalanceRetrievalError> {
+        let total_balance: u128 = coins.iter().map(|coin| u128::from(coin.balance)).sum();
+        Ok(Self {
+            total_balance: u64::try_from(total_balance).context("total balance exceeds u64")?,
+            coin_balance: CoinBalance::Coins(coins),
+        })
+    }
+
+    /// Returns the number of coin objects.
+    pub fn coin_object_count(&self) -> usize {
+        match self.coin_balance {
+            CoinBalance::Count(count) => count,
+            CoinBalance::Coins(ref coins) => coins.len(),
+        }
+    }
+
+    /// Returns the total balance.
+    pub fn total_balance(&self) -> u64 {
+        self.total_balance
+    }
+
+    /// Take the coins if they are available.
+    pub fn coins(self) -> Option<Vec<Coin>> {
+        match self.coin_balance {
+            CoinBalance::Count(_) => None,
+            CoinBalance::Coins(coins) => Some(coins),
+        }
+    }
+}
+
+impl TryFrom<sui_sdk::rpc_types::Balance> for Balance {
+    type Error = BalanceRetrievalError;
+
+    fn try_from(balance: sui_sdk::rpc_types::Balance) -> Result<Self, BalanceRetrievalError> {
+        Ok(Self {
+            coin_balance: CoinBalance::Count(balance.coin_object_count),
+            total_balance: u64::try_from(balance.total_balance)
+                .context("total balance exceeds u64")?,
+        })
+    }
+}