datum_cdc/

checkpoint.rs

use std::{
    collections::BTreeMap,
    fs,
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};

use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

use crate::{CdcError, CdcOffset, CdcResult, PgLsn};

pub(crate) enum FeedbackCommand {
    Checkpoint(CdcOffset),
    Stop,
    DropSlot {
        force: bool,
        reply: tokio::sync::oneshot::Sender<CdcResult<()>>,
    },
}

/// Durable checkpoint store used to resume a source after process restart.
pub trait CdcCheckpointStore: Send + Sync + 'static {
    fn load(&self, slot: &str) -> CdcResult<Option<CdcOffset>>;
    fn store(&self, offset: &CdcOffset) -> CdcResult<()>;
}

/// File-backed checkpoint store.
///
/// If the path is a directory, one JSON file per slot is written under it. If
/// the path is a file, that exact file is used.
#[derive(Debug, Clone)]
pub struct FileCheckpointStore {
    path: Arc<PathBuf>,
}

impl FileCheckpointStore {
    #[must_use]
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: Arc::new(path.into()),
        }
    }

    fn slot_path(&self, slot: &str) -> PathBuf {
        if self.path.extension().is_some() {
            return (*self.path).clone();
        }
        self.path.join(format!("{slot}.json"))
    }
}

impl CdcCheckpointStore for FileCheckpointStore {
    fn load(&self, slot: &str) -> CdcResult<Option<CdcOffset>> {
        let path = self.slot_path(slot);
        if !path.exists() {
            return Ok(None);
        }
        let bytes = fs::read(&path)?;
        let record: CheckpointRecord = serde_json::from_slice(&bytes)?;
        Ok(Some(record.offset))
    }

