Avoid cloning bundles and use Arc<Inner> in cell.

Author: akonring

timeboost-sequencer/Cargo.toml

@@ -33,5 +33,5 @@ tracing = { workspace = true }
 ark-std = { workspace = true }
 bs58 = { workspace = true }
 futures = { workspace = true }
-test-utils = { path = "../test-utils", features = ["ports"]}
+test-utils = { path = "../test-utils", features = ["ports"] }
 timeboost-utils = { path = "../timeboost-utils" }

timeboost-sequencer/src/decrypt.rs

@@ -1,28 +1,28 @@
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
+use std::result::Result as StdResult;
+use std::sync::Arc;
+
 use ark_std::{UniformRand, rand::thread_rng};
 use bon::Builder;
 use bytes::{BufMut, Bytes, BytesMut};
 use cliquenet::overlay::{Data, DataError, NetworkDown, Overlay};
 use cliquenet::{
     AddressableCommittee, MAX_MESSAGE_SIZE, Network, NetworkError, NetworkMetrics, Role,
 };
-use rayon::iter::ParallelIterator;
-use rayon::prelude::*;
-
 use multisig::{Committee, CommitteeId, PublicKey};
 use parking_lot::RwLock;
+use rayon::iter::ParallelIterator;
+use rayon::prelude::*;
 use sailfish::types::{Evidence, Round, RoundNumber};
 use serde::{Deserialize, Serialize};
-use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
-use std::result::Result as StdResult;
-use std::sync::Arc;
 use timeboost_config::DECRYPTER_PORT_OFFSET;
 use timeboost_crypto::prelude::{
     DkgDecKey, LabeledDkgDecKey, Plaintext, ThresholdCiphertext, ThresholdDecShare,
     ThresholdEncError, ThresholdEncScheme, ThresholdScheme, Vess, VssSecret,
 };
 use timeboost_types::{
-    AccumulatorMode, DkgAccumulator, DkgBundle, DkgSubset, InclusionList, KeyStore, KeyStoreVec,
-    ThresholdKey, ThresholdKeyCell,
+    AccumulatorMode, DkgAccumulator, DkgBundle, DkgSubset, DkgSubsetRef, InclusionList, KeyStore,
+    KeyStoreVec, ThresholdKey, ThresholdKeyCell,
 };
 use tokio::spawn;
 use tokio::sync::mpsc::{Receiver, Sender, channel};
