Stream<CommitteeInfo> and separating pubsub backend handle

Author: alxiong

timeboost-contract/src/provider.rs

@@ -1,13 +1,21 @@
 //! Helper functions to build Ethereum [providers](https://docs.rs/alloy/latest/alloy/providers/trait.Provider.html)
 //! Partial Credit: <https://github.com/EspressoSystems/espresso-network/tree/main/contracts/rust/deployer>
 
+use std::{ops::Deref, pin::Pin};
+
 use alloy::{
+    eips::BlockNumberOrTag,
     network::EthereumWallet,
-    providers::ProviderBuilder,
+    primitives::Address,
+    providers::{Provider, ProviderBuilder},
+    rpc::types::{Filter, Log},
     signers::local::{LocalSignerError, MnemonicBuilder, PrivateKeySigner, coins_bip39::English},
-    transports::http::reqwest::Url,
+    sol_types::SolEvent,
+    transports::{http::reqwest::Url, ws::WsConnect},
 };
-use timeboost_types::HttpProviderWithWallet;
+use futures::{Stream, StreamExt};
+use timeboost_types::{HttpProvider, HttpProviderWithWallet};
+use tracing::error;
 
 /// Build a local signer from wallet mnemonic and account index
 pub fn build_signer(
@@ -31,3 +39,102 @@ pub fn build_provider(
     let wallet = EthereumWallet::from(signer);
     Ok(ProviderBuilder::new().wallet(wallet).connect_http(url))
 }
+
+/// A PubSub service (with backend handle), disconnect on drop.
+#[derive(Clone)]
+pub struct PubSubProvider(HttpProvider);
+
+impl Deref for PubSubProvider {
+    type Target = HttpProvider;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl PubSubProvider {
+    pub async fn new(ws_url: Url) -> anyhow::Result<Self> {
+        let ws = WsConnect::new(ws_url);
+        let provider = ProviderBuilder::new()
+            .connect_pubsub_with(ws)
+            .await
+            .map_err(|err| {
+                error!(?err, "event pubsub failed to start");
+                err
+            })?;
+        Ok(Self(provider))
+    }
+
+    pub fn inner(&self) -> &HttpProvider {
+        &self.0
+    }
+
+    /// create an event stream of event type `E`, subscribing since `from_block` on `contract`
+    pub async fn event_stream<E: SolEvent>(
+        &self,
+        contract: Address,
+        from_block: BlockNumberOrTag,
+    ) -> anyhow::Result<Pin<Box<dyn Stream<Item = Log<E>>>>> {
+        let filter = Filter::new()
+            .address(contract)
+            .event(E::SIGNATURE)
+            .from_block(from_block);
+
+        let events = self
+            .subscribe_logs(&filter)
+            .await
+            .map_err(|err| {
+                error!(?err, "pubsub subscription failed");
+                err
+            })?
+            .into_stream();
+
+        let validated = events.filter_map(|log| async move {
+            let Ok(event) = log.log_decode_validate::<E>() else {
+                error!("fail to parse CommitteeCreated event log");
+                return None;
+            };
+            Some(event)
+        });
+
+        Ok(Box::pin(validated))
+    }
+
+    /// Returns the smallest block number whose timestamp is >= `target_ts` through binary search.
+    /// Useful to determine `from_block` input of `Self::event_strea()` subscription.
+    pub async fn get_block_number_by_timestamp(
+        &self,
+        target_ts: u64,
+    ) -> anyhow::Result<Option<u64>> {
+        let latest = self.get_block_number().await?;
+        let mut lo: u64 = 0;
+        let mut hi: u64 = latest;
+
+        while lo <= hi {
+            let mid = lo + (hi - lo) / 2;
+
+            let block = match self
+                .0
+                .get_block_by_number(BlockNumberOrTag::Number(mid))
+                .await?
+            {
+                Some(b) => b,
+                None => {
+                    lo = mid + 1;
+                    continue;
+                }
+            };
+
+            if block.header.timestamp >= target_ts {
+                if mid == 0 {
+                    return Ok(Some(0));
+                }
+                hi = mid - 1;
+            } else {
+                lo = mid + 1;
+            }
+        }
+
+        // At this point, `lo` is the smallest index with ts >= target_ts (if any)
+        if lo > latest { Ok(None) } else { Ok(Some(lo)) }
+    }
+}

timeboost/src/binaries/timeboost.rs

@@ -94,8 +94,20 @@ async fn main() -> Result<()> {
 
     let pubkey = sign_keypair.public_key();
 
+    // optionally fetch previous committee info
+    let cid: u64 = cli.committee_id.into();
+    let prev_comm = if cid > 0u64 {
+        let c = &node_config.chain.parent;
+        let prev_comm =
+            CommitteeInfo::fetch(c.rpc_url.clone(), c.key_manager_contract, cid - 1).await?;
+        Some(prev_comm)
+    } else {
+        None
+    };
+
     let config = TimeboostConfig::builder()
         .sailfish_committee(sailfish_committee)
+        .maybe_prev_committee(prev_comm)
         .decrypt_committee(decrypt_committee)
         .certifier_committee(certifier_committee)
         .sign_keypair(sign_keypair)

timeboost/src/committee.rs

@@ -1,15 +1,11 @@
 //! Syncing committee info from the KeyManager contract
 
 use std::pin::Pin;
-use std::task::{Context, Poll};
 
 use alloy::{
     eips::BlockNumberOrTag,
     primitives::Address,
     providers::{Provider, ProviderBuilder},
-    rpc::types::Filter,
-    sol_types::SolEvent,
-    transports::ws::WsConnect,
 };
 use anyhow::{Context as AnyhowContext, Result};
 use cliquenet::AddressableCommittee;
@@ -18,11 +14,15 @@ use itertools::{Itertools, izip};
 use multisig::{Committee, CommitteeId, x25519};
 use timeboost_config::{CERTIFIER_PORT_OFFSET, DECRYPTER_PORT_OFFSET, ParentChain};
 use timeboost_contract::KeyManager::{self, CommitteeCreated};
+use timeboost_contract::provider::PubSubProvider;
 use timeboost_crypto::prelude::DkgEncKey;
 use timeboost_types::{KeyStore, Timestamp};
 use tracing::error;
 use url::Url;
 
+/// Type alias for the committee stream
+pub type NewCommitteeStream = Pin<Box<dyn Stream<Item = CommitteeInfo>>>;
+
 /// The committee info stored on the KeyManager contract, a subset of [`CommitteeConfig`]
 /// Keys and hosts are ordered in the same as they were registered (with KeyId from 0..n)
 #[derive(Debug, Clone)]
@@ -39,7 +39,14 @@ impl CommitteeInfo {
     /// Fetch the committee info for `committee_id` from `key_manager_addr` on chain
     pub async fn fetch(rpc: Url, key_manager_addr: Address, committee_id: u64) -> Result<Self> {
         let provider = ProviderBuilder::new().connect_http(rpc);
+        Self::fetch_with(provider, key_manager_addr, committee_id).await
+    }
 
+    pub(crate) async fn fetch_with(
+        provider: impl Provider,
+        key_manager_addr: Address,
+        committee_id: u64,
+    ) -> Result<Self> {
         let contract = KeyManager::new(key_manager_addr, &provider);
         let c = contract.getCommitteeById(committee_id).call().await?;
 
@@ -71,6 +78,10 @@ impl CommitteeInfo {
         })
     }
 
+    pub fn id(&self) -> CommitteeId {
+        self.id
+    }
+
     pub fn effective_timestamp(&self) -> Timestamp {
         self.timestamp
     }
@@ -138,64 +149,44 @@ impl CommitteeInfo {
                 .map(|(i, k)| (i as u8, k.clone())),
         )
     }
-}
 
-/// An pubsub-provider-holding event stream. (the pubsub will close on drop)
-pub struct NewCommitteeStream {
-    _provider: Box<dyn Provider>,
-    inner: Pin<Box<dyn Stream<Item = u64> + Send>>,
-}
-
-impl Stream for NewCommitteeStream {
-    type Item = u64;
-
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        self.inner.as_mut().poll_next(cx)
-    }
-}
-
-impl NewCommitteeStream {
-    pub async fn create(config: &ParentChain) -> Result<Self> {
-        // setup the websocket for contract event stream
-        let ws = WsConnect::new(config.ws_url.clone());
-        // spawn the pubsub service (and backend) and the frontend is registered at the provider
-        let provider = ProviderBuilder::new()
-            .connect_pubsub_with(ws)
+    /// subscribe an event stream
+    pub async fn new_committee_stream(
+        provider: &PubSubProvider,
+        start_ts: Timestamp,
+        config: &ParentChain,
+    ) -> Result<NewCommitteeStream> {
+        let from_block = provider
+            .get_block_number_by_timestamp(start_ts.into())
+            .await?
+            .unwrap_or_default();
+        let events = provider
+            .event_stream::<CommitteeCreated>(
+                config.key_manager_contract,
+                BlockNumberOrTag::Number(from_block),
+            )
             .await
-            .map_err(|err| {
-                error!(?err, "event pubsub failed to start");
-                err
+            .map_err(|e| {
+                error!("Failed to create CommitteeCreated stream: {:?}", e);
+                e
             })?;
 
-        let chain_id = config.id;
-        let tag = if chain_id == 31337 || chain_id == 1337 {
-            // local test chain, we start scanning from the genesis
-            BlockNumberOrTag::Number(0)
-        } else {
-            config.block_tag
-        };
-
-        let filter = Filter::new()
-            .address(config.key_manager_contract)
-            .event(KeyManager::CommitteeCreated::SIGNATURE)
-            .from_block(tag);
-        let events = provider
-            .subscribe_logs(&filter)
-            .await
-            .map_err(|err| {
-                error!(?err, "pubsub subscription failed");
-                err
-            })?
-            .into_stream();
-
-        let validated = events.filter_map(|log| async move {
-            log.log_decode_validate::<CommitteeCreated>()
-                .ok()
-                .map(|v| v.data().id)
+        let provider = provider.clone();
+        let key_manager_contract = config.key_manager_contract;
+        let s = events.filter_map(move |log| {
+            let provider = provider.clone();
+            async move {
+                let id = log.data().id;
+                match CommitteeInfo::fetch_with(provider.inner(), key_manager_contract, id).await {
+                    Ok(comm_info) => Some(comm_info),
+                    Err(_) => {
+                        error!(committee_id = %id, "fail to fetch new CommitteeInfo");
+                        None
+                    }
+                }
+            }
         });
-        Ok(Self {
-            _provider: Box::new(provider),
-            inner: Box::pin(validated),
-        })
+
+        Ok(Box::pin(s))
     }
 }

timeboost/src/conf.rs

@@ -7,20 +7,19 @@ use timeboost_crypto::prelude::DkgDecKey;
 use timeboost_sequencer::SequencerConfig;
 use timeboost_types::{KeyStore, ThresholdKeyCell};
 
+use crate::committee::CommitteeInfo;
+
 #[derive(Debug, Clone, Builder)]
 pub struct TimeboostConfig {
     /// The sailfish peers that this node will connect to.
     pub(crate) sailfish_committee: AddressableCommittee,
 
-    /// Previous sailfish peers
-    pub(crate) prev_sailfish_committee: Option<AddressableCommittee>,
+    /// Previous committee info stored on chain
+    pub(crate) prev_committee: Option<CommitteeInfo>,
 
     /// The decrypt peers that this node will connect to.
     pub(crate) decrypt_committee: AddressableCommittee,
 
-    /// Previous decrypt peers
-    pub(crate) prev_decrypt_committee: Option<(AddressableCommittee, KeyStore)>,
-
     /// The block certifier peers that this node will connect to.
     pub(crate) certifier_committee: AddressableCommittee,
 
@@ -80,8 +79,14 @@ impl TimeboostConfig {
             .sailfish_committee(self.sailfish_committee.clone())
             .decrypt_committee((self.decrypt_committee.clone(), self.key_store.clone()))
             .recover(self.recover)
-            .maybe_previous_sailfish_committee(self.prev_sailfish_committee.clone())
-            .maybe_previous_decrypt_committee(self.prev_decrypt_committee.clone())
+            .maybe_previous_sailfish_committee(
+                self.prev_committee.as_ref().map(|c| c.sailfish_committee()),
+            )
+            .maybe_previous_decrypt_committee(
+                self.prev_committee
+                    .as_ref()
+                    .map(|c| (c.decrypt_committee(), c.dkg_key_store())),
+            )
             .leash_len(self.leash_len)
             .threshold_dec_key(self.threshold_dec_key.clone())
             .chain_config(self.chain_config.clone())