Merge pull request #442 from EspressoSystems/tw/some-tweaks

Some minor tweaks and updates

Author: twittner

Cargo.toml

@@ -71,7 +71,7 @@ nohash-hasher = "0.2"
 parking_lot = "0.12.3"
 portpicker = "0.1.1"
 prometheus = "0.14"
-prost = "0.13.5"
+prost = "0.14.1"
 quickcheck = "1.0.3"
 rand = "0.9"
 rayon = "1.10"
@@ -95,8 +95,9 @@ tokio-stream = "0.1.17"
 tokio-tungstenite = { version = "0.27.0", features = ["rustls-tls-webpki-roots", "url"] }
 tokio-util = "0.7.15"
 toml = "0.8.19"
-tonic = "0.13.1"
-tonic-build = { version = "0.13.1", features = ["prost"] }
+tonic = "0.14.1"
+tonic-prost = "0.14.1"
+tonic-prost-build = "0.14.1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter", "json"] }
 turmoil = "0.6.4"

tests/src/tests/timeboost/block_order.rs

@@ -74,8 +74,7 @@ async fn block_order() {
         rxs.push(rx)
     }
 
-    // wait until DKG is done
-    enc_key.wait().await;
+    enc_key.read().await;
     tracing::info!("DKG done");
 
     tasks.spawn(gen_bundles(enc_key, bcast.clone()));

tests/src/tests/timeboost/transaction_order.rs

@@ -73,8 +73,7 @@ async fn transaction_order() {
         rxs.push(rx)
     }
 
-    // wait until DKG is done
-    enc_key.wait().await;
+    enc_key.read().await;
     tracing::info!("DKG done");
 
     tasks.spawn(gen_bundles(enc_key, bcast.clone()));

timeboost-crypto/src/prelude.rs

@@ -78,9 +78,6 @@ pub type ThresholdDecKeyShare = <DecryptionScheme as ThresholdEncScheme>::KeySha
 
 /// `ThresholdEncKeyCell` is a thread-safe container for an optional `ThresholdEncKey`
 /// that allows asynchronous notification when the key is set.
