kevy_embedded/
store_persist.rs

1//! Durability methods on [`Store`] — `BGREWRITEAOF` and `SAVE` — plus
2//! their per-shard helpers. Extracted from `store.rs` to keep that file
3//! under the 500-LOC project ceiling. Behaviour is unchanged from the
4//! pre-split layout; this module hosts long-running disk paths
5//! separately from the hot lock/dispatch surface in `store.rs`.
6
7use std::io;
8use std::sync::RwLock;
9use std::time::Instant;
10
11use kevy_persist::RewriteStats;
12
13use crate::metric::KevyMetric;
14use crate::store::{Inner, Store, lock_write};
15
16impl Store {
17    /// `BGREWRITEAOF`: rebuild every shard's AOF from current state.
18    /// Synchronous. Returns the summed stats (`None` if persistence is off /
19    /// no shard rewrote).
20    pub fn rewrite_aof(&self) -> io::Result<Option<RewriteStats>> {
21        let mut agg: Option<RewriteStats> = None;
22        for shard in self.shards.iter() {
23            let start = Instant::now();
24            // Phase 1 (locked): freeze the COW view + start the tee —
25            // O(n)-shallow, no serialization under the lock.
26            let (view, tmp, before_bytes) = {
27                let mut g = lock_write(shard);
28                let Inner { store, aof, .. } = &mut *g;
29                let Some(aof) = aof else { continue };
30                if aof.is_rewriting() {
31                    continue;
32                }
33                let before = aof.size_bytes();
34                let view = store.collect_snapshot();
35                (view, aof.begin_view_rewrite()?, before)
36            };
37            // Phase 2 (unlocked): serialize + fsync the compacted log.
38            let keys = match kevy_persist::dump_aof(&tmp, &view) {
39                Ok((keys, _)) => keys,
40                Err(e) => {
41                    let mut g = lock_write(shard);
42                    if let Some(aof) = &mut g.aof {
43                        aof.abort_concurrent_rewrite();
44                    }
45                    let _ = std::fs::remove_file(&tmp);
46                    return Err(e);
47                }
48            };
49            // Phase 3 (locked): append the tee'd diff and swap.
50            let mut g = lock_write(shard);
51            let Some(aof) = &mut g.aof else { continue };
52            let stats = match aof.finish_concurrent_rewrite(&tmp, keys) {
53                Ok(s) => s,
54                Err(e) => {
55                    aof.abort_concurrent_rewrite();
56                    let _ = std::fs::remove_file(&tmp);
57                    return Err(e);
58                }
59            };
60            if let Some(sink) = &self.config.metric_sink {
61                sink.emit(KevyMetric::Rewrite {
62                    keys: stats.keys,
63                    before_bytes,
64                    after_bytes: stats.bytes,
65                    elapsed_ms: start.elapsed().as_millis() as u64,
66                });
67            }
68            let acc = agg.get_or_insert(RewriteStats { keys: 0, bytes: 0 });
69            acc.keys += stats.keys;
70            acc.bytes += stats.bytes;
71        }
72        Ok(agg)
73    }
74
75    /// Snapshot every shard to its `dump-{i}.rdb` (single shard: the configured
76    /// name), atomically. `Ok(false)` when persistence is disabled.
77    pub fn save_snapshot(&self) -> io::Result<bool> {
78        let Some(dir) = self.config.data_dir.as_ref() else {
79            return Ok(false);
80        };
81        let n = self.shards.len();
82        for (i, shard) in self.shards.iter().enumerate() {
83            let name = if n == 1 {
84                self.config.snapshot_filename.clone()
85            } else {
86                kevy_persist::layout::snapshot_file(i)
87            };
88            save_shard_snapshot(shard, &dir.join(name))?;
89        }
90        Ok(true)
91    }
92}
93
94/// Save one shard's snapshot with the snapshot+log contract intact:
95/// after a successful save the AOF holds **only post-collect writes**,
96/// so a restart replays them over the snapshot without double-applying
97/// history (non-idempotent commands like RPUSH duplicated before this).
98///
99/// Phase 1 (write lock): freeze the COW view + start the AOF tee — no
100/// write may land between the two (the tee atomicity contract). Phase 2
101/// (unlocked): serialize the view to the snapshot's durable tmp.
102/// Phase 3 (write lock): commit — snapshot rename and tee'd AOF reset
103/// adjacent, so the snapshot/log commit window stays microseconds.
104pub(crate) fn save_shard_snapshot(
105    shard: &RwLock<Inner>,
106    path: &std::path::Path,
107) -> io::Result<()> {
108    let (view, reset_tmp) = freeze_for_save(shard)?;
109    let tmp = match kevy_persist::write_snapshot_tmp(&view, path) {
110        Ok(t) => t,
111        Err(e) => {
112            if reset_tmp.is_some()
113                && let Some(aof) = &mut lock_write(shard).aof
114            {
115                aof.abort_concurrent_rewrite();
116            }
117            return Err(e);
118        }
119    };
120    let mut g = lock_write(shard);
121    std::fs::rename(&tmp, path)?;
122    if let (Some(reset), Some(aof)) = (reset_tmp, &mut g.aof) {
123        let swap = kevy_persist::write_aof_base(&reset)
124            .and_then(|()| aof.finish_concurrent_rewrite(&reset, 0));
125        if let Err(e) = swap {
126            aof.abort_concurrent_rewrite();
127            let _ = std::fs::remove_file(&reset);
128            return Err(e);
129        }
130    }
131    Ok(())
132}
133
134/// Phase-1 helper: collect the view and start the tee under one write
135/// lock. A racing background auto-rewrite owns the tee; it runs its
136/// slow half off-lock and finishes in milliseconds, so wait it out
137/// (bounded) rather than saving a snapshot whose log would double-
138/// apply on replay.
139fn freeze_for_save(
140    shard: &RwLock<Inner>,
141) -> io::Result<(kevy_store::SnapshotView, Option<std::path::PathBuf>)> {
142    for _ in 0..2000 {
143        {
144            let mut g = lock_write(shard);
145            let Inner { store, aof, .. } = &mut *g;
146            match aof {
147                Some(a) if a.is_rewriting() => {} // racing rewrite — retry
148                Some(a) => {
149                    let view = store.collect_snapshot();
150                    return Ok((view, Some(a.begin_view_rewrite()?)));
151                }
152                None => return Ok((store.collect_snapshot(), None)),
153            }
154        }
155        std::thread::sleep(std::time::Duration::from_millis(5));
156    }
157    Err(io::Error::new(
158        io::ErrorKind::TimedOut,
159        "kevy-embedded: AOF rewrite still in flight after 10s; snapshot aborted",
160    ))
161}
kevy_embedded/store_persist.rs

kevy_embedded/
store_persist.rs
