armdb 0.14.2

//! Durable transaction-id ceiling.
//!
//! Ids are handed out from a reserved window; the window's upper bound is
//! persisted before any id inside it is used, so a restart — including one
//! after a checkpoint deleted every record — always continues above every id
//! that may already exist on disk.

use std::io::Write;
use std::path::{Path, PathBuf};

use crate::error::{DbError, DbResult};
use crate::io::direct::sync_dir;

const MAGIC: [u8; 4] = *b"ATXS";
const VERSION: u16 = 1;
const STATE_LEN: usize = 16; // magic | version u16 | pad u16 | reserved u64

#[derive(Debug)]
pub(crate) struct TxIdAllocator {
    dir: PathBuf,
    path: PathBuf,
    /// The next id to hand out.
    next: u64,
    /// Highest id that may already have been handed out (persisted).
    reserved: u64,
    window: u64,
}

impl TxIdAllocator {
    pub(crate) fn open(dir: &Path, window: u64) -> DbResult<Self> {
        let path = dir.join("state");
        let reserved = match std::fs::read(&path) {
            Ok(bytes) => decode_state(&bytes)?,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => 0,
            Err(e) => return Err(DbError::Io(e)),
        };
        Ok(Self {
            dir: dir.to_path_buf(),
            path,
            next: reserved + 1,
            reserved,
            window: window.max(1),
        })
    }

    pub(crate) fn next_id(&mut self) -> DbResult<u64> {
        if self.next > self.reserved {
            let new_ceiling = ceiling_for(self.next, self.window)?;
            self.persist(new_ceiling)?;
            self.reserved = new_ceiling;
        }
        let id = self.next;
        // Checked because this counter is what guarantees a `tx_id` is never
        // handed out twice, and an encrypted log uses `tx_id` as its GCM nonce:
        // wrapping here would repeat a nonce under the same key.
        self.next = id.checked_add(1).ok_or(EXHAUSTED)?;
        Ok(id)
    }

    /// Raise the floor above an id that is already on disk.
    ///
    /// `state` is a separate file from the segments, so a partial restore (or an
    /// over-eager cleanup) can lose the ceiling while sealed records survive.
    /// Reissuing an id a recovered record already carries would overwrite that
    /// record in the pending map — and, under encryption, repeat its nonce.
    pub(crate) fn ensure_above(&mut self, highest_on_disk: u64) -> DbResult<()> {
        if highest_on_disk < self.next {
            return Ok(());
        }
        let next = highest_on_disk.checked_add(1).ok_or(EXHAUSTED)?;
        let ceiling = ceiling_for(next, self.window)?;
        // Persist before moving `next`: a crash between the two must leave the
        // allocator no lower than it is now.
        self.persist(ceiling)?;
        self.reserved = ceiling;
        self.next = next;
        Ok(())
    }

    fn persist(&self, ceiling: u64) -> DbResult<()> {
        let mut bytes = Vec::with_capacity(STATE_LEN);
        bytes.extend_from_slice(&MAGIC);
        bytes.extend_from_slice(&VERSION.to_le_bytes());
        bytes.extend_from_slice(&0u16.to_le_bytes());
        bytes.extend_from_slice(&ceiling.to_le_bytes());

        // Same tmp + rename + directory fsync discipline as `engine.rs:36-46`.
        let tmp = self.path.with_extension("tmp");
        let mut f = std::fs::File::create(&tmp)?;
        f.write_all(&bytes)?;
        f.sync_data()?;
        drop(f);
        std::fs::rename(&tmp, &self.path)?;
        sync_dir(&self.dir)?;
        Ok(())
    }
}

const EXHAUSTED: DbError = DbError::Internal("txlog transaction ids are exhausted");

fn ceiling_for(next: u64, window: u64) -> DbResult<u64> {
    next.checked_add(window - 1).ok_or(EXHAUSTED)
}

fn decode_state(bytes: &[u8]) -> DbResult<u64> {
    if bytes.len() != STATE_LEN || bytes[0..4] != MAGIC {
        return Err(DbError::FormatMismatch(
            "txlog state file is corrupt".into(),
        ));
    }
    let version = u16::from_le_bytes([bytes[4], bytes[5]]);
    if version != VERSION {
        return Err(DbError::FormatMismatch(format!(
            "txlog state version {version}; this build writes {VERSION}"
        )));
    }
    Ok(u64::from_le_bytes([
        bytes[8], bytes[9], bytes[10], bytes[11], bytes[12], bytes[13], bytes[14], bytes[15],
    ]))
}

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

    #[test]
    fn txlog_state_ids_are_monotonic_within_a_session() {
        let dir = tempfile::tempdir().unwrap();
        let mut alloc = TxIdAllocator::open(dir.path(), 4096).unwrap();
        let ids: Vec<u64> = (0..5).map(|_| alloc.next_id().unwrap()).collect();
        assert_eq!(ids, vec![1, 2, 3, 4, 5]);
    }

    #[test]
    fn txlog_state_never_reuses_an_id_after_restart() {
        let dir = tempfile::tempdir().unwrap();
        let last = {
            let mut alloc = TxIdAllocator::open(dir.path(), 4096).unwrap();
            let mut last = 0;
            for _ in 0..3 {
                last = alloc.next_id().unwrap();
            }
            last
        };
        // Restart with no records at all — the situation a checkpoint leaves.
        let mut alloc = TxIdAllocator::open(dir.path(), 4096).unwrap();
        let next = alloc.next_id().unwrap();
        assert!(
            next > last + 4090,
            "next id {next} must clear the reserved window, last handed out {last}"
        );
    }

    #[test]
    fn txlog_state_floor_rises_above_a_recovered_id() {
        let dir = tempfile::tempdir().unwrap();
        let mut alloc = TxIdAllocator::open(dir.path(), 16).unwrap();
        // Segments survived a restore that lost `state`: tx 40 is already sealed.
        alloc.ensure_above(40).unwrap();
        let next = alloc.next_id().unwrap();
        assert!(
            next > 40,
            "an id already on disk must never be handed out again, got {next}"
        );

        // The raised floor is durable, not just in memory.
        drop(alloc);
        let mut alloc = TxIdAllocator::open(dir.path(), 16).unwrap();
        assert!(alloc.next_id().unwrap() > next);
    }

    #[test]
    fn txlog_state_floor_never_moves_backwards() {
        let dir = tempfile::tempdir().unwrap();
        let mut alloc = TxIdAllocator::open(dir.path(), 4096).unwrap();
        let first = alloc.next_id().unwrap();
        alloc.ensure_above(0).unwrap();
        assert_eq!(
            alloc.next_id().unwrap(),
            first + 1,
            "a floor below the allocator's position changes nothing"
        );
    }

    #[test]
    fn txlog_state_survives_a_torn_state_file() {
        let dir = tempfile::tempdir().unwrap();
        {
            let mut alloc = TxIdAllocator::open(dir.path(), 16).unwrap();
            alloc.next_id().unwrap();
        }
        std::fs::write(dir.path().join("state"), b"junk").unwrap();
        let err =
            TxIdAllocator::open(dir.path(), 16).expect_err("a corrupt state file fails closed");
        assert!(format!("{err}").contains("txlog state"), "{err}");
    }
}