Skip to main content

aft/blob_store/
mod.rs

1//! Immutable, content-addressed payload storage shared by a repository family.
2//!
3//! Each [`BlobStore`] owns one SQLite connection for exactly one plane.  The
4//! caller supplies the family (`artifact_key`); this module maps it to the
5//! fixed on-disk layout and never derives a family from a checkout path.
6
7use std::collections::HashMap;
8use std::error::Error;
9use std::fmt;
10use std::fs;
11use std::path::{Path, PathBuf};
12use std::sync::{Mutex, MutexGuard, OnceLock};
13use std::time::{Duration, SystemTime, UNIX_EPOCH};
14
15use rusqlite::{params, Connection, ErrorCode, OptionalExtension, TransactionBehavior};
16
17use crate::db::lifecycle::{SqliteStore, TrackedConnection};
18
19/// The SQLite busy wait used by every blob-store connection.
20pub const BUSY_TIMEOUT_MS: u64 = 5_000;
21/// SQLite's documented connection default.  The blob store reads this value
22/// after opening instead of overriding it so future SQLite defaults are caught.
23pub const DEFAULT_WAL_AUTOCHECKPOINT_PAGES: i64 = 1_000;
24
25#[derive(Clone, Copy)]
26struct BlobDurabilityState {
27    dirty: bool,
28}
29
30static BLOB_DURABILITY: OnceLock<Mutex<HashMap<PathBuf, BlobDurabilityState>>> = OnceLock::new();
31static BLOB_DURABILITY_BARRIER: Mutex<()> = Mutex::new(());
32
33fn durability_states() -> &'static Mutex<HashMap<PathBuf, BlobDurabilityState>> {
34    BLOB_DURABILITY.get_or_init(|| Mutex::new(HashMap::new()))
35}
36
37fn register_blob_database(path: &Path) {
38    durability_states()
39        .lock()
40        .unwrap_or_else(std::sync::PoisonError::into_inner)
41        .entry(path.to_path_buf())
42        .or_insert(BlobDurabilityState { dirty: true });
43}
44
45fn mark_blob_database_dirty(path: &Path) {
46    durability_states()
47        .lock()
48        .unwrap_or_else(std::sync::PoisonError::into_inner)
49        .entry(path.to_path_buf())
50        .and_modify(|state| state.dirty = true)
51        .or_insert(BlobDurabilityState { dirty: true });
52}
53
54pub(crate) fn blob_database_needs_durability(path: &Path) -> bool {
55    let mut states = durability_states()
56        .lock()
57        .unwrap_or_else(std::sync::PoisonError::into_inner);
58    states
59        .entry(path.to_path_buf())
60        .or_insert(BlobDurabilityState { dirty: true })
61        .dirty
62}
63
64pub(crate) fn mark_blob_database_durable(path: &Path) {
65    durability_states()
66        .lock()
67        .unwrap_or_else(std::sync::PoisonError::into_inner)
68        .entry(path.to_path_buf())
69        .and_modify(|state| state.dirty = false)
70        .or_insert(BlobDurabilityState { dirty: false });
71}
72
73pub(crate) fn publication_durability_barrier() -> MutexGuard<'static, ()> {
74    BLOB_DURABILITY_BARRIER
75        .lock()
76        .unwrap_or_else(std::sync::PoisonError::into_inner)
77}
78
79/// Payload encodings and the key producers that name them are pinned together.
80/// A payload format change must update the corresponding producer string in the
81/// same edit, because producer versions are components of content-address keys.
82pub const SEMANTIC_PAYLOAD_SCHEMA: u32 = 1;
83pub const SEMANTIC_PRODUCER_VERSION: &str = "semantic-v1";
84pub const CALLGRAPH_PAYLOAD_SCHEMA: u32 = 1;
85pub const CALLGRAPH_PRODUCER_VERSION: &str = "callgraph-v1";
86
87const BLOB_SCHEMA: &str = r#"
88CREATE TABLE IF NOT EXISTS blob_payloads (
89    full_key BLOB NOT NULL PRIMARY KEY CHECK(length(full_key) = 32),
90    payload BLOB NOT NULL,
91    payload_digest BLOB NOT NULL CHECK(length(payload_digest) = 32),
92    payload_schema INTEGER NOT NULL,
93    created_at_ms INTEGER NOT NULL DEFAULT 0
94) WITHOUT ROWID;
95-- WITHOUT ROWID stores payload bytes in the primary-key tree. Membership probes
96-- need a separate, narrow tree to avoid reading payload overflow pages.
97CREATE INDEX IF NOT EXISTS blob_membership ON blob_payloads(full_key);
98CREATE TABLE IF NOT EXISTS blob_quarantine (
99    full_key BLOB NOT NULL PRIMARY KEY CHECK(length(full_key) = 32)
100) WITHOUT ROWID;
101"#;
102
103/// The two repository-family blob planes.  Trigram data is per-view derived
104/// state and deliberately is not represented here.
105#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
106pub enum BlobPlane {
107    Semantic,
108    Callgraph,
109}
110
111impl BlobPlane {
112    pub const fn as_str(self) -> &'static str {
113        match self {
114            Self::Semantic => "semantic",
115            Self::Callgraph => "callgraph",
116        }
117    }
118
119    const fn payload_schema(self) -> u32 {
120        match self {
121            Self::Semantic => SEMANTIC_PAYLOAD_SCHEMA,
122            Self::Callgraph => CALLGRAPH_PAYLOAD_SCHEMA,
123        }
124    }
125}
126
127/// A complete, fixed-size database key.  Its bytes are the BLAKE3 digest of a
128/// length-delimited canonical encoding of the semantic or callgraph key tuple.
129#[derive(Clone, Debug, Eq, PartialEq, Hash)]
130pub struct FullKey {
131    bytes: [u8; 32],
132    plane: BlobPlane,
133}
134
135impl FullKey {
136    /// Returns the bytes stored in SQLite's `full_key` BLOB column.
137    pub fn as_bytes(&self) -> &[u8; 32] {
138        &self.bytes
139    }
140
141    /// The only plane where this key can be stored.
142    pub const fn plane(&self) -> BlobPlane {
143        self.plane
144    }
145
146    /// Stable lower-case hexadecimal form suitable for logs and telemetry.
147    pub fn to_hex(&self) -> String {
148        let mut hex = String::with_capacity(64);
149        for byte in self.bytes {
150            use std::fmt::Write;
151            let _ = write!(hex, "{byte:02x}");
152        }
153        hex
154    }
155}
156
157impl fmt::Display for FullKey {
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        f.write_str(&self.to_hex())
160    }
161}
162
163/// A semantic blob identity.  `rel_path` remains a key component even when the
164/// source bytes are identical, so path-specific embedding text cannot be reused
165/// under a different path by accident.
166#[derive(Clone, Debug, Eq, PartialEq, Hash)]
167pub struct SemanticKey {
168    source_digest: [u8; 32],
169    rel_path: Vec<u8>,
170    chunker_version: String,
171    embed_template_version: String,
172    model_fingerprint: String,
173}
174
175impl SemanticKey {
176    pub fn from_bytes(
177        bytes: &[u8],
178        rel_path: &[u8],
179        chunker_version: impl Into<String>,
180        embed_template_version: impl Into<String>,
181        model_fingerprint: impl Into<String>,
182    ) -> Self {
183        Self {
184            source_digest: *blake3::hash(bytes).as_bytes(),
185            rel_path: rel_path.to_vec(),
186            chunker_version: chunker_version.into(),
187            embed_template_version: embed_template_version.into(),
188            model_fingerprint: model_fingerprint.into(),
189        }
190    }
191
192    /// Builds a key using this release's semantic producer version for both
193    /// producer components.  Callers that independently version the chunker
194    /// and template should use [`Self::from_bytes`] instead.
195    pub fn for_current(
196        bytes: &[u8],
197        rel_path: &[u8],
198        model_fingerprint: impl Into<String>,
199    ) -> Self {
200        Self::from_bytes(
201            bytes,
202            rel_path,
203            SEMANTIC_PRODUCER_VERSION,
204            SEMANTIC_PRODUCER_VERSION,
205            model_fingerprint,
206        )
207    }
208
209    pub fn full_key(&self) -> FullKey {
210        full_key(
211            BlobPlane::Semantic,
212            b"aft/blob-store/semantic/v1",
213            &[
214                &self.source_digest,
215                &self.rel_path,
216                self.chunker_version.as_bytes(),
217                self.embed_template_version.as_bytes(),
218                self.model_fingerprint.as_bytes(),
219            ],
220        )
221    }
222
223    pub fn source_digest(&self) -> &[u8; 32] {
224        &self.source_digest
225    }
226}
227
228/// A callgraph blob identity.  Unlike semantic keys, the relative path is not a
229/// component: equal source bytes with the same language and extractor version
230/// share one parse extraction across all paths in the family.
231#[derive(Clone, Debug, Eq, PartialEq, Hash)]
232pub struct CallgraphKey {
233    source_digest: [u8; 32],
234    language: String,
235    extractor_version: String,
236}
237
238impl CallgraphKey {
239    pub fn from_bytes(
240        bytes: &[u8],
241        language: impl Into<String>,
242        extractor_version: impl Into<String>,
243    ) -> Self {
244        Self {
245            source_digest: *blake3::hash(bytes).as_bytes(),
246            language: language.into(),
247            extractor_version: extractor_version.into(),
248        }
249    }
250
251    /// Uses this release's extractor producer version.  `language = "config"`
252    /// is intentionally accepted because configuration inputs are callgraph
253    /// blobs too.
254    pub fn for_current(bytes: &[u8], language: impl Into<String>) -> Self {
255        Self::from_bytes(bytes, language, CALLGRAPH_PRODUCER_VERSION)
256    }
257
258    pub fn full_key(&self) -> FullKey {
259        full_key(
260            BlobPlane::Callgraph,
261            b"aft/blob-store/callgraph/v1",
262            &[
263                &self.source_digest,
264                self.language.as_bytes(),
265                self.extractor_version.as_bytes(),
266            ],
267        )
268    }
269
270    pub fn source_digest(&self) -> &[u8; 32] {
271        &self.source_digest
272    }
273}
274
275fn full_key(plane: BlobPlane, domain: &[u8], fields: &[&[u8]]) -> FullKey {
276    let mut hasher = blake3::Hasher::new();
277    hasher.update(domain);
278    for field in fields {
279        hasher.update(&(field.len() as u64).to_be_bytes());
280        hasher.update(field);
281    }
282    FullKey {
283        bytes: *hasher.finalize().as_bytes(),
284        plane,
285    }
286}
287
288/// The only public put outcomes.  A durable report is impossible for the last
289/// three outcomes because [`PutReport::durable`] derives from this enum.
290#[derive(Clone, Copy, Debug, Eq, PartialEq)]
291pub enum PutOutcome {
292    Inserted,
293    Reused,
294    Quarantined,
295    Failed,
296    QuotaExceeded,
297}
298
299/// The result of a successful put attempt.
300#[derive(Clone, Copy, Debug, Eq, PartialEq)]
301pub struct PutReport {
302    pub outcome: PutOutcome,
303    pub durable: bool,
304}
305
306impl PutReport {
307    fn new(outcome: PutOutcome) -> Self {
308        Self {
309            durable: matches!(outcome, PutOutcome::Inserted | PutOutcome::Reused),
310            outcome,
311        }
312    }
313}
314
315/// The pragma values read back from the open connection after configuration.
316#[derive(Clone, Debug, Eq, PartialEq)]
317pub struct BlobStorePragmas {
318    pub journal_mode: String,
319    pub synchronous: i64,
320    pub busy_timeout_ms: i64,
321    pub foreign_keys: i64,
322    pub wal_autocheckpoint_pages: i64,
323}
324
325/// Space and row usage for one immutable family/plane store.
326#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
327pub struct BlobUsage {
328    pub rows: u64,
329    pub payload_bytes: u64,
330    pub file_bytes: u64,
331}
332
333/// A breaker implementation is notified exactly once after a corrupt database
334/// has been moved aside and a clean replacement has opened successfully; the
335/// notification records that recovery event.
336pub trait BlobStoreBreaker {
337    fn record_corruption_death(&self, artifact_key: &str, plane: BlobPlane);
338}
339
340#[derive(Debug)]
341pub enum BlobStoreError {
342    Io(std::io::Error),
343    Sqlite(rusqlite::Error),
344    InvalidArtifactKey(String),
345    PragmaMismatch {
346        name: &'static str,
347        expected: String,
348        actual: String,
349    },
350    PlaneKeyMismatch {
351        store_plane: BlobPlane,
352        key_plane: BlobPlane,
353    },
354}
355
356impl fmt::Display for BlobStoreError {
357    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
358        match self {
359            Self::Io(error) => write!(f, "blob-store I/O error: {error}"),
360            Self::Sqlite(error) => write!(f, "blob-store SQLite error: {error}"),
361            Self::InvalidArtifactKey(key) => write!(f, "invalid artifact key `{key}`"),
362            Self::PlaneKeyMismatch {
363                store_plane,
364                key_plane,
365            } => write!(
366                f,
367                "a {} key cannot be stored in the {} plane",
368                key_plane.as_str(),
369                store_plane.as_str()
370            ),
371            Self::PragmaMismatch {
372                name,
373                expected,
374                actual,
375            } => write!(
376                f,
377                "blob-store PRAGMA {name} was `{actual}`, expected `{expected}`"
378            ),
379        }
380    }
381}
382
383impl Error for BlobStoreError {
384    fn source(&self) -> Option<&(dyn Error + 'static)> {
385        match self {
386            Self::Io(error) => Some(error),
387            Self::Sqlite(error) => Some(error),
388            Self::InvalidArtifactKey(_)
389            | Self::PragmaMismatch { .. }
390            | Self::PlaneKeyMismatch { .. } => None,
391        }
392    }
393}
394
395impl From<std::io::Error> for BlobStoreError {
396    fn from(error: std::io::Error) -> Self {
397        Self::Io(error)
398    }
399}
400
401impl From<rusqlite::Error> for BlobStoreError {
402    fn from(error: rusqlite::Error) -> Self {
403        Self::Sqlite(error)
404    }
405}
406
407/// One SQLite-backed immutable payload plane for one repository family.
408#[derive(Debug)]
409pub struct BlobStore {
410    artifact_key: String,
411    plane: BlobPlane,
412    path: PathBuf,
413    pragmas: BlobStorePragmas,
414    connection: TrackedConnection,
415}
416
417impl BlobStore {
418    /// Opens `<storage>/blobs/<artifact_key>/<plane>.sqlite` and creates its
419    /// schema under `BEGIN IMMEDIATE`.  Existing corrupt SQLite headers are
420    /// preserved beside the database before a clean empty store is created.
421    pub fn open(
422        storage: &Path,
423        artifact_key: impl Into<String>,
424        plane: BlobPlane,
425    ) -> Result<Self, BlobStoreError> {
426        Self::open_with_optional_breaker(storage, artifact_key.into(), plane, None)
427    }
428
429    /// Like [`Self::open`], but invokes the supplied breaker exactly once after
430    /// corruption recovery has completed successfully.
431    pub fn open_with_breaker(
432        storage: &Path,
433        artifact_key: impl Into<String>,
434        plane: BlobPlane,
435        breaker: &dyn BlobStoreBreaker,
436    ) -> Result<Self, BlobStoreError> {
437        Self::open_with_optional_breaker(storage, artifact_key.into(), plane, Some(breaker))
438    }
439
440    fn open_with_optional_breaker(
441        storage: &Path,
442        artifact_key: String,
443        plane: BlobPlane,
444        breaker: Option<&dyn BlobStoreBreaker>,
445    ) -> Result<Self, BlobStoreError> {
446        validate_artifact_key(&artifact_key)?;
447        let path = storage
448            .join("blobs")
449            .join(&artifact_key)
450            .join(format!("{}.sqlite", plane.as_str()));
451        if let Some(parent) = path.parent() {
452            fs::create_dir_all(parent)?;
453        }
454
455        match Self::open_at(&artifact_key, plane, path.clone()) {
456            Ok(store) => Ok(store),
457            Err(error) if is_corrupt_database_error(&error) && path.exists() => {
458                let corrupt_path = move_corrupt_database_aside(&path)?;
459                log::warn!(
460                    "blob store database at {} was corrupt; moved it to {}",
461                    path.display(),
462                    corrupt_path.display()
463                );
464                let store = Self::open_at(&artifact_key, plane, path)?;
465                if let Some(breaker) = breaker {
466                    breaker.record_corruption_death(&artifact_key, plane);
467                }
468                Ok(store)
469            }
470            Err(error) => Err(error),
471        }
472    }
473
474    fn open_at(
475        artifact_key: &str,
476        plane: BlobPlane,
477        path: PathBuf,
478    ) -> Result<Self, BlobStoreError> {
479        let mut connection = TrackedConnection::open(&path, SqliteStore::BlobStore)?;
480        configure_connection(&connection)?;
481        ensure_schema(&mut connection)?;
482        let pragmas = read_and_assert_pragmas(&connection)?;
483        register_blob_database(&path);
484        Ok(Self {
485            artifact_key: artifact_key.to_owned(),
486            plane,
487            path,
488            pragmas,
489            connection,
490        })
491    }
492
493    pub fn artifact_key(&self) -> &str {
494        &self.artifact_key
495    }
496
497    pub const fn plane(&self) -> BlobPlane {
498        self.plane
499    }
500
501    pub fn path(&self) -> &Path {
502        &self.path
503    }
504
505    pub fn pragmas(&self) -> &BlobStorePragmas {
506        &self.pragmas
507    }
508
509    /// Returns current payload rows plus the SQLite file allocation.  Quota
510    /// admission uses this before an insert so a rejected payload never enters
511    /// the immutable store.
512    pub fn usage(&self) -> Result<BlobUsage, BlobStoreError> {
513        let (rows, payload_bytes): (u64, u64) = self.connection.query_row(
514            "SELECT COUNT(*), COALESCE(SUM(length(payload)), 0) FROM blob_payloads",
515            [],
516            |row| Ok((row.get(0)?, row.get(1)?)),
517        )?;
518        let page_count: u64 = self
519            .connection
520            .pragma_query_value(None, "page_count", |row| row.get(0))?;
521        let page_size: u64 = self
522            .connection
523            .pragma_query_value(None, "page_size", |row| row.get(0))?;
524        Ok(BlobUsage {
525            rows,
526            payload_bytes,
527            file_bytes: page_count.saturating_mul(page_size),
528        })
529    }
530
531    /// Inserts a payload once.  An existing row is never updated, even if its
532    /// bytes fail a later integrity check; callers must use a new producer key.
533    pub fn put(&mut self, full_key: &FullKey, payload: &[u8]) -> Result<PutReport, BlobStoreError> {
534        self.ensure_key_plane(full_key)?;
535        let _durability = publication_durability_barrier();
536        let tx = self
537            .connection
538            .transaction_with_behavior(TransactionBehavior::Immediate)?;
539        let quarantined = tx
540            .query_row(
541                "SELECT 1 FROM blob_quarantine WHERE full_key = ?1",
542                params![full_key.as_bytes().as_slice()],
543                |_| Ok(()),
544            )
545            .optional()?
546            .is_some();
547        if quarantined {
548            tx.commit()?;
549            return Ok(PutReport::new(PutOutcome::Quarantined));
550        }
551
552        let payload_digest = blake3::hash(payload);
553        let inserted = tx.execute(
554            "INSERT INTO blob_payloads
555             (full_key, payload, payload_digest, payload_schema, created_at_ms)
556             VALUES (?1, ?2, ?3, ?4, ?5)
557             ON CONFLICT(full_key) DO NOTHING",
558            params![
559                full_key.as_bytes().as_slice(),
560                payload,
561                payload_digest.as_bytes().as_slice(),
562                i64::from(self.plane.payload_schema()),
563                unix_millis_now(),
564            ],
565        )?;
566        tx.commit()?;
567        if inserted == 1 {
568            mark_blob_database_dirty(&self.path);
569        }
570        Ok(PutReport::new(if inserted == 1 {
571            PutOutcome::Inserted
572        } else {
573            PutOutcome::Reused
574        }))
575    }
576
577    /// Reads a payload only after verifying its stored digest and schema.  A
578    /// malformed row is indistinguishable from a miss to consumers, but remains
579    /// on disk for quarantine/forensics rather than being rewritten in place.
580    pub fn get(&self, full_key: &FullKey) -> Result<Option<Vec<u8>>, BlobStoreError> {
581        self.ensure_key_plane(full_key)?;
582        let row = self
583            .connection
584            .query_row(
585                "SELECT payload, payload_digest, payload_schema
586                 FROM blob_payloads WHERE full_key = ?1",
587                params![full_key.as_bytes().as_slice()],
588                |row| {
589                    Ok((
590                        row.get::<_, Vec<u8>>(0)?,
591                        row.get::<_, Vec<u8>>(1)?,
592                        row.get::<_, i64>(2)?,
593                    ))
594                },
595            )
596            .optional()?;
597        let Some((payload, payload_digest, payload_schema)) = row else {
598            return Ok(None);
599        };
600
601        let digest_matches = payload_digest.as_slice() == blake3::hash(&payload).as_bytes();
602        let schema_matches = payload_schema == i64::from(self.plane.payload_schema());
603        if digest_matches && schema_matches {
604            return Ok(Some(payload));
605        }
606
607        let reason = match (digest_matches, schema_matches) {
608            (false, false) => "payload digest and schema mismatch",
609            (false, true) => "payload digest mismatch",
610            (true, false) => "payload schema mismatch",
611            (true, true) => unreachable!("matching payload was returned above"),
612        };
613        log::warn!(
614            "blob store rejected committed payload for key {} in {}/{}: {}",
615            full_key,
616            self.artifact_key,
617            self.plane.as_str(),
618            reason
619        );
620        Ok(None)
621    }
622
623    /// Records a deterministic failure without modifying the immutable payload
624    /// table.  Subsequent puts report `quarantined` until a changed key is used.
625    pub fn quarantine(&mut self, full_key: &FullKey) -> Result<(), BlobStoreError> {
626        self.ensure_key_plane(full_key)?;
627        let _durability = publication_durability_barrier();
628        let inserted = self.connection.execute(
629            "INSERT INTO blob_quarantine (full_key) VALUES (?1)
630             ON CONFLICT(full_key) DO NOTHING",
631            params![full_key.as_bytes().as_slice()],
632        )?;
633        if inserted == 1 {
634            mark_blob_database_dirty(&self.path);
635        }
636        Ok(())
637    }
638
639    fn ensure_key_plane(&self, full_key: &FullKey) -> Result<(), BlobStoreError> {
640        if full_key.plane() == self.plane {
641            Ok(())
642        } else {
643            Err(BlobStoreError::PlaneKeyMismatch {
644                store_plane: self.plane,
645                key_plane: full_key.plane(),
646            })
647        }
648    }
649}
650
651fn validate_artifact_key(artifact_key: &str) -> Result<(), BlobStoreError> {
652    if artifact_key.is_empty()
653        || artifact_key == "."
654        || artifact_key == ".."
655        || artifact_key.contains(['/', '\\', '\0'])
656    {
657        return Err(BlobStoreError::InvalidArtifactKey(artifact_key.to_owned()));
658    }
659    Ok(())
660}
661
662fn configure_connection(connection: &Connection) -> Result<(), BlobStoreError> {
663    // Set the wait policy before WAL attempts to acquire the journal lock so
664    // concurrent first-open callers wait instead of failing immediately.
665    connection.busy_timeout(Duration::from_millis(BUSY_TIMEOUT_MS))?;
666    connection.pragma_update(None, "foreign_keys", "OFF")?;
667    // Switching a rollback-journal file to WAL needs an exclusive lock, and
668    // SQLite skips the busy handler on that upgrade when another connection
669    // is mid-switch (it would risk a deadlock), so two first-openers racing on
670    // a fresh file see SQLITE_BUSY straight away. Wait it out ourselves within
671    // the same budget the busy handler would have used.
672    retry_while_busy(Duration::from_millis(BUSY_TIMEOUT_MS), || {
673        connection.pragma_update(None, "journal_mode", "WAL")
674    })?;
675    connection.pragma_update(None, "synchronous", "NORMAL")?;
676    Ok(())
677}
678
679/// Run `operation` until it stops failing with SQLITE_BUSY/SQLITE_LOCKED or
680/// `budget` elapses; the last error is returned when the budget runs out.
681/// Shared with the view store, whose pointer database has the same first-open
682/// WAL-switch race.
683pub(crate) fn retry_while_busy<T>(
684    budget: Duration,
685    mut operation: impl FnMut() -> rusqlite::Result<T>,
686) -> rusqlite::Result<T> {
687    let deadline = std::time::Instant::now() + budget;
688    let mut backoff = Duration::from_millis(1);
689    loop {
690        match operation() {
691            Err(rusqlite::Error::SqliteFailure(error, _))
692                if matches!(
693                    error.code,
694                    ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked
695                ) && std::time::Instant::now() < deadline =>
696            {
697                std::thread::sleep(backoff);
698                backoff = (backoff * 2).min(Duration::from_millis(50));
699            }
700            result => return result,
701        }
702    }
703}
704
705fn ensure_schema(connection: &mut Connection) -> Result<(), BlobStoreError> {
706    let tx = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
707    tx.execute_batch(BLOB_SCHEMA)?;
708    let has_created_at = tx
709        .prepare("PRAGMA table_info(blob_payloads)")?
710        .query_map([], |row| row.get::<_, String>(1))?
711        .collect::<Result<Vec<_>, _>>()?
712        .iter()
713        .any(|column| column == "created_at_ms");
714    if !has_created_at {
715        // Existing immutable rows predate per-row timestamps and are eligible for
716        // the normal age-floor policy immediately after this migration.
717        tx.execute(
718            "ALTER TABLE blob_payloads ADD COLUMN created_at_ms INTEGER NOT NULL DEFAULT 0",
719            [],
720        )?;
721    }
722    tx.commit()?;
723    Ok(())
724}
725
726fn unix_millis_now() -> u64 {
727    SystemTime::now()
728        .duration_since(UNIX_EPOCH)
729        .unwrap_or_default()
730        .as_millis() as u64
731}
732
733fn read_and_assert_pragmas(connection: &Connection) -> Result<BlobStorePragmas, BlobStoreError> {
734    let pragmas = BlobStorePragmas {
735        journal_mode: connection.pragma_query_value(None, "journal_mode", |row| row.get(0))?,
736        synchronous: connection.pragma_query_value(None, "synchronous", |row| row.get(0))?,
737        busy_timeout_ms: connection.pragma_query_value(None, "busy_timeout", |row| row.get(0))?,
738        foreign_keys: connection.pragma_query_value(None, "foreign_keys", |row| row.get(0))?,
739        // Do not write this pragma: preserving SQLite's default is part of the
740        // storage contract and querying it catches accidental future overrides.
741        wal_autocheckpoint_pages: connection.pragma_query_value(
742            None,
743            "wal_autocheckpoint",
744            |row| row.get(0),
745        )?,
746    };
747    assert_pragma("journal_mode", "wal", &pragmas.journal_mode)?;
748    assert_pragma("synchronous", "1", &pragmas.synchronous.to_string())?;
749    assert_pragma(
750        "busy_timeout",
751        &BUSY_TIMEOUT_MS.to_string(),
752        &pragmas.busy_timeout_ms.to_string(),
753    )?;
754    assert_pragma("foreign_keys", "0", &pragmas.foreign_keys.to_string())?;
755    assert_pragma(
756        "wal_autocheckpoint",
757        &DEFAULT_WAL_AUTOCHECKPOINT_PAGES.to_string(),
758        &pragmas.wal_autocheckpoint_pages.to_string(),
759    )?;
760    Ok(pragmas)
761}
762
763fn assert_pragma(name: &'static str, expected: &str, actual: &str) -> Result<(), BlobStoreError> {
764    if actual.eq_ignore_ascii_case(expected) {
765        Ok(())
766    } else {
767        Err(BlobStoreError::PragmaMismatch {
768            name,
769            expected: expected.to_owned(),
770            actual: actual.to_owned(),
771        })
772    }
773}
774
775fn is_corrupt_database_error(error: &BlobStoreError) -> bool {
776    matches!(
777        error,
778        BlobStoreError::Sqlite(rusqlite::Error::SqliteFailure(sqlite_error, _))
779            if matches!(sqlite_error.code, ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase)
780    )
781}
782
783fn move_corrupt_database_aside(path: &Path) -> Result<PathBuf, BlobStoreError> {
784    let timestamp = SystemTime::now()
785        .duration_since(UNIX_EPOCH)
786        .unwrap_or_default()
787        .as_secs();
788    let file_name = path
789        .file_name()
790        .and_then(|name| name.to_str())
791        .ok_or_else(|| BlobStoreError::InvalidArtifactKey(path.display().to_string()))?;
792    let destination = path.with_file_name(format!("{file_name}.corrupt-{timestamp}"));
793    fs::rename(path, &destination)?;
794    for suffix in ["-wal", "-shm"] {
795        let sidecar = PathBuf::from(format!("{}{suffix}", path.display()));
796        if sidecar.exists() {
797            fs::rename(
798                &sidecar,
799                PathBuf::from(format!("{}{suffix}", destination.display())),
800            )?;
801        }
802    }
803    Ok(destination)
804}
805
806#[cfg(test)]
807mod tests {
808    use super::*;
809
810    /// Eight racing first-openers on a fresh file trip the WAL-switch BUSY that
811    /// the busy handler does not cover. Without the retry this sees 1-3
812    /// failures per 60 rounds on an idle laptop; 150 rounds make the red
813    /// reliable while keeping the test under two seconds.
814    #[test]
815    fn concurrent_first_opens_never_see_busy() {
816        let mut failures = Vec::new();
817        for round in 0..150 {
818            let dir = tempfile::tempdir().expect("tempdir");
819            let handles: Vec<_> = (0..8)
820                .map(|_| {
821                    let storage = dir.path().to_path_buf();
822                    std::thread::spawn(move || {
823                        BlobStore::open(&storage, "fam", BlobPlane::Semantic).map(|_| ())
824                    })
825                })
826                .collect();
827            for handle in handles {
828                if let Err(error) = handle.join().expect("opener thread") {
829                    failures.push(format!("round {round}: {error}"));
830                }
831            }
832        }
833        assert!(
834            failures.is_empty(),
835            "concurrent first-open must wait out the WAL switch: {failures:?}"
836        );
837    }
838
839    #[test]
840    fn payload_schema_and_producer_version_pairs_are_pinned() {
841        assert_eq!(
842            [
843                (SEMANTIC_PAYLOAD_SCHEMA, SEMANTIC_PRODUCER_VERSION),
844                (CALLGRAPH_PAYLOAD_SCHEMA, CALLGRAPH_PRODUCER_VERSION),
845            ],
846            [(1, "semantic-v1"), (1, "callgraph-v1")],
847            "a payload encoding change must bump its producer key version in the same edit"
848        );
849    }
850
851    #[test]
852    fn semantic_paths_are_distinct_while_callgraph_content_reuses() {
853        let bytes = b"same source";
854        let semantic_a = SemanticKey::for_current(bytes, b"src/a.rs", "model-a").full_key();
855        let semantic_b = SemanticKey::for_current(bytes, b"src/b.rs", "model-a").full_key();
856        let callgraph_a = CallgraphKey::for_current(bytes, "rust").full_key();
857        let callgraph_b = CallgraphKey::for_current(bytes, "rust").full_key();
858
859        assert_ne!(semantic_a, semantic_b);
860        assert_eq!(callgraph_a, callgraph_b);
861    }
862
863    #[test]
864    fn config_is_a_valid_callgraph_language() {
865        let key = CallgraphKey::for_current(b"[package]", "config");
866        assert_ne!(key.source_digest(), &[0; 32]);
867    }
868
869    #[test]
870    fn abandoned_insert_transaction_leaves_no_partial_payload_row() {
871        let directory = tempfile::tempdir().expect("create temporary storage");
872        let mut store = BlobStore::open(directory.path(), "family-a", BlobPlane::Semantic)
873            .expect("open blob store");
874        let key = SemanticKey::for_current(b"source", b"src/lib.rs", "model-a").full_key();
875        let payload = b"payload";
876        let payload_digest = blake3::hash(payload);
877
878        {
879            let tx = store
880                .connection
881                .transaction_with_behavior(TransactionBehavior::Immediate)
882                .expect("start payload transaction");
883            tx.execute(
884                "INSERT INTO blob_payloads (full_key, payload, payload_digest, payload_schema)
885                 VALUES (?1, ?2, ?3, ?4)",
886                params![
887                    key.as_bytes().as_slice(),
888                    payload,
889                    payload_digest.as_bytes().as_slice(),
890                    i64::from(SEMANTIC_PAYLOAD_SCHEMA),
891                ],
892            )
893            .expect("stage payload row");
894            // Dropping an uncommitted SQLite transaction models a process that
895            // dies after staging a row but before the put transaction commits.
896        }
897
898        assert_eq!(store.get(&key).expect("read after aborted put"), None);
899    }
900
901    #[test]
902    fn usage_counts_rows_payloads_and_sqlite_pages() {
903        let directory = tempfile::tempdir().expect("create temporary storage");
904        let mut store = BlobStore::open(directory.path(), "family-a", BlobPlane::Semantic)
905            .expect("open blob store");
906        let key = SemanticKey::for_current(b"source", b"src/lib.rs", "model-a").full_key();
907        store.put(&key, b"payload").expect("insert payload");
908
909        let usage = store.usage().expect("read usage");
910        assert_eq!(usage.rows, 1);
911        assert_eq!(usage.payload_bytes, 7);
912        assert!(usage.file_bytes >= usage.payload_bytes);
913    }
914
915    #[test]
916    fn only_inserted_and_reused_are_durable() {
917        for outcome in [
918            PutOutcome::Inserted,
919            PutOutcome::Reused,
920            PutOutcome::Quarantined,
921            PutOutcome::Failed,
922            PutOutcome::QuotaExceeded,
923        ] {
924            assert_eq!(
925                PutReport::new(outcome).durable,
926                matches!(outcome, PutOutcome::Inserted | PutOutcome::Reused)
927            );
928        }
929    }
930}