address comments

Author: lukeiannucci

Cargo.toml

@@ -72,7 +72,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"
@@ -96,8 +96,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"

robusta/src/lib.rs

@@ -303,7 +303,7 @@ mod tests {
         let clt = Client::new(cfg.clone());
         let height = clt.height().await.unwrap();
         let mut watcher = Watcher::new(cfg, height, None);
-        let header = watcher.next().await;
+        let header = watcher.next().await.unwrap_right();
         assert_eq!(u64::from(height), header.height());
     }
 }

robusta/src/multiwatcher.rs

@@ -4,15 +4,18 @@ use std::{
         Arc,
         atomic::{AtomicU64, Ordering},
     },
+    time::Duration,
 };
 
 use crate::{Config, Height, Watcher};
+use either::Either;
 use espresso_types::{Header, NamespaceId};
 use futures::{StreamExt, stream::SelectAll};
 use tokio::{
     spawn,
     sync::{Barrier, mpsc},
     task::JoinHandle,
+    time::sleep,
 };
 use tokio_stream::wrappers::ReceiverStream;
 use tracing::{debug, warn};
@@ -66,16 +69,47 @@ impl Multiwatcher {
             let lower_bound = lower_bound.clone();
             watchers.push(spawn(async move {
                 let i = Id(i);
-                let mut w = Watcher::new(c, height, nsid);
                 loop {
-                    let h = w.next().await;
-                    if h.height() <= lower_bound.load(Ordering::Relaxed) {
-                        continue;
+                    let height = lower_bound.load(Ordering::Relaxed);
+                    let mut w = Watcher::new(c.clone(), height, nsid);
+                    let mut expected = height + 1;
+                    loop {
+                        match w.next().await {
+                            Either::Right(hdr) => {
+                                if hdr.height() > expected {
+                                    warn!(
+                                        url      = %c.wss_base_url,
+                                        height   = %hdr.height(),
+                                        expected = %expected,
+                                        "unexpected block height"
+                                    );
+                                    break;
+                                }
+                                expected += 1;
+                                if hdr.height() <= lower_bound.load(Ordering::Relaxed) {
+                                    continue;
+                                }
+                                if tx.send((i, hdr)).await.is_err() {
+                                    return;
+                                }
+                            }
+                            Either::Left(height) => {
+                                if *height > expected {
+                                    warn!(
+                                        url      = %c.wss_base_url,
+                                        height   = %height,
+                                        expected = %expected,
+                                        "unexpected block height"
+                                    );
+                                    break;
+                                }
+                                expected += 1;
+                            }
+                        }
+                        barrier.wait().await;
                     }
-                    if tx.send((i, h)).await.is_err() {
-                        break;
-                    }
-                    barrier.wait().await;
+                    drop(w);
+                    sleep(Duration::from_secs(3)).await // wait a little before re-connecting
                 }
             }));
         }

robusta/src/types.rs

@@ -1,4 +1,8 @@
-use std::{borrow::Cow, fmt};
+use std::{
+    borrow::Cow,
+    fmt,
+    ops::{Add, AddAssign, Deref},
+};
 
 use bon::Builder;
 use data_encoding::BASE64URL_NOPAD;
@@ -42,6 +46,42 @@ macro_rules! Primitive {
                 self.0.fmt(f)
             }
         }
+
+        impl Deref for $name {
+            type Target = $t;
+
+            fn deref(&self) -> &$t {
+                &self.0
+            }
+        }
+
+        impl Add for $name {
+            type Output = Self;
+
+            fn add(self, rhs: Self) -> Self::Output {
+                Self(self.0 + rhs.0)
+            }
+        }
+
+        impl Add<$t> for $name {
+            type Output = Self;
+
+            fn add(self, rhs: $t) -> Self {
+                Self(self.0 + rhs)
+            }
+        }
+
+        impl AddAssign for $name {
+            fn add_assign(&mut self, rhs: Self) {
+                *self = *self + rhs
+            }
+        }
+
+        impl AddAssign<$t> for $name {
+            fn add_assign(&mut self, rhs: $t) {
+                *self = *self + rhs
+            }
+        }
     };
 }
 

robusta/src/watcher.rs

@@ -1,6 +1,7 @@
 use std::str::from_utf8;
 
 use bytes::Bytes;
+use either::Either;
 use espresso_types::{Header, NamespaceId};
 use futures::{SinkExt, StreamExt};
 use reqwest::header::LOCATION;
@@ -37,7 +38,7 @@ impl Watcher {
         }
     }
 
-    pub async fn next(&mut self) -> Header {
+    pub async fn next(&mut self) -> Either<Height, Header> {
         'main: loop {
             let ws = if let Some(w) = &mut self.websocket {
                 w
@@ -74,16 +75,16 @@ impl Watcher {
                     Some(Ok(Message::Text(text))) => {
                         match serde_json::from_str::<Header>(text.as_str()) {
                             Ok(hdr) => {
+                                self.height = hdr.height().into();
                                 if let Some(id) = &self.namespace {
                                     if hdr.ns_table().find_ns_id(id).is_some() {
-                                        self.height = hdr.height().into();
-                                        return hdr;
+                                        return Either::Right(hdr);
                                     } else {
                                         debug!(height = %hdr.height(), "namespace id not found");
+                                        return Either::Left(self.height);
                                     }
                                 } else {
-                                    self.height = hdr.height().into();
-                                    return hdr;
+                                    return Either::Right(hdr);
                                 }
                             }
                             Err(err) => {

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/feldman.rs

@@ -6,6 +6,7 @@ use ark_serialize::{CanonicalSerialize, SerializationError, serialize_to_vec};
 use ark_std::marker::PhantomData;
 use ark_std::rand::Rng;
 use derive_more::{Deref, From, IntoIterator};
+use multisig::Committee;
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
@@ -33,6 +34,10 @@ impl FeldmanVssPublicParam {
         Self { t, n }
     }
 
+    pub fn from(c: &Committee) -> Self {
+        Self::new(c.one_honest_threshold(), c.size())
+    }
+
     pub fn threshold(&self) -> usize {
         self.t.get()
     }

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;
         }
     }
 }