linera-exporter 0.15.21

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

use std::{
    collections::HashSet,
    marker::PhantomData,
    sync::{atomic::AtomicU64, Arc},
};

use futures::future::try_join_all;
#[cfg(with_metrics)]
use linera_base::prometheus_util::MeasureLatency as _;
use linera_base::{
    crypto::CryptoHash,
    data_types::{Blob, BlockHeight},
    identifiers::{BlobId, ChainId},
};
use linera_chain::types::ConfirmedBlockCertificate;
use linera_sdk::{ensure, views::View};
use linera_storage::{Arc as CacheArc, Storage};
use linera_views::{
    batch::Batch, context::Context, log_view::LogView, store::WritableKeyValueStore as _,
};
use mini_moka::unsync::Cache as LfuCache;
use quick_cache::{sync::Cache as FifoCache, Weighter};
use tokio::sync::RwLock;

#[cfg(with_metrics)]
use crate::metrics;
use crate::{
    common::{BlockId, CanonicalBlock, ExporterError, LiteBlockId},
    config::{DestinationId, LimitsConfig},
    state::{BlockExporterStateView, DestinationStates},
};

pub(super) struct ExporterStorage<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    shared_storage: SharedStorage<S::BlockExporterContext, S>,
}

type BlobCache = FifoCache<BlobId, Arc<Blob>, BlobCacheWeighter>;
type BlockCache = FifoCache<CryptoHash, CacheArc<ConfirmedBlockCertificate>, BlockCacheWeighter>;

struct SharedStorage<C, S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    storage: S,
    destination_states: DestinationStates,
    shared_canonical_state: CanonicalState<C>,
    blobs_cache: Arc<BlobCache>,
    blocks_cache: Arc<BlockCache>,
}

pub(super) struct BlockProcessorStorage<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    blob_state_cache: LfuCache<BlobId, ()>,
    chain_states_cache: LfuCache<ChainId, LiteBlockId>,
    shared_storage: SharedStorage<S::BlockExporterContext, S>,
    // Handle on the persistent storage where the exporter state is pushed to periodically.
    exporter_state_view: BlockExporterStateView<<S as Storage>::BlockExporterContext>,
}

impl<C, S> SharedStorage<C, S>
where
    C: Context + Send + Sync + 'static,
    S: Storage + Clone + Send + Sync + 'static,
{
    fn new(
        storage: S,
        state_context: LogView<C, CanonicalBlock>,
        destination_states: DestinationStates,
        limits: LimitsConfig,
    ) -> Self {
        let shared_canonical_state = CanonicalState::new(state_context);
        let blobs_cache = Arc::new(FifoCache::with_weighter(
            limits.blob_cache_items_capacity as usize,
            (limits.blob_cache_weight_mb as u64) * 1024 * 1024,
            CacheWeighter::default(),
        ));
        let blocks_cache = Arc::new(FifoCache::with_weighter(
            limits.block_cache_items_capacity as usize,
            (limits.block_cache_weight_mb as u64) * 1024 * 1024,
            CacheWeighter::default(),
        ));

        Self {
            storage,
            shared_canonical_state,
            blobs_cache,
            blocks_cache,
            destination_states,
        }
    }

    async fn get_block(
        &self,
        hash: CryptoHash,
    ) -> Result<CacheArc<ConfirmedBlockCertificate>, ExporterError> {
        match self.blocks_cache.get_value_or_guard_async(&hash).await {
            Ok(value) => Ok(value),
            Err(guard) => {
                #[cfg(with_metrics)]
                metrics::GET_CERTIFICATE_HISTOGRAM.measure_latency();
                let block = self
                    .storage
                    .read_certificate(hash)
                    .await?
                    .ok_or_else(|| ExporterError::ReadCertificateError(hash))?;
                guard.insert(block.clone()).ok();
                Ok(block)
            }
        }
    }

    async fn get_blob(&self, blob_id: BlobId) -> Result<Arc<Blob>, ExporterError> {
        match self.blobs_cache.get_value_or_guard_async(&blob_id).await {
            Ok(blob) => Ok(blob),
            Err(guard) => {
                #[cfg(with_metrics)]
                metrics::GET_BLOB_HISTOGRAM.measure_latency();
                let blob = self.storage.read_blob(blob_id).await?.unwrap().into_std();
                guard.insert(blob.clone()).ok();
                Ok(blob)
            }
        }
    }

    async fn get_blobs(&self, blobs: &[BlobId]) -> Result<Vec<Arc<Blob>>, ExporterError> {
        let tasks = blobs.iter().map(|id| self.get_blob(*id));
        let results = try_join_all(tasks).await?;
        Ok(results)
    }

    /// Enqueues a block into the in-memory buffer.
    /// Only acquires the buffer lock; no I/O is performed.
    #[allow(unused_variables)]
    async fn push_block(&self, block: CanonicalBlock) {
        let count = self.shared_canonical_state.push(block).await;
        #[cfg(with_metrics)]
        metrics::CANONICAL_STATE_HEIGHT.set(count as i64);
    }

    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            shared_canonical_state: self.shared_canonical_state.clone(),
            blobs_cache: self.blobs_cache.clone(),
            blocks_cache: self.blocks_cache.clone(),
            destination_states: self.destination_states.clone(),
        }
    }
}

