archiver/
replay_guard.rs

1//! Protection against JetStream replaying messages the archiver already wrote.
2//!
3//! JetStream keeps a durable consumer's progress (delivered and ack floor sequences) in
4//! a small state file next to the stream. The server replaces that file by writing a
5//! temporary file and renaming it, without an fsync. After a power cut the file can be
6//! left empty or missing. The server then recovers the consumer with no progress, and
7//! with deliver policy `all` it delivers the whole stream again. A consumer whose state
8//! file is corrupt is dropped at recovery, and the archiver creates it again, also from
9//! the start of the stream.
10//!
11//! The archiver therefore keeps its own checkpoint per stream and subject: the highest
12//! stream sequence such that every message on the subject up to it has been acked, which
13//! means it is inside a closed and fsynced Parquet file. The value is the consumer's ack
14//! floor as reported by the server, saved with an fsync of the file and its directory, so
15//! it survives the same power cut. A message at or below the checkpoint is acked without
16//! being written again.
17//!
18//! The checkpoint is only trusted while it clearly belongs to the same stream: the stream
19//! name, its creation time and the subject must match, and the stream's last sequence must
20//! not be below the checkpoint (a recreated or rewound stream reuses sequence numbers).
21//! Otherwise it is ignored and every message is written, as without the guard.
22
23use serde::{Deserialize, Serialize};
24use std::{
25    fs, io,
26    io::Write,
27    path::{Path, PathBuf},
28};
29
30/// Saved archiver progress for one stream and subject.
31#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
32pub struct Checkpoint {
33    /// JetStream stream name.
34    pub stream: String,
35    /// Stream creation time in Unix nanoseconds; changes if the stream is recreated.
36    pub stream_created_unix_ns: i64,
37    /// Subject the channel consumer filters on.
38    pub subject: String,
39    /// Every message on `subject` with a stream sequence up to this one is archived.
40    pub archived_through_seq: u64,
41}
42
43/// Turns a stream name or subject into a safe file name part.
44///
45/// ASCII letters, digits, `-`, `_` and `.` are kept, anything else becomes `_`. A
46/// leading `.` (or an empty result) gets a `_` prefix so no hidden or special name is
47/// produced.
48fn file_token(raw: &str) -> String {
49    let mut out: String = raw
50        .chars()
51        .map(|c| {
52            if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' {
53                c
54            } else {
55                '_'
56            }
57        })
58        .collect();
59    if out.is_empty() || out.starts_with('.') {
60        out.insert(0, '_');
61    }
62    out
63}
64
65/// Path of the checkpoint file for one stream and subject.
66///
67/// # Examples
68///
69/// ```text
70/// checkpoint_path("/state", "labjack", "avenars.i69.mu1.lj2.live.ch08")
71///     -> "/state/labjack/avenars.i69.mu1.lj2.live.ch08.json"
72/// ```
73pub fn checkpoint_path(state_dir: &Path, stream: &str, subject: &str) -> PathBuf {
74    state_dir
75        .join(file_token(stream))
76        .join(format!("{}.json", file_token(subject)))
77}
78
79/// Reads a checkpoint file.
80///
81/// # Returns
82///
83/// `None` if the file does not exist.
84///
85/// # Errors
86///
87/// Returns an error if the file cannot be read or does not parse.
88pub fn load(path: &Path) -> io::Result<Option<Checkpoint>> {
89    match fs::read(path) {
90        Ok(bytes) => serde_json::from_slice(&bytes)
91            .map(Some)
92            .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)),
93        Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
94        Err(err) => Err(err),
95    }
96}
97
98/// Writes a checkpoint durably.
99///
100/// The JSON is written to a temporary file that is fsynced, renamed over the checkpoint
101/// and followed by an fsync of the directory, so after a power cut the file holds either
102/// the old or the new checkpoint, never an empty one.
103///
104/// # Errors
105///
106/// Returns an error if the directory cannot be created or any write, sync or rename
107/// fails.
108pub fn save(path: &Path, checkpoint: &Checkpoint) -> io::Result<()> {
109    let dir = path
110        .parent()
111        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "checkpoint has no parent"))?;
112    fs::create_dir_all(dir)?;
113    let tmp = path.with_extension("json.tmp");
114    {
115        let mut file = fs::File::create(&tmp)?;
116        file.write_all(&serde_json::to_vec_pretty(checkpoint)?)?;
117        file.sync_all()?;
118    }
119    fs::rename(&tmp, path)?;
120    fs::File::open(dir)?.sync_all()
121}
122
123/// Decides whether a saved checkpoint can be used for the stream as it is now.
124///
125/// # Arguments
126///
127/// * `checkpoint` - Saved checkpoint.
128/// * `stream` - Current stream name.
129/// * `stream_created_unix_ns` - Current stream creation time.
130/// * `stream_last_seq` - Current last sequence of the stream.
131/// * `subject` - Current channel subject.
132///
133/// # Returns
134///
135/// The trusted floor, or the reason the checkpoint is ignored.
136pub fn trusted_floor(
137    checkpoint: &Checkpoint,
138    stream: &str,
139    stream_created_unix_ns: i64,
140    stream_last_seq: u64,
141    subject: &str,
142) -> Result<u64, String> {
143    if checkpoint.stream != stream || checkpoint.subject != subject {
144        return Err(format!(
145            "it is for stream '{}' subject {}",
146            checkpoint.stream, checkpoint.subject
147        ));
148    }
149    if checkpoint.stream_created_unix_ns != stream_created_unix_ns {
150        return Err("the stream was recreated since it was written".to_string());
151    }
152    if checkpoint.archived_through_seq > stream_last_seq {
153        return Err(format!(
154            "it is at sequence {} but the stream ends at {}, so sequences may have been reused",
155            checkpoint.archived_through_seq, stream_last_seq
156        ));
157    }
158    Ok(checkpoint.archived_through_seq)
159}
160
161/// Skips messages at or below a trusted floor and counts them for logging.
162#[derive(Debug, Default)]
163pub struct ReplayGuard {
164    floor: u64,
165    skipped: u64,
166    first_skipped: u64,
167    last_skipped: u64,
168}
169
170/// A finished run of skipped messages, for one log line.
171#[derive(Debug, PartialEq, Eq)]
172pub struct SkipRun {
173    /// Number of messages skipped.
174    pub count: u64,
175    /// Lowest stream sequence skipped.
176    pub first: u64,
177    /// Highest stream sequence skipped.
178    pub last: u64,
179}
180
181impl ReplayGuard {
182    /// Creates a guard that skips stream sequences `1..=floor`. A floor of 0 skips
183    /// nothing.
184    pub fn new(floor: u64) -> Self {
185        Self {
186            floor,
187            ..Self::default()
188        }
189    }
190
191    /// Current floor.
192    pub fn floor(&self) -> u64 {
193        self.floor
194    }
195
196    /// Raises the floor; a lower value is ignored.
197    pub fn raise(&mut self, floor: u64) {
198        self.floor = self.floor.max(floor);
199    }
200
201    /// Returns `true` if the message at `stream_seq` is already archived and must be
202    /// acked without writing. Sequence 0 (unknown) is never skipped.
203    pub fn is_archived(&mut self, stream_seq: u64) -> bool {
204        if stream_seq == 0 || stream_seq > self.floor {
205            return false;
206        }
207        if self.skipped == 0 {
208            self.first_skipped = stream_seq;
209        }
210        self.first_skipped = self.first_skipped.min(stream_seq);
211        self.last_skipped = self.last_skipped.max(stream_seq);
212        self.skipped += 1;
213        true
214    }
215
216    /// Number of messages skipped in the current run.
217    pub fn skipped_in_run(&self) -> u64 {
218        self.skipped
219    }
220
221    /// Ends the current run of skipped messages and returns it, if there was one.
222    pub fn take_run(&mut self) -> Option<SkipRun> {
223        if self.skipped == 0 {
224            return None;
225        }
226        let run = SkipRun {
227            count: self.skipped,
228            first: self.first_skipped,
229            last: self.last_skipped,
230        };
231        self.skipped = 0;
232        self.first_skipped = 0;
233        self.last_skipped = 0;
234        Some(run)
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    fn cp(seq: u64) -> Checkpoint {
243        Checkpoint {
244            stream: "labjack".into(),
245            stream_created_unix_ns: 42,
246            subject: "avenars.s.b.src.live.ch08".into(),
247            archived_through_seq: seq,
248        }
249    }
250
251    #[test]
252    fn checkpoint_round_trips_and_missing_file_is_none() {
253        let dir = std::env::temp_dir().join(format!("replay-guard-{}", uuid::Uuid::new_v4()));
254        let path = checkpoint_path(&dir, "labjack", "avenars.s.b.src.live.ch08");
255        assert_eq!(load(&path).unwrap(), None);
256        save(&path, &cp(7)).unwrap();
257        save(&path, &cp(9)).unwrap();
258        assert_eq!(load(&path).unwrap(), Some(cp(9)));
259        assert!(!path.with_extension("json.tmp").exists());
260        fs::write(&path, b"").unwrap();
261        assert!(
262            load(&path).is_err(),
263            "an empty file is an error, not a zero floor"
264        );
265        fs::remove_dir_all(dir).unwrap();
266    }
267
268    #[test]
269    fn file_names_stay_inside_the_state_dir() {
270        let p = checkpoint_path(Path::new("/s"), "../x", "a/b.*.>");
271        assert_eq!(p, PathBuf::from("/s/_.._x/a_b._._.json"));
272        assert_eq!(file_token(""), "_");
273    }
274
275    #[test]
276    fn checkpoint_is_trusted_only_for_the_same_stream_generation() {
277        let s = "avenars.s.b.src.live.ch08";
278        assert_eq!(trusted_floor(&cp(100), "labjack", 42, 500, s), Ok(100));
279        assert_eq!(trusted_floor(&cp(100), "labjack", 42, 100, s), Ok(100));
280        assert!(trusted_floor(&cp(100), "other", 42, 500, s).is_err());
281        assert!(trusted_floor(&cp(100), "labjack", 42, 500, "x.ch09").is_err());
282        assert!(trusted_floor(&cp(100), "labjack", 43, 500, s).is_err());
283        assert!(trusted_floor(&cp(100), "labjack", 42, 99, s).is_err());
284    }
285
286    #[test]
287    fn guard_skips_up_to_the_floor_and_reports_runs() {
288        let mut g = ReplayGuard::new(10);
289        assert!(!g.is_archived(0));
290        assert!(g.is_archived(3));
291        assert!(g.is_archived(10));
292        assert!(!g.is_archived(11));
293        assert_eq!(
294            g.take_run(),
295            Some(SkipRun {
296                count: 2,
297                first: 3,
298                last: 10
299            })
300        );
301        assert_eq!(g.take_run(), None);
302        g.raise(5);
303        assert_eq!(g.floor(), 10);
304        g.raise(20);
305        assert!(g.is_archived(11));
306        assert!(!ReplayGuard::new(0).is_archived(1));
307    }
308}
archiver/replay_guard.rs

archiver/
replay_guard.rs
