From 88e0624ef07635f5a5d854688a93d309e9fd0fb6 Mon Sep 17 00:00:00 2001 From: Haider-Ali-DS Date: Fri, 12 Jul 2024 22:48:42 +0500 Subject: [PATCH] add funds checking chronicle functionality unregisters if chronicle have less funds before submitting transaction check chronicle fund in start --- chronicle/src/lib.rs | 2 ++ chronicle/src/main.rs | 4 ++++ chronicle/src/mock.rs | 4 ++++ chronicle/src/tasks/spawner.rs | 34 ++++++++++++++++++++++++++++++++-- primitives/src/lib.rs | 2 ++ tc-subxt/src/lib.rs | 12 ++++++++++++ 6 files changed, 56 insertions(+), 2 deletions(-) diff --git a/chronicle/src/lib.rs b/chronicle/src/lib.rs index e5c3b559fe..752b2cd74b 100644 --- a/chronicle/src/lib.rs +++ b/chronicle/src/lib.rs @@ -33,6 +33,7 @@ pub struct ChronicleConfig { pub timechain_keyfile: PathBuf, pub target_url: String, pub target_keyfile: PathBuf, + pub target_min_balance: u128, pub tss_keyshare_cache: PathBuf, } @@ -60,6 +61,7 @@ pub async fn run_chronicle( let task_spawner_params = TaskSpawnerParams { tss: tss_tx, blockchain: chain, + min_balance: config.target_min_balance, network: subchain, network_id: config.network_id, url: config.target_url, diff --git a/chronicle/src/main.rs b/chronicle/src/main.rs index 29e88e9333..2a30a0ae9a 100644 --- a/chronicle/src/main.rs +++ b/chronicle/src/main.rs @@ -22,6 +22,9 @@ pub struct ChronicleArgs { /// key file for connector wallet #[clap(long)] pub target_keyfile: PathBuf, + /// key file for connector wallet + #[clap(long)] + pub target_min_balance: u128, /// Metadata version to use to connect to timechain node. #[clap(long)] pub timechain_metadata: Option, @@ -48,6 +51,7 @@ impl ChronicleArgs { network_id: self.network_id, network_keyfile: self.network_keyfile, network_port: self.network_port, + target_min_balance: self.target_min_balance, timechain_metadata: self.timechain_metadata.unwrap_or_default(), timechain_url: self.timechain_url, timechain_keyfile: self.timechain_keyfile, diff --git a/chronicle/src/mock.rs b/chronicle/src/mock.rs index fc3aa23fa7..f01823a6c9 100644 --- a/chronicle/src/mock.rs +++ b/chronicle/src/mock.rs @@ -386,6 +386,10 @@ impl Runtime for Mock { Ok(()) } + async fn submit_unregister_member(&self) -> Result<()> { + Ok(()) + } + async fn submit_heartbeat(&self) -> Result<()> { Ok(()) } diff --git a/chronicle/src/tasks/spawner.rs b/chronicle/src/tasks/spawner.rs index 7f4f92b2c5..9f356aa47a 100644 --- a/chronicle/src/tasks/spawner.rs +++ b/chronicle/src/tasks/spawner.rs @@ -8,6 +8,7 @@ use rosetta_client::Wallet; use rosetta_config_ethereum::{query::GetLogs, CallResult, FilterBlockOption}; use rosetta_core::{BlockOrIdentifier, ClientEvent}; use schnorr_evm::VerifyingKey; +use std::time::Duration; use std::{ future::Future, num::NonZeroU64, @@ -22,11 +23,13 @@ use time_primitives::{ }; use time_primitives::{IGateway, Msg}; use tokio::sync::Mutex; +use tokio::time::sleep; #[derive(Clone)] pub struct TaskSpawnerParams { pub tss: mpsc::Sender, pub blockchain: String, + pub min_balance: u128, pub network: String, pub network_id: NetworkId, pub url: String, @@ -38,6 +41,7 @@ pub struct TaskSpawnerParams { pub struct TaskSpawner { tss: mpsc::Sender, wallet: Arc, + wallet_min_balance: u128, wallet_guard: Arc>, substrate: S, network_id: NetworkId, @@ -58,13 +62,29 @@ where ) .await?, ); - Ok(Self { + + let spawner = Self { tss: params.tss, + wallet_min_balance: params.min_balance, wallet, wallet_guard: Arc::new(Mutex::new(())), substrate: params.substrate, network_id: params.network_id, - }) + }; + + while !spawner.is_balance_available().await? { + sleep(Duration::from_secs(10)).await; + tracing::warn!("Chronicle balance is too low, retrying..."); + } + + Ok(spawner) + } + + /// + /// Checks if wallet have enough balance + async fn is_balance_available(&self) -> Result { + let balance = self.wallet.balance().await?; + Ok(balance > self.wallet_min_balance) } /// @@ -242,6 +262,16 @@ where } async fn write(self, task_id: TaskId, function: Function) -> Result<()> { + let is_balance_available = self.is_balance_available().await?; + if !is_balance_available { + // unregister member + if let Err(e) = self.substrate.submit_unregister_member().await { + tracing::error!("Failed to unregister member: {:?}", e); + }; + tracing::warn!("Chronicle balance too low, exiting"); + std::process::exit(1); + } + let submission = async move { match function { Function::EvmDeploy { bytecode } => { diff --git a/primitives/src/lib.rs b/primitives/src/lib.rs index 9368931f75..251792d5c2 100644 --- a/primitives/src/lib.rs +++ b/primitives/src/lib.rs @@ -183,6 +183,8 @@ pub trait Runtime: Clone + Send + Sync + 'static { stake_amount: u128, ) -> Result<()>; + async fn submit_unregister_member(&self) -> Result<()>; + async fn submit_heartbeat(&self) -> Result<()>; async fn submit_commitment( diff --git a/tc-subxt/src/lib.rs b/tc-subxt/src/lib.rs index 320dbd6ae3..1c1e0fff4d 100644 --- a/tc-subxt/src/lib.rs +++ b/tc-subxt/src/lib.rs @@ -48,6 +48,7 @@ pub trait TxSubmitter: Clone + Send + Sync + 'static { pub enum Tx { RegisterMember { network: NetworkId, peer_id: PeerId, stake_amount: u128 }, + UnregisterMember, Heartbeat, Commitment { shard_id: ShardId, commitment: Commitment, proof_of_knowledge: ProofOfKnowledge }, CreateTask { task: TaskDescriptorParams }, @@ -129,6 +130,10 @@ impl SubxtWorker { ); self.create_signed_payload(&payload).await }, + Tx::UnregisterMember => { + let payload = metadata::tx().members().unregister_member(); + self.create_signed_payload(&payload).await + }, Tx::Heartbeat => { let payload = metadata::tx().members().send_heartbeat(); self.create_signed_payload(&payload).await @@ -589,6 +594,13 @@ impl Runtime for SubxtClient { Ok(()) } + async fn submit_unregister_member(&self) -> Result<()> { + let (tx, rx) = oneshot::channel(); + self.tx.unbounded_send((Tx::UnregisterMember, tx))?; + rx.await?; + Ok(()) + } + async fn submit_heartbeat(&self) -> Result<()> { let (tx, rx) = oneshot::channel(); self.tx.unbounded_send((Tx::Heartbeat, tx))?;