tempo_consensus/follow/upstream/

actor.rs

use std::{sync::Arc, time::Duration};

use commonware_consensus::{Reporter, types::Height};
use commonware_runtime::{Clock, ContextCell, Metrics, Spawner, spawn_cell};
use eyre::{Report, WrapErr as _};
use futures::{
    FutureExt as _, StreamExt as _,
    future::{BoxFuture, Either},
    stream::{self, Fuse, FusedStream},
};
use jsonrpsee::{
    core::{client, client::Subscription},
    ws_client::{WsClient, WsClientBuilder},
};
use rand_08::Rng as _;
use tempo_node::rpc::consensus::{CertifiedBlock, Event, Query, TempoConsensusApiClient};
use tempo_telemetry_util::display_duration;
use tokio::{
    select,
    sync::{mpsc, oneshot},
};
use tracing::{debug, debug_span, instrument, warn, warn_span};

use crate::utils::OptionFuture;

pub(super) type EventStream =
    Either<stream::Empty<Result<Event, serde_json::Error>>, Fuse<Subscription<Event>>>;

const RECONNECT_BACKOFF_FACTOR: u64 = 2;
const RECONNECT_MAX_BACKOFF: Duration = Duration::from_secs(20);
const RECONNECT_JITTER: Duration = Duration::from_secs(1);

/// Manages the connection to the upstream node.
///
/// This actor holds the websocket connection to the upstream node, reconnecting
/// it if necessary.
pub(crate) struct Actor<TContext> {
    pub(super) context: ContextCell<TContext>,
    pub(super) connection: Option<Arc<WsClient>>,
    pub(super) mailbox: mpsc::UnboundedReceiver<super::ingress::Message>,
    pub(super) url: &'static str,
    pub(super) pending_connect: OptionFuture<BoxFuture<'static, (u64, eyre::Result<WsClient>)>>,
    pub(super) pending_stream:
        OptionFuture<BoxFuture<'static, Result<Subscription<Event>, client::Error>>>,
    pub(super) event_stream: EventStream,
    /// Requests for blocks while the actor is trying to establish a connection.
    pub(super) waiters: Vec<(Height, oneshot::Sender<Option<CertifiedBlock>>)>,
}

impl<TContext> Actor<TContext>
where
    TContext: Clock + Metrics + Spawner,
{
    pub(crate) fn start(
        mut self,
        reporter: impl Reporter<Activity = Event>,
    ) -> commonware_runtime::Handle<()> {
        spawn_cell!(self.context, self.run(reporter))
    }

    async fn run(mut self, mut reporter: impl Reporter<Activity = Event>) {
        loop {
            self.reconnect_or_resubscribe();
            self.drain_waiters();

            select!(
                biased;

                (attempts, client) = &mut self.pending_connect => {
                    match client {
                        Ok(client) => {
                            let client = Arc::new(client);
                            self.connection.replace(client);
                        }
                        Err(reason) => {
                            let reconnect_in = reconnect_delay(attempts);
                            warn_span!("reconnect").in_scope(|| warn!(
                                %reason,
                                attempts,
                                reconnect_in = %display_duration(reconnect_in),
                                url = self.url,
                                "connecting to upstream node failed, attempting again",
                            ));
                            self.pending_connect.replace({
                                let context = self.context.clone();
                                let url = self.url;
                                async move {
                                    context.sleep(reconnect_in).await;
                                    connect(url, attempts.saturating_add(1)).await
                                }.boxed()
                            });
                        }
                    }
                }

                stream = &mut self.pending_stream => {
                    match stream {
                        Ok(stream) => {
                        debug_span!("consensus_event_subscription")
                            .in_scope(|| debug!("subscription for consensus events established"));
                            self.event_stream = active_event_stream(stream);
                        }
                        Err(error) => {
                            warn_span!("event_subscription").in_scope(|| warn!(
                                reason = %Report::new(error),
                                "failed subscribing to events; reconnecting to upstream node"
                            ));
                            self.connection.take();
                            self.event_stream = inactive_event_stream();
                        }
                    }
                }

                event = self.event_stream.next(), if !self.event_stream.is_terminated() => {
                    match event {
                        Some(Ok(event)) => {
                            debug_span!("consensus_event").in_scope(|| debug!(
                                ?event, "received consensus event, forwarding to reporter"
                            ));
                            reporter.report(event).await;
                        }
                        Some(Err(error)) => {
                            warn_span!("event").in_scope(|| warn!(
                                %error,
                                "event stream encountered an error",
                            ));
                            self.event_stream = inactive_event_stream();
                        }
                        None => {
                            warn_span!("event_subscription").in_scope(|| warn!(
                                url = self.url,
                                "event stream terminated",
                            ));
                            self.event_stream = inactive_event_stream();
                        }
                    }
                }

                Some(msg) = self.mailbox.recv() => {
                    match msg {
                        super::ingress::Message::GetFinalization { height, response, } => {
                            self.waiters.push((height, response));
                        }
                    }
                }
            );
        }
    }

    #[instrument(skip_all)]
    fn reconnect_or_resubscribe(&mut self) {
        if self.pending_connect.is_some() || self.pending_stream.is_some() {
            return;
        }

        let Some(client) = self.connection.clone() else {
            self.pending_connect.replace(connect(self.url, 1));
            return;
        };

        if !self.event_stream.is_terminated() {
            return;
        }

        if client.is_connected() {
            self.pending_stream.replace(subscribe(client));
        } else {
            warn!(url = self.url, "upstream client disconnected, reconnecting");
            self.connection.take();
            self.pending_connect.replace(connect(self.url, 1));
        }
    }

    /// Drains the waiters by fetching the finalizations they are waiting for.
    ///
    /// Only executes if a client is present and connected.
    fn drain_waiters(&mut self) {
        if self.pending_connect.is_some()
            || self.pending_stream.is_some()
            || self.event_stream.is_terminated()
        {
            return;
        }

        let Some(client) = &self.connection else {
            return;
        };
        if !client.is_connected() {
            return;
        }

        for (height, response) in self.waiters.drain(..) {
            let client = client.clone();
            self.context
                .with_label("get_finalization")
                .spawn(move |_| get_finalization(client, height, response));
        }
    }
}