impl<S> ExporterStorage<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    fn new(shared_storage: SharedStorage<S::BlockExporterContext, S>) -> Self {
        Self { shared_storage }
    }

    pub(crate) async fn get_block_with_blob_ids(
        &self,
        index: usize,
    ) -> Result<(CacheArc<ConfirmedBlockCertificate>, Vec<BlobId>), ExporterError> {
        let block = self
            .shared_storage
            .shared_canonical_state
            .get(index)
            .await?;

        Ok((
            self.shared_storage.get_block(block.block_hash).await?,
            block.blobs.into(),
        ))
    }

    pub(crate) async fn get_block_with_blobs(
        &self,
        index: usize,
    ) -> Result<(CacheArc<ConfirmedBlockCertificate>, Vec<Arc<Blob>>), ExporterError> {
        let canonical_block = self
            .shared_storage
            .shared_canonical_state
            .get(index)
            .await?;

        let block_task = self.shared_storage.get_block(canonical_block.block_hash);
        let blobs_task = self.shared_storage.get_blobs(&canonical_block.blobs);

        let (block, blobs) = tokio::try_join!(block_task, blobs_task)?;
        Ok((block, blobs))
    }

    pub(crate) async fn get_blob(&self, blob_id: BlobId) -> Result<Arc<Blob>, ExporterError> {
        self.shared_storage.get_blob(blob_id).await
    }

    pub(crate) fn load_destination_state(&self, id: &DestinationId) -> Arc<AtomicU64> {
        self.shared_storage.destination_states.load_state(id)
    }

    pub(crate) fn clone(&self) -> Self {
        ExporterStorage::new(self.shared_storage.clone())
    }

    pub(crate) async fn get_latest_index(&self) -> usize {
        self.shared_storage
            .shared_canonical_state
            .latest_index()
            .await
    }
}

