armdb 0.11.2

//! One transaction-log segment file plus the directory-durability discipline
//! around it: create → fsync file → fsync dir, and unlink → fsync dir.

use std::fs::{File, OpenOptions};
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
#[cfg(test)]
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use crate::error::{DbError, DbResult};
use crate::io::direct::{fsync, sync_dir};
use crate::txlog::record::{Decoded, TxRecord, decode_record_with};

/// Write only this many bytes and then fail (tests only).
#[cfg(test)]
pub(crate) static FAIL_APPEND_AFTER: AtomicUsize = AtomicUsize::new(usize::MAX);
/// Make the truncate-back repair fail (tests only).
#[cfg(test)]
pub(crate) static FAIL_TRUNCATE: AtomicBool = AtomicBool::new(false);
/// Make `Segment::sync` fail (tests only).
#[cfg(test)]
pub(crate) static FAIL_SYNC: AtomicBool = AtomicBool::new(false);
/// Counts successful `Segment::sync` calls (tests only).
#[cfg(test)]
pub(crate) static SEGMENT_SYNCS: AtomicUsize = AtomicUsize::new(0);
/// Make the file/directory fsync inside `Segment::create` fail (tests only).
#[cfg(test)]
pub(crate) static FAIL_CREATE_SYNC: AtomicBool = AtomicBool::new(false);
/// Make `Segment::remove` fail for this segment id (tests only); `u32::MAX`
/// disables it.
#[cfg(test)]
pub(crate) static FAIL_REMOVE_SEGMENT: std::sync::atomic::AtomicU32 =
    std::sync::atomic::AtomicU32::new(u32::MAX);

pub(crate) struct Segment {
    id: u32,
    path: PathBuf,
    pub(crate) file: File,
    len: u64,
}

pub(crate) enum AppendOutcome {
    Written {
        #[allow(dead_code)] // physical offset; part 2 may use for diagnostics
        offset: u64,
    },
    /// The write failed and the segment was repaired back to its previous end.
    FailedAndTruncated(DbError),
    /// The write failed *and* the repair failed — the log must reject further
    /// commits until the database is reopened.
    FailedAndPoisoned(DbError),
}

pub(crate) struct ScanResult {
    pub records: Vec<TxRecord>,
    /// Offset of the first byte that did not decode, if any.
    pub torn_at: Option<u64>,
}

pub(crate) fn segment_path(dir: &Path, id: u32) -> PathBuf {
    dir.join(format!("{id:06}.log"))
}

pub(crate) fn list_segment_ids(dir: &Path) -> DbResult<Vec<u32>> {
    let mut ids = Vec::new();
    for entry in std::fs::read_dir(dir)? {
        let entry = entry?;
        let name = entry.file_name();
        let Some(name) = name.to_str() else { continue };
        let Some(stem) = name.strip_suffix(".log") else {
            continue;
        };
        if let Ok(id) = stem.parse::<u32>() {
            ids.push(id);
        }
    }
    ids.sort_unstable();
    Ok(ids)
}

impl Segment {
    pub(crate) fn create(dir: &Path, id: u32) -> DbResult<Self> {
        let path = segment_path(dir, id);
        let file = OpenOptions::new()
            .create_new(true)
            .read(true)
            .write(true)
            .open(&path)?;
        // The file must exist durably before any record in it counts as
        // committed: fsync the file, then its directory entry.
        //
        // If either step fails the file has to go. Segment ids come from the
        // active segment, so the next rotation asks for this same id, gets
        // `AlreadyExists` from `create_new`, and every further append and
        // checkpoint fails until the database is reopened. The file is not
        // durable and holds no record, so removing it loses nothing.
        let durable = (|| -> DbResult<()> {
            #[cfg(test)]
            if FAIL_CREATE_SYNC.load(Ordering::SeqCst) {
                return Err(DbError::Io(std::io::Error::other(
                    "injected create sync failure",
                )));
            }
            fsync(&file)?;
            sync_dir(dir)
        })();
        if let Err(e) = durable {
            drop(file);
            if let Err(cleanup) = std::fs::remove_file(&path) {
                tracing::error!(
                    segment = id,
                    "removing an undurable txlog segment failed: {cleanup}"
                );
            }
            return Err(e);
        }
        Ok(Self {
            id,
            path,
            file,
            len: 0,
        })
    }

    pub(crate) fn open(dir: &Path, id: u32) -> DbResult<Self> {
        let path = segment_path(dir, id);
        let file = OpenOptions::new().read(true).write(true).open(&path)?;
        let len = file.metadata()?.len();
        Ok(Self {
            id,
            path,
            file,
            len,
        })
    }

    pub(crate) fn id(&self) -> u32 {
        self.id
    }

