sui_rpc_loadgen/payload/
get_all_balances.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use 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}