Make TSS startup checks more forgiving (#934)

* Add exponential backoff when connecting to Substrate node

* Split out code for getting threshold account ID from mnemonic check

* Make TSS balance check non-fatal

* Move code for checking node connection and TSS balance to a helper

* Move account check out of `setup_mnemonic`

* Get rid of import warnings

* Update lockfile

* TaploFmt

* Use `threshold_account_id()` in tests

* Use correct account in test

* Remove `dbg!` statement

crates/threshold-signature-server/Cargo.toml

@@ -23,6 +23,7 @@ reqwest-eventsource="0.6"
 serde_derive       ="1.0.147"
 synedrion          ={ git="https://github.com/entropyxyz/synedrion", rev="25373111cbb01e1a25d8a5c5bb8f4652c725b3f1" }
 strum              ="0.26.2"
+backoff            ={ version="0.4.0", features=["tokio"] }
 
 # Async
 futures="0.3"

crates/threshold-signature-server/src/helpers/launch.rs

@@ -241,21 +241,14 @@ pub struct StartupArgs {
     pub mnemonic_file: Option<PathBuf>,
 }
 
-pub async fn has_mnemonic(kv: &KvManager) -> (bool, String) {
+pub async fn has_mnemonic(kv: &KvManager) -> bool {
     let exists = kv.kv().exists(FORBIDDEN_KEY_MNEMONIC).await.expect("issue querying DB");
-    let mut account_id = "".to_string();
+
     if exists {
         tracing::debug!("Existing mnemonic found in keystore.");
-        let mnemonic = kv.kv().get(FORBIDDEN_KEYS[0]).await.expect("Issue getting mnemonic");
-        let pair = <sr25519::Pair as Pair>::from_phrase(
-            &String::from_utf8(mnemonic).expect("Issue converting mnemonic to string"),
-            None,
-        )
-        .expect("Issue converting mnemonic to pair");
-        account_id = AccountId32::new(pair.0.public().into()).to_ss58check();
     }
 
-    (exists, account_id)
+    exists
 }
 
 pub fn development_mnemonic(validator_name: &Option<ValidatorName>) -> bip39::Mnemonic {
@@ -275,8 +268,8 @@ pub fn development_mnemonic(validator_name: &Option<ValidatorName>) -> bip39::Mn
         .expect("Unable to parse given mnemonic.")
 }
 
-pub async fn setup_mnemonic(kv: &KvManager, mnemonic: bip39::Mnemonic) -> String {
-    if has_mnemonic(kv).await.0 {
+pub async fn setup_mnemonic(kv: &KvManager, mnemonic: bip39::Mnemonic) {
+    if has_mnemonic(kv).await {
         tracing::warn!("Deleting account related keys from KVDB.");
 
         kv.kv()
@@ -341,7 +334,16 @@ pub async fn setup_mnemonic(kv: &KvManager, mnemonic: bip39::Mnemonic) -> String
     fs::write(".entropy/account_id", format!("{id}")).expect("Failed to write account_id file");
 
     tracing::debug!("Starting process with account ID: `{id}`");
-    id.to_ss58check()
+}
+
+pub async fn threshold_account_id(kv: &KvManager) -> String {
+    let mnemonic = kv.kv().get(FORBIDDEN_KEY_MNEMONIC).await.expect("Issue getting mnemonic");
+    let pair = <sr25519::Pair as Pair>::from_phrase(
+        &String::from_utf8(mnemonic).expect("Issue converting mnemonic to string"),
+        None,
+    )
+    .expect("Issue converting mnemonic to pair");
+    AccountId32::new(pair.0.public().into()).to_ss58check()
 }
 
 pub async fn setup_latest_block_number(kv: &KvManager) -> Result<(), KvError> {
@@ -392,3 +394,63 @@ pub async fn setup_only(kv: &KvManager) {
 
     println!("{}", output);
 }
+
+pub async fn check_node_prerequisites(url: &str, account_id: &str) {
+    use crate::chain_api::{get_api, get_rpc};
+
+    let connect_to_substrate_node = || async {
+        tracing::info!("Attempting to establish connection to Substrate node at `{}`", url);
+
+        let api = get_api(url).await.map_err(|_| {
+            Err::<(), String>("Unable to connect to Substrate chain API".to_string())
+        })?;
+
+        let rpc = get_rpc(url)
+            .await
+            .map_err(|_| Err("Unable to connect to Substrate chain RPC".to_string()))?;
+
+        Ok((api, rpc))
+    };
+
+    // Note: By default this will wait 15 minutes before it stops retry attempts.
+    let backoff = backoff::ExponentialBackoff::default();
+    match backoff::future::retry(backoff, connect_to_substrate_node).await {
+        Ok((api, rpc)) => {
+            tracing::info!("Sucessfully connected to Substrate node!");
+
+            tracing::info!("Checking balance of threshold server AccountId `{}`", &account_id);
+            let balance_query = crate::validator::api::check_balance_for_fees(
+                &api,
+                &rpc,
+                account_id.to_string(),
+                entropy_shared::MIN_BALANCE,
+            )
+            .await
+            .map_err(|_| Err::<bool, String>("Failed to get balance of account.".to_string()));
+
+            match balance_query {
+                Ok(has_minimum_balance) => {
+                    if has_minimum_balance {
+                        tracing::info!(
+                            "The account `{}` has enough funds for submitting extrinsics.",
+                            &account_id
+                        )
+                    } else {
+                        tracing::warn!(
+                            "The account `{}` does not meet the minimum balance of `{}`",
+                            &account_id,
+                            entropy_shared::MIN_BALANCE,
+                        )
+                    }
+                },
+                Err(_) => {
+                    tracing::warn!("Unable to query the account balance of `{}`", &account_id)
+                },
+            }
+        },
+        Err(_err) => {
+            tracing::error!("Unable to establish connection with Substrate node at `{}`", url);
+            panic!("Unable to establish connection with Substrate node.");
+        },
+    }
+}

crates/threshold-signature-server/src/main.rs

@@ -17,15 +17,12 @@ use std::{net::SocketAddr, str::FromStr};
 
 use clap::Parser;
 
-use entropy_shared::MIN_BALANCE;
 use entropy_tss::{
     app,
-    chain_api::{get_api, get_rpc},
     launch::{
         development_mnemonic, load_kv_store, setup_latest_block_number, setup_mnemonic, setup_only,
         Configuration, StartupArgs, ValidatorName,
     },
-    validator::api::check_balance_for_fees,
     AppState,
 };
 
@@ -91,17 +88,16 @@ async fn main() {
             })
         });
 
-    let account_id = if let Some(mnemonic) = user_mnemonic {
+    if let Some(mnemonic) = user_mnemonic {
         setup_mnemonic(&kv_store, mnemonic).await
     } else if cfg!(test) || validator_name.is_some() {
         setup_mnemonic(&kv_store, development_mnemonic(&validator_name)).await
     } else {
-        let (has_mnemonic, account_id) = entropy_tss::launch::has_mnemonic(&kv_store).await;
+        let has_mnemonic = entropy_tss::launch::has_mnemonic(&kv_store).await;
         assert!(
             has_mnemonic,
             "No mnemonic provided. Please provide one or use a development account."
         );
-        account_id
     };
 
     setup_latest_block_number(&kv_store).await.expect("Issue setting up Latest Block Number");
@@ -112,14 +108,13 @@ async fn main() {
     if args.setup_only {
         setup_only(&kv_store).await;
     } else {
-        let api = get_api(&app_state.configuration.endpoint).await.expect("Error getting api");
-        let rpc = get_rpc(&app_state.configuration.endpoint).await.expect("Error getting rpc");
-        let has_fee_balance = check_balance_for_fees(&api, &rpc, account_id.clone(), MIN_BALANCE)
-            .await
-            .expect("Error in check balance");
-        if !has_fee_balance {
-            panic!("threshold account needs balance: {:?}", account_id);
-        }
+        let account_id = entropy_tss::launch::threshold_account_id(&kv_store).await;
+        entropy_tss::launch::check_node_prerequisites(
+            &app_state.configuration.endpoint,
+            &account_id,
+        )
+        .await;
+
         let listener = tokio::net::TcpListener::bind(&addr)
             .await
             .expect("Unable to bind to given server address.");

crates/threshold-signature-server/src/user/tests.rs

@@ -111,8 +111,9 @@ use crate::{
     get_signer,
     helpers::{
         launch::{
-            development_mnemonic, load_kv_store, setup_mnemonic, Configuration, ValidatorName,
-            DEFAULT_BOB_MNEMONIC, DEFAULT_CHARLIE_MNEMONIC, DEFAULT_ENDPOINT, DEFAULT_MNEMONIC,
+            development_mnemonic, load_kv_store, setup_mnemonic, threshold_account_id,
+            Configuration, ValidatorName, DEFAULT_BOB_MNEMONIC, DEFAULT_CHARLIE_MNEMONIC,
+            DEFAULT_ENDPOINT, DEFAULT_MNEMONIC,
         },
         signing::Hasher,
         substrate::{query_chain, submit_transaction},
@@ -143,9 +144,16 @@ async fn test_get_signer_does_not_throw_err() {
     initialize_test_logger().await;
     clean_tests();
 
+    let pair = <sr25519::Pair as Pair>::from_phrase(crate::helpers::launch::DEFAULT_MNEMONIC, None)
+        .expect("Issue converting mnemonic to pair");
+    let expected_account_id = AccountId32::new(pair.0.public().into()).to_ss58check();
+
     let kv_store = load_kv_store(&None, None).await;
-    let account = setup_mnemonic(&kv_store, development_mnemonic(&None)).await;
-    assert_eq!(account, "5DACCJgQV6sHoYUKfTGEimddFxe16NJXgkzHZ3RC9QCBShMH");
+    setup_mnemonic(&kv_store, development_mnemonic(&None)).await;
+    development_mnemonic(&None).to_string();
+    let account = threshold_account_id(&kv_store).await;
+
+    assert_eq!(account, expected_account_id);
     get_signer(&kv_store).await.unwrap();
     clean_tests();
 }