impl<S> BlockProcessorStorage<S>
where
    S: Storage + Clone + Send + Sync + 'static,
{
    pub(super) async fn load(
        storage: S,
        id: u32,
        destinations: Vec<DestinationId>,
        limits: LimitsConfig,
    ) -> Result<(Self, ExporterStorage<S>), ExporterError> {
        let context = storage.block_exporter_context(id).await?;
        let (view, canonical_state, destination_states) =
            BlockExporterStateView::initiate(context, destinations).await?;

        let chain_states_cache_capacity =
            ((limits.auxiliary_cache_size_mb / 3) as u64 * 1024 * 1024)
                / (size_of::<CryptoHash>() + size_of::<LiteBlockId>()) as u64;
        let chain_states_cache = LfuCache::builder()
            .max_capacity(chain_states_cache_capacity)
            .build();

        let blob_state_cache_capacity = ((limits.auxiliary_cache_size_mb / 3) as u64 * 1024 * 1024)
            / (size_of::<BlobId>() as u64);
        let blob_state_cache = LfuCache::builder()
            .max_capacity(blob_state_cache_capacity)
            .build();

        let shared_storage =
            SharedStorage::new(storage, canonical_state, destination_states, limits);
        let exporter_storage = ExporterStorage::new(shared_storage.clone());

        Ok((
            Self {
                shared_storage,
                chain_states_cache,
                exporter_state_view: view,
                blob_state_cache,
            },
            exporter_storage,
        ))
    }

    pub(super) async fn get_block(
        &self,
        hash: CryptoHash,
    ) -> Result<CacheArc<ConfirmedBlockCertificate>, ExporterError> {
        self.shared_storage.get_block(hash).await
    }

    pub(super) async fn get_blob(&self, blob: BlobId) -> Result<Arc<Blob>, ExporterError> {
        self.shared_storage.get_blob(blob).await
    }

    pub(super) async fn is_blob_indexed(&mut self, blob: BlobId) -> Result<bool, ExporterError> {
        match self.blob_state_cache.get(&blob) {
            Some(_) => Ok(true),
            None => self.exporter_state_view.is_blob_indexed(blob).await,
        }
    }

    pub(super) async fn is_block_indexed(
        &mut self,
        block_id: &BlockId,
    ) -> Result<bool, ExporterError> {
        if let Some(status) = self.chain_states_cache.get(&block_id.chain_id) {
            return Ok(status.height >= block_id.height);
        }

        if let Some(status) = self
            .exporter_state_view
            .get_chain_status(&block_id.chain_id)
            .await?
        {
            let result = status.height >= block_id.height;
            self.chain_states_cache.insert(block_id.chain_id, status);
            return Ok(result);
        }

        Err(ExporterError::UnprocessedChain)
    }

    pub(super) async fn index_chain(&mut self, block_id: &BlockId) -> Result<(), ExporterError> {
        ensure!(
            block_id.height == BlockHeight::ZERO,
            ExporterError::BadInitialization
        );
        self.exporter_state_view.initialize_chain(*block_id).await?;
        self.chain_states_cache
            .insert(block_id.chain_id, (*block_id).into());

        Ok(())
    }

    pub(super) async fn index_block(&mut self, block_id: &BlockId) -> Result<bool, ExporterError> {
        if block_id.height == BlockHeight::ZERO {
            self.index_chain(block_id).await?;
            return Ok(true);
        }

        if self.exporter_state_view.index_block(*block_id).await? {
            self.chain_states_cache
                .insert(block_id.chain_id, (*block_id).into());
            return Ok(true);
        }

        Ok(false)
    }

    pub(super) fn index_blob(&mut self, blob: BlobId) -> Result<(), ExporterError> {
        self.exporter_state_view.index_blob(blob)?;
        self.blob_state_cache.insert(blob, ());
        Ok(())
    }

    /// Enqueues a block into the shared canonical state buffer.
    /// Only acquires the buffer lock; no I/O is performed.
    pub(super) async fn push_block(&self, block: CanonicalBlock) {
        self.shared_storage.push_block(block).await
    }

    pub(super) fn new_committee(&mut self, committee_destinations: HashSet<DestinationId>) {
        committee_destinations.into_iter().for_each(|id| {
            let state = match self.shared_storage.destination_states.get(&id) {
                None => {
                    tracing::info!(id=?id, "adding new committee member");
                    #[cfg(with_metrics)]
                    {
                        metrics::DESTINATION_STATE_COUNTER
                            .with_label_values(&[id.address()])
                            .reset();
                    }
                    Arc::new(AtomicU64::new(0))
                }
                Some(state) => state.clone(),
            };
            self.shared_storage.destination_states.insert(id, state);
        });
    }

    pub(super) fn set_latest_committee_blob(&mut self, blob_id: BlobId) {
        self.exporter_state_view.set_latest_committee_blob(blob_id);
    }

    pub(super) fn get_latest_committee_blob(&self) -> Option<BlobId> {
        self.exporter_state_view.get_latest_committee_blob()
    }

    pub(super) async fn save(&mut self) -> Result<(), ExporterError> {
        let mut batch = Batch::new();

        self.shared_storage
            .shared_canonical_state
            .flush(&mut batch)
            .await?;

        self.exporter_state_view
            .set_destination_states(self.shared_storage.destination_states.clone());

        self.exporter_state_view.pre_save(&mut batch)?;
        #[cfg(with_metrics)]
        metrics::SAVE_HISTOGRAM.measure_latency();
        if let Err(e) = self
            .exporter_state_view
            .context()
            .store()
            .write_batch(batch)
            .await
        {
            Err(ExporterError::ViewError(e.into()))?;
        };
        self.exporter_state_view.post_save();

        Ok(())
    }
}

/// In-memory buffer of canonical blocks awaiting flush to the LogView.
struct CanonicalBuffer {
    /// Total count of canonical blocks (persisted + buffered).
    count: usize,
    /// Blocks not yet flushed to the LogView.
    items: Vec<CanonicalBlock>,
}