fn connect(url: &'static str, attempts: u64) -> BoxFuture<'static, (u64, eyre::Result<WsClient>)> {
    async move {
        (
            attempts,
            WsClientBuilder::default()
                .build(&url)
                .await
                .map_err(Report::new),
        )
    }
    .boxed()
}

fn subscribe(
    client: Arc<WsClient>,
) -> BoxFuture<'static, Result<Subscription<Event>, client::Error>> {
    async move { client.subscribe_events().await }.boxed()
}

pub(super) fn inactive_event_stream() -> EventStream {
    Either::Left(stream::empty())
}

fn active_event_stream(stream: Subscription<Event>) -> EventStream {
    Either::Right(stream.fuse())
}

fn reconnect_delay(attempts: u64) -> Duration {
    reconnect_backoff(attempts) + random_jitter()
}

fn reconnect_backoff(attempts: u64) -> Duration {
    let backoff_secs = attempts.saturating_mul(RECONNECT_BACKOFF_FACTOR);
    let backoff = Duration::from_secs(backoff_secs);

    backoff.min(RECONNECT_MAX_BACKOFF)
}

fn random_jitter() -> Duration {
    let max_jitter_millis = RECONNECT_JITTER.as_millis() as u64;
    Duration::from_millis(rand_08::thread_rng().gen_range(0..=max_jitter_millis))
}

#[instrument(skip_all, fields(%height), err)]
async fn get_finalization(
    client: Arc<WsClient>,
    height: Height,
    response: oneshot::Sender<Option<CertifiedBlock>>,
) -> eyre::Result<()> {
    // TODO: right now, the response channel would just drop and an error
    // emitted here. Should this failure be propagated upstream?
    let finalization = client
        .get_finalization(Query::Height(height.get()))
        .await
        .wrap_err("failed getting finalization")?;
    response
        .send(Some(finalization))
        .map_err(|_| eyre::eyre!("receiver went away"))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn reconnect_backoff_linearly_increases_and_caps() {
        assert_eq!(reconnect_backoff(0), Duration::from_secs(0));
        assert_eq!(reconnect_backoff(1), Duration::from_secs(2));
        assert_eq!(reconnect_backoff(2), Duration::from_secs(4));
        assert_eq!(reconnect_backoff(3), Duration::from_secs(6));
        assert_eq!(reconnect_backoff(4), Duration::from_secs(8));
        assert_eq!(reconnect_backoff(5), Duration::from_secs(10));
        assert_eq!(reconnect_backoff(10), RECONNECT_MAX_BACKOFF);
        assert_eq!(reconnect_backoff(u64::MAX), RECONNECT_MAX_BACKOFF);
    }
}