zakura-rpc 8.0.0

//! Implements `Indexer` methods on the `IndexerRPC` type

use std::{collections::HashSet, pin::Pin, time::Duration};

use futures::Stream;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Response, Status};
use tower::util::ServiceExt;

use tracing::Span;
use zakura_chain::{block, chain_tip::ChainTip, serialization::BytesInDisplayOrder};
use zakura_node_services::mempool::MempoolChangeKind;
use zakura_state::{
    constants::MAX_NON_FINALIZED_CHAIN_FORKS, ReadRequest, ReadResponse, ReadState,
};

use super::{
    indexer_server::Indexer, server::IndexerRPC, BlockAndHash, BlockAndHeight, BlockHashAndHeight,
    BlockRangeRequest, BlockRequest, Empty, MempoolChangeMessage, NonFinalizedStateChangeRequest,
    BLOCK_HASH_BYTE_LEN, BLOCK_HEIGHT_BYTE_LEN,
};

/// The maximum number of messages that can be queued to be streamed to a client.
const RESPONSE_BUFFER_SIZE: usize = 64;

/// How long to wait for a backpressured send before treating the consumer as hung
/// and dropping the subscription.
///
/// All indexer streams apply backpressure so a slow consumer doesn't miss
/// notifications, but without a bound a consumer whose connection is half-open (dead
/// TCP not yet detected) would block the listener task indefinitely.
const SEND_TIMEOUT: Duration = Duration::from_secs(60);

/// The maximum number of blocks read from the state in a single request while
/// streaming a block range.
///
/// Reading [`RESPONSE_BUFFER_SIZE`] blocks per read means one state read fills
/// at most one response buffer, so a `get_block_range` stream holds at most two
/// buffers' worth of blocks while it waits for a slow consumer.
//
// The cast is safe because `RESPONSE_BUFFER_SIZE` is a small constant that
// fits in a `u32`.
#[cfg(feature = "indexer")]
const BLOCK_RANGE_CHUNK_SIZE: u32 = RESPONSE_BUFFER_SIZE as u32;

