linera_rpc/simple/

server.rs

// Copyright (c) Zefchain Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::{pin::Pin, sync::Arc};

use async_trait::async_trait;
use futures::{channel::mpsc, lock::Mutex, Stream, StreamExt as _};
use linera_base::{data_types::Blob, identifiers::ChainId, time::Duration};
use linera_core::{
    data_types::CrossChainRequest,
    node::NodeError,
    worker::{NetworkActions, Notification, WorkerError, WorkerState},
    JoinSetExt as _, ProcessConfirmedBlockMode,
};
use linera_storage::Storage;
use tokio::{sync, sync::oneshot, task::JoinSet};
use tokio_stream::wrappers::BroadcastStream;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, instrument};

use super::transport::{MessageHandler, ServerHandle, TransportProtocol};
use crate::{
    config::{CrossChainConfig, ShardId, ValidatorInternalNetworkPreConfig},
    cross_chain_message_queue, RpcMessage,
};

/// A server handling RPC requests over a simple (UDP or TCP) transport.
#[derive(Clone)]
pub struct Server<S>
where
    S: Storage,
{
    network: ValidatorInternalNetworkPreConfig<TransportProtocol>,
    host: String,
    port: u16,
    state: WorkerState<S>,
    shard_id: ShardId,
    cross_chain_config: CrossChainConfig,
    // Stats
    packets_processed: u64,
    user_errors: u64,
}

impl<S> Server<S>
where
    S: Storage,
{
    /// Creates a new server with the given network configuration and worker state.
    pub fn new(
        network: ValidatorInternalNetworkPreConfig<TransportProtocol>,
        host: String,
        port: u16,
        state: WorkerState<S>,
        shard_id: ShardId,
        cross_chain_config: CrossChainConfig,
    ) -> Self {
        Self {
            network,
            host,
            port,
            state,
            shard_id,
            cross_chain_config,
            packets_processed: 0,
            user_errors: 0,
        }
    }

    /// Returns the number of packets processed so far.
    pub fn packets_processed(&self) -> u64 {
        self.packets_processed
    }

    /// Returns the number of user errors encountered so far.
    pub fn user_errors(&self) -> u64 {
        self.user_errors
    }
}

impl<S> Server<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    #[expect(clippy::too_many_arguments)]
    async fn forward_cross_chain_queries(
        nickname: String,
        network: ValidatorInternalNetworkPreConfig<TransportProtocol>,
        cross_chain_max_retries: u32,
        cross_chain_retry_delay: Duration,
        cross_chain_max_backoff: Duration,
        cross_chain_sender_delay: Duration,
        cross_chain_sender_failure_rate: f32,
        this_shard: ShardId,
        receiver: mpsc::Receiver<(CrossChainRequest, ShardId)>,
    ) {
        let pool = Arc::new(Mutex::new(
            network
                .protocol
                .make_outgoing_connection_pool()
                .await
                .expect("Initialization should not fail"),
        ));
        let handle_request = move |shard_id, request| {
            let pool = pool.clone();
            let shard = network.shard(shard_id);
            let remote_address = format!("{}:{}", shard.host, shard.port);
            let message = RpcMessage::CrossChainRequest(Box::new(request));
            async move {
                pool.lock()
                    .await
                    .send_message_to(message.clone(), &remote_address)
                    .await?;
                anyhow::Result::<_, anyhow::Error>::Ok(())
            }
        };
        cross_chain_message_queue::forward_cross_chain_queries(
            nickname,
            cross_chain_max_retries,
            cross_chain_retry_delay,
            cross_chain_max_backoff,
            cross_chain_sender_delay,
            cross_chain_sender_failure_rate,
            this_shard,
            receiver,
            handle_request,
        )
        .await;
    }

    /// Spawns the server, returning a handle to track its completion.
    pub fn spawn(
        mut self,
        shutdown_signal: CancellationToken,
        join_set: &mut JoinSet<()>,
    ) -> ServerHandle {
        info!(
            "Listening to {:?} traffic on {}:{}",
            self.network.protocol, self.host, self.port
        );
        let address = (self.host.clone(), self.port);

        let (cross_chain_sender, cross_chain_receiver) =
            mpsc::channel(self.cross_chain_config.queue_size);

        let (notification_sender, _) = sync::broadcast::channel(1000);

        // Give the worker a shard-routing sender for cross-chain requests generated
        // outside the normal `NetworkActions` return path (specifically, the
        // `RevertConfirm`s emitted after resetting a corrupted chain).
        {
            let routing_network = self.network.clone();
            let routing_sender = cross_chain_sender.clone();
            self.state = self
                .state
                .clone()
                .with_outbound_cross_chain_sender(Arc::new(move |request| {
                    let shard_id = routing_network.get_shard_id(request.target_chain_id());
                    if let Err(error) = routing_sender.clone().try_send((request, shard_id)) {
                        tracing::error!(%error, "dropping cross-chain request");
                    }
                }));
        }

        join_set.spawn_task(Self::forward_cross_chain_queries(
            self.state.nickname().to_string(),
            self.network.clone(),
            self.cross_chain_config.max_retries,
            Duration::from_millis(self.cross_chain_config.retry_delay_ms),
            Duration::from_millis(self.cross_chain_config.max_backoff_ms),
            Duration::from_millis(self.cross_chain_config.sender_delay_ms),
            self.cross_chain_config.sender_failure_rate,
            self.shard_id,
            cross_chain_receiver,
        ));

        let protocol = self.network.protocol;
        let state = RunningServerState {
            server: self,
            cross_chain_sender,
            notification_sender,
        };
        // Launch server for the appropriate protocol.
        protocol.spawn_server(address, state, shutdown_signal, join_set)
    }
}

