sui_rpc_loadgen/payload/
get_all_balances.rs1use crate::payload::{GetAllBalances, ProcessPayload, RpcCommandProcessor, SignerInfo};
5use anyhow::Result;
6use async_trait::async_trait;
7use futures::future::join_all;
8use sui_json_rpc_types::Balance;
9use sui_sdk::SuiClient;
10use sui_types::base_types::SuiAddress;
11
12use super::validation::chunk_entities;
13
14#[async_trait]
15impl<'a> ProcessPayload<'a, &'a GetAllBalances> for RpcCommandProcessor {
16 async fn process(
17 &'a self,
18 op: &'a GetAllBalances,
19 _signer_info: &Option<SignerInfo>,
20 ) -> Result<()> {
21 if op.addresses.is_empty() {
22 panic!("No addresses provided, skipping query");
23 }
24 let clients = self.get_clients().await?;
25 let chunked = chunk_entities(&op.addresses, Some(op.chunk_size));
26 for chunk in chunked {
27 let mut tasks = Vec::new();
28 for address in chunk {
29 for client in clients.iter() {
30 let owner_address = address;
31 let task = async move {
32 get_all_balances(client, owner_address).await.unwrap();
33 };
34 tasks.push(task);
35 }
36 }
37 join_all(tasks).await;
38 }
39 Ok(())
40 }
41}
42
43async fn get_all_balances(client: &SuiClient, owner_address: SuiAddress) -> Result<Vec<Balance>> {
44 let balances = client
45 .coin_read_api()
46 .get_all_balances(owner_address)
47 .await
48 .unwrap();
49 Ok(balances)
50}