linera_rpc/grpc/

client.rs

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

use std::{
    collections::BTreeSet,
    fmt,
    future::Future,
    iter,
    sync::{
        atomic::{AtomicU32, Ordering},
        Arc,
    },
};

use futures::{future, stream, StreamExt};
use linera_base::{
    crypto::CryptoHash,
    data_types::{BlobContent, BlockHeight, NetworkDescription},
    ensure,
    identifiers::{BlobId, ChainId, EventId},
    time::{Duration, Instant},
};
use linera_chain::{
    data_types::{self},
    types::{
        self, Certificate, ConfirmedBlock, ConfirmedBlockCertificate, GenericCertificate,
        LiteCertificate, Timeout, ValidatedBlock,
    },
};
#[cfg(with_metrics)]
mod metrics {
    use std::sync::LazyLock;

    use linera_base::prometheus_util::register_int_counter_vec;
    use prometheus::IntCounterVec;

    pub static VALIDATOR_SUBSCRIPTION_ERRORS: LazyLock<IntCounterVec> = LazyLock::new(|| {
        register_int_counter_vec(
            "validator_subscription_errors",
            "Number of notification subscription stream errors per validator",
            &["address"],
        )
    });
}

use linera_core::{
    data_types::{CertificatesByHeightRequest, ChainInfoResponse},
    node::{BlobStream, CrossChainMessageDelivery, NodeError, NotificationStream, ValidatorNode},
    worker::Notification,
};
use linera_storage::Arc as CacheArc;
use linera_version::VersionInfo;
use tonic::{Code, IntoRequest, Request, Status};
use tracing::{debug, instrument, trace, Level};

use super::{
    api::{self, validator_node_client::ValidatorNodeClient, SubscriptionRequest},
    transport, GRPC_MAX_MESSAGE_SIZE,
};
#[cfg(feature = "opentelemetry")]
use crate::propagation::{get_context_with_traffic_type, inject_context};
use crate::{
    grpc::api::RawCertificate, HandleConfirmedCertificateRequest, HandleLiteCertRequest,
    HandleTimeoutCertificateRequest, HandleValidatedCertificateRequest,
};

#[derive(Clone)]
pub struct GrpcClient {
    address: String,
    client: ValidatorNodeClient<transport::Channel>,
    retry_delay: Duration,
    max_retries: u32,
    max_backoff: Duration,
    /// Shared across all `GrpcClient` instances created by the same `GrpcNodeProvider`.
    /// Tracks when each validator address last had a subscription failure, so that
    /// other chains don't independently retry the same dead validator.
    subscription_cooldowns: Arc<papaya::HashMap<String, Instant>>,
}

impl GrpcClient {
    pub fn new(
        address: String,
        channel: transport::Channel,
        retry_delay: Duration,
        max_retries: u32,
        max_backoff: Duration,
        subscription_cooldowns: Arc<papaya::HashMap<String, Instant>>,
    ) -> Self {
        let client = ValidatorNodeClient::new(channel)
            .max_encoding_message_size(GRPC_MAX_MESSAGE_SIZE)
            .max_decoding_message_size(GRPC_MAX_MESSAGE_SIZE);
        Self {
            address,
            client,
            retry_delay,
            max_retries,
            max_backoff,
            subscription_cooldowns,
        }
    }

    pub fn address(&self) -> &str {
        &self.address
    }

    /// Returns whether this gRPC status means the server stream should be reconnected to, or not.
    /// Logs a warning on unexpected status codes.
    fn is_retryable(status: &Status) -> bool {
        match status.code() {
            Code::DeadlineExceeded | Code::Aborted | Code::Unavailable | Code::Unknown => {
                trace!("gRPC request interrupted: {status:?}; retrying");
                true
            }
            Code::Ok | Code::Cancelled | Code::ResourceExhausted => {
                trace!("Unexpected gRPC status: {status:?}; retrying");
                true
            }
            Code::Internal if status.message().contains("h2 protocol error") => {
                // HTTP/2 connection reset errors are transient network issues, not real
                // internal errors. This happens when the server restarts and the
                // connection is forcibly closed.
                trace!("gRPC connection reset: {status:?}; retrying");
                true
            }
            Code::Internal if status.message().contains("502 Bad Gateway") => {
                // When a proxy/ingress returns HTTP 502 (e.g. during rolling restarts),
                // tonic's frame decoder fails on the non-gRPC response body before the
                // HTTP-to-gRPC status mapping can run, producing Code::Internal instead
                // of Code::Unavailable. Per the gRPC spec, HTTP 502 maps to UNAVAILABLE
                // which is retryable. This works around tonic#2365.
                trace!("gRPC proxy error (502): {status:?}; retrying");
                true
            }
            Code::NotFound => false, // This code is used if e.g. the validator is missing blobs.
            Code::InvalidArgument
            | Code::AlreadyExists
            | Code::PermissionDenied
            | Code::FailedPrecondition
            | Code::OutOfRange
            | Code::Unimplemented
            | Code::Internal
            | Code::DataLoss
            | Code::Unauthenticated => {
                trace!("Unexpected gRPC status: {status:?}");
                false
            }
        }
    }