#[derive(Clone)]
struct RunningServerState<S>
where
    S: Storage,
{
    server: Server<S>,
    cross_chain_sender: mpsc::Sender<(CrossChainRequest, ShardId)>,
    notification_sender: sync::broadcast::Sender<Notification>,
}

#[async_trait]
impl<S> MessageHandler for RunningServerState<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    #[instrument(
        target = "simple_server",
        skip_all,
        fields(
            nickname = self.server.state.nickname(),
            chain_id = ?message.target_chain_id()
        )
    )]
    async fn handle_message(&mut self, message: RpcMessage) -> Option<RpcMessage> {
        let reply = match message {
            RpcMessage::BlockProposal(message) => {
                let (result, actions) = self.server.state.handle_block_proposal(*message).await;
                // Dispatch actions whether or not the proposal was accepted: a
                // rejected proposal can still advance the manager's `current_round`
                // (via `update_signed_proposal` on the `HasIncompatibleConfirmedVote`
                // recovery path), and subscribers need the resulting `NewRound`
                // notification.
                self.handle_network_actions(actions);
                match result {
                    Ok(info) => Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info)))),
                    Err(error) => {
                        self.log_error(&error, "Failed to handle block proposal");
                        Err(error.into())
                    }
                }
            }
            RpcMessage::LiteCertificate(request) => {
                let (sender, receiver) = request
                    .wait_for_outgoing_messages
                    .then(oneshot::channel)
                    .unzip();
                match Box::pin(
                    self.server
                        .state
                        .handle_lite_certificate(request.certificate, sender),
                )
                .await
                {
                    Ok((info, actions)) => {
                        // Cross-shard requests
                        self.handle_network_actions(actions);
                        if let Some(receiver) = receiver {
                            if let Err(e) = receiver.await {
                                error!("Failed to wait for message delivery: {e}");
                            }
                        }
                        // Response
                        Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info))))
                    }
                    Err(error) => {
                        let nickname = self.server.state.nickname();
                        if let WorkerError::MissingCertificateValue = &error {
                            debug!(nickname, %error, "Failed to handle lite certificate");
                        } else {
                            error!(nickname, %error, "Failed to handle lite certificate");
                        }
                        Err(error.into())
                    }
                }
            }
            RpcMessage::TimeoutCertificate(request) => {
                match self
                    .server
                    .state
                    .handle_timeout_certificate(request.certificate)
                    .await
                {
                    Ok((info, actions)) => {
                        // Cross-shard requests
                        self.handle_network_actions(actions);
                        // Response
                        Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info))))
                    }
                    Err(error) => {
                        self.log_error(&error, "Failed to handle timeout certificate");
                        Err(error.into())
                    }
                }
            }
            RpcMessage::ValidatedCertificate(request) => {
                match self
                    .server
                    .state
                    .handle_validated_certificate(request.certificate)
                    .await
                {
                    Ok((info, actions)) => {
                        // Cross-shard requests
                        self.handle_network_actions(actions);
                        // Response
                        Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info))))
                    }
                    Err(error) => {
                        self.log_error(&error, "Failed to handle validated certificate");
                        Err(error.into())
                    }
                }
            }
            RpcMessage::ConfirmedCertificate(request) => {
                let (sender, receiver) = request
                    .wait_for_outgoing_messages
                    .then(oneshot::channel)
                    .unzip();
                match self
                    .server
                    .state
                    .handle_confirmed_certificate(
                        request.certificate,
                        ProcessConfirmedBlockMode::Auto,
                        sender,
                    )
                    .await
                {
                    Ok((info, actions)) => {
                        // Cross-shard requests
                        self.handle_network_actions(actions);
                        if let Some(receiver) = receiver {
                            if let Err(e) = receiver.await {
                                error!("Failed to wait for message delivery: {e}");
                            }
                        }
                        // Response
                        Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info))))
                    }
                    Err(error) => {
                        self.log_error(&error, "Failed to handle confirmed certificate");
                        Err(error.into())
                    }
                }
            }
            RpcMessage::ChainInfoQuery(message) => {
                match self.server.state.handle_chain_info_query(*message).await {
                    Ok(info) => Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info)))),
                    Err(error) => {
                        self.log_error(&error, "Failed to handle chain info query");
                        Err(error.into())
                    }
                }
            }
            RpcMessage::CrossChainRequest(request) => {
                match self.server.state.handle_cross_chain_request(*request).await {
                    Ok(actions) => {
                        self.handle_network_actions(actions);
                    }
                    Err(error) => {
                        self.log_error(&error, "Failed to handle cross-chain request");
                    }
                }
                // No user to respond to.
                Ok(None)
            }
            RpcMessage::DownloadPendingBlob(request) => {
                let (chain_id, blob_id) = *request;
                match self
                    .server
                    .state
                    .download_pending_blob(chain_id, blob_id)
                    .await
                {
                    Ok(blob) => Ok(Some(RpcMessage::DownloadPendingBlobResponse(Box::new(
                        blob.content().clone(),
                    )))),
                    Err(error) => {
                        self.log_error(&error, "Failed to handle pending blob request");
                        Err(error.into())
                    }
                }
            }
            RpcMessage::HandlePendingBlob(request) => {
                let (chain_id, blob_content) = *request;
                match self
                    .server
                    .state
                    .handle_pending_blob(chain_id, Blob::new(blob_content))
                    .await
                {
                    Ok(info) => Ok(Some(RpcMessage::ChainInfoResponse(Box::new(info)))),
                    Err(error) => {
                        self.log_error(&error, "Failed to handle pending blob");
                        Err(error.into())
                    }
                }
            }

            RpcMessage::VersionInfoQuery => {
                Ok(Some(RpcMessage::VersionInfoResponse(Box::default())))
            }

            RpcMessage::SubscribeNotifications(_) | RpcMessage::Notification(_) => {
                // Subscriptions are handled at the transport level, not here.
                Err(NodeError::UnexpectedMessage)
            }

            RpcMessage::Vote(_)
            | RpcMessage::Error(_)
            | RpcMessage::ChainInfoResponse(_)
            | RpcMessage::VersionInfoResponse(_)
            | RpcMessage::NetworkDescriptionQuery
            | RpcMessage::NetworkDescriptionResponse(_)
            | RpcMessage::ShardInfoQuery(_)
            | RpcMessage::ShardInfoResponse(_)
            | RpcMessage::DownloadBlob(_)
            | RpcMessage::DownloadBlobs(_)
            | RpcMessage::DownloadBlobResponse(_)
            | RpcMessage::DownloadPendingBlobResponse(_)
            | RpcMessage::DownloadConfirmedBlock(_)
            | RpcMessage::DownloadConfirmedBlockResponse(_)
            | RpcMessage::BlobLastUsedBy(_)
            | RpcMessage::BlobLastUsedByResponse(_)
            | RpcMessage::BlobLastUsedByCertificate(_)
            | RpcMessage::BlobLastUsedByCertificateResponse(_)
            | RpcMessage::MissingBlobIds(_)
            | RpcMessage::MissingBlobIdsResponse(_)
            | RpcMessage::EventBlockHeights(_)
            | RpcMessage::EventBlockHeightsResponse(_)
            | RpcMessage::DownloadCertificates(_)
            | RpcMessage::DownloadCertificatesResponse(_)
            | RpcMessage::UploadBlob(_)
            | RpcMessage::UploadBlobResponse(_)
            | RpcMessage::DownloadCertificatesByHeights(_, _)
            | RpcMessage::DownloadCertificatesByHeightsResponse(_) => {
                Err(NodeError::UnexpectedMessage)
            }
        };

        self.server.packets_processed += 1;
        // We allow this because `is_multiple_of` is still unstable in our MSRV.
        #[allow(unknown_lints)]
        #[expect(clippy::manual_is_multiple_of)]
        if self.server.packets_processed % 5000 == 0 {
            debug!(
                "[{}] {}:{} (shard {}) has processed {} packets",
                self.server.state.nickname(),
                self.server.host,
                self.server.port,
                self.server.shard_id,
                self.server.packets_processed
            );
        }

        match reply {
            Ok(x) => x,
            Err(error) => {
                // TODO(#459): Make it a warning or an error again.
                debug!(
                    "[{}] User query failed: {}",
                    self.server.state.nickname(),
                    error
                );
                self.server.user_errors += 1;
                Some(error.into())
            }
        }
    }

    async fn handle_subscribe(
        &mut self,
        chains: Vec<ChainId>,
    ) -> Option<Pin<Box<dyn Stream<Item = RpcMessage> + Send>>> {
        RunningServerState::subscribe_to_notifications(self, chains).await
    }
}

