iota_rpc_loadgen/payload/
get_all_balances.rs

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