feat: make sign queue use less shared state

chain-signatures/contract/src/primitives.rs

@@ -3,7 +3,7 @@ use k256::Scalar;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::serde::{Deserialize, Serialize};
 use near_sdk::{AccountId, BorshStorageKey, CryptoHash, NearToken, PublicKey};
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{btree_map, BTreeMap, HashMap, HashSet};
 
 pub mod hpke {
     pub type PublicKey = [u8; 32];
@@ -125,8 +125,8 @@ impl Default for Participants {
 impl From<Candidates> for Participants {
     fn from(candidates: Candidates) -> Self {
         let mut participants = Participants::new();
-        for (account_id, candidate_info) in candidates.iter() {
-            participants.insert(account_id.clone(), candidate_info.clone().into());
+        for (account_id, candidate_info) in candidates.into_iter() {
+            participants.insert(account_id, candidate_info.into());
         }
         participants
     }
@@ -162,10 +162,14 @@ impl Participants {
         self.participants.get(account_id)
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = (&AccountId, &ParticipantInfo)> {
+    pub fn iter(&self) -> btree_map::Iter<'_, AccountId, ParticipantInfo> {
         self.participants.iter()
     }
 
+    pub fn iter_mut(&mut self) -> btree_map::IterMut<'_, AccountId, ParticipantInfo> {
+        self.participants.iter_mut()
+    }
+
     pub fn keys(&self) -> impl Iterator<Item = &AccountId> {
         self.participants.keys()
     }
@@ -179,6 +183,33 @@ impl Participants {
     }
 }
 
+impl<'a> IntoIterator for &'a Participants {
+    type Item = (&'a AccountId, &'a ParticipantInfo);
+    type IntoIter = btree_map::Iter<'a, AccountId, ParticipantInfo>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.participants.iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a mut Participants {
+    type Item = (&'a AccountId, &'a mut ParticipantInfo);
+    type IntoIter = btree_map::IterMut<'a, AccountId, ParticipantInfo>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.participants.iter_mut()
+    }
+}
+
+impl IntoIterator for Participants {
+    type Item = (AccountId, ParticipantInfo);
+    type IntoIter = btree_map::IntoIter<AccountId, ParticipantInfo>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.participants.into_iter()
+    }
+}
+
 #[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Debug, Clone)]
 pub struct Candidates {
     pub candidates: BTreeMap<AccountId, CandidateInfo>,
@@ -213,9 +244,40 @@ impl Candidates {
         self.candidates.get(account_id)
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = (&AccountId, &CandidateInfo)> {
+    pub fn iter(&self) -> btree_map::Iter<'_, AccountId, CandidateInfo> {
         self.candidates.iter()
     }
+
+    pub fn iter_mut(&mut self) -> btree_map::IterMut<'_, AccountId, CandidateInfo> {
+        self.candidates.iter_mut()
+    }
+}
+
+impl<'a> IntoIterator for &'a Candidates {
+    type Item = (&'a AccountId, &'a CandidateInfo);
+    type IntoIter = btree_map::Iter<'a, AccountId, CandidateInfo>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.candidates.iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a mut Candidates {
+    type Item = (&'a AccountId, &'a mut CandidateInfo);
+    type IntoIter = btree_map::IterMut<'a, AccountId, CandidateInfo>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.candidates.iter_mut()
+    }
+}
+
+impl IntoIterator for Candidates {
+    type Item = (AccountId, CandidateInfo);
+    type IntoIter = btree_map::IntoIter<AccountId, CandidateInfo>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.candidates.into_iter()
+    }
 }
 
 #[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Debug)]

chain-signatures/node/src/cli.rs

