Repository: tangle-network/gadget (language: rust)

Issue:
[TASK] Implement chain query functionality for active runner roles



Discussion / hints:

Fix patch:
diff --git a/common/src/client.rs b/common/src/client.rs
index f8e9b57e8..803d953ac 100644
--- a/common/src/client.rs
+++ b/common/src/client.rs
@@ -188,6 +188,16 @@ pub trait ClientWithApi: Client {
         at: [u8; 32],
         address: AccountId32,
     ) -> Result<Option<Vec<u8>>, crate::Error>;
+
+    /// Queries restaker's roles
+    ///
+    /// # Returns
+    /// List of roles enabled for restaker
+    async fn query_restaker_roles(
+        &self,
+        at: [u8; 32],
+        address: AccountId32,
+    ) -> Result<Vec<roles::RoleType>, crate::Error>;
 }
 
 impl<C: ClientWithApi> JobsClient<C> {
diff --git a/shell/src/shell.rs b/shell/src/shell.rs
index 3e62f9c10..0592df053 100644
--- a/shell/src/shell.rs
+++ b/shell/src/shell.rs
@@ -1,5 +1,6 @@
-use std::sync::Arc;
+use std::collections::{HashMap, HashSet};
 use std::time::Duration;
+use std::{hash::Hash, sync::Arc};
 
 use crate::{
     keystore::KeystoreContainer,
@@ -12,8 +13,11 @@ use gadget_common::{
     full_protocol::NodeInput,
     keystore::{ECDSAKeyStore, InMemoryBackend},
 };
+use gadget_core::gadget::substrate::Client;
+use parity_scale_codec::Encode;
 use sp_core::{ecdsa, ed25519, sr25519, ByteArray, Pair};
 use sp_keystore::Keystore;
+use tangle_subxt::tangle_runtime::api::runtime_types::tangle_primitives::roles::tss::ThresholdSignatureRoleType;
 
 use crate::config::ShellConfig;
 use crate::network::gossip::GossipHandle;
@@ -23,7 +27,7 @@ use dfns_cggmp21_protocol::constants::{
     DFNS_CGGMP21_KEYROTATE_PROTOCOL_NAME, DFNS_CGGMP21_SIGNING_PROTOCOL_NAME,
 };
 use gadget_common::keystore::KeystoreBackend;
-use tangle_subxt::subxt;
+use tangle_subxt::{subxt, tangle_runtime::api::runtime_types::tangle_primitives::roles::RoleType};
 
 /// The version of the shell
 pub const AGENT_VERSION: &str = "tangle/gadget-shell/1.0.0";
@@ -34,12 +38,7 @@ pub const CLIENT_VERSION: &str = env!("CARGO_PKG_VERSION");
 pub async fn run_forever(config: ShellConfig) -> color_eyre::Result<()> {
     let (role_key, acco_key) = load_keys_from_keystore(&config.keystore)?;
     let network_key = ed25519::Pair::from_seed(&config.node_key);
-    let subxt_client =
-        subxt::OnlineClient::<subxt::PolkadotConfig>::from_url(&config.subxt.endpoint).await?;
     let logger = DebugLogger::default();
-    let pair_signer = PairSigner::new(acco_key.clone());
-    let pallet_tx_submitter =
-        SubxtPalletSubmitter::with_client(subxt_client.clone(), pair_signer, logger.clone());
     let wrapped_keystore = ECDSAKeyStore::new(InMemoryBackend::new(), role_key.clone());
 
     let libp2p_key = libp2p::identity::Keypair::ed25519_from_bytes(network_key.to_raw_vec())
@@ -63,21 +62,8 @@ pub async fn run_forever(config: ShellConfig) -> color_eyre::Result<()> {
     logger.debug("Successfully initialized network, now waiting for bootnodes to connect ...");
     wait_for_connection_to_bootnodes(&config, &networks, &logger).await?;
 
-    let pallet_tx = Arc::new(pallet_tx_submitter);
-
-    let dfns_cggmp21_protocol = start_protocol(
-        vec![
-            TangleRuntime::new(subxt_client.clone()),
-            TangleRuntime::new(subxt_client.clone()),
-            TangleRuntime::new(subxt_client.clone()),
-            TangleRuntime::new(subxt_client.clone()),
-        ],
-        networks,
-        acco_key.public(),
-        logger,
-        pallet_tx,
-        wrapped_keystore,
-    );
+    let protocols =
+        start_required_protocols(&config.subxt, networks, acco_key, logger, wrapped_keystore);
 
     let ctrl_c = tokio::signal::ctrl_c();
 
@@ -85,9 +71,9 @@ pub async fn run_forever(config: ShellConfig) -> color_eyre::Result<()> {
         _ = ctrl_c => {
             tracing::info!("Received Ctrl-C, shutting down");
         }
-        dfns = dfns_cggmp21_protocol => {
-            if let Err(e) = dfns {
-                tracing::error!(?e, "DFNS-CGGMP21 protocol failed");
+        res = protocols => {
+            if let Err(e) = res {
+                tracing::error!(error = ?e, "Protocols watcher task unexpectedly shutdown");
             }
         }
 
@@ -99,34 +85,121 @@ pub async fn run_forever(config: ShellConfig) -> color_eyre::Result<()> {
     Ok(())
 }
 
-pub async fn start_protocol<C, KBE>(
-    clients: Vec<C>,
+pub fn start_protocol_by_role<KBE>(
+    role: RoleType,
+    runtime: TangleRuntime,
     networks: Vec<GossipHandle>,
     account_id: sr25519::Public,
     logger: DebugLogger,
     pallet_tx: Arc<SubxtPalletSubmitter<TangleConfig, PairSigner<TangleConfig>>>,
     keystore: ECDSAKeyStore<KBE>,
-) -> color_eyre::Result<()>
+) -> color_eyre::Result<tokio::task::AbortHandle>
 where
-    C: ClientWithApi + 'static,
     KBE: KeystoreBackend,
 {
-    let node_input = NodeInput {
-        clients,
-        account_id,
-        logger,
-        pallet_tx,
-        keystore,
-        node_index: 0,
-        additional_params: (),
-        prometheus_config: PrometheusConfig::Disabled,
-        networks,
+    use RoleType::*;
+    use ThresholdSignatureRoleType::*;
+    let handle = match role {
+        Tss(DfnsCGGMP21Stark) | Tss(DfnsCGGMP21Secp256r1) | Tss(DfnsCGGMP21Secp256k1) => {
+            let node_input = NodeInput {
+                clients: vec![
+                    TangleRuntime::new(runtime.client()),
+                    TangleRuntime::new(runtime.client()),
+                    TangleRuntime::new(runtime.client()),
+                    TangleRuntime::new(runtime.client()),
+                ],
+                account_id,
+                logger,
+                pallet_tx,
+                keystore,
+                node_index: 0,
+                additional_params: (),
+                prometheus_config: PrometheusConfig::Disabled,
+                networks,
+            };
+
+            tokio::spawn(dfns_cggmp21_protocol::setup_node(node_input))
+        }
+        _ => {
+            return Err(color_eyre::eyre::eyre!(
+                "Role {:?} is not supported by the shell",
+                role
+            ))
+        }
     };
 
-    // Wait for the protocol to finish
-    dfns_cggmp21_protocol::setup_node(node_input).await;
+    Ok(handle.abort_handle())
+}
+
+pub async fn start_required_protocols<KBE>(
+    subxt_config: &crate::config::SubxtConfig,
+    networks: Vec<GossipHandle>,
+    acco_key: sr25519::Pair,
+    logger: DebugLogger,
+    keystore: ECDSAKeyStore<KBE>,
+) -> color_eyre::Result<()>
+where
+    KBE: KeystoreBackend,
+{
+    let sub_account_id = subxt::utils::AccountId32(acco_key.public().0);
+    // Create a loop that listens to new finality notifications,
+    // queries the chain for the current restaking roles,
+    // and then starts the required protocols based on the roles.
+    let mut current_roles = Vec::new();
+    let mut running_protocols = HashMap::<HashedRoleTypeWrapper, tokio::task::AbortHandle>::new();
+
+    let subxt_client =
+        subxt::OnlineClient::<subxt::PolkadotConfig>::from_url(&subxt_config.endpoint).await?;
 
-    Err(color_eyre::eyre::eyre!("Protocol finished unexpectedly"))
+    let pair_signer = PairSigner::new(acco_key.clone());
+    let pallet_tx_submitter =
+        SubxtPalletSubmitter::with_client(subxt_client.clone(), pair_signer, logger.clone());
+    let pallet_tx = Arc::new(pallet_tx_submitter);
+    let runtime = TangleRuntime::new(subxt_client);
+    while let Some(notification) = runtime.get_next_finality_notification().await {
+        let roles = runtime
+            .query_restaker_roles(notification.hash, sub_account_id.clone())
+            .await?;
+        if roles == current_roles {
+            continue;
+        }
+        let diff = vec_diff(
+            current_roles.iter().cloned().map(HashedRoleTypeWrapper),
+            roles.iter().cloned().map(HashedRoleTypeWrapper),
+        );
+        if diff.is_empty() {
+            continue;
+        }
+        for d in diff {
+            match d {
+                Diff::Added(role) => {
+                    let handle = start_protocol_by_role(
+                        role.0.clone(),
+                        TangleRuntime::new(runtime.client()),
+                        networks.clone(),
+                        acco_key.public(),
+                        logger.clone(),
+                        pallet_tx.clone(),
+                        keystore.clone(),
+                    )?;
+                    running_protocols.insert(role, handle);
+                }
+                Diff::Removed(role) => {
+                    logger.debug(format!("Trying to stop protocol for role {:?}", role.0));
+                    let maybe_handle = running_protocols.remove(&role);
+                    if let Some(handle) = maybe_handle {
+                        handle.abort();
+                        logger.warn(format!(
+                            "Aborted protocol for role {:?}. Reason: Role Removed from profile",
+                            role.0
+                        ));
+                    }
+                }
+            }
+        }
+        current_roles = roles;
+    }
+    Ok(())
 }
 
 pub fn load_keys_from_keystore(
@@ -206,3 +279,36 @@ pub async fn wait_for_connection_to_bootnodes(
 
     Ok(())
 }
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+enum Diff<T> {
+    Added(T),
+    Removed(T),
+}
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+struct HashedRoleTypeWrapper(RoleType);
+
+impl Hash for HashedRoleTypeWrapper {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.0.encode().hash(state);
+    }
+}
+
+fn vec_diff<T: PartialEq + Eq + Hash + Clone, I: Iterator<Item = T>>(a: I, b: I) -> Vec<Diff<T>> {
+    let a_set = HashSet::<T, std::hash::RandomState>::from_iter(a);
+    let b_set = HashSet::from_iter(b);
+    // find the elements that are in a but not in b (removed)
+    let removed = a_set
+        .difference(&b_set)
+        .cloned()
+        .map(Diff::Removed)
+        .collect::<Vec<_>>();
+    // find the elements that are in b but not in a (added)
+    let added = b_set
+        .difference(&a_set)
+        .cloned()
+        .map(Diff::Added)
+        .collect::<Vec<_>>();
+    [removed, added].concat()
+}
diff --git a/shell/src/tangle.rs b/shell/src/tangle.rs
index 6380ee620..0649ea2f5 100644
--- a/shell/src/tangle.rs
+++ b/shell/src/tangle.rs
@@ -59,6 +59,10 @@ impl TangleRuntime {
         }
     }
 
+    pub fn client(&self) -> subxt::OnlineClient<PolkadotConfig> {
+        self.client.clone()
+    }
+
     /// Initialize the TangleRuntime instance by listening for finality notifications.
     /// This method must be called before using the instance.
     async fn initialize(&self) -> Result<()> {
@@ -226,6 +230,22 @@ impl ClientWithApi for TangleRuntime {
             .map_err(|e| gadget_common::Error::ClientError { err: e.to_string() })
             .await
     }
+
+    async fn query_restaker_roles(
+        &self,
+        at: [u8; 32],
+        address: AccountId32,
+    ) -> std::result::Result<Vec<RoleType>, gadget_common::Error> {
+        let storage_address = api::storage().roles().account_roles_mapping(address);
+        let block_ref = BlockRef::from_hash(sp_core::hash::H256::from_slice(&at));
+        self.client
+            .storage()
+            .at(block_ref)
+            .fetch_or_default(&storage_address)
+            .map_ok(|v| v.0)
+            .map_err(|e| gadget_common::Error::ClientError { err: e.to_string() })
+            .await
+    }
 }
 
 #[cfg(test)]


Test patch:
diff --git a/test-utils/src/mock.rs b/test-utils/src/mock.rs
index 9a2a3776c..a01b068d8 100644
--- a/test-utils/src/mock.rs
+++ b/test-utils/src/mock.rs
@@ -941,6 +941,14 @@ pub mod mock_wrapper_client {
             })
             .await
         }
+
+        async fn query_restaker_roles(
+            &self,
+            _at: [u8; 32],
+            _address: AccountId32,
+        ) -> Result<Vec<roles::RoleType>, gadget_common::Error> {
+            Ok(Default::default())
+        }
     }
 
     impl<R: ProvideRuntimeApi<B>, B: Block> ProvideRuntimeApi<B> for MockClient<R, B> {