    /// End of the last complete record: where the next append lands.
    /// Deliberately not called `len` — `clippy::len_without_is_empty` would
    /// then demand an `is_empty` that means nothing here.
    pub(crate) fn end_offset(&self) -> u64 {
        self.len
    }

    pub(crate) fn append(&mut self, bytes: &[u8]) -> DbResult<AppendOutcome> {
        let offset = self.len;
        match self.write_all_at(bytes, offset) {
            Ok(()) => {
                self.len += bytes.len() as u64;
                Ok(AppendOutcome::Written { offset })
            }
            Err(e) => match self.truncate_back(offset) {
                Ok(()) => Ok(AppendOutcome::FailedAndTruncated(DbError::Io(e))),
                Err(repair) => {
                    tracing::error!(
                        segment = self.id,
                        "txlog append failed and the repair failed: {repair}"
                    );
                    Ok(AppendOutcome::FailedAndPoisoned(DbError::Io(e)))
                }
            },
        }
    }

    fn write_all_at(&self, buf: &[u8], offset: u64) -> std::io::Result<()> {
        #[cfg(test)]
        {
            let limit = FAIL_APPEND_AFTER.load(Ordering::SeqCst);
            if limit < buf.len() {
                self.file.write_all_at(&buf[..limit], offset)?;
                return Err(std::io::Error::other("injected txlog append failure"));
            }
        }
        self.file.write_all_at(buf, offset)
    }

    /// Truncate the segment file back to `offset` and fsync (repair torn tail).
    pub(crate) fn truncate_to(&mut self, offset: u64) -> DbResult<()> {
        self.truncate_back(offset)
    }

    fn truncate_back(&mut self, offset: u64) -> DbResult<()> {
        #[cfg(test)]
        if FAIL_TRUNCATE.load(Ordering::SeqCst) {
            return Err(DbError::Io(std::io::Error::other(
                "injected txlog truncate failure",
            )));
        }
        self.file.set_len(offset)?;
        fsync(&self.file)?;
        self.len = offset;
        Ok(())
    }

    pub(crate) fn sync(&self) -> DbResult<()> {
        #[cfg(test)]
        if FAIL_SYNC.load(Ordering::SeqCst) {
            return Err(DbError::Io(std::io::Error::other(
                "injected txlog sync failure",
            )));
        }
        fsync(&self.file)?;
        #[cfg(test)]
        SEGMENT_SYNCS.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }

    pub(crate) fn scan(&self, cipher: Option<&crate::txlog::LogCipher>) -> DbResult<ScanResult> {
        let mut buf = vec![0u8; self.len as usize];
        self.file.read_exact_at(&mut buf, 0)?;
        let mut records = Vec::new();
        let mut at = 0usize;
        while at < buf.len() {
            match decode_record_with(&buf[at..], cipher)? {
                Decoded::Record { record, len } => {
                    records.push(record);
                    at += len;
                }
                Decoded::Torn => {
                    return Ok(ScanResult {
                        records,
                        torn_at: Some(at as u64),
                    });
                }
            }
        }
        Ok(ScanResult {
            records,
            torn_at: None,
        })
    }