@@ -191,7 +191,7 @@ pub fn run(cmd: Cli) -> anyhow::Result<()> {
             mesh_options,
             message_options,
         } => {
-            let sign_queue = Arc::new(RwLock::new(SignQueue::new()));
+            let (sign_tx, sign_rx) = SignQueue::channel();
             let rt = tokio::runtime::Builder::new_multi_thread()
                 .enable_all()
                 .build()?;
@@ -221,7 +221,7 @@ pub fn run(cmd: Cli) -> anyhow::Result<()> {
                 &indexer_options,
                 &mpc_contract_id,
                 &account_id,
-                &sign_queue,
+                sign_tx,
                 app_data_storage,
                 rpc_client.clone(),
             )?;
@@ -257,7 +257,7 @@ pub fn run(cmd: Cli) -> anyhow::Result<()> {
                 rpc_client.clone(),
                 signer,
                 receiver,
-                sign_queue,
+                sign_rx,
                 key_storage,
                 triple_storage,
                 presignature_storage,

chain-signatures/node/src/http_client.rs

@@ -1,4 +1,4 @@
-use crate::protocol::contract::primitives::{ParticipantInfo, Participants};
+use crate::protocol::contract::primitives::Participants;
 use crate::protocol::message::SignedMessage;
 use crate::protocol::MpcMessage;
 use cait_sith::protocol::Participant;
@@ -95,7 +95,7 @@ pub async fn send_encrypted<U: IntoUrl>(
 // TODO: add in retry logic either in struct or at call site.
 // TODO: add check for participant list to see if the messages to be sent are still valid.
 pub struct MessageQueue {
-    deque: VecDeque<(ParticipantInfo, MpcMessage, Instant)>,
+    deque: VecDeque<(Participant, MpcMessage, Instant)>,
     seen_counts: HashSet<String>,
     message_options: Options,
 }
@@ -117,16 +117,21 @@ impl MessageQueue {
         self.deque.is_empty()
     }
 
-    pub fn push(&mut self, info: ParticipantInfo, msg: MpcMessage) {
-        self.deque.push_back((info, msg, Instant::now()));
+    pub fn push(&mut self, node: Participant, msg: MpcMessage) {
+        self.deque.push_back((node, msg, Instant::now()));
+    }
+
+    pub fn extend(&mut self, other: impl IntoIterator<Item = (Participant, MpcMessage)>) {
+        self.deque
+            .extend(other.into_iter().map(|(i, msg)| (i, msg, Instant::now())));
     }
 
     pub async fn send_encrypted(
         &mut self,
         from: Participant,
         sign_sk: &near_crypto::SecretKey,
         client: &Client,
-        participants: &Participants,
+        active: &Participants,
         cfg: &ProtocolConfig,
     ) -> Vec<SendError> {
         let mut failed = VecDeque::new();
@@ -136,21 +141,22 @@ impl MessageQueue {
         let outer = Instant::now();
         let uncompacted = self.deque.len();
         let mut encrypted = HashMap::new();
-        while let Some((info, msg, instant)) = self.deque.pop_front() {
+        while let Some((id, msg, instant)) = self.deque.pop_front() {
             if instant.elapsed() > timeout(&msg, cfg) {
                 errors.push(SendError::Timeout(format!(
-                    "{} message has timed out: {info:?}",
+                    "{} message has timed out for node={id:?}",
                     msg.typename(),
                 )));
                 continue;
             }
 
-            if !participants.contains_key(&Participant::from(info.id)) {
-                let counter = participant_counter.entry(info.id).or_insert(0);
+            let Some(info) = active.get(&id) else {
+                let counter = participant_counter.entry(id).or_insert(0);
                 *counter += 1;
-                failed.push_back((info, msg, instant));
+                failed.push_back((id, msg, instant));
                 continue;
-            }
+            };
+
             let encrypted_msg = match SignedMessage::encrypt(&msg, from, sign_sk, &info.cipher_pk) {
                 Ok(encrypted) => encrypted,
                 Err(err) => {
@@ -159,15 +165,15 @@ impl MessageQueue {
                 }
             };
             let encrypted = encrypted.entry(info.id).or_insert_with(Vec::new);
-            encrypted.push((encrypted_msg, (info, msg, instant)));
+            encrypted.push((encrypted_msg, (id, msg, instant)));
         }
 
         let mut compacted = 0;
         for (id, encrypted) in encrypted {
             for partition in partition_ciphered_256kb(encrypted) {
                 let (encrypted_partition, msgs): (Vec<_>, Vec<_>) = partition.into_iter().unzip();
                 // guaranteed to unwrap due to our previous loop check:
-                let info = participants.get(&Participant::from(id)).unwrap();
+                let info = active.get(&Participant::from(id)).unwrap();
                 let account_id = &info.account_id;
                 let number_of_messages = encrypted_partition.len() as f64;
 
@@ -231,7 +237,7 @@ impl MessageQueue {
 /// Encrypted message with a reference to the old message. Only the ciphered portion of this
 /// type will be sent over the wire, while the original message is kept just in case things
 /// go wrong somewhere and the message needs to be requeued to be sent later.
-type EncryptedMessage = (Ciphered, (ParticipantInfo, MpcMessage, Instant));
+type EncryptedMessage = (Ciphered, (Participant, MpcMessage, Instant));
 
 fn partition_ciphered_256kb(encrypted: Vec<EncryptedMessage>) -> Vec<Vec<EncryptedMessage>> {
     let mut result = Vec::new();

chain-signatures/node/src/indexer.rs

@@ -1,4 +1,4 @@
-use crate::protocol::{SignQueue, SignRequest};
+use crate::protocol::SignRequest;
 use crate::storage::app_data_storage::AppDataStorage;
 use crypto_shared::{derive_epsilon, ScalarExt};
 use k256::Scalar;
@@ -13,7 +13,7 @@ use std::ops::Mul;
 use std::sync::Arc;
 use std::thread::JoinHandle;
 use std::time::{Duration, Instant};
-use tokio::sync::RwLock;
+use tokio::sync::{mpsc, RwLock};
 
 /// Configures indexer.
 #[derive(Debug, Clone, clap::Parser)]
@@ -169,7 +169,7 @@ impl Indexer {
 struct Context {
     mpc_contract_id: AccountId,
     node_account_id: AccountId,
-    queue: Arc<RwLock<SignQueue>>,
+    sign_tx: mpsc::Sender<SignRequest>,
     indexer: Indexer,
 }
 
@@ -267,14 +267,20 @@ async fn handle_block(
 
     // Add the requests after going through the whole block to avoid partial processing if indexer fails somewhere.
     // This way we can revisit the same block if we failed while not having added the requests partially.
-    let mut queue = ctx.queue.write().await;
     for request in pending_requests {
-        queue.add(request);
+        tracing::info!(
+            request_id = ?near_primitives::hash::CryptoHash(request.request_id),
+            payload = hex::encode(request.request.payload.to_bytes()),
+            entropy = hex::encode(request.entropy),
+            "new sign request"
+        );
+        if let Err(err) = ctx.sign_tx.send(request).await {
+            tracing::error!(?err, "failed to send the sign request into sign queue");
+        }
         crate::metrics::NUM_SIGN_REQUESTS
             .with_label_values(&[ctx.node_account_id.as_str()])
             .inc();
     }
-    drop(queue);
 
     let log_indexing_interval = 1000;
     if block.block_height() % log_indexing_interval == 0 {
@@ -292,7 +298,7 @@ pub fn run(
     options: &Options,
     mpc_contract_id: &AccountId,
     node_account_id: &AccountId,
-    queue: &Arc<RwLock<SignQueue>>,
+    sign_tx: mpsc::Sender<SignRequest>,
     app_data_storage: AppDataStorage,
     rpc_client: near_fetch::Client,
 ) -> anyhow::Result<(JoinHandle<anyhow::Result<()>>, Indexer)> {
@@ -308,7 +314,7 @@ pub fn run(
     let context = Context {
         mpc_contract_id: mpc_contract_id.clone(),
         node_account_id: node_account_id.clone(),
-        queue: queue.clone(),
+        sign_tx,
         indexer: indexer.clone(),
     };
 

chain-signatures/node/src/mesh/connection.rs

@@ -71,18 +71,19 @@ impl Pool {
             }
         }
 
-        let connections = self.connections.read().await.clone(); // Clone connections for iteration
-        let mut join_set = JoinSet::new();
+        let connections = {
+            let conn = self.connections.read().await;
+            conn.clone()
+        };
 
         // Spawn tasks for each participant
-        for (participant, info) in connections.iter() {
-            let participant = *participant;
-            let info = info.clone();
-            let self_clone = Arc::clone(&self); // Clone Arc for use inside tasks
+        let mut join_set = JoinSet::new();
+        for (participant, info) in connections.into_iter() {
+            let pool = Arc::clone(&self);
 
             join_set.spawn(async move {
-                match self_clone.fetch_participant_state(&info).await {
-                    Ok(state) => match self_clone.send_empty_msg(&participant, &info).await {
+                match pool.fetch_participant_state(&info).await {
+                    Ok(state) => match pool.send_empty_msg(&participant, &info).await {
                         Ok(()) => Ok((participant, state, info)),
                         Err(e) => {
                             tracing::warn!(
@@ -139,19 +140,19 @@ impl Pool {
             }
         }
 
-        let connections = self.potential_connections.read().await;
-
-        let mut join_set = JoinSet::new();
+        let connections = {
+            let conn = self.potential_connections.read().await;
+            conn.clone()
+        };
 
         // Spawn tasks for each participant
-        for (participant, info) in connections.iter() {
-            let participant = *participant;
-            let info = info.clone();
-            let self_clone = Arc::clone(&self); // Clone Arc for use inside tasks
+        let mut join_set = JoinSet::new();
+        for (participant, info) in connections.into_iter() {
+            let pool = Arc::clone(&self); // Clone Arc for use inside tasks
 
             join_set.spawn(async move {
-                match self_clone.fetch_participant_state(&info).await {
-                    Ok(state) => match self_clone.send_empty_msg(&participant, &info).await {
+                match pool.fetch_participant_state(&info).await {
+                    Ok(state) => match pool.send_empty_msg(&participant, &info).await {
                         Ok(()) => Ok((participant, state, info)),
                         Err(e) => {
                             tracing::warn!(