Merge pull request #66 from sideprotocol/websocket_reconnect

Optimize web socket connection and DKG procession

Author: Lazy Luis

Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 edition = "2021"
 name = "shuttler"
-version = "2.0.0-rc3"
+version = "2.0.0-rc4"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

Cargo.toml

@@ -14,7 +14,7 @@ use crate::{
     helper::{
         encoding::to_base64,
         now,
-        store::{DefaultStore, MemStore, Store},
+        store::{DefaultStore, Store},
     }, protocols::refresh::RefreshInput,
 };
 
@@ -210,8 +210,12 @@ type SignatureShareStore =
 type SignerNonceStore = DefaultStore<String, BTreeMap<Index, round1::SigningNonces>>;
 
 pub type Round1Store =
-    MemStore<String, BTreeMap<Identifier, Vec<frost_adaptor_signature::keys::dkg::round1::Package>>>;
-pub type Round2Store = MemStore<String, BTreeMap<Identifier, Vec<Vec<u8>>>>;
+    DefaultStore<String, BTreeMap<Identifier, Vec<frost_adaptor_signature::keys::dkg::round1::Package>>>;
+pub type Round2Store = DefaultStore<String, BTreeMap<Identifier, Vec<Vec<u8>>>>;
+
+
+pub type Round1SecetStore = DefaultStore<String, Vec<frost_adaptor_signature::keys::dkg::round1::SecretPackage>>;
+pub type Round2SecetStore = DefaultStore<String, Vec<frost_adaptor_signature::keys::dkg::round2::SecretPackage>>;
 
 pub struct Context {
     pub swarm: Swarm<ShuttlerBehaviour>,
@@ -232,6 +236,8 @@ pub struct Context {
     // dkg stores
     pub db_round1: Round1Store,
     pub db_round2: Round2Store,
+    pub sec_round1: Round1SecetStore,
+    pub sec_round2: Round2SecetStore,
 }
 
 impl Context {
@@ -265,13 +271,24 @@ impl Context {
             commitment_store: CommitmentStore::new(conf.get_database_with_name("commitments")),
             signature_store: SignatureShareStore::new(conf.get_database_with_name("signature_shares")),
             general_store: DefaultStore::new(conf.get_database_with_name("general")),
+
+            db_round1: Round1Store::new(conf.get_database_with_name("round1")),
+            db_round2: Round2Store::new(conf.get_database_with_name("round2")),
+            sec_round1: Round1SecetStore::new(conf.get_database_with_name("sec_round1")),
+            sec_round2: Round2SecetStore::new(conf.get_database_with_name("sec_round2")),
+
             conf,
 
-            db_round1: Round1Store::new(),
-            db_round2: Round2Store::new(),
         }
     }
 
+    pub fn clean_dkg_cache(&self, task_id: &String) {
+        self.db_round1.remove(task_id);
+        self.db_round2.remove(task_id);
+        self.sec_round1.remove(task_id);
+        self.sec_round2.remove(task_id);
+    }
+
     pub fn clean_task_cache(&self, task_id: &String) {
         self.task_store.remove(task_id);
         self.nonce_store.remove(task_id);

src/apps/core.rs

@@ -3,7 +3,6 @@
 use cosmrs::Any;
 use frost_adaptor_signature::VerifyingKey;
 use side_proto::side::tss::{MsgCompleteDkg, MsgCompleteRefreshing, MsgSubmitSignatures, SigningType};
-use tracing::debug;
 
 use crate::config::{VaultKeypair, APP_NAME_LENDING};
 use crate::helper::encoding::{from_base64, hash, pubkey_to_identifier};
@@ -64,7 +63,6 @@ impl DKGAdaptor for KeygenHander {
         match event {
             SideEvent::BlockEvent(events) => {
                 if events.contains_key("initiate_dkg.id") {
-                    println!("Events: {:?}", events);
 
                     let live_peers = mem_store::alive_participants();
 
@@ -75,16 +73,27 @@ impl DKGAdaptor for KeygenHander {
                         .zip(events.get("initiate_dkg.batch_size")?) {
 
                             let mut participants = vec![];
+                            let mut down_peers = vec![];
+                            let mut names = vec![];
                             for p in ps.split(",") {
                                 if let Ok(keybytes) = from_base64(p) {
                                     let identifier = pubkey_to_identifier(&keybytes);
                                     // not have enough participants
+                                    let moniker = mem_store::get_participant_moniker(&identifier);
                                     if !live_peers.contains(&identifier) {
-                                        break;
-                                    }
+                                        down_peers.push(moniker.clone());
+                                    } 
+                                    names.push(moniker);
+                                    
                                     participants.push(identifier);
                                 }
                             };
+
+                            tracing::debug!("Task {} has {} offline participants {:?} {:?}, threshold {}", id, down_peers.len(), down_peers, names, t);
+                            if down_peers.len() > 0 {
+                                continue;
+                            }
+
                             if let Ok(threshold) = t.parse() {
                                 if threshold as usize * 3 >= participants.len() * 2  {
                                     if let Ok(batch_size) = b.parse() {
@@ -116,7 +125,7 @@ impl DKGAdaptor for KeygenHander {
             pub_keys.push(hexkey);
         });
 
-        debug!("Oracle pubkey >>>: {:?}", pub_keys);
+        // debug!("Oracle pubkey >>>: {:?}", pub_keys);
 
         // save dkg id and keys for refresh
         ctx.general_store.save(&format!("{}", task.id).as_str(), &pub_keys.join(","));

src/apps/lending.rs

@@ -528,7 +528,7 @@ pub async fn send_withdraw_tx(
     info!("Submitting withdrawal tx: {:?}", msg);
 
     let any_msg = Any::from_msg(&msg).unwrap();
-    send_cosmos_transaction(relayer.config(), any_msg).await
+    send_cosmos_transaction(&relayer.identifier, relayer.config(), any_msg).await
 }
 
 pub async fn send_deposit_tx(
@@ -549,7 +549,7 @@ pub async fn send_deposit_tx(
     info!("Submitting deposit tx: {:?}", msg);
 
     let any_msg = Any::from_msg(&msg).unwrap();
-    send_cosmos_transaction(&relayer.config(), any_msg).await
+    send_cosmos_transaction(&relayer.identifier, &relayer.config(), any_msg).await
 }
 
 pub(crate) fn get_last_scanned_height(relayer: &Relayer) -> u64 {
@@ -629,7 +629,7 @@ pub async fn submit_fee_rate_to_side(relayer: &Relayer, fee_rate: i64) {
 
     info!("Submitting fee rate: {:?}", msg_submit_fee_rate);
     let any_msg = Any::from_msg(&msg_submit_fee_rate).unwrap();
-    match send_cosmos_transaction(relayer.config(), any_msg).await {
+    match send_cosmos_transaction(&relayer.identifier, relayer.config(), any_msg).await {
         Ok(resp) => {
             let tx_response = resp.into_inner().tx_response.unwrap();
             if tx_response.code != 0 {

src/apps/relayer/bridge.rs

@@ -404,7 +404,7 @@ pub async fn send_deposit_tx(
     info!("submit deposit tx to side: {:?}", msg);
 
     let any_msg = Any::from_msg(&msg).unwrap();
-    send_cosmos_transaction(&relayer.config(), any_msg).await
+    send_cosmos_transaction(&relayer.identifier, &relayer.config(), any_msg).await
 }
 
 pub async fn handle_cet(relayer: &Relayer, loan_id: String, cet_type: String) {

src/apps/relayer/lending.rs

@@ -3,8 +3,9 @@
 use std::time::Duration;
 
 use bitcoincore_rpc::{Auth, Client};
+use ed25519_compact::SecretKey;
 use futures::join;
-use crate::{config::Config, helper::{client_fee_provider::FeeProviderClient, client_ordinals::OrdinalsClient}};
+use crate::{config::Config, helper::{client_fee_provider::FeeProviderClient, client_ordinals::OrdinalsClient, encoding::pubkey_to_identifier}};
 
 pub mod bridge;
 pub mod lending;
@@ -18,6 +19,7 @@ pub struct Relayer {
     pub fee_provider_client: FeeProviderClient,
     pub db_relayer: sled::Db,
     pub ticker: tokio::time::Interval,
+    pub identifier: frost_adaptor_signature::Identifier,
 }
 
 impl Relayer {
@@ -40,6 +42,20 @@ impl Relayer {
         let db_relayer = sled::open(conf.get_database_with_name("relayer")).expect("Counld not create database!");
         let ticker = tokio::time::interval(Duration::from_secs(conf.loop_interval as u64));
 
+
+        // load private key from priv_validator_key_path
+        let priv_validator_key = conf.load_validator_key();
+
+        let mut b = priv_validator_key
+            .priv_key
+            .ed25519_signing_key()
+            .unwrap()
+            .as_bytes()
+            .to_vec();
+        b.extend(priv_validator_key.pub_key.to_bytes());
+        let node_key = SecretKey::new(b.as_slice().try_into().unwrap());
+        let identifier = pubkey_to_identifier(node_key.public_key().as_slice());
+
         Self {
             // priv_validator_key: validator_key,
             bitcoin_client,
@@ -48,6 +64,7 @@ impl Relayer {
             config: conf,
             db_relayer,
             ticker,
+            identifier
         }
     }
 

src/apps/relayer/mod.rs

@@ -165,10 +165,11 @@ impl<'a> Shuttler<'a> {
         // Tx Sender for Tx Quene
         let (tx_sender, tx_receiver) = std::sync::mpsc::channel::<Any>();
         let conf2 = conf.clone();
+        let identifier2 = identifier.clone();
         spawn(async move {
             while let Ok(message) = tx_receiver.recv() {
-                println!("Received: {:?}", message);
-                match send_cosmos_transaction(&conf2, message).await {
+                // println!("Received: {:?}", message);
+                match send_cosmos_transaction(&identifier2, &conf2, message).await {
                     Ok(resp) => {
                         if let Some(inner) = resp.into_inner().tx_response {
                             debug!("Submited {}, {}, {}", inner.txhash, inner.code, inner.raw_log)
@@ -179,7 +180,6 @@ impl<'a> Shuttler<'a> {
             }
         });
 
-
         // Common Setting: Context and Heart Beat
         let mut context = Context::new(swarm, tx_sender, identifier, node_key, conf.clone()); 
         let mut ticker = tokio::time::interval_at(get_next_full_hour(), Duration::from_secs(5 * 60));
@@ -206,18 +206,31 @@ impl<'a> Shuttler<'a> {
 
         loop {
             select! {
-                Some(msg) = client.receive_message() => {
-                    // tracing::info!("msg {:?}", msg);
-                    if self.handle_block_event(&mut context, msg) {
-                        tracing::error!("websocket connection closed, reconnecting...");
-                        tokio::time::sleep(Duration::from_secs(5)).await;
-                        if client.reconnect().await.is_err() {
-                            tracing::error!("Failed to reconnect to websocket");
-                            break;
-                        } else {
-                            tracing::info!("Reconnected to websocket");
+                Some(recv) = client.receive_message() => {
+                    match recv {
+                        Ok(msg) => {
+                            if self.handle_block_event(&mut context, msg) {
+                                tracing::error!("websocket connection closed, will reconnect in 5s.");
+                                tokio::time::sleep(Duration::from_secs(5)).await;
+                                if client.reconnect().await.is_err() {
+                                    tracing::error!("Failed to reconnect to websocket");
+                                    break;
+                                } else {
+                                    tracing::info!("Reconnected to websocket");
+                                }
+                            };
+                        },
+                        Err(e) => {
+                            tracing::error!("websocket error: {:?}, will reconnect in 5s", e);
+                            tokio::time::sleep(Duration::from_secs(5)).await;
+                            if client.reconnect().await.is_err() {
+                                tracing::error!("Failed to reconnect to websocket");
+                                break;
+                            } else {
+                                tracing::info!("Reconnected to websocket");
+                            }
                         }
-                    };
+                    }
                 }
                 // recv = sidechain_event_stream.next() => {
                 //     match recv {
@@ -248,7 +261,7 @@ impl<'a> Shuttler<'a> {
                         // info!(" @@(Received) Discovered new peer: {peer_id} with info: {connection_id} {:?}", info);
                         info.listen_addrs.iter().for_each(|addr| {
                             if !addr.to_string().starts_with("/ip4/127.0.0.1") {
-                                tracing::debug!("Discovered: {addr}/p2p/{peer_id}");
+                                // tracing::debug!("Discovered: {addr}/p2p/{peer_id}");
                                 context.swarm
                                     .behaviour_mut()
                                     .kad
@@ -276,7 +289,10 @@ impl<'a> Shuttler<'a> {
                         } else {
                             let _ = context.swarm.disconnect_peer_id(peer_id);
                         }
-                        info!("Connected peers {:?}", context.swarm.connected_peers().collect::<Vec<_>>());
+                        
+                        info!("Connected to {}", peer_id)
+
+                        // info!("Connected peers {:?}", context.swarm.connected_peers().map(|i|).collect::<Vec<_>>());
                     },
                     SwarmEvent::ConnectionClosed { peer_id, cause, .. } => {
                         info!("Disconnected {peer_id}: {:?}", cause);
@@ -336,7 +352,7 @@ impl<'a> Shuttler<'a> {
             },
             tokio_tungstenite::tungstenite::Message::Close(_close) => {
                 tracing::error!("connection closed");
-                return false
+                return true
             },
             _ => return false
         };
@@ -350,7 +366,7 @@ impl<'a> Shuttler<'a> {
             tendermint_rpc::event::v0_38::DeEventData::NewBlock { block, result_finalize_block , ..} => {
                 if let Some(b) = block { 
                     let height = b.header.height.value();
-                    debug!("Received New block: #{:?}, {:?}", height, mem_store::alive_participants_monikers());
+                    debug!("Block: #{:?}, offline: {:?}", height, mem_store::offline_participants_monikers());
                     sending_heart_beat(ctx, height);
                 }
                 if let Some(finalize_block) = result_finalize_block {
@@ -420,8 +436,7 @@ impl<'a> Shuttler<'a> {
                         tasks.push(task);
                     }
                 };
-            });
-            
+            });    
         }
 
         if let Some(lending) = self.apps.iter().find(|a| a.name() == APP_NAME_LENDING) {

src/apps/shuttler.rs

@@ -60,7 +60,7 @@ impl Candidate {
         params.allowed_dkg_participants.iter().for_each(|v| {
             if let Ok(pk ) = from_base64(&v.consensus_pubkey) {
                 let id = pubkey_to_identifier(&pk);
-                debug!("added {:?} in white list", id);
+                debug!("Access granted to {} {}", v.consensus_pubkey, v.moniker);
                 mem_store::add_moniker(&id, v.moniker.clone());
                 self.peers.push(identifier_to_peer_id(&id ));
                 self.identifiers.push( id );