wait enough DKG to finish in yapper

Author: alxiong

timeboost-crypto/src/lib.rs

@@ -120,7 +120,7 @@ pub struct CombKey<C: CurveGroup> {
 }
 
 #[serde_as]
-#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, From)]
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, From, Hash)]
 pub struct PublicKey<C: CurveGroup> {
     #[serde_as(as = "crate::SerdeAs")]
     key: C,

timeboost-crypto/src/prelude.rs

@@ -87,8 +87,6 @@ pub struct ThresholdEncKeyCell {
     notify: Arc<Notify>,
 }
 
-impl ThresholdEncKeyCell {}
-
 impl ThresholdEncKeyCell {
     pub fn new() -> Self {
         Self::default()

yapper/src/tx.rs

@@ -1,6 +1,6 @@
 use futures::future::join_all;
 use reqwest::{Client, Url};
-use std::time::Duration;
+use std::{collections::HashMap, time::Duration};
 use timeboost::types::BundleVariant;
 use timeboost_crypto::prelude::{ThresholdEncKey, ThresholdEncKeyCell};
 use timeboost_utils::load_generation::{make_bundle, tps_to_millis};
@@ -85,26 +85,75 @@ async fn fetch_encryption_key(client: &Client, enckey_url: &Url) -> Option<Thres
     }
 }
 
+/// helper struct to keep track of sufficient quorum of DKG keys
+struct ThresholdEncKeyCellAccumulator {
+    client: Client,
+    // DKG results on individual node
+    results: HashMap<Url, Option<ThresholdEncKey>>,
+    // (t+1)-agreed upon encryption key
+    output: ThresholdEncKeyCell,
+    // threshold for the accumulator to be considered as matured / finalized
+    threshold: usize,
+}
+
+impl ThresholdEncKeyCellAccumulator {
+    /// give a list of TimeboostApi's endpoint to query `/enckey` status
+    pub fn new(client: Client, urls: impl Iterator<Item = Url>) -> Self {
+        let results: HashMap<Url, Option<ThresholdEncKey>> = urls.map(|url| (url, None)).collect();
+        let output = ThresholdEncKeyCell::new();
+        let threshold = results.len().div_ceil(3);
+        Self {
+            client,
+            results,
+            output,
+            threshold,
+        }
+    }
+
+    /// try to get the threshold encryption key, only available after a threshold of nodes
+    /// finish their DKG processes.
+    pub async fn enc_key(&mut self) -> &ThresholdEncKeyCell {
+        // if result is already available, directly return
+        if self.output.get_ref().is_some() {
+            &self.output
+        } else {
+            // first update DKG status for yet-finished nodes
+            for (url, res) in self.results.iter_mut() {
+                if res.is_none() {
+                    *res = fetch_encryption_key(&self.client, url).await;
+                }
+            }
+
+            // count for each unique enc key from DKG results of different nodes
+            let mut counts: HashMap<ThresholdEncKey, usize> = HashMap::new();
+            for v in self.results.values().flatten() {
+                *counts.entry(v.to_owned()).or_insert(0) += 1;
+            }
+
+            // (t+1)-agreed enc_key is the output
+            for (v, c) in counts.iter() {
+                if *c >= self.threshold {
+                    self.output.set(v.to_owned());
+                }
+            }
+            &self.output
+        }
+    }
+}
+
 pub async fn yap(addresses: &[Address], tps: u32) -> Result<()> {
     let c = Client::builder().timeout(Duration::from_secs(1)).build()?;
     let urls = setup_urls(addresses)?;
 
     let mut interval = interval(Duration::from_millis(tps_to_millis(tps)));
     interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 
-    let enc_key = ThresholdEncKeyCell::new();
-    let enckey_url = &urls.first().expect("urls shouldn't be empty").2;
-    loop {
-        if enc_key.get_ref().is_none() {
-            tracing::debug!("DKG ongoing ...");
-            if let Some(k) = fetch_encryption_key(&c, enckey_url).await {
-                tracing::info!("DKG done");
-                enc_key.set(k)
-            }
-        }
+    let mut acc =
+        ThresholdEncKeyCellAccumulator::new(c.clone(), urls.iter().map(|url| url.2.clone()));
 
+    loop {
         // create a bundle for next `interval.tick()`, then send this bundle to each node
-        let Ok(b) = make_bundle(&enc_key) else {
+        let Ok(b) = make_bundle(acc.enc_key().await) else {
             warn!("failed to generate bundle");
             continue;
         };