From 6a70c344f0b1eb5e3e6ee5d6bddebcbd14822d23 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 19 May 2026 15:45:32 +0800 Subject: [PATCH 1/6] feat(jni,io): default-open Session reuse + ObjectStoreRegistry single-flight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Coalesce concurrent cold `Dataset.open` calls in the JNI layer so 144 parallel opens against the same URI reuse one `ObjectStore` instead of each rebuilding the credential chain and HTTP client. `BlockingDataset::open` without an explicit `Session` now constructs a fresh per-call `Session` whose `ObjectStoreRegistry` is the process-wide `GLOBAL_OBJECT_STORE_REGISTRY` (a `LazyLock>`). Sharing only the registry — not the `Session` — keeps the asymmetric caching shape: the registry coalesces expensive `ObjectStore` builds (credential probe, IMDS, TLS handshake), while metadata/index caches remain per-call so callers do not share eviction policy or cache size with other tenants. Bare-URI opens (empty storage options, no provider, no namespace commit handler) collapse onto a single cache entry per URI: the first caller's resolved default-credential chain becomes the credentials used by every subsequent caller for the lifetime of that `Arc`. This cross-tenant credential bleed is intentional under the bare-URI invariant; callers that need per-tenant isolation can opt out via `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING={1,true,yes}`, pass tenant-distinguishing storage options, or supply an explicit `Session`. The opt-out env var is parsed once into an `AtomicU8` tri-state (`UNINIT/ENABLED/DISABLED`) on first read, with a truthy whitelist (`1|true|yes` — anything else falls back to enabled and warns). A `DisableSharingTestGuard` RAII helper resets the atomic for unit tests so the env-var probe can be exercised deterministically. `ObjectStoreRegistry::get_store` adds a per-key build-lock so concurrent cold builds for the same `CacheKey` serialize behind one in-flight `build()` instead of all racing the credential chain. The lock map is `Mutex>>>`; each caller acquires its slot via a `BuildLockSession` RAII guard whose `Drop` releases the inner mutex first, then opportunistically removes the HashMap entry if it observes `strong_count == 1`. This ordering guarantees that panics, errors, and cancellations all reclaim the build-lock entry without leaking a permanently-locked key. Failures are not cached — when a cold build returns `Err`, the next waiter retries from scratch (one-Err-per-N-waiters retry contract). An opportunistic sweep gated by `cold_attempt_nth.is_multiple_of(64)` amortizes HashMap shrinkage across cold-path attempts so a long-lived process does not accumulate stale entries from short-lived URIs. The `CacheKey` uses `sanitized_authority` derived from `Url::host()` (IPv6 brackets preserved, port preserved, userinfo stripped) so URIs that differ only in embedded credentials still collapse onto the same cache slot. - `lance-io`: hit/miss, single-flight coalescing, panic-path RAII cleanup, error-path RAII cleanup, `sanitized_authority` over IPv6 / userinfo / non-default port, opportunistic sweep counter. - `lance-jni`: `parse_disable_value` truthy whitelist + warn path, `DisableSharingTestGuard` ordering, end-to-end coalescing test that drives 144 concurrent default-opens through the JNI layer and asserts a single `ObjectStore` is built. `java/README.md` documents the `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING` env var, the bare-URI credential-bleed invariant, and the three opt-out paths (env var, tenant-distinguishing storage options, explicit `Session`). --- java/README.md | 10 + java/lance-jni/Cargo.lock | 1 + java/lance-jni/Cargo.toml | 5 + java/lance-jni/src/blocking_dataset.rs | 345 ++++++++++ java/lance-jni/src/lib.rs | 69 ++ rust/lance-io/src/object_store/providers.rs | 675 ++++++++++++++++++-- 6 files changed, 1052 insertions(+), 53 deletions(-) diff --git a/java/README.md b/java/README.md index b49a4892527..5f4c7f70ef0 100644 --- a/java/README.md +++ b/java/README.md @@ -212,6 +212,16 @@ JVM engine connectors can be built using the Lance Java SDK. Here are some conne * [Flink Lance connector](https://github.com/lancedb/lance-flink) * [Trino Lance connector](https://github.com/lancedb/lance-trino) +## Configuration + +Environment variables read by the JNI layer at runtime: + +| Variable | Default | Effect | +|---|---|---| +| `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING` | unset (sharing on) | Set to `1`/`true`/`yes` to give every default `Dataset.open` (no explicit `Session`) a fresh `ObjectStoreRegistry`. Disables the process-wide registry that coalesces concurrent cold builds for the same URI. | + +Default-open paths share an `ObjectStoreRegistry` so concurrent opens against the same URI reuse one `ObjectStore` instead of each rebuilding the credential chain and HTTP client. Bare-URI opens (empty storage options, no provider, no namespace commit handler) collapse onto a single cache entry per URI: the first caller's resolved default credentials become the credentials used by every subsequent caller for the lifetime of that `ObjectStore`. Callers that need cross-tenant isolation under bare URIs should either pass an explicit `Session`, supply tenant-distinguishing storage options, or set the variable above to opt out entirely. + ## Contributing From the codebase dimension, the lance project is a multiple-lang project. All Java-related code is located in the `java` directory. diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 5be46e293d2..f395d562027 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4103,6 +4103,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "url", "uuid", ] diff --git a/java/lance-jni/Cargo.toml b/java/lance-jni/Cargo.toml index c51df5769c7..18f36475528 100644 --- a/java/lance-jni/Cargo.toml +++ b/java/lance-jni/Cargo.toml @@ -59,6 +59,11 @@ roaring = "0.11.4" prost-types = "0.14.1" chrono = "0.4.41" +[dev-dependencies] +# Test-only: builds Url values for the JNI default-open coalescing test in +# src/blocking_dataset.rs. Workspace pins this same version transitively. +url = "2.5.7" + [profile.dev] debug = "line-tables-only" incremental = false diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index caf837b371a..1ed20a26ba4 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -68,6 +68,155 @@ use uuid::Uuid; pub const NATIVE_DATASET: &str = "nativeDatasetHandle"; +/// Reads `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING` once at first access and +/// caches the result. +/// +/// Only `"1" | "true" | "yes"` (case-insensitive) disable sharing. Empty / unset +/// keeps sharing enabled — `export VAR=` does NOT activate the escape hatch. +/// Logs a warning the first time the value is read when sharing is disabled, +/// and a separate warning when the variable is set to a value the parser does +/// not recognize. +/// +/// Backed by `AtomicU8` rather than `OnceLock` so unit tests can override +/// the resolved value via `DisableSharingTestGuard` without spawning a fresh +/// process per case. +/// +/// State encoding: +/// - 0 = uninitialized (read env on next access) +/// - 1 = sharing enabled +/// - 2 = sharing disabled +const SHARING_UNINIT: u8 = 0; +const SHARING_ENABLED: u8 = 1; +const SHARING_DISABLED: u8 = 2; + +static DISABLE_DEFAULT_REGISTRY_SHARING: std::sync::atomic::AtomicU8 = + std::sync::atomic::AtomicU8::new(SHARING_UNINIT); + +fn disable_default_registry_sharing() -> bool { + use std::sync::atomic::Ordering; + match DISABLE_DEFAULT_REGISTRY_SHARING.load(Ordering::Relaxed) { + SHARING_ENABLED => false, + SHARING_DISABLED => true, + _ => { + // First read: resolve from env, warn if appropriate, then publish. + // Under a race multiple threads may each observe UNINIT and emit + // duplicate `log::warn!` lines — the env read is idempotent so the + // resolved bool is the same. Cheap duplication is the price of an + // `AtomicU8` that tests can override. + let raw = std::env::var("LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING").ok(); + let disabled = match raw.as_deref().map(parse_disable_value) { + Some(Some(parsed)) => parsed, + Some(None) => { + // Set but not recognized — warn so a misspelled escape hatch + // (e.g. `=on`, `=y`) surfaces instead of silently keeping + // sharing enabled and looking like the env var "did nothing". + log::warn!( + "LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING={:?} is unrecognized; \ + keeping sharing enabled. Use 1/true/yes to disable.", + raw.as_deref().unwrap_or(""), + ); + false + } + None => false, + }; + if disabled { + log::warn!( + "LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING is set; JNI default-open will \ + use a fresh ObjectStoreRegistry per call (single-flight coalescing \ + disabled)." + ); + } + DISABLE_DEFAULT_REGISTRY_SHARING.store( + if disabled { + SHARING_DISABLED + } else { + SHARING_ENABLED + }, + Ordering::Relaxed, + ); + disabled + } + } +} + +/// Parse a `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING` env value. +/// +/// Returns (after trim + ASCII lowercase): +/// - `Some(true)` for `1`/`true`/`yes` — operator opted out of sharing. +/// - `Some(false)` for empty, whitespace-only, `0`/`false`/`no` — explicit keep-default. +/// - `None` for anything else (typos, `on`, `y`, numeric noise) — caller +/// should warn and fall back to default to keep misconfigurations visible. +/// +/// Pure helper extracted from [`disable_default_registry_sharing`] so the +/// truthiness rules can be unit-tested without manipulating process env vars. +fn parse_disable_value(raw: &str) -> Option { + // `trim()` canonicalizes whitespace-only input to `""`, so the empty-string + // arm catches both literal `""` and ` `/`\n` etc. — keeping the doc + // contract honest with a single match arm. + match raw.trim().to_ascii_lowercase().as_str() { + "1" | "true" | "yes" => Some(true), + "" | "0" | "false" | "no" => Some(false), + _ => None, + } +} + +/// Pick the `ObjectStoreRegistry` for the JNI default-open path. +/// +/// When sharing is enabled (the default), every default-open call reuses the +/// process-wide [`crate::GLOBAL_OBJECT_STORE_REGISTRY`] so concurrent cold +/// builds for the same URI coalesce on its single-flight. When the +/// `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING` escape hatch is set, each call +/// gets a fresh registry — pre-PR isolation at the cost of single-flight. +/// +/// Extracted into a pure helper so the selection logic is unit-testable +/// without spinning up a JVM. +fn select_default_open_registry( + disable_sharing: bool, +) -> Arc { + if disable_sharing { + Arc::new(lance_io::object_store::ObjectStoreRegistry::default()) + } else { + crate::GLOBAL_OBJECT_STORE_REGISTRY.clone() + } +} + +#[cfg(test)] +/// RAII guard that restores `DISABLE_DEFAULT_REGISTRY_SHARING` to its prior +/// value on drop. Tests should always go through this guard so they cannot +/// leak state into sibling tests run in the same process. +/// +/// Each guard captures and restores its own snapshot, so nested guards compose +/// correctly: LIFO drop order means the outer guard's snapshot always wins, +/// returning the flag to the value seen before the outermost `set` call. +struct DisableSharingTestGuard { + prior: u8, +} + +#[cfg(test)] +impl DisableSharingTestGuard { + fn set(disabled: bool) -> Self { + use std::sync::atomic::Ordering; + let prior = DISABLE_DEFAULT_REGISTRY_SHARING.load(Ordering::Relaxed); + DISABLE_DEFAULT_REGISTRY_SHARING.store( + if disabled { + SHARING_DISABLED + } else { + SHARING_ENABLED + }, + Ordering::Relaxed, + ); + Self { prior } + } +} + +#[cfg(test)] +impl Drop for DisableSharingTestGuard { + fn drop(&mut self) { + use std::sync::atomic::Ordering; + DISABLE_DEFAULT_REGISTRY_SHARING.store(self.prior, Ordering::Relaxed); + } +} + impl FromJObjectWithEnv for JObject<'_> { fn extract_object(&self, env: &mut JNIEnv<'_>) -> Result { let id = env.get_u32_from_method(self, "getId")?; @@ -178,6 +327,21 @@ impl BlockingDataset { storage_options_accessor: accessor, ..Default::default() }; + + // Default-open path: share the process-wide registry so concurrent + // opens for the same URI coalesce on single-flight; each call still + // gets its own `Session`. Tenant-isolation contract is documented on + // [`crate::GLOBAL_OBJECT_STORE_REGISTRY`]; opt out via + // `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING=1`. + let session = session.or_else(|| { + let registry = select_default_open_registry(disable_default_registry_sharing()); + Some(Arc::new(LanceSession::new( + index_cache_size_bytes as usize, + metadata_cache_size_bytes as usize, + registry, + ))) + }); + let params = ReadParams { index_cache_size_bytes: index_cache_size_bytes as usize, metadata_cache_size_bytes: metadata_cache_size_bytes as usize, @@ -3723,3 +3887,184 @@ fn inner_get_zonemap_stats<'local>( Ok(array_list) } + +#[cfg(test)] +mod default_open_registry_tests { + use super::*; + + /// Default-open path must reuse the process-wide + /// `GLOBAL_OBJECT_STORE_REGISTRY` so concurrent opens can coalesce on its + /// single-flight. Compare by `Arc::ptr_eq` — a fresh registry would have + /// a distinct allocation. + #[test] + fn select_default_open_registry_reuses_global_when_sharing_enabled() { + let registry = select_default_open_registry(false); + assert!( + Arc::ptr_eq(®istry, &crate::GLOBAL_OBJECT_STORE_REGISTRY), + "sharing-enabled path must hand back the GLOBAL_OBJECT_STORE_REGISTRY Arc" + ); + } + + /// Escape-hatch path must hand back a *fresh* registry per call so the + /// pre-PR isolation behavior is preserved when an operator opts out. + #[test] + fn select_default_open_registry_returns_fresh_when_disabled() { + let r1 = select_default_open_registry(true); + let r2 = select_default_open_registry(true); + assert!( + !Arc::ptr_eq(&r1, &crate::GLOBAL_OBJECT_STORE_REGISTRY), + "disabled path must NOT alias the global registry" + ); + assert!( + !Arc::ptr_eq(&r1, &r2), + "disabled path must allocate a new registry per call" + ); + } + + /// `disable_default_registry_sharing()` must honor the in-process override + /// hook. This guards against future refactors that re-introduce a once-cell + /// that bypasses the AtomicU8 state. + #[test] + fn disable_flag_honors_test_override() { + // RAII guard restores prior state at end-of-scope so this test cannot + // leak SHARING_ENABLED/DISABLED state into sibling tests that rely on + // env-driven first-read resolution. + let _g = DisableSharingTestGuard::set(true); + assert!(disable_default_registry_sharing()); + + let _g2 = DisableSharingTestGuard::set(false); + assert!(!disable_default_registry_sharing()); + } + + /// Truthy spellings (case + whitespace insensitive) must parse to `Some(true)`. + /// These are the values that opt the operator OUT of registry sharing — + /// getting any of them wrong silently re-enables coalescing in a config + /// that asked for isolation, so accept-list correctness is load-bearing. + #[test] + fn parse_disable_value_accepts_truthy() { + for raw in [ + "1", "true", "yes", "TRUE", "Yes", "True", " 1 ", "\ttrue\n", + ] { + assert_eq!( + parse_disable_value(raw), + Some(true), + "expected {raw:?} to parse as Some(true)", + ); + } + } + + /// Documented falsy spellings — empty, whitespace-only, `0`/`false`/`no` — + /// must parse to `Some(false)`. These are explicit keep-default opts; + /// distinguishing them from unrecognized noise lets the caller skip the + /// warning log. + #[test] + fn parse_disable_value_accepts_falsy() { + for raw in ["", " ", "\t\n", "0", "false", "no", "FALSE", " No "] { + assert_eq!( + parse_disable_value(raw), + Some(false), + "expected {raw:?} to parse as Some(false)", + ); + } + } + + /// Anything outside the accept-list — typos, `on`, `y`, numeric noise — + /// must parse to `None` so [`disable_default_registry_sharing`] can warn + /// the operator that their escape hatch is being silently ignored. + #[test] + fn parse_disable_value_returns_none_for_unrecognized() { + for raw in ["off", "disable", "2", "1.0", "true!", "y", "on", "enable"] { + assert_eq!( + parse_disable_value(raw), + None, + "expected {raw:?} to parse as None (unrecognized)", + ); + } + } + + /// End-to-end pin for PR#1's two-piece contract at the JNI boundary: + /// `select_default_open_registry` must hand back a registry whose + /// `get_store` coalesces concurrent cold builds for the same URI. + /// A regression in single-flight (dropped from `get_store`) would + /// surface here as `provider.builds > 1` or callers observing + /// different `Arc` instances. + /// + /// The complementary "sharing-on returns the process-wide registry" + /// invariant is covered by + /// `select_default_open_registry_reuses_global_when_sharing_enabled`, + /// so this test deliberately uses the sharing-disabled path to take a + /// fresh registry — registering a one-off provider on + /// `GLOBAL_OBJECT_STORE_REGISTRY` would leak that scheme's `Arc` for + /// the rest of the test binary's lifetime. + /// + /// Provider counts invocations and sleeps before returning so concurrent + /// callers reliably queue on the build lock before the first build + /// completes — without the delay the winner returns before any contention + /// develops and the test would pass even if single-flight were broken. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn jni_default_open_coalesces_concurrent_cold_builds() { + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::Duration; + + use lance_core::Result as LanceResult; + use lance_io::object_store::ObjectStore as LanceObjectStore; + use lance_io::object_store::providers::memory::MemoryStoreProvider; + use lance_io::object_store::{ObjectStoreParams, ObjectStoreProvider}; + use url::Url; + + #[derive(Debug, Default)] + struct CountingProvider { + builds: AtomicU64, + } + + #[async_trait::async_trait] + impl ObjectStoreProvider for CountingProvider { + async fn new_store( + &self, + base_path: Url, + params: &ObjectStoreParams, + ) -> LanceResult { + self.builds.fetch_add(1, Ordering::Relaxed); + tokio::time::sleep(Duration::from_millis(20)).await; + MemoryStoreProvider.new_store(base_path, params).await + } + } + + let provider = Arc::new(CountingProvider::default()); + // Take a fresh registry via the sharing-disabled path so the + // test-only `pr1-jni-coalesce` provider never touches + // `GLOBAL_OBJECT_STORE_REGISTRY`. + let registry = select_default_open_registry(true); + registry.insert("pr1-jni-coalesce", provider.clone()); + + let url = Url::parse("pr1-jni-coalesce://x").unwrap(); + let params = ObjectStoreParams::default(); + + let n = 8; + let mut handles = Vec::with_capacity(n); + for _ in 0..n { + let registry = registry.clone(); + let url = url.clone(); + let params = params.clone(); + handles.push(tokio::spawn(async move { + registry.get_store(url, ¶ms).await.unwrap() + })); + } + let mut stores = Vec::with_capacity(n); + for h in handles { + stores.push(h.await.expect("task must not panic")); + } + + assert_eq!( + provider.builds.load(Ordering::Relaxed), + 1, + "single-flight must collapse N concurrent cold misses into one provider call", + ); + for store in &stores[1..] { + assert!( + Arc::ptr_eq(&stores[0], store), + "every coalesced waiter must observe the same Arc", + ); + } + } +} diff --git a/java/lance-jni/src/lib.rs b/java/lance-jni/src/lib.rs index 37eeff66693..87d33d3f94b 100644 --- a/java/lance-jni/src/lib.rs +++ b/java/lance-jni/src/lib.rs @@ -85,6 +85,75 @@ pub static RT: LazyLock = LazyLock::new(|| { .expect("Failed to create tokio runtime") }); +/// Process-wide [`lance_io::object_store::ObjectStoreRegistry`] used for JNI +/// default-open paths. +/// +/// When the Java caller does not supply an explicit session, the JNI open +/// path constructs a per-call session that shares this registry. Sharing the +/// registry across calls allows the registry's per-key single-flight to +/// coalesce concurrent cold builds for the same URI, and lets long-lived +/// `ObjectStore` strong references be reused across opens — both of which +/// turn what would otherwise be a thundering herd into a cheap weak-Arc +/// upgrade. +/// +/// # Why the registry is shared but the `Session` is not +/// +/// The JNI default-open path is intentionally asymmetric: the +/// `ObjectStoreRegistry` is process-global, but each open builds a fresh +/// `Session` (which owns the metadata/index caches). This shape is chosen +/// because the two layers cache fundamentally different things: +/// +/// - **Registry → `Arc`**: an HTTP/S3 client, credential chain, +/// and connection pool. Building one is the *expensive* operation +/// (credential probe, IMDS round-trip, TLS handshake) — so this is what +/// the 144-concurrent-open regression was made of, and what the global +/// registry exists to coalesce. +/// +/// The cache key is derived from the provider-specific store prefix +/// (typically scheme + authority — e.g. `s3://bucket` — but providers +/// such as Hugging Face fold `repo_id` in instead) plus the relevant +/// fields of `ObjectStoreParams` (block size, dynamic +/// `storage_options_accessor`'s `provider_id()`, etc.). It does **not** +/// incorporate auth headers, STS tokens, namespace identity, or any +/// bearer credentials. Tenant isolation under sharing therefore relies +/// entirely on callers providing a key-distinguishing input — typically +/// non-empty `storage_options`, a `storage_options_provider` whose +/// `provider_id()` carries tenant identity, or an explicit `session`. +/// +/// Bare-URI opens (empty `storage_options`, no provider, no namespace +/// commit-handler) collapse onto a single cache entry per URI: the first +/// caller's resolved default-credential chain becomes the credentials +/// used by every subsequent caller for the lifetime of that +/// `Arc`. Callers who need cross-tenant isolation under +/// bare URIs MUST opt out via +/// `LANCE_JNI_DISABLE_DEFAULT_REGISTRY_SHARING=1`; the resolved bool is +/// consulted on every default-open path. +/// +/// - **Session → metadata/index caches**: query-shaped, sized by +/// `index_cache_size_bytes` and `metadata_cache_size_bytes` from each +/// open's `ReadParams`. Sharing a Session across opens would force every +/// caller to pick the same cache size, would make eviction policy a +/// cross-tenant policy decision, and would let one tenant's hot dataset +/// evict another's. None of those are problems we want to take on inside +/// the JNI bridge — Java callers that want metadata-cache reuse can build +/// their own [`lance::session::Session`] and pass it in explicitly via +/// `BlockingDataset::open` with `session: Some(...)`. +/// +/// # Lifetime +/// +/// This static lives for the lifetime of the process. JVM unload (e.g. via +/// `System.exit`) on most platforms exits the host process, so the +/// registry is dropped along with it; the JNI library is not designed to +/// be unloaded and re-loaded within a single process. Embedders that +/// genuinely need per-JVM isolation — multiple JVMs in one address space +/// or hot-reload of the Lance native library — should construct their own +/// `Session` per JVM and pass it explicitly via +/// `BlockingDataset::open(..., session: Some(...))`, bypassing this +/// static entirely. +pub(crate) static GLOBAL_OBJECT_STORE_REGISTRY: LazyLock< + Arc, +> = LazyLock::new(|| Arc::new(lance_io::object_store::ObjectStoreRegistry::default())); + fn set_timestamp_precision(builder: &mut env_logger::Builder) { if let Ok(timestamp_precision) = env::var("LANCE_LOG_TS_PRECISION") { match timestamp_precision.as_str() { diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index 5e7963e5e97..dbb2b8c57bd 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -10,13 +10,14 @@ use std::{ }; use object_store::path::Path; -use url::Url; +use tokio::sync::Mutex as AsyncMutex; +use url::{Host, Url}; use crate::object_store::WrappingObjectStore; use crate::object_store::uri_to_url; use super::{ObjectStore, ObjectStoreParams, tracing::ObjectStoreTracingExt}; -use lance_core::error::{Error, LanceOptionExt, Result}; +use lance_core::error::{Error, Result}; #[cfg(feature = "aws")] pub mod aws; @@ -40,6 +41,15 @@ pub mod tos; #[async_trait::async_trait] pub trait ObjectStoreProvider: std::fmt::Debug + Sync + Send { + /// Construct a new object store for the given base path and params. + /// + /// **Reentry warning**: implementations MUST NOT recursively call + /// [`ObjectStoreRegistry::get_store`] for the same `(base_path, params)` + /// key. The registry coalesces concurrent cold builds via a per-key + /// async lock held across the call to `new_store`; a re-entrant call + /// would deadlock waiting on the lock the current task already holds. + /// Calling `get_store` for a *different* key (e.g. an underlying delegate + /// store) is safe. async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result; /// Extract the path relative to the base of the store. @@ -71,16 +81,60 @@ pub trait ObjectStoreProvider: std::fmt::Debug + Sync + Send { /// this will be something like 'az$account_name@container' /// /// Providers should override this if they have special requirements like Azure's. + /// + /// # Userinfo / IPv6 handling + /// + /// Overrides that build the prefix from URL authority MUST use + /// [`sanitized_authority`] (or equivalent) instead of `Url::authority()` — + /// the latter includes any embedded `userinfo@`, which would leak + /// credentials into the cache key. The default impl below already does + /// this. Azure is the only intentional exception: it parses `userinfo` + /// position as the container name, not credentials. fn calculate_object_store_prefix( &self, url: &Url, _storage_options: Option<&HashMap>, ) -> Result { - Ok(format!("{}${}", url.scheme(), url.authority())) + Ok(format!("{}${}", url.scheme(), sanitized_authority(url))) + } +} + +/// Build a cache-key authority component from a URL with two safety properties: +/// +/// 1. Strips any `userinfo@` segment so URL-embedded credentials never land in +/// the cache key (which is logged via `Debug` and observed by anything that +/// calls `calculate_object_store_prefix`). +/// 2. Wraps IPv6 hosts in brackets so `[::1]:9000` is unambiguous — +/// `Url::host_str` returns `"::1"` without brackets, which combined with a +/// port would produce the ambiguous string `"::1:9000"`. +/// +/// Returns just `host[:port]`. Use this in any +/// [`ObjectStoreProvider::calculate_object_store_prefix`] override that would +/// otherwise reach for `Url::authority()`. +pub fn sanitized_authority(url: &Url) -> String { + let host = match url.host() { + Some(Host::Ipv6(addr)) => format!("[{}]", addr), + Some(Host::Ipv4(addr)) => addr.to_string(), + Some(Host::Domain(d)) => d.to_string(), + None => String::new(), + }; + match url.port() { + Some(p) => format!("{}:{}", host, p), + None => host, } } +type CacheKey = (String, ObjectStoreParams); +type BuildLockMap = HashMap>>; + /// Statistics for the object store registry cache. +/// +/// Exposed for in-tree diagnostic consumers (e.g. the +/// `concurrent_open_bench` example, which prints hit/miss counts to prove +/// registry sharing across opens). `#[doc(hidden)]` because this is an +/// observability surface, not a stable API — field names and types may +/// change without a semver bump. +#[doc(hidden)] #[derive(Debug, Clone, Default)] pub struct ObjectStoreRegistryStats { /// Number of cache hits (store was already cached and reused). @@ -89,6 +143,20 @@ pub struct ObjectStoreRegistryStats { pub misses: u64, /// Number of currently active object stores in the cache. pub active_stores: usize, + /// Number of cache keys with an in-flight or recently-finished cold + /// build whose RAII cleanup has not yet acquired the build-locks mutex. + /// + /// Steady-state should be 0 — non-zero indicates either an in-progress + /// thundering herd (expected during cold start) or a stuck builder. + pub build_locks: usize, + /// Number of cold builds that returned an error from the underlying + /// provider (`ObjectStoreProvider::new_store`). + /// + /// Per-call, not per-key: single-flight serializes failures via the + /// per-key lock but does not cache the `Err` result. N waiters racing + /// a failing upstream each retry in turn and bump this counter, so a + /// stuck failure grows roughly linearly with attempts. + pub build_failures: u64, } /// A registry of object store providers. @@ -109,19 +177,63 @@ pub struct ObjectStoreRegistryStats { /// Use [`Self::empty()`] to create an empty registry, with no providers registered. /// /// The registry also caches object stores that are currently in use. It holds -/// weak references to the object stores, so they are not held onto. If an object -/// store is no longer in use, it will be removed from the cache on the next -/// call to either [`Self::active_stores()`] or [`Self::get_store()`]. +/// weak references to the object stores, so they are not held onto. Stale +/// entries are reclaimed lazily: +/// - [`Self::active_stores()`] sweeps the whole map. +/// - [`Self::get_store()`] evicts the queried key on miss, and opportunistically +/// prunes other dead entries on every cold-build insert (bounded work, keyed +/// off in-use opens). #[derive(Debug)] pub struct ObjectStoreRegistry { providers: RwLock>>, // Cache of object stores currently in use. We use a weak reference so the // cache itself doesn't keep them alive if no object store is actually using // it. - active_stores: RwLock>>, + active_stores: RwLock>>, + // Per-key async build locks coalesce concurrent cold builds for the same + // key into one `provider.new_store` call (avoids the thundering herd on + // S3-style providers whose client construction is expensive). + build_locks: std::sync::Mutex, // Cache statistics hits: AtomicU64, misses: AtomicU64, + build_failures: AtomicU64, +} + +/// RAII session for the cold-build path: holds the per-key async lock guard +/// and reclaims the `build_locks` HashMap entry on drop, in that order. +/// +/// `Drop::drop` explicitly releases `guard` *before* running +/// `cleanup_build_lock_if_idle`, so cleanup observes `strong_count == 1` +/// (only the HashMap entry) and can remove the lock entry safely. Encoding +/// the ordering in this Drop body — rather than relying on the declaration +/// order of two separate `let` bindings at the call site — keeps the +/// invariant locally checkable: a future refactor of the cold path cannot +/// accidentally invert it. +/// +/// Drop runs on every exit (success, error, panic, cancellation), so the +/// `build_locks` entry can never leak if the cold-build path is interrupted +/// between acquiring the lock and finishing the build. +/// +/// Owns the `CacheKey` (rather than borrowing) so the session's drop is not +/// tied to the lifetime of a local variable. +struct BuildLockSession<'a> { + registry: &'a ObjectStoreRegistry, + cache_key: CacheKey, + guard: Option>, +} + +impl Drop for BuildLockSession<'_> { + fn drop(&mut self) { + // 1. Release the per-key async lock first by dropping the + // `OwnedMutexGuard`. This decrements the per-key + // `Arc>` strong count. + self.guard.take(); + // 2. Reclaim the `build_locks` entry. With the guard gone, only the + // HashMap entry holds an Arc, so cleanup observes count == 1 and + // safely removes it. + self.registry.cleanup_build_lock_if_idle(&self.cache_key); + } } impl ObjectStoreRegistry { @@ -133,8 +245,10 @@ impl ObjectStoreRegistry { Self { providers: RwLock::new(HashMap::new()), active_stores: RwLock::new(HashMap::new()), + build_locks: std::sync::Mutex::new(HashMap::new()), hits: AtomicU64::new(0), misses: AtomicU64::new(0), + build_failures: AtomicU64::new(0), } } @@ -142,7 +256,7 @@ impl ObjectStoreRegistry { pub fn get_provider(&self, scheme: &str) -> Option> { self.providers .read() - .expect("ObjectStoreRegistry lock poisoned") + .unwrap_or_else(|p| p.into_inner()) .get(scheme) .cloned() } @@ -151,12 +265,18 @@ impl ObjectStoreRegistry { /// /// Calling this will also clean up any weak references to object stores that /// are no longer valid. + /// + /// **Caution**: the returned `Arc` instances may carry + /// per-tenant credential state (S3 access keys, Azure SAS tokens, etc.). + /// Callers that hand these stores out across trust boundaries are + /// responsible for filtering — the registry itself does not know which + /// store belongs to which tenant. pub fn active_stores(&self) -> Vec> { let mut found_inactive = false; let output = self .active_stores .read() - .expect("ObjectStoreRegistry lock poisoned") + .unwrap_or_else(|p| p.into_inner()) .values() .filter_map(|weak| match weak.upgrade() { Some(store) => Some(store), @@ -172,7 +292,7 @@ impl ObjectStoreRegistry { let mut cache_lock = self .active_stores .write() - .expect("ObjectStoreRegistry lock poisoned"); + .unwrap_or_else(|p| p.into_inner()); cache_lock.retain(|_, weak| weak.upgrade().is_some()); } output @@ -180,19 +300,35 @@ impl ObjectStoreRegistry { /// Get cache statistics for monitoring and debugging. /// - /// Returns the number of cache hits, misses, and currently active stores. - /// This is useful for detecting configuration issues that cause excessive - /// cache misses (e.g., storage options that vary per-request). + /// Returns counters for hits, misses, currently active stores, + /// in-flight build locks, and accumulated build failures — see + /// [`ObjectStoreRegistryStats`]. Useful for detecting configuration + /// issues that cause excessive cache misses (e.g., storage options + /// that vary per-request). + /// + /// `#[doc(hidden)]`: paired with [`ObjectStoreRegistryStats`] — this is + /// an unstable observability surface kept `pub` only to support the + /// in-tree `concurrent_open_bench` example. + #[doc(hidden)] pub fn stats(&self) -> ObjectStoreRegistryStats { let active_stores = self .active_stores .read() - .map(|s| s.values().filter(|w| w.strong_count() > 0).count()) - .unwrap_or(0); + .unwrap_or_else(|p| p.into_inner()) + .values() + .filter(|w| w.strong_count() > 0) + .count(); + let build_locks = self + .build_locks + .lock() + .unwrap_or_else(|p| p.into_inner()) + .len(); ObjectStoreRegistryStats { hits: self.hits.load(Ordering::Relaxed), misses: self.misses.load(Ordering::Relaxed), active_stores, + build_locks, + build_failures: self.build_failures.load(Ordering::Relaxed), } } @@ -210,6 +346,15 @@ impl ObjectStoreRegistry { /// If the object store is already in use, it will return a strong reference /// to the object store. If the object store is not in use, it will create a /// new object store and return a strong reference to it. + /// + /// Concurrent cold builds for the same key are coalesced via a per-key + /// async lock: the first task builds the store, all others wait then + /// re-check the cache and observe the freshly populated entry. + /// + /// On build *failure*, the cache is not populated, so each waiter retries + /// in turn (serialized — not parallel-amplified, but not deduplicated + /// either). Transient errors thus surface to operators rather than being + /// masked by stale-error reuse. pub async fn get_store( &self, base_path: Url, @@ -230,39 +375,65 @@ impl ObjectStoreRegistry { provider.calculate_object_store_prefix(&base_path, params.storage_options())?; let cache_key = (cache_path.clone(), params.clone()); - // Check if we have a cached store for this base path and params - { - let maybe_store = self - .active_stores - .read() - .ok() - .expect_ok()? - .get(&cache_key) - .cloned(); - if let Some(store) = maybe_store { - if let Some(store) = store.upgrade() { - self.hits.fetch_add(1, Ordering::Relaxed); - return Ok(store); - } else { - // Remove the weak reference if it is no longer valid - let mut cache_lock = self - .active_stores - .write() - .expect("ObjectStoreRegistry lock poisoned"); - if let Some(store) = cache_lock.get(&cache_key) - && store.upgrade().is_none() - { - // Remove the weak reference if it is no longer valid - cache_lock.remove(&cache_key); - } - } - } + // Fast path: cache hit avoids both the std mutex and the async lock. + if let Some(store) = self.lookup_cached(&cache_key) { + self.hits.fetch_add(1, Ordering::Relaxed); + return Ok(store); } - self.misses.fetch_add(1, Ordering::Relaxed); + // Cold path: per-key single-flight. The RAII session guarantees the + // build_locks entry is GC'd on every exit (success, error, panic, + // cancellation), preventing unbounded HashMap growth. Drop ordering + // (guard released *before* cleanup) is encoded in + // `BuildLockSession::drop`, so the cold path can't accidentally + // invert it via `let` reordering. + let lock = self.acquire_build_lock(&cache_key); + let guard = lock.lock_owned().await; + let _session = BuildLockSession { + registry: self, + cache_key: cache_key.clone(), + guard: Some(guard), + }; - let mut store = provider.new_store(base_path, params).await?; + // Re-check after acquiring the lock — coalesced waiters become hits. + if let Some(store) = self.lookup_cached(&cache_key) { + self.hits.fetch_add(1, Ordering::Relaxed); + log::debug!( + "ObjectStoreRegistry: coalesced wait hit for cache_key path={:?}", + cache_key.0, + ); + return Ok(store); + } + + // `fetch_add` returns the prior value; `+ 1` is this attempt's + // 1-indexed sequence number. Use it (not a separate counter) to + // gate the opportunistic sweep below: cadence is approximate + // because cold attempts that fail still bump `misses` but skip + // the sweep block, so a failure at a multiple-of-N sequence + // skips that sweep cycle entirely; the next candidate is the + // next multiple-of-N cold attempt that succeeds. That's fine — + // the sweep is a footprint trim, not a correctness invariant. + // `wrapping_add` only silences debug-build overflow checks; at + // 10M cold opens/sec the wrap is >58 millennia away. + let cold_attempt_nth = self.misses.fetch_add(1, Ordering::Relaxed).wrapping_add(1); + log::debug!( + "ObjectStoreRegistry: cold build starting for cache_key path={:?}", + cache_key.0, + ); + let mut store = match provider.new_store(base_path, params).await { + Ok(s) => s, + Err(e) => { + self.build_failures.fetch_add(1, Ordering::Relaxed); + // Intentionally no log here: provider error Display impls can + // surface raw URLs whose username or query-string component + // may carry credentials (SAS tokens, access keys), and `Url`'s + // own `Display` only masks the password. The full error goes + // to the caller via `Err(e)`; operators wanting an aggregate + // signal can scrape `stats().build_failures`. + return Err(e); + } + }; store.inner = store.inner.traced(); #[cfg(feature = "metrics")] @@ -276,19 +447,121 @@ impl ObjectStoreRegistry { if let Some(wrapper) = ¶ms.object_store_wrapper { store.inner = wrapper.wrap(&cache_path, store.inner); } - - // Always wrap with IO tracking store.inner = store.io_tracker.wrap("", store.inner); + let cached = Arc::new(store); + // Amortized opportunistic sweep: gate the O(n) `retain` behind a + // mod-N counter so a burst of cold builds with distinct URIs costs + // O(n) work in aggregate, not O(n²). N=64 is small enough that the + // map's footprint stays roughly proportional to live opens, large + // enough that bursty cold-open paths don't pay sweep cost on every + // insert. Per-key `lookup_cached` and full sweep via `active_stores()` + // remain unchanged; this is purely a per-insert cost lever. + // + // Cadence is "every 64th cold attempt" (1-indexed via the pre-bump + // above), so the first sweep fires at the 64th, not the 0th, and + // a near-empty map never pays the retain cost on its very first + // cold build. + const SWEEP_INTERVAL: u64 = 64; + let should_sweep = cold_attempt_nth.is_multiple_of(SWEEP_INTERVAL); + { + let mut cache_lock = self + .active_stores + .write() + .unwrap_or_else(|p| p.into_inner()); + // Safe under the write lock: no other writer can resurrect a weak + // ref between `retain` and `insert`, and `insert` for `cache_key` + // overwrites whatever `retain` left (live or dead) for that key. + if should_sweep { + cache_lock.retain(|_, weak| weak.upgrade().is_some()); + } + cache_lock.insert(cache_key, Arc::downgrade(&cached)); + } + Ok(cached) + } - let store = Arc::new(store); + /// Acquire (or create) the per-key async build lock for `cache_key`. + /// + /// On poison: recovers via `into_inner()` rather than bricking the + /// registry. The protected data is just a HashMap of build-lock Arcs; + /// any stale entry left by a panic is reclaimed by `BuildLockSession`. + fn acquire_build_lock(&self, cache_key: &CacheKey) -> Arc> { + let mut locks = self + .build_locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + locks + .entry(cache_key.clone()) + .or_insert_with(|| Arc::new(AsyncMutex::new(()))) + .clone() + } + /// Drop the build-lock entry for `cache_key` if no waiters remain. + /// + /// Strong-count == 1 means only the HashMap entry references the Arc, so + /// no concurrent task is using or queued on the lock. The HashMap mutex + /// makes this atomic: a new caller can't clone the Arc between the count + /// read and the removal because `acquire_build_lock` takes the same mutex. + /// A waiter parked inside `lock_owned()` holds its own Arc clone in the + /// future, so it counts toward `strong_count` and keeps the entry alive + /// until it wakes. + /// + /// Best-effort: if a task drops without calling this (e.g. between + /// dropping the build lock and entering this function), the entry is + /// reclaimed by the next caller for the same key, who finds count == 1 + /// and removes it. + /// + /// Bound: per-key cleanup is sufficient — no periodic full-map sweep is + /// needed. The map's worst-case size is the number of distinct cache keys + /// with an in-flight builder *plus* keys whose builder has finished but + /// whose `BuildLockSession::drop` has not yet acquired the std mutex. + /// Both terms are bounded by concurrently outstanding cold opens, which + /// the underlying I/O concurrency limits already cap. RAII guarantees + /// every successful `acquire_build_lock` is paired with exactly one + /// `cleanup_build_lock_if_idle`, on every exit path. + fn cleanup_build_lock_if_idle(&self, cache_key: &CacheKey) { + let mut locks = self + .build_locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(entry) = locks.get(cache_key) + && Arc::strong_count(entry) == 1 { - // Insert the store into the cache - let mut cache_lock = self.active_stores.write().ok().expect_ok()?; - cache_lock.insert(cache_key, Arc::downgrade(&store)); + locks.remove(cache_key); } + } + + #[cfg(test)] + fn build_locks_len(&self) -> usize { + self.build_locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .len() + } - Ok(store) + /// Look up a cached store by key, evicting the entry if its weak ref is + /// no longer upgradable. + fn lookup_cached(&self, cache_key: &CacheKey) -> Option> { + let maybe_weak = self + .active_stores + .read() + .unwrap_or_else(|p| p.into_inner()) + .get(cache_key) + .cloned(); + let weak = maybe_weak?; + if let Some(store) = weak.upgrade() { + return Some(store); + } + // Stale weak; evict under the write lock. + let mut cache_lock = self + .active_stores + .write() + .unwrap_or_else(|p| p.into_inner()); + if let Some(weak) = cache_lock.get(cache_key) + && weak.upgrade().is_none() + { + cache_lock.remove(cache_key); + } + None } /// Calculate the datastore prefix based on the URI and the storage options. @@ -361,8 +634,10 @@ impl Default for ObjectStoreRegistry { Self { providers: RwLock::new(providers), active_stores: RwLock::new(HashMap::new()), + build_locks: std::sync::Mutex::new(HashMap::new()), hits: AtomicU64::new(0), misses: AtomicU64::new(0), + build_failures: AtomicU64::new(0), } } } @@ -373,7 +648,7 @@ impl ObjectStoreRegistry { pub fn insert(&self, scheme: &str, provider: Arc) { self.providers .write() - .expect("ObjectStoreRegistry lock poisoned") + .unwrap_or_else(|p| p.into_inner()) .insert(scheme.into(), provider); } } @@ -438,6 +713,54 @@ mod tests { assert!(Arc::ptr_eq(&store_scoped, &store_plain)); } + /// `userinfo@host` (URL-embedded credentials) MUST be stripped from the + /// cache-key prefix so the secret never lands in HashMap Debug logs or + /// downstream state. Two URLs that differ only in `userinfo` collapse to + /// the same prefix; multi-tenant isolation has to come from + /// `ObjectStoreParams` (storage_options / accessor), not URL embedding. + #[test] + fn test_calculate_object_store_prefix_strips_userinfo() { + let provider = DummyProvider; + let with_creds = Url::parse("dummy://user:s3cret@blah/path").unwrap(); + let without = Url::parse("dummy://blah/path").unwrap(); + let with_prefix = provider + .calculate_object_store_prefix(&with_creds, None) + .unwrap(); + assert_eq!("dummy$blah", with_prefix); + assert_eq!( + with_prefix, + provider + .calculate_object_store_prefix(&without, None) + .unwrap(), + ); + // Defense in depth: assert the secret literally cannot appear in the + // key, regardless of formatting changes. + assert!(!with_prefix.contains('@')); + assert!(!with_prefix.contains("s3cret")); + assert!(!with_prefix.contains("user")); + } + + #[test] + fn test_calculate_object_store_prefix_keeps_port() { + let provider = DummyProvider; + let url = Url::parse("dummy://host:9000/path").unwrap(); + assert_eq!( + "dummy$host:9000", + provider.calculate_object_store_prefix(&url, None).unwrap() + ); + } + + #[test] + fn test_sanitized_authority_brackets_ipv6() { + // Url::host_str returns "::1" without brackets; combined with a port, + // the naive `format!("{}:{}", host, port)` would produce the + // ambiguous string "::1:9000". sanitized_authority restores brackets. + let url = Url::parse("dummy://[::1]:9000/path").unwrap(); + assert_eq!("[::1]:9000", sanitized_authority(&url)); + let url_no_port = Url::parse("dummy://[2001:db8::1]/path").unwrap(); + assert_eq!("[2001:db8::1]", sanitized_authority(&url_no_port)); + } + #[test] fn test_calculate_object_store_scheme_not_found() { let registry = ObjectStoreRegistry::empty(); @@ -518,7 +841,253 @@ mod tests { ); } - // Same params returns same instance + // The cache reused the entry for the second call (hits == 1 above); + // when the same params hit the same cache slot, the returned Arcs + // point to the same underlying ObjectStore. assert!(Arc::ptr_eq(&stores[0], &stores[1])); } + + /// A provider that counts `new_store` invocations and yields explicitly, + /// so concurrent callers reliably contend on the registry's build lock. + #[derive(Debug, Default)] + struct CountingProvider { + builds: AtomicU64, + } + + #[async_trait::async_trait] + impl ObjectStoreProvider for CountingProvider { + async fn new_store( + &self, + base_path: Url, + params: &ObjectStoreParams, + ) -> Result { + self.builds.fetch_add(1, Ordering::Relaxed); + // Force a yield + delay so other tasks reach `get_store` and queue + // on the build lock before this build completes. + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + memory::MemoryStoreProvider + .new_store(base_path, params) + .await + } + } + + /// Concurrent cold-misses for the same key must coalesce into a single + /// `provider.new_store` invocation, with all callers receiving the same Arc. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_get_store_coalesces_concurrent_misses() { + let provider = Arc::new(CountingProvider::default()); + let registry = Arc::new(ObjectStoreRegistry::empty()); + registry.insert("counting", provider.clone()); + + let url = Url::parse("counting://test").unwrap(); + let params = ObjectStoreParams::default(); + + let n = 16; + let mut handles = Vec::with_capacity(n); + for _ in 0..n { + let registry = registry.clone(); + let url = url.clone(); + let params = params.clone(); + handles.push(tokio::spawn(async move { + registry.get_store(url, ¶ms).await.unwrap() + })); + } + let mut stores = Vec::with_capacity(n); + for h in handles { + stores.push(h.await.unwrap()); + } + + assert_eq!( + provider.builds.load(Ordering::Relaxed), + 1, + "single-flight must collapse N concurrent misses into 1 build" + ); + // All concurrent callers share the same cached Arc — single-flight + // produced exactly one ObjectStore and the registry handed it out + // to every waiter. + for s in &stores[1..] { + assert!(Arc::ptr_eq(&stores[0], s)); + } + let s = registry.stats(); + assert_eq!((s.misses, s.hits as usize, s.active_stores), (1, n - 1, 1)); + } + + /// After a successful build, the per-key build lock must be removed so + /// `build_locks` does not grow unbounded across many distinct keys. + #[tokio::test] + async fn test_get_store_cleans_build_locks_after_use() { + let registry = ObjectStoreRegistry::default(); + let params = ObjectStoreParams::default(); + + for i in 0..5 { + let url = Url::parse(&format!("memory://cleanup-{}", i)).unwrap(); + let _store = registry.get_store(url, ¶ms).await.unwrap(); + } + assert_eq!( + registry.build_locks_len(), + 0, + "build_locks must be empty once all builds settle" + ); + } + + /// A provider whose `new_store` panics, used to verify the RAII cleanup + /// guard reclaims build_locks entries on panic. + #[derive(Debug, Default)] + struct PanickingProvider; + + #[async_trait::async_trait] + impl ObjectStoreProvider for PanickingProvider { + async fn new_store( + &self, + _base_path: Url, + _params: &ObjectStoreParams, + ) -> Result { + panic!("boom from provider") + } + } + + /// A provider that returns an `Err` after a yield, simulating a normal + /// (non-panic) build failure such as bad credentials or a network error. + /// Counts invocations so the test can assert the single-flight collapsed + /// N concurrent failures into 1 underlying call. + #[derive(Debug, Default)] + struct FailingProvider { + builds: AtomicU64, + } + + #[async_trait::async_trait] + impl ObjectStoreProvider for FailingProvider { + async fn new_store( + &self, + _base_path: Url, + _params: &ObjectStoreParams, + ) -> Result { + self.builds.fetch_add(1, Ordering::Relaxed); + // Yield + delay so concurrent callers reach the build lock and + // queue behind us — this is the configuration that exercises + // the coalesced-failure path in `get_store`. + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + Err(Error::invalid_input("simulated cold-build failure")) + } + } + + /// Build-Err (non-panic) path under coalesced contention asserts the + /// shipped behavior of PR#1's single-flight: failures **serialize via + /// the per-key lock but are not cached**. Each of N waiters acquires + /// the lock in turn, re-checks the (still-empty) cache, and retries + /// the build. The lock keeps failures from *fanning out in parallel* + /// (so an upstream that costs T ms sees ~N·T total, not T concurrent), + /// but it does not collapse them into one call. + /// + /// Concretely this test pins: + /// 1. `provider.new_store` fires once per caller (no Err coalescing), + /// 2. every concurrent caller surfaces an `Err`, + /// 3. `build_failures` increments per failed call (not per key), + /// 4. `active_stores` stays 0 — failed builds must never be cached, + /// 5. `build_locks` drains back to 0 via the RAII cleanup, even on + /// the Err return path. + /// + /// Caching errors (with TTL or until-eviction) is a separate design + /// point. This test exists to lock in the current contract so a future + /// "let's cache Err for 100ms" change is forced to update this guard + /// deliberately rather than land silently. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_get_store_serializes_build_failures_without_caching() { + let provider = Arc::new(FailingProvider::default()); + let registry = Arc::new(ObjectStoreRegistry::empty()); + registry.insert("failing", provider.clone()); + + let url = Url::parse("failing://x").unwrap(); + let params = ObjectStoreParams::default(); + + let n = 8; + let mut handles = Vec::with_capacity(n); + for _ in 0..n { + let registry = registry.clone(); + let url = url.clone(); + let params = params.clone(); + handles.push(tokio::spawn(async move { + registry.get_store(url, ¶ms).await + })); + } + + let mut err_count = 0; + for h in handles { + let result = h.await.expect("task must not panic"); + assert!(result.is_err(), "every waiter must surface Err"); + err_count += 1; + } + assert_eq!(err_count, n, "all N callers must observe an error"); + + assert_eq!( + provider.builds.load(Ordering::Relaxed) as usize, + n, + "shipped behavior: failures are not cached, so each waiter retries the build", + ); + + let s = registry.stats(); + assert_eq!( + s.build_failures as usize, n, + "build_failures counts every failed call (per-call, not per-key)", + ); + assert_eq!( + s.active_stores, 0, + "failed builds must never populate active_stores", + ); + assert_eq!( + s.build_locks, 0, + "RAII cleanup must drain build_locks even on Err path", + ); + } + + /// If `new_store` panics, the RAII guard must still remove the build_locks + /// entry — otherwise a single bad URI would leak a lock per failed open. + /// + /// Relies on the default `panic = "unwind"` profile: the panic unwinds + /// the spawned task and `Drop` runs on the way out, which is what triggers + /// `BuildLockSession`. With `panic = "abort"`, the process dies before + /// cleanup can run — that's a stricter regime where this test (and the + /// leak it guards against) becomes moot. + #[tokio::test] + async fn test_get_store_cleans_build_locks_after_panic() { + let registry = Arc::new(ObjectStoreRegistry::empty()); + registry.insert("boom", Arc::new(PanickingProvider)); + let url = Url::parse("boom://x").unwrap(); + let params = ObjectStoreParams::default(); + + let registry_for_task = registry.clone(); + let result = + tokio::task::spawn(async move { registry_for_task.get_store(url, ¶ms).await }) + .await; + assert!( + result.as_ref().err().is_some_and(|e| e.is_panic()), + "task must propagate the panic" + ); + assert_eq!( + registry.build_locks_len(), + 0, + "RAII guard must clean up the build_locks entry on panic" + ); + } + + /// Regression: two callers resolving the same cache key must receive + /// the same cached `Arc`. Single-flight correctness depends + /// on this — see `test_get_store_coalesces_concurrent_misses` for the + /// concurrent variant. + #[tokio::test] + async fn test_get_store_returns_shared_arc_on_hit() { + let registry = ObjectStoreRegistry::default(); + let params = ObjectStoreParams::default(); + let url = Url::parse("memory://shared-arc").unwrap(); + + let a = registry.get_store(url.clone(), ¶ms).await.unwrap(); + let b = registry.get_store(url, ¶ms).await.unwrap(); + + let s = registry.stats(); + assert_eq!((s.misses, s.hits, s.active_stores), (1, 1, 1)); + assert!( + Arc::ptr_eq(&a, &b), + "cache hit must return the same Arc as the cold-build caller" + ); + } } From dc4a1103c77d5890bf5fa088aa2870ddbd467b82 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 19 May 2026 21:45:51 +0800 Subject: [PATCH 2/6] chore: retrigger CI for flaky test_create_index_nulls::case_4_ivf_hnsw_pq From 40cea2ef72a9cdee0853783e36a517cdee94d7e4 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 22 Jun 2026 13:05:44 +0800 Subject: [PATCH 3/6] fix(io): address review comments on default-open Session reuse - ObjectStoreRegistryStats: add #[non_exhaustive] so future field additions stay source-compatible (review comment 1). - Sanitize URL authority into the cache-key prefix across every registered authority-based provider (huggingface, shared-memory prefix + backend key, and the unknown-scheme fallback), so embedded userinfo can no longer leak into cache keys or the cache-key debug logs (review comment 2). - Make per-key build-lock cleanup cancellation-safe via an explicit waiter count on BuildLockEntry: acquire increments, the RAII BuildLockSession (now constructed before the lock_owned() await) decrements on every unwinding exit. A waiter cancelled while queued no longer orphans its build_locks entry; cleanup no longer depends on Arc::strong_count or async drop order (review comment 3). Adds tests for each: userinfo stripping (huggingface + shared-memory), backend-routing consistency, and the cancelled-while-queued waiter (white-box guard-less drop + end-to-end waiter-count observation). --- rust/lance-io/src/object_store.rs | 9 +- rust/lance-io/src/object_store/providers.rs | 311 ++++++++++++++---- .../src/object_store/providers/huggingface.rs | 31 +- .../object_store/providers/shared_memory.rs | 53 ++- 4 files changed, 334 insertions(+), 70 deletions(-) diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 15905ff4e69..abb388d2cae 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -1126,7 +1126,14 @@ impl ObjectStore { .calculate_object_store_prefix(&location, storage_options) .unwrap(), None => { - let store_prefix = format!("{}${}", location.scheme(), location.authority()); + // `sanitized_authority`, not `Url::authority()`: strip any + // `userinfo@` so credentials don't leak into the prefix (which + // is warn-logged below) for unknown-scheme URLs. + let store_prefix = format!( + "{}${}", + location.scheme(), + providers::sanitized_authority(&location) + ); log::warn!( "Guessing that object store prefix is {}, since object store scheme is not found in registry.", store_prefix diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index dbb2b8c57bd..0ea0829a275 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -125,7 +125,25 @@ pub fn sanitized_authority(url: &Url) -> String { } type CacheKey = (String, ObjectStoreParams); -type BuildLockMap = HashMap>>; +type BuildLockMap = HashMap; + +/// A per-key cold-build lock plus an explicit count of the tasks currently +/// holding or queued on it. +/// +/// The count is the cleanup signal — not `Arc::strong_count` of the mutex. +/// `acquire_build_lock` increments it (under the `build_locks` std mutex) the +/// instant a task claims the entry; [`BuildLockSession::drop`] decrements it on +/// every exit path. Both happen with no intervening `.await`, so the +/// increment/decrement is perfectly paired regardless of when the pending +/// `lock_owned()` future releases its own Arc — which is what makes +/// cancellation-while-queued safe without relying on async drop order. +#[derive(Debug)] +struct BuildLockEntry { + mutex: Arc>, + /// Number of tasks holding or queued on `mutex`. The entry is removed when + /// this reaches zero. + waiters: usize, +} /// Statistics for the object store registry cache. /// @@ -136,6 +154,11 @@ type BuildLockMap = HashMap>>; /// change without a semver bump. #[doc(hidden)] #[derive(Debug, Clone, Default)] +// Adding fields to this struct must stay source-compatible: external code may +// read it via `stats()` and `..Default::default()`-construct it. `#[non_exhaustive]` +// forbids struct-literal construction outside this crate, so future field +// additions don't break downstream builds. +#[non_exhaustive] pub struct ObjectStoreRegistryStats { /// Number of cache hits (store was already cached and reused). pub hits: u64, @@ -143,8 +166,8 @@ pub struct ObjectStoreRegistryStats { pub misses: u64, /// Number of currently active object stores in the cache. pub active_stores: usize, - /// Number of cache keys with an in-flight or recently-finished cold - /// build whose RAII cleanup has not yet acquired the build-locks mutex. + /// Number of cache keys with at least one task currently holding or queued + /// on the per-key cold-build lock. /// /// Steady-state should be 0 — non-zero indicates either an in-progress /// thundering herd (expected during cold start) or a stuck builder. @@ -200,23 +223,27 @@ pub struct ObjectStoreRegistry { build_failures: AtomicU64, } -/// RAII session for the cold-build path: holds the per-key async lock guard -/// and reclaims the `build_locks` HashMap entry on drop, in that order. +/// RAII session for the cold-build path. On drop it releases the per-key async +/// lock guard, then calls `release_build_lock` to drop this task's +/// [`BuildLockEntry::waiters`] hold (which removes the entry at zero). Releasing +/// the guard first lets a queued waiter wake before the build-locks map is +/// touched. /// -/// `Drop::drop` explicitly releases `guard` *before* running -/// `cleanup_build_lock_if_idle`, so cleanup observes `strong_count == 1` -/// (only the HashMap entry) and can remove the lock entry safely. Encoding -/// the ordering in this Drop body — rather than relying on the declaration -/// order of two separate `let` bindings at the call site — keeps the -/// invariant locally checkable: a future refactor of the cold path cannot -/// accidentally invert it. +/// It is constructed *before* awaiting `lock_owned()`, with no `.await` between +/// `acquire_build_lock`'s increment and the construction. So a waiter cancelled +/// while still *queued* on `lock_owned()` — guard never assigned — still drops a +/// live session and decrements; without that ordering the entry would be +/// orphaned until the next open for the same key. (Cleanup tracks the explicit +/// count, not `Arc::strong_count`, so it is also immune to the drop order of the +/// pending `lock_owned()` future — see [`BuildLockEntry`].) /// -/// Drop runs on every exit (success, error, panic, cancellation), so the -/// `build_locks` entry can never leak if the cold-build path is interrupted -/// between acquiring the lock and finishing the build. +/// Drop runs on every unwinding exit — success, error, unwinding panic, +/// cancellation. (`panic = "abort"` skips destructors, but the process is dying, +/// so there is no surviving registry to leak into; see +/// `test_get_store_cleans_build_locks_after_panic`.) /// -/// Owns the `CacheKey` (rather than borrowing) so the session's drop is not -/// tied to the lifetime of a local variable. +/// `guard` is `None` until the lock is acquired. The owned `CacheKey` keeps the +/// drop independent of any local variable's lifetime. struct BuildLockSession<'a> { registry: &'a ObjectStoreRegistry, cache_key: CacheKey, @@ -225,14 +252,11 @@ struct BuildLockSession<'a> { impl Drop for BuildLockSession<'_> { fn drop(&mut self) { - // 1. Release the per-key async lock first by dropping the - // `OwnedMutexGuard`. This decrements the per-key - // `Arc>` strong count. + // Release the async lock first so a queued waiter can wake before we + // touch the build-locks map, then drop this task's waiter hold (which + // removes the entry once the last waiter leaves). self.guard.take(); - // 2. Reclaim the `build_locks` entry. With the guard gone, only the - // HashMap entry holds an Arc, so cleanup observes count == 1 and - // safely removes it. - self.registry.cleanup_build_lock_if_idle(&self.cache_key); + self.registry.release_build_lock(&self.cache_key); } } @@ -381,19 +405,20 @@ impl ObjectStoreRegistry { return Ok(store); } - // Cold path: per-key single-flight. The RAII session guarantees the - // build_locks entry is GC'd on every exit (success, error, panic, - // cancellation), preventing unbounded HashMap growth. Drop ordering - // (guard released *before* cleanup) is encoded in - // `BuildLockSession::drop`, so the cold path can't accidentally - // invert it via `let` reordering. + // Cold path: per-key single-flight. `acquire_build_lock` registers this + // task as a waiter; the RAII `BuildLockSession` releases that hold on + // every exit, bounding the map. The session is built before the + // `lock_owned()` await (see `BuildLockSession`) so a waiter cancelled + // while queued still cleans up. let lock = self.acquire_build_lock(&cache_key); - let guard = lock.lock_owned().await; - let _session = BuildLockSession { + let mut _session = BuildLockSession { registry: self, cache_key: cache_key.clone(), - guard: Some(guard), + guard: None, }; + // Cancellation point: `_session` already exists, so its Drop cleans up + // even if we are cancelled here before the guard is assigned. + _session.guard = Some(lock.lock_owned().await); // Re-check after acquiring the lock — coalesced waiters become hits. if let Some(store) = self.lookup_cached(&cache_key) { @@ -479,54 +504,52 @@ impl ObjectStoreRegistry { Ok(cached) } - /// Acquire (or create) the per-key async build lock for `cache_key`. + /// Acquire (or create) the per-key async build lock for `cache_key`, + /// registering this task as a waiter. + /// + /// Increments [`BuildLockEntry::waiters`] under the `build_locks` std + /// mutex; the caller MUST pair this with exactly one `release_build_lock`, + /// which the RAII [`BuildLockSession`] guarantees on every exit path. /// /// On poison: recovers via `into_inner()` rather than bricking the - /// registry. The protected data is just a HashMap of build-lock Arcs; + /// registry. The protected data is just a HashMap of build-lock entries; /// any stale entry left by a panic is reclaimed by `BuildLockSession`. fn acquire_build_lock(&self, cache_key: &CacheKey) -> Arc> { let mut locks = self .build_locks .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - locks + let entry = locks .entry(cache_key.clone()) - .or_insert_with(|| Arc::new(AsyncMutex::new(()))) - .clone() + .or_insert_with(|| BuildLockEntry { + mutex: Arc::new(AsyncMutex::new(())), + waiters: 0, + }); + entry.waiters += 1; + entry.mutex.clone() } - /// Drop the build-lock entry for `cache_key` if no waiters remain. + /// Release this task's hold on the per-key build lock for `cache_key`, + /// removing the entry once the last waiter leaves. /// - /// Strong-count == 1 means only the HashMap entry references the Arc, so - /// no concurrent task is using or queued on the lock. The HashMap mutex - /// makes this atomic: a new caller can't clone the Arc between the count - /// read and the removal because `acquire_build_lock` takes the same mutex. - /// A waiter parked inside `lock_owned()` holds its own Arc clone in the - /// future, so it counts toward `strong_count` and keeps the entry alive - /// until it wakes. + /// Decrements [`BuildLockEntry::waiters`] and removes the entry at zero. The + /// `build_locks` std mutex makes the decrement-and-maybe-remove atomic + /// against concurrent `acquire_build_lock` calls, so a new waiter can't be + /// lost between the two. /// - /// Best-effort: if a task drops without calling this (e.g. between - /// dropping the build lock and entering this function), the entry is - /// reclaimed by the next caller for the same key, who finds count == 1 - /// and removes it. - /// - /// Bound: per-key cleanup is sufficient — no periodic full-map sweep is - /// needed. The map's worst-case size is the number of distinct cache keys - /// with an in-flight builder *plus* keys whose builder has finished but - /// whose `BuildLockSession::drop` has not yet acquired the std mutex. - /// Both terms are bounded by concurrently outstanding cold opens, which - /// the underlying I/O concurrency limits already cap. RAII guarantees - /// every successful `acquire_build_lock` is paired with exactly one - /// `cleanup_build_lock_if_idle`, on every exit path. - fn cleanup_build_lock_if_idle(&self, cache_key: &CacheKey) { + /// The map's worst-case size is the number of distinct cache keys with an + /// outstanding cold open, which the I/O concurrency limits already cap — + /// RAII pairs every `acquire_build_lock` with exactly one `release_build_lock`. + fn release_build_lock(&self, cache_key: &CacheKey) { let mut locks = self .build_locks .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(entry) = locks.get(cache_key) - && Arc::strong_count(entry) == 1 - { - locks.remove(cache_key); + if let Some(entry) = locks.get_mut(cache_key) { + entry.waiters -= 1; + if entry.waiters == 0 { + locks.remove(cache_key); + } } } @@ -538,6 +561,18 @@ impl ObjectStoreRegistry { .len() } + /// Current waiter count for `cache_key`, or `None` if no entry exists. + /// Lets tests observe the per-key acquire/release count directly — the + /// signal that drives cleanup — rather than only the map's key count. + #[cfg(test)] + fn build_lock_waiters(&self, cache_key: &CacheKey) -> Option { + self.build_locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .get(cache_key) + .map(|entry| entry.waiters) + } + /// Look up a cached store by key, evicting the entry if its weak ref is /// no longer upgradable. fn lookup_cached(&self, cache_key: &CacheKey) -> Option> { @@ -1070,6 +1105,154 @@ mod tests { ); } + /// A waiter cancelled while still *queued* on the per-key build lock — + /// before its `OwnedMutexGuard` is acquired — must still reclaim the + /// `build_locks` entry on drop. This pins that state directly: the + /// [`BuildLockSession`] is built with `guard: None`, exactly as it exists + /// for a task cancelled inside `lock_owned().await`. The lingering + /// `Arc>` (`_held`) stands in for the Arc that the pending + /// `lock_owned()` future keeps alive — cleanup must not be fooled by it, + /// because it keys off the explicit waiter count, not `Arc::strong_count`. + #[tokio::test] + async fn test_build_lock_session_reclaims_entry_when_cancelled_before_guard() { + let registry = ObjectStoreRegistry::empty(); + let cache_key = ("memory$queued".to_string(), ObjectStoreParams::default()); + + // Mirror the cold path: register a waiter and take the lock Arc, but + // never acquire the guard. + let _held = registry.acquire_build_lock(&cache_key); + assert_eq!( + registry.build_locks_len(), + 1, + "acquiring the build lock must register the entry" + ); + + let session = BuildLockSession { + registry: ®istry, + cache_key, + guard: None, + }; + drop(session); + + assert_eq!( + registry.build_locks_len(), + 0, + "dropping the session for a still-queued (guard-less) waiter must \ + reclaim the entry, even while another Arc to the lock is alive" + ); + } + + /// A provider that signals when its build starts and blocks until released, + /// pinning a builder in flight so other callers queue behind it. + #[derive(Debug)] + struct BlockingProvider { + started: Arc, + release: Arc, + } + + #[async_trait::async_trait] + impl ObjectStoreProvider for BlockingProvider { + async fn new_store( + &self, + base_path: Url, + params: &ObjectStoreParams, + ) -> Result { + self.started.notify_one(); + self.release.notified().await; + memory::MemoryStoreProvider + .new_store(base_path, params) + .await + } + } + + /// End-to-end cancellation safety: a waiter cancelled while queued behind + /// an in-flight builder must release its hold on the build lock. The test + /// observes the per-key waiter count directly so the cancellation's effect + /// is visible *before* the builder finishes — otherwise the builder's own + /// cleanup would mask whether the cancelled waiter ever decremented. + /// + /// Sequence: builder A registers (count 1) and parks; waiter B queues + /// (count 2); B is cancelled mid-wait → count must drop back to 1; A + /// completes → entry fully reclaimed. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_get_store_cleans_build_locks_when_queued_waiter_cancelled() { + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let provider = Arc::new(BlockingProvider { + started: started.clone(), + release: release.clone(), + }); + let registry = Arc::new(ObjectStoreRegistry::empty()); + registry.insert("blocking", provider); + + let url = Url::parse("blocking://same").unwrap(); + let params = ObjectStoreParams::default(); + // The cache key the cold path computes for this URL: default prefix is + // `scheme$sanitized_authority`, paired with the (default) params. + let cache_key = ("blocking$same".to_string(), params.clone()); + + // Builder A acquires the lock and parks inside `new_store`. + let (ra, ua, pa) = (registry.clone(), url.clone(), params.clone()); + let builder = tokio::spawn(async move { ra.get_store(ua, &pa).await }); + started.notified().await; + assert_eq!( + registry.build_lock_waiters(&cache_key), + Some(1), + "builder A must be registered as the sole waiter" + ); + + // Waiter B targets the same key, so it queues on `lock_owned()` behind + // A (it cannot acquire the guard — A holds it). Wait deterministically + // for B to register rather than racing a fixed sleep; a regression that + // never queues fails here loudly instead of hanging. + let (rb, ub, pb) = (registry.clone(), url.clone(), params.clone()); + let waiter = tokio::spawn(async move { rb.get_store(ub, &pb).await }); + let mut queued = false; + for _ in 0..200 { + if registry.build_lock_waiters(&cache_key) == Some(2) { + queued = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!( + queued, + "waiter B did not queue on the build lock within timeout" + ); + + // Cancel B while it is parked, then confirm its session decremented the + // count back to 1. This is the assertion the old strong-count cleanup + // could not satisfy: a waiter cancelled mid-await left its hold behind. + waiter.abort(); + assert!( + waiter.await.unwrap_err().is_cancelled(), + "waiter must be cancelled, not completed" + ); + assert_eq!( + registry.build_lock_waiters(&cache_key), + Some(1), + "B's cancellation must decrement the waiter count, leaving only A" + ); + + // Release A; it finishes the build and drops its session, draining the + // entry entirely. + release.notify_one(); + assert!( + builder.await.unwrap().is_ok(), + "builder must complete successfully" + ); + assert_eq!( + registry.build_lock_waiters(&cache_key), + None, + "after the builder completes, the build-lock entry must be reclaimed" + ); + assert_eq!( + registry.build_locks_len(), + 0, + "no build-lock entries may leak", + ); + } + /// Regression: two callers resolving the same cache key must receive /// the same cached `Arc`. Single-flight correctness depends /// on this — see `test_get_store_coalesces_concurrent_misses` for the diff --git a/rust/lance-io/src/object_store/providers/huggingface.rs b/rust/lance-io/src/object_store/providers/huggingface.rs index cfef0440068..9107553cbea 100644 --- a/rust/lance-io/src/object_store/providers/huggingface.rs +++ b/rust/lance-io/src/object_store/providers/huggingface.rs @@ -10,6 +10,7 @@ use object_store_opendal::OpendalStore; use opendal::{Operator, services::Huggingface}; use url::Url; +use super::sanitized_authority; use crate::object_store::dynamic_opendal::DynamicOpenDalStore; use crate::object_store::parse_hf_repo_id; use crate::object_store::{ @@ -234,7 +235,15 @@ impl ObjectStoreProvider for HuggingfaceStoreProvider { _storage_options: Option<&HashMap>, ) -> Result { let repo_id = parse_hf_repo_id(url)?; - Ok(format!("{}${}@{}", url.scheme(), url.authority(), repo_id)) + // `sanitized_authority`, not `Url::authority()`: the latter keeps any + // embedded `userinfo@`, which would leak credentials into the cache key + // and the registry's cache-key debug logs. + Ok(format!( + "{}${}@{}", + url.scheme(), + sanitized_authority(url), + repo_id + )) } } @@ -404,6 +413,26 @@ mod tests { assert_eq!(prefix, "hf$datasets@acme/repo"); } + /// URL-embedded credentials (`userinfo@`) must never reach the cache-key + /// prefix — it is emitted via `Debug` and the registry's cache-key debug + /// logs. A URL that differs from the plain form only in `userinfo` must + /// collapse to the same prefix. + #[test] + fn calculate_prefix_strips_userinfo() { + let provider = HuggingfaceStoreProvider; + let with_creds = Url::parse("hf://user:s3cret@datasets/acme/repo/path").unwrap(); + let prefix = provider + .calculate_object_store_prefix(&with_creds, None) + .unwrap(); + assert_eq!(prefix, "hf$datasets@acme/repo"); + // Defense in depth: the secret literally cannot appear in the key. + assert!( + !prefix.contains("s3cret"), + "prefix leaked the token: {prefix}" + ); + assert!(!prefix.contains("user"), "prefix leaked userinfo: {prefix}"); + } + #[test] fn parse_invalid_url_errors() { let url = Url::parse("hf://datasets/only-owner").unwrap(); diff --git a/rust/lance-io/src/object_store/providers/shared_memory.rs b/rust/lance-io/src/object_store/providers/shared_memory.rs index 5e7ef4e8dbf..d55b7b1a0b0 100644 --- a/rust/lance-io/src/object_store/providers/shared_memory.rs +++ b/rust/lance-io/src/object_store/providers/shared_memory.rs @@ -6,6 +6,7 @@ use std::{ sync::{Arc, LazyLock, Mutex}, }; +use super::sanitized_authority; use crate::object_store::{ ObjectStore, ObjectStoreParams, ObjectStoreProvider, providers::memory::MemoryStoreProvider, }; @@ -13,11 +14,15 @@ use lance_core::error::Result; use object_store::{memory::InMemory, path::Path}; use url::Url; -/// Process-global pool of in-memory backends keyed by URL authority. +/// Process-global pool of in-memory backends keyed by sanitized URL authority. /// /// Different authorities map to different backends (act as "buckets"); same /// authority across any caller in the process resolves to the same `Arc`. -/// The pool grows for the lifetime of the process — entries are never evicted. +/// The key uses [`sanitized_authority`] (host[:port], userinfo stripped) so it +/// agrees with the registry cache-key prefix in `calculate_object_store_prefix` +/// — userinfo never distinguishes a backend, just as it never distinguishes a +/// cache entry. The pool grows for the lifetime of the process — entries are +/// never evicted. static SHARED_BACKENDS: LazyLock>>> = LazyLock::new(|| Mutex::new(HashMap::new())); @@ -25,7 +30,7 @@ fn shared_backend_for(url: &Url) -> Arc { SHARED_BACKENDS .lock() .expect("SHARED_BACKENDS mutex poisoned") - .entry(url.authority().to_string()) + .entry(sanitized_authority(url)) .or_insert_with(|| Arc::new(InMemory::new())) .clone() } @@ -67,7 +72,10 @@ impl ObjectStoreProvider for SharedMemoryStoreProvider { url: &Url, _storage_options: Option<&HashMap>, ) -> Result { - Ok(format!("shared-memory${}", url.authority())) + // `sanitized_authority`, not `Url::authority()`: the prefix is the + // registry cache key (logged via Debug / the cache-key debug lines), so + // any embedded `userinfo@` must be stripped. Matches `shared_backend_for`. + Ok(format!("shared-memory${}", sanitized_authority(url))) } } @@ -146,4 +154,41 @@ mod tests { assert_ne!(a, b); assert_eq!(a, "shared-memory$x"); } + + /// URL-embedded credentials must never reach the cache-key prefix — it is + /// the registry cache key, logged via the cache-key debug lines. A URL + /// differing only in `userinfo` collapses to the same prefix. + #[test] + fn calculate_prefix_strips_userinfo() { + let provider = SharedMemoryStoreProvider::default(); + let prefix = provider + .calculate_object_store_prefix( + &Url::parse("shared-memory://user:s3cret@bucket/p").unwrap(), + None, + ) + .unwrap(); + assert_eq!(prefix, "shared-memory$bucket"); + assert!( + !prefix.contains("s3cret"), + "prefix leaked the secret: {prefix}" + ); + assert!(!prefix.contains("user"), "prefix leaked userinfo: {prefix}"); + } + + /// `shared_backend_for` keys on the sanitized authority too, so two URLs + /// differing only in `userinfo` resolve to the *same* backend bytes — + /// keeping backend routing consistent with the cache-key prefix. + #[tokio::test] + async fn userinfo_does_not_isolate_backend() { + let (writer, _) = store_for("shared-memory://creds-bucket/").await; + writer + .inner + .put(&Path::from("f"), PutPayload::from_static(b"v")) + .await + .unwrap(); + + let (reader, _) = store_for("shared-memory://user:pw@creds-bucket/").await; + let bytes = reader.inner.get(&Path::from("f")).await.unwrap(); + assert_eq!(bytes.bytes().await.unwrap(), Bytes::from_static(b"v")); + } } From a62e5779135720168f19fc136f273748b3a648f1 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 22 Jun 2026 17:38:58 +0800 Subject: [PATCH 4/6] fix(io): sanitize authority in goosefs provider after merging main MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merging main brought in the new GooseFS object-store provider, whose calculate_object_store_prefix used raw url.authority() — re-introducing the userinfo-in-cache-key leak that was fixed for the other authority-based providers. Route it through sanitized_authority (host[:port] preserved, userinfo stripped) and add a userinfo-stripping test. Also sync java/lance-jni/Cargo.lock (rayon) to match the merged manifests so --locked builds stay consistent. --- .../src/object_store/providers/goosefs.rs | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/rust/lance-io/src/object_store/providers/goosefs.rs b/rust/lance-io/src/object_store/providers/goosefs.rs index 5d2a648522b..e10ee2e48ab 100644 --- a/rust/lance-io/src/object_store/providers/goosefs.rs +++ b/rust/lance-io/src/object_store/providers/goosefs.rs @@ -8,6 +8,7 @@ use object_store_opendal::OpendalStore; use opendal::{Operator, services::GooseFs}; use url::Url; +use super::sanitized_authority; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, @@ -215,12 +216,21 @@ impl ObjectStoreProvider for GooseFsStoreProvider { ) -> Result { // If a custom `goosefs_root` is provided, include it in the prefix so // that stores built with different roots don't accidentally collide. + // + // `sanitized_authority`, not `Url::authority()`: the prefix is the + // registry cache key (logged via Debug / the cache-key debug lines), so + // any embedded `userinfo@` must be stripped. host[:port] is preserved. let opts = StorageOptions(storage_options.cloned().unwrap_or_default()); let root = Self::resolve_root(&opts); if root == "/" { - Ok(format!("{}${}", url.scheme(), url.authority())) + Ok(format!("{}${}", url.scheme(), sanitized_authority(url))) } else { - Ok(format!("{}${}#{}", url.scheme(), url.authority(), root)) + Ok(format!( + "{}${}#{}", + url.scheme(), + sanitized_authority(url), + root + )) } } } @@ -337,6 +347,23 @@ mod tests { assert_ne!(default_prefix, custom_prefix); } + /// URL-embedded credentials must never reach the cache-key prefix — it is + /// the registry cache key, logged via the cache-key debug lines. host:port + /// is preserved so distinct clusters still get separate caches. + #[test] + fn test_calculate_object_store_prefix_strips_userinfo() { + let provider = GooseFsStoreProvider; + + let url = Url::parse("goosefs://user:s3cret@myhost:9200/data").unwrap(); + let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + assert_eq!(prefix, "goosefs$myhost:9200"); + assert!( + !prefix.contains("s3cret"), + "prefix leaked the secret: {prefix}" + ); + assert!(!prefix.contains("user"), "prefix leaked userinfo: {prefix}"); + } + #[test] fn test_resolve_master_addr_from_url() { let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap(); From 45eb72a1346b21624fdd378bcc3a81e5da6ad506 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Wed, 15 Jul 2026 19:33:44 +0800 Subject: [PATCH 5/6] fix(io,jni): address CodeRabbit review on default-open Session reuse lance-io/providers.rs: - rename `_session` -> `session` (it is read via `session.guard`) - drop `wrapping_add` in sweep scheduling; gate on the 0-based `fetch_add` value directly (`% N == N-1`), no self-owned counter arithmetic to overflow - `release_build_lock`: `debug_assert!` the entry exists and its waiter count is positive, `saturating_sub` + missing-entry no-op as the release-build fallback instead of silently underflowing - fix stale FailingProvider doc comment (failures serialize + retry per caller, they do not collapse into one call) - assert the exact `InvalidInput` variant and message in the serialized-failure test instead of bare `is_err()` lance-jni/blocking_dataset.rs: - make the test-only sharing override thread-local so parallel tests cannot observe each other's temporary setting - parameterize the `parse_disable_value` cases with `rstest` - replace the duplicated registry unit test with `open_wires_selected_registry_into_session`, which drives the real `BlockingDataset::open` path and pins that the selected registry is wired into the Session (global when sharing on, fresh when off); provider single-flight stays covered in providers.rs Co-Authored-By: Claude Fable 5 --- java/lance-jni/Cargo.lock | 44 ++- java/lance-jni/Cargo.toml | 9 +- java/lance-jni/src/blocking_dataset.rs | 300 ++++++++++---------- rust/lance-io/src/object_store/providers.rs | 72 +++-- 4 files changed, 255 insertions(+), 170 deletions(-) diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index f395d562027..524f249bb66 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -2576,6 +2576,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + [[package]] name = "futures-util" version = "0.3.32" @@ -4100,10 +4106,11 @@ dependencies = [ "prost", "prost-types", "roaring", + "rstest", "serde", "serde_json", + "tempfile", "tokio", - "url", "uuid", ] @@ -5798,6 +5805,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + [[package]] name = "reqsign-aliyun-oss" version = "3.1.0" @@ -6102,6 +6115,35 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if 1.0.4", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn", + "unicode-ident", +] + [[package]] name = "rust-ini" version = "0.21.3" diff --git a/java/lance-jni/Cargo.toml b/java/lance-jni/Cargo.toml index 18f36475528..cb279b85b07 100644 --- a/java/lance-jni/Cargo.toml +++ b/java/lance-jni/Cargo.toml @@ -60,9 +60,12 @@ prost-types = "0.14.1" chrono = "0.4.41" [dev-dependencies] -# Test-only: builds Url values for the JNI default-open coalescing test in -# src/blocking_dataset.rs. Workspace pins this same version transitively. -url = "2.5.7" +# Test-only: parameterizes the `parse_disable_value` classification cases in +# src/blocking_dataset.rs. Matches the version the top-level workspace pins. +rstest = "0.26.1" +# Test-only: temp dir for the default-open registry integration test in +# src/blocking_dataset.rs. Matches the version the top-level workspace pins. +tempfile = "3" [profile.dev] debug = "line-tables-only" diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 1ed20a26ba4..40ee08b7d55 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -77,9 +77,12 @@ pub const NATIVE_DATASET: &str = "nativeDatasetHandle"; /// and a separate warning when the variable is set to a value the parser does /// not recognize. /// -/// Backed by `AtomicU8` rather than `OnceLock` so unit tests can override -/// the resolved value via `DisableSharingTestGuard` without spawning a fresh -/// process per case. +/// Backed by `AtomicU8` so the env var is read and resolved at most once per +/// process (lazy first-read caching), rather than on every default-open. Unit +/// tests do NOT touch this atomic: they override the resolved value through a +/// test-only per-thread thread-local (see [`DisableSharingTestGuard`]), so +/// tests running in parallel on different threads cannot observe each other's +/// temporary settings. /// /// State encoding: /// - 0 = uninitialized (read env on next access) @@ -92,7 +95,25 @@ const SHARING_DISABLED: u8 = 2; static DISABLE_DEFAULT_REGISTRY_SHARING: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(SHARING_UNINIT); +// Test-only per-thread override for `disable_default_registry_sharing`. +// +// `DisableSharingTestGuard` sets this instead of the process-global +// `DISABLE_DEFAULT_REGISTRY_SHARING` atomic, so parallel tests on different +// threads never race on a shared setting and concurrent guards cannot restore +// a sibling's snapshot. +#[cfg(test)] +thread_local! { + static DISABLE_SHARING_TEST_OVERRIDE: std::cell::Cell> = + const { std::cell::Cell::new(None) }; +} + fn disable_default_registry_sharing() -> bool { + // Test-only per-thread override wins so unit tests never mutate the + // process-global atomic below. Compiled out entirely in production. + #[cfg(test)] + if let Some(over) = DISABLE_SHARING_TEST_OVERRIDE.with(|cell| cell.get()) { + return over; + } use std::sync::atomic::Ordering; match DISABLE_DEFAULT_REGISTRY_SHARING.load(Ordering::Relaxed) { SHARING_ENABLED => false, @@ -181,30 +202,26 @@ fn select_default_open_registry( } #[cfg(test)] -/// RAII guard that restores `DISABLE_DEFAULT_REGISTRY_SHARING` to its prior -/// value on drop. Tests should always go through this guard so they cannot -/// leak state into sibling tests run in the same process. +/// RAII guard that overrides the sharing flag for the *current thread* and +/// restores its prior per-thread value on drop. Tests should always go through +/// this guard so they cannot leak state into sibling tests run in the same +/// process. +/// +/// The override is a thread-local (see [`DISABLE_SHARING_TEST_OVERRIDE`]), not +/// the process-global atomic, so tests running in parallel on different threads +/// never observe each other's temporary settings. /// /// Each guard captures and restores its own snapshot, so nested guards compose /// correctly: LIFO drop order means the outer guard's snapshot always wins, /// returning the flag to the value seen before the outermost `set` call. struct DisableSharingTestGuard { - prior: u8, + prior: Option, } #[cfg(test)] impl DisableSharingTestGuard { fn set(disabled: bool) -> Self { - use std::sync::atomic::Ordering; - let prior = DISABLE_DEFAULT_REGISTRY_SHARING.load(Ordering::Relaxed); - DISABLE_DEFAULT_REGISTRY_SHARING.store( - if disabled { - SHARING_DISABLED - } else { - SHARING_ENABLED - }, - Ordering::Relaxed, - ); + let prior = DISABLE_SHARING_TEST_OVERRIDE.with(|cell| cell.replace(Some(disabled))); Self { prior } } } @@ -212,8 +229,7 @@ impl DisableSharingTestGuard { #[cfg(test)] impl Drop for DisableSharingTestGuard { fn drop(&mut self) { - use std::sync::atomic::Ordering; - DISABLE_DEFAULT_REGISTRY_SHARING.store(self.prior, Ordering::Relaxed); + DISABLE_SHARING_TEST_OVERRIDE.with(|cell| cell.set(self.prior)); } } @@ -3891,6 +3907,7 @@ fn inner_get_zonemap_stats<'local>( #[cfg(test)] mod default_open_registry_tests { use super::*; + use rstest::rstest; /// Default-open path must reuse the process-wide /// `GLOBAL_OBJECT_STORE_REGISTRY` so concurrent opens can coalesce on its @@ -3923,12 +3940,12 @@ mod default_open_registry_tests { /// `disable_default_registry_sharing()` must honor the in-process override /// hook. This guards against future refactors that re-introduce a once-cell - /// that bypasses the AtomicU8 state. + /// that bypasses the per-thread test override. #[test] fn disable_flag_honors_test_override() { - // RAII guard restores prior state at end-of-scope so this test cannot - // leak SHARING_ENABLED/DISABLED state into sibling tests that rely on - // env-driven first-read resolution. + // RAII guard restores the prior per-thread override at end-of-scope, so + // this test cannot leak its setting into sibling tests sharing this + // thread, and — being thread-local — never races tests on other threads. let _g = DisableSharingTestGuard::set(true); assert!(disable_default_registry_sharing()); @@ -3936,134 +3953,131 @@ mod default_open_registry_tests { assert!(!disable_default_registry_sharing()); } - /// Truthy spellings (case + whitespace insensitive) must parse to `Some(true)`. - /// These are the values that opt the operator OUT of registry sharing — - /// getting any of them wrong silently re-enables coalescing in a config - /// that asked for isolation, so accept-list correctness is load-bearing. - #[test] - fn parse_disable_value_accepts_truthy() { - for raw in [ - "1", "true", "yes", "TRUE", "Yes", "True", " 1 ", "\ttrue\n", - ] { - assert_eq!( - parse_disable_value(raw), - Some(true), - "expected {raw:?} to parse as Some(true)", - ); - } - } - - /// Documented falsy spellings — empty, whitespace-only, `0`/`false`/`no` — - /// must parse to `Some(false)`. These are explicit keep-default opts; - /// distinguishing them from unrecognized noise lets the caller skip the - /// warning log. - #[test] - fn parse_disable_value_accepts_falsy() { - for raw in ["", " ", "\t\n", "0", "false", "no", "FALSE", " No "] { - assert_eq!( - parse_disable_value(raw), - Some(false), - "expected {raw:?} to parse as Some(false)", - ); - } + /// Classification contract for [`parse_disable_value`], one case per input + /// so a regression names the exact spelling that broke. + /// + /// - Truthy (`Some(true)`): case- and whitespace-insensitive `1`/`true`/`yes`. + /// These opt the operator OUT of registry sharing, so accept-list + /// correctness is load-bearing — a miss silently re-enables coalescing in a + /// config that asked for isolation. + /// - Falsy (`Some(false)`): empty, whitespace-only, `0`/`false`/`no` — explicit + /// keep-default opts, distinguished from noise so the caller can skip warning. + /// - Unrecognized (`None`): typos, `on`, `y`, numeric noise — so + /// [`disable_default_registry_sharing`] warns that the escape hatch is ignored. + #[rstest] + #[case::one("1", Some(true))] + #[case::true_lower("true", Some(true))] + #[case::yes("yes", Some(true))] + #[case::true_upper("TRUE", Some(true))] + #[case::yes_title("Yes", Some(true))] + #[case::true_title("True", Some(true))] + #[case::padded_one(" 1 ", Some(true))] + #[case::whitespace_true("\ttrue\n", Some(true))] + #[case::empty("", Some(false))] + #[case::space(" ", Some(false))] + #[case::tab_newline("\t\n", Some(false))] + #[case::zero("0", Some(false))] + #[case::false_lower("false", Some(false))] + #[case::no("no", Some(false))] + #[case::false_upper("FALSE", Some(false))] + #[case::padded_no(" No ", Some(false))] + #[case::off("off", None)] + #[case::disable("disable", None)] + #[case::two("2", None)] + #[case::one_point_zero("1.0", None)] + #[case::true_bang("true!", None)] + #[case::y("y", None)] + #[case::on("on", None)] + #[case::enable("enable", None)] + fn parse_disable_value_classifies(#[case] raw: &str, #[case] expected: Option) { + assert_eq!( + parse_disable_value(raw), + expected, + "unexpected classification for {raw:?}", + ); } - /// Anything outside the accept-list — typos, `on`, `y`, numeric noise — - /// must parse to `None` so [`disable_default_registry_sharing`] can warn - /// the operator that their escape hatch is being silently ignored. + /// Integration pin for the core behavior change of this PR: a default open + /// (`session = None`) must build its `Session` on the registry chosen by + /// `select_default_open_registry`, i.e. the process-wide + /// `GLOBAL_OBJECT_STORE_REGISTRY` when sharing is enabled and a fresh + /// registry when the escape hatch disables it. + /// + /// This closes the gap the unit tests leave open: they pin + /// `select_default_open_registry` in isolation, but nothing asserts that + /// `BlockingDataset::open` actually *wires that registry into the Session*. + /// A regression that stops threading the selected registry into + /// `LanceSession::new` — or hardcodes a fresh one — passes every other test + /// but fails here. Provider-level single-flight itself is covered by + /// `test_get_store_coalesces_concurrent_misses` in `providers.rs`, so this + /// test deliberately checks registry *identity* (an `Arc::ptr_eq`) rather + /// than re-testing coalescing: identity is what the JNI layer owns, needs + /// no custom scheme provider, and so never leaks state into + /// `GLOBAL_OBJECT_STORE_REGISTRY`. #[test] - fn parse_disable_value_returns_none_for_unrecognized() { - for raw in ["off", "disable", "2", "1.0", "true!", "y", "on", "enable"] { - assert_eq!( - parse_disable_value(raw), + fn open_wires_selected_registry_into_session() { + use std::sync::Arc; + + use arrow_array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchemaDef}; + + // Write a tiny local dataset (pure Rust, no JVM) to open below. + let tmp = tempfile::tempdir().expect("tempdir"); + let uri = tmp.path().join("ds").to_str().unwrap().to_string(); + let schema = Arc::new(ArrowSchemaDef::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + BlockingDataset::write(reader, &uri, None).expect("write dataset"); + + // `open` with no session and no key-distinguishing input takes the + // default path. The `DisableSharingTestGuard` override is read + // synchronously inside `open` (outside `RT.block_on`), so it applies. + let open_default = |uri: &str| { + BlockingDataset::open( + uri, None, - "expected {raw:?} to parse as None (unrecognized)", - ); - } - } - - /// End-to-end pin for PR#1's two-piece contract at the JNI boundary: - /// `select_default_open_registry` must hand back a registry whose - /// `get_store` coalesces concurrent cold builds for the same URI. - /// A regression in single-flight (dropped from `get_store`) would - /// surface here as `provider.builds > 1` or callers observing - /// different `Arc` instances. - /// - /// The complementary "sharing-on returns the process-wide registry" - /// invariant is covered by - /// `select_default_open_registry_reuses_global_when_sharing_enabled`, - /// so this test deliberately uses the sharing-disabled path to take a - /// fresh registry — registering a one-off provider on - /// `GLOBAL_OBJECT_STORE_REGISTRY` would leak that scheme's `Arc` for - /// the rest of the test binary's lifetime. - /// - /// Provider counts invocations and sleeps before returning so concurrent - /// callers reliably queue on the build lock before the first build - /// completes — without the delay the winner returns before any contention - /// develops and the test would pass even if single-flight were broken. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn jni_default_open_coalesces_concurrent_cold_builds() { - use std::sync::atomic::{AtomicU64, Ordering}; - use std::time::Duration; - - use lance_core::Result as LanceResult; - use lance_io::object_store::ObjectStore as LanceObjectStore; - use lance_io::object_store::providers::memory::MemoryStoreProvider; - use lance_io::object_store::{ObjectStoreParams, ObjectStoreProvider}; - use url::Url; - - #[derive(Debug, Default)] - struct CountingProvider { - builds: AtomicU64, - } - - #[async_trait::async_trait] - impl ObjectStoreProvider for CountingProvider { - async fn new_store( - &self, - base_path: Url, - params: &ObjectStoreParams, - ) -> LanceResult { - self.builds.fetch_add(1, Ordering::Relaxed); - tokio::time::sleep(Duration::from_millis(20)).await; - MemoryStoreProvider.new_store(base_path, params).await - } - } + None, + 0, + 0, + HashMap::new(), + HashMap::new(), + None, + None, + None, + None, + None, + false, + ) + .expect("open dataset") + }; - let provider = Arc::new(CountingProvider::default()); - // Take a fresh registry via the sharing-disabled path so the - // test-only `pr1-jni-coalesce` provider never touches - // `GLOBAL_OBJECT_STORE_REGISTRY`. - let registry = select_default_open_registry(true); - registry.insert("pr1-jni-coalesce", provider.clone()); - - let url = Url::parse("pr1-jni-coalesce://x").unwrap(); - let params = ObjectStoreParams::default(); - - let n = 8; - let mut handles = Vec::with_capacity(n); - for _ in 0..n { - let registry = registry.clone(); - let url = url.clone(); - let params = params.clone(); - handles.push(tokio::spawn(async move { - registry.get_store(url, ¶ms).await.unwrap() - })); - } - let mut stores = Vec::with_capacity(n); - for h in handles { - stores.push(h.await.expect("task must not panic")); + // Sharing enabled (default): the Session must reuse the global registry. + { + let _g = DisableSharingTestGuard::set(false); + let ds = open_default(&uri); + let registry = ds.inner.session().store_registry(); + assert!( + Arc::ptr_eq(®istry, &crate::GLOBAL_OBJECT_STORE_REGISTRY), + "default open must build its Session on GLOBAL_OBJECT_STORE_REGISTRY", + ); } - assert_eq!( - provider.builds.load(Ordering::Relaxed), - 1, - "single-flight must collapse N concurrent cold misses into one provider call", - ); - for store in &stores[1..] { + // Escape hatch on: the Session must get a fresh, non-global registry. + { + let _g = DisableSharingTestGuard::set(true); + let ds = open_default(&uri); + let registry = ds.inner.session().store_registry(); assert!( - Arc::ptr_eq(&stores[0], store), - "every coalesced waiter must observe the same Arc", + !Arc::ptr_eq(®istry, &crate::GLOBAL_OBJECT_STORE_REGISTRY), + "disabled-sharing open must NOT alias the global registry", ); } } diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index 0ea0829a275..bc4e0a31cc9 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -411,14 +411,14 @@ impl ObjectStoreRegistry { // `lock_owned()` await (see `BuildLockSession`) so a waiter cancelled // while queued still cleans up. let lock = self.acquire_build_lock(&cache_key); - let mut _session = BuildLockSession { + let mut session = BuildLockSession { registry: self, cache_key: cache_key.clone(), guard: None, }; - // Cancellation point: `_session` already exists, so its Drop cleans up + // Cancellation point: `session` already exists, so its Drop cleans up // even if we are cancelled here before the guard is assigned. - _session.guard = Some(lock.lock_owned().await); + session.guard = Some(lock.lock_owned().await); // Re-check after acquiring the lock — coalesced waiters become hits. if let Some(store) = self.lookup_cached(&cache_key) { @@ -430,17 +430,17 @@ impl ObjectStoreRegistry { return Ok(store); } - // `fetch_add` returns the prior value; `+ 1` is this attempt's - // 1-indexed sequence number. Use it (not a separate counter) to - // gate the opportunistic sweep below: cadence is approximate - // because cold attempts that fail still bump `misses` but skip - // the sweep block, so a failure at a multiple-of-N sequence - // skips that sweep cycle entirely; the next candidate is the - // next multiple-of-N cold attempt that succeeds. That's fine — - // the sweep is a footprint trim, not a correctness invariant. - // `wrapping_add` only silences debug-build overflow checks; at - // 10M cold opens/sec the wrap is >58 millennia away. - let cold_attempt_nth = self.misses.fetch_add(1, Ordering::Relaxed).wrapping_add(1); + // `fetch_add` returns the prior value: this attempt's 0-based + // sequence index. Use it (not a separate counter) to gate the + // opportunistic sweep below: cadence is approximate because cold + // attempts that fail still bump `misses` but skip the sweep block, + // so a failure at a sweep-eligible index skips that sweep cycle + // entirely; the next candidate is the next sweep-eligible cold + // attempt that succeeds. That's fine — the sweep is a footprint + // trim, not a correctness invariant. The raw `fetch_add` value is + // used as-is (no `+ 1`), so there is no counter arithmetic of our + // own that could overflow. + let cold_attempt_index = self.misses.fetch_add(1, Ordering::Relaxed); log::debug!( "ObjectStoreRegistry: cold build starting for cache_key path={:?}", cache_key.0, @@ -482,12 +482,13 @@ impl ObjectStoreRegistry { // insert. Per-key `lookup_cached` and full sweep via `active_stores()` // remain unchanged; this is purely a per-insert cost lever. // - // Cadence is "every 64th cold attempt" (1-indexed via the pre-bump - // above), so the first sweep fires at the 64th, not the 0th, and - // a near-empty map never pays the retain cost on its very first - // cold build. + // Cadence is "every 64th cold attempt": with a 0-based index the + // eligible attempts are indices 63, 127, 191, … so the first sweep + // fires at the 64th cold build, not the very first, and a + // near-empty map never pays the retain cost on its first cold + // build. const SWEEP_INTERVAL: u64 = 64; - let should_sweep = cold_attempt_nth.is_multiple_of(SWEEP_INTERVAL); + let should_sweep = cold_attempt_index % SWEEP_INTERVAL == SWEEP_INTERVAL - 1; { let mut cache_lock = self .active_stores @@ -545,11 +546,28 @@ impl ObjectStoreRegistry { .build_locks .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); + // RAII pairs every `acquire_build_lock` with exactly one release, so + // the entry must exist and its waiter count must be positive here. A + // violation means the acquire/release accounting is broken — surface + // it loudly in debug/test, and fall back safely (saturating, no-op on + // a missing entry) in release so a bug can't underflow `usize` or + // panic in production. if let Some(entry) = locks.get_mut(cache_key) { - entry.waiters -= 1; + debug_assert!( + entry.waiters > 0, + "release_build_lock: waiter count already zero for an existing entry — \ + acquire/release pairing is broken", + ); + entry.waiters = entry.waiters.saturating_sub(1); if entry.waiters == 0 { locks.remove(cache_key); } + } else { + debug_assert!( + false, + "release_build_lock: no build-lock entry for cache_key — \ + release called without a matching acquire", + ); } } @@ -983,8 +1001,8 @@ mod tests { /// A provider that returns an `Err` after a yield, simulating a normal /// (non-panic) build failure such as bad credentials or a network error. - /// Counts invocations so the test can assert the single-flight collapsed - /// N concurrent failures into 1 underlying call. + /// Counts invocations so the test can assert failures are serialized but + /// retried once per caller. #[derive(Debug, Default)] struct FailingProvider { builds: AtomicU64, @@ -1049,7 +1067,15 @@ mod tests { let mut err_count = 0; for h in handles { let result = h.await.expect("task must not panic"); - assert!(result.is_err(), "every waiter must surface Err"); + let err = result.expect_err("every waiter must surface Err"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput from the failing provider, got: {err:?}", + ); + assert!( + err.to_string().contains("simulated cold-build failure"), + "error must carry the provider's message, got: {err}", + ); err_count += 1; } assert_eq!(err_count, n, "all N callers must observe an error"); From 6d5317b8d6ea194da7d5ce9af93056c7b9dd65b1 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 16 Jul 2026 00:09:35 +0800 Subject: [PATCH 6/6] docs(io): make StorageOptionsProvider::provider_id tenant-isolation contract a MUST The registry keys its ObjectStore cache solely on `provider_id()`, so it is the only isolation boundary between providers. Spell out the security consequence on the trait method: a provider that vends principal-specific credentials MUST encode that identity in its ID, because two providers returning the same ID share one cached ObjectStore (and thus one resolved credential chain). Doc-only; no behavior change. Co-Authored-By: Claude Fable 5 --- rust/lance-io/src/object_store/storage_options.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/rust/lance-io/src/object_store/storage_options.rs b/rust/lance-io/src/object_store/storage_options.rs index fa5f4995afd..e9158889ea7 100644 --- a/rust/lance-io/src/object_store/storage_options.rs +++ b/rust/lance-io/src/object_store/storage_options.rs @@ -98,6 +98,17 @@ pub trait StorageOptionsProvider: Send + Sync + fmt::Debug { /// For example: `"namespace[dir(root=/data)],table[db$schema$table1]"` /// /// The ID should uniquely identify the provider's configuration. + /// + /// # Security + /// + /// The ObjectStore registry keys its cache on this ID, so it is the *only* + /// isolation boundary between providers. If this provider vends credentials + /// that differ by tenant, user, role, or any other principal, `provider_id()` + /// **MUST** encode that distinguishing identity. Two providers that return the + /// same ID share one cached `ObjectStore` — and therefore one resolved + /// credential chain — so a constant ID across principals causes the first + /// caller's credentials to be reused for every subsequent caller + /// (cross-tenant credential sharing). fn provider_id(&self) -> String; }