    fn store(&self, offset: &CdcOffset) -> CdcResult<()> {
        let path = self.slot_path(&offset.slot);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let tmp = path.with_extension("json.tmp");
        let bytes = serde_json::to_vec_pretty(&CheckpointRecord {
            offset: offset.clone(),
        })?;
        fs::write(&tmp, bytes)?;
        fs::rename(tmp, path)?;
        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct CheckpointRecord {
    offset: CdcOffset,
}

/// In-memory checkpoint store useful for tests and embedded ephemeral readers.
#[derive(Debug, Default)]
pub struct MemoryCheckpointStore {
    offsets: Mutex<BTreeMap<String, CdcOffset>>,
}

impl MemoryCheckpointStore {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl CdcCheckpointStore for MemoryCheckpointStore {
    fn load(&self, slot: &str) -> CdcResult<Option<CdcOffset>> {
        Ok(self
            .offsets
            .lock()
            .expect("checkpoint store poisoned")
            .get(slot)
            .cloned())
    }

    fn store(&self, offset: &CdcOffset) -> CdcResult<()> {
        self.offsets
            .lock()
            .expect("checkpoint store poisoned")
            .insert(offset.slot.clone(), offset.clone());
        Ok(())
    }
}

/// Handle used by downstream code after it has durably accepted an event.
#[derive(Clone)]
pub struct CdcCheckpointHandle {
    slot: String,
    store: Option<Arc<dyn CdcCheckpointStore>>,
    commands: mpsc::UnboundedSender<FeedbackCommand>,
}

impl CdcCheckpointHandle {
    pub(crate) fn new(
        slot: String,
        store: Option<Arc<dyn CdcCheckpointStore>>,
        commands: mpsc::UnboundedSender<FeedbackCommand>,
    ) -> Self {
        Self {
            slot,
            store,
            commands,
        }
    }

    /// Persist and announce a downstream checkpoint.
    ///
    /// The replication task will advance PostgreSQL feedback only after all
    /// events in the next contiguous transaction have been checkpointed.
    pub fn checkpoint(&self, offset: CdcOffset) -> CdcResult<()> {
        if offset.slot != self.slot {
            return Err(CdcError::Checkpoint(format!(
                "checkpoint for slot {:?} cannot be applied to slot {:?}",
                offset.slot, self.slot
            )));
        }
        if let Some(store) = &self.store {
            store.store(&offset)?;
        }
        self.commands
            .send(FeedbackCommand::Checkpoint(offset))
            .map_err(|_| CdcError::ChannelClosed)
    }

    pub async fn checkpoint_async(&self, offset: CdcOffset) -> CdcResult<()> {
        self.checkpoint(offset)
    }
}

#[derive(Debug)]
pub(crate) struct FeedbackCoordinator {
    slot: String,
    watermark: PgLsn,
    pending: BTreeMap<PgLsn, PendingTx>,
}

#[derive(Debug)]
struct PendingTx {
    seen: Vec<bool>,
    seen_count: usize,
}

impl FeedbackCoordinator {
    pub(crate) fn new(slot: String, watermark: PgLsn) -> Self {
        Self {
            slot,
            watermark,
            pending: BTreeMap::new(),
        }
    }

    pub(crate) fn register_tx(&mut self, tx_end_lsn: PgLsn, event_count: u32) -> CdcResult<()> {
        if event_count == 0 || tx_end_lsn <= self.watermark {
            return Ok(());
        }
        let event_count = usize::try_from(event_count).map_err(|_| {
            CdcError::Checkpoint("transaction event count does not fit usize".into())
        })?;
        self.pending
            .entry(tx_end_lsn)
            .or_insert_with(|| PendingTx::new(event_count));
        Ok(())
    }

    pub(crate) fn checkpoint(&mut self, offset: &CdcOffset) -> CdcResult<Option<PgLsn>> {
        if offset.slot != self.slot || offset.tx_end_lsn <= self.watermark {
            return Ok(None);
        }
        self.register_tx(offset.tx_end_lsn, offset.event_count)?;
        let Some(tx) = self.pending.get_mut(&offset.tx_end_lsn) else {
            return Ok(None);
        };
        tx.mark(offset.event_index)?;
        let mut advanced = None;
        while let Some((lsn, tx)) = self.pending.first_key_value() {
            if !tx.is_complete() {
                break;
            }
            let lsn = *lsn;
            self.watermark = lsn;
            advanced = Some(lsn);
            self.pending.remove(&lsn);
        }
        Ok(advanced)
    }

    #[cfg(test)]
    #[must_use]
    pub(crate) fn watermark(&self) -> PgLsn {
        self.watermark
    }
}

impl PendingTx {
    fn new(event_count: usize) -> Self {
        Self {
            seen: vec![false; event_count],
            seen_count: 0,
        }
    }

    fn mark(&mut self, event_index: u32) -> CdcResult<()> {
        let event_index = usize::try_from(event_index)
            .map_err(|_| CdcError::Checkpoint("event index does not fit usize".into()))?;
        let Some(seen) = self.seen.get_mut(event_index) else {
            return Err(CdcError::Checkpoint(format!(
                "event index {event_index} outside transaction event count {}",
                self.seen.len()
            )));
        };
        if !*seen {
            *seen = true;
            self.seen_count += 1;
        }
        Ok(())
    }

    fn is_complete(&self) -> bool {
        self.seen_count == self.seen.len()
    }
}

#[allow(dead_code)]
fn _assert_path_send_sync(_: &Path) {}

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

    fn offset(slot: &str, tx: u64, index: u32, count: u32) -> CdcOffset {
        CdcOffset {
            slot: slot.to_owned(),
            tx_end_lsn: PgLsn::from_u64(tx),
            commit_lsn: PgLsn::from_u64(tx - 1),
            xid: 7,
            event_index: index,
            event_count: count,
        }
    }

    #[test]
    fn feedback_waits_for_complete_contiguous_transactions() {
        let mut feedback = FeedbackCoordinator::new("slot".to_owned(), PgLsn::ZERO);
        feedback
            .register_tx(PgLsn::from_u64(10), 2)
            .expect("register tx 10");
        feedback
            .register_tx(PgLsn::from_u64(20), 1)
            .expect("register tx 20");

        assert_eq!(
            feedback.checkpoint(&offset("slot", 20, 0, 1)).unwrap(),
            None
        );
        assert_eq!(
            feedback.checkpoint(&offset("slot", 10, 0, 2)).unwrap(),
            None
        );
        assert_eq!(
            feedback.checkpoint(&offset("slot", 10, 1, 2)).unwrap(),
            Some(PgLsn::from_u64(20))
        );
        assert_eq!(feedback.watermark(), PgLsn::from_u64(20));
    }
}