    async fn delegate<F, Fut, R, S>(
        &self,
        f: F,
        request: impl TryInto<R> + fmt::Debug + Clone,
        handler: &str,
    ) -> Result<S, NodeError>
    where
        F: Fn(ValidatorNodeClient<transport::Channel>, Request<R>) -> Fut,
        Fut: Future<Output = Result<tonic::Response<S>, Status>>,
        R: IntoRequest<R> + Clone,
    {
        let mut retry_count = 0;
        let request_inner = request.try_into().map_err(|_| NodeError::GrpcError {
            error: "could not convert request to proto".to_string(),
        })?;
        loop {
            #[allow(unused_mut)]
            let mut request = Request::new(request_inner.clone());
            // Inject OpenTelemetry context (trace context + baggage) into gRPC metadata.
            // This uses get_context_with_traffic_type() to also check the LINERA_TRAFFIC_TYPE
            // environment variable, allowing benchmark tools to mark their traffic as synthetic.
            #[cfg(feature = "opentelemetry")]
            inject_context(&get_context_with_traffic_type(), request.metadata_mut());
            match f(self.client.clone(), request).await {
                Err(s) if Self::is_retryable(&s) && retry_count < self.max_retries => {
                    let delay = crate::jittered_backoff_delay(
                        self.retry_delay,
                        retry_count,
                        self.max_backoff,
                    );
                    retry_count += 1;
                    linera_base::time::timer::sleep(delay).await;
                    continue;
                }
                Err(s) => {
                    return Err(NodeError::GrpcError {
                        error: format!("remote request [{handler}] failed with status: {s:?}"),
                    });
                }
                Ok(result) => return Ok(result.into_inner()),
            };
        }
    }

    fn try_into_chain_info(
        result: api::ChainInfoResult,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        let inner = result.inner.ok_or_else(|| NodeError::GrpcError {
            error: "missing body from response".to_string(),
        })?;
        match inner {
            api::chain_info_result::Inner::ChainInfoResponse(response) => {
                Ok(response.try_into().map_err(|err| NodeError::GrpcError {
                    error: format!("failed to unmarshal response: {err}"),
                })?)
            }
            api::chain_info_result::Inner::Error(error) => Err(bincode::deserialize(&error)
                .map_err(|err| NodeError::GrpcError {
                    error: format!("failed to unmarshal error message: {err}"),
                })?),
        }
    }
}

impl TryFrom<api::PendingBlobResult> for BlobContent {
    type Error = NodeError;

    fn try_from(result: api::PendingBlobResult) -> Result<Self, Self::Error> {
        let inner = result.inner.ok_or_else(|| NodeError::GrpcError {
            error: "missing body from response".to_string(),
        })?;
        match inner {
            api::pending_blob_result::Inner::Blob(blob) => {
                Ok(blob.try_into().map_err(|err| NodeError::GrpcError {
                    error: format!("failed to unmarshal response: {err}"),
                })?)
            }
            api::pending_blob_result::Inner::Error(error) => Err(bincode::deserialize(&error)
                .map_err(|err| NodeError::GrpcError {
                    error: format!("failed to unmarshal error message: {err}"),
                })?),
        }
    }
}

macro_rules! client_delegate {
    ($self:ident, $handler:ident, $req:ident) => {{
        debug!(
            handler = stringify!($handler),
            request = ?$req,
            "sending gRPC request"
        );
        $self
            .delegate(
                |mut client, req| async move { client.$handler(req).await },
                $req,
                stringify!($handler),
            )
            .await
    }};
}

impl ValidatorNode for GrpcClient {
    type NotificationStream = NotificationStream;