@@ -728,21 +728,31 @@ impl Worker {
             .ok_or(DecrypterError::NoCommittee(self.current))?;
 
         let prev = (guard.len() == 2).then(|| guard.last().clone());
-
-        subsets.insert(src, res.subset.to_owned());
         let committee = current.committee();
         let threshold: usize = committee.one_honest_threshold().into();
 
-        let mut counts = HashMap::new();
-        for subset in subsets.values() {
-            *counts.entry(subset).or_insert(0) += 1;
+        subsets.insert(src, res.subset);
+
+        let mut subset_groups: HashMap<DkgSubsetRef, Vec<&PublicKey>> = HashMap::new();
+        let mut threshold_subset = None;
+
+        for (pk, subset) in subsets.iter() {
+            let pks = subset_groups.entry(subset.as_ref()).or_default();
+            pks.push(pk);
+
+            if pks.len() >= threshold {
+                threshold_subset = Some((subset.as_ref(), pks[0]));
+                break;
+            }
         }
 
-        if let Some((&subset, _)) = counts.iter().find(|(_, count)| **count >= threshold) {
-            let acc = DkgAccumulator::from_subset(current.clone(), subset.to_owned());
+        if let Some((subset_ref, pk)) = threshold_subset {
+            let subset = subsets.to_owned().remove(pk).expect("subset exists in map");
+            let acc = DkgAccumulator::from_subset(current.clone(), subset);
             let mode = acc.mode().clone();
             self.tracker.insert(committee.id(), acc);
-            let dec_key = subset
+
+            let dec_key = subset_ref
                 .extract_key(current, &self.dkg_sk, prev.as_ref())
                 .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
 
@@ -830,7 +840,7 @@ impl Worker {
         if matches!(mode, AccumulatorMode::Dkg) {
             if let Some(subset) = acc.try_finalize() {
                 let dec_key = subset
-                    .extract_key(&key_store.clone(), &self.dkg_sk, None)
+                    .extract_key(&key_store, &self.dkg_sk, None)
                     .map_err(|e| DecrypterError::Dkg(e.to_string()))?;
                 self.dec_key.set(dec_key);
                 self.state = WorkerState::Running;

timeboost-types/src/decryption.rs

@@ -131,10 +131,15 @@ impl ThresholdKey {
 
 /// `DecryptionKeyCell` is a thread-safe container for an optional `DecryptionKey`
 /// that allows asynchronous notification when the key is set.
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Default)]
 pub struct ThresholdKeyCell {
-    key: Arc<RwLock<Option<ThresholdKey>>>,
-    notify: Arc<Notify>,
+    inner: Arc<ThresholdKeyCellInner>,
+}
+
+#[derive(Debug, Default)]
+struct ThresholdKeyCellInner {
+    key: RwLock<Option<ThresholdKey>>,
+    notify: Notify,
 }
 
 impl ThresholdKeyCell {
@@ -143,25 +148,25 @@ impl ThresholdKeyCell {
     }
 
     pub fn set(&self, key: ThresholdKey) {
-        *self.key.write() = Some(key);
-        self.notify.notify_waiters();
+        *self.inner.key.write() = Some(key);
+        self.inner.notify.notify_waiters();
     }
 
     pub fn get(&self) -> Option<ThresholdKey> {
-        (*self.key.read()).clone()
+        (*self.inner.key.read()).clone()
     }
 
     pub fn enc_key(&self) -> Option<ThresholdEncKey> {
         self.get().map(|sk| sk.pubkey)
     }
 
     pub fn get_ref(&self) -> impl Deref<Target = Option<ThresholdKey>> {
-        self.key.read()
+        self.inner.key.read()
     }
 
     pub async fn read(&self) -> ThresholdKey {
         loop {
-            let fut = self.notify.notified();
+            let fut = self.inner.notify.notified();
             if let Some(k) = self.get() {
                 return k;
             }
@@ -170,6 +175,14 @@ impl ThresholdKeyCell {
     }
 }
 
+impl Clone for ThresholdKeyCell {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
 /// A small, non-empty collection of KeyStores.
 #[derive(Debug, Default, Clone)]
 #[allow(clippy::len_without_is_empty)]
@@ -280,10 +293,10 @@ impl<const N: usize> From<KeyStore> for KeyStoreVec<N> {
 
 /// A `KeyStore` with committee information and public keys used in the DKG or key resharing
 #[derive(Debug, Clone)]
-pub struct KeyStore(Arc<Inner>);
+pub struct KeyStore(Arc<KeyStoreInner>);
 
 #[derive(Debug)]
-struct Inner {
+struct KeyStoreInner {
     committee: Committee,
     keys: BTreeMap<KeyId, DkgEncKey>,
 }
@@ -294,7 +307,7 @@ impl KeyStore {
         I: IntoIterator<Item = (T, DkgEncKey)>,
         T: Into<KeyId>,
     {
-        let this = Self(Arc::new(Inner {
+        let this = Self(Arc::new(KeyStoreInner {
             committee: c,
             keys: keys.into_iter().map(|(i, k)| (i.into(), k)).collect(),
         }));
@@ -446,20 +459,19 @@ impl DkgAccumulator {
     }
 
     /// Try to finalize the accumulator into a subset if enough bundles are collected.
-    pub fn try_finalize(&mut self) -> Option<DkgSubset> {
+    /// Returns a reference to the internal data to avoid cloning the bundles.
+    pub fn try_finalize(&self) -> Option<DkgSubsetRef<'_>> {
         if self.complete {
             let combkey = match &self.mode {
                 AccumulatorMode::Dkg => None,
                 AccumulatorMode::Resharing(combkey) => Some(combkey.clone()),
             };
 
-            let subset = DkgSubset {
+            Some(DkgSubsetRef {
                 committee_id: self.committee().id(),
-                bundles: self.bundles.clone(),
+                bundles: &self.bundles,
                 combkey,
-            };
-
-            Some(subset)
+            })
         } else {
             None
         }
@@ -488,6 +500,14 @@ pub struct DkgSubset {
     combkey: Option<ThresholdCombKey>,
 }
 
+/// A reference-based version of DkgSubset to avoid cloning bundles.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct DkgSubsetRef<'a> {
+    committee_id: CommitteeId,
+    bundles: &'a [DkgBundle],
+    combkey: Option<ThresholdCombKey>,
+}
+
 impl DkgSubset {
     /// Create a new subset with DKG bundles.
     pub fn new_dkg(committee_id: CommitteeId, bundles: Vec<DkgBundle>) -> Self {
@@ -536,6 +556,27 @@ impl DkgSubset {
         self.combkey.is_some()
     }
 
+    /// Convert this DkgSubset to a DkgSubsetRef.
+    pub fn as_ref(&self) -> DkgSubsetRef<'_> {
+        DkgSubsetRef {
+            committee_id: self.committee_id,
+            bundles: &self.bundles,
+            combkey: self.combkey.as_ref().cloned(),
+        }
+    }
+
+    /// Extract the new threshold decryption key from the subset.
+    pub fn extract_key(
+        &self,
+        curr: &KeyStore,
+        dkg_sk: &LabeledDkgDecKey,
+        prev: Option<&KeyStore>,
+    ) -> anyhow::Result<ThresholdKey> {
+        self.as_ref().extract_key(curr, dkg_sk, prev)
+    }
+}
+
+impl<'a> DkgSubsetRef<'a> {
     /// Extract the new threshold decryption key from the subset.
     pub fn extract_key(
         &self,
@@ -547,7 +588,7 @@ impl DkgSubset {
 
         match &self.combkey {
             None => {
-                let mut dealings_iter = ResultIter::new(self.bundles().iter().map(|b| {
+                let mut dealings_iter = ResultIter::new(self.bundles.iter().map(|b| {
                     vess.decrypt_share(curr.committee(), dkg_sk, b.vess_ct(), DKG_AAD)
                         .map(|s| (s, b.comm().clone()))
                 }));
@@ -568,7 +609,7 @@ impl DkgSubset {
                 };
 
                 let dealings: Vec<_> = self
-                    .bundles()
+                    .bundles
                     .iter()
                     .map(|b| {
                         let node_idx = b.origin().0.into();

timeboost-types/src/lib.rs

@@ -22,8 +22,8 @@ pub use bundle::{
 pub use bytes::Bytes;
 pub use candidate_list::{CandidateList, CandidateListBytes};
 pub use decryption::{
-    AccumulatorMode, DkgAccumulator, DkgSubset, KeyStore, KeyStoreVec, ResultIter, ThresholdKey,
-    ThresholdKeyCell,
+    AccumulatorMode, DkgAccumulator, DkgSubset, DkgSubsetRef, KeyStore, KeyStoreVec, ResultIter,
+    ThresholdKey, ThresholdKeyCell,
 };
 pub use delayed_inbox::DelayedInboxIndex;
 pub use inclusion_list::InclusionList;