From 9292ff42a9987328a4e50502a7f4734eeaf98ea2 Mon Sep 17 00:00:00 2001 From: aaltshuler Date: Sat, 4 Jul 2026 17:50:43 +0300 Subject: [PATCH] =?UTF-8?q?refactor(engine):=20one=20dataset-open=20chokep?= =?UTF-8?q?oint=20=E2=80=94=20open=5Fdataset(uri,=20VersionResolution,=20s?= =?UTF-8?q?ession,=20wrapper)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unify the two instrumented openers (open_dataset_tracked latest-only, open_table_dataset pinned-only) into a single instrumentation::open_dataset whose version resolution (Latest | At(v)) is an explicit typed parameter, and route every production Dataset open in the engine through it — including the previously-raw opens in the branch ops (delete/list/force_delete_branch), manifest branch delete, version-metadata resolution, recovery classification/restore, the recovery-audit table, and schema-apply version GC (whose forbidden-api-allow sentinel is now unnecessary and removed). One chokepoint means three properties hold uniformly instead of per-path: record_open feeds the cost probes on every open (the raw sites were invisible to the gates), the per-query IO wrapper is attached with the right class (manifest vs table) everywhere, and the shared per-graph Session parameter exists on every path — threading it into the write-side owners lands separately. Behavior-preserving: every call site keeps its existing version resolution, session argument, and wrapper class; in-source tests keep raw opens (out of scope for the chokepoint by design). --- crates/omnigraph/src/db/manifest.rs | 24 +++++-- crates/omnigraph/src/db/manifest/layout.rs | 4 +- crates/omnigraph/src/db/manifest/metadata.rs | 10 ++- crates/omnigraph/src/db/manifest/namespace.rs | 10 ++- crates/omnigraph/src/db/manifest/recovery.rs | 42 +++++++---- .../src/db/omnigraph/schema_apply.rs | 11 +-- crates/omnigraph/src/db/recovery_audit.rs | 19 ++++- crates/omnigraph/src/instrumentation.rs | 70 ++++++++++--------- crates/omnigraph/src/runtime_cache.rs | 8 ++- crates/omnigraph/src/table_store.rs | 36 +++++++--- 10 files changed, 154 insertions(+), 80 deletions(-) diff --git a/crates/omnigraph/src/db/manifest.rs b/crates/omnigraph/src/db/manifest.rs index e887e45a..f9bfdfb8 100644 --- a/crates/omnigraph/src/db/manifest.rs +++ b/crates/omnigraph/src/db/manifest.rs @@ -256,13 +256,19 @@ impl SubTableEntry { // matches the physical layout the namespace path resolved, without the // per-open `__manifest` scan. let location = table_uri_for_path(root_uri, &self.table_path, self.table_branch.as_deref()); - // Route through the instrumented data-table opener (Fix 3). With no - // session this is exactly the Fix-2 `from_uri(location).with_version`. - // This is the uncached fallback (a snapshot with no read caches); the + // Route through the one opener (Fix 3). With no session this is exactly + // the Fix-2 `from_uri(location).with_version`. This is the uncached + // fallback (a snapshot detached from its graph's read caches); the // cached path (`Snapshot::open` → handle cache) calls the same opener on // a miss with the shared session, so both paths count on the per-query // `table_wrapper`. - crate::instrumentation::open_table_dataset(&location, self.table_version, None).await + crate::instrumentation::open_dataset( + &location, + crate::instrumentation::VersionResolution::At(self.table_version), + None, + crate::instrumentation::table_wrapper(), + ) + .await } } @@ -586,9 +592,13 @@ impl ManifestCoordinator { pub async fn delete_branch(&mut self, name: &str) -> Result<()> { let uri = manifest_uri(&self.root_uri); - let mut ds = Dataset::open(&uri) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let mut ds = crate::instrumentation::open_dataset( + &uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::manifest_wrapper(), + ) + .await?; ds.delete_branch(name) .await .map_err(|e| OmniError::Lance(e.to_string()))?; diff --git a/crates/omnigraph/src/db/manifest/layout.rs b/crates/omnigraph/src/db/manifest/layout.rs index 12894a70..37af5236 100644 --- a/crates/omnigraph/src/db/manifest/layout.rs +++ b/crates/omnigraph/src/db/manifest/layout.rs @@ -21,8 +21,10 @@ pub(crate) fn manifest_uri(root: &str) -> String { pub(super) async fn open_manifest_dataset(root_uri: &str, branch: Option<&str>) -> Result { let uri = manifest_uri(root_uri.trim_end_matches('/')); - let dataset = crate::instrumentation::open_dataset_tracked( + let dataset = crate::instrumentation::open_dataset( &uri, + crate::instrumentation::VersionResolution::Latest, + None, crate::instrumentation::manifest_wrapper(), ) .await?; diff --git a/crates/omnigraph/src/db/manifest/metadata.rs b/crates/omnigraph/src/db/manifest/metadata.rs index d84db34b..5e7f8a49 100644 --- a/crates/omnigraph/src/db/manifest/metadata.rs +++ b/crates/omnigraph/src/db/manifest/metadata.rs @@ -229,9 +229,13 @@ pub(super) async fn table_version_metadata_for_state( version: u64, ) -> Result { let full_path = format!("{}/{}", root_uri.trim_end_matches('/'), table_path); - let ds = Dataset::open(&full_path) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let ds = crate::instrumentation::open_dataset( + &full_path, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let ds = match branch { Some(branch) => ds .checkout_branch(branch) diff --git a/crates/omnigraph/src/db/manifest/namespace.rs b/crates/omnigraph/src/db/manifest/namespace.rs index a684b4d7..b3ab7522 100644 --- a/crates/omnigraph/src/db/manifest/namespace.rs +++ b/crates/omnigraph/src/db/manifest/namespace.rs @@ -107,9 +107,13 @@ impl StagedTableNamespace { } async fn open_head(&self) -> Result { - Dataset::open(&self.table_uri()) - .await - .map_err(|e| OmniError::Lance(e.to_string())) + crate::instrumentation::open_dataset( + &self.table_uri(), + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::manifest_wrapper(), + ) + .await } async fn open_version(&self, version: u64) -> Result { diff --git a/crates/omnigraph/src/db/manifest/recovery.rs b/crates/omnigraph/src/db/manifest/recovery.rs index cc9c501a..2c224a63 100644 --- a/crates/omnigraph/src/db/manifest/recovery.rs +++ b/crates/omnigraph/src/db/manifest/recovery.rs @@ -36,7 +36,6 @@ use std::collections::HashMap; -use lance::Dataset; use serde::{Deserialize, Serialize}; use tracing::warn; @@ -737,9 +736,13 @@ pub(crate) async fn restore_table_to_version( branch: Option<&str>, target_version: u64, ) -> Result<()> { - let head = Dataset::open(table_path) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let head = crate::instrumentation::open_dataset( + table_path, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let head = match branch { Some(b) if b != "main" => head .checkout_branch(b) @@ -1729,9 +1732,13 @@ async fn roll_forward_all( continue; } let dataset_uri = format!("{}/{}", root_uri.trim_end_matches('/'), reg.table_path); - let head_ds = Dataset::open(&dataset_uri) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let head_ds = crate::instrumentation::open_dataset( + &dataset_uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let head_ds = match reg.table_branch.as_deref() { Some(b) if b != "main" => head_ds .checkout_branch(b) @@ -1826,9 +1833,13 @@ async fn push_table_update( updates: &mut Vec, expected: &mut HashMap, ) -> Result { - let ds = Dataset::open(table_path) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let ds = crate::instrumentation::open_dataset( + table_path, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let ds = match branch { Some(b) if b != "main" => ds .checkout_branch(b) @@ -1924,9 +1935,13 @@ pub(crate) async fn has_schema_apply_sidecar( /// HEAD version. Recovery uses this so feature-branch sidecars classify /// against the feature-branch's Lance HEAD, not main's. async fn open_lance_head(table_path: &str, branch: Option<&str>) -> Result { - let ds = Dataset::open(table_path) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let ds = crate::instrumentation::open_dataset( + table_path, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let ds = match branch { Some(b) if b != "main" => ds .checkout_branch(b) @@ -1977,6 +1992,7 @@ pub(crate) fn new_sidecar( mod tests { use super::*; use crate::storage::ObjectStorageAdapter; + use lance::Dataset; use crate::table_store::TableStore; use arrow_array::{Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; diff --git a/crates/omnigraph/src/db/omnigraph/schema_apply.rs b/crates/omnigraph/src/db/omnigraph/schema_apply.rs index 364f5a43..a9638653 100644 --- a/crates/omnigraph/src/db/omnigraph/schema_apply.rs +++ b/crates/omnigraph/src/db/omnigraph/schema_apply.rs @@ -751,10 +751,13 @@ where async fn cleanup_dataset_old_versions(db: &Omnigraph, full_uri: &str) -> Result<()> { use chrono::Utc; use lance::dataset::cleanup::CleanupPolicy; - // forbidden-api-allow: maintenance (Hard-drop version GC) opens the dataset to run cleanup_old_versions. - let ds = lance::Dataset::open(full_uri) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let ds = crate::instrumentation::open_dataset( + full_uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let policy = CleanupPolicy { before_timestamp: Some(Utc::now()), before_version: None, diff --git a/crates/omnigraph/src/db/recovery_audit.rs b/crates/omnigraph/src/db/recovery_audit.rs index 8247c1bc..60137fe8 100644 --- a/crates/omnigraph/src/db/recovery_audit.rs +++ b/crates/omnigraph/src/db/recovery_audit.rs @@ -101,7 +101,14 @@ impl RecoveryAudit { /// with no dataset yet (it is created lazily on the first append). pub(crate) async fn open(root_uri: &str) -> Result { let root = root_uri.trim_end_matches('/').to_string(); - let dataset = Dataset::open(&recoveries_uri(&root)).await.ok(); + let dataset = crate::instrumentation::open_dataset( + &recoveries_uri(&root), + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::manifest_wrapper(), + ) + .await + .ok(); Ok(Self { root_uri: root, dataset, @@ -195,9 +202,15 @@ async fn create_recoveries_dataset(root_uri: &str) -> Result { // variant, not the display string (not a Lance API contract). Same // discipline as `commit_graph.rs`'s create-or-open; pinned by // `lance_surface_guards.rs::lance_error_dataset_already_exists_variant_exists`. - Err(lance::Error::DatasetAlreadyExists { .. }) => Dataset::open(&uri) + Err(lance::Error::DatasetAlreadyExists { .. }) => { + crate::instrumentation::open_dataset( + &uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::manifest_wrapper(), + ) .await - .map_err(|open_err| OmniError::Lance(open_err.to_string())), + } Err(err) => Err(OmniError::Lance(err.to_string())), } } diff --git a/crates/omnigraph/src/instrumentation.rs b/crates/omnigraph/src/instrumentation.rs index 861592ec..0f693d29 100644 --- a/crates/omnigraph/src/instrumentation.rs +++ b/crates/omnigraph/src/instrumentation.rs @@ -252,47 +252,49 @@ pub(crate) fn record_scan_staged_combined() { }); } -/// Open a Lance dataset at `uri`, attaching `wrapper` (for IO counting) when -/// present. With no wrapper this is exactly `Dataset::open(uri)`. The wrapper is -/// set via `ObjectStoreParams` on the builder so the open itself is counted -/// (`Dataset::with_object_store_wrappers` only wraps an already-open store). -pub(crate) async fn open_dataset_tracked( +/// Which version [`open_dataset`] resolves. +/// +/// `Latest` re-resolves the dataset's current head (the substrate's cheap +/// latest-location probe); `At(v)` is a list-free pinned open. The choice is +/// a correctness decision — strict read-modify-write ops need `Latest`, +/// snapshot reads need `At(v)` — so it is an explicit parameter of the one +/// opener rather than a property of which helper a caller happened to reach. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum VersionResolution { + Latest, + At(u64), +} + +/// THE dataset-open chokepoint. Every engine `Dataset` open routes through +/// here so three things hold uniformly, on every path: +/// +/// 1. `record_open` feeds the per-query cost probes — an open that bypasses +/// this function is invisible to the cost gates. +/// 2. The per-query IO `wrapper` (manifest- or table-class) is set via +/// `ObjectStoreParams` on the builder, so the open itself is counted +/// (`Dataset::with_object_store_wrappers` only wraps an already-open +/// store). No wrapper (production) adds nothing. +/// 3. The shared per-graph `Session` (LanceDB's one-session-per-connection +/// pattern; warms Lance's metadata/index caches across opens) is attached +/// whenever the caller has one. `None` is for genuinely session-less +/// contexts (a `Snapshot` detached from its graph's read caches) — owners +/// that hold a session (`TableStore`, the handle cache) pass it +/// unconditionally, so it cannot be silently dropped on one path. +pub(crate) async fn open_dataset( uri: &str, + version: VersionResolution, + session: Option<&Arc>, wrapper: Option>, ) -> Result { record_open(uri); - let result = match wrapper { - None => Dataset::open(uri).await, - Some(wrapper) => { - DatasetBuilder::from_uri(uri) - .with_store_params(ObjectStoreParams { - object_store_wrapper: Some(wrapper), - ..Default::default() - }) - .load() - .await - } - }; - result.map_err(|e| OmniError::Lance(e.to_string())) -} - -/// Open a data-table dataset at `location` pinned to `version` — the cache-miss -/// path of the data-read boundary (`SubTableEntry::open`). Attaches the shared -/// per-graph `Session` (warms metadata/index caches across opens, LanceDB's -/// one-session-per-connection pattern) and the per-query `table_wrapper` (for IO -/// counting) when present. With neither, this is exactly the Fix-2 -/// `from_uri(location).with_version(version)` open. -pub(crate) async fn open_table_dataset( - location: &str, - version: u64, - session: Option<&Arc>, -) -> Result { - record_open(location); - let mut builder = DatasetBuilder::from_uri(location).with_version(version); + let mut builder = DatasetBuilder::from_uri(uri); + if let VersionResolution::At(version) = version { + builder = builder.with_version(version); + } if let Some(session) = session { builder = builder.with_session(session.clone()); } - if let Some(wrapper) = table_wrapper() { + if let Some(wrapper) = wrapper { builder = builder.with_store_params(ObjectStoreParams { object_store_wrapper: Some(wrapper), ..Default::default() diff --git a/crates/omnigraph/src/runtime_cache.rs b/crates/omnigraph/src/runtime_cache.rs index 7023a925..84e5de8c 100644 --- a/crates/omnigraph/src/runtime_cache.rs +++ b/crates/omnigraph/src/runtime_cache.rs @@ -264,7 +264,13 @@ impl TableHandleCache { // Miss: open without holding the lock (the open is async IO). A concurrent // double-miss opens twice and one wins the insert — correct (the dataset // at a version is immutable) and rare. - let ds = crate::instrumentation::open_table_dataset(location, version, session).await?; + let ds = crate::instrumentation::open_dataset( + location, + crate::instrumentation::VersionResolution::At(version), + session, + crate::instrumentation::table_wrapper(), + ) + .await?; let mut inner = self.inner.lock().await; if let Some(existing) = inner.entries.get(&key).cloned() { return Ok(existing); diff --git a/crates/omnigraph/src/table_store.rs b/crates/omnigraph/src/table_store.rs index 4a695704..166b3733 100644 --- a/crates/omnigraph/src/table_store.rs +++ b/crates/omnigraph/src/table_store.rs @@ -202,12 +202,14 @@ impl TableStore { dataset_uri: &str, branch: Option<&str>, ) -> Result { - // Direct open by URI (O(1) latest-resolution). Routed through the tracked + // Direct open by URI (O(1) latest-resolution). Routed through the one // opener so a cost test counts it via the per-query `table_wrapper` // (no-op in production — the task-local is unset, so this is exactly // `Dataset::open(uri)`). - let ds = crate::instrumentation::open_dataset_tracked( + let ds = crate::instrumentation::open_dataset( dataset_uri, + crate::instrumentation::VersionResolution::Latest, + None, crate::instrumentation::table_wrapper(), ) .await?; @@ -237,9 +239,13 @@ impl TableStore { } pub async fn delete_branch(&self, dataset_uri: &str, branch: &str) -> Result<()> { - let mut ds = Dataset::open(dataset_uri) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let mut ds = crate::instrumentation::open_dataset( + dataset_uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; ds.delete_branch(branch) .await .map_err(|e| OmniError::Lance(e.to_string())) @@ -250,9 +256,13 @@ impl TableStore { /// set to find orphaned per-table forks. `main`/default is not a named /// branch and never appears here. pub async fn list_branches(&self, dataset_uri: &str) -> Result> { - let ds = Dataset::open(dataset_uri) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let ds = crate::instrumentation::open_dataset( + dataset_uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; let branches = ds .list_branches() .await @@ -274,9 +284,13 @@ impl TableStore { /// Used by the eager best-effort reclaim in `cleanup_deleted_branch_tables` /// and the `cleanup` orphan reconciler. pub async fn force_delete_branch(&self, dataset_uri: &str, branch: &str) -> Result<()> { - let mut ds = Dataset::open(dataset_uri) - .await - .map_err(|e| OmniError::Lance(e.to_string()))?; + let mut ds = crate::instrumentation::open_dataset( + dataset_uri, + crate::instrumentation::VersionResolution::Latest, + None, + crate::instrumentation::table_wrapper(), + ) + .await?; match ds.force_delete_branch(branch).await { Ok(()) => Ok(()), Err(lance::Error::RefNotFound { .. }) | Err(lance::Error::NotFound { .. }) => Ok(()),