    pub(crate) fn remove(self, dir: &Path) -> DbResult<()> {
        #[cfg(test)]
        if FAIL_REMOVE_SEGMENT.load(Ordering::SeqCst) == self.id {
            return Err(DbError::Io(std::io::Error::other(
                "injected segment remove failure",
            )));
        }
        let path = self.path.clone();
        drop(self);
        std::fs::remove_file(&path)?;
        sync_dir(dir)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::os::unix::fs::FileExt;
    use std::sync::atomic::Ordering;

    use super::*;
    use crate::txlog::record::{OpKind, TxOp, TxRecord, VersionTag, encode_record};

    fn rec(tx_id: u64) -> TxRecord {
        TxRecord {
            tx_id,
            ops: vec![TxOp {
                collection_id: 1,
                name: "t".to_string(),
                schema_fingerprint: 42,
                key: tx_id.to_le_bytes().to_vec(),
                kind: OpKind::Put,
                value: vec![7; 16],
                version: VersionTag::Bitcask { gsn: tx_id },
            }],
        }
    }

    fn encoded(tx_id: u64) -> Vec<u8> {
        let mut b = Vec::new();
        encode_record(&rec(tx_id), &mut b).unwrap();
        b
    }

    #[test]
    fn txlog_segment_append_then_scan_round_trips() {
        let dir = tempfile::tempdir().unwrap();
        let mut seg = Segment::create(dir.path(), 1).unwrap();
        for id in 1..=3 {
            assert!(matches!(
                seg.append(&encoded(id)).unwrap(),
                AppendOutcome::Written { .. }
            ));
        }
        let scanned = seg.scan(None).unwrap();
        assert_eq!(
            scanned.records.iter().map(|r| r.tx_id).collect::<Vec<_>>(),
            vec![1, 2, 3]
        );
        assert!(scanned.torn_at.is_none());
    }

    #[test]
    fn txlog_segment_reopen_sees_prior_records_and_appends_after_them() {
        let dir = tempfile::tempdir().unwrap();
        {
            let mut seg = Segment::create(dir.path(), 1).unwrap();
            seg.append(&encoded(1)).unwrap();
        }
        let mut seg = Segment::open(dir.path(), 1).unwrap();
        seg.append(&encoded(2)).unwrap();
        let ids: Vec<u64> = seg
            .scan(None)
            .unwrap()
            .records
            .iter()
            .map(|r| r.tx_id)
            .collect();
        assert_eq!(ids, vec![1, 2]);
    }

    #[test]
    fn txlog_segment_partial_append_is_truncated_and_the_next_append_is_clean() {
        let dir = tempfile::tempdir().unwrap();
        let mut seg = Segment::create(dir.path(), 1).unwrap();
        seg.append(&encoded(1)).unwrap();
        let good_len = seg.end_offset();

        FAIL_APPEND_AFTER.store(10, Ordering::SeqCst);
        let outcome = seg.append(&encoded(2)).unwrap();
        FAIL_APPEND_AFTER.store(usize::MAX, Ordering::SeqCst);
        assert!(matches!(outcome, AppendOutcome::FailedAndTruncated(_)));
        assert_eq!(seg.end_offset(), good_len, "the torn bytes must be gone");

        seg.append(&encoded(3)).unwrap();
        let ids: Vec<u64> = seg
            .scan(None)
            .unwrap()
            .records
            .iter()
            .map(|r| r.tx_id)
            .collect();
        assert_eq!(ids, vec![1, 3], "no torn record between the two good ones");
    }

    #[test]
    fn txlog_segment_reports_a_torn_tail() {
        let dir = tempfile::tempdir().unwrap();
        let mut seg = Segment::create(dir.path(), 1).unwrap();
        seg.append(&encoded(1)).unwrap();
        let good_len = seg.end_offset();
        // Simulate a crash mid-append: raw bytes with no marker.
        seg.file.write_all_at(&[0xAB; 40], good_len).unwrap();

        let seg = Segment::open(dir.path(), 1).unwrap();
        let scanned = seg.scan(None).unwrap();
        assert_eq!(scanned.records.len(), 1);
        assert_eq!(scanned.torn_at, Some(good_len));
    }

    #[test]
    fn txlog_segment_failed_truncate_poisons() {
        let dir = tempfile::tempdir().unwrap();
        let mut seg = Segment::create(dir.path(), 1).unwrap();
        FAIL_APPEND_AFTER.store(10, Ordering::SeqCst);
        FAIL_TRUNCATE.store(true, Ordering::SeqCst);
        let outcome = seg.append(&encoded(1)).unwrap();
        FAIL_APPEND_AFTER.store(usize::MAX, Ordering::SeqCst);
        FAIL_TRUNCATE.store(false, Ordering::SeqCst);
        assert!(matches!(outcome, AppendOutcome::FailedAndPoisoned(_)));
    }

    #[test]
    fn txlog_segment_ids_are_listed_in_order() {
        let dir = tempfile::tempdir().unwrap();
        for id in [3u32, 1, 2] {
            Segment::create(dir.path(), id).unwrap();
        }
        std::fs::write(dir.path().join("state"), b"not a segment").unwrap();
        assert_eq!(list_segment_ids(dir.path()).unwrap(), vec![1, 2, 3]);
    }

    /// A segment that never became durable must not stay on disk: the id is
    /// derived from the active segment, so the next rotation asks for the same
    /// one, `create_new` reports `AlreadyExists`, and every further append and
    /// checkpoint fails until the database is reopened.
    #[test]
    fn txlog_segment_create_removes_the_file_when_it_cannot_be_made_durable() {
        let dir = tempfile::tempdir().unwrap();
        FAIL_CREATE_SYNC.store(true, Ordering::SeqCst);
        let err = match Segment::create(dir.path(), 7) {
            Err(e) => e,
            Ok(_) => panic!("create must fail when it cannot be made durable"),
        };
        FAIL_CREATE_SYNC.store(false, Ordering::SeqCst);
        assert!(
            format!("{err}").contains("injected create sync failure"),
            "{err}"
        );

        assert!(
            !segment_path(dir.path(), 7).exists(),
            "an undurable segment must be removed, not left behind"
        );
        Segment::create(dir.path(), 7).expect("the id must be free for the next rotation");
    }
}