Add Multiwatcher.

Author: twittner

robusta/Cargo.toml

@@ -8,6 +8,7 @@ rust-version.workspace = true
 [dependencies]
 bincode = { workspace = true }
 bon = { workspace = true }
+bytes = { workspace = true }
 data-encoding = { workspace = true }
 either = { workspace = true }
 espresso-types = { workspace = true }
@@ -21,6 +22,7 @@ serde_json = { workspace = true }
 thiserror = { workspace = true }
 timeboost-types = { path = "../timeboost-types" }
 tokio = { workspace = true }
+tokio-stream = { workspace = true }
 tokio-tungstenite = { workspace = true }
 tracing = { workspace = true }
 url = { workspace = true }

robusta/src/config.rs

@@ -1,9 +1,9 @@
-use std::{iter::repeat, time::Duration};
+use std::{iter::repeat, num::NonZeroUsize, time::Duration};
 
 use bon::Builder;
-use url::{ParseError, Url};
+use url::Url;
 
-const NUM_DELAYS: usize = 5;
+const NUM_DELAYS: NonZeroUsize = NonZeroUsize::new(5).expect("5 > 0");
 
 #[derive(Debug, Clone, Builder)]
 pub struct Config {
@@ -12,25 +12,29 @@ pub struct Config {
     pub(crate) label: String,
 
     /// Espresso network base URL.
-    #[builder(with = |s: &str| -> Result<_, ParseError> { Url::parse(s) })]
     pub(crate) base_url: Url,
 
     /// Espresso network websocket base URL.
-    #[builder(with = |s: &str| -> Result<_, ParseError> { Url::parse(s) })]
     pub(crate) wss_base_url: Url,
 
     /// The sequence of delays between successive requests.
     ///
     /// The last value is repeated forever.
     #[builder(default = [1, 3, 5, 10, 15])]
-    pub(crate) delays: [u8; NUM_DELAYS],
+    pub(crate) delays: [u8; NUM_DELAYS.get()],
 }
 
 impl Config {
+    pub fn with_websocket_base_url(&self, u: Url) -> Self {
+        let mut c = self.clone();
+        c.wss_base_url = u;
+        c
+    }
+
     pub fn delay_iter(&self) -> impl Iterator<Item = Duration> + use<> {
         self.delays
             .into_iter()
-            .chain(repeat(self.delays[NUM_DELAYS - 1]))
+            .chain(repeat(self.delays[NUM_DELAYS.get() - 1]))
             .map(|n| Duration::from_secs(n.into()))
     }
 }

robusta/src/lib.rs

@@ -1,4 +1,5 @@
 mod config;
+mod multiwatcher;
 mod types;
 mod watcher;
 
@@ -18,8 +19,9 @@ use tracing::{debug, warn};
 
 use crate::types::{TX, TaggedBase64, TransactionsWithProof, VidCommonResponse};
 
+pub use crate::multiwatcher::Multiwatcher;
 pub use crate::types::Height;
-pub use crate::watcher::{WatchError, watch};
+pub use crate::watcher::{WatchError, Watcher};
 pub use config::{Config, ConfigBuilder};
 pub use espresso_types;
 
@@ -276,10 +278,7 @@ fn deserialize<T: DeserializeOwned>(d: &[u8]) -> Result<T, Error> {
 
 #[cfg(test)]
 mod tests {
-    use futures::StreamExt;
-    use tokio::pin;
-
-    use super::{Client, Config};
+    use super::{Client, Config, Watcher};
 
     #[tokio::test]
     async fn decaf_smoke() {
@@ -288,17 +287,23 @@ mod tests {
             .try_init();
 
         let cfg = Config::builder()
-            .base_url("https://query.decaf.testnet.espresso.network/v1/")
-            .unwrap()
-            .wss_base_url("wss://query.decaf.testnet.espresso.network/v1/")
-            .unwrap()
+            .base_url(
+                "https://query.decaf.testnet.espresso.network/v1/"
+                    .parse()
+                    .unwrap(),
+            )
+            .wss_base_url(
+                "wss://query.decaf.testnet.espresso.network/v1/"
+                    .parse()
+                    .unwrap(),
+            )
             .label("decaf_smoke")
             .build();
 
         let clt = Client::new(cfg.clone());
         let height = clt.height().await.unwrap();
-        let headers = super::watch(&cfg, height, None).await.unwrap();
-        pin!(headers);
-        assert_eq!(u64::from(height), headers.next().await.unwrap().height());
+        let mut watcher = Watcher::new(cfg, height, None);
+        let header = watcher.next().await;
+        assert_eq!(u64::from(height), header.height());
     }
 }

robusta/src/multiwatcher.rs

@@ -0,0 +1,100 @@
+use std::collections::{BTreeMap, HashMap, HashSet};
+
+use crate::{Config, Height, Watcher};
+use espresso_types::{Header, NamespaceId};
+use futures::{StreamExt, stream::SelectAll};
+use tokio::{spawn, sync::mpsc, task::JoinHandle};
+use tokio_stream::wrappers::ReceiverStream;
+use tracing::{debug, warn};
+use url::Url;
+
+#[derive(Debug)]
+pub struct Multiwatcher {
+    height: Height,
+    threshold: usize,
+    watchers: Vec<JoinHandle<()>>,
+    headers: BTreeMap<Height, HashMap<Header, HashSet<Id>>>,
+    stream: SelectAll<ReceiverStream<(Id, Header)>>,
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+struct Id(usize);
+
+impl Drop for Multiwatcher {
+    fn drop(&mut self) {
+        for w in &self.watchers {
+            w.abort();
+        }
+    }
+}
+
+impl Multiwatcher {
+    pub fn new<C, H, I, N>(configs: C, height: H, nsid: N, threshold: usize) -> Self
+    where
+        C: IntoIterator<Item = Config>,
+        H: Into<Height>,
+        I: IntoIterator<Item = Url>,
+        N: Into<NamespaceId>,
+    {
+        let height = height.into();
+        let nsid = nsid.into();
+        let mut stream = SelectAll::new();
+        let mut watchers = Vec::new();
+        for (i, c) in configs.into_iter().enumerate() {
+            let (tx, rx) = mpsc::channel(32);
+            stream.push(ReceiverStream::new(rx));
+            watchers.push(spawn(async move {
+                let id = Id(i);
+                let mut w = Watcher::new(c, height, nsid);
+                while tx.send((id, w.next().await)).await.is_ok() {}
+            }));
+        }
+        Self {
+            height,
+            threshold,
+            stream,
+            watchers,
+            headers: BTreeMap::from_iter([(height, HashMap::new())]),
+        }
+    }
+
+    pub async fn next(&mut self) -> Option<Header> {
+        loop {
+            let (i, hdr) = self.stream.next().await?;
+            let h = Height::from(hdr.height());
+            if Some(h) < self.headers.first_entry().map(|e| *e.key()) {
+                debug!(%h, "ignoring header below minimum height");
+                continue;
+            }
+            if self.has_voted(h, i) {
+                warn!(%h, "source sent multiple headers for same height");
+                continue;
+            }
+            let votes = self.headers.entry(h).or_default();
+            if let Some(ids) = votes.get(&hdr)
+                && ids.len() + 1 >= self.threshold
+            {
+                self.gc(h);
+                return Some(hdr);
+            }
+            votes.entry(hdr).or_default().insert(i);
+        }
+    }
+
+    fn has_voted(&self, height: Height, id: Id) -> bool {
+        let Some(m) = self.headers.get(&height) else {
+            return false;
+        };
+        for v in m.values() {
+            if v.contains(&id) {
+                return true;
+            }
+        }
+        false
+    }
+
+    fn gc(&mut self, height: Height) {
+        self.headers.retain(|h, _| *h >= height);
+        self.height = height;
+    }
+}

robusta/src/types.rs

@@ -19,7 +19,9 @@ pub(crate) struct VidCommonResponse {
 
 macro_rules! Primitive {
     ($name:ident, $t:ty) => {
-        #[derive(Debug, Copy, Clone, Deserialize, Serialize)]
+        #[derive(
+            Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize,
+        )]
         #[serde(transparent)]
         pub struct $name($t);