From c736bca8edc96efa603d580472b7d26c0d7224c7 Mon Sep 17 00:00:00 2001
From: Evan B <ebatsell@gmail.com>
Date: Fri, 3 Jan 2025 15:16:28 -0500
Subject: [PATCH] Emit VRT supply in cranker (#174)

Co-authored-by: Coach Chuck <169060940+coachchucksol@users.noreply.github.com>
---
 Cargo.lock              |  2 ++
 crankers/Cargo.toml     |  2 ++
 crankers/src/metrics.rs | 64 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0b03a43f..338996ba 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2581,6 +2581,8 @@ dependencies = [
  "solana-rpc-client",
  "solana-rpc-client-api",
  "solana-sdk",
+ "spl-associated-token-account",
+ "spl-token",
  "thiserror",
  "tokio",
 ]
diff --git a/crankers/Cargo.toml b/crankers/Cargo.toml
index 73b76c3c..6e98bd96 100644
--- a/crankers/Cargo.toml
+++ b/crankers/Cargo.toml
@@ -33,5 +33,7 @@ solana-metrics = { workspace = true }
 solana-rpc-client = { workspace = true }
 solana-rpc-client-api = { workspace = true }
 solana-sdk = { workspace = true }
+spl-associated-token-account = { workspace = true }
+spl-token = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
diff --git a/crankers/src/metrics.rs b/crankers/src/metrics.rs
index c769a0f8..5792f2fc 100644
--- a/crankers/src/metrics.rs
+++ b/crankers/src/metrics.rs
@@ -1,8 +1,12 @@
+use std::collections::HashMap;
+
 use jito_jsm_core::get_epoch;
 use jito_vault_core::config::Config;
 use solana_metrics::datapoint_info;
 use solana_rpc_client::nonblocking::rpc_client::RpcClient;
-use solana_sdk::signature::Keypair;
+use solana_sdk::{program_pack::Pack, pubkey::Pubkey, signature::Keypair};
+use spl_associated_token_account::get_associated_token_address;
+use spl_token::state::{Account as TokenAccount, Mint};
 
 use crate::vault_handler::VaultHandler;
 
@@ -50,6 +54,64 @@ pub async fn emit_vault_metrics(
         })
         .count() as i64;
 
+    let vrt_mint_pubkeys: Vec<Pubkey> = vaults.iter().map(|(_, vault)| vault.vrt_mint).collect();
+    let vrt_mint_accounts = rpc_client.get_multiple_accounts(&vrt_mint_pubkeys).await?;
+    let vrt_mint_map: HashMap<Pubkey, Mint> = vrt_mint_pubkeys
+        .into_iter()
+        .zip(vrt_mint_accounts.into_iter())
+        .filter_map(|(pubkey, account)| account.map(|acc| (pubkey, acc)))
+        .map(|(pubkey, account)| {
+            let mint = Mint::unpack(&account.data).expect("Failed to unpack Mint");
+            (pubkey, mint)
+        })
+        .collect();
+
+    let st_ata_pubkeys: Vec<Pubkey> = vaults
+        .iter()
+        .map(|(vault_address, vault)| {
+            get_associated_token_address(vault_address, &vault.supported_mint)
+        })
+        .collect();
+
+    let st_ata_accounts = rpc_client.get_multiple_accounts(&st_ata_pubkeys).await?;
+    let st_ata_map: HashMap<Pubkey, TokenAccount> = st_ata_pubkeys
+        .into_iter()
+        .zip(st_ata_accounts.into_iter())
+        .filter_map(|(pubkey, account)| {
+            account.map(|acc| {
+                (
+                    pubkey,
+                    TokenAccount::unpack(&acc.data).expect("Failed to unpack TokenAccount"),
+                )
+            })
+        })
+        .collect();
+
+    for (address, vault) in vaults.iter() {
+        let vrt_mint = vrt_mint_map
+            .get(&vault.vrt_mint)
+            .ok_or_else(|| anyhow::anyhow!("Mint not found in map"))?;
+
+        let st_deposit_account: &TokenAccount = st_ata_map
+            .get(&get_associated_token_address(
+                address,
+                &vault.supported_mint,
+            ))
+            .ok_or_else(|| anyhow::anyhow!("ST deposit account not found in map"))?;
+
+        datapoint_info!(
+            "restaking-vault-supply",
+            "vault" => address.to_string(),
+            "vrt_mint" => vault.vrt_mint.to_string(),
+            ("slot", slot as i64, i64),
+            ("slot_index", slot_index as i64, i64),
+            ("vrt_supply_internal", vault.vrt_supply() as i64, i64),
+            ("vrt_supply_external", vrt_mint.supply as i64, i64),
+            ("st_supply_internal", vault.tokens_deposited() as i64, i64),
+            ("st_supply_external", st_deposit_account.amount as i64, i64),
+        );
+    }
+
     datapoint_info!(
         "restaking-vault-stats",
         ("slot", slot as i64, i64),