|
| 1 | +use std::{collections::HashSet, future::Future, marker::PhantomData}; |
| 2 | + |
| 3 | +use alloy_rpc_types_beacon::BlsPublicKey; |
| 4 | +use cb_crypto::{ |
| 5 | + manager::{Signer, SigningManager}, |
| 6 | + service::SigningService, |
| 7 | + types::SignRequest, |
| 8 | +}; |
| 9 | +use cb_pbs::{ |
| 10 | + BuilderApi, BuilderApiState, BuilderEvent, BuilderState, DefaultBuilderApi, PbsService, |
| 11 | +}; |
| 12 | +use tokio::sync::{ |
| 13 | + broadcast, |
| 14 | + mpsc::{self, unbounded_channel, UnboundedSender}, |
| 15 | +}; |
| 16 | + |
| 17 | +pub type SignRequestSender = mpsc::UnboundedSender<SignRequest>; |
| 18 | + |
| 19 | +pub struct Runner<S: BuilderApiState = (), T: BuilderApi<S> = DefaultBuilderApi> { |
| 20 | + state: BuilderState<S>, |
| 21 | + commit_ids: HashSet<String>, |
| 22 | + hooks_ids: HashSet<String>, |
| 23 | + sign_manager: SigningManager, |
| 24 | + |
| 25 | + notif_tx: SignRequestSender, |
| 26 | + notif_rx: mpsc::UnboundedReceiver<SignRequest>, |
| 27 | + _marker: PhantomData<T>, |
| 28 | +} |
| 29 | + |
| 30 | +impl<S: BuilderApiState, T: BuilderApi<S>> Runner<S, T> { |
| 31 | + pub fn new(state: BuilderState<S>) -> Self { |
| 32 | + let (notif_tx, notif_rx) = unbounded_channel(); |
| 33 | + |
| 34 | + // TODO: move this in run + spawn only if needed |
| 35 | + let mut sign_manager = SigningManager::new(state.chain); |
| 36 | + sign_manager.add_consensus_signer(Signer::new_random()); |
| 37 | + |
| 38 | + Self { |
| 39 | + state, |
| 40 | + commit_ids: HashSet::new(), |
| 41 | + hooks_ids: HashSet::new(), |
| 42 | + sign_manager, |
| 43 | + notif_tx, |
| 44 | + notif_rx, |
| 45 | + _marker: PhantomData, |
| 46 | + } |
| 47 | + } |
| 48 | + |
| 49 | + pub fn add_commitment<F, R>(&mut self, commit_id: impl Into<String>, commitment: F) |
| 50 | + where |
| 51 | + F: FnOnce(UnboundedSender<SignRequest>, Vec<BlsPublicKey>) -> R + 'static, |
| 52 | + R: Future<Output = eyre::Result<()>> + Send + 'static, |
| 53 | + { |
| 54 | + let id = commit_id.into(); |
| 55 | + |
| 56 | + if !self.commit_ids.insert(id.clone()) { |
| 57 | + eprintln!("Commitments ids need to be unique, found duplicate: {id}"); |
| 58 | + std::process::exit(1); |
| 59 | + } |
| 60 | + |
| 61 | + // move to vector and spawn after signing service |
| 62 | + tokio::spawn(commitment(self.notif_tx.clone(), self.sign_manager.consensus_pubkeys())); |
| 63 | + } |
| 64 | + |
| 65 | + pub fn add_boost_hook<F, R>(&mut self, hook_id: impl Into<String>, hook: F) |
| 66 | + where |
| 67 | + F: FnOnce(broadcast::Receiver<BuilderEvent>) -> R + 'static, |
| 68 | + R: Future<Output = eyre::Result<()>> + Send + 'static, |
| 69 | + { |
| 70 | + let id = hook_id.into(); |
| 71 | + |
| 72 | + if !self.hooks_ids.insert(id.clone()) { |
| 73 | + eprintln!("Hook ids need to be unique, found duplicate: {id}"); |
| 74 | + std::process::exit(1); |
| 75 | + } |
| 76 | + |
| 77 | + // move to vector and spawn after signing service |
| 78 | + tokio::spawn(hook(self.state.subscribe_events())); |
| 79 | + } |
| 80 | + |
| 81 | + pub async fn run(self) -> eyre::Result<()> { |
| 82 | + // start signature service |
| 83 | + if !self.commit_ids.is_empty() { |
| 84 | + let sign_service = SigningService::new(self.sign_manager, self.notif_rx); |
| 85 | + tokio::spawn(sign_service.run()); |
| 86 | + } |
| 87 | + |
| 88 | + // TODO: start commitments and hooks here |
| 89 | + |
| 90 | + // start boost service |
| 91 | + PbsService::run::<S, T>(self.state).await; |
| 92 | + |
| 93 | + Ok(()) |
| 94 | + } |
| 95 | +} |
0 commit comments