linera_client/
chain_listener.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    collections::{btree_map::Entry, BTreeMap, BTreeSet},
6    sync::Arc,
7    time::Duration,
8};
9
10use futures::{
11    future::{join_all, select_all},
12    lock::Mutex,
13    Future, FutureExt as _, StreamExt,
14};
15use linera_base::{
16    crypto::{CryptoHash, Signer},
17    data_types::{ChainDescription, Epoch, Timestamp},
18    identifiers::{AccountOwner, BlobType, ChainId},
19    task::NonBlockingFuture,
20};
21use linera_core::{
22    client::{AbortOnDrop, ChainClient, ChainClientError, ListeningMode},
23    node::NotificationStream,
24    worker::{Notification, Reason},
25    Environment,
26};
27use linera_storage::{Clock as _, Storage as _};
28use tokio_util::sync::CancellationToken;
29use tracing::{debug, info, instrument, warn, Instrument as _};
30
31use crate::{
32    wallet::{UserChain, Wallet},
33    Error,
34};
35
36#[derive(Debug, Default, Clone, clap::Args, serde::Serialize)]
37pub struct ChainListenerConfig {
38    /// Do not create blocks automatically to receive incoming messages. Instead, wait for
39    /// an explicit mutation `processInbox`.
40    #[arg(
41        long = "listener-skip-process-inbox",
42        env = "LINERA_LISTENER_SKIP_PROCESS_INBOX"
43    )]
44    pub skip_process_inbox: bool,
45
46    /// Wait before processing any notification (useful for testing).
47    #[arg(
48        long = "listener-delay-before-ms",
49        default_value = "0",
50        env = "LINERA_LISTENER_DELAY_BEFORE"
51    )]
52    pub delay_before_ms: u64,
53
54    /// Wait after processing any notification (useful for rate limiting).
55    #[arg(
56        long = "listener-delay-after-ms",
57        default_value = "0",
58        env = "LINERA_LISTENER_DELAY_AFTER"
59    )]
60    pub delay_after_ms: u64,
61}
62
63type ContextChainClient<C> = ChainClient<<C as ClientContext>::Environment>;
64
65#[cfg_attr(not(web), trait_variant::make(Send))]
66#[allow(async_fn_in_trait)]
67pub trait ClientContext {
68    type Environment: linera_core::Environment;
69
70    fn wallet(&self) -> &Wallet;
71
72    fn storage(&self) -> &<Self::Environment as linera_core::Environment>::Storage;
73
74    fn client(&self) -> &Arc<linera_core::client::Client<Self::Environment>>;
75
76    /// Gets the timing sender for benchmarking, if available.
77    #[cfg(not(web))]
78    fn timing_sender(
79        &self,
80    ) -> Option<tokio::sync::mpsc::UnboundedSender<(u64, linera_core::client::TimingType)>>;
81
82    #[cfg(web)]
83    fn timing_sender(
84        &self,
85    ) -> Option<tokio::sync::mpsc::UnboundedSender<(u64, linera_core::client::TimingType)>> {
86        None
87    }
88
89    fn make_chain_client(&self, chain_id: ChainId) -> ChainClient<Self::Environment> {
90        let chain = self
91            .wallet()
92            .get(chain_id)
93            .cloned()
94            .unwrap_or_else(|| UserChain::make_other(chain_id, Timestamp::from(0)));
95        self.client().create_chain_client(
96            chain_id,
97            chain.block_hash,
98            chain.next_block_height,
99            chain.pending_proposal,
100            chain.owner,
101            self.timing_sender(),
102        )
103    }
104
105    async fn update_wallet_for_new_chain(
106        &mut self,
107        chain_id: ChainId,
108        owner: Option<AccountOwner>,
109        timestamp: Timestamp,
110        epoch: Epoch,
111    ) -> Result<(), Error>;
112
113    async fn update_wallet(&mut self, client: &ContextChainClient<Self>) -> Result<(), Error>;
114}
115
116#[allow(async_fn_in_trait)]
117pub trait ClientContextExt: ClientContext {
118    fn clients(&self) -> Vec<ContextChainClient<Self>> {
119        let chain_ids = self.wallet().chain_ids();
120        let mut clients = vec![];
121        for chain_id in chain_ids {
122            clients.push(self.make_chain_client(chain_id));
123        }
124        clients
125    }
126}
127
128impl<T: ClientContext> ClientContextExt for T {}
129
130/// A chain client together with the stream of notifications from the local node.
131///
132/// A background task listens to the validators and updates the local node, so any updates to
133/// this chain will trigger a notification. The background task is terminated when this gets
134/// dropped.
135struct ListeningClient<C: ClientContext> {
136    /// The chain client.
137    client: ContextChainClient<C>,
138    /// The abort handle for the task that listens to the validators.
139    abort_handle: AbortOnDrop,
140    /// The listening task's join handle.
141    join_handle: NonBlockingFuture<()>,
142    /// The stream of notifications from the local node.
143    notification_stream: Arc<Mutex<NotificationStream>>,
144    /// This is only `< u64::MAX` when the client is waiting for a timeout to process the inbox.
145    timeout: Timestamp,
146    /// The mode of listening to this chain.
147    listening_mode: ListeningMode,
148}
149
150impl<C: ClientContext> ListeningClient<C> {
151    fn new(
152        client: ContextChainClient<C>,
153        abort_handle: AbortOnDrop,
154        join_handle: NonBlockingFuture<()>,
155        notification_stream: NotificationStream,
156        listening_mode: ListeningMode,
157    ) -> Self {
158        Self {
159            client,
160            abort_handle,
161            join_handle,
162            #[allow(clippy::arc_with_non_send_sync)] // Only `Send` with `futures-util/alloc`.
163            notification_stream: Arc::new(Mutex::new(notification_stream)),
164            timeout: Timestamp::from(u64::MAX),
165            listening_mode,
166        }
167    }
168
169    async fn stop(self) {
170        drop(self.abort_handle);
171        if let Err(error) = self.join_handle.await {
172            warn!("Failed to join listening task: {error:?}");
173        }
174    }
175}
176
177/// A `ChainListener` is a process that listens to notifications from validators and reacts
178/// appropriately.
179pub struct ChainListener<C: ClientContext> {
180    context: Arc<Mutex<C>>,
181    storage: <C::Environment as Environment>::Storage,
182    config: Arc<ChainListenerConfig>,
183    listening: BTreeMap<ChainId, ListeningClient<C>>,
184    /// Map from publishing chain to subscriber chains.
185    /// Events emitted on the _publishing chain_ are of interest to the _subscriber chains_.
186    event_subscribers: BTreeMap<ChainId, BTreeSet<ChainId>>,
187    cancellation_token: CancellationToken,
188}
189
190impl<C: ClientContext> ChainListener<C> {
191    /// Creates a new chain listener given client chains.
192    pub fn new(
193        config: ChainListenerConfig,
194        context: Arc<Mutex<C>>,
195        storage: <C::Environment as Environment>::Storage,
196        cancellation_token: CancellationToken,
197    ) -> Self {
198        Self {
199            storage,
200            context,
201            config: Arc::new(config),
202            listening: Default::default(),
203            event_subscribers: Default::default(),
204            cancellation_token,
205        }
206    }
207
208    /// Runs the chain listener.
209    #[instrument(skip(self))]
210    pub async fn run(mut self) -> Result<impl Future<Output = Result<(), Error>>, Error> {
211        let chain_ids = {
212            let guard = self.context.lock().await;
213            let admin_chain_id = guard.wallet().genesis_admin_chain();
214            guard
215                .make_chain_client(admin_chain_id)
216                .synchronize_from_validators()
217                .await?;
218            BTreeMap::from_iter(
219                guard
220                    .wallet()
221                    .chain_ids()
222                    .into_iter()
223                    .chain([admin_chain_id])
224                    .map(|chain_id| (chain_id, ListeningMode::FullChain)),
225            )
226        };
227
228        Ok(async {
229            self.listen_recursively(chain_ids).await?;
230            loop {
231                match self.next_action().await? {
232                    Action::ProcessInbox(chain_id) => self.maybe_process_inbox(chain_id).await?,
233                    Action::Notification(notification) => {
234                        self.process_notification(notification).await?
235                    }
236                    Action::Stop => break,
237                }
238            }
239            join_all(self.listening.into_values().map(|client| client.stop())).await;
240            Ok(())
241        })
242    }
243
244    /// Processes a notification, updating local chains and validators as needed.
245    async fn process_notification(&mut self, notification: Notification) -> Result<(), Error> {
246        Self::sleep(self.config.delay_before_ms).await;
247        let Some(listening_mode) = self
248            .listening
249            .get(&notification.chain_id)
250            .map(|listening_client| &listening_client.listening_mode)
251        else {
252            warn!(
253                ?notification,
254                "ChainListener::process_notification: got a notification without listening to the chain"
255            );
256            return Ok(());
257        };
258
259        match &notification.reason {
260            Reason::NewIncomingBundle { .. } => {
261                self.maybe_process_inbox(notification.chain_id).await?;
262            }
263            Reason::NewRound { .. } => self.update_validators(&notification).await?,
264            Reason::NewBlock { hash, .. } => {
265                if matches!(listening_mode, ListeningMode::EventsOnly(_)) {
266                    debug!("ChainListener::process_notification: ignoring notification due to listening mode");
267                    return Ok(());
268                }
269                self.update_wallet(notification.chain_id).await?;
270                self.add_new_chains(*hash).await?;
271                let publishers = self
272                    .update_event_subscriptions(notification.chain_id)
273                    .await?;
274                if !publishers.is_empty() {
275                    self.listen_recursively(publishers).await?;
276                    self.maybe_process_inbox(notification.chain_id).await?;
277                }
278                self.process_new_events(notification.chain_id).await?;
279            }
280            Reason::NewEvents { event_streams, .. } => {
281                let should_process = match listening_mode {
282                    ListeningMode::FullChain => true,
283                    ListeningMode::EventsOnly(relevant_events) => {
284                        relevant_events.intersection(event_streams).count() != 0
285                    }
286                };
287                if !should_process {
288                    debug!(
289                        ?notification,
290                        ?listening_mode,
291                        "ChainListener::process_notification: ignoring notification due to no relevant events",
292                    );
293                    return Ok(());
294                }
295                self.process_new_events(notification.chain_id).await?;
296            }
297        }
298        Self::sleep(self.config.delay_after_ms).await;
299        Ok(())
300    }
301
302    /// If any new chains were created by the given block, and we have a key pair for them,
303    /// add them to the wallet and start listening for notifications.
304    async fn add_new_chains(&mut self, hash: CryptoHash) -> Result<(), Error> {
305        let block = self
306            .storage
307            .read_confirmed_block(hash)
308            .await?
309            .ok_or(ChainClientError::MissingConfirmedBlock(hash))?
310            .into_block();
311        let blobs = block.created_blobs().into_iter();
312        let new_chains = blobs
313            .filter_map(|(blob_id, blob)| {
314                if blob_id.blob_type == BlobType::ChainDescription {
315                    let chain_desc: ChainDescription = bcs::from_bytes(blob.content().bytes())
316                        .expect("ChainDescription should deserialize correctly");
317                    let owners = chain_desc.config().ownership.all_owners().cloned();
318                    Some((ChainId(blob_id.hash), owners.collect::<Vec<_>>()))
319                } else {
320                    None
321                }
322            })
323            .collect::<Vec<_>>();
324        if new_chains.is_empty() {
325            return Ok(());
326        }
327        let mut new_ids = BTreeMap::new();
328        let mut context_guard = self.context.lock().await;
329        for (new_chain_id, owners) in new_chains {
330            for chain_owner in owners {
331                if context_guard
332                    .client()
333                    .signer()
334                    .contains_key(&chain_owner)
335                    .await
336                    .map_err(ChainClientError::signer_failure)?
337                {
338                    context_guard
339                        .update_wallet_for_new_chain(
340                            new_chain_id,
341                            Some(chain_owner),
342                            block.header.timestamp,
343                            block.header.epoch,
344                        )
345                        .await?;
346                    new_ids.insert(new_chain_id, ListeningMode::FullChain);
347                }
348            }
349        }
350        drop(context_guard);
351        self.listen_recursively(new_ids).await?;
352        Ok(())
353    }
354
355    /// Processes the inboxes of all chains that are subscribed to `chain_id`.
356    async fn process_new_events(&mut self, chain_id: ChainId) -> Result<(), Error> {
357        let Some(subscribers) = self.event_subscribers.get(&chain_id).cloned() else {
358            return Ok(());
359        };
360        for subscriber_id in subscribers {
361            self.maybe_process_inbox(subscriber_id).await?;
362        }
363        Ok(())
364    }
365
366    /// Starts listening for notifications about the given chains, and any chains that publish
367    /// event streams those chains are subscribed to.
368    async fn listen_recursively(
369        &mut self,
370        mut chain_ids: BTreeMap<ChainId, ListeningMode>,
371    ) -> Result<(), Error> {
372        while let Some((chain_id, listening_mode)) = chain_ids.pop_first() {
373            for (new_chain_id, new_listening_mode) in self.listen(chain_id, listening_mode).await? {
374                match chain_ids.entry(new_chain_id) {
375                    Entry::Vacant(vacant) => {
376                        vacant.insert(new_listening_mode);
377                    }
378                    Entry::Occupied(mut occupied) => {
379                        occupied.get_mut().extend(Some(new_listening_mode));
380                    }
381                }
382            }
383        }
384
385        Ok(())
386    }
387
388    /// Starts listening for notifications about the given chain.
389    ///
390    /// Returns all publishing chains, that we also need to listen to.
391    async fn listen(
392        &mut self,
393        chain_id: ChainId,
394        mut listening_mode: ListeningMode,
395    ) -> Result<BTreeMap<ChainId, ListeningMode>, Error> {
396        if self
397            .listening
398            .get(&chain_id)
399            .is_some_and(|existing_client| existing_client.listening_mode >= listening_mode)
400        {
401            return Ok(BTreeMap::new());
402        }
403        listening_mode.extend(
404            self.listening
405                .get(&chain_id)
406                .map(|existing_client| existing_client.listening_mode.clone()),
407        );
408        let client = self.context.lock().await.make_chain_client(chain_id);
409        let (listener, abort_handle, notification_stream) =
410            client.listen(listening_mode.clone()).await?;
411        let join_handle = linera_base::task::spawn(listener.in_current_span());
412        let listening_client = ListeningClient::new(
413            client,
414            abort_handle,
415            join_handle,
416            notification_stream,
417            listening_mode,
418        );
419        self.listening.insert(chain_id, listening_client);
420        let publishing_chains = self.update_event_subscriptions(chain_id).await?;
421        self.maybe_process_inbox(chain_id).await?;
422        Ok(publishing_chains)
423    }
424
425    /// Updates the event subscribers map, and returns all publishing chains we need to listen to.
426    async fn update_event_subscriptions(
427        &mut self,
428        chain_id: ChainId,
429    ) -> Result<BTreeMap<ChainId, ListeningMode>, Error> {
430        let listening_client = self.listening.get_mut(&chain_id).expect("missing client");
431        if !listening_client.client.is_tracked() {
432            return Ok(BTreeMap::new());
433        }
434        let publishing_chains: BTreeMap<_, _> = listening_client
435            .client
436            .event_stream_publishers()
437            .await?
438            .into_iter()
439            .map(|(chain_id, streams)| (chain_id, ListeningMode::EventsOnly(streams)))
440            .collect();
441        for publisher_id in publishing_chains.keys() {
442            self.event_subscribers
443                .entry(*publisher_id)
444                .or_default()
445                .insert(chain_id);
446        }
447        Ok(publishing_chains)
448    }
449
450    /// Returns the next notification or timeout to process.
451    async fn next_action(&mut self) -> Result<Action, Error> {
452        loop {
453            let (timeout_chain_id, timeout) = self.next_timeout()?;
454            let notification_futures = self
455                .listening
456                .values_mut()
457                .map(|client| {
458                    let stream = client.notification_stream.clone();
459                    Box::pin(async move { stream.lock().await.next().await })
460                })
461                .collect::<Vec<_>>();
462            futures::select! {
463                () = self.cancellation_token.cancelled().fuse() => {
464                    return Ok(Action::Stop);
465                }
466                () = self.storage.clock().sleep_until(timeout).fuse() => {
467                    return Ok(Action::ProcessInbox(timeout_chain_id));
468                }
469                (maybe_notification, index, _) = select_all(notification_futures).fuse() => {
470                    let Some(notification) = maybe_notification else {
471                        let chain_id = *self.listening.keys().nth(index).unwrap();
472                        self.listening.remove(&chain_id);
473                        warn!("Notification stream for {chain_id} closed");
474                        continue;
475                    };
476                    return Ok(Action::Notification(notification));
477                }
478            }
479        }
480    }
481
482    /// Returns the next timeout to process, and the chain to which it applies.
483    fn next_timeout(&self) -> Result<(ChainId, Timestamp), Error> {
484        let (chain_id, client) = self
485            .listening
486            .iter()
487            .min_by_key(|(_, client)| client.timeout)
488            .expect("No chains left to listen to");
489        Ok((*chain_id, client.timeout))
490    }
491
492    /// Updates the validators about the chain.
493    async fn update_validators(&self, notification: &Notification) -> Result<(), Error> {
494        let chain_id = notification.chain_id;
495        let listening_client = self.listening.get(&chain_id).expect("missing client");
496        if let Err(error) = listening_client.client.update_validators(None).await {
497            warn!(
498                "Failed to update validators about the local chain after \
499                 receiving {notification:?} with error: {error:?}"
500            );
501        }
502        Ok(())
503    }
504
505    /// Updates the wallet based on the client for this chain.
506    async fn update_wallet(&self, chain_id: ChainId) -> Result<(), Error> {
507        let client = &self
508            .listening
509            .get(&chain_id)
510            .expect("missing client")
511            .client;
512        self.context.lock().await.update_wallet(client).await?;
513        Ok(())
514    }
515
516    /// Processes the inbox, unless `skip_process_inbox` is set.
517    ///
518    /// If no block can be produced because we are not the round leader, a timeout is returned
519    /// for when to retry; otherwise `u64::MAX` is returned.
520    ///
521    /// The wallet is persisted with any blocks that processing the inbox added. An error
522    /// is returned if persisting the wallet fails.
523    async fn maybe_process_inbox(&mut self, chain_id: ChainId) -> Result<(), Error> {
524        if self.config.skip_process_inbox {
525            debug!("Not processing inbox for {chain_id:.8} due to listener configuration");
526            return Ok(());
527        }
528        let listening_client = self.listening.get_mut(&chain_id).expect("missing client");
529        if !listening_client.client.is_tracked() {
530            debug!("Not processing inbox for non-tracked chain {chain_id:.8}");
531            return Ok(());
532        }
533        debug!("Processing inbox for {chain_id:.8}");
534        listening_client.timeout = Timestamp::from(u64::MAX);
535        match listening_client
536            .client
537            .process_inbox_without_prepare()
538            .await
539        {
540            Err(ChainClientError::CannotFindKeyForChain(chain_id)) => {
541                debug!(%chain_id, "Cannot find key for chain");
542            }
543            Err(error) => warn!(%error, "Failed to process inbox."),
544            Ok((certs, None)) => info!("Done processing inbox. {} blocks created.", certs.len()),
545            Ok((certs, Some(new_timeout))) => {
546                info!(
547                    "{} blocks created. Will try processing the inbox later based \
548                     on the given round timeout: {new_timeout:?}",
549                    certs.len(),
550                );
551                listening_client.timeout = new_timeout.timestamp;
552            }
553        }
554        let mut context_guard = self.context.lock().await;
555        context_guard
556            .update_wallet(&listening_client.client)
557            .await?;
558        Ok(())
559    }
560
561    /// Sleeps for the given number of milliseconds, if greater than 0.
562    async fn sleep(delay_ms: u64) {
563        if delay_ms > 0 {
564            linera_base::time::timer::sleep(Duration::from_millis(delay_ms)).await;
565        }
566    }
567}
568
569enum Action {
570    ProcessInbox(ChainId),
571    Notification(Notification),
572    Stop,
573}
linera_client/chain_listener.rs

linera_client/
chain_listener.rs
