sqlite_graphrag/storage/
utils.rs

1//! Storage utility helpers shared across the storage sub-modules.
2
3use crate::constants::{MAX_SQLITE_BUSY_RETRIES, SQLITE_BUSY_BASE_DELAY_MS};
4use crate::errors::AppError;
5use rusqlite::ErrorCode;
6use std::thread;
7use std::time::Duration;
8
9/// Returns `true` when `err` wraps an `SQLITE_BUSY` (or `SQLITE_LOCKED`)
10/// condition reported by rusqlite.
11///
12/// Both `SQLITE_BUSY` (`ErrorCode::DatabaseBusy`) and `SQLITE_LOCKED`
13/// (`ErrorCode::DatabaseLocked`) indicate that the write cannot proceed
14/// immediately due to WAL concurrency.  We treat both as transient and
15/// eligible for retry.
16pub fn is_sqlite_busy(err: &AppError) -> bool {
17    match err {
18        AppError::Database(rusqlite::Error::SqliteFailure(e, _)) => {
19            e.code == ErrorCode::DatabaseBusy || e.code == ErrorCode::DatabaseLocked
20        }
21        _ => false,
22    }
23}
24
25/// Executes `op` up to `MAX_SQLITE_BUSY_RETRIES` times with exponential
26/// backoff whenever the operation fails with `SQLITE_BUSY` / `SQLITE_LOCKED`.
27///
28/// Delay schedule (base = `SQLITE_BUSY_BASE_DELAY_MS`):
29/// - attempt 1 → `base` ms
30/// - attempt 2 → `base * 2` ms
31/// - attempt 3 → `base * 4` ms
32/// - attempt 4 → `base * 8` ms
33/// - attempt 5 → `base * 16` ms
34///
35/// After all retries are exhausted the last `SQLITE_BUSY` error is converted
36/// to [`AppError::DbBusy`] so callers can route on exit-code `15`.
37pub fn with_busy_retry<F>(op: F) -> Result<(), AppError>
38where
39    F: Fn() -> Result<(), AppError>,
40{
41    for attempt in 0..MAX_SQLITE_BUSY_RETRIES {
42        match op() {
43            Ok(()) => return Ok(()),
44            Err(e) if is_sqlite_busy(&e) => {
45                // Exponential backoff: base_ms * 2^attempt
46                let delay_ms = SQLITE_BUSY_BASE_DELAY_MS * (1u64 << attempt);
47                thread::sleep(Duration::from_millis(delay_ms));
48            }
49            Err(other) => return Err(other),
50        }
51    }
52
53    // All retries exhausted — convert to DbBusy for stable exit-code 15.
54    Err(AppError::DbBusy(format!(
55        "SQLITE_BUSY after {MAX_SQLITE_BUSY_RETRIES} retries"
56    )))
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use std::sync::atomic::{AtomicU32, Ordering};
63    use std::sync::Arc;
64
65    /// Helper that builds a fake `AppError::Database` wrapping
66    /// `SQLITE_BUSY` (error code 5) so that `is_sqlite_busy` can be tested
67    /// without needing a live SQLite connection.
68    fn make_busy_error() -> AppError {
69        // rusqlite::Error::SqliteFailure requires a `ffi::Error` + optional msg.
70        // We construct it via the public `rusqlite::ffi` interface.
71        let ffi_err = rusqlite::ffi::Error {
72            code: ErrorCode::DatabaseBusy,
73            extended_code: 5,
74        };
75        AppError::Database(rusqlite::Error::SqliteFailure(ffi_err, None))
76    }
77
78    fn make_locked_error() -> AppError {
79        let ffi_err = rusqlite::ffi::Error {
80            code: ErrorCode::DatabaseLocked,
81            extended_code: 6,
82        };
83        AppError::Database(rusqlite::Error::SqliteFailure(ffi_err, None))
84    }
85
86    #[test]
87    fn is_sqlite_busy_detecta_database_busy() {
88        assert!(is_sqlite_busy(&make_busy_error()));
89    }
90
91    #[test]
92    fn is_sqlite_busy_detecta_database_locked() {
93        assert!(is_sqlite_busy(&make_locked_error()));
94    }
95
96    #[test]
97    fn is_sqlite_busy_rejeita_outros_erros() {
98        let err = AppError::Validation("campo inválido".into());
99        assert!(!is_sqlite_busy(&err));
100    }
101
102    #[test]
103    fn with_busy_retry_propagates_non_busy_error() {
104        let calls = Arc::new(AtomicU32::new(0));
105        let calls_clone = Arc::clone(&calls);
106
107        let result = with_busy_retry(|| {
108            calls_clone.fetch_add(1, Ordering::SeqCst);
109            Err(AppError::Validation("campo x".into()))
110        });
111
112        // Non-busy errors must propagate immediately without retrying.
113        assert_eq!(calls.load(Ordering::SeqCst), 1);
114        assert!(matches!(result, Err(AppError::Validation(_))));
115    }
116
117    #[test]
118    fn with_busy_retry_succeeds_on_third_attempt() {
119        let calls = Arc::new(AtomicU32::new(0));
120        let calls_clone = Arc::clone(&calls);
121
122        // Fail twice with SQLITE_BUSY, succeed on the third call.
123        let result = with_busy_retry(|| {
124            let n = calls_clone.fetch_add(1, Ordering::SeqCst);
125            if n < 2 {
126                Err(make_busy_error())
127            } else {
128                Ok(())
129            }
130        });
131
132        assert_eq!(calls.load(Ordering::SeqCst), 3);
133        assert!(result.is_ok(), "expected Ok after 3rd attempt");
134    }
135
136    #[test]
137    fn with_busy_retry_returns_db_busy_after_all_retries() {
138        let calls = Arc::new(AtomicU32::new(0));
139        let calls_clone = Arc::clone(&calls);
140
141        let result = with_busy_retry(|| {
142            calls_clone.fetch_add(1, Ordering::SeqCst);
143            Err(make_busy_error())
144        });
145
146        assert_eq!(
147            calls.load(Ordering::SeqCst),
148            MAX_SQLITE_BUSY_RETRIES,
149            "must attempt exactly MAX_SQLITE_BUSY_RETRIES times"
150        );
151        assert!(
152            matches!(result, Err(AppError::DbBusy(_))),
153            "must convert to DbBusy after exhausting retries"
154        );
155    }
156}
sqlite_graphrag/storage/utils.rs

sqlite_graphrag/storage/
utils.rs