-///
-/// Internally, it uses an `RwLock<Option<ThresholdEncKey>>` to guard the key,
-/// and a `Notify` to wake up tasks waiting for the key to become available.
 #[derive(Debug, Clone, Default)]
 pub struct ThresholdEncKeyCell {
     key: Arc<RwLock<Option<ThresholdEncKey>>>,
@@ -105,18 +102,13 @@ impl ThresholdEncKeyCell {
         self.key.read()
     }
 
-    /// Asynchronously waits for the key to become available, then returns it.
-    ///
-    /// If the key is already present, it is returned immediately.
-    /// Otherwise, the current task is suspended until `set()` is called.
-    ///
-    /// The returned key is a clone of the stored key.
-    pub async fn wait(&self) -> ThresholdEncKey {
+    pub async fn read(&self) -> ThresholdEncKey {
         loop {
+            let fut = self.notify.notified();
             if let Some(k) = self.get() {
                 return k;
             }
-            self.notify.notified().await;
+            fut.await;
         }
     }
 }

timeboost-proto/Cargo.toml

@@ -9,6 +9,7 @@ rust-version.workspace = true
 prost = { workspace = true }
 thiserror = { workspace = true }
 tonic = { workspace = true }
+tonic-prost = { workspace = true }
 
 [build-dependencies]
-tonic-build = { workspace = true }
+tonic-prost-build = { workspace = true }

timeboost-proto/build.rs

@@ -1,10 +1,11 @@
 use std::error::Error;
 
 fn main() -> Result<(), Box<dyn Error>> {
-    tonic_build::configure()
+    tonic_prost_build::configure()
         .build_server(true)
         .out_dir("src")
-        .bytes([".block.Block.payload", ".inclusion.Transaction.encoded_txn"])
+        .bytes(".block.Block.payload")
+        .bytes(".inclusion.Transaction.encoded_txn")
         .compile_protos(
             &[
                 "protos/inclusion_list.proto",

timeboost-sequencer/src/decrypt.rs

@@ -10,7 +10,7 @@ use multisig::{Committee, CommitteeId, PublicKey};
 use parking_lot::RwLock;
 use sailfish::types::{CommitteeVec, Evidence, Round, RoundNumber};
 use serde::{Deserialize, Serialize};
-use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use std::result::Result as StdResult;
 use std::sync::Arc;
 use timeboost_crypto::prelude::{LabeledDkgDecKey, ThresholdEncKeyCell, Vess, Vss};
@@ -232,18 +232,13 @@ impl Decrypter {
         Ok(())
     }
 
-    /// Send the received DKG bundles to worker
-    pub async fn enqueue_dkg(
-        &self,
-        pending_dkgs: &mut VecDeque<DkgBundle>,
-    ) -> StdResult<(), DecrypterDown> {
-        while let Some(b) = pending_dkgs.pop_front() {
-            self.worker_tx
-                .send(Command::Dkg(b))
-                .await
-                .map_err(|_| DecrypterDown(()))?;
-            debug!(node = %self.label, "enqueued one dkg bundle");
-        }
+    /// Send the received DKG bundle to worker
+    pub async fn enqueue_dkg(&self, dkg: DkgBundle) -> StdResult<(), DecrypterDown> {
+        self.worker_tx
+            .send(Command::Dkg(dkg))
+            .await
+            .map_err(|_| DecrypterDown(()))?;
+        debug!(node = %self.label, "enqueued one dkg bundle");
         Ok(())
     }
 
@@ -1621,11 +1616,12 @@ mod tests {
 
         // enqueuing them all to decrypters
         for decrypter in decrypters.iter_mut() {
-            let mut pending_dkgs = dkg_bundles.clone();
-            decrypter
-                .enqueue_dkg(&mut pending_dkgs)
-                .await
-                .expect("DKG bundles should be enqueued successfully");
+            for dkg in dkg_bundles.clone() {
+                decrypter
+                    .enqueue_dkg(dkg)
+                    .await
+                    .expect("DKG bundles should be enqueued successfully");
+            }
         }
     }
 

timeboost-sequencer/src/lib.rs

@@ -261,23 +261,22 @@ impl Task {
     // processing its actions continues unhindered.
     async fn go(mut self) -> Result<()> {
         let mut pending = None;
-        let mut pending_dkgs = VecDeque::new();
+        let mut dkg_bundles = VecDeque::new();
         let mut candidates = Candidates::new();
 
         if !self.sailfish.is_init() {
             let actions = self.sailfish.init();
             candidates = self.execute(actions).await?;
+        }
 
-            // DKG dealing generation
-            // TODO: move/copy to main loop when resharing
-            if let Some(bundle) = self.decrypter.gen_dkg_bundle() {
-                self.bundles.add_bundles(once(BundleVariant::Dkg(bundle)));
-            }
+        // TODO: move/copy to main loop when resharing
+        if let Some(bundle) = self.decrypter.gen_dkg_bundle() {
+            self.bundles.add_bundles(once(BundleVariant::Dkg(bundle)));
         }
 
         loop {
             if pending.is_none() {
-                while let Some(ilist) = self.next_inclusion(&mut candidates, &mut pending_dkgs) {
+                while let Some(ilist) = self.next_inclusion(&mut candidates, &mut dkg_bundles) {
                     if !self.decrypter.has_capacity() {
                         pending = Some(ilist);
                         break;
@@ -288,12 +287,15 @@ impl Task {
                 }
             }
 
-            // always sync DKG bundles
-            self.next_dkg(&mut candidates, &mut pending_dkgs);
-            if !pending_dkgs.is_empty() {
-                tracing::debug!(num_bundles = %pending_dkgs.len(), "enqueuing dkg bundles");
-                if let Err(err) = self.decrypter.enqueue_dkg(&mut pending_dkgs).await {
-                    error!(node = %self.label, %err, "dkg enqueue error");
+            if pending.is_none() {
+                while let Some(dkg) = dkg_bundles.pop_front() {
+                    if !self.decrypter.has_capacity() {
+                        dkg_bundles.push_front(dkg);
+                        break;
+                    }
+                    if let Err(err) = self.decrypter.enqueue_dkg(dkg).await {
+                        error!(node = %self.label, %err, "dkg enqueue error");
+                    }
                 }
             }
 
@@ -435,7 +437,7 @@ impl Task {
             // preprocess the candidate list to pull out the DKG bundles first
             for cl in lists.iter() {
                 if let Some(dkg) = cl.dkg_bundle() {
-                    pending_dkgs.push_back(dkg);
+                    pending_dkgs.push_back(dkg.clone());
                 }
             }
             // then process it to construct the next inclusion list
@@ -456,19 +458,6 @@ impl Task {
         }
         None
     }
-
-    /// Handle candidate lists and "pull" out the DKG bundles, it won't touch or drop or consume
-    /// regular/priority bundles, but only consume/take the DKG bundles inside `candidates` then
-    /// append them to `pending_dkgs`.
-    fn next_dkg(&mut self, candidates: &mut Candidates, pending_dkgs: &mut VecDeque<DkgBundle>) {
-        for (_, _, list) in candidates.iter_mut() {
-            for cl in list.iter_mut() {
-                if let Some(dkg) = cl.take_dkg_bundle() {
-                    pending_dkgs.push_back(dkg);
-                }
-            }
-        }
-    }
 }
 
 #[derive(Debug, thiserror::Error)]

timeboost-sequencer/src/queue.rs

@@ -191,17 +191,14 @@ impl DataSource for BundleQueue {
         inner.set_time(time);
 
         if r.is_genesis() || inner.mode.is_passive() {
-            // allow DKG in passive mode (first 8 rounds from genesis)
-            let candidate_list = CandidateList::builder(Timestamp::now(), inner.index)
-                .with_dkg(inner.dkg.clone())
+            return CandidateList::builder(Timestamp::now(), inner.index)
+                .with_dkg(inner.dkg.take())
                 .finish()
                 .try_into()
                 .unwrap_or_else(|err| {
                     error!(%err, "candidate list serialization error");
                     CandidateListBytes::default()
                 });
-            inner.dkg = None;
-            return candidate_list;
         }
 
         let mut size_budget = inner.max_len;
@@ -229,19 +226,16 @@ impl DataSource for BundleQueue {
             regular.push(b.clone())
         }
 
-        let candidate_list = CandidateList::builder(inner.time, inner.index)
+        CandidateList::builder(inner.time, inner.index)
             .with_priority_bundles(priority)
             .with_regular_bundles(regular)
-            .with_dkg(inner.dkg.clone())
+            .with_dkg(inner.dkg.take())
             .finish()
             .try_into()
             .unwrap_or_else(|err| {
                 error!(%err, "candidate list serialization error");
                 CandidateListBytes::default()
-            });
-
-        inner.dkg = None;
-        candidate_list
+            })
     }
 }
 

timeboost-types/src/candidate_list.rs

@@ -1,5 +1,5 @@
+use std::ops::Deref;
 use std::sync::Arc;
-use std::{ops::Deref, sync::RwLock};
 
 use bytes::{BufMut, Bytes, BytesMut};
 use committable::{Commitment, Committable, RawCommitmentBuilder};
@@ -20,7 +20,7 @@ struct Inner {
     index: DelayedInboxIndex,
     priority: Vec<SignedPriorityBundle>,
     regular: Vec<Bundle>,
-    dkg: RwLock<Option<DkgBundle>>,
+    dkg: Option<DkgBundle>,
 }
 
 #[derive(Debug)]
@@ -29,7 +29,7 @@ pub struct Builder {
     index: DelayedInboxIndex,
     priority: Vec<SignedPriorityBundle>,
     regular: Vec<Bundle>,
-    dkg: RwLock<Option<DkgBundle>>,
+    dkg: Option<DkgBundle>,
 }
 
 impl Builder {
@@ -44,7 +44,7 @@ impl Builder {
     }
 
     pub fn with_dkg(mut self, d: Option<DkgBundle>) -> Self {
-        self.dkg = RwLock::new(d);
+        self.dkg = d;
         self
     }
 
@@ -69,7 +69,7 @@ impl CandidateList {
             index: i.into(),
             regular: Vec::new(),
             priority: Vec::new(),
-            dkg: RwLock::new(None),
+            dkg: None,
         }
     }
 
@@ -112,26 +112,8 @@ impl CandidateList {
         self.0.index
     }
 
-    pub fn dkg_bundle(&self) -> Option<DkgBundle> {
-        match self.0.dkg.read() {
-            Ok(dkg) => dkg.clone(),
-            Err(poisoned) => {
-                tracing::error!("RwLock on DkgBundle poisoned during read");
-                let dkg = poisoned.into_inner();
-                dkg.clone()
-            }
-        }
-    }
-
-    pub fn take_dkg_bundle(&self) -> Option<DkgBundle> {
-        match self.0.dkg.write() {
-            Ok(mut dkg) => dkg.take(),
-            Err(poisoned) => {
-                tracing::error!("RwLock on DkgBundle poisoned during write");
-                let mut dkg = poisoned.into_inner();
-                dkg.take()
-            }
-        }
+    pub fn dkg_bundle(&self) -> Option<&DkgBundle> {
+        self.0.dkg.as_ref()
     }
 }
 
@@ -140,7 +122,7 @@ impl Committable for CandidateList {
         let mut builder = RawCommitmentBuilder::new("CandidateList")
             .u64_field("time", self.0.time.into())
             .u64_field("index", self.0.index.into())
-            .optional("dkg", &self.dkg_bundle())
+            .optional("dkg", &self.0.dkg)
             .u64_field("priority", self.0.priority.len() as u64);
         builder = self
             .0