tensor_eigen/commands/
fees.rs

1use crate::{spinner::pb_with_len, FEE_SHARDS};
2
3use super::*;
4
5use std::{fs::File, str::FromStr};
6
7use solana_sdk::{
8    instruction::Instruction, pubkey, signer::Signer, system_instruction, transaction::Transaction,
9};
10
11const TFEE_PROGRAM_ID: Pubkey = pubkey!("TFEEgwDP6nn1s8mMX2tTNPPz8j2VomkphLUmyxKm17A");
12
13pub struct FeeParams {
14    pub keypair_path: Option<PathBuf>,
15    pub rpc_url: Option<String>,
16}
17
18pub fn generate_fee_shards() -> Result<()> {
19    let mut shards = vec![];
20
21    for i in 0..=255 {
22        let shard: &[u8] = &[i];
23        shards.push(
24            Pubkey::find_program_address(&[b"fee_vault", shard], &TFEE_PROGRAM_ID)
25                .0
26                .to_string(),
27        );
28    }
29
30    let file = File::create("fee_vault_shards.json")?;
31    serde_json::to_writer_pretty(&file, &shards)?;
32
33    Ok(())
34}
35
36pub fn fund_shards(args: FeeParams) -> Result<()> {
37    let config = CliConfig::new(args.keypair_path, args.rpc_url)?;
38
39    let rent_exempt_lamports = config.client.get_minimum_balance_for_rent_exemption(0)?;
40
41    // Convert FEE_SHARDS to Pubkeys
42    let shard_pubkeys: Vec<Pubkey> = FEE_SHARDS
43        .iter()
44        .filter_map(|s| Pubkey::from_str(s).ok())
45        .collect();
46
47    // Check balances and create transfer instructions only for underfunded shards
48    let mut instructions: Vec<Instruction> = Vec::new();
49    let pb = pb_with_len("shards checked", shard_pubkeys.len() as u64)?;
50
51    for pubkey in &shard_pubkeys {
52        let balance = config.client.get_balance(pubkey)?;
53        if balance < rent_exempt_lamports {
54            instructions.push(system_instruction::transfer(
55                &config.keypair.pubkey(),
56                pubkey,
57                rent_exempt_lamports - balance,
58            ));
59        }
60        pb.inc(1);
61    }
62    pb.finish_with_message("Finished checking shard balances");
63
64    if instructions.is_empty() {
65        println!("All shards are already funded.");
66        return Ok(());
67    }
68
69    // Pack instructions into transactions (15 instructions per transaction)
70    for chunk in instructions.chunks(15) {
71        let transaction = Transaction::new_signed_with_payer(
72            chunk,
73            Some(&config.keypair.pubkey()),
74            &[&config.keypair],
75            config.client.get_latest_blockhash()?,
76        );
77
78        config
79            .client
80            .send_and_confirm_transaction_with_spinner(&transaction)?;
81    }
82
83    Ok(())
84}
85
86pub fn get_shard_balances(args: FeeParams) -> Result<()> {
87    let config = CliConfig::new(args.keypair_path, args.rpc_url)?;
88
89    let shard_pubkeys: Vec<Pubkey> = FEE_SHARDS
90        .iter()
91        .filter_map(|s| Pubkey::from_str(s).ok())
92        .collect();
93
94    let rent_exempt_lamports = config.client.get_minimum_balance_for_rent_exemption(0)?;
95
96    let mut fully_funded_count = 0;
97
98    for pubkey in &shard_pubkeys {
99        let balance = config.client.get_balance(pubkey)?;
100        let status = if balance >= rent_exempt_lamports {
101            fully_funded_count += 1;
102            "✓"
103        } else {
104            "✗"
105        };
106        println!("{} {} {}", pubkey, balance, status);
107    }
108
109    if fully_funded_count == shard_pubkeys.len() {
110        println!("All shards fully funded!");
111    } else {
112        println!(
113            "{}/{} shards fully funded.",
114            fully_funded_count,
115            shard_pubkeys.len()
116        );
117    }
118
119    Ok(())
120}