chio_store_sqlite/budget_store/
replication.rs

1use super::*;
2
3/// Quorum-witness identity of a stored budget mutation event: its `event_seq`,
4/// origin `authority_id`, and origin `lease_epoch`.
5pub type BudgetEventWitness = (u64, Option<String>, Option<u64>);
6
7/// Initialize the replication sequence counter from existing rows on first open.
8///
9/// Uses an IMMEDIATE transaction, which acquires a write lock before any reads
10/// or writes occur. In SQLite WAL mode, IMMEDIATE transactions are serialized:
11/// concurrent reads can proceed, but no two processes can both hold IMMEDIATE
12/// (or EXCLUSIVE) transactions simultaneously. This means two processes calling
13/// `initialize_budget_replication_seq` concurrently will be serialized by
14/// SQLite's locking protocol -- the second caller blocks until the first commits,
15/// then runs with the updated seq floor already in place. No additional
16/// application-level locking is required.
17pub(super) fn initialize_budget_replication_seq(
18    connection: &mut Connection,
19) -> Result<(), BudgetStoreError> {
20    let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
21    let mut next_seq = current_budget_replication_seq(&transaction)?
22        .max(max_budget_usage_seq(&transaction)?)
23        .max(max_budget_mutation_event_seq(&transaction)?);
24    let mut statement = transaction.prepare(
25        r#"
26        SELECT rowid
27        FROM capability_grant_budgets
28        WHERE seq <= 0
29        ORDER BY updated_at ASC, capability_id ASC, grant_index ASC
30        "#,
31    )?;
32    let pending = statement
33        .query_map([], |row| row.get::<_, i64>(0))?
34        .collect::<Result<Vec<_>, _>>()?;
35    drop(statement);
36    for rowid in pending {
37        next_seq = next_seq.saturating_add(1);
38        transaction.execute(
39            "UPDATE capability_grant_budgets SET seq = ?1 WHERE rowid = ?2",
40            params![budget_u64_to_sqlite(next_seq, "seq")?, rowid],
41        )?;
42    }
43
44    let existing_event_seq_count = transaction.query_row(
45        "SELECT COUNT(*) FROM budget_mutation_events WHERE event_seq IS NOT NULL AND event_seq > 0",
46        [],
47        |row| row.get::<_, i64>(0),
48    )?;
49    if existing_event_seq_count <= 0 {
50        let mut statement = transaction.prepare(
51            r#"
52            SELECT rowid
53            FROM budget_mutation_events
54            ORDER BY rowid ASC
55            "#,
56        )?;
57        let pending = statement
58            .query_map([], |row| row.get::<_, i64>(0))?
59            .collect::<Result<Vec<_>, _>>()?;
60        drop(statement);
61        let mut event_seq = 0u64;
62        for rowid in pending {
63            event_seq = event_seq.saturating_add(1);
64            transaction.execute(
65                "UPDATE budget_mutation_events SET event_seq = ?1 WHERE rowid = ?2",
66                params![budget_u64_to_sqlite(event_seq, "event_seq")?, rowid],
67            )?;
68        }
69        next_seq = next_seq.max(event_seq);
70    } else {
71        let mut statement = transaction.prepare(
72            r#"
73            SELECT rowid
74            FROM budget_mutation_events
75            WHERE event_seq IS NULL OR event_seq <= 0
76            ORDER BY rowid ASC
77            "#,
78        )?;
79        let pending = statement
80            .query_map([], |row| row.get::<_, i64>(0))?
81            .collect::<Result<Vec<_>, _>>()?;
82        drop(statement);
83        for rowid in pending {
84            next_seq = next_seq.saturating_add(1);
85            transaction.execute(
86                "UPDATE budget_mutation_events SET event_seq = ?1 WHERE rowid = ?2",
87                params![budget_u64_to_sqlite(next_seq, "event_seq")?, rowid],
88            )?;
89        }
90    }
91    set_budget_replication_seq(&transaction, next_seq)?;
92    transaction.commit()?;
93    Ok(())
94}
95
96pub(super) fn allocate_budget_replication_seq(
97    transaction: &rusqlite::Transaction<'_>,
98) -> Result<u64, BudgetStoreError> {
99    let current = current_budget_replication_seq(transaction)?
100        .max(max_budget_usage_seq(transaction)?)
101        .max(max_budget_mutation_event_seq(transaction)?);
102    let next_seq = current.saturating_add(1);
103    set_budget_replication_seq(transaction, next_seq)?;
104    Ok(next_seq)
105}
106
107pub(super) fn raise_budget_replication_seq_floor(
108    transaction: &rusqlite::Transaction<'_>,
109    seq: u64,
110) -> Result<(), BudgetStoreError> {
111    let current = current_budget_replication_seq(transaction)?;
112    if seq > current {
113        set_budget_replication_seq(transaction, seq)?;
114    }
115    Ok(())
116}
117
118fn current_budget_replication_seq(
119    transaction: &rusqlite::Transaction<'_>,
120) -> Result<u64, BudgetStoreError> {
121    let next_seq = transaction.query_row(
122        "SELECT next_seq FROM budget_replication_meta WHERE singleton = 1",
123        [],
124        |row| budget_u64_from_row(row, 0, "next_seq"),
125    )?;
126    Ok(next_seq)
127}
128
129fn max_budget_usage_seq(transaction: &rusqlite::Transaction<'_>) -> Result<u64, BudgetStoreError> {
130    let max_seq = transaction.query_row(
131        "SELECT COALESCE(MAX(seq), 0) FROM capability_grant_budgets",
132        [],
133        |row| budget_u64_from_row(row, 0, "seq"),
134    )?;
135    Ok(max_seq)
136}
137
138fn max_budget_mutation_event_seq(
139    transaction: &rusqlite::Transaction<'_>,
140) -> Result<u64, BudgetStoreError> {
141    let max_seq = transaction.query_row(
142        "SELECT COALESCE(MAX(event_seq), 0) FROM budget_mutation_events",
143        [],
144        |row| budget_u64_from_row(row, 0, "event_seq"),
145    )?;
146    Ok(max_seq)
147}
148
149fn set_budget_replication_seq(
150    transaction: &rusqlite::Transaction<'_>,
151    seq: u64,
152) -> Result<(), BudgetStoreError> {
153    transaction.execute(
154        "UPDATE budget_replication_meta SET next_seq = ?1 WHERE singleton = 1",
155        params![budget_u64_to_sqlite(seq, "next_seq")?],
156    )?;
157    Ok(())
158}
159
160impl SqliteBudgetStore {
161    /// Highest budget mutation event_seq, or 0 when empty. Mirrors the private
162    /// max_budget_mutation_event_seq helper but is a public head read for the
163    /// status path.
164    pub fn max_mutation_event_seq(&self) -> Result<u64, BudgetStoreError> {
165        let mut connection = self.connection()?;
166        let transaction = self.begin_read(&mut connection)?;
167        let seq = transaction.query_row(
168            "SELECT COALESCE(MAX(event_seq), 0) FROM budget_mutation_events",
169            [],
170            |row| budget_u64_from_row(row, 0, "event_seq"),
171        )?;
172        transaction.rollback()?;
173        Ok(seq)
174    }
175
176    /// Highest mutation event_seq written under one origin authority, or 0 when
177    /// none. The cluster budget-write handler reads this immediately after a
178    /// local (leader) write to build the write's quorum-witness token: it is
179    /// always >= the just-written event's own seq (a concurrent same-origin
180    /// write can only raise it), so the per-origin contiguous witness can only
181    /// under-count witnesses, never over-count one (fail-closed).
182    pub fn max_mutation_event_seq_for_authority(
183        &self,
184        authority_id: &str,
185    ) -> Result<u64, BudgetStoreError> {
186        let mut connection = self.connection()?;
187        let transaction = self.begin_read(&mut connection)?;
188        let seq = transaction.query_row(
189            "SELECT COALESCE(MAX(event_seq), 0) FROM budget_mutation_events WHERE authority_id = ?1",
190            rusqlite::params![authority_id],
191            |row| budget_u64_from_row(row, 0, "event_seq"),
192        )?;
193        transaction.rollback()?;
194        Ok(seq)
195    }
196
197    /// The exact event_seq of the mutation event written under `event_id`, or
198    /// None if no such event exists (or it predates seq assignment). The
199    /// cluster budget-write handler waits on THIS write's own event_seq, looked
200    /// up by the write's event_id, instead of MAX(event_seq) for the authority:
201    /// a concurrent same-authority commit (or an idempotent retry while later
202    /// same-origin events already exist) can raise that MAX above this write's
203    /// seq, making the quorum wait target the wrong (higher) seq and roll back a
204    /// write that itself reached quorum. Looked up by the unique event_id, this
205    /// is race-free and, for an idempotent retry, returns the ORIGINAL event's
206    /// seq.
207    pub fn mutation_event_seq_for_event_id(
208        &self,
209        event_id: &str,
210    ) -> Result<Option<u64>, BudgetStoreError> {
211        let mut connection = self.connection()?;
212        let transaction = self.begin_read(&mut connection)?;
213        let seq: Option<Option<i64>> = transaction
214            .query_row(
215                "SELECT event_seq FROM budget_mutation_events WHERE event_id = ?1",
216                rusqlite::params![event_id],
217                |row| row.get::<_, Option<i64>>(0),
218            )
219            .optional()?;
220        transaction.rollback()?;
221        seq.flatten()
222            .map(|value| {
223                u64::try_from(value).map_err(|_| {
224                    BudgetStoreError::Invariant(
225                        "budget mutation event has a negative event sequence".to_string(),
226                    )
227                })
228            })
229            .transpose()
230    }
231
232    pub fn mutation_event_for_event_id(
233        &self,
234        event_id: &str,
235    ) -> Result<Option<BudgetMutationRecord>, BudgetStoreError> {
236        let mut connection = self.connection()?;
237        let transaction = self.begin_read(&mut connection)?;
238        let event = Self::load_projected_mutation_event(&transaction, event_id)?;
239        transaction.rollback()?;
240        Ok(event)
241    }
242
243    pub fn usage_projection_for_event_id(
244        &self,
245        event_id: &str,
246    ) -> Result<Option<BudgetUsageRecord>, BudgetStoreError> {
247        let mut connection = self.connection()?;
248        let transaction = self.begin_read(&mut connection)?;
249        let event = Self::load_mutation_event(&transaction, event_id)?;
250        let Some(event) = event else {
251            transaction.rollback()?;
252            return Ok(None);
253        };
254        let Some(usage_seq) = event.usage_seq else {
255            transaction.rollback()?;
256            return Ok(None);
257        };
258        if usage_seq == 0 || usage_seq > event.event_seq {
259            return Err(BudgetStoreError::Invariant(format!(
260                "budget event `{event_id}` has an invalid usage sequence {usage_seq}"
261            )));
262        }
263        let usage = transaction
264            .query_row(
265                r#"
266                SELECT recorded_at, invocation_count_after,
267                       total_cost_exposed_after, total_cost_realized_spend_after
268                FROM budget_mutation_events
269                WHERE capability_id = ?1 AND grant_index = ?2
270                  AND event_seq = ?3 AND usage_seq = event_seq
271                "#,
272                params![
273                    &event.capability_id,
274                    i64::from(event.grant_index),
275                    budget_u64_to_sqlite(usage_seq, "usage_seq")?,
276                ],
277                |row| {
278                    Ok(BudgetUsageRecord {
279                        capability_id: event.capability_id.clone(),
280                        grant_index: event.grant_index,
281                        invocation_count: budget_u32_from_row(row, 1, "invocation_count_after")?,
282                        updated_at: row.get(0)?,
283                        seq: usage_seq,
284                        total_cost_exposed: budget_u64_from_row(
285                            row,
286                            2,
287                            "total_cost_exposed_after",
288                        )?,
289                        total_cost_realized_spend: budget_u64_from_row(
290                            row,
291                            3,
292                            "total_cost_realized_spend_after",
293                        )?,
294                    })
295                },
296            )
297            .optional()?
298            .ok_or_else(|| {
299                BudgetStoreError::Invariant(format!(
300                    "budget event `{event_id}` references missing usage sequence {usage_seq}"
301                ))
302            })?;
303        if usage.invocation_count != event.invocation_count_after
304            || usage.total_cost_exposed != event.total_cost_exposed_after
305            || usage.total_cost_realized_spend != event.total_cost_realized_spend_after
306        {
307            return Err(BudgetStoreError::Invariant(format!(
308                "budget event `{event_id}` usage projection counters changed"
309            )));
310        }
311        transaction.rollback()?;
312        Ok(Some(usage))
313    }
314
315    /// The quorum-witness identity of the mutation event stored under `event_id`:
316    /// `(event_seq, authority_id, lease_epoch)`, or None when no such event exists
317    /// or it predates seq assignment. The witness must target the event's STORED
318    /// origin authority, not the current lease: an idempotent retry after
319    /// leadership moved re-reads the already-written event, and peers advertise it
320    /// under its ORIGINAL authority, so keying the wait on the current leader would
321    /// look under the wrong origin and time out a write that already committed.
322    /// A null-seq (legacy) row returns None so the caller
323    /// falls back to the authority MAX rather than witnessing on seq 0.
324    pub fn mutation_event_witness_for_event_id(
325        &self,
326        event_id: &str,
327    ) -> Result<Option<BudgetEventWitness>, BudgetStoreError> {
328        let mut connection = self.connection()?;
329        let transaction = self.begin_read(&mut connection)?;
330        let row = transaction
331            .query_row(
332                "SELECT event_seq, authority_id, lease_epoch FROM budget_mutation_events WHERE event_id = ?1",
333                rusqlite::params![event_id],
334                |row| {
335                    let seq: Option<i64> = row.get(0)?;
336                    let authority_id: Option<String> = row.get(1)?;
337                    let lease_epoch: Option<i64> = row.get(2)?;
338                    Ok((seq, authority_id, lease_epoch))
339                },
340            )
341            .optional()?;
342        let witness = row
343            .map(
344                |(seq, authority_id, lease_epoch)| -> Result<_, BudgetStoreError> {
345                    let Some(seq) = seq else {
346                        return Ok(None);
347                    };
348                    let seq = u64::try_from(seq).map_err(|_| {
349                        BudgetStoreError::Invariant(
350                            "budget mutation witness has a negative event sequence".to_string(),
351                        )
352                    })?;
353                    let lease_epoch = lease_epoch
354                        .map(|epoch| {
355                            u64::try_from(epoch).map_err(|_| {
356                                BudgetStoreError::Invariant(
357                                    "budget mutation witness has a negative lease epoch"
358                                        .to_string(),
359                                )
360                            })
361                        })
362                        .transpose()?;
363                    Ok(Some((seq, authority_id, lease_epoch)))
364                },
365            )
366            .transpose()?
367            .flatten();
368        transaction.rollback()?;
369        Ok(witness)
370    }
371
372    /// Record snapshot-carried abandoned event sequences.
373    pub fn record_abandoned_event_seqs(&self, seqs: &[u64]) -> Result<(), BudgetStoreError> {
374        self.require_standalone_mutation("abandoned event sequence import")?;
375        if seqs.is_empty() {
376            return Ok(());
377        }
378        let mut connection = self.connection()?;
379        let transaction = self.begin_write(&mut connection)?;
380        for &seq in seqs {
381            if seq == 0 {
382                continue;
383            }
384            transaction.execute(
385                "INSERT OR IGNORE INTO budget_abandoned_event_seqs(seq) VALUES (?1)",
386                rusqlite::params![budget_u64_to_sqlite(seq, "seq")?],
387            )?;
388        }
389        // Tombstone inserts only fill holes. Preserve the proven watermark and
390        // any snapshot-installed per-origin heads so the next scan can advance.
391        transaction.commit()?;
392        Ok(())
393    }
394
395    /// Record snapshot-carried abandoned event sequences as inclusive ranges.
396    pub fn record_abandoned_event_seq_ranges(
397        &self,
398        ranges: &[(u64, u64)],
399    ) -> Result<(), BudgetStoreError> {
400        self.require_standalone_mutation("abandoned event sequence range import")?;
401        if ranges.is_empty() {
402            return Ok(());
403        }
404        let mut connection = self.connection()?;
405        let transaction = self.begin_write(&mut connection)?;
406        Self::insert_abandoned_event_seq_ranges(&transaction, ranges)?;
407        // As above, adding covered slots cannot invalidate an acknowledged prefix.
408        transaction.commit()?;
409        Ok(())
410    }
411
412    pub(super) fn insert_abandoned_event_seq_ranges(
413        transaction: &rusqlite::Transaction<'_>,
414        ranges: &[(u64, u64)],
415    ) -> Result<(), BudgetStoreError> {
416        let mut previous_end = 0;
417        for &(start, end) in ranges {
418            if start == 0 || end < start || start <= previous_end {
419                return Err(BudgetStoreError::Invariant(
420                    "abandoned budget sequence ranges must be positive, ordered, and non-overlapping"
421                        .to_string(),
422                ));
423            }
424            let start = budget_u64_to_sqlite(start, "range_start_seq")?;
425            let end = budget_u64_to_sqlite(end, "range_end_seq")?;
426            let overlaps_event: bool = transaction.query_row(
427                "SELECT EXISTS(SELECT 1 FROM budget_mutation_events WHERE event_seq BETWEEN ?1 AND ?2)",
428                rusqlite::params![start, end],
429                |row| row.get::<_, i64>(0).map(|value| value != 0),
430            )?;
431            if overlaps_event {
432                return Err(BudgetStoreError::Invariant(format!(
433                    "abandoned budget sequence range {start}..={end} overlaps a mutation event"
434                )));
435            }
436            transaction.execute(
437                r#"
438                INSERT OR IGNORE INTO budget_abandoned_event_seqs(seq)
439                WITH RECURSIVE run(s) AS (
440                    SELECT ?1
441                    UNION ALL
442                    SELECT s + 1 FROM run WHERE s < ?2
443                )
444                SELECT s FROM run
445                "#,
446                rusqlite::params![start, end],
447            )?;
448            previous_end = end as u64;
449        }
450        Ok(())
451    }
452}
453
454pub(super) fn unix_now() -> i64 {
455    SystemTime::now()
456        .duration_since(UNIX_EPOCH)
457        .map(|duration| duration.as_secs() as i64)
458        .unwrap_or(0)
459}
chio_store_sqlite/budget_store/replication.rs

chio_store_sqlite/budget_store/
replication.rs
