refactor: extract out storage logic

Author: georgeh0

rust/core/src/engine/app.rs

@@ -72,14 +72,9 @@ impl<Prof: EngineProfile> App<Prof> {
         let app_reg = AppRegistration::new(name, &env)?;
 
         // TODO: This database initialization logic should happen lazily on first call to `update()`.
-        let db = {
-            let mut wtxn = env.db_env().write_txn()?;
-            let db = env.db_env().create_database(&mut wtxn, Some(name))?;
-            wtxn.commit()?;
-            db
-        };
+        let store = env.create_app_store(name)?;
 
-        let app_ctx = AppContext::new(env, db, app_reg, max_inflight_components);
+        let app_ctx = AppContext::new(env, store, app_reg, max_inflight_components);
         let root_component = Component::new(app_ctx, StablePath::root(), None);
         crate::telemetry::track("app_create");
         Ok(Self { root_component })
@@ -216,12 +211,12 @@ impl<Prof: EngineProfile> App<Prof> {
                 handle.ready().await?;
 
                 // Clear the database
-                let db = root_component.app_ctx().db().clone();
+                let store = root_component.app_ctx().store().clone();
                 root_component
                     .app_ctx()
                     .env()
                     .txn_batcher()
-                    .run(move |wtxn| Ok(db.clear(wtxn)?))
+                    .run(move |wtxn| crate::state_store::ops::clear_all(wtxn, &store))
                     .await?;
 
                 info!("App dropped successfully");

rust/core/src/engine/context.rs

@@ -18,11 +18,12 @@ use crate::state::target_state_path::TargetStatePath;
 use crate::{
     engine::environment::{AppRegistration, Environment},
     state::stable_path::StablePath,
+    state_store::Store,
 };
 
 struct AppContextInner<Prof: EngineProfile> {
     env: Environment<Prof>,
-    db: db_schema::Database,
+    store: Store,
     app_reg: AppRegistration<Prof>,
     id_sequencer_manager: IdSequencerManager,
     inflight_semaphore: Option<Arc<tokio::sync::Semaphore>>,
@@ -53,7 +54,7 @@ pub struct AppContext<Prof: EngineProfile> {
 impl<Prof: EngineProfile> AppContext<Prof> {
     pub fn new(
         env: Environment<Prof>,
-        db: db_schema::Database,
+        store: Store,
         app_reg: AppRegistration<Prof>,
         max_inflight_components: Option<usize>,
     ) -> Self {
@@ -62,7 +63,7 @@ impl<Prof: EngineProfile> AppContext<Prof> {
         Self {
             inner: Arc::new(AppContextInner {
                 env,
-                db,
+                store,
                 app_reg,
                 id_sequencer_manager: IdSequencerManager::new(),
                 inflight_semaphore,
@@ -113,8 +114,8 @@ impl<Prof: EngineProfile> AppContext<Prof> {
         &self.inner.env
     }
 
-    pub fn db(&self) -> &db_schema::Database {
-        &self.inner.db
+    pub fn store(&self) -> &Store {
+        &self.inner.store
     }
 
     pub fn app_reg(&self) -> &AppRegistration<Prof> {
@@ -152,7 +153,7 @@ impl<Prof: EngineProfile> AppContext<Prof> {
         let key = key.unwrap_or(&default_key);
         self.inner
             .id_sequencer_manager
-            .next_id(self.inner.env.txn_batcher(), &self.inner.db, key)
+            .next_id(self.inner.env.txn_batcher(), &self.inner.store, key)
             .await
     }
 }

rust/core/src/engine/environment.rs

@@ -1,6 +1,8 @@
 use crate::{
-    engine::profile::EngineProfile, engine::target_state::TargetStateProviderRegistry,
-    engine::txn_batcher::TxnBatcher, prelude::*,
+    engine::profile::EngineProfile,
+    engine::target_state::TargetStateProviderRegistry,
+    prelude::*,
+    state_store::{ReadTxn, TxnBatcher},
 };
 
 use cocoindex_utils::fingerprint::Fingerprint;
@@ -133,13 +135,26 @@ impl<Prof: EngineProfile> Environment<Prof> {
         &self.inner.db_env
     }
 
+    /// Create the per-app sub-database for this environment and wrap it in a
+    /// `Store`. Hides the underlying LMDB sub-database creation from
+    /// `App::new`.
+    pub fn create_app_store(&self, app_name: &str) -> Result<crate::state_store::Store> {
+        let mut wtxn = self.inner.db_env.write_txn()?;
+        let db = self
+            .inner
+            .db_env
+            .create_database(&mut wtxn, Some(app_name))?;
+        wtxn.commit()?;
+        Ok(crate::state_store::Store::new(db))
+    }
+
     /// Open an LMDB read transaction with automatic retry on `MDB_READERS_FULL`.
     ///
     /// Two-phase strategy:
     /// 1. Retry with a short timeout — handles transient reader slot contention.
     /// 2. If phase 1 times out, call `clear_stale_readers()` to reclaim slots
     ///    from dead processes, then retry indefinitely.
-    pub async fn read_txn(&self) -> Result<heed::RoTxn<'_, heed::WithoutTls>> {
+    pub async fn read_txn(&self) -> Result<ReadTxn<'_>> {
         let db_env = self.db_env();
         let try_read_txn = || async {
             match db_env.read_txn() {
@@ -156,7 +171,7 @@ impl<Prof: EngineProfile> Environment<Prof> {
 
         // Phase 1: short timeout for transient concurrency.
         match retryable::run(&try_read_txn, &LMDB_READ_TXN_RETRY_PHASE1).await {
-            Ok(txn) => return Ok(txn),
+            Ok(txn) => return Ok(ReadTxn::new(txn)),
             Err(e) if !e.is_retryable => return Err(e.into()),
             Err(_) => {}
         }
@@ -168,6 +183,7 @@ impl<Prof: EngineProfile> Environment<Prof> {
         }
         retryable::run(&try_read_txn, &LMDB_READ_TXN_RETRY_PHASE2)
             .await
+            .map(ReadTxn::new)
             .map_err(Into::into)
     }