kevy_embedded/

store_persist.rs

//! Durability methods on [`Store`] — `BGREWRITEAOF` and `SAVE` — plus
//! their per-shard helpers. Extracted from `store.rs` to keep that file
//! under the 500-LOC project ceiling. Behaviour is unchanged from the
//! pre-split layout; this module hosts long-running disk paths
//! separately from the hot lock/dispatch surface in `store.rs`.

use std::io;
use std::sync::RwLock;
use std::time::Instant;

use kevy_persist::RewriteStats;

use crate::metric::KevyMetric;
use crate::store::{Inner, Store, lock_write};

impl Store {
    /// Durability barrier (v2.1): flush + `fdatasync` every shard's
    /// AOF now, regardless of the configured `appendfsync` policy. On
    /// `Ok(())`, every write acknowledged before this call is on
    /// stable storage. The `EverySec` serving-store idiom:
    ///
    /// ```ignore
    /// store.atomic(|c| { /* critical write */ Ok(()) })?;
    /// store.fsync_aof()?; // durable-on-ack for THIS block only
    /// ```
    ///
    /// Cost: one `fdatasync` per dirty shard; a no-op on clean shards.
    /// Under `appendfsync = always` it is a no-op (already durable).
    pub fn fsync_aof(&self) -> io::Result<()> {
        for shard in self.shards.iter() {
            let mut g = lock_write(shard);
            if let Some(aof) = &mut g.aof {
                aof.sync_now()?;
            }
        }
        Ok(())
    }

    /// `BGREWRITEAOF`: rebuild every shard's AOF from current state.
    /// Synchronous (despite the RESP name). Returns the summed stats
    /// (`None` if persistence is off / no shard rewrote). Shards mid
    /// rewrite are skipped. Emits [`KevyMetric::Rewrite`] per shard.
    pub fn rewrite_aof(&self) -> io::Result<Option<RewriteStats>> {
        let mut agg: Option<RewriteStats> = None;
        for shard in self.shards.iter() {
            if let Some(stats) = self.rewrite_one_shard(shard)? {
                let acc = agg.get_or_insert(RewriteStats { keys: 0, bytes: 0 });
                acc.keys += stats.keys;
                acc.bytes += stats.bytes;
            }
        }
        Ok(agg)
    }

    /// One shard's synchronous three-phase rewrite. `Ok(None)` =
    /// skipped (persistence off / already mid-rewrite).
    fn rewrite_one_shard(&self, shard: &RwLock<Inner>) -> io::Result<Option<RewriteStats>> {
        let start = Instant::now();
        // Phase 1 (locked): freeze the COW view + start the tee —
        // O(n)-shallow, no serialization under the lock.
        let (view, tmp, before_bytes) = {
            let mut g = lock_write(shard);
            let Inner { store, aof, .. } = &mut *g;
            let Some(aof) = aof else { return Ok(None) };
            if aof.is_rewriting() {
                return Ok(None);
            }
            let before = aof.size_bytes();
            let view = store.collect_snapshot();
            (view, aof.begin_view_rewrite()?, before)
        };
        // Phase 2 (unlocked): serialize + fsync the compacted log.
        let keys = match kevy_persist::dump_aof(&tmp, &view) {
            Ok((keys, _)) => keys,
            Err(e) => {
                let mut g = lock_write(shard);
                if let Some(aof) = &mut g.aof {
                    aof.abort_concurrent_rewrite();
                }
                let _ = std::fs::remove_file(&tmp);
                return Err(e);
            }
        };
        // Phase 3 (locked): append the tee'd diff and swap.
        let mut g = lock_write(shard);
        let Some(aof) = &mut g.aof else { return Ok(None) };
        let stats = match aof.finish_concurrent_rewrite(&tmp, keys) {
            Ok(s) => s,
            Err(e) => {
                aof.abort_concurrent_rewrite();
                let _ = std::fs::remove_file(&tmp);
                return Err(e);
            }
        };
        if let Some(sink) = &self.config.metric_sink {
            sink.emit(KevyMetric::Rewrite {
                keys: stats.keys,
                before_bytes,
                after_bytes: stats.bytes,
                elapsed_ms: start.elapsed().as_millis() as u64,
            });
        }
        Ok(Some(stats))
    }

    /// Snapshot every shard to its `dump-{i}.rdb` (single shard: the configured
    /// name), atomically. `Ok(false)` when persistence is disabled.
    pub fn save_snapshot(&self) -> io::Result<bool> {
        let Some(dir) = self.config.data_dir.as_ref() else {
            return Ok(false);
        };
        let n = self.shards.len();
        for (i, shard) in self.shards.iter().enumerate() {
            let name = if n == 1 {
                self.config.snapshot_filename.clone()
            } else {
                kevy_persist::layout::snapshot_file(i)
            };
            save_shard_snapshot(shard, &dir.join(name))?;
        }
        Ok(true)
    }
}

/// Save one shard's snapshot with the snapshot+log contract intact:
/// after a successful save the AOF holds **only post-collect writes**,
/// so a restart replays them over the snapshot without double-applying
/// history (non-idempotent commands like RPUSH duplicated before this).
///
/// Phase 1 (write lock): freeze the COW view + start the AOF tee — no
/// write may land between the two (the tee atomicity contract). Phase 2
/// (unlocked): serialize the view to the snapshot's durable tmp.
/// Phase 3 (write lock): commit — snapshot rename and tee'd AOF reset
/// adjacent, so the snapshot/log commit window stays microseconds.
pub(crate) fn save_shard_snapshot(
    shard: &RwLock<Inner>,
    path: &std::path::Path,
) -> io::Result<()> {
    let (view, reset_tmp) = freeze_for_save(shard)?;
    let tmp = match kevy_persist::write_snapshot_tmp(&view, path) {
        Ok(t) => t,
        Err(e) => {
            if reset_tmp.is_some()
                && let Some(aof) = &mut lock_write(shard).aof
            {
                aof.abort_concurrent_rewrite();
            }
            return Err(e);
        }
    };
    let mut g = lock_write(shard);
    std::fs::rename(&tmp, path)?;
    if let (Some(reset), Some(aof)) = (reset_tmp, &mut g.aof) {
        let swap = kevy_persist::write_aof_base(&reset)
            .and_then(|()| aof.finish_concurrent_rewrite(&reset, 0));
        if let Err(e) = swap {
            aof.abort_concurrent_rewrite();
            let _ = std::fs::remove_file(&reset);
            return Err(e);
        }
    }
    Ok(())
}

/// Phase-1 helper: collect the view and start the tee under one write
/// lock. A racing background auto-rewrite owns the tee; it runs its
/// slow half off-lock and finishes in milliseconds, so wait it out
/// (bounded) rather than saving a snapshot whose log would double-
/// apply on replay.
fn freeze_for_save(
    shard: &RwLock<Inner>,
) -> io::Result<(kevy_store::SnapshotView, Option<std::path::PathBuf>)> {
    for _ in 0..2000 {
        {
            let mut g = lock_write(shard);
            let Inner { store, aof, .. } = &mut *g;
            match aof {
                Some(a) if a.is_rewriting() => {} // racing rewrite — retry
                Some(a) => {
                    let view = store.collect_snapshot();
                    return Ok((view, Some(a.begin_view_rewrite()?)));
                }
                None => return Ok((store.collect_snapshot(), None)),
            }
        }
        std::thread::sleep(std::time::Duration::from_millis(5));
    }
    Err(io::Error::new(
        io::ErrorKind::TimedOut,
        "kevy-embedded: AOF rewrite still in flight after 10s; snapshot aborted",
    ))
}