#[tonic::async_trait]
impl<ReadStateService, Tip> Indexer for IndexerRPC<ReadStateService, Tip>
where
    ReadStateService: ReadState,
    Tip: ChainTip + Clone + Send + Sync + 'static,
{
    type ChainTipChangeStream =
        Pin<Box<dyn Stream<Item = Result<BlockHashAndHeight, Status>> + Send>>;
    type NonFinalizedStateChangeStream =
        Pin<Box<dyn Stream<Item = Result<BlockAndHash, Status>> + Send>>;
    type MempoolChangeStream =
        Pin<Box<dyn Stream<Item = Result<MempoolChangeMessage, Status>> + Send>>;
    type GetBlockRangeStream = Pin<Box<dyn Stream<Item = Result<BlockAndHeight, Status>> + Send>>;

    async fn chain_tip_change(
        &self,
        _: tonic::Request<Empty>,
    ) -> Result<Response<Self::ChainTipChangeStream>, Status> {
        let span = Span::current();
        let (response_sender, response_receiver) = tokio::sync::mpsc::channel(RESPONSE_BUFFER_SIZE);
        let response_stream = ReceiverStream::new(response_receiver);
        let mut chain_tip_change = self.chain_tip_change.clone();

        tokio::spawn(async move {
            // Notify the client of chain tip changes until the channel is closed
            while let Ok(()) = chain_tip_change.best_tip_changed().await {
                let Some((tip_height, tip_hash)) = chain_tip_change.best_tip_height_and_hash()
                else {
                    continue;
                };

                let send = response_sender.send(Ok(BlockHashAndHeight::new(tip_hash, tip_height)));
                match tokio::time::timeout(SEND_TIMEOUT, send).await {
                    Ok(Ok(())) => {}
                    Ok(Err(_)) => {
                        span.in_scope(|| {
                            tracing::info!("client disconnected, dropping chain_tip_change task");
                        });
                        return;
                    }
                    Err(_) => {
                        span.in_scope(|| {
                            tracing::warn!(
                                "slow consumer, dropping chain_tip_change stream after \
                                 send timed out"
                            );
                        });
                        return;
                    }
                }
            }

            span.in_scope(|| {
                tracing::warn!("chain_tip_change channel has closed");
            });

            let _ = response_sender
                .send(Err(Status::unavailable(
                    "chain_tip_change channel has closed",
                )))
                .await;
        });

        Ok(Response::new(Box::pin(response_stream)))
    }

    async fn non_finalized_state_change(
        &self,
        request: tonic::Request<NonFinalizedStateChangeRequest>,
    ) -> Result<Response<Self::NonFinalizedStateChangeStream>, Status> {
        let span = Span::current();
        let read_state = self.read_state.clone();
        let (response_sender, response_receiver) = tokio::sync::mpsc::channel(RESPONSE_BUFFER_SIZE);
        let response_stream = ReceiverStream::new(response_receiver);

        // The caller may provide the hashes of the chain tips it already has so
        // the server only streams blocks after those tips. Malformed hashes are
        // rejected up front.
        let known_chain_tips = decode_known_chain_tips(request.into_inner().chain_tip_hashes)?;

        tokio::spawn(async move {
            let mut non_finalized_state_change = match read_state
                .oneshot(ReadRequest::NonFinalizedBlocksListener { known_chain_tips })
                .await
            {
                Ok(ReadResponse::NonFinalizedBlocksListener(listener)) => listener.unwrap(),
                Ok(_) => unreachable!("unexpected response type from ReadStateService"),
                Err(error) => {
                    span.in_scope(|| {
                        tracing::error!(
                            ?error,
                            "failed to subscribe to non-finalized state changes"
                        );
                    });

                    let _ = response_sender
                        .send(Err(Status::unavailable(
                            "failed to subscribe to non-finalized state changes",
                        )))
                        .await;
                    return;
                }
            };

            // Notify the client of new blocks until the channel is closed.
            //
            // This uses `send().await` to apply backpressure to the
            // non-finalized state listener rather than dropping blocks for a
            // slow consumer. A send error means the client disconnected; a
            // timed-out send means the consumer is hung. In both cases the task
            // ends rather than blocking forever.
            loop {
                // A full listener buffer means the state-side task is blocked
                // sending into it, so it may already have missed non-finalized
                // state updates. The stream can no longer guarantee
                // completeness, so drop the subscription instead of silently
                // missing blocks.
                if non_finalized_state_change.capacity() == 0 {
                    span.in_scope(|| {
                        tracing::warn!(
                            "slow consumer, dropping non_finalized_state_change stream after \
                             buffer filled"
                        );
                    });
                    return;
                }

                let Some((hash, block)) = non_finalized_state_change.recv().await else {
                    break;
                };

                let send = response_sender.send(Ok(BlockAndHash::new(hash, block)));
                match tokio::time::timeout(SEND_TIMEOUT, send).await {
                    Ok(Ok(())) => {}
                    Ok(Err(_)) => {
                        span.in_scope(|| {
                            tracing::info!(
                                "client disconnected, dropping non_finalized_state_change task"
                            );
                        });
                        return;
                    }
                    Err(_) => {
                        span.in_scope(|| {
                            tracing::warn!(
                                "slow consumer, dropping non_finalized_state_change stream after \
                                 send timed out"
                            );
                        });
                        return;
                    }
                }
            }

            span.in_scope(|| {
                tracing::warn!("non-finalized state change channel has closed");
            });

            let _ = response_sender
                .send(Err(Status::unavailable(
                    "non-finalized state change channel has closed",
                )))
                .await;
        });

        Ok(Response::new(Box::pin(response_stream)))
    }

    async fn mempool_change(
        &self,
        _: tonic::Request<Empty>,
    ) -> Result<Response<Self::MempoolChangeStream>, Status> {
        let span = Span::current();
        let (response_sender, response_receiver) = tokio::sync::mpsc::channel(RESPONSE_BUFFER_SIZE);
        let response_stream = ReceiverStream::new(response_receiver);
        let mut mempool_change = self.mempool_change.subscribe();

        tokio::spawn(async move {
            // Notify the client of chain tip changes until the channel is closed
            while let Ok(change) = mempool_change.recv().await {
                for tx_id in change.tx_ids() {
                    span.in_scope(|| {
                        tracing::debug!("mempool change: {:?}", change);
                    });

                    let msg = Ok(MempoolChangeMessage {
                        change_type: match change.kind() {
                            MempoolChangeKind::Added => 0,
                            MempoolChangeKind::Invalidated => 1,
                            MempoolChangeKind::Mined => 2,
                        },
                        tx_hash: tx_id.mined_id().bytes_in_display_order().to_vec(),
                        auth_digest: tx_id
                            .auth_digest()
                            .map(|d| d.bytes_in_display_order().to_vec())
                            .unwrap_or_default(),
                    });

                    let send = response_sender.send(msg);
                    match tokio::time::timeout(SEND_TIMEOUT, send).await {
                        Ok(Ok(())) => {}
                        Ok(Err(_)) => {
                            span.in_scope(|| {
                                tracing::info!("client disconnected, dropping mempool_change task");
                            });
                            return;
                        }
                        Err(_) => {
                            span.in_scope(|| {
                                tracing::warn!(
                                    "slow consumer, dropping mempool_change stream after \
                                     send timed out"
                                );
                            });
                            return;
                        }
                    }
                }
            }

            span.in_scope(|| {
                tracing::warn!("mempool_change channel has closed");
            });

            let _ = response_sender
                .send(Err(Status::unavailable(
                    "mempool_change channel has closed",
                )))
                .await;
        });

        Ok(Response::new(Box::pin(response_stream)))
    }

    async fn get_block(
        &self,
        request: tonic::Request<BlockRequest>,
    ) -> Result<Response<BlockAndHash>, Status> {
        // The request carries a single `hash_or_height` byte string: a 32-byte
        // block hash in display order, or a 4-byte big-endian block height. The
        // length tells the two apart.
        let hash_or_height = request.into_inner().hash_or_height;
        let hash_or_height = match hash_or_height.len() {
            BLOCK_HASH_BYTE_LEN => {
                zakura_state::HashOrHeight::Hash(hash_from_display_bytes(hash_or_height)?)
            }
            BLOCK_HEIGHT_BYTE_LEN => {
                let height = u32::from_be_bytes(
                    hash_or_height
                        .try_into()
                        .expect("length was validated by BLOCK_HEIGHT_BYTE_LEN"),
                );
                let height = block::Height::try_from(height).map_err(|_| {
                    Status::invalid_argument(format!("block height out of range: {height}"))
                })?;
                zakura_state::HashOrHeight::Height(height)
            }
            len => {
                return Err(Status::invalid_argument(format!(
                    "block request must be a {BLOCK_HASH_BYTE_LEN}-byte hash or a \
                     {BLOCK_HEIGHT_BYTE_LEN}-byte height, got {len} bytes"
                )));
            }
        };

        match self
            .read_state
            .clone()
            .oneshot(ReadRequest::Block(hash_or_height))
            .await
        {
            Ok(ReadResponse::Block(Some(block))) => {
                Ok(Response::new(BlockAndHash::new(block.hash(), block)))
            }
            Ok(ReadResponse::Block(None)) => Err(Status::not_found("block not found")),
            Ok(_) => unreachable!("unexpected response type from ReadStateService"),
            Err(error) => Err(Status::unavailable(format!(
                "failed to read block: {error}"
            ))),
        }
    }

    async fn get_block_range(
        &self,
        request: tonic::Request<BlockRangeRequest>,
    ) -> Result<Response<Self::GetBlockRangeStream>, Status> {
        // Serving raw block ranges needs the raw read path in
        // `zakura-state/indexer`, so the method sits behind this crate's
        // `indexer` feature and is not implemented without it.
        #[cfg(not(feature = "indexer"))]
        {
            let _ = request;
            return Err(Status::unimplemented(
                "GetBlockRange requires the `indexer` compile-time feature",
            ));
        }

        #[cfg(feature = "indexer")]
        {
            let BlockRangeRequest {
                start_height,
                end_height,
            } = request.into_inner();

            // Invalid ranges are rejected up front, without touching the state.
            let block::Height(start) = block::Height::try_from(start_height).map_err(|_| {
                Status::invalid_argument(format!("start height out of range: {start_height}"))
            })?;
            let block::Height(end) = block::Height::try_from(end_height).map_err(|_| {
                Status::invalid_argument(format!("end height out of range: {end_height}"))
            })?;

            if end < start {
                return Err(Status::invalid_argument(format!(
                    "end height {end_height} must not be below start height {start_height}"
                )));
            }

            let span = Span::current();
            let read_state = self.read_state.clone();
            let (response_sender, response_receiver) =
                tokio::sync::mpsc::channel(RESPONSE_BUFFER_SIZE);
            let response_stream = ReceiverStream::new(response_receiver);

            tokio::spawn(async move {
                // Serve the range in bounded chunks so the task never buffers more
                // than one chunk of blocks beyond the response buffer. When the
                // state returns fewer blocks than a chunk asked for, it stopped at
                // the finalized tip or before a missing block body, so the stream
                // ends cleanly there, following the `BlocksByHeightRange`
                // convention.
                let mut next = start;
                while next <= end {
                    // The subtraction and increment never overflow: `next <= end`,
                    // and `end` is a valid `Height`, which is far below `u32::MAX`.
                    let count = (end - next + 1).min(BLOCK_RANGE_CHUNK_SIZE);

                    let blocks = match read_state
                        .clone()
                        .oneshot(ReadRequest::RawBlocksByHeightRange {
                            start: block::Height(next),
                            count,
                        })
                        .await
                    {
                        Ok(ReadResponse::RawBlocks(blocks)) => blocks,
                        Ok(_) => unreachable!("unexpected response type from ReadStateService"),
                        Err(error) => {
                            span.in_scope(|| {
                                tracing::error!(
                                    ?error,
                                    "failed to read blocks for get_block_range"
                                );
                            });

                            // A stalled consumer must not pin this task, so
                            // the error send is bounded like the data sends.
                            let send = response_sender.send(Err(Status::unavailable(format!(
                                "failed to read blocks: {error}"
                            ))));
                            let _ = tokio::time::timeout(SEND_TIMEOUT, send).await;
                            return;
                        }
                    };

                    let block_count = u32::try_from(blocks.len())
                        .expect("the state returns at most `count` blocks, which is a u32");

                    for (height, data) in blocks {
                        let send = response_sender.send(Ok(BlockAndHeight::new(height, data)));
                        match tokio::time::timeout(SEND_TIMEOUT, send).await {
                            Ok(Ok(())) => {}
                            Ok(Err(_)) => {
                                span.in_scope(|| {
                                    tracing::info!(
                                        "client disconnected, dropping get_block_range task"
                                    );
                                });
                                return;
                            }
                            Err(_) => {
                                span.in_scope(|| {
                                    tracing::warn!(
                                        "slow consumer, dropping get_block_range stream after \
                                         send timed out"
                                    );
                                });
                                return;
                            }
                        }
                    }

                    if block_count < count {
                        break;
                    }

                    next += count;
                }

                // Dropping the sender ends the stream cleanly at the end of the
                // served range.
            });

            Ok(Response::new(Box::pin(response_stream)))
        }
    }
}

