Keep prepared search readers warm and verify abrupt indexing recovery

crates/searchi-service/src/store.rs

@@ -45,7 +45,8 @@ pub(crate) struct Store {
     connection: Connection,
     lock_path: PathBuf,
     dimensions: usize,
-    ann_cache: Option<(i64, Index)>,
+    ann_cache: Option<(i64, Arc<Index>)>,
+    prepared_graphs: crate::PreparedGraphs,
     eligibility_cache: Option<EligibilityCache>,
 }
 
@@ -68,6 +69,19 @@ thread_local! { static HASH_CALLS:std::cell::Cell<usize>=const {std::cell::Cell:
 thread_local! { static MEDIA_STAT_CALLS:std::cell::Cell<usize>=const {std::cell::Cell::new(0)}; }
 
 impl Store {
+    pub(crate) fn share_prepared_graphs(&mut self, graphs: &crate::PreparedGraphs) -> Result<()> {
+        graphs.bind(
+            Path::new(
+                self.connection
+                    .path()
+                    .context("Database path unavailable")?,
+            ),
+            self.dimensions,
+        )?;
+        self.prepared_graphs = graphs.clone();
+        Ok(())
+    }
+
     pub(crate) fn active_fingerprint(&self) -> Result<Option<String>> {
         Ok(self.connection.query_row("SELECT g.fingerprint FROM semantic_meta m JOIN semantic_generations g ON g.id=m.active_generation WHERE m.singleton=1",[],|r|r.get(0)).optional()?)
     }
@@ -170,6 +184,7 @@ impl Store {
             lock_path,
             dimensions,
             ann_cache: None,
+            prepared_graphs: crate::PreparedGraphs::default(),
             eligibility_cache: None,
         })
     }
@@ -328,6 +343,16 @@ impl Store {
         // before acquiring the write transaction. The revision check below
         // still closes races with indexed changes after this validation.
         validate_sources()?;
+        let fingerprint: String = self.connection.query_row(
+            "SELECT fingerprint FROM semantic_generations WHERE id=?1",
+            [generation],
+            |row| row.get(0),
+        )?;
+        let ann = Arc::new(ann);
+        // Readers observing SQL publication already have a prepared immutable graph.
+        let publication =
+            self.prepared_graphs
+                .publish(generation, fingerprint, checksum.clone(), ann.clone())?;
         let transaction = self
             .connection
             .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
@@ -366,6 +391,7 @@ impl Store {
         transaction.execute_batch("UPDATE semantic_meta SET dirty_floor=max(dirty_floor,corpus_revision-100000) WHERE singleton=1;
             DELETE FROM semantic_dirty_paths WHERE revision<(SELECT dirty_floor FROM semantic_meta WHERE singleton=1);")?;
         transaction.commit()?;
+        publication.commit();
         self.ann_cache = Some((generation, ann));
         self.eligibility_cache = None;
         Ok(())
@@ -505,6 +531,23 @@ impl Store {
             return Ok(Vec::new());
         };
         trace.end(generation_span, Default::default());
+        if !self
+            .ann_cache
+            .as_ref()
+            .is_some_and(|(generation, _)| *generation == active)
+        {
+            let prepared = trace.timed("semantic.prepared_graph_lookup", parent, |_, _| {
+                let checksum: String = snapshot.query_row(
+                    "SELECT ann_sha256 FROM semantic_generations WHERE id=?1",
+                    [active],
+                    |row| row.get(0),
+                )?;
+                self.prepared_graphs.get(active, fingerprint, &checksum)
+            })?;
+            if let Some(graph) = prepared {
+                self.ann_cache = Some((active, graph));
+            }
+        }
         let ann_load = trace.begin("semantic.ann_load", parent);
         trace.counter(
             ann_load,
@@ -553,7 +596,7 @@ impl Store {
             ann.reserve_capacity_and_threads(ann.size().max(1), 1)?;
             trace.counter(deserialize_span, "chunks", ann.size() as u64);
             trace.end(deserialize_span, Default::default());
-            self.ann_cache = Some((active, ann));
+            self.ann_cache = Some((active, Arc::new(ann)));
         }
         trace.end(ann_load, Default::default());
         let filter_span = trace.begin("semantic.filter_normalization", parent);
@@ -1417,6 +1460,94 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn failed_sql_activation_removes_prepared_pending_and_preserves_live_reader() -> Result<()> {
+        // Storage publication invariant only; real Google coverage is a separate integration.
+        let directory = tempfile::tempdir()?;
+        let root = directory.path().join("docs");
+        std::fs::create_dir(&root)?;
+        std::fs::write(root.join("notes.txt"), "original relevant content")?;
+        let database = directory.path().join("index.db");
+        let mut lexical = SearchIndex::open(&database)?;
+        lexical.index_folder(&root)?;
+        let passage = lexical.passage_batch(0, 1)?.remove(0);
+        let windows = [TextWindow {
+            text: passage.text.clone(),
+            start_byte: 0,
+            end_byte: passage.text.len(),
+        }];
+        let mut vector = vec![0.0; DIMENSIONS];
+        vector[0] = 1.0;
+        let policy = "test-storage-publication";
+        let graphs = crate::PreparedGraphs::default();
+        let mut writer = Store::open(&database, DIMENSIONS)?;
+        writer.share_prepared_graphs(&graphs)?;
+        let active = writer.begin(policy)?;
+        writer.insert(active, &passage, 0, &windows, &[vector.clone()])?;
+        writer.activate(active)?;
+        let checksum: String = writer.connection.query_row(
+            "SELECT ann_sha256 FROM semantic_generations WHERE id=?1",
+            [active],
+            |r| r.get(0),
+        )?;
+        let mut reader = Store::open(&database, DIMENSIONS)?;
+        reader.share_prepared_graphs(&graphs)?;
+        assert_eq!(
+            reader
+                .search(&vector, policy, 10, &SearchOptions::default())?
+                .len(),
+            1
+        );
+        let retained = reader.ann_cache.as_ref().unwrap().1.clone();
+        let pending = writer.begin(policy)?;
+        writer.insert(pending, &passage, 0, &windows, &[vector.clone()])?;
+        let prepared = writer.build_ann(pending)?;
+        let mut serialized = vec![0; prepared.serialized_length()];
+        prepared.save_to_buffer(&mut serialized)?;
+        let pending_checksum = format!("{:x}", Sha256::digest(&serialized));
+        writer.connection.execute_batch("CREATE TRIGGER reject_activation BEFORE UPDATE OF active_generation ON semantic_meta BEGIN SELECT RAISE(ABORT,'activation rejected'); END")?;
+        assert!(writer
+            .activate(pending)
+            .unwrap_err()
+            .to_string()
+            .contains("activation rejected"));
+        assert_eq!(
+            writer.connection.query_row(
+                "SELECT active_generation FROM semantic_meta WHERE singleton=1",
+                [],
+                |r| r.get::<_, i64>(0)
+            )?,
+            active
+        );
+        assert!(graphs.get(pending, policy, &pending_checksum)?.is_none());
+        assert!(Arc::ptr_eq(
+            &retained,
+            &graphs.get(active, policy, &checksum)?.unwrap()
+        ));
+        assert_eq!(
+            reader
+                .search(&vector, policy, 10, &SearchOptions::default())?
+                .len(),
+            1
+        );
+        assert!(Arc::ptr_eq(
+            &retained,
+            &reader.ann_cache.as_ref().unwrap().1
+        ));
+        writer
+            .connection
+            .execute_batch("DROP TRIGGER reject_activation")?;
+        writer.activate(pending)?;
+        assert_eq!(
+            reader
+                .search(&vector, policy, 10, &SearchOptions::default())?
+                .len(),
+            1
+        );
+        assert_eq!(reader.ann_cache.as_ref().unwrap().0, pending);
+        Ok(())
+    }
+
     #[test]
     fn generation_persists_failed_staging_preserves_active_and_changes_hide_stale_vectors(
     ) -> Result<()> {