liminal-server 0.24.5

//! The two read-counting stores of the checkpoint-batch pins: row reads for
//! the containment pin and history-node reads for the restore pin. Each
//! delegates every call, counting the reads its pin is about.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use liminal::durability::{
    DurabilityError, DurableStore, RoutedBatchEntry, RoutedBatchItem, ScanCursor, ScanPage,
    StoredEntry,
};

use super::{is_node_key, is_row_key};

/// Row reads on one conversation's store, for the containment pin.
#[derive(Debug)]
pub(super) struct RowReads {
    pub(super) inner: Arc<dyn DurableStore>,
    pub(super) reads: AtomicUsize,
}

#[async_trait::async_trait]
impl DurableStore for RowReads {
    async fn apply_routed_batch(
        &self,
        route_key: &str,
        items: Vec<RoutedBatchItem>,
    ) -> Result<(), DurabilityError> {
        self.inner.apply_routed_batch(route_key, items).await
    }

    async fn apply_routed_entries(
        &self,
        route_key: &str,
        entries: Vec<RoutedBatchEntry>,
    ) -> Result<(), DurabilityError> {
        self.inner.apply_routed_entries(route_key, entries).await
    }

    async fn get_routed(
        &self,
        route_key: &str,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, DurabilityError> {
        if is_row_key(key) {
            self.reads.fetch_add(1, Ordering::SeqCst);
        }
        self.inner.get_routed(route_key, key).await
    }

    async fn append(
        &self,
        stream_key: &str,
        payload: Vec<u8>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner.append(stream_key, payload, expected_seq).await
    }

    async fn append_all(
        &self,
        stream_key: &str,
        payloads: Vec<Vec<u8>>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner
            .append_all(stream_key, payloads, expected_seq)
            .await
    }

    async fn read_from(
        &self,
        stream_key: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Vec<StoredEntry>, DurabilityError> {
        let entries = self.inner.read_from(stream_key, offset, limit).await?;
        self.reads.fetch_add(entries.len(), Ordering::SeqCst);
        Ok(entries)
    }

    async fn cas(&self, key: &str, old_value: u64, new_value: u64) -> Result<(), DurabilityError> {
        self.inner.cas(key, old_value, new_value).await
    }

    async fn read_stream_next_seq(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
        self.inner.read_stream_next_seq(stream_key).await
    }

    async fn read_value(&self, key: &str) -> Result<Option<u64>, DurabilityError> {
        self.inner.read_value(key).await
    }

    async fn scan_routed(
        &self,
        route_key: &str,
        from: &[u8],
        to: &[u8],
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, DurabilityError> {
        self.inner.scan_routed(route_key, from, to).await
    }

    async fn remove_routed_range(
        &self,
        route_key: &str,
        from: &[u8],
        to: &[u8],
    ) -> Result<u64, DurabilityError> {
        self.inner.remove_routed_range(route_key, from, to).await
    }

    async fn scan_page(
        &self,
        prefix: &str,
        cursor: Option<ScanCursor>,
        limit: usize,
    ) -> Result<ScanPage, DurabilityError> {
        self.inner.scan_page(prefix, cursor, limit).await
    }

    async fn flush(&self) -> Result<(), DurabilityError> {
        self.inner.flush().await
    }
}

/// History-node point reads on one store, for the restore pin: every
/// `get_routed` whose key is a checkpoint history node.
#[derive(Debug)]
pub(super) struct NodeReads {
    /// The store every call is delegated to.
    pub(super) inner: Arc<dyn DurableStore>,
    /// `get_routed` calls on a key under [`NODE_KEY_PREFIX`].
    pub(super) reads: AtomicUsize,
}

#[async_trait::async_trait]
impl DurableStore for NodeReads {
    async fn apply_routed_batch(
        &self,
        route_key: &str,
        items: Vec<RoutedBatchItem>,
    ) -> Result<(), DurabilityError> {
        self.inner.apply_routed_batch(route_key, items).await
    }

    async fn apply_routed_entries(
        &self,
        route_key: &str,
        entries: Vec<RoutedBatchEntry>,
    ) -> Result<(), DurabilityError> {
        self.inner.apply_routed_entries(route_key, entries).await
    }

    async fn get_routed(
        &self,
        route_key: &str,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, DurabilityError> {
        if is_node_key(key) {
            self.reads.fetch_add(1, Ordering::SeqCst);
        }
        self.inner.get_routed(route_key, key).await
    }

    async fn append(
        &self,
        stream_key: &str,
        payload: Vec<u8>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner.append(stream_key, payload, expected_seq).await
    }

    async fn append_all(
        &self,
        stream_key: &str,
        payloads: Vec<Vec<u8>>,
        expected_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner
            .append_all(stream_key, payloads, expected_seq)
            .await
    }

    async fn read_from(
        &self,
        stream_key: &str,
        offset: u64,
        limit: usize,
    ) -> Result<Vec<StoredEntry>, DurabilityError> {
        self.inner.read_from(stream_key, offset, limit).await
    }

    async fn cas(&self, key: &str, old_value: u64, new_value: u64) -> Result<(), DurabilityError> {
        self.inner.cas(key, old_value, new_value).await
    }

    async fn read_stream_next_seq(&self, stream_key: &str) -> Result<Option<u64>, DurabilityError> {
        self.inner.read_stream_next_seq(stream_key).await
    }

    async fn read_value(&self, key: &str) -> Result<Option<u64>, DurabilityError> {
        self.inner.read_value(key).await
    }

    async fn scan_routed(
        &self,
        route_key: &str,
        from: &[u8],
        to: &[u8],
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, DurabilityError> {
        self.inner.scan_routed(route_key, from, to).await
    }

    async fn remove_routed_range(
        &self,
        route_key: &str,
        from: &[u8],
        to: &[u8],
    ) -> Result<u64, DurabilityError> {
        self.inner.remove_routed_range(route_key, from, to).await
    }

    async fn scan_page(
        &self,
        prefix: &str,
        cursor: Option<ScanCursor>,
        limit: usize,
    ) -> Result<ScanPage, DurabilityError> {
        self.inner.scan_page(prefix, cursor, limit).await
    }

    async fn flush(&self) -> Result<(), DurabilityError> {
        self.inner.flush().await
    }
}