/// Canonical state split into a fast in-memory buffer and a persistent log.
///
/// The buffer and log use separate locks so that `push` (buffer-only) is never
/// blocked by slow persistent reads, and readers checking the buffer don't hold
/// a lock across I/O.
struct CanonicalState<C> {
    buffer: Arc<RwLock<CanonicalBuffer>>,
    log: Arc<RwLock<LogView<C, CanonicalBlock>>>,
}

impl<C> Clone for CanonicalState<C> {
    fn clone(&self) -> Self {
        Self {
            buffer: self.buffer.clone(),
            log: self.log.clone(),
        }
    }
}

impl<C> CanonicalState<C>
where
    C: Context + Send + Sync + 'static,
{
    fn new(state_context: LogView<C, CanonicalBlock>) -> Self {
        let count = state_context.count();
        Self {
            buffer: Arc::new(RwLock::new(CanonicalBuffer {
                count,
                items: Vec::new(),
            })),
            log: Arc::new(RwLock::new(state_context)),
        }
    }

    /// Returns the total number of canonical blocks (persisted + buffered).
    async fn latest_index(&self) -> usize {
        self.buffer.read().await.count
    }

    /// Retrieves a canonical block by index.
    ///
    /// Checks the in-memory buffer first (fast, no I/O). On a buffer miss,
    /// reads from the persistent LogView under a separate lock,
    /// avoiding holding any lock across I/O on the hot path for `push`.
    async fn get(&self, index: usize) -> Result<CanonicalBlock, ExporterError> {
        {
            let buf = self.buffer.read().await;
            let buffer_start = buf.count - buf.items.len();
            if index >= buffer_start {
                return buf
                    .items
                    .get(index - buffer_start)
                    .cloned()
                    .ok_or(ExporterError::UnprocessedBlock);
            }
        }

        #[cfg(with_metrics)]
        metrics::GET_CANONICAL_BLOCK_HISTOGRAM.measure_latency();
        let log = self.log.read().await;
        log.get(index).await?.ok_or(ExporterError::UnprocessedBlock)
    }

    /// Enqueues a block into the buffer. Returns the new total count.
    /// Only acquires the buffer lock; no I/O is performed.
    async fn push(&self, value: CanonicalBlock) -> usize {
        let mut buf = self.buffer.write().await;
        buf.items.push(value);
        buf.count += 1;
        buf.count
    }

    /// Drains the buffer into the LogView and serializes changes into `batch`.
    ///
    /// Both locks are held during the drain-and-push sequence to prevent
    /// readers from seeing a gap where a block is neither in the buffer
    /// nor yet in the log.
    async fn flush(&self, batch: &mut Batch) -> Result<(), ExporterError> {
        let mut buf = self.buffer.write().await;
        let mut log = self.log.write().await;

        for value in buf.items.drain(..) {
            log.push(value);
        }

        log.pre_save(batch)?;
        log.post_save();

        Ok(())
    }
}

#[derive(Clone)]
struct CacheWeighter<Q, V> {
    key: PhantomData<Q>,
    value: PhantomData<V>,
}

impl Weighter<BlobId, Arc<Blob>> for BlobCacheWeighter {
    fn weight(&self, _key: &BlobId, val: &Arc<Blob>) -> u64 {
        (size_of::<BlobId>()
            + size_of::<Arc<Blob>>()
            + 2 * size_of::<usize>() // two reference counts in Arc, just a micro-optimization
            + size_of::<Blob>()
            + val.bytes().len()) as u64
    }
}

impl Weighter<CryptoHash, CacheArc<ConfirmedBlockCertificate>> for BlockCacheWeighter {
    fn weight(&self, _key: &CryptoHash, _val: &CacheArc<ConfirmedBlockCertificate>) -> u64 {
        (size_of::<CryptoHash>()
            + 2 * size_of::<usize>()
            + size_of::<CacheArc<ConfirmedBlockCertificate>>()
            + 1_000_000) as u64 // maximum block size in testnet resource control policy
    }
}

impl<Q, V> Default for CacheWeighter<Q, V> {
    fn default() -> Self {
        Self {
            key: PhantomData,
            value: PhantomData,
        }
    }
}

type BlobCacheWeighter = CacheWeighter<BlobId, Arc<Blob>>;
type BlockCacheWeighter = CacheWeighter<CryptoHash, CacheArc<ConfirmedBlockCertificate>>;