liminal/durability/

head_checkpoint_stop.rs

//! Clean-stop checkpoint ownership at the shared store's append boundary.
use super::{DurabilityError, DurableStore, ScanCursor, ScanPage, StoredEntry};
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};

/// Tracks the streams appended through this store for the final shutdown checkpoint.
#[derive(Debug)]
pub struct HeadCheckpointStop {
    inner: Arc<dyn DurableStore>,
    touched: Mutex<BTreeMap<String, u64>>,
    stopping: AtomicBool,
    generation: AtomicU64,
}

impl HeadCheckpointStop {
    /// Wraps the shared store before any server stream is constructed.
    #[must_use]
    pub fn new(inner: Arc<dyn DurableStore>) -> Self {
        Self {
            inner,
            touched: Mutex::default(),
            stopping: AtomicBool::new(false),
            generation: AtomicU64::new(0),
        }
    }

    /// Records every head still owed at a clean stop, before the final store flush.
    ///
    /// # Errors
    /// Propagates the store's checkpoint refusal.
    pub async fn checkpoint_heads(&self) -> Result<(), DurabilityError> {
        self.stopping.store(true, Ordering::Release);
        self.advance_touched_heads().await
    }

    fn touched(&self) -> Result<MutexGuard<'_, BTreeMap<String, u64>>, DurabilityError> {
        self.touched.lock().map_err(|_| {
            DurabilityError::ConfigError("clean-stop head tracking lock poisoned".to_owned())
        })
    }

    fn note_append(&self, stream: &str) -> Result<(), DurabilityError> {
        // A receipt stays unique after its stream leaves the pending map, so
        // another checkpoint cannot clear a later append of the same stream.
        let generation = self
            .generation
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |previous| {
                previous.checked_add(1)
            })
            .map_err(|_| {
                DurabilityError::ConfigError(
                    "clean-stop head tracking generation exhausted".to_owned(),
                )
            })?;
        self.touched()?.insert(stream.to_owned(), generation);
        Ok(())
    }

    async fn advance_touched_heads(&self) -> Result<(), DurabilityError> {
        let touched = self.touched()?.clone();
        for (stream, generation) in touched {
            if let Some(next) = self.inner.read_stream_next_seq(&stream).await? {
                let key = super::head_checkpoint_key(&stream);
                let recorded = self.inner.read_value(&key).await?.unwrap_or(0);
                if recorded > next {
                    return Err(DurabilityError::ConfigError(format!(
                        "clean-stop head {recorded} exceeds stream {stream} head {next}"
                    )));
                }
                if recorded < next {
                    self.inner.cas(&key, recorded, next).await?;
                }
            }
            let mut current = self.touched()?;
            if current.get(&stream) == Some(&generation) {
                current.remove(&stream);
            }
        }
        Ok(())
    }
}

#[async_trait::async_trait]
impl DurableStore for HeadCheckpointStop {
    fn reclaim_spine(&self) -> Option<Arc<haematite::EventStore>> {
        self.inner.reclaim_spine()
    }
    fn max_request_bytes(&self) -> Option<usize> {
        self.inner.max_request_bytes()
    }
    fn store_counts(&self) -> Option<super::StoreCountsHandle> {
        self.inner.store_counts()
    }
    async fn truncate_before(
        &self,
        stream_key: &str,
        before_seq: u64,
    ) -> Result<u64, DurabilityError> {
        self.inner.truncate_before(stream_key, before_seq).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 apply_routed_batch(
        &self,
        route_key: &str,
        items: Vec<crate::durability::RoutedBatchItem>,
    ) -> Result<(), crate::durability::DurabilityError> {
        self.inner.apply_routed_batch(route_key, items).await
    }

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

    async fn get_routed(
        &self,
        route_key: &str,
        key: &[u8],
    ) -> Result<Option<Vec<u8>>, crate::durability::DurabilityError> {
        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.note_append(stream_key)?;
        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.note_append(stream_key)?;
        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 read_at(
        &self,
        stream_key: &str,
        sequence: u64,
    ) -> Result<Option<StoredEntry>, DurabilityError> {
        self.inner.read_at(stream_key, sequence).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_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> {
        // Final-flush observations can append through this same store. Their
        // own durable flush must record the newly appended head too.
        if self.stopping.load(Ordering::Acquire) {
            self.advance_touched_heads().await?;
        }
        self.inner.flush().await
    }
}