/// Decodes the chain tip hashes from a [`NonFinalizedStateChangeRequest`] into
/// a set of [`block::Hash`]es.
///
/// Each hash is expected to be 32 bytes in display order, matching the encoding
/// used when the server streams [`BlockAndHash`] messages back to the caller.
///
/// # Errors
///
/// Returns an [`invalid_argument`](Status::invalid_argument) status if there are
/// more hashes than the non-finalized state can hold chains
/// ([`MAX_NON_FINALIZED_CHAIN_FORKS`]), or if any hash is not exactly 32 bytes.
fn decode_known_chain_tips(chain_tip_hashes: Vec<Vec<u8>>) -> Result<HashSet<block::Hash>, Status> {
    // The non-finalized state holds at most
    // `MAX_NON_FINALIZED_CHAIN_FORKS` chains, so a caller can never legitimately
    // have more chain tips than that. Bound the untrusted input up front rather
    // than allocating a set sized by the request.
    if chain_tip_hashes.len() > MAX_NON_FINALIZED_CHAIN_FORKS {
        return Err(Status::invalid_argument(format!(
            "too many chain tip hashes: got {}, expected at most {MAX_NON_FINALIZED_CHAIN_FORKS}",
            chain_tip_hashes.len(),
        )));
    }

    chain_tip_hashes
        .into_iter()
        .map(hash_from_display_bytes)
        .collect()
}

