rivet/state/
keyset_range.rs

1//! Parallel-keyset crash-recovery ranges (`keyset_range`), feat/parallel-keyset
2//! iteration 2.
3//!
4//! A parallel keyset run partitions the key into N stable ROW-percentile ranges
5//! sampled once at open. Recovery is **per-range** (coarse): the boundaries are
6//! persisted here at open (keyed by run_id) so a resume reloads the SAME ranges
7//! rather than re-sampling a possibly-changed table (which would move the
8//! boundaries and leave a gap). Each worker flips only its OWN `(export_name,
9//! range_index)` row to `done=1` at completion — in the SAME transaction that
10//! records its parts to `file_log` — so the checkpoint is atomic: a crash before
11//! that commit leaves the range `done=0` (re-read from `lo` on resume, its parts
12//! overwritten by stable run_id-based names), a crash after it leaves the range
13//! `done=1` (skipped on resume, its parts rehydrated from `file_log`).
14
15use crate::error::Result;
16
17use super::{StateRef, StateStore, open_connection, pg_sql};
18
19/// One persisted range of a parallel keyset run.
20#[derive(Debug, Clone)]
21pub struct KeysetRangeRow {
22    pub range_index: i64,
23    /// Exclusive lower bound (`None` = -∞, the first range).
24    pub lo: Option<String>,
25    /// Inclusive upper bound (`None` = +∞, the last range).
26    pub hi: Option<String>,
27    pub done: bool,
28}
29
30/// A part a worker committed for its range — the `file_log` payload written
31/// atomically with the range's `done` flip.
32pub struct KeysetRangePart {
33    pub file_name: String,
34    pub rows: i64,
35    pub bytes: i64,
36}
37
38impl StateStore {
39    /// Persist the sampled ranges of a FRESH parallel keyset run (all `done=0`),
40    /// replacing any prior set for this export. Called once at open, before the
41    /// workers spawn; the boundaries are then immutable for the run's lifetime.
42    pub fn persist_keyset_ranges(
43        &self,
44        export_name: &str,
45        run_id: &str,
46        ranges: &[(Option<String>, Option<String>)],
47    ) -> Result<()> {
48        let now = chrono::Utc::now().to_rfc3339();
49        self.execute(
50            "DELETE FROM keyset_range WHERE export_name = ?1",
51            &[export_name.into()],
52        )?;
53        for (idx, (lo, hi)) in ranges.iter().enumerate() {
54            self.execute(
55                "INSERT INTO keyset_range \
56                 (export_name, run_id, range_index, lo, hi, done, updated_at) \
57                 VALUES (?1, ?2, ?3, ?4, ?5, 0, ?6)",
58                &[
59                    export_name.into(),
60                    run_id.into(),
61                    (idx as i64).into(),
62                    lo.clone().into(),
63                    hi.clone().into(),
64                    now.as_str().into(),
65                ],
66            )?;
67        }
68        Ok(())
69    }
70
71    /// Load the persisted ranges for a resuming run, ordered by `range_index`.
72    /// Filtered by `run_id` so a stale set from a superseded run is ignored.
73    pub fn load_keyset_ranges(
74        &self,
75        export_name: &str,
76        run_id: &str,
77    ) -> Result<Vec<KeysetRangeRow>> {
78        let sql = "SELECT range_index, lo, hi, done FROM keyset_range \
79                   WHERE export_name = ?1 AND run_id = ?2 ORDER BY range_index";
80        self.query(sql, &[export_name.into(), run_id.into()], |r| {
81            KeysetRangeRow {
82                range_index: r.i64(0),
83                lo: r.opt_text(1),
84                hi: r.opt_text(2),
85                done: r.i64(3) != 0,
86            }
87        })
88    }
89
90    /// Clear an export's persisted ranges. Called post-finalize (via
91    /// `finalize_keyset_anchor`) once the complete manifest is written — a no-op
92    /// for any export that never ran parallel keyset.
93    pub fn clear_keyset_ranges(&self, export_name: &str) -> Result<()> {
94        self.execute(
95            "DELETE FROM keyset_range WHERE export_name = ?1",
96            &[export_name.into()],
97        )?;
98        Ok(())
99    }
100
101    /// Atomically COMMIT a completed range from a worker thread: record every part
102    /// to `file_log` AND flip the range's `done=1`, in ONE transaction, over a
103    /// fresh connection (the live `StateStore` is not `Sync`, so workers reconnect
104    /// per the ADR-0011 `*_at_ref` pattern). The transaction is the checkpoint
105    /// boundary — a crash before it commits leaves the range `done=0` with no
106    /// `file_log` rows (re-read on resume), after it leaves both durable (skipped +
107    /// rehydrated on resume). Only the range's OWN row is touched, so concurrent
108    /// workers never contend.
109    pub fn commit_keyset_range_at_ref(
110        state_ref: &StateRef,
111        run_id: &str,
112        export_name: &str,
113        range_index: i64,
114        parts: &[KeysetRangePart],
115        format: &str,
116        compression: Option<&str>,
117    ) -> Result<()> {
118        let now = chrono::Utc::now().to_rfc3339();
119        let file_sql = "INSERT INTO file_log \
120             (run_id, export_name, file_name, row_count, bytes, format, compression, created_at) \
121             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)";
122        // Scope the `done` flip by run_id (the table's PK is (export_name,
123        // range_index) — run_id is NOT in it), so a run can never flip a row
124        // belonging to ANOTHER run's recovery set left behind by a crash (H1).
125        let done_sql = "UPDATE keyset_range SET done = 1, updated_at = ?1 \
126             WHERE export_name = ?2 AND range_index = ?3 AND run_id = ?4";
127        match state_ref {
128            StateRef::Sqlite(db_path) => {
129                let mut conn = open_connection(db_path)?;
130                let tx = conn.transaction()?;
131                for p in parts {
132                    tx.execute(
133                        file_sql,
134                        rusqlite::params![
135                            run_id,
136                            export_name,
137                            p.file_name,
138                            p.rows,
139                            p.bytes,
140                            format,
141                            compression,
142                            now
143                        ],
144                    )?;
145                }
146                tx.execute(
147                    done_sql,
148                    rusqlite::params![now, export_name, range_index, run_id],
149                )?;
150                tx.commit()?;
151            }
152            StateRef::Postgres(url) => {
153                let mut client = super::connect_pg(url)?;
154                let mut tx = client.transaction()?;
155                for p in parts {
156                    tx.execute(
157                        &pg_sql(file_sql),
158                        &[
159                            &run_id,
160                            &export_name,
161                            &p.file_name,
162                            &p.rows,
163                            &p.bytes,
164                            &format,
165                            &compression,
166                            &now,
167                        ],
168                    )?;
169                }
170                tx.execute(
171                    &pg_sql(done_sql),
172                    &[&now, &export_name, &range_index, &run_id],
173                )?;
174                tx.commit()?;
175            }
176        }
177        Ok(())
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn store() -> StateStore {
186        StateStore::open_in_memory().expect("in-memory store")
187    }
188
189    #[test]
190    fn persist_then_load_round_trips_ranges_all_not_done() {
191        let s = store();
192        let ranges = vec![
193            (None, Some("k0500".to_string())),
194            (Some("k0500".to_string()), Some("k1000".to_string())),
195            (Some("k1000".to_string()), None),
196        ];
197        s.persist_keyset_ranges("exp", "run-1", &ranges).unwrap();
198        let loaded = s.load_keyset_ranges("exp", "run-1").unwrap();
199        assert_eq!(loaded.len(), 3);
200        assert_eq!(loaded[0].lo, None);
201        assert_eq!(loaded[0].hi.as_deref(), Some("k0500"));
202        assert_eq!(loaded[2].hi, None);
203        assert!(loaded.iter().all(|r| !r.done), "fresh ranges are not done");
204    }
205
206    #[test]
207    fn load_with_a_different_run_id_returns_nothing() {
208        // A stale set from a superseded run must not leak into a fresh run's load.
209        let s = store();
210        s.persist_keyset_ranges("exp", "run-1", &[(None, None)])
211            .unwrap();
212        assert!(s.load_keyset_ranges("exp", "run-2").unwrap().is_empty());
213    }
214
215    #[test]
216    fn commit_range_marks_only_its_own_row_done_and_records_parts() {
217        // `_at_ref` reconnects (workers are not `Sync`), so an in-memory store won't
218        // do: `open_connection(":memory:")` is a DISTINCT empty DB. Use a file store.
219        let dir = tempfile::tempdir().unwrap();
220        let cfg = dir.path().join("rivet.yaml");
221        std::fs::write(&cfg, "# test").unwrap();
222        let s = StateStore::open(cfg.to_str().unwrap()).unwrap();
223        let ranges = vec![
224            (None, Some("k5".to_string())),
225            (Some("k5".to_string()), None),
226        ];
227        s.persist_keyset_ranges("exp", "run-1", &ranges).unwrap();
228
229        StateStore::commit_keyset_range_at_ref(
230            s.state_ref(),
231            "run-1",
232            "exp",
233            1,
234            &[KeysetRangePart {
235                file_name: "exp_run-1_pk_w1_0.parquet".to_string(),
236                rows: 42,
237                bytes: 100,
238            }],
239            "parquet",
240            Some("zstd"),
241        )
242        .unwrap();
243
244        let loaded = s.load_keyset_ranges("exp", "run-1").unwrap();
245        assert!(!loaded[0].done, "range 0 untouched");
246        assert!(loaded[1].done, "range 1 committed → done");
247        // The part landed in file_log under the run_id (rehydration source).
248        let files = s.list_files_for_run("run-1").unwrap();
249        assert_eq!(files.len(), 1);
250        assert_eq!(files[0].file_name, "exp_run-1_pk_w1_0.parquet");
251        assert_eq!(files[0].row_count, 42);
252    }
253
254    #[test]
255    fn clear_removes_all_ranges_for_the_export() {
256        let s = store();
257        s.persist_keyset_ranges("exp", "run-1", &[(None, None)])
258            .unwrap();
259        s.clear_keyset_ranges("exp").unwrap();
260        assert!(s.load_keyset_ranges("exp", "run-1").unwrap().is_empty());
261    }
262}
rivet/state/keyset_range.rs

rivet/state/
keyset_range.rs