impl<S> RunningServerState<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    async fn subscribe_to_notifications(
        &self,
        chains: Vec<ChainId>,
    ) -> Option<Pin<Box<dyn Stream<Item = RpcMessage> + Send>>> {
        let receiver = self.notification_sender.subscribe();
        let stream = BroadcastStream::new(receiver).filter_map(move |result| {
            let chains = chains.clone();
            async move {
                match result {
                    Ok(notification) if chains.contains(&notification.chain_id) => {
                        Some(RpcMessage::Notification(Box::new(notification)))
                    }
                    _ => None,
                }
            }
        });
        Some(Box::pin(stream))
    }
}

impl<S> RunningServerState<S>
where
    S: Storage + Send,
{
    fn handle_network_actions(&mut self, actions: NetworkActions) {
        for request in actions.cross_chain_requests {
            let shard_id = self.server.network.get_shard_id(request.target_chain_id());
            debug!(
                "[{}] Scheduling cross-chain query: {} -> {}",
                self.server.state.nickname(),
                self.server.shard_id,
                shard_id
            );
            if let Err(error) = self.cross_chain_sender.try_send((request, shard_id)) {
                error!(%error, "dropping cross-chain request");
                break;
            }
        }
        for notification in actions.notifications {
            debug!("Scheduling notification query");
            if let Err(error) = self.notification_sender.send(notification) {
                debug!(%error, "dropping notification (no receivers)");
            }
        }
    }

    fn log_error(&self, error: &WorkerError, context: &str) {
        let nickname = self.server.state.nickname();
        if error.is_local() {
            error!(nickname, %error, "{}", context);
        } else {
            debug!(nickname, %error, "{}", context);
        }
    }
}