    fn address(&self) -> String {
        self.address.clone()
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn handle_block_proposal(
        &self,
        proposal: data_types::BlockProposal,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        GrpcClient::try_into_chain_info(client_delegate!(self, handle_block_proposal, proposal)?)
    }

    #[instrument(target = "grpc_client", skip_all, fields(address = self.address))]
    async fn handle_lite_certificate(
        &self,
        certificate: types::LiteCertificate<'_>,
        delivery: CrossChainMessageDelivery,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        let wait_for_outgoing_messages = delivery.wait_for_outgoing_messages();
        let request = HandleLiteCertRequest {
            certificate,
            wait_for_outgoing_messages,
        };
        GrpcClient::try_into_chain_info(client_delegate!(self, handle_lite_certificate, request)?)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn handle_confirmed_certificate(
        &self,
        certificate: CacheArc<GenericCertificate<ConfirmedBlock>>,
        delivery: CrossChainMessageDelivery,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        let wait_for_outgoing_messages: bool = delivery.wait_for_outgoing_messages();
        let request = HandleConfirmedCertificateRequest {
            certificate: CacheArc::unwrap_or_clone(certificate),
            wait_for_outgoing_messages,
        };
        GrpcClient::try_into_chain_info(client_delegate!(
            self,
            handle_confirmed_certificate,
            request
        )?)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn handle_validated_certificate(
        &self,
        certificate: GenericCertificate<ValidatedBlock>,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        let request = HandleValidatedCertificateRequest { certificate };
        GrpcClient::try_into_chain_info(client_delegate!(
            self,
            handle_validated_certificate,
            request
        )?)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn handle_timeout_certificate(
        &self,
        certificate: GenericCertificate<Timeout>,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        let request = HandleTimeoutCertificateRequest { certificate };
        GrpcClient::try_into_chain_info(client_delegate!(
            self,
            handle_timeout_certificate,
            request
        )?)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn handle_chain_info_query(
        &self,
        query: linera_core::data_types::ChainInfoQuery,
    ) -> Result<linera_core::data_types::ChainInfoResponse, NodeError> {
        GrpcClient::try_into_chain_info(client_delegate!(self, handle_chain_info_query, query)?)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn subscribe(&self, chains: Vec<ChainId>) -> Result<Self::NotificationStream, NodeError> {
        let retry_delay = self.retry_delay;
        let max_retries = self.max_retries;
        let max_backoff = self.max_backoff;
        let address = self.address.clone();
        let subscription_cooldowns = self.subscription_cooldowns.clone();

        // Fast-fail if another subscription to this address recently failed.
        // Prevents N chains from independently retrying the same dead validator.
        {
            let pinned = subscription_cooldowns.pin();
            if let Some(&last_failure) = pinned.get(&address) {
                if last_failure.elapsed() < max_backoff {
                    return Err(NodeError::SubscriptionFailed {
                        status: format!(
                            "validator {address} on cooldown after recent subscription failure"
                        ),
                    });
                }
            }
        }

        // Use shared atomic counter so unfold can reset it on successful reconnection.
        let retry_count = Arc::new(AtomicU32::new(0));
        let subscription_request = SubscriptionRequest {
            chain_ids: chains.into_iter().map(|chain| chain.into()).collect(),
        };
        let mut client = self.client.clone();

        // Make the first connection attempt before returning from this method.
        let mut stream = Some(
            client
                .subscribe(subscription_request.clone())
                .await
                .map_err(|status| {
                    subscription_cooldowns
                        .pin()
                        .insert(address.clone(), Instant::now());
                    NodeError::SubscriptionFailed {
                        status: status.to_string(),
                    }
                })?
                .into_inner(),
        );

        // A stream of `Result<grpc::Notification, tonic::Status>` that keeps calling
        // `client.subscribe(request)` endlessly and without delay.
        let retry_count_for_unfold = retry_count.clone();
        let cooldowns_for_unfold = subscription_cooldowns.clone();
        let address_for_unfold = address.clone();
        let endlessly_retrying_notification_stream = stream::unfold((), move |()| {
            let mut client = client.clone();
            let subscription_request = subscription_request.clone();
            let mut stream = stream.take();
            let retry_count = retry_count_for_unfold.clone();
            let cooldowns = cooldowns_for_unfold.clone();
            let cooldown_address = address_for_unfold.clone();
            async move {
                let stream = if let Some(stream) = stream.take() {
                    future::Either::Right(stream)
                } else {
                    match client.subscribe(subscription_request.clone()).await {
                        Err(err) => future::Either::Left(stream::iter(iter::once(Err(err)))),
                        Ok(response) => {
                            // Reset retry count on successful reconnection.
                            retry_count.store(0, Ordering::Relaxed);
                            cooldowns.pin().remove(&cooldown_address);
                            trace!("Successfully reconnected subscription stream");
                            future::Either::Right(response.into_inner())
                        }
                    }
                };
                Some((stream, ()))
            }
        })
        .flatten();

        let span = tracing::info_span!("notification stream");
        #[cfg(with_metrics)]
        let address_for_metrics = self.address.clone();
        let cooldowns_for_take_while = subscription_cooldowns;
        let address_for_take_while = self.address.clone();
        // The stream of `Notification`s that inserts increasing delays after retriable errors, and
        // terminates after unexpected or fatal errors.
        let notification_stream = endlessly_retrying_notification_stream
            .map(|result| {
                Option::<Notification>::try_from(result?).map_err(|err| {
                    let message = format!("Could not deserialize notification: {err}");
                    tonic::Status::new(Code::Internal, message)
                })
            })
            .take_while(move |result| {
                let Err(status) = result else {
                    retry_count.store(0, Ordering::Relaxed);
                    return future::Either::Left(future::ready(true));
                };

                #[cfg(with_metrics)]
                metrics::VALIDATOR_SUBSCRIPTION_ERRORS
                    .with_label_values(&[&address_for_metrics])
                    .inc();

                let current_retry_count = retry_count.load(Ordering::Relaxed);
                if !span.in_scope(|| Self::is_retryable(status))
                    || current_retry_count >= max_retries
                {
                    cooldowns_for_take_while
                        .pin()
                        .insert(address_for_take_while.clone(), Instant::now());
                    return future::Either::Left(future::ready(false));
                }
                let delay =
                    crate::jittered_backoff_delay(retry_delay, current_retry_count, max_backoff);
                retry_count.fetch_add(1, Ordering::Relaxed);
                future::Either::Right(async move {
                    linera_base::time::timer::sleep(delay).await;
                    true
                })
            })
            .filter_map(move |result| {
                future::ready(match result {
                    Ok(notification @ Some(_)) => notification,
                    Ok(None) => None,
                    Err(err) => {
                        debug!(%address, "{}", err);
                        None
                    }
                })
            });

        Ok(Box::pin(notification_stream))
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn get_version_info(&self) -> Result<VersionInfo, NodeError> {
        let req = ();
        Ok(client_delegate!(self, get_version_info, req)?.into())
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn get_network_description(&self) -> Result<NetworkDescription, NodeError> {
        let req = ();
        Ok(client_delegate!(self, get_network_description, req)?.try_into()?)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn upload_blob(&self, content: BlobContent) -> Result<BlobId, NodeError> {
        Ok(client_delegate!(self, upload_blob, content)?.try_into()?)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn download_blob(&self, blob_id: BlobId) -> Result<BlobContent, NodeError> {
        Ok(client_delegate!(self, download_blob, blob_id)?.try_into()?)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn download_blobs(&self, blob_ids: Vec<BlobId>) -> Result<BlobStream, NodeError> {
        debug!(
            handler = "download_blobs",
            num_blobs = blob_ids.len(),
            "sending gRPC request"
        );
        let request = api::BlobIds::try_from(blob_ids)?;
        let stream = self
            .client
            .clone()
            .download_blobs(request)
            .await
            .map_err(|status| NodeError::GrpcError {
                error: status.to_string(),
            })?
            .into_inner();
        let blob_stream = stream.map(|result| match result {
            Ok(proto_blob) => BlobContent::try_from(proto_blob).map_err(NodeError::from),
            Err(status) => Err(NodeError::GrpcError {
                error: status.to_string(),
            }),
        });
        Ok(Box::pin(blob_stream))
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn download_pending_blob(
        &self,
        chain_id: ChainId,
        blob_id: BlobId,
    ) -> Result<BlobContent, NodeError> {
        let req = (chain_id, blob_id);
        client_delegate!(self, download_pending_blob, req)?.try_into()
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn handle_pending_blob(
        &self,
        chain_id: ChainId,
        blob: BlobContent,
    ) -> Result<ChainInfoResponse, NodeError> {
        let req = (chain_id, blob);
        GrpcClient::try_into_chain_info(client_delegate!(self, handle_pending_blob, req)?)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn download_certificate(
        &self,
        hash: CryptoHash,
    ) -> Result<ConfirmedBlockCertificate, NodeError> {
        ConfirmedBlockCertificate::try_from(Certificate::try_from(client_delegate!(
            self,
            download_certificate,
            hash
        )?)?)
        .map_err(|_| NodeError::UnexpectedCertificateValue)
    }

    #[instrument(target = "grpc_client", skip_all, err(level = Level::DEBUG), fields(address = self.address))]
    async fn download_certificates(
        &self,
        hashes: Vec<CryptoHash>,
    ) -> Result<Vec<ConfirmedBlockCertificate>, NodeError> {
        let mut missing_hashes = hashes;
        let mut certs_collected = Vec::with_capacity(missing_hashes.len());
        while !missing_hashes.is_empty() {
            // Macro doesn't compile if we pass `missing_hashes.clone()` directly to `client_delegate!`.
            let missing = missing_hashes.clone();
            let mut received: Vec<_> = Vec::<Certificate>::try_from(client_delegate!(
                self,
                download_certificates,
                missing
            )?)?
            .into_iter()
            .map(|cert| {
                ConfirmedBlockCertificate::try_from(cert)
                    .map_err(|_| NodeError::UnexpectedCertificateValue)
            })
            .collect::<Result<_, _>>()?;

            // In the case of the server not returning any certificates, we break the loop.
            if received.is_empty() {
                break;
            }

            // Honest validator should return certificates in the same order as the requested hashes.
            missing_hashes = missing_hashes[received.len()..].to_vec();
            certs_collected.append(&mut received);
        }
        ensure!(
            missing_hashes.is_empty(),
            NodeError::MissingCertificates(missing_hashes)
        );
        Ok(certs_collected)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn download_certificates_by_heights(
        &self,
        chain_id: ChainId,
        heights: Vec<BlockHeight>,
    ) -> Result<Vec<ConfirmedBlockCertificate>, NodeError> {
        let mut missing = heights.into_iter().collect::<BTreeSet<_>>();
        let mut certs_collected = vec![];
        while !missing.is_empty() {
            let request = CertificatesByHeightRequest {
                chain_id,
                heights: missing.iter().copied().collect(),
            };
            let mut received: Vec<_> =
                client_delegate!(self, download_raw_certificates_by_heights, request)?
                    .certificates
                    .into_iter()
                    .map(
                        |RawCertificate {
                             lite_certificate,
                             confirmed_block,
                         }| {
                            let cert = bcs::from_bytes::<LiteCertificate>(&lite_certificate)
                                .map_err(|_| NodeError::UnexpectedCertificateValue)?;

                            let block = bcs::from_bytes::<ConfirmedBlock>(&confirmed_block)
                                .map_err(|_| NodeError::UnexpectedCertificateValue)?;

                            cert.with_value(block)
                                .ok_or(NodeError::UnexpectedCertificateValue)
                        },
                    )
                    .collect::<Result<_, _>>()?;

            if received.is_empty() {
                break;
            }

            // Remove only the heights we actually received from missing set.
            for cert in &received {
                missing.remove(&cert.inner().height());
            }
            certs_collected.append(&mut received);
        }
        certs_collected.sort_by_key(|cert| cert.inner().height());
        Ok(certs_collected)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn blob_last_used_by(&self, blob_id: BlobId) -> Result<CryptoHash, NodeError> {
        Ok(client_delegate!(self, blob_last_used_by, blob_id)?.try_into()?)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn blob_last_used_by_certificate(
        &self,
        blob_id: BlobId,
    ) -> Result<ConfirmedBlockCertificate, NodeError> {
        Ok(client_delegate!(self, blob_last_used_by_certificate, blob_id)?.try_into()?)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn event_block_heights(
        &self,
        event_ids: Vec<EventId>,
    ) -> Result<Vec<Option<BlockHeight>>, NodeError> {
        let request = api::EventBlockHeightsRequest::from(event_ids);
        Ok(client_delegate!(self, event_block_heights, request)?.try_into()?)
    }

    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn missing_blob_ids(&self, blob_ids: Vec<BlobId>) -> Result<Vec<BlobId>, NodeError> {
        Ok(client_delegate!(self, missing_blob_ids, blob_ids)?.try_into()?)
    }

    #[expect(
        clippy::cast_possible_truncation,
        reason = "shard counts are bounded by validator config and fit in usize on supported targets"
    )]
    #[instrument(target = "grpc_client", skip(self), err(level = Level::DEBUG), fields(address = self.address))]
    async fn get_shard_info(
        &self,
        chain_id: ChainId,
    ) -> Result<linera_core::data_types::ShardInfo, NodeError> {
        let response = client_delegate!(self, get_shard_info, chain_id)?;
        Ok(linera_core::data_types::ShardInfo {
            shard_id: response.shard_id as usize,
            total_shards: response.total_shards as usize,
        })
    }
}