miden-validator 0.17.0-rc.3

use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use miden_node_proto::generated as grpc;
use miden_node_proto::generated::validator::BlockSubscriptionResponse;
use miden_node_tracing::{ErrorReport, error, info, miden_instrument, miden_span_record};
use miden_protocol::block::{BlockNumber, SignedBlock};
use miden_protocol::utils::serde::Deserializable;
use tokio::sync::OwnedRwLockWriteGuard;
use tokio_stream::wrappers::ReceiverStream;
use tonic::Status;
use tonic::codegen::tokio_stream::Stream;

use super::ValidatorService;
use crate::COMPONENT;

type BlockStream =
    Pin<Box<dyn Stream<Item = tonic::Result<BlockSubscriptionResponse>> + Send + 'static>>;

struct BackupBlockStream {
    inner: ReceiverStream<tonic::Result<BlockSubscriptionResponse>>,
    _guard: OwnedRwLockWriteGuard<()>,
}

impl Stream for BackupBlockStream {
    type Item = tonic::Result<BlockSubscriptionResponse>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.get_mut().inner).poll_next(cx)
    }
}

#[tonic::async_trait]
impl grpc::server::validator_api::BlockSubscription for ValidatorService {
    type Input = grpc::validator::BlockSubscriptionRequest;
    type Item = BlockSubscriptionResponse;
    type ItemStream = BlockStream;

    fn decode(request: grpc::validator::BlockSubscriptionRequest) -> tonic::Result<Self::Input> {
        Ok(request)
    }

    fn encode(item: Self::Item) -> tonic::Result<Self::Item> {
        Ok(item)
    }

    #[miden_instrument(
        target = COMPONENT,
        name = "validator.block_subscription",
        err,
    )]
    async fn handle(
        &self,
        request: Self::Input,
        _metadata: &tonic::metadata::MetadataMap,
        _extensions: &tonic::codegen::http::Extensions,
    ) -> tonic::Result<Self::ItemStream> {
        miden_span_record!(block.from = request.block_from);

        let committed_tip = *self.committed_tip.borrow();
        if request.block_from > committed_tip.as_u32() {
            return Err(Status::out_of_range(
                "subscriber's requested starting block should be <= the committed chain tip",
            ));
        }

        // Hold the exclusive backup lock for the entire lifetime of the stream. While a backup
        // subscription is active no other RPCs may run, and vice versa.
        let guard = Arc::clone(&self.serve_lock).try_write_owned().map_err(|_| {
            Status::resource_exhausted("cannot stream backup while validator is serving requests")
        })?;

        let from = BlockNumber::from(request.block_from);
        // The tip should never move since we are in recovery mode and therefore there is no active
        // sequencer.
        miden_span_record!(tip.number = committed_tip);

        let (tx, rx) = tokio::sync::mpsc::channel(32);

        tokio::spawn({
            let store = self.block_store.clone();
            let db = self.db.reader();
            async move {
                let mut previous_config_commitment = None;
                for block in from.as_u32()..=committed_tip.as_u32() {
                    let response = match store.load_block(block.into()).await {
                        Ok(Some(bytes)) => match SignedBlock::read_from_bytes(&bytes) {
                            Ok(signed_block) => {
                                let commitment = signed_block.header().protocol_config_commitment();
                                let protocol_config = if previous_config_commitment
                                    == Some(commitment)
                                {
                                    Ok(None)
                                } else {
                                    match db.load_protocol_config(commitment).await {
                                        Ok(Some(config)) => {
                                            previous_config_commitment = Some(commitment);
                                            Ok(Some((&config).into()))
                                        },
                                        Ok(None) => Err(tonic::Status::internal(format!(
                                            "protocol config {commitment} not found"
                                        ))),
                                        Err(err) => Err(tonic::Status::internal(
                                            err.as_report_context("failed to load protocol config"),
                                        )),
                                    }
                                };
                                protocol_config.map(|protocol_config| BlockSubscriptionResponse {
                                    block: Some(signed_block.into()),
                                    committed_chain_tip: committed_tip.as_u32(),
                                    protocol_config,
                                })
                            },
                            Err(err) => Err(tonic::Status::internal(
                                err.as_report_context("failed to decode backed-up block"),
                            )),
                        },
                        Ok(None) => {
                            Err(tonic::Status::not_found(format!("block {block} not found")))
                        },
                        Err(err) => Err(tonic::Status::internal(
                            err.as_report_context("failed to load block"),
                        )),
                    }
                    .inspect_err(|err| {
                        error!(
                            err,
                            "failed to load block in validator recovery stream",
                            block.number = block
                        );
                    });

                    // Errors are not recoverable so we abort the stream after informing the client.
                    //
                    // Also exit if the client closed the stream.
                    //
                    // Note that the condition ordering is deliberate; otherwise `is_err` would short-circuit
                    // and prevent the sending of the error response.
                    let is_err = response.is_err();
                    if tx.send(response).await.is_err() || is_err {
                        info!("validator recovery stream closing");
                        return;
                    }
                }
            }
        });

        Ok(Box::pin(BackupBlockStream {
            inner: ReceiverStream::new(rx),
            _guard: guard,
        }))
    }
}