/// Decodes a block hash in display order, rejecting wrong-length input.
fn hash_from_display_bytes(hash: Vec<u8>) -> Result<block::Hash, Status> {
    let bytes: [u8; BLOCK_HASH_BYTE_LEN] = hash.try_into().map_err(|hash: Vec<u8>| {
        Status::invalid_argument(format!(
            "invalid block hash length: expected {BLOCK_HASH_BYTE_LEN} bytes, got {}",
            hash.len()
        ))
    })?;

    Ok(block::Hash::from_bytes_in_display_order(&bytes))
}

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

    fn hash(byte: u8) -> block::Hash {
        block::Hash::from_bytes_in_display_order(&[byte; 32])
    }

    #[test]
    fn decode_known_chain_tips_round_trips_display_order() {
        let hashes = [hash(1), hash(2), hash(3)];
        let encoded = hashes
            .iter()
            .map(|h| h.bytes_in_display_order().to_vec())
            .collect();

        let decoded = decode_known_chain_tips(encoded).expect("valid hashes should decode");

        assert_eq!(decoded, hashes.into_iter().collect());
    }

    #[test]
    fn decode_known_chain_tips_accepts_empty() {
        assert!(decode_known_chain_tips(Vec::new())
            .expect("empty input should decode")
            .is_empty());
    }

    #[test]
    fn decode_known_chain_tips_dedups() {
        let encoded = vec![
            hash(7).bytes_in_display_order().to_vec(),
            hash(7).bytes_in_display_order().to_vec(),
        ];

        let decoded = decode_known_chain_tips(encoded).expect("duplicate hashes should decode");

        assert_eq!(decoded, std::iter::once(hash(7)).collect());
    }

    #[test]
    fn decode_known_chain_tips_rejects_wrong_length() {
        let status = decode_known_chain_tips(vec![vec![0; 31]])
            .expect_err("a 31-byte hash should be rejected");

        assert_eq!(status.code(), Code::InvalidArgument);
    }

    #[test]
    fn decode_known_chain_tips_rejects_too_many() {
        // This limit is 10, so it fits in a `u8`.
        let encoded = (0..=MAX_NON_FINALIZED_CHAIN_FORKS as u8)
            .map(|b| hash(b).bytes_in_display_order().to_vec())
            .collect();

        let status = decode_known_chain_tips(encoded)
            .expect_err("more than MAX_NON_FINALIZED_CHAIN_FORKS hashes should be rejected");

        assert_eq!(status.code(), Code::InvalidArgument);
    }
}