From 02fbe212e8c14b4e878e4f4a1a24f4769a494c41 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 22:56:40 +0000 Subject: [PATCH] fix(rust): normalize AWS credentials in built-in provider --- Cargo.lock | 1 - Cargo.toml | 4 + rust/lancedb/src/database/namespace.rs | 11 +- rust/lancedb/src/io/object_store.rs | 1287 ++------ rust/lancedb/src/table.rs | 66 +- vendor/lance-io/Cargo.toml | 74 + vendor/lance-io/LANCEDB_PATCH.md | 13 + vendor/lance-io/README.md | 9 + vendor/lance-io/src/ffi.rs | 60 + vendor/lance-io/src/lib.rs | 387 +++ vendor/lance-io/src/local.rs | 331 +++ vendor/lance-io/src/object_reader.rs | 464 +++ vendor/lance-io/src/object_store.rs | 1879 ++++++++++++ .../src/object_store/dynamic_credentials.rs | 411 +++ .../src/object_store/dynamic_opendal.rs | 368 +++ .../lance-io/src/object_store/list_retry.rs | 403 +++ vendor/lance-io/src/object_store/metrics.rs | 1821 ++++++++++++ vendor/lance-io/src/object_store/providers.rs | 520 ++++ .../src/object_store/providers/aws.rs | 1531 ++++++++++ .../src/object_store/providers/azure.rs | 674 +++++ .../src/object_store/providers/gcp.rs | 262 ++ .../src/object_store/providers/goosefs.rs | 379 +++ .../src/object_store/providers/huggingface.rs | 440 +++ .../src/object_store/providers/local.rs | 139 + .../src/object_store/providers/memory.rs | 84 + .../src/object_store/providers/oss.rs | 285 ++ .../object_store/providers/shared_memory.rs | 149 + .../src/object_store/providers/tencent.rs | 122 + .../src/object_store/providers/tos.rs | 300 ++ .../src/object_store/storage_options.rs | 1244 ++++++++ .../lance-io/src/object_store/test_utils.rs | 25 + vendor/lance-io/src/object_store/throttle.rs | 2097 +++++++++++++ vendor/lance-io/src/object_store/tracing.rs | 401 +++ vendor/lance-io/src/object_writer.rs | 806 +++++ vendor/lance-io/src/scheduler.rs | 2585 +++++++++++++++++ vendor/lance-io/src/scheduler/lite.rs | 955 ++++++ vendor/lance-io/src/spill.rs | 543 ++++ vendor/lance-io/src/stream.rs | 76 + vendor/lance-io/src/testing.rs | 51 + vendor/lance-io/src/traits.rs | 177 ++ vendor/lance-io/src/uring.rs | 84 + vendor/lance-io/src/uring/current_thread.rs | 430 +++ .../src/uring/current_thread_future.rs | 102 + vendor/lance-io/src/uring/future.rs | 46 + vendor/lance-io/src/uring/reader.rs | 301 ++ vendor/lance-io/src/uring/requests.rs | 54 + vendor/lance-io/src/uring/tests.rs | 392 +++ vendor/lance-io/src/uring/thread.rs | 396 +++ vendor/lance-io/src/utils.rs | 291 ++ vendor/lance-io/src/utils/tracking_store.rs | 583 ++++ 50 files changed, 22988 insertions(+), 1125 deletions(-) create mode 100644 vendor/lance-io/Cargo.toml create mode 100644 vendor/lance-io/LANCEDB_PATCH.md create mode 100644 vendor/lance-io/README.md create mode 100644 vendor/lance-io/src/ffi.rs create mode 100644 vendor/lance-io/src/lib.rs create mode 100644 vendor/lance-io/src/local.rs create mode 100644 vendor/lance-io/src/object_reader.rs create mode 100644 vendor/lance-io/src/object_store.rs create mode 100644 vendor/lance-io/src/object_store/dynamic_credentials.rs create mode 100644 vendor/lance-io/src/object_store/dynamic_opendal.rs create mode 100644 vendor/lance-io/src/object_store/list_retry.rs create mode 100644 vendor/lance-io/src/object_store/metrics.rs create mode 100644 vendor/lance-io/src/object_store/providers.rs create mode 100644 vendor/lance-io/src/object_store/providers/aws.rs create mode 100644 vendor/lance-io/src/object_store/providers/azure.rs create mode 100644 vendor/lance-io/src/object_store/providers/gcp.rs create mode 100644 vendor/lance-io/src/object_store/providers/goosefs.rs create mode 100644 vendor/lance-io/src/object_store/providers/huggingface.rs create mode 100644 vendor/lance-io/src/object_store/providers/local.rs create mode 100644 vendor/lance-io/src/object_store/providers/memory.rs create mode 100644 vendor/lance-io/src/object_store/providers/oss.rs create mode 100644 vendor/lance-io/src/object_store/providers/shared_memory.rs create mode 100644 vendor/lance-io/src/object_store/providers/tencent.rs create mode 100644 vendor/lance-io/src/object_store/providers/tos.rs create mode 100644 vendor/lance-io/src/object_store/storage_options.rs create mode 100644 vendor/lance-io/src/object_store/test_utils.rs create mode 100644 vendor/lance-io/src/object_store/throttle.rs create mode 100644 vendor/lance-io/src/object_store/tracing.rs create mode 100644 vendor/lance-io/src/object_writer.rs create mode 100644 vendor/lance-io/src/scheduler.rs create mode 100644 vendor/lance-io/src/scheduler/lite.rs create mode 100644 vendor/lance-io/src/spill.rs create mode 100644 vendor/lance-io/src/stream.rs create mode 100644 vendor/lance-io/src/testing.rs create mode 100644 vendor/lance-io/src/traits.rs create mode 100644 vendor/lance-io/src/uring.rs create mode 100644 vendor/lance-io/src/uring/current_thread.rs create mode 100644 vendor/lance-io/src/uring/current_thread_future.rs create mode 100644 vendor/lance-io/src/uring/future.rs create mode 100644 vendor/lance-io/src/uring/reader.rs create mode 100644 vendor/lance-io/src/uring/requests.rs create mode 100644 vendor/lance-io/src/uring/tests.rs create mode 100644 vendor/lance-io/src/uring/thread.rs create mode 100644 vendor/lance-io/src/utils.rs create mode 100644 vendor/lance-io/src/utils/tracking_store.rs diff --git a/Cargo.lock b/Cargo.lock index 34826e2d6..3b1e74919 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5206,7 +5206,6 @@ dependencies = [ [[package]] name = "lance-io" version = "11.0.0-beta.2" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.2#35da5d920159b49d1b53032652f7615ab699c160" dependencies = [ "arrow", "arrow-array", diff --git a/Cargo.toml b/Cargo.toml index a879d1f8b..3ca51a022 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = ["rust/lancedb", "nodejs", "python"] +exclude = ["vendor/lance-io"] resolver = "2" [workspace.package] @@ -67,6 +68,9 @@ regex = "1.10" semver = "1.0.25" chrono = "0.4" +[patch."https://github.com/lance-format/lance.git"] +lance-io = { path = "vendor/lance-io" } + [profile.ci] debug = "line-tables-only" inherits = "dev" diff --git a/rust/lancedb/src/database/namespace.rs b/rust/lancedb/src/database/namespace.rs index f46ad715d..5e2856b3a 100644 --- a/rust/lancedb/src/database/namespace.rs +++ b/rust/lancedb/src/database/namespace.rs @@ -34,7 +34,7 @@ use crate::database::read_freshness::{ FreshnessBaselines, ReadFreshnessContextProvider, TableFreshness, }; use crate::error::{Error, Result}; -use crate::io::object_store::{atomic_aws_session, install_atomic_aws_provider}; +use crate::io::object_store::atomic_aws_session; use crate::table::{NativeTable, map_namespace_lance_error}; use lance::dataset::WriteMode; @@ -102,9 +102,6 @@ impl LanceNamespaceDatabase { session: Option>, namespace_client_pushdown_operations: HashSet, ) -> Self { - if let Some(session) = &session { - install_atomic_aws_provider(session); - } // Client is pre-built, so we can't install the freshness provider here; // baselines are still tracked for a uniform bump path. let delimiter = resolve_delimiter(&namespace_client_properties); @@ -157,9 +154,9 @@ impl LanceNamespaceDatabase { pushdown_operations: HashSet, new_table_config: NewTableConfig, ) -> Result { - // Namespace construction needs a protected session even when the connection did not - // supply one. Keep the original option separately so per-operation sessions retain - // precedence when tables are opened or created later. + // Namespace construction needs a shared session even when the connection did not supply + // one. Keep the original option separately so per-operation sessions retain precedence + // when tables are opened or created later. let builder_session = atomic_aws_session(session.clone()); let mut builder = ConnectBuilder::new(ns_impl); for (key, value) in ns_properties.clone() { diff --git a/rust/lancedb/src/io/object_store.rs b/rust/lancedb/src/io/object_store.rs index cddbd28ea..c8aaff406 100644 --- a/rust/lancedb/src/io/object_store.rs +++ b/rust/lancedb/src/io/object_store.rs @@ -5,84 +5,24 @@ use std::{collections::HashMap, fmt::Formatter, sync::Arc}; -#[cfg(feature = "aws")] -use std::{ - ops::Range, - sync::{LazyLock, Mutex, Weak}, -}; - -#[cfg(feature = "aws")] -use bytes::Bytes; use futures::{StreamExt, TryFutureExt, stream::BoxStream}; -#[cfg(feature = "aws")] -use futures::{TryStreamExt, stream}; use lance::io::{ObjectStoreParams, WrappingObjectStore}; -#[cfg(feature = "aws")] -use lance_io::object_store::{ - ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry, StorageOptions, - providers::aws::build_aws_credential, -}; use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; +#[cfg(feature = "aws")] +use object_store::aws::AmazonS3ConfigKey; use object_store::{ CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, UploadPart, path::Path, }; #[cfg(feature = "aws")] -use object_store::{ - CredentialProvider, RenameOptions, StaticCredentialProvider, - aws::{AmazonS3ConfigKey, AwsCredential}, -}; -#[cfg(feature = "aws")] use std::str::FromStr; -#[cfg(feature = "aws")] -use tokio::sync::RwLock as TokioRwLock; use async_trait::async_trait; #[cfg(test)] pub mod io_tracking; -#[cfg(feature = "aws")] -fn explicit_aws_credential( - storage_options: &HashMap, -) -> lance_core::Result> { - let aws_options = storage_options - .iter() - .filter_map(|(key, value)| { - AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) - .ok() - .map(|key| (key, value)) - }) - .collect::>(); - - let key_id = aws_options.get(&AmazonS3ConfigKey::AccessKeyId); - let secret_key = aws_options.get(&AmazonS3ConfigKey::SecretAccessKey); - let token = aws_options.get(&AmazonS3ConfigKey::Token); - if key_id.is_none() && secret_key.is_none() && token.is_none() { - return Ok(None); - } - let (Some(key_id), Some(secret_key)) = (key_id, secret_key) else { - return Err(lance_core::Error::invalid_input( - "Explicit AWS credentials require both aws_access_key_id and aws_secret_access_key", - )); - }; - - Ok(Some(AwsCredential { - key_id: (*key_id).clone(), - secret_key: (*secret_key).clone(), - token: token.map(|token| (*token).clone()), - })) -} - -#[cfg(feature = "aws")] -fn has_aws_credential_member(storage_options: &HashMap) -> bool { - storage_options.keys().any(|key| { - AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) - .is_ok_and(|key| is_aws_credential_key(&key)) - }) -} - #[cfg(feature = "aws")] fn is_aws_credential_key(key: &AmazonS3ConfigKey) -> bool { matches!( @@ -104,552 +44,14 @@ pub(crate) fn is_aws_credential_option(_key: &str) -> bool { false } -#[cfg(feature = "aws")] -fn canonical_noncredential_options( - storage_options: &HashMap, -) -> HashMap { - storage_options - .iter() - .filter_map( - |(key, value)| match AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) { - Ok(config_key) if is_aws_credential_key(&config_key) => None, - Ok(config_key) => Some((config_key.as_ref().to_string(), value.clone())), - Err(_) => Some((key.clone(), value.clone())), - }, - ) - .collect() -} - -#[cfg(feature = "aws")] -fn insert_aws_credential(options: &mut HashMap, credential: AwsCredential) { - options.insert( - AmazonS3ConfigKey::AccessKeyId.as_ref().to_string(), - credential.key_id, - ); - options.insert( - AmazonS3ConfigKey::SecretAccessKey.as_ref().to_string(), - credential.secret_key, - ); - if let Some(token) = credential.token { - options.insert(AmazonS3ConfigKey::Token.as_ref().to_string(), token); - } else { - // Lance's environment merge treats an empty value as an explicit sentinel, while - // OpenDAL ignores an empty session token. This blocks a foreign ambient token without - // changing the semantics of long-lived key/secret credentials. - options.insert(AmazonS3ConfigKey::Token.as_ref().to_string(), String::new()); - } -} - -/// Merge an OpenDAL configuration without ever combining two AWS credential families. -#[cfg(feature = "aws")] -fn atomic_opendal_options( - base_options: &HashMap, - dynamic_options: &HashMap, - credential: Option, - environment: impl IntoIterator, -) -> lance_core::Result> { - let mut options = canonical_noncredential_options(base_options); - for (key, value) in environment { - match AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) { - Ok(config_key) if is_aws_credential_key(&config_key) => {} - Ok(config_key) => { - options - .entry(config_key.as_ref().to_string()) - .or_insert(value); - } - Err(_) => {} - } - } - options.extend(canonical_noncredential_options(dynamic_options)); - - let credential = match credential { - Some(credential) => Some(credential), - None if has_aws_credential_member(dynamic_options) => { - explicit_aws_credential(dynamic_options)? - } - None => explicit_aws_credential(base_options)?, - }; - if let Some(credential) = credential { - insert_aws_credential(&mut options, credential); - // OpenDAL must not run another credential lookup after an explicit family wins. - options.insert("disable_config_load".to_string(), "true".to_string()); - } - Ok(options) -} - -#[cfg(feature = "aws")] -#[derive(Debug)] -struct AtomicAccessorAwsCredentialProvider { - accessor: Arc, - fallback: Option, -} - -#[cfg(feature = "aws")] -#[async_trait] -impl CredentialProvider for AtomicAccessorAwsCredentialProvider { - type Credential = AwsCredential; - - async fn get_credential(&self) -> object_store::Result> { - let options = self - .accessor - .get_storage_options() - .await - .map_err(|error| Error::Generic { - store: "AtomicAwsCredentialProvider", - source: Box::new(error), - })? - .0; - match explicit_aws_credential(&options).map_err(|error| Error::Generic { - store: "AtomicAwsCredentialProvider", - source: Box::new(error), - })? { - Some(credential) => Ok(Arc::new(credential)), - None => match &self.fallback { - Some(fallback) => fallback.get_credential().await, - None => Err(Error::Generic { - store: "AtomicAwsCredentialProvider", - source: "Explicit AWS credentials require both aws_access_key_id and aws_secret_access_key".into(), - }), - }, - } - } -} - -#[cfg(feature = "aws")] -#[derive(Debug, Clone)] -struct CachedProviderStore { - config: HashMap, - store: Arc, -} - -/// Store that refreshes credentials by rebuilding through the registered provider. -/// -/// Re-entering the original provider preserves custom encryption, authorization, wrapping, and -/// backend behavior while still letting built-in OpenDAL stores consume refreshed credentials. -#[cfg(feature = "aws")] -#[derive(Clone)] -struct AtomicProviderStore { - provider: Arc, - base_path: url::Url, - base_params: ObjectStoreParams, - base_options: Arc>, - accessor: Option>, - aws_credentials: Option, - cache: Arc>>, -} - -#[cfg(feature = "aws")] -impl std::fmt::Debug for AtomicProviderStore { - fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("AtomicProviderStore") - .field("base_path", &self.base_path) - .field("accessor", &self.accessor) - .finish() - } -} - -#[cfg(feature = "aws")] -impl std::fmt::Display for AtomicProviderStore { - fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { - write!(formatter, "AtomicProviderStore({})", self.base_path) - } -} - -#[cfg(feature = "aws")] -impl AtomicProviderStore { - async fn current_config(&self) -> lance_core::Result> { - let dynamic_options = match &self.accessor { - Some(accessor) if accessor.has_provider() => accessor.get_storage_options().await?.0, - _ => HashMap::new(), - }; - let credential = match &self.aws_credentials { - Some(provider) => Some({ - let credential = provider - .get_credential() - .await - .map_err(|error| lance_core::Error::io_source(Box::new(error)))?; - AwsCredential { - key_id: credential.key_id.clone(), - secret_key: credential.secret_key.clone(), - token: credential.token.clone(), - } - }), - None => None, - }; - atomic_opendal_options( - &self.base_options, - &dynamic_options, - credential, - std::env::vars_os().filter_map(|(key, value)| { - Some((key.into_string().ok()?, value.into_string().ok()?)) - }), - ) - } - - async fn build_store( - &self, - config: &HashMap, - ) -> lance_core::Result { - let mut params = self.base_params.clone(); - params.aws_credentials = None; - set_storage_options(&mut params, config.clone(), None); - self.provider - .new_store(self.base_path.clone(), ¶ms) - .await - } - - async fn initialize_store(&self) -> lance_core::Result { - let config = self.current_config().await?; - let store = self.build_store(&config).await?; - *self.cache.write().await = Some(CachedProviderStore { - config, - store: store.inner.clone(), - }); - Ok(store) - } - - async fn current_store(&self) -> lance_core::Result> { - let config = self.current_config().await?; - - { - let cache = self.cache.read().await; - if let Some(cached) = cache.as_ref() - && cached.config == config - { - return Ok(cached.store.clone()); - } - } - - let store = self.build_store(&config).await?.inner; - let mut cache = self.cache.write().await; - if let Some(cached) = cache.as_ref() - && cached.config == config - { - return Ok(cached.store.clone()); - } - *cache = Some(CachedProviderStore { - config, - store: store.clone(), - }); - Ok(store) - } - - fn map_store_error(error: lance_core::Error) -> Error { - Error::Generic { - store: "AtomicProviderStore", - source: Box::new(error), - } - } -} - -#[cfg(feature = "aws")] -#[async_trait] -impl ObjectStore for AtomicProviderStore { - async fn put_opts( - &self, - location: &Path, - payload: PutPayload, - options: PutOptions, - ) -> Result { - self.current_store() - .await - .map_err(Self::map_store_error)? - .put_opts(location, payload, options) - .await - } - - async fn put_multipart_opts( - &self, - location: &Path, - options: PutMultipartOptions, - ) -> Result> { - self.current_store() - .await - .map_err(Self::map_store_error)? - .put_multipart_opts(location, options) - .await - } - - async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { - self.current_store() - .await - .map_err(Self::map_store_error)? - .get_opts(location, options) - .await - } - - async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> Result> { - self.current_store() - .await - .map_err(Self::map_store_error)? - .get_ranges(location, ranges) - .await - } - - fn delete_stream( - &self, - locations: BoxStream<'static, Result>, - ) -> BoxStream<'static, Result> { - let this = self.clone(); - stream::once(async move { - let store = this.current_store().await.map_err(Self::map_store_error)?; - Ok::<_, Error>((store, locations)) - }) - .map_ok(|(store, locations)| store.delete_stream(locations)) - .try_flatten() - .boxed() - } - - fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { - let prefix = prefix.cloned(); - let this = self.clone(); - stream::once(async move { this.current_store().await.map_err(Self::map_store_error) }) - .map_ok(move |store| store.list(prefix.as_ref())) - .try_flatten() - .boxed() - } - - async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { - self.current_store() - .await - .map_err(Self::map_store_error)? - .list_with_delimiter(prefix) - .await - } - - async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { - self.current_store() - .await - .map_err(Self::map_store_error)? - .copy_opts(from, to, options) - .await - } - - async fn rename_opts(&self, from: &Path, to: &Path, options: RenameOptions) -> Result<()> { - self.current_store() - .await - .map_err(Self::map_store_error)? - .rename_opts(from, to, options) - .await - } -} - -#[cfg(feature = "aws")] -#[derive(Debug)] -struct AtomicAwsStoreProvider { - inner: Arc, -} - -#[cfg(feature = "aws")] -impl AtomicAwsStoreProvider { - const CACHE_GENERATION: &'static str = "lancedb-atomic-aws-v1"; - - fn generated_prefix( - &self, - url: &url::Url, - storage_options: Option<&HashMap>, - ) -> lance_core::Result { - self.inner - .calculate_object_store_prefix(url, storage_options) - .map(|prefix| format!("{prefix}${}", Self::CACHE_GENERATION)) - } - - async fn new_store_inner( - &self, - base_path: url::Url, - params: &ObjectStoreParams, - ) -> lance_core::Result { - let storage_options = params.storage_options().cloned().unwrap_or_default(); - let use_opendal = storage_options - .get("use_opendal") - .is_some_and(|value| value == "true"); - - if use_opendal { - let has_dynamic_options = params - .storage_options_accessor - .as_ref() - .is_some_and(|accessor| accessor.has_provider()); - if params.aws_credentials.is_none() - && !has_dynamic_options - && explicit_aws_credential(&storage_options)?.is_none() - { - return self.inner.new_store(base_path, params).await; - } - - let dynamic_store = AtomicProviderStore { - provider: self.inner.clone(), - base_path, - base_params: params.clone(), - base_options: Arc::new(storage_options.clone()), - accessor: params.storage_options_accessor.clone(), - aws_credentials: params.aws_credentials.clone(), - cache: Arc::new(TokioRwLock::new(None)), - }; - let mut store = dynamic_store.initialize_store().await?; - - // Static explicit credentials need no runtime wrapper, so an unknown provider's - // returned store remains pointer-identical. Dynamic authorities rebuild through that - // same provider whenever their normalized credential configuration changes. - if has_dynamic_options || params.aws_credentials.is_some() { - store.inner = Arc::new(dynamic_store); - } - return Ok(store); - } - - if params.aws_credentials.is_some() { - return self.inner.new_store(base_path, params).await; - } - - let Some(accessor) = params.storage_options_accessor.as_ref() else { - return self.inner.new_store(base_path, params).await; - }; - let credential_provider: object_store::aws::AwsCredentialProvider = - if accessor.has_provider() { - // Validate the currently vended family first. A complete dynamic family replaces - // the whole static family, while a provider returning no AWS options must never - // make a partial static family fall through to Lance's environment-merged map. - let current_options = accessor.get_storage_options().await?.0; - let current_credential = explicit_aws_credential(¤t_options)?; - let static_credential = explicit_aws_credential(&storage_options); - if current_credential.is_none() { - static_credential - .as_ref() - .map_err(|error| lance_core::Error::invalid_input(error.to_string()))?; - } - - let s3_options = storage_options - .iter() - .filter_map(|(key, value)| { - AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) - .ok() - .map(|key| (key, value.clone())) - }) - .collect::>(); - let provider_scheme = - StorageOptions::new(storage_options.clone()).aws_provider_scheme()?; - let region = s3_options.get(&AmazonS3ConfigKey::Region).cloned(); - let fallback = if static_credential.is_ok() { - Some( - build_aws_credential( - params.s3_credentials_refresh_offset, - None, - Some(&s3_options), - region, - None, - provider_scheme, - ) - .await? - .0, - ) - } else { - None - }; - Arc::new(AtomicAccessorAwsCredentialProvider { - accessor: accessor.clone(), - fallback, - }) - } else if let Some(credential) = explicit_aws_credential(&storage_options)? { - Arc::new(StaticCredentialProvider::new(credential)) - } else { - return self.inner.new_store(base_path, params).await; - }; - - // This allocation occurs only after the registry cache miss. Cache identity therefore - // remains the semantic identity of the original storage-options accessor. - let mut atomic_params = params.clone(); - atomic_params.aws_credentials = Some(credential_provider); - self.inner.new_store(base_path, &atomic_params).await - } -} - -#[cfg(feature = "aws")] -#[async_trait] -impl ObjectStoreProvider for AtomicAwsStoreProvider { - async fn new_store( - &self, - base_path: url::Url, - params: &ObjectStoreParams, - ) -> lance_core::Result { - let store_prefix = self.generated_prefix(&base_path, params.storage_options())?; - let mut store = self.new_store_inner(base_path, params).await?; - // The registry cache and the returned store must use the same identity. Lance compares - // these values when resolving external blob bases. - store.store_prefix = store_prefix; - Ok(store) - } - - fn extract_path(&self, url: &url::Url) -> lance_core::Result { - self.inner.extract_path(url) - } - - fn calculate_object_store_prefix( - &self, - url: &url::Url, - storage_options: Option<&HashMap>, - ) -> lance_core::Result { - self.generated_prefix(url, storage_options) - } -} - -#[cfg(feature = "aws")] -static ATOMIC_AWS_REGISTRIES: LazyLock>>> = - LazyLock::new(|| Mutex::new(Vec::new())); - -/// Install the credential-safe S3 provider once on a session's shared object-store registry. -#[cfg(feature = "aws")] -fn install_atomic_aws_provider_inner(session: &lance::session::Session) { - let registry = session.store_registry(); - let mut installed = ATOMIC_AWS_REGISTRIES - .lock() - .expect("atomic AWS registry lock poisoned"); - installed.retain(|entry| entry.strong_count() > 0); - if installed - .iter() - .filter_map(Weak::upgrade) - .any(|entry| Arc::ptr_eq(&entry, ®istry)) - { - return; - } - - for scheme in ["s3", "s3+ddb"] { - if let Some(inner) = registry.get_provider(scheme) { - registry.insert(scheme, Arc::new(AtomicAwsStoreProvider { inner })); - } - } - installed.push(Arc::downgrade(®istry)); -} - -#[cfg(feature = "aws")] -pub(crate) fn install_atomic_aws_provider(session: &lance::session::Session) { - install_atomic_aws_provider_inner(session); -} - -#[cfg(not(feature = "aws"))] -pub(crate) fn install_atomic_aws_provider(_session: &lance::session::Session) {} - -/// Select or create a session and protect its registered AWS providers. +/// Select a supplied session or create the default Lance session. pub(crate) fn atomic_aws_session( session: Option>, ) -> Arc { - match session { - Some(session) => { - install_atomic_aws_provider(&session); - session - } - None => { - let session = Arc::new(lance::session::Session::default()); - #[cfg(feature = "aws")] - install_atomic_aws_provider_inner(&session); - session - } - } + session.unwrap_or_else(|| Arc::new(lance::session::Session::default())) } /// Apply storage options to object store parameters. -/// -/// Credential providers are deliberately installed by [`AtomicAwsStoreProvider`] only after a -/// registry cache miss, preserving semantic cache reuse for identical option maps. pub(crate) fn set_storage_options( params: &mut ObjectStoreParams, storage_options: HashMap, @@ -849,34 +251,21 @@ impl WrappingObjectStore for MirroringObjectStoreWrapper { #[cfg(all(test, feature = "aws"))] mod credential_tests { use super::*; - use lance_io::object_store::providers::aws::{AwsStoreProvider, build_aws_credential}; - use std::sync::{ - Mutex, - atomic::{AtomicBool, AtomicUsize, Ordering}, + use lance_io::object_store::{ + ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry, StorageOptions, + providers::aws::{AwsStoreProvider, merge_atomic_aws_environment}, }; - use std::time::Duration; - - #[derive(Debug)] - struct RecordingProvider { - saw_atomic_credentials: Arc, - } - - #[async_trait] - impl ObjectStoreProvider for RecordingProvider { - async fn new_store( - &self, - _base_path: url::Url, - params: &ObjectStoreParams, - ) -> lance_core::Result { - self.saw_atomic_credentials - .store(params.aws_credentials.is_some(), Ordering::SeqCst); - Err(lance_core::Error::invalid_input("recorded test request")) - } - } + use object_store::{ + StaticCredentialProvider, + aws::{AwsCredential, AwsCredentialProvider}, + memory::InMemory, + }; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; #[derive(Debug)] struct RotatingOptionsProvider { fetches: Arc, + custom_ordered: bool, } #[async_trait] @@ -891,11 +280,15 @@ mod credential_tests { "aws_secret_access_key".to_string(), "refreshed-secret".to_string(), ), + ( + "custom_ordered".to_string(), + self.custom_ordered.to_string(), + ), ]))) } fn provider_id(&self) -> String { - "rotating-test-provider".to_string() + format!("rotating-test-provider-{}", self.custom_ordered) } } @@ -918,31 +311,17 @@ mod credential_tests { } } - #[derive(Debug)] - struct CustomPathProvider; - - #[async_trait] - impl ObjectStoreProvider for CustomPathProvider { - async fn new_store( - &self, - _base_path: url::Url, - _params: &ObjectStoreParams, - ) -> lance_core::Result { - Err(lance_core::Error::invalid_input("unused test provider")) - } - - fn extract_path(&self, _url: &url::Url) -> lance_core::Result { - Ok(Path::from("custom/tenant/path")) - } - } - struct CustomStoreProvider { + expected_accessor: Arc, + expected_credentials: AwsCredentialProvider, marker: Arc, + constructions: Arc, + saw_original_inputs: Arc, } impl std::fmt::Debug for CustomStoreProvider { fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { - // Diagnostic output must never grant authority to replace a custom provider's store. + // Diagnostic output must never grant built-in AWS provider capabilities. formatter.write_str("AwsStoreProvider") } } @@ -954,266 +333,75 @@ mod credential_tests { base_path: url::Url, params: &ObjectStoreParams, ) -> lance_core::Result { + let accessor = params.storage_options_accessor.as_ref().ok_or_else(|| { + lance_core::Error::invalid_input("custom provider lost dynamic accessor") + })?; + if !accessor.has_provider() || !Arc::ptr_eq(accessor, &self.expected_accessor) { + return Err(lance_core::Error::invalid_input( + "custom provider lost dynamic accessor", + )); + } + let credentials = params.aws_credentials.as_ref().ok_or_else(|| { + lance_core::Error::invalid_input("custom provider lost AWS credential provider") + })?; + if !Arc::ptr_eq(credentials, &self.expected_credentials) { + return Err(lance_core::Error::invalid_input( + "custom provider lost AWS credential provider", + )); + } + + self.saw_original_inputs.store(true, Ordering::SeqCst); + self.constructions.fetch_add(1, Ordering::SeqCst); + let current_options = accessor.get_storage_options().await?.0; let mut store = AwsStoreProvider.new_store(base_path, params).await?; store.inner = self.marker.clone(); + store.list_is_lexically_ordered = current_options + .get("custom_ordered") + .is_none_or(|value| value == "true"); Ok(store) } } - #[derive(Debug, PartialEq, Eq)] - struct ObservedCredential { - key_id: String, - token: Option, - } - - #[derive(Debug)] - struct ResolvingProvider { - resolved_credential: Arc>>, - } - - #[async_trait] - impl ObjectStoreProvider for ResolvingProvider { - async fn new_store( - &self, - _base_path: url::Url, - params: &ObjectStoreParams, - ) -> lance_core::Result { - let storage_options = params - .storage_options() - .cloned() - .unwrap_or_default() - .into_iter() - .filter_map(|(key, value)| { - AmazonS3ConfigKey::from_str(&key) - .ok() - .map(|key| (key, value)) - }) - .collect::>(); - let (provider, _) = build_aws_credential( - Duration::from_secs(60), - params.aws_credentials.clone(), - Some(&storage_options), - Some("us-east-1".to_string()), - params.storage_options_accessor.clone(), - None, - ) - .await?; - let credential = provider.get_credential().await?; - *self.resolved_credential.lock().unwrap() = Some(ObservedCredential { - key_id: credential.key_id.clone(), - token: credential.token.clone(), - }); - Err(lance_core::Error::invalid_input("recorded test request")) - } - } - - #[tokio::test] - async fn explicit_aws_credentials_do_not_inherit_an_ambient_session_token() { - let storage_options = HashMap::from([ - ("aws_access_key_id".to_string(), "explicit-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "explicit-secret".to_string(), - ), - ]); - let resolved_credential = Arc::new(Mutex::new(None)); - let provider = AtomicAwsStoreProvider { - inner: Arc::new(ResolvingProvider { - resolved_credential: resolved_credential.clone(), - }), - }; - - provider - .new_store( - url::Url::parse("s3://bucket/table").unwrap(), - &object_store_params_from_storage_options(storage_options), - ) - .await - .unwrap_err(); - - assert_eq!( - *resolved_credential.lock().unwrap(), - Some(ObservedCredential { - key_id: "explicit-key".to_string(), - token: None, - }) - ); - } - - #[test] - fn opendal_explicit_credentials_exclude_ambient_token() { - let storage_options = HashMap::from([ - ("aws_access_key_id".to_string(), "explicit-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "explicit-secret".to_string(), - ), - ("use_opendal".to_string(), "true".to_string()), - ]); - let environment = [ - ( - "AWS_SESSION_TOKEN".to_string(), - "lambda-execution-role-token".to_string(), - ), - ("AWS_REGION".to_string(), "us-east-1".to_string()), - ]; - let params = object_store_params_from_storage_options(storage_options.clone()); - - let options = - atomic_opendal_options(&storage_options, &HashMap::new(), None, environment).unwrap(); - - assert!(params.aws_credentials.is_none()); - assert_eq!(options.get("aws_access_key_id").unwrap(), "explicit-key"); - assert_eq!( - options.get("aws_secret_access_key").unwrap(), - "explicit-secret" - ); - assert_eq!(options.get("aws_session_token").unwrap(), ""); - assert_eq!(options.get("aws_region").unwrap(), "us-east-1"); - assert_eq!(options.get("disable_config_load").unwrap(), "true"); - } - - #[test] - fn opendal_preserves_an_explicit_session_token() { - let storage_options = HashMap::from([ - ("aws_access_key_id".to_string(), "explicit-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "explicit-secret".to_string(), - ), - ( - "aws_session_token".to_string(), - "explicit-token".to_string(), - ), - ]); - - let options = atomic_opendal_options( - &storage_options, - &HashMap::new(), - None, - [("AWS_SESSION_TOKEN".to_string(), "ambient-token".to_string())], - ) - .unwrap(); - - assert_eq!(options.get("aws_session_token").unwrap(), "explicit-token"); - } - - #[test] - fn opendal_dynamic_credential_family_replaces_the_entire_static_family() { - let base_options = HashMap::from([ - ("aws_access_key_id".to_string(), "base-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "base-secret".to_string(), - ), - ("aws_session_token".to_string(), "base-token".to_string()), - ]); - let dynamic_options = HashMap::from([ - ("aws_access_key_id".to_string(), "dynamic-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "dynamic-secret".to_string(), - ), - ]); - - let options = - atomic_opendal_options(&base_options, &dynamic_options, None, std::iter::empty()) - .unwrap(); - - assert_eq!(options.get("aws_access_key_id").unwrap(), "dynamic-key"); - assert_eq!( - options.get("aws_secret_access_key").unwrap(), - "dynamic-secret" - ); - assert_eq!(options.get("aws_session_token").unwrap(), ""); - } - - #[test] - fn wholly_ambient_credentials_still_use_the_default_chain() { - let options = atomic_opendal_options( - &HashMap::new(), - &HashMap::new(), - None, - [ - ("AWS_ACCESS_KEY_ID".to_string(), "ambient-key".to_string()), + fn dynamic_opendal_params( + fetches: Arc, + custom_ordered: bool, + ) -> ( + ObjectStoreParams, + Arc, + AwsCredentialProvider, + ) { + let accessor = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + HashMap::from([ + ("aws_access_key_id".to_string(), "expired-key".to_string()), ( - "AWS_SECRET_ACCESS_KEY".to_string(), - "ambient-secret".to_string(), + "aws_secret_access_key".to_string(), + "expired-secret".to_string(), ), - ("AWS_SESSION_TOKEN".to_string(), "ambient-token".to_string()), - ], - ) - .unwrap(); - - assert!(!options.contains_key("aws_access_key_id")); - assert!(!options.contains_key("aws_secret_access_key")); - assert!(!options.contains_key("aws_session_token")); - assert!(!options.contains_key("disable_config_load")); - } - - #[test] - fn partial_explicit_credentials_are_rejected() { - let error = explicit_aws_credential(&HashMap::from([( - "aws_access_key_id".to_string(), - "explicit-key".to_string(), - )])) - .unwrap_err(); - - assert!(error.to_string().contains("require both")); - } - - #[tokio::test] - async fn public_namespace_connection_installs_the_atomic_provider() { - let saw_atomic_credentials = Arc::new(AtomicBool::new(false)); - let registry = Arc::new(ObjectStoreRegistry::default()); - registry.insert( - "s3", - Arc::new(RecordingProvider { - saw_atomic_credentials: saw_atomic_credentials.clone(), + ("expires_at_millis".to_string(), "0".to_string()), + ("use_opendal".to_string(), "true".to_string()), + ("aws_region".to_string(), "us-east-1".to_string()), + ("custom_ordered".to_string(), "true".to_string()), + ]), + Arc::new(RotatingOptionsProvider { + fetches, + custom_ordered, }), - ); - let session = Arc::new(lance::session::Session::new(16, 16, registry.clone())); - let root = tempfile::tempdir().unwrap(); - - crate::connect_namespace( - "dir", - HashMap::from([( - "root".to_string(), - root.path().to_string_lossy().into_owned(), - )]), + )); + let credentials: AwsCredentialProvider = + Arc::new(StaticCredentialProvider::new(AwsCredential { + key_id: "provider-key".to_string(), + secret_key: "provider-secret".to_string(), + token: None, + })); + ( + ObjectStoreParams { + aws_credentials: Some(credentials.clone()), + storage_options_accessor: Some(accessor.clone()), + ..Default::default() + }, + accessor, + credentials, ) - .storage_options([ - ("aws_access_key_id", "explicit-key"), - ("aws_secret_access_key", "explicit-secret"), - ]) - .session(session) - .execute() - .await - .unwrap(); - - // DirectoryNamespaceBuilder constructs fresh params with only a static accessor. - let params = ObjectStoreParams { - storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( - HashMap::from([ - ("aws_access_key_id".to_string(), "explicit-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "explicit-secret".to_string(), - ), - ]), - ))), - ..Default::default() - }; - - let error = registry - .get_provider("s3") - .unwrap() - .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) - .await - .unwrap_err(); - - assert!(error.to_string().contains("recorded test request")); - assert!(saw_atomic_credentials.load(Ordering::SeqCst)); } #[tokio::test] @@ -1234,6 +422,7 @@ mod credential_tests { ]), Arc::new(RotatingOptionsProvider { fetches: fetches.clone(), + custom_ordered: true, }), ), )), @@ -1252,6 +441,73 @@ mod credential_tests { assert_eq!(fetches.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn opendal_preserves_custom_provider_dynamic_accessor() { + let fetches = Arc::new(AtomicUsize::new(0)); + let (params, accessor, credentials) = dynamic_opendal_params(fetches.clone(), false); + let marker: Arc = Arc::new(InMemory::new()); + let saw_original_inputs = Arc::new(AtomicBool::new(false)); + let registry = Arc::new(ObjectStoreRegistry::default()); + registry.insert( + "s3", + Arc::new(CustomStoreProvider { + expected_accessor: accessor, + expected_credentials: credentials, + marker: marker.clone(), + constructions: Arc::new(AtomicUsize::new(0)), + saw_original_inputs: saw_original_inputs.clone(), + }), + ); + let session = atomic_aws_session(Some(Arc::new(lance::session::Session::new( + 16, + 16, + registry.clone(), + )))); + + let store = session + .store_registry() + .get_provider("s3") + .unwrap() + .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) + .await + .unwrap(); + + assert!(saw_original_inputs.load(Ordering::SeqCst)); + assert!(Arc::ptr_eq(&store.inner, &marker)); + assert_eq!(fetches.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn opendal_refresh_preserves_provider_metadata() { + let fetches = Arc::new(AtomicUsize::new(0)); + let (params, accessor, credentials) = dynamic_opendal_params(fetches.clone(), false); + let marker: Arc = Arc::new(InMemory::new()); + let constructions = Arc::new(AtomicUsize::new(0)); + let registry = Arc::new(ObjectStoreRegistry::default()); + registry.insert( + "s3", + Arc::new(CustomStoreProvider { + expected_accessor: accessor, + expected_credentials: credentials, + marker, + constructions: constructions.clone(), + saw_original_inputs: Arc::new(AtomicBool::new(false)), + }), + ); + + let store = registry + .get_provider("s3") + .unwrap() + .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) + .await + .unwrap(); + let _ = store.inner.list(None).next().await; + + assert!(!store.list_is_lexically_ordered); + assert_eq!(constructions.load(Ordering::SeqCst), 1); + assert_eq!(fetches.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn non_aws_dynamic_options_cannot_complete_a_partial_static_family() { let params = ObjectStoreParams { @@ -1260,6 +516,7 @@ mod credential_tests { HashMap::from([ ("aws_access_key_id".to_string(), "explicit-key".to_string()), ("expires_at_millis".to_string(), "0".to_string()), + ("use_opendal".to_string(), "true".to_string()), ]), Arc::new(NonAwsOptionsProvider), ), @@ -1267,68 +524,14 @@ mod credential_tests { ..Default::default() }; - let error = AtomicAwsStoreProvider { - inner: Arc::new(AwsStoreProvider), - } - .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) - .await - .unwrap_err(); + let error = AwsStoreProvider + .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) + .await + .unwrap_err(); assert!(error.to_string().contains("require both")); } - #[tokio::test] - async fn complete_dynamic_credentials_replace_a_partial_static_family() { - let fetches = Arc::new(AtomicUsize::new(0)); - let resolved_credential = Arc::new(Mutex::new(None)); - let params = ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - StorageOptionsAccessor::with_initial_and_provider( - HashMap::from([ - ("aws_access_key_id".to_string(), "stale-key".to_string()), - ("expires_at_millis".to_string(), "0".to_string()), - ]), - Arc::new(RotatingOptionsProvider { - fetches: fetches.clone(), - }), - ), - )), - ..Default::default() - }; - - AtomicAwsStoreProvider { - inner: Arc::new(ResolvingProvider { - resolved_credential: resolved_credential.clone(), - }), - } - .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) - .await - .unwrap_err(); - - assert_eq!(fetches.load(Ordering::SeqCst), 1); - assert_eq!( - *resolved_credential.lock().unwrap(), - Some(ObservedCredential { - key_id: "refreshed-key".to_string(), - token: None, - }) - ); - } - - #[test] - fn wrapper_delegates_custom_path_extraction() { - let provider = AtomicAwsStoreProvider { - inner: Arc::new(CustomPathProvider), - }; - - assert_eq!( - provider - .extract_path(&url::Url::parse("s3://bucket/original/path").unwrap()) - .unwrap(), - Path::from("custom/tenant/path") - ); - } - fn local_s3_options() -> HashMap { HashMap::from([ ("aws_access_key_id".to_string(), "explicit-key".to_string()), @@ -1342,76 +545,36 @@ mod credential_tests { ]) } - #[tokio::test] - async fn installing_the_wrapper_does_not_reuse_a_preexisting_store() { - let registry = Arc::new(ObjectStoreRegistry::default()); - let params = object_store_params_from_storage_options(local_s3_options()); - let url = url::Url::parse("s3://bucket/table").unwrap(); - let before = registry.get_store(url.clone(), ¶ms).await.unwrap(); + #[test] + fn explicit_aws_credentials_do_not_inherit_an_ambient_session_token() { + let mut options = StorageOptions::new(HashMap::from([ + ("aws_access_key_id".to_string(), "explicit-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "explicit-secret".to_string(), + ), + ])); - let session = lance::session::Session::new(16, 16, registry.clone()); - install_atomic_aws_provider(&session); - let after = registry.get_store(url, ¶ms).await.unwrap(); - - assert!( - !Arc::ptr_eq(&before, &after), - "the wrapper cache generation must isolate pre-install stores" + merge_atomic_aws_environment( + &mut options, + [ + ("AWS_SESSION_TOKEN".to_string(), "ambient-token".to_string()), + ("AWS_REGION".to_string(), "us-east-1".to_string()), + ], ); - } - #[tokio::test] - async fn opendal_preserves_custom_provider_store_behavior() { - let marker: Arc = Arc::new(object_store::memory::InMemory::new()); - let registry = Arc::new(ObjectStoreRegistry::default()); - registry.insert( - "s3", - Arc::new(CustomStoreProvider { - marker: marker.clone(), - }), + assert_eq!(options.0.get("aws_access_key_id").unwrap(), "explicit-key"); + assert_eq!( + options.0.get("aws_secret_access_key").unwrap(), + "explicit-secret" ); - let session = lance::session::Session::new(16, 16, registry.clone()); - install_atomic_aws_provider(&session); - let mut options = local_s3_options(); - options.insert("use_opendal".to_string(), "true".to_string()); - - let store = registry - .get_provider("s3") - .unwrap() - .new_store( - url::Url::parse("s3://bucket/table").unwrap(), - &object_store_params_from_storage_options(options), - ) - .await - .unwrap(); - - assert!( - Arc::ptr_eq(&store.inner, &marker), - "the wrapper must not discard custom provider store behavior" - ); - } - - #[tokio::test] - async fn wrapper_store_prefix_matches_registry_identity() { - let registry = Arc::new(ObjectStoreRegistry::default()); - let session = lance::session::Session::new(16, 16, registry.clone()); - install_atomic_aws_provider(&session); - let uri = "s3://bucket/table"; - let url = url::Url::parse(uri).unwrap(); - let params = object_store_params_from_storage_options(local_s3_options()); - - let store = registry.get_store(url, ¶ms).await.unwrap(); - let registry_prefix = registry - .calculate_object_store_prefix(uri, params.storage_options()) - .unwrap(); - - assert_eq!(store.store_prefix, registry_prefix); + assert!(!options.0.contains_key("aws_session_token")); + assert_eq!(options.0.get("aws_region").unwrap(), "us-east-1"); } #[tokio::test] async fn identical_explicit_options_reuse_the_session_store() { let registry = Arc::new(ObjectStoreRegistry::default()); - let session = lance::session::Session::new(16, 16, registry.clone()); - install_atomic_aws_provider(&session); let url = url::Url::parse("s3://bucket/table").unwrap(); let first_params = object_store_params_from_storage_options(local_s3_options()); let second_params = object_store_params_from_storage_options(local_s3_options()); @@ -1428,87 +591,27 @@ mod credential_tests { ); } - #[tokio::test] - async fn dynamic_storage_options_provider_remains_the_credential_authority() { - let fetches = Arc::new(AtomicUsize::new(0)); - let resolved_credential = Arc::new(Mutex::new(None)); - let provider = AtomicAwsStoreProvider { - inner: Arc::new(ResolvingProvider { - resolved_credential: resolved_credential.clone(), - }), - }; - let params = ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - StorageOptionsAccessor::with_initial_and_provider( - HashMap::from([ - ("aws_access_key_id".to_string(), "expired-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "expired-secret".to_string(), - ), - ("expires_at_millis".to_string(), "0".to_string()), - ]), - Arc::new(RotatingOptionsProvider { - fetches: fetches.clone(), - }), - ), - )), - ..Default::default() - }; - - provider - .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) - .await - .unwrap_err(); - - assert_eq!(fetches.load(Ordering::SeqCst), 1); - assert_eq!( - *resolved_credential.lock().unwrap(), - Some(ObservedCredential { - key_id: "refreshed-key".to_string(), - token: None, - }) - ); + #[test] + fn aws_credential_options_are_one_merge_family() { + assert!(is_aws_credential_option("aws_access_key_id")); + assert!(is_aws_credential_option("AWS_SECRET_ACCESS_KEY")); + assert!(is_aws_credential_option("aws_session_token")); + assert!(!is_aws_credential_option("aws_region")); } - #[tokio::test] - async fn caller_supplied_aws_provider_remains_the_credential_authority() { - let resolved_credential = Arc::new(Mutex::new(None)); - let provider = AtomicAwsStoreProvider { - inner: Arc::new(ResolvingProvider { - resolved_credential: resolved_credential.clone(), - }), - }; - let params = ObjectStoreParams { - aws_credentials: Some(Arc::new(StaticCredentialProvider::new(AwsCredential { - key_id: "provider-key".to_string(), - secret_key: "provider-secret".to_string(), - token: None, - }))), - storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( - HashMap::from([ - ("aws_access_key_id".to_string(), "option-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "option-secret".to_string(), - ), - ]), - ))), - ..Default::default() - }; + #[test] + fn dynamic_params_keep_the_original_opaque_authorities() { + let fetches = Arc::new(AtomicUsize::new(0)); + let (params, accessor, credentials) = dynamic_opendal_params(fetches, true); - provider - .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) - .await - .unwrap_err(); - - assert_eq!( - *resolved_credential.lock().unwrap(), - Some(ObservedCredential { - key_id: "provider-key".to_string(), - token: None, - }) - ); + assert!(Arc::ptr_eq( + params.storage_options_accessor.as_ref().unwrap(), + &accessor + )); + assert!(Arc::ptr_eq( + params.aws_credentials.as_ref().unwrap(), + &credentials + )); } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index cb7b05d87..48282e965 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3580,6 +3580,7 @@ mod tests { #[cfg(feature = "aws")] use lance_io::object_store::{ ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry, + StorageOptionsAccessor, }; use tempfile::tempdir; @@ -3593,7 +3594,8 @@ mod tests { #[cfg(feature = "aws")] #[derive(Debug)] struct RecordingS3Provider { - saw_atomic_credentials: Arc, + expected_accessor: Arc, + saw_original_params: Arc, } #[cfg(feature = "aws")] @@ -3604,37 +3606,49 @@ mod tests { _base_path: url::Url, params: &ObjectStoreParams, ) -> lance_core::Result { - self.saw_atomic_credentials - .store(params.aws_credentials.is_some(), Ordering::SeqCst); + self.saw_original_params.store( + params.aws_credentials.is_none() + && params + .storage_options_accessor + .as_ref() + .is_some_and(|accessor| Arc::ptr_eq(accessor, &self.expected_accessor)), + Ordering::SeqCst, + ); Err(lance_core::Error::invalid_input("recorded test request")) } } #[cfg(feature = "aws")] - fn recording_s3_session() -> (Arc, Arc) { - let saw_atomic_credentials = Arc::new(AtomicBool::new(false)); + fn recording_s3_session( + expected_accessor: Arc, + ) -> (Arc, Arc) { + let saw_original_params = Arc::new(AtomicBool::new(false)); let registry = Arc::new(ObjectStoreRegistry::default()); registry.insert( "s3", Arc::new(RecordingS3Provider { - saw_atomic_credentials: saw_atomic_credentials.clone(), + expected_accessor, + saw_original_params: saw_original_params.clone(), }), ); ( Arc::new(lance::session::Session::new(16, 16, registry)), - saw_atomic_credentials, + saw_original_params, ) } #[cfg(feature = "aws")] - fn explicit_s3_store_params() -> ObjectStoreParams { - crate::io::object_store::object_store_params_from_storage_options(HashMap::from([ - ("aws_access_key_id".to_string(), "explicit-key".to_string()), - ( - "aws_secret_access_key".to_string(), - "explicit-secret".to_string(), - ), - ])) + fn explicit_s3_store_params() -> (ObjectStoreParams, Arc) { + let params = + crate::io::object_store::object_store_params_from_storage_options(HashMap::from([ + ("aws_access_key_id".to_string(), "explicit-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "explicit-secret".to_string(), + ), + ])); + let accessor = params.storage_options_accessor.as_ref().unwrap().clone(); + (params, accessor) } #[test] @@ -3700,11 +3714,12 @@ mod tests { #[cfg(feature = "aws")] #[tokio::test] - async fn direct_native_open_installs_the_atomic_provider() { - let (session, saw_atomic_credentials) = recording_s3_session(); + async fn direct_native_open_preserves_custom_provider_params() { + let (store_options, accessor) = explicit_s3_store_params(); + let (session, saw_original_params) = recording_s3_session(accessor); let params = ReadParams { session: Some(session), - store_options: Some(explicit_s3_store_params()), + store_options: Some(store_options), ..Default::default() }; @@ -3724,18 +3739,19 @@ mod tests { assert!(error.to_string().contains("recorded test request")); assert!( - saw_atomic_credentials.load(Ordering::SeqCst), - "the public direct open path must install the wrapper on its session" + saw_original_params.load(Ordering::SeqCst), + "the public direct open path must preserve custom provider parameters" ); } #[cfg(feature = "aws")] #[tokio::test] - async fn direct_native_create_installs_the_atomic_provider() { - let (session, saw_atomic_credentials) = recording_s3_session(); + async fn direct_native_create_preserves_custom_provider_params() { + let (store_params, accessor) = explicit_s3_store_params(); + let (session, saw_original_params) = recording_s3_session(accessor); let params = WriteParams { session: Some(session), - store_params: Some(explicit_s3_store_params()), + store_params: Some(store_params), ..Default::default() }; let batch = make_test_batches(); @@ -3757,8 +3773,8 @@ mod tests { assert!(error.to_string().contains("recorded test request")); assert!( - saw_atomic_credentials.load(Ordering::SeqCst), - "the public direct create path must install the wrapper on its session" + saw_original_params.load(Ordering::SeqCst), + "the public direct create path must preserve custom provider parameters" ); } diff --git a/vendor/lance-io/Cargo.toml b/vendor/lance-io/Cargo.toml new file mode 100644 index 000000000..da800133d --- /dev/null +++ b/vendor/lance-io/Cargo.toml @@ -0,0 +1,74 @@ +[package] +name = "lance-io" +version = "11.0.0-beta.2" +edition = "2024" +authors = ["Lance Devs "] +license = "Apache-2.0" +repository = "https://github.com/lance-format/lance" +readme = "README.md" +description = "I/O utilities for Lance" +keywords = ["data-format", "data-science", "machine-learning", "apache-arrow"] +categories = ["database-implementations", "data-structures"] +rust-version = "1.91.0" +autobenches = false +autotests = false + +[dependencies] +object_store = "0.13.2" +opendal = { version = "0.58.1", optional = true } +object_store_opendal = { version = "0.58", optional = true } +lance-arrow = { version = "=11.0.0-beta.2", tag = "v11.0.0-beta.2", git = "https://github.com/lance-format/lance.git" } +lance-core = { version = "=11.0.0-beta.2", tag = "v11.0.0-beta.2", git = "https://github.com/lance-format/lance.git" } +lance-namespace = { version = "=11.0.0-beta.2", tag = "v11.0.0-beta.2", git = "https://github.com/lance-format/lance.git" } +arrow = { version = "58.0.0", features = ["ffi"] } +arrow-array = "58.0.0" +arrow-schema = "58.0.0" +async-trait = "0.1" +aws-config = { version = "1.2.0", optional = true } +aws-credential-types = { version = "1.2.0", optional = true } +byteorder = "1.5" +bytes = "1.11.1" +chrono = { version = "0.4.41", default-features = false, features = ["std", "now", "serde"] } +futures = "0.3" +http = "1.1.0" +log = "0.4" +metrics = { version = "0.24", optional = true } +moka = { version = "0.12", features = ["future"] } +pin-project = "1.0" +prost = "0.14.1" +serde = { version = "1", features = ["derive"] } +tokio = { version = "1.23", features = ["rt-multi-thread", "macros", "fs", "sync"] } +tracing = "0.1" +url = "2.5.7" +path_abs = "0.5" +rand = "0.9.1" +tempfile = "3" + +[target.'cfg(target_os = "linux")'.dependencies] +io-uring = "0.7" + +[dev-dependencies] +lance-testing = { version = "=11.0.0-beta.2", tag = "v11.0.0-beta.2", git = "https://github.com/lance-format/lance.git" } +test-log = "0.2.15" +mockall = "0.14.0" +rstest = "0.26.1" +mock_instant = "0.6.0" +tokio = { version = "1.23", features = ["test-util"] } +tracing-mock = "=0.1.0-beta.3" +metrics-util = "0.19" + +[features] +default = ["aws", "azure", "gcp"] +metrics = ["dep:metrics"] +gcs-test = [] +goosefs-test = [] +gcp = ["object_store/gcp", "dep:opendal", "opendal/services-gcs", "dep:object_store_opendal"] +aws = ["object_store/aws", "dep:aws-config", "dep:aws-credential-types", "dep:opendal", "opendal/services-s3", "dep:object_store_opendal"] +azure = ["object_store/azure", "dep:opendal", "opendal/services-azblob", "opendal/services-azdls", "dep:object_store_opendal"] +oss = ["dep:opendal", "opendal/services-oss", "dep:object_store_opendal"] +goosefs = ["dep:opendal", "opendal/services-goosefs", "dep:object_store_opendal"] +tencent = ["dep:opendal", "opendal/services-cos", "dep:object_store_opendal"] +huggingface = ["dep:opendal", "opendal/services-huggingface", "dep:object_store_opendal"] +tos = ["dep:opendal", "opendal/services-tos", "dep:object_store_opendal"] +tos-test = ["tos"] +test-util = [] diff --git a/vendor/lance-io/LANCEDB_PATCH.md b/vendor/lance-io/LANCEDB_PATCH.md new file mode 100644 index 000000000..87aaafa3e --- /dev/null +++ b/vendor/lance-io/LANCEDB_PATCH.md @@ -0,0 +1,13 @@ +# LanceDB patch provenance + +This directory vendors `lance-io` 11.0.0-beta.2 from Lance commit +`35da5d920159b49d1b53032652f7615ab699c160`. + +`Cargo.toml` uses the equivalent standalone dependency metadata from the published crate. Upstream +benchmark and integration-test targets are omitted because this copy is compiled only as a patched +dependency; the library sources are otherwise retained. + +The local patch makes AWS credential-family merging atomic before backend selection and teaches +the built-in OpenDAL S3 path to refresh credential-only storage options. Keeping the change inside +`AwsStoreProvider` leaves arbitrary registry providers and their complete `ObjectStore` results +untouched. Remove this patch when the same behavior is available in the pinned Lance release. diff --git a/vendor/lance-io/README.md b/vendor/lance-io/README.md new file mode 100644 index 000000000..bcefbf38a --- /dev/null +++ b/vendor/lance-io/README.md @@ -0,0 +1,9 @@ +# lance-io + +`lance-io` is an internal sub-crate, containing various utilities for +reading and writing data. It includes reader/writer traits that +define what Lance expects from a filesystem, encoders and decoders to +convert to/from Arrow data and various layouts, and misc. utilities such +as routines for reading protobuf data from files. + +**Important Note**: This crate is **not intended for external usage**. diff --git a/vendor/lance-io/src/ffi.rs b/vendor/lance-io/src/ffi.rs new file mode 100644 index 000000000..950977e3f --- /dev/null +++ b/vendor/lance-io/src/ffi.rs @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::RecordBatch; +use arrow_schema::{ArrowError, SchemaRef}; +use futures::StreamExt; +use lance_core::Result; + +use crate::stream::RecordBatchStream; + +#[pin_project::pin_project] +struct RecordBatchIteratorAdaptor { + schema: SchemaRef, + + #[pin] + stream: S, + + handle: tokio::runtime::Handle, +} + +impl RecordBatchIteratorAdaptor { + fn new(stream: S, schema: SchemaRef, handle: tokio::runtime::Handle) -> Self { + Self { + schema, + stream, + handle, + } + } +} + +impl arrow::record_batch::RecordBatchReader + for RecordBatchIteratorAdaptor +{ + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Iterator for RecordBatchIteratorAdaptor { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + self.handle + .block_on(async { self.stream.next().await }) + .map(|r| r.map_err(|e| ArrowError::ExternalError(Box::new(e)))) + } +} + +/// Wrap a [`RecordBatchStream`] into an [FFI_ArrowArrayStream]. +pub fn to_ffi_arrow_array_stream( + stream: impl RecordBatchStream + std::marker::Unpin + 'static, + handle: tokio::runtime::Handle, +) -> Result { + let schema = stream.schema(); + let arrow_stream = RecordBatchIteratorAdaptor::new(stream, schema, handle); + let reader = FFI_ArrowArrayStream::new(Box::new(arrow_stream)); + + Ok(reader) +} diff --git a/vendor/lance-io/src/lib.rs b/vendor/lance-io/src/lib.rs new file mode 100644 index 000000000..b6bdc404e --- /dev/null +++ b/vendor/lance-io/src/lib.rs @@ -0,0 +1,387 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors +#![recursion_limit = "512"] +use std::{ + ops::{Range, RangeFrom, RangeFull, RangeTo}, + sync::Arc, +}; + +use arrow::datatypes::UInt32Type; +use arrow_array::{PrimitiveArray, UInt32Array}; + +use lance_core::{Error, Result}; + +pub mod ffi; +pub mod local; +pub mod object_reader; +pub mod object_store; +pub mod object_writer; +pub mod scheduler; +pub mod spill; +pub mod stream; +#[cfg(test)] +pub mod testing; +pub mod traits; +#[cfg(target_os = "linux")] +pub mod uring; +pub mod utils; + +pub use scheduler::{bytes_read_counter, iops_counter}; + +/// Defines a selection of rows to read from a file/batch +#[derive(Debug, Clone, PartialEq, Default)] +pub enum ReadBatchParams { + /// Select a contiguous range of rows + Range(Range), + /// Select multiple contiguous ranges of rows + Ranges(Arc<[Range]>), + /// Select all rows (this is the default) + #[default] + RangeFull, + /// Select all rows up to a given index + RangeTo(RangeTo), + /// Select all rows starting at a given index + RangeFrom(RangeFrom), + /// Select scattered non-contiguous rows + Indices(UInt32Array), +} + +impl std::fmt::Display for ReadBatchParams { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + Self::Range(r) => write!(f, "Range({}..{})", r.start, r.end), + Self::Ranges(ranges) => { + let mut ranges_str = ranges.iter().fold(String::new(), |mut acc, r| { + acc.push_str(&format!("{}..{}", r.start, r.end)); + acc.push(','); + acc + }); + // Remove the trailing comma + if !ranges_str.is_empty() { + ranges_str.pop(); + } + write!(f, "Ranges({})", ranges_str) + } + Self::RangeFull => write!(f, "RangeFull"), + Self::RangeTo(r) => write!(f, "RangeTo({})", r.end), + Self::RangeFrom(r) => write!(f, "RangeFrom({})", r.start), + Self::Indices(indices) => { + let mut indices_str = indices.values().iter().fold(String::new(), |mut acc, v| { + acc.push_str(&v.to_string()); + acc.push(','); + acc + }); + if !indices_str.is_empty() { + indices_str.pop(); + } + write!(f, "Indices({})", indices_str) + } + } + } +} + +impl From<&[u32]> for ReadBatchParams { + fn from(value: &[u32]) -> Self { + Self::Indices(UInt32Array::from_iter_values(value.iter().copied())) + } +} + +impl From for ReadBatchParams { + fn from(value: UInt32Array) -> Self { + Self::Indices(value) + } +} + +impl From for ReadBatchParams { + fn from(_: RangeFull) -> Self { + Self::RangeFull + } +} + +impl From> for ReadBatchParams { + fn from(r: Range) -> Self { + Self::Range(r) + } +} + +impl From> for ReadBatchParams { + fn from(r: RangeTo) -> Self { + Self::RangeTo(r) + } +} + +impl From> for ReadBatchParams { + fn from(r: RangeFrom) -> Self { + Self::RangeFrom(r) + } +} + +impl From<&Self> for ReadBatchParams { + fn from(params: &Self) -> Self { + params.clone() + } +} + +impl ReadBatchParams { + /// Validate that the selection is valid given the length of the batch + pub fn valid_given_len(&self, len: usize) -> bool { + match self { + Self::Indices(indices) => indices.iter().all(|i| i.unwrap_or(0) < len as u32), + Self::Range(r) => r.start < len && r.end <= len, + Self::Ranges(ranges) => ranges.iter().all(|r| r.end <= len as u64), + Self::RangeFull => true, + Self::RangeTo(r) => r.end <= len, + Self::RangeFrom(r) => r.start < len, + } + } + + /// Slice the selection + /// + /// For example, given ReadBatchParams::RangeFull and slice(10, 20), the output will be + /// ReadBatchParams::Range(10..20) + /// + /// Given ReadBatchParams::Range(10..20) and slice(5, 3), the output will be + /// ReadBatchParams::Range(15..18) + /// + /// Given ReadBatchParams::RangeTo(20) and slice(10, 5), the output will be + /// ReadBatchParams::Range(10..15) + /// + /// Given ReadBatchParams::RangeFrom(20) and slice(10, 5), the output will be + /// ReadBatchParams::Range(30..35) + /// + /// Given ReadBatchParams::Indices([1, 3, 5, 7, 9]) and slice(1, 3), the output will be + /// ReadBatchParams::Indices([3, 5, 7]) + /// + /// You cannot slice beyond the bounds of the selection and an attempt to do so will + /// return an error. + pub fn slice(&self, start: usize, length: usize) -> Result { + let out_of_bounds = |size: usize| { + Err(Error::invalid_input_source( + format!( + "Cannot slice from {} with length {} given a selection of size {}", + start, length, size + ) + .into(), + )) + }; + + match self { + Self::Indices(indices) => { + if start + length > indices.len() { + return out_of_bounds(indices.len()); + } + Ok(Self::Indices(indices.slice(start, length))) + } + Self::Range(r) => { + if (r.start + start + length) > r.end { + return out_of_bounds(r.end - r.start); + } + Ok(Self::Range((r.start + start)..(r.start + start + length))) + } + Self::Ranges(ranges) => { + let mut new_ranges = Vec::with_capacity(ranges.len()); + let mut to_skip = start as u64; + let mut to_take = length as u64; + let mut total_num_rows = 0; + for r in ranges.as_ref() { + let num_rows = r.end - r.start; + total_num_rows += num_rows; + if to_skip > num_rows { + to_skip -= num_rows; + continue; + } + let new_start = r.start + to_skip; + let to_take_this_range = (num_rows - to_skip).min(to_take); + new_ranges.push(new_start..(new_start + to_take_this_range)); + to_skip = 0; + to_take -= to_take_this_range; + if to_take == 0 { + break; + } + } + if to_take > 0 { + out_of_bounds(total_num_rows as usize) + } else { + Ok(Self::Ranges(new_ranges.into())) + } + } + Self::RangeFull => Ok(Self::Range(start..(start + length))), + Self::RangeTo(range) => { + if start + length > range.end { + return out_of_bounds(range.end); + } + Ok(Self::Range(start..(start + length))) + } + Self::RangeFrom(r) => { + // No way to validate out_of_bounds, assume caller will do so + Ok(Self::Range((r.start + start)..(r.start + start + length))) + } + } + } + + /// Convert a read range into a vector of row offsets + /// + /// RangeFull and RangeFrom are unbounded and cannot be converted into row offsets + /// and any attempt to do so will return an error. Call slice first + pub fn to_offsets(&self) -> Result> { + match self { + Self::Indices(indices) => Ok(indices.clone()), + Self::Range(r) => Ok(UInt32Array::from(Vec::from_iter( + r.start as u32..r.end as u32, + ))), + Self::Ranges(ranges) => { + let num_rows = ranges + .iter() + .map(|r| (r.end - r.start) as usize) + .sum::(); + let mut offsets = Vec::with_capacity(num_rows); + for r in ranges.as_ref() { + offsets.extend(r.start as u32..r.end as u32); + } + Ok(UInt32Array::from(offsets)) + } + Self::RangeFull => Err(Error::invalid_input("cannot materialize RangeFull")), + Self::RangeTo(r) => Ok(UInt32Array::from(Vec::from_iter(0..r.end as u32))), + Self::RangeFrom(_) => Err(Error::invalid_input("cannot materialize RangeFrom")), + } + } + + pub fn iter_offset_ranges<'a>( + &'a self, + ) -> Result> + Send + 'a>> { + match self { + Self::Indices(indices) => Ok(Box::new(indices.values().iter().map(|i| *i..(*i + 1)))), + Self::Range(r) => Ok(Box::new(std::iter::once(r.start as u32..r.end as u32))), + Self::Ranges(ranges) => Ok(Box::new( + ranges.iter().map(|r| r.start as u32..r.end as u32), + )), + Self::RangeFull => Err(Error::invalid_input("cannot materialize RangeFull")), + Self::RangeTo(r) => Ok(Box::new(std::iter::once(0..r.end as u32))), + Self::RangeFrom(_) => Err(Error::invalid_input("cannot materialize RangeFrom")), + } + } + + /// Convert a read range into a vector of row ranges + pub fn to_ranges(&self) -> Result>> { + match self { + Self::Indices(indices) => Ok(indices + .values() + .iter() + .map(|i| *i as u64..(*i + 1) as u64) + .collect()), + Self::Range(r) => Ok(vec![r.start as u64..r.end as u64]), + Self::Ranges(ranges) => Ok(ranges.to_vec()), + Self::RangeFull => Err(Error::invalid_input("cannot materialize RangeFull")), + Self::RangeTo(r) => Ok(vec![0..r.end as u64]), + Self::RangeFrom(_) => Err(Error::invalid_input("cannot materialize RangeFrom")), + } + } + + /// Same thing as to_offsets but the caller knows the total number of rows in the file + /// + /// This makes it possible to materialize RangeFull / RangeFrom + pub fn to_offsets_total(&self, total: u32) -> PrimitiveArray { + match self { + Self::Indices(indices) => indices.clone(), + Self::Range(r) => UInt32Array::from_iter_values(r.start as u32..r.end as u32), + Self::Ranges(ranges) => { + let num_rows = ranges + .iter() + .map(|r| (r.end - r.start) as usize) + .sum::(); + let mut offsets = Vec::with_capacity(num_rows); + for r in ranges.as_ref() { + offsets.extend(r.start as u32..r.end as u32); + } + UInt32Array::from(offsets) + } + Self::RangeFull => UInt32Array::from_iter_values(0_u32..total), + Self::RangeTo(r) => UInt32Array::from_iter_values(0..r.end as u32), + Self::RangeFrom(r) => UInt32Array::from_iter_values(r.start as u32..total), + } + } +} + +#[cfg(test)] +mod test { + use std::ops::{RangeFrom, RangeTo}; + + use arrow_array::UInt32Array; + + use crate::ReadBatchParams; + + #[test] + fn test_params_slice() { + let params = ReadBatchParams::Ranges(vec![0..15, 20..40].into()); + let sliced = params.slice(10, 10).unwrap(); + assert_eq!(sliced, ReadBatchParams::Ranges(vec![10..15, 20..25].into())); + } + + #[test] + fn test_params_to_offsets() { + let check = |params: ReadBatchParams, base_offset, length, expected: Vec| { + let offsets = params + .slice(base_offset, length) + .unwrap() + .to_offsets() + .unwrap(); + let expected = UInt32Array::from(expected); + assert_eq!(offsets, expected); + }; + + check(ReadBatchParams::RangeFull, 0, 100, (0..100).collect()); + check(ReadBatchParams::RangeFull, 50, 100, (50..150).collect()); + check( + ReadBatchParams::RangeFrom(RangeFrom { start: 500 }), + 0, + 100, + (500..600).collect(), + ); + check( + ReadBatchParams::RangeFrom(RangeFrom { start: 500 }), + 100, + 100, + (600..700).collect(), + ); + check( + ReadBatchParams::RangeTo(RangeTo { end: 800 }), + 0, + 100, + (0..100).collect(), + ); + check( + ReadBatchParams::RangeTo(RangeTo { end: 800 }), + 200, + 100, + (200..300).collect(), + ); + check( + ReadBatchParams::Indices(UInt32Array::from(vec![1, 3, 5, 7, 9])), + 0, + 2, + vec![1, 3], + ); + check( + ReadBatchParams::Indices(UInt32Array::from(vec![1, 3, 5, 7, 9])), + 2, + 2, + vec![5, 7], + ); + + let check_error = |params: ReadBatchParams, base_offset, length| { + assert!(params.slice(base_offset, length).is_err()); + }; + + check_error(ReadBatchParams::Indices(UInt32Array::from(vec![1])), 0, 2); + check_error(ReadBatchParams::Indices(UInt32Array::from(vec![1])), 1, 1); + check_error(ReadBatchParams::Range(0..10), 5, 6); + check_error(ReadBatchParams::RangeTo(RangeTo { end: 10 }), 5, 6); + + assert!(ReadBatchParams::RangeFull.to_offsets().is_err()); + assert!( + ReadBatchParams::RangeFrom(RangeFrom { start: 10 }) + .to_offsets() + .is_err() + ); + } +} diff --git a/vendor/lance-io/src/local.rs b/vendor/lance-io/src/local.rs new file mode 100644 index 000000000..31b66e4bc --- /dev/null +++ b/vendor/lance-io/src/local.rs @@ -0,0 +1,331 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Optimized local I/Os + +use std::fs::File; +use std::io::{ErrorKind, Read, SeekFrom}; +use std::ops::Range; +use std::sync::Arc; + +// TODO: Clean up windows/unix stuff +#[cfg(unix)] +use std::os::unix::fs::FileExt; +#[cfg(windows)] +use std::os::windows::fs::FileExt; + +use async_trait::async_trait; +use bytes::{Bytes, BytesMut}; +use futures::future::BoxFuture; +use lance_core::deepsize::DeepSizeOf; +use lance_core::{Error, Result}; +use object_store::path::Path; +use tokio::io::AsyncSeekExt; +use tokio::sync::OnceCell; +use tracing::instrument; + +use crate::object_reader::stream_local_range; +use crate::object_store::DEFAULT_LOCAL_IO_PARALLELISM; +use crate::object_writer::WriteResult; +use crate::traits::{ByteStream, Reader, Writer}; +use crate::utils::tracking_store::IOTracker; + +/// Convert an [`object_store::path::Path`] to a [`std::path::Path`]. +pub fn to_local_path(path: &Path) -> String { + if cfg!(windows) { + path.to_string() + } else { + format!("/{path}") + } +} + +/// Recursively remove a directory, specified by [`object_store::path::Path`]. +pub fn remove_dir_all(path: &Path) -> Result<()> { + let local_path = to_local_path(path); + std::fs::remove_dir_all(local_path).map_err(|err| match err.kind() { + ErrorKind::NotFound => Error::not_found(path.to_string()), + _ => Error::from(err), + })?; + Ok(()) +} + +/// Copy a file from one location to another, supporting cross-filesystem copies. +/// +/// Unlike hard links, this function works across filesystem boundaries. +pub fn copy_file(from: &Path, to: &Path) -> Result<()> { + let from_path = to_local_path(from); + let to_path = to_local_path(to); + + // Ensure the parent directory exists + if let Some(parent) = std::path::Path::new(&to_path).parent() { + std::fs::create_dir_all(parent).map_err(Error::from)?; + } + + std::fs::copy(&from_path, &to_path).map_err(|err| match err.kind() { + ErrorKind::NotFound => Error::not_found(from.to_string()), + _ => Error::from(err), + })?; + Ok(()) +} + +/// Await a filesystem operation running on a blocking thread, flattening the +/// join and IO errors into a single `object_store` error. +/// +/// Deliberately not written as `handle.await?` at the call sites: a `JoinError` +/// means the operation panicked, and short-circuiting on it would skip the +/// caller's metrics recording for exactly the failure worth counting. +pub(crate) async fn join_local_io( + handle: tokio::task::JoinHandle>, +) -> object_store::Result { + match handle.await { + Ok(result) => result.map_err(|err| object_store::Error::Generic { + store: "LocalFileSystem", + source: err.into(), + }), + Err(err) => Err(err.into()), + } +} + +/// Object reader for local file system. +#[derive(Debug)] +pub struct LocalObjectReader { + /// File handler. + file: Arc, + + /// Fie path. + path: Path, + + /// Known size of the file. This is either passed in on construction or + /// cached on the first metadata call. + size: OnceCell, + + /// Block size, in bytes. + block_size: usize, + + /// IO tracker for monitoring read operations. + io_tracker: Arc, +} + +impl DeepSizeOf for LocalObjectReader { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // Skipping `file` as it should just be a file handle + self.path.as_ref().deep_size_of_children(context) + } +} + +impl LocalObjectReader { + pub async fn open_local_path( + path: impl AsRef, + block_size: usize, + known_size: Option, + ) -> Result> { + let path = path.as_ref().to_owned(); + let object_store_path = Path::from_filesystem_path(&path)?; + Self::open(&object_store_path, block_size, known_size).await + } + + /// Open a local object reader, with default prefetch size. + /// + /// For backward compatibility with existing code that doesn't need tracking. + #[instrument(level = "debug")] + pub async fn open( + path: &Path, + block_size: usize, + known_size: Option, + ) -> Result> { + Self::open_with_tracker(path, block_size, known_size, Default::default()).await + } + + /// Open a local object reader with optional IO tracking. + #[instrument(level = "debug")] + pub(crate) async fn open_with_tracker( + path: &Path, + block_size: usize, + known_size: Option, + io_tracker: Arc, + ) -> Result> { + let path = path.clone(); + let local_path = to_local_path(&path); + tokio::task::spawn_blocking(move || { + let file = File::open(&local_path).map_err(|e| match e.kind() { + ErrorKind::NotFound => Error::not_found(path.to_string()), + _ => e.into(), + })?; + let size = OnceCell::new_with(known_size); + Ok(Box::new(Self { + file: Arc::new(file), + block_size, + size, + path, + io_tracker, + }) as Box) + }) + .await? + } +} + +impl Reader for LocalObjectReader { + fn path(&self) -> &Path { + &self.path + } + + fn block_size(&self) -> usize { + self.block_size + } + + fn io_parallelism(&self) -> usize { + DEFAULT_LOCAL_IO_PARALLELISM + } + + /// Returns the file size. + fn size(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { + let file = self.file.clone(); + self.size + .get_or_try_init(|| async move { + // The metadata lookup is this reader's equivalent of the HEAD + // request a cloud reader makes to learn the object size. + let metrics = self.io_tracker.begin_io("head"); + let result = + join_local_io(tokio::task::spawn_blocking(move || file.metadata())).await; + metrics.record(&result, 0); + Ok(result?.len() as usize) + }) + .await + .cloned() + }) + } + + /// Reads a range of data. + #[instrument(level = "debug", skip(self))] + fn get_range(&self, range: Range) -> BoxFuture<'static, object_store::Result> { + let file = self.file.clone(); + let io_tracker = self.io_tracker.clone(); + let path = self.path.clone(); + let num_bytes = range.len() as u64; + let range_u64 = (range.start as u64)..(range.end as u64); + + Box::pin(async move { + let metrics = io_tracker.begin_io("get"); + let result = join_local_io(tokio::task::spawn_blocking(move || { + let mut buf = BytesMut::with_capacity(range.len()); + // Safety: `buf` is set with appropriate capacity above. It is + // written to below and we check all data is initialized at that point. + unsafe { buf.set_len(range.len()) }; + #[cfg(unix)] + file.read_exact_at(buf.as_mut(), range.start as u64)?; + #[cfg(windows)] + read_exact_at(file, buf.as_mut(), range.start as u64)?; + + Ok(buf.freeze()) + })) + .await; + + metrics.record(&result, num_bytes); + if result.is_ok() { + io_tracker.record_read("get_range", path, num_bytes, Some(range_u64)); + } + + result + }) + } + + /// Reads the entire file. + #[instrument(level = "debug", skip(self))] + fn get_all(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { + let mut file = self.file.clone(); + let io_tracker = self.io_tracker.clone(); + let path = self.path.clone(); + + let metrics = io_tracker.begin_io("get"); + let result = join_local_io(tokio::task::spawn_blocking(move || { + let mut buf = Vec::new(); + file.read_to_end(buf.as_mut())?; + Ok(Bytes::from(buf)) + })) + .await; + + let num_bytes = result.as_ref().map_or(0, |bytes| bytes.len() as u64); + metrics.record(&result, num_bytes); + if let Ok(bytes) = &result { + io_tracker.record_read("get_all", path, bytes.len() as u64, None); + } + + result + }) + } + + fn get_stream(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { + let size = self.size().await?; + Ok(stream_local_range( + self.file.clone(), + self.path.clone(), + self.io_tracker.clone(), + 0..size, + self.block_size.max(8 * 1024), + )) + }) + } + + fn get_range_stream( + &self, + range: Range, + ) -> BoxFuture<'_, object_store::Result> { + let file = self.file.clone(); + let path = self.path.clone(); + let io_tracker = self.io_tracker.clone(); + let chunk_size = self.block_size.max(8 * 1024); + Box::pin(async move { + Ok(stream_local_range( + file, path, io_tracker, range, chunk_size, + )) + }) + } +} + +#[cfg(windows)] +pub(crate) fn read_exact_at( + file: Arc, + mut buf: &mut [u8], + mut offset: u64, +) -> std::io::Result<()> { + let expected_len = buf.len(); + while !buf.is_empty() { + match file.seek_read(buf, offset) { + Ok(0) => break, + Ok(n) => { + let tmp = buf; + buf = &mut tmp[n..]; + offset += n as u64; + } + Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {} + Err(e) => return Err(e), + } + } + if !buf.is_empty() { + Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + format!( + "failed to fill whole buffer. Expected {} bytes, got {}", + expected_len, offset + ), + )) + } else { + Ok(()) + } +} + +#[async_trait] +impl Writer for tokio::fs::File { + async fn tell(&mut self) -> Result { + Ok(self.seek(SeekFrom::Current(0)).await? as usize) + } + + async fn shutdown(&mut self) -> Result { + let size = self.seek(SeekFrom::Current(0)).await? as usize; + tokio::io::AsyncWriteExt::shutdown(self).await?; + Ok(WriteResult { size, e_tag: None }) + } +} diff --git a/vendor/lance-io/src/object_reader.rs b/vendor/lance-io/src/object_reader.rs new file mode 100644 index 000000000..000686aaf --- /dev/null +++ b/vendor/lance-io/src/object_reader.rs @@ -0,0 +1,464 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::fs::File; +use std::ops::Range; +use std::sync::Arc; + +use crate::local::join_local_io; +#[cfg(windows)] +use crate::local::read_exact_at; +#[cfg(unix)] +use std::os::unix::fs::FileExt; + +use bytes::Bytes; +use futures::{ + FutureExt, + future::{BoxFuture, Shared}, + stream::{self, StreamExt}, +}; +use lance_core::deepsize::DeepSizeOf; +use lance_core::{Error, Result, error::CloneableError}; +use object_store::ObjectStoreExt; +use object_store::{GetOptions, GetResult, ObjectStore, Result as OSResult, path::Path}; +use tokio::sync::OnceCell; +use tracing::instrument; + +use crate::{ + object_store::DEFAULT_CLOUD_IO_PARALLELISM, + traits::{ByteStream, Reader}, +}; + +trait StaticGetRange { + fn path(&self) -> &Path; + fn get_range(&self) -> BoxFuture<'static, OSResult>; +} + +/// A wrapper around an object store and a path that implements a static +/// get_range method by assuming self is stored in an Arc. +struct GetRequest { + object_store: Arc, + path: Path, + options: GetOptions, +} + +impl StaticGetRange for Arc { + fn path(&self) -> &Path { + &self.path + } + + fn get_range(&self) -> BoxFuture<'static, OSResult> { + let store_and_path = self.clone(); + Box::pin(async move { + store_and_path + .object_store + .get_opts(&store_and_path.path, store_and_path.options.clone()) + .await + }) + } +} + +/// Object Reader +/// +/// Object Store + Base Path +#[derive(Debug)] +pub struct CloudObjectReader { + // Object Store. + pub object_store: Arc, + // File path + pub path: Path, + // File size, if known. + size: OnceCell, + + block_size: usize, + download_retry_count: usize, +} + +impl DeepSizeOf for CloudObjectReader { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // Skipping object_store because there is no easy way to do that and it shouldn't be too big + self.path.as_ref().deep_size_of_children(context) + } +} + +impl CloudObjectReader { + /// Create an ObjectReader from URI + pub fn new( + object_store: Arc, + path: Path, + block_size: usize, + known_size: Option, + download_retry_count: usize, + ) -> Result { + Ok(Self { + object_store, + path, + size: OnceCell::new_with(known_size), + block_size, + download_retry_count, + }) + } +} + +// Retries for the initial request are handled by object store, but +// there are no retries for failures that occur during the streaming +// of the response body. Thus we add an outer retry loop here. +async fn do_with_retry<'a, O>(f: impl Fn() -> BoxFuture<'a, OSResult> + Clone) -> OSResult { + let mut retries = 3; + loop { + let f = f.clone(); + match f().await { + Ok(val) => return Ok(val), + Err(err) => { + if retries == 0 { + return Err(err); + } + retries -= 1; + } + } + } +} + +// We have a separate retry loop here. This is because object_store does not +// attempt retries on downloads that fail during streaming of the response body. +// +// However, this failure is pretty common (e.g. timeout) and we want to retry in these +// situations. In addition, we provide additional logging information in these +// failures cases. +async fn do_get_with_outer_retry( + download_retry_count: usize, + get_request: Arc, + desc: impl Fn() -> String, +) -> OSResult { + let mut retries = download_retry_count; + loop { + let get_request_clone = get_request.clone(); + let get_result = do_with_retry(move || get_request_clone.get_range()).await?; + match get_result.bytes().await { + Ok(bytes) => return Ok(bytes), + Err(err) => { + if retries == 0 { + log::warn!( + "Failed to download {} from {} after {} attempts. This may indicate that cloud storage is overloaded or your timeout settings are too restrictive. Error details: {:?}", + desc(), + get_request.path(), + download_retry_count, + err + ); + return Err(err); + } + log::debug!( + "Retrying {} from {} (remaining retries: {}). Error details: {:?}", + desc(), + get_request.path(), + retries, + err + ); + retries -= 1; + } + } + } +} + +impl Reader for CloudObjectReader { + fn path(&self) -> &Path { + &self.path + } + + fn block_size(&self) -> usize { + self.block_size + } + + fn io_parallelism(&self) -> usize { + DEFAULT_CLOUD_IO_PARALLELISM + } + + /// Object/File Size. + fn size(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { + self.size + .get_or_try_init(|| async move { + let meta = + do_with_retry(|| Box::pin(self.object_store.head(&self.path))).await?; + Ok(meta.size as usize) + }) + .await + .cloned() + }) + } + + #[instrument(level = "debug", skip(self))] + fn get_range(&self, range: Range) -> BoxFuture<'static, OSResult> { + let object_store = self.object_store.clone(); + let path = self.path.clone(); + let get_range = Range { + start: range.start as u64, + end: range.end as u64, + }; + Box::pin(async move { + let bytes = do_with_retry(move || { + let object_store = object_store.clone(); + let path = path.clone(); + let get_range = get_range.clone(); + Box::pin(async move { object_store.get_ranges(&path, &[get_range]).await }) + }) + .await?; + + bytes + .into_iter() + .next() + .ok_or_else(|| object_store::Error::Generic { + store: "CloudObjectReader", + source: "get_ranges returned no bytes".into(), + }) + }) + } + + #[instrument(level = "debug", skip_all)] + fn get_all(&self) -> BoxFuture<'_, OSResult> { + let get_request = Arc::new(GetRequest { + object_store: self.object_store.clone(), + path: self.path.clone(), + options: GetOptions::default(), + }); + Box::pin(async move { + do_get_with_outer_retry(self.download_retry_count, get_request, || { + "read_all".to_string() + }) + .await + }) + } + + fn get_stream(&self) -> BoxFuture<'_, OSResult> { + let get_request = Arc::new(GetRequest { + object_store: self.object_store.clone(), + path: self.path.clone(), + options: GetOptions::default(), + }); + Box::pin(async move { + let get_request_clone = get_request.clone(); + let get_result = do_with_retry(move || get_request_clone.get_range()).await?; + Ok(get_result.into_stream()) + }) + } + + fn get_range_stream(&self, range: Range) -> BoxFuture<'_, OSResult> { + let get_request = Arc::new(GetRequest { + object_store: self.object_store.clone(), + path: self.path.clone(), + options: GetOptions { + range: Some( + Range { + start: range.start as u64, + end: range.end as u64, + } + .into(), + ), + ..Default::default() + }, + }); + Box::pin(async move { + let get_request_clone = get_request.clone(); + let get_result = do_with_retry(move || get_request_clone.get_range()).await?; + Ok(get_result.into_stream()) + }) + } +} + +#[derive(Debug)] +pub struct SmallReaderInner { + path: Path, + size: usize, + state: std::sync::Mutex, +} + +/// A reader for a file so small, we just eagerly read it all into memory. +/// +/// When created, it represents a future that will read the whole file into memory. +/// +/// On the first read call, it will start the read. Multiple threads can call read at the same time. +/// +/// Once the read is complete, any thread can call read again to get the result. +#[derive(Clone, Debug)] +pub struct SmallReader { + inner: Arc, +} + +enum SmallReaderState { + Loading(Shared>>), + Finished(std::result::Result), +} + +impl std::fmt::Debug for SmallReaderState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Loading(_) => write!(f, "Loading"), + Self::Finished(Ok(data)) => { + write!(f, "Finished({} bytes)", data.len()) + } + Self::Finished(Err(err)) => { + write!(f, "Finished({})", err.0) + } + } + } +} + +impl SmallReader { + pub fn new( + store: Arc, + path: Path, + download_retry_count: usize, + size: usize, + ) -> Self { + let path_ref = path.clone(); + let state = SmallReaderState::Loading( + Box::pin(async move { + let object_reader = + CloudObjectReader::new(store, path_ref, 0, None, download_retry_count) + .map_err(CloneableError)?; + object_reader + .get_all() + .await + .map_err(|err| CloneableError(Error::from(err))) + }) + .boxed() + .shared(), + ); + Self { + inner: Arc::new(SmallReaderInner { + path, + size, + state: std::sync::Mutex::new(state), + }), + } + } +} + +impl SmallReaderInner { + async fn wait(&self) -> OSResult { + let future = { + let state = self.state.lock().unwrap(); + match &*state { + SmallReaderState::Loading(future) => future.clone(), + SmallReaderState::Finished(result) => { + return result.clone().map_err(|err| err.0.into()); + } + } + }; + + let result = future.await; + let result_to_return = result.clone().map_err(|err| err.0.into()); + let mut state = self.state.lock().unwrap(); + if matches!(*state, SmallReaderState::Loading(_)) { + *state = SmallReaderState::Finished(result); + } + result_to_return + } +} + +impl Reader for SmallReader { + fn path(&self) -> &Path { + &self.inner.path + } + + fn block_size(&self) -> usize { + 64 * 1024 + } + + fn io_parallelism(&self) -> usize { + 1024 + } + + /// Object/File Size. + fn size(&self) -> BoxFuture<'_, OSResult> { + let size = self.inner.size; + Box::pin(async move { Ok(size) }) + } + + fn get_range(&self, range: Range) -> BoxFuture<'static, OSResult> { + let inner = self.inner.clone(); + Box::pin(async move { + let bytes = inner.wait().await?; + let start = range.start; + let end = range.end; + if start >= bytes.len() || end > bytes.len() { + return Err(object_store::Error::Generic { + store: "memory", + source: format!( + "Invalid range {}..{} for object of size {} bytes", + start, + end, + bytes.len() + ) + .into(), + }); + } + Ok(bytes.slice(range)) + }) + } + + fn get_all(&self) -> BoxFuture<'_, OSResult> { + Box::pin(async move { self.inner.wait().await }) + } +} + +pub(crate) fn stream_local_range( + file: Arc, + path: Path, + io_tracker: Arc, + range: Range, + chunk_size: usize, +) -> ByteStream { + stream::try_unfold( + (file, path, io_tracker, range.start, range.end), + move |state| async move { + let (file, path, io_tracker, start, end) = state; + if start >= end { + return Ok(None); + } + + let next = (start + chunk_size).min(end); + let file_clone = file.clone(); + let path_clone = path.clone(); + let num_bytes = (next - start) as u64; + let metrics = io_tracker.begin_io("get"); + let result = join_local_io(tokio::task::spawn_blocking(move || { + let mut buf = bytes::BytesMut::with_capacity(next - start); + // Safety: buffer capacity matches the exact number of bytes we read below. + unsafe { buf.set_len(next - start) }; + #[cfg(unix)] + file_clone.read_exact_at(buf.as_mut(), start as u64)?; + #[cfg(windows)] + read_exact_at(file_clone, buf.as_mut(), start as u64)?; + Ok::<_, std::io::Error>(buf.freeze()) + })) + .await; + metrics.record(&result, num_bytes); + let bytes = result?; + + io_tracker.record_read( + "get_range_stream", + path_clone, + num_bytes, + Some(start as u64..next as u64), + ); + + Ok(Some((bytes, (file, path, io_tracker, next, end)))) + }, + ) + .boxed() +} + +impl DeepSizeOf for SmallReader { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + let mut size = self.inner.path.as_ref().deep_size_of_children(context); + + if let Ok(guard) = self.inner.state.try_lock() + && let SmallReaderState::Finished(Ok(data)) = &*guard + { + size += data.len(); + } + + size + } +} diff --git a/vendor/lance-io/src/object_store.rs b/vendor/lance-io/src/object_store.rs new file mode 100644 index 000000000..2517d5932 --- /dev/null +++ b/vendor/lance-io/src/object_store.rs @@ -0,0 +1,1879 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extend [object_store::ObjectStore] functionalities + +use std::borrow::Cow; +use std::collections::HashMap; +use std::ops::Range; +use std::pin::Pin; +use std::str::FromStr; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use bytes::Bytes; +use chrono::{DateTime, Utc}; +use futures::{FutureExt, Stream}; +use futures::{StreamExt, TryStreamExt, future, stream::BoxStream}; +use lance_core::deepsize::DeepSizeOf; +use lance_core::error::LanceOptionExt; +use lance_core::utils::parse::str_is_truthy; +use list_retry::ListRetryStream; +use object_store::DynObjectStore; +use object_store::ObjectStoreExt as OSObjectStoreExt; +#[cfg(feature = "aws")] +use object_store::aws::AwsCredentialProvider; +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +use object_store::{ClientOptions, HeaderMap, HeaderValue}; +use object_store::{ListResult, ObjectMeta, ObjectStore as OSObjectStore, path::Path}; +use providers::local::FileStoreProvider; +use providers::memory::MemoryStoreProvider; +use tokio::io::AsyncWriteExt; +use url::Url; + +use super::local::LocalObjectReader; +#[cfg(target_os = "linux")] +use crate::uring::{UringCurrentThreadReader, UringReader}; +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +pub(crate) mod dynamic_credentials; +#[cfg(any( + feature = "aws", + feature = "oss", + feature = "huggingface", + feature = "tos" +))] +pub(crate) mod dynamic_opendal; +mod list_retry; +#[cfg(feature = "metrics")] +pub mod metrics; +pub mod providers; +pub mod storage_options; +#[cfg(test)] +pub(crate) mod test_utils; +pub mod throttle; +mod tracing; +use crate::object_reader::SmallReader; +use crate::object_writer::{LocalWriter, WriteResult}; +use crate::traits::{WriteExt, Writer}; +use crate::utils::tracking_store::{IOTracker, IoStats}; +use crate::{object_reader::CloudObjectReader, object_writer::ObjectWriter, traits::Reader}; +use lance_core::{Error, Result}; + +// Local disks tend to do fine with a few threads +// Note: the number of threads here also impacts the number of files +// we need to read in some situations. So keeping this at 8 keeps the +// RAM on our scanner down. +pub const DEFAULT_LOCAL_IO_PARALLELISM: usize = 8; +// Cloud disks often need many many threads to saturate the network +pub const DEFAULT_CLOUD_IO_PARALLELISM: usize = 64; + +const DEFAULT_LOCAL_BLOCK_SIZE: usize = 4 * 1024; // 4KB block size +#[cfg(any( + feature = "aws", + feature = "gcp", + feature = "azure", + feature = "oss", + feature = "tencent", + feature = "huggingface", + feature = "tos", + feature = "goosefs", +))] +const DEFAULT_CLOUD_BLOCK_SIZE: usize = 64 * 1024; // 64KB block size + +pub static DEFAULT_MAX_IOP_SIZE: std::sync::LazyLock = std::sync::LazyLock::new(|| { + std::env::var("LANCE_MAX_IOP_SIZE") + .map(|val| val.parse().unwrap()) + .unwrap_or(16 * 1024 * 1024) +}); + +pub const DEFAULT_DOWNLOAD_RETRY_COUNT: usize = 3; + +pub use providers::{ObjectStoreProvider, ObjectStoreRegistry}; +pub use storage_options::{ + BASE_SCOPED_OPTION_PREFIX, BaseScopedStorageOptionsProvider, EXPIRES_AT_MILLIS_KEY, + LanceNamespaceStorageOptionsProvider, REFRESH_OFFSET_MILLIS_KEY, StorageOptionsAccessor, + StorageOptionsProvider, has_base_scoped_options, parse_base_scoped_key, + resolve_base_scoped_options, +}; + +#[async_trait] +pub trait ObjectStoreExt { + /// Returns true if the file exists. + async fn exists(&self, path: &Path) -> Result; + + /// Read all files (start from base directory) recursively + /// + /// unmodified_since can be specified to only return files that have not been modified since the given time. + fn read_dir_all<'a, 'b>( + &'a self, + dir_path: impl Into<&'b Path> + Send, + unmodified_since: Option>, + ) -> BoxStream<'a, Result>; +} + +#[async_trait] +impl ObjectStoreExt for O { + fn read_dir_all<'a, 'b>( + &'a self, + dir_path: impl Into<&'b Path> + Send, + unmodified_since: Option>, + ) -> BoxStream<'a, Result> { + let output = self.list(Some(dir_path.into())).map_err(|e| e.into()); + if let Some(unmodified_since_val) = unmodified_since { + output + .try_filter(move |file| future::ready(file.last_modified <= unmodified_since_val)) + .boxed() + } else { + output.boxed() + } + } + + async fn exists(&self, path: &Path) -> Result { + match self.head(path).await { + Ok(_) => Ok(true), + Err(object_store::Error::NotFound { path: _, source: _ }) => Ok(false), + Err(e) => Err(e.into()), + } + } +} + +/// Wraps [ObjectStore](object_store::ObjectStore) +#[derive(Debug, Clone)] +pub struct ObjectStore { + // Inner object store + pub inner: Arc, + scheme: String, + block_size: usize, + max_iop_size: u64, + /// Whether to use constant size upload parts for multipart uploads. This + /// is only necessary for Cloudflare R2. + pub use_constant_size_upload_parts: bool, + /// Whether we can assume that the list of files is lexically ordered. This + /// is true for object stores, but not for local filesystems. + pub list_is_lexically_ordered: bool, + io_parallelism: usize, + /// Number of times to retry a failed download + download_retry_count: usize, + /// IO tracker for monitoring read/write operations + io_tracker: IOTracker, + /// The datastore prefix that uniquely identifies this object store. It encodes information + /// which usually cannot be found in the URL such as Azure account name. The prefix plus the + /// path uniquely identifies any object inside the store. + pub store_prefix: String, +} + +impl DeepSizeOf for ObjectStore { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // We aren't counting `inner` here which is problematic but an ObjectStore + // shouldn't be too big. The only exception might be the write cache but, if + // the writer cache has data, it means we're using it somewhere else that isn't + // a cache and so that doesn't really count. + self.scheme.deep_size_of_children(context) + self.block_size.deep_size_of_children(context) + } +} + +impl std::fmt::Display for ObjectStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "ObjectStore({})", self.scheme) + } +} + +pub trait WrappingObjectStore: std::fmt::Debug + Send + Sync { + /// Wrap an object store with additional functionality + /// + /// The store_prefix is a string which uniquely identifies the object + /// store being wrapped. + fn wrap(&self, store_prefix: &str, original: Arc) -> Arc; +} + +#[derive(Debug, Clone)] +pub struct ChainedWrappingObjectStore { + wrappers: Vec>, +} + +impl ChainedWrappingObjectStore { + pub fn new(wrappers: Vec>) -> Self { + Self { wrappers } + } + + pub fn add_wrapper(&mut self, wrapper: Arc) { + self.wrappers.push(wrapper); + } +} + +impl WrappingObjectStore for ChainedWrappingObjectStore { + fn wrap(&self, store_prefix: &str, original: Arc) -> Arc { + self.wrappers + .iter() + .fold(original, |acc, wrapper| wrapper.wrap(store_prefix, acc)) + } +} + +/// Parameters to create an [ObjectStore] +/// +#[derive(Debug, Clone)] +pub struct ObjectStoreParams { + pub block_size: Option, + #[deprecated(note = "Implement an ObjectStoreProvider instead")] + pub object_store: Option<(Arc, Url)>, + /// Refresh offset for AWS credentials when using the legacy AWS credentials path. + /// For StorageOptionsAccessor, use `refresh_offset_millis` storage option instead. + pub s3_credentials_refresh_offset: Duration, + #[cfg(feature = "aws")] + pub aws_credentials: Option, + pub object_store_wrapper: Option>, + /// Unified storage options accessor with caching and automatic refresh + /// + /// Provides storage options and optionally a dynamic provider for automatic + /// credential refresh. Use `StorageOptionsAccessor::with_static_options()` for static + /// options or `StorageOptionsAccessor::with_initial_and_provider()` for dynamic refresh. + pub storage_options_accessor: Option>, + /// Use constant size upload parts for multipart uploads. Only necessary + /// for Cloudflare R2, which doesn't support variable size parts. When this + /// is false, max upload size is 2.5TB. When this is true, the max size is + /// 50GB. + pub use_constant_size_upload_parts: bool, + pub list_is_lexically_ordered: Option, +} + +impl Default for ObjectStoreParams { + fn default() -> Self { + #[allow(deprecated)] + Self { + object_store: None, + block_size: None, + s3_credentials_refresh_offset: Duration::from_secs(60), + #[cfg(feature = "aws")] + aws_credentials: None, + object_store_wrapper: None, + storage_options_accessor: None, + use_constant_size_upload_parts: false, + list_is_lexically_ordered: None, + } + } +} + +impl ObjectStoreParams { + /// Get the StorageOptionsAccessor from the params + pub fn get_accessor(&self) -> Option> { + self.storage_options_accessor.clone() + } + + /// Get storage options from the accessor, if any + /// + /// Returns the initial storage options from the accessor without triggering refresh. + pub fn storage_options(&self) -> Option<&HashMap> { + self.storage_options_accessor + .as_ref() + .and_then(|a| a.initial_storage_options()) + } + + /// Resolve these params for a single base path scope. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to one registered base path; see + /// [`StorageOptionsAccessor::scoped_to_base`]. Returns the params unchanged + /// when the storage options contain no base-scoped entries. + pub fn scoped_to_base(&self, base_id: Option) -> Cow<'_, Self> { + let Some(accessor) = &self.storage_options_accessor else { + return Cow::Borrowed(self); + }; + let scoped = accessor.scoped_to_base(base_id); + if Arc::ptr_eq(&scoped, accessor) { + Cow::Borrowed(self) + } else { + Cow::Owned(Self { + storage_options_accessor: Some(scoped), + ..self.clone() + }) + } + } +} + +fn wrapper_allocation_ptr(wrapper: &Arc) -> *const () { + // Trait object pointers include vtable metadata, which is not stable across codegen units. + // Cache identity must follow the Arc allocation instead. + Arc::as_ptr(wrapper) as *const () +} + +// We implement hash for caching +impl std::hash::Hash for ObjectStoreParams { + #[allow(deprecated)] + fn hash(&self, state: &mut H) { + // For hashing, we use pointer values for ObjectStore, S3 credentials, wrapper + self.block_size.hash(state); + if let Some((store, url)) = &self.object_store { + Arc::as_ptr(store).hash(state); + url.hash(state); + } + self.s3_credentials_refresh_offset.hash(state); + #[cfg(feature = "aws")] + if let Some(aws_credentials) = &self.aws_credentials { + Arc::as_ptr(aws_credentials).hash(state); + } + if let Some(wrapper) = &self.object_store_wrapper { + wrapper_allocation_ptr(wrapper).hash(state); + } + if let Some(accessor) = &self.storage_options_accessor { + accessor.accessor_id().hash(state); + } + self.use_constant_size_upload_parts.hash(state); + self.list_is_lexically_ordered.hash(state); + } +} + +// We implement eq for caching +impl Eq for ObjectStoreParams {} +impl PartialEq for ObjectStoreParams { + #[allow(deprecated)] + fn eq(&self, other: &Self) -> bool { + #[cfg(feature = "aws")] + if self.aws_credentials.is_some() != other.aws_credentials.is_some() { + return false; + } + + // For equality, we use pointer comparison for ObjectStore, S3 credentials, wrapper + // For accessor, we use accessor_id() for semantic equality + self.block_size == other.block_size + && self + .object_store + .as_ref() + .map(|(store, url)| (Arc::as_ptr(store), url)) + == other + .object_store + .as_ref() + .map(|(store, url)| (Arc::as_ptr(store), url)) + && self.s3_credentials_refresh_offset == other.s3_credentials_refresh_offset + && self + .object_store_wrapper + .as_ref() + .map(wrapper_allocation_ptr) + == other + .object_store_wrapper + .as_ref() + .map(wrapper_allocation_ptr) + && self + .storage_options_accessor + .as_ref() + .map(|a| a.accessor_id()) + == other + .storage_options_accessor + .as_ref() + .map(|a| a.accessor_id()) + && self.use_constant_size_upload_parts == other.use_constant_size_upload_parts + && self.list_is_lexically_ordered == other.list_is_lexically_ordered + } +} + +/// Convert a URI string or local path to a URL +/// +/// This function handles both proper URIs (with schemes like `file://`, `s3://`, etc.) +/// and plain local filesystem paths. On Windows, it correctly handles drive letters +/// that might be parsed as URL schemes. +/// +/// # Examples +/// +/// ``` +/// # use lance_io::object_store::uri_to_url; +/// // URIs are preserved +/// let url = uri_to_url("s3://bucket/path").unwrap(); +/// assert_eq!(url.scheme(), "s3"); +/// +/// // Local paths are converted to file:// URIs +/// # #[cfg(unix)] +/// let url = uri_to_url("/tmp/data").unwrap(); +/// # #[cfg(unix)] +/// assert_eq!(url.scheme(), "file"); +/// ``` +pub fn uri_to_url(uri: &str) -> Result { + match Url::parse(uri) { + Ok(url) if url.scheme().len() == 1 && cfg!(windows) => { + // On Windows, the drive is parsed as a scheme + local_path_to_url(uri) + } + Ok(url) => Ok(url), + Err(_) => local_path_to_url(uri), + } +} + +fn expand_path(str_path: impl AsRef) -> Result { + let str_path = str_path.as_ref(); + let expanded = expand_tilde_path(str_path).unwrap_or_else(|| str_path.into()); + + let mut expanded_path = path_abs::PathAbs::new(expanded) + .unwrap() + .as_path() + .to_path_buf(); + // path_abs::PathAbs::new(".") returns an empty string. + if let Some(s) = expanded_path.as_path().to_str() + && s.is_empty() + { + expanded_path = std::env::current_dir()?; + } + + Ok(expanded_path) +} + +fn expand_tilde_path(path: &str) -> Option { + let home_dir = std::env::home_dir()?; + if path == "~" { + return Some(home_dir); + } + if let Some(stripped) = path.strip_prefix("~/") { + return Some(home_dir.join(stripped)); + } + #[cfg(windows)] + if let Some(stripped) = path.strip_prefix("~\\") { + return Some(home_dir.join(stripped)); + } + + None +} + +fn local_path_to_url(str_path: &str) -> Result { + let expanded_path = expand_path(str_path)?; + + Url::from_directory_path(expanded_path).map_err(|_| { + Error::invalid_input_source(format!("Invalid table location: '{}'", str_path).into()) + }) +} + +#[cfg(feature = "huggingface")] +fn parse_hf_repo_id(url: &Url) -> Result { + // Accept forms with repo type prefix (models/datasets/spaces) or legacy without. + let mut segments: Vec = Vec::new(); + if let Some(host) = url.host_str() { + segments.push(host.to_string()); + } + segments.extend( + url.path() + .trim_start_matches('/') + .split('/') + .map(|s| s.to_string()), + ); + + if segments.len() < 2 { + return Err(Error::invalid_input( + "Huggingface URL must contain at least owner and repo", + )); + } + + let repo_type_candidates = ["models", "datasets", "spaces"]; + let (owner, repo_with_rev) = if repo_type_candidates.contains(&segments[0].as_str()) { + if segments.len() < 3 { + return Err(Error::invalid_input( + "Huggingface URL missing owner/repo after repo type", + )); + } + (segments[1].as_str(), segments[2].as_str()) + } else { + (segments[0].as_str(), segments[1].as_str()) + }; + + let repo = repo_with_rev + .split_once('@') + .map(|(r, _)| r) + .unwrap_or(repo_with_rev); + Ok(format!("{owner}/{repo}")) +} + +impl ObjectStore { + /// Parse from a string URI. + /// + /// Returns the ObjectStore instance and the absolute path to the object. + /// + /// This uses the default [ObjectStoreRegistry] to find the object store. To + /// allow for potential re-use of object store instances, it's recommended to + /// create a shared [ObjectStoreRegistry] and pass that to [Self::from_uri_and_params]. + pub async fn from_uri(uri: &str) -> Result<(Arc, Path)> { + let registry = Arc::new(ObjectStoreRegistry::default()); + + Self::from_uri_and_params(registry, uri, &ObjectStoreParams::default()).await + } + + /// Parse from a string URI. + /// + /// Returns the ObjectStore instance and the absolute path to the object. + pub async fn from_uri_and_params( + registry: Arc, + uri: &str, + params: &ObjectStoreParams, + ) -> Result<(Arc, Path)> { + #[allow(deprecated)] + if let Some((store, path)) = params.object_store.as_ref() { + let mut inner = store.clone(); + let store_prefix = + registry.calculate_object_store_prefix(uri, params.storage_options())?; + + let mut io_tracker = IOTracker::default(); + meter_store(&mut inner, &mut io_tracker, &store_prefix); + + if let Some(wrapper) = params.object_store_wrapper.as_ref() { + inner = wrapper.wrap(&store_prefix, inner); + } + + // Always wrap with IO tracking + let tracked_store = io_tracker.wrap("", inner); + + let store = Self { + inner: tracked_store, + scheme: path.scheme().to_string(), + block_size: params.block_size.unwrap_or(64 * 1024), + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or_default(), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT, + io_tracker, + store_prefix, + }; + let path = Path::parse(path.path())?; + return Ok((Arc::new(store), path)); + } + let url = uri_to_url(uri)?; + + let store = registry.get_store(url.clone(), params).await?; + // We know the scheme is valid if we got a store back. + let provider = registry.get_provider(url.scheme()).expect_ok()?; + let path = provider.extract_path(&url)?; + + Ok((store, path)) + } + + /// Extract the path component from a URI without initializing the object store. + /// + /// This is a synchronous operation that only parses the URI and extracts the path, + /// without creating or initializing any object store instance. + /// + /// # Arguments + /// + /// * `registry` - The object store registry to get the provider + /// * `uri` - The URI to extract the path from + /// + /// # Returns + /// + /// The extracted path component + pub fn extract_path_from_uri(registry: Arc, uri: &str) -> Result { + let url = uri_to_url(uri)?; + let provider = registry + .get_provider(url.scheme()) + .ok_or_else(|| Error::invalid_input(format!("Unknown scheme: {}", url.scheme())))?; + provider.extract_path(&url) + } + + #[deprecated(note = "Use `from_uri` instead")] + pub fn from_path(str_path: &str) -> Result<(Arc, Path)> { + Self::from_uri_and_params( + Arc::new(ObjectStoreRegistry::default()), + str_path, + &Default::default(), + ) + .now_or_never() + .unwrap() + } + + /// Local object store. + pub fn local() -> Self { + let provider = FileStoreProvider; + provider + .new_store(Url::parse("file:///").unwrap(), &Default::default()) + .now_or_never() + .unwrap() + .unwrap() + } + + /// Create a in-memory object store directly for testing. + pub fn memory() -> Self { + let provider = MemoryStoreProvider; + provider + .new_store(Url::parse("memory:///").unwrap(), &Default::default()) + .now_or_never() + .unwrap() + .unwrap() + } + + /// Returns true if the object store pointed to a local file system. + pub fn is_local(&self) -> bool { + self.scheme == "file" || self.scheme == "file+uring" + } + + pub fn is_cloud(&self) -> bool { + if self.is_local() || self.scheme == "memory" || self.scheme == "shared-memory" { + return false; + } + true + } + + /// Whether this object store prefers the lite scheduler. + /// + /// The lite scheduler is designed for backends like io_uring where + /// tasks should only be polled when the consumer polls them. + pub fn prefers_lite_scheduler(&self) -> bool { + self.scheme == "file+uring" + } + + pub fn scheme(&self) -> &str { + &self.scheme + } + + pub fn block_size(&self) -> usize { + self.block_size + } + + pub fn max_iop_size(&self) -> u64 { + self.max_iop_size + } + + /// The amount of parallelism to use for I/O operations. + /// + /// Honors the `LANCE_IO_THREADS` override when set, otherwise the store's configured value. + /// Always at least 1: callers feed this straight into `buffered` / `buffer_unordered`, and a + /// window of 0 makes those streams never poll their input — e.g. a metadata-only `count_rows` + /// would hang rather than return. + pub fn io_parallelism(&self) -> usize { + std::env::var("LANCE_IO_THREADS") + .map(|val| val.parse::().unwrap()) + .unwrap_or(self.io_parallelism) + .max(1) + } + + /// Get the IO tracker for this object store + /// + /// The IO tracker can be used to get statistics about read/write operations + /// performed on this object store. + pub fn io_tracker(&self) -> &IOTracker { + &self.io_tracker + } + + /// Get a snapshot of current IO statistics without resetting counters + /// + /// Returns the current IO statistics without modifying the internal state. + /// Use this when you need to check stats without resetting them. + pub fn io_stats_snapshot(&self) -> IoStats { + self.io_tracker.stats() + } + + /// Get incremental IO statistics since the last call to this method + /// + /// Returns the accumulated statistics since the last call and resets the + /// counters to zero. This is useful for tracking IO operations between + /// different stages of processing. + pub fn io_stats_incremental(&self) -> IoStats { + self.io_tracker.incremental_stats() + } + + /// Open a file for path. + /// + /// Parameters + /// - ``path``: Absolute path to the file. + pub async fn open(&self, path: &Path) -> Result> { + match self.scheme.as_str() { + "file" => { + LocalObjectReader::open_with_tracker( + path, + self.block_size, + None, + Arc::new(self.io_tracker.clone()), + ) + .await + } + #[cfg(target_os = "linux")] + "file+uring" => { + // Check if current-thread mode enabled + let use_current_thread = std::env::var("LANCE_URING_CURRENT_THREAD") + .map(|v| str_is_truthy(&v)) + .unwrap_or(false); + + if use_current_thread { + UringCurrentThreadReader::open( + path, + self.block_size, + None, + Arc::new(self.io_tracker.clone()), + ) + .await + } else { + UringReader::open( + path, + self.block_size, + None, + Arc::new(self.io_tracker.clone()), + ) + .await + } + } + _ => Ok(Box::new(CloudObjectReader::new( + self.inner.clone(), + path.clone(), + self.block_size, + None, + self.download_retry_count, + )?)), + } + } + + /// Open a reader for a file with known size. + /// + /// This size may either have been retrieved from a list operation or + /// cached metadata. By passing in the known size, we can skip a HEAD / metadata + /// call. + pub async fn open_with_size(&self, path: &Path, known_size: usize) -> Result> { + // If we know the file is really small, we can read the whole thing + // as a single request. + if known_size <= self.block_size { + return Ok(Box::new(SmallReader::new( + self.inner.clone(), + path.clone(), + self.download_retry_count, + known_size, + ))); + } + + match self.scheme.as_str() { + "file" => { + LocalObjectReader::open_with_tracker( + path, + self.block_size, + Some(known_size), + Arc::new(self.io_tracker.clone()), + ) + .await + } + #[cfg(target_os = "linux")] + "file+uring" => { + // Check if current-thread mode enabled + let use_current_thread = std::env::var("LANCE_URING_CURRENT_THREAD") + .map(|v| str_is_truthy(&v)) + .unwrap_or(false); + + if use_current_thread { + UringCurrentThreadReader::open( + path, + self.block_size, + Some(known_size), + Arc::new(self.io_tracker.clone()), + ) + .await + } else { + UringReader::open( + path, + self.block_size, + Some(known_size), + Arc::new(self.io_tracker.clone()), + ) + .await + } + } + _ => Ok(Box::new(CloudObjectReader::new( + self.inner.clone(), + path.clone(), + self.block_size, + Some(known_size), + self.download_retry_count, + )?)), + } + } + + /// Create an [ObjectWriter] from local [std::path::Path] + pub async fn create_local_writer(path: &std::path::Path) -> Result { + let object_store = Self::local(); + let absolute_path = expand_path(path.to_string_lossy())?; + let os_path = Path::from_absolute_path(absolute_path)?; + ObjectWriter::new(&object_store, &os_path).await + } + + /// Open an [Reader] from local [std::path::Path] + pub async fn open_local(path: &std::path::Path) -> Result> { + let object_store = Self::local(); + let absolute_path = expand_path(path.to_string_lossy())?; + let os_path = Path::from_absolute_path(absolute_path)?; + object_store.open(&os_path).await + } + + /// Create a new file. + pub async fn create(&self, path: &Path) -> Result> { + match self.scheme.as_str() { + "file" => { + let local_path = super::local::to_local_path(path); + let local_path = std::path::PathBuf::from(&local_path); + if let Some(parent) = local_path.parent() { + tokio::fs::create_dir_all(parent).await?; + } + let parent = local_path + .parent() + .expect("file path must have parent") + .to_owned(); + let named_temp = + tokio::task::spawn_blocking(move || tempfile::NamedTempFile::new_in(parent)) + .await + .map_err(|e| Error::io(format!("spawn_blocking failed: {}", e)))??; + let (std_file, temp_path) = named_temp.into_parts(); + let file = tokio::fs::File::from_std(std_file); + Ok(Box::new(LocalWriter::new( + file, + path.clone(), + temp_path, + Arc::new(self.io_tracker.clone()), + ))) + } + _ => Ok(Box::new(ObjectWriter::new(self, path).await?)), + } + } + + /// A helper function to create a file and write content to it. + pub async fn put(&self, path: &Path, content: &[u8]) -> Result { + let mut writer = self.create(path).await?; + writer.write_all(content).await?; + Writer::shutdown(writer.as_mut()).await + } + + pub async fn delete(&self, path: &Path) -> Result<()> { + self.inner.delete(path).await?; + Ok(()) + } + + /// AWS S3 and GCS reject a single-shot server-side copy whose source is + /// larger than this; such sources are streamed through a multipart write. + const MAX_SINGLE_COPY_BYTES: u64 = 5 * 1024 * 1024 * 1024; // 5 GiB + + pub async fn copy(&self, from: &Path, to: &Path) -> Result<()> { + // S3 and GCS cap single-shot server-side copies at 5 GiB and object_store + // does not fall back to a multipart copy for larger sources + // (https://github.com/apache/arrow-rs-object-store/issues/563). Azure and + // other blob stores don't have this limit, so we only pay for the fallback + // (an extra size lookup) on S3 and GCS. + let multipart_copy_fallback = matches!(self.scheme.as_str(), "s3" | "s3+ddb" | "gs"); + self.copy_impl( + from, + to, + multipart_copy_fallback, + Self::MAX_SINGLE_COPY_BYTES, + ) + .await + } + + /// Copy `from` to `to`. When `multipart_copy_fallback` is set, a source + /// larger than `max_single_copy` is streamed through a multipart write + /// instead of a single-shot server-side copy. Both are parameters so tests + /// can drive the streaming path without a multi-gigabyte fixture or an S3 + /// endpoint. + async fn copy_impl( + &self, + from: &Path, + to: &Path, + multipart_copy_fallback: bool, + max_single_copy: u64, + ) -> Result<()> { + if self.is_local() { + // Use std::fs::copy for local filesystem to support cross-filesystem copies + let metrics = self.io_tracker.begin_io("copy"); + let result = super::local::copy_file(from, to); + metrics.record(&result, 0); + return result; + } + if multipart_copy_fallback { + // Reuse the reader for both the size lookup (a single cached HEAD) + // and the streamed copy, avoiding a separate HEAD request. + let reader = self.open(from).await?; + if reader.size().await? as u64 > max_single_copy { + let mut writer = self.create(to).await?; + writer.copy_from_reader(reader.as_ref()).await?; + Writer::shutdown(writer.as_mut()).await?; + return Ok(()); + } + } + Ok(self.inner.copy(from, to).await?) + } + + /// Read a directory (start from base directory) and returns all sub-paths in the directory. + pub async fn read_dir(&self, dir_path: impl Into) -> Result> { + let path = dir_path.into(); + let path = Path::parse(&path)?; + let output = self.inner.list_with_delimiter(Some(&path)).await?; + Ok(output + .common_prefixes + .iter() + .chain(output.objects.iter().map(|o| &o.location)) + .filter_map(|s| s.filename().map(|f| f.to_string())) + .collect()) + } + + /// Non-recursive, path-segment delimited list of a single directory level. + /// + /// Unlike [`Self::list`], which recurses into the entire subtree, this returns + /// only the immediate children of `prefix`: the child "directories" as + /// [`ListResult::common_prefixes`] and the direct child files as + /// [`ListResult::objects`]. + pub async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + Ok(self.inner.list_with_delimiter(prefix).await?) + } + + pub fn list( + &self, + path: Option, + ) -> Pin> + Send>> { + Box::pin(ListRetryStream::new(self.inner.clone(), path, 5).map(|m| m.map_err(|e| e.into()))) + } + + /// Read all files (start from base directory) recursively + /// + /// unmodified_since can be specified to only return files that have not been modified since the given time. + pub fn read_dir_all<'a, 'b>( + &'a self, + dir_path: impl Into<&'b Path> + Send, + unmodified_since: Option>, + ) -> BoxStream<'a, Result> { + self.inner.read_dir_all(dir_path, unmodified_since) + } + + /// Remove a directory recursively. + pub async fn remove_dir_all(&self, dir_path: impl Into) -> Result<()> { + let path = dir_path.into(); + let path = Path::parse(&path)?; + + if self.is_local() { + // The local file system provider needs to delete both files and directories. + // Counted as a single delete request, matching how `delete_stream` + // counts one batched request regardless of how many paths it removes. + let metrics = self.io_tracker.begin_io("delete"); + let result = super::local::remove_dir_all(&path); + metrics.record(&result, 0); + return result; + } + let sub_entries = self + .inner + .list(Some(&path)) + .map(|m| m.map(|meta| meta.location)) + .boxed(); + self.inner + .delete_stream(sub_entries) + .try_collect::>() + .await?; + if self.scheme == "file-object-store" { + // file-object-store tries to do everything as similarly as possible to the remote + // object stores. But we still have to delete the directory entries afterwards. + return super::local::remove_dir_all(&path); + } + Ok(()) + } + + pub fn remove_stream<'a>( + &'a self, + locations: BoxStream<'a, Result>, + ) -> BoxStream<'a, Result> { + let store = Arc::clone(&self.inner); + locations + .and_then(move |location| { + let store = Arc::clone(&store); + async move { + store.delete(&location).await?; + Ok(location) + } + }) + .boxed() + } + + /// Check a file exists. + pub async fn exists(&self, path: &Path) -> Result { + match self.inner.head(path).await { + Ok(_) => Ok(true), + Err(object_store::Error::NotFound { path: _, source: _ }) => Ok(false), + Err(e) => Err(e.into()), + } + } + + /// Get file size. + pub async fn size(&self, path: &Path) -> Result { + Ok(self.inner.head(path).await?.size) + } + + /// Convenience function to open a reader and read all the bytes + pub async fn read_one_all(&self, path: &Path) -> Result { + let reader = self.open(path).await?; + Ok(reader.get_all().await?) + } + + /// Convenience function open a reader and make a single request + /// + /// If you will be making multiple requests to the path it is more efficient to call [`Self::open`] + /// and then call [`Reader::get_range`] multiple times. + pub async fn read_one_range(&self, path: &Path, range: Range) -> Result { + let reader = self.open(path).await?; + Ok(reader.get_range(range).await?) + } +} + +/// Options that can be set for multiple object stores +#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy)] +pub enum LanceConfigKey { + /// Number of times to retry a download that fails + DownloadRetryCount, +} + +impl FromStr for LanceConfigKey { + type Err = Error; + + fn from_str(s: &str) -> std::result::Result { + match s.to_ascii_lowercase().as_str() { + "download_retry_count" => Ok(Self::DownloadRetryCount), + _ => Err(Error::invalid_input_source( + format!("Invalid LanceConfigKey: {}", s).into(), + )), + } + } +} + +#[derive(Clone, Debug, Default)] +pub struct StorageOptions(pub HashMap); + +impl StorageOptions { + /// Create a new instance of [`StorageOptions`] + pub fn new(options: HashMap) -> Self { + let mut options = options; + if let Ok(value) = std::env::var("AZURE_STORAGE_ALLOW_HTTP") { + options.insert("allow_http".into(), value); + } + if let Ok(value) = std::env::var("AZURE_STORAGE_USE_HTTP") { + options.insert("allow_http".into(), value); + } + if let Ok(value) = std::env::var("AWS_ALLOW_HTTP") { + options.insert("allow_http".into(), value); + } + if let Ok(value) = std::env::var("OBJECT_STORE_CLIENT_MAX_RETRIES") { + options.insert("client_max_retries".into(), value); + } + if let Ok(value) = std::env::var("OBJECT_STORE_CLIENT_RETRY_TIMEOUT") { + options.insert("client_retry_timeout".into(), value); + } + Self(options) + } + + /// Denotes if unsecure connections via http are allowed + pub fn allow_http(&self) -> bool { + self.0.iter().any(|(key, value)| { + key.to_ascii_lowercase().contains("allow_http") & str_is_truthy(value) + }) + } + + /// Number of times to retry a download that fails + pub fn download_retry_count(&self) -> usize { + self.0 + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("download_retry_count")) + .map(|(_, value)| value.parse::().unwrap_or(3)) + .unwrap_or(3) + } + + /// Max retry times to set in RetryConfig for object store client + pub fn client_max_retries(&self) -> usize { + self.0 + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("client_max_retries")) + .and_then(|(_, value)| value.parse::().ok()) + .unwrap_or(3) + } + + /// Seconds of timeout to set in RetryConfig for object store client + pub fn client_retry_timeout(&self) -> u64 { + self.0 + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("client_retry_timeout")) + .and_then(|(_, value)| value.parse::().ok()) + .unwrap_or(180) + } + + pub fn get(&self, key: &str) -> Option<&String> { + self.0.get(key) + } + + /// Build [`ClientOptions`] with default headers extracted from `headers.*` keys. + /// + /// Keys prefixed with `headers.` are parsed into HTTP headers. For example, + /// `headers.x-ms-version = 2023-11-03` results in a default header + /// `x-ms-version: 2023-11-03`. + /// + /// Returns an error if any `headers.*` key has an invalid header name or value. + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + pub fn client_options(&self) -> Result { + let mut headers = HeaderMap::new(); + for (key, value) in &self.0 { + if let Some(header_name) = key.strip_prefix("headers.") { + let name = header_name + .parse::() + .map_err(|e| { + Error::invalid_input(format!("invalid header name '{header_name}': {e}")) + })?; + let val = HeaderValue::from_str(value).map_err(|e| { + Error::invalid_input(format!("invalid header value for '{header_name}': {e}")) + })?; + headers.insert(name, val); + } + } + let mut client_options = ClientOptions::default(); + if !headers.is_empty() { + client_options = client_options.with_default_headers(headers); + } + Ok(client_options) + } + + /// Get the expiration time in milliseconds since epoch, if present + pub fn expires_at_millis(&self) -> Option { + self.0 + .get(EXPIRES_AT_MILLIS_KEY) + .and_then(|s| s.parse::().ok()) + } +} + +impl From> for StorageOptions { + fn from(value: HashMap) -> Self { + Self::new(value) + } +} + +static DEFAULT_OBJECT_STORE_REGISTRY: std::sync::LazyLock = + std::sync::LazyLock::new(ObjectStoreRegistry::default); + +impl ObjectStore { + #[allow(clippy::too_many_arguments)] + pub fn new( + mut store: Arc, + location: Url, + block_size: Option, + wrapper: Option>, + use_constant_size_upload_parts: bool, + list_is_lexically_ordered: bool, + io_parallelism: usize, + download_retry_count: usize, + storage_options: Option<&HashMap>, + ) -> Self { + let scheme = location.scheme(); + let block_size = block_size.unwrap_or_else(|| infer_block_size(scheme)); + let store_prefix = match DEFAULT_OBJECT_STORE_REGISTRY.get_provider(scheme) { + Some(provider) => provider + .calculate_object_store_prefix(&location, storage_options) + .unwrap(), + None => { + let store_prefix = format!("{}${}", location.scheme(), location.authority()); + log::warn!( + "Guessing that object store prefix is {}, since object store scheme is not found in registry.", + store_prefix + ); + store_prefix + } + }; + let mut io_tracker = IOTracker::default(); + meter_store(&mut store, &mut io_tracker, &store_prefix); + + let store = match wrapper { + Some(wrapper) => wrapper.wrap(&store_prefix, store), + None => store, + }; + + // Always wrap with IO tracking + let tracked_store = io_tracker.wrap("", store); + + Self { + inner: tracked_store, + scheme: scheme.into(), + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts, + list_is_lexically_ordered, + io_parallelism, + download_retry_count, + io_tracker, + store_prefix, + } + } +} + +/// Wrap `inner` so its operations publish metrics labelled by `store_prefix`, +/// and label `io_tracker` with the same prefix so the local reads and writes +/// that bypass `inner` publish under it too. +/// +/// The two go together on purpose: a store metered on one path but not the other +/// would report a partial picture that reads like a complete one. Every +/// constructor that hands an [`ObjectStore`] to a caller must route its `inner` +/// through here, or through nothing at all. +#[cfg(feature = "metrics")] +fn meter_store(inner: &mut Arc, io_tracker: &mut IOTracker, store_prefix: &str) { + use crate::object_store::metrics::ObjectStoreMetricsExt; + io_tracker.set_metrics_base(store_prefix); + *inner = inner.clone().metered(store_prefix.to_owned()); +} + +#[cfg(not(feature = "metrics"))] +fn meter_store( + _inner: &mut Arc, + _io_tracker: &mut IOTracker, + _store_prefix: &str, +) { +} + +fn infer_block_size(scheme: &str) -> usize { + // Block size: On local file systems, we use 4KB block size. On cloud + // object stores, we use 64KB block size. This is generally the largest + // block size where we don't see a latency penalty. + match scheme { + "file" => 4 * 1024, + _ => 64 * 1024, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use async_trait::async_trait; + use bytes::Bytes; + use lance_core::utils::tempfile::{TempStdDir, TempStdFile, TempStrDir}; + use object_store::memory::InMemory; + use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, PutMultipartOptions, + PutOptions, PutPayload, PutResult, Result as OSResult, + }; + use rstest::rstest; + use std::env::set_current_dir; + use std::fmt::{Display, Formatter}; + use std::fs::{create_dir_all, write}; + use std::ops::Range; + use std::path::Path as StdPath; + use std::sync::atomic::{AtomicBool, Ordering}; + + /// Write test content to file. + fn write_to_file(path_str: &str, contents: &str) -> std::io::Result<()> { + let path = expand_path(path_str).map_err(std::io::Error::other)?; + std::fs::create_dir_all(path.parent().unwrap())?; + write(path, contents) + } + + async fn read_from_store(store: &ObjectStore, path: &Path) -> Result { + let test_file_store = store.open(path).await.unwrap(); + let size = test_file_store.size().await.unwrap(); + let bytes = test_file_store.get_range(0..size).await.unwrap(); + let contents = String::from_utf8(bytes.to_vec()).unwrap(); + Ok(contents) + } + + #[test] + fn test_io_parallelism_clamped_to_nonzero() { + // `io_parallelism()` feeds `buffered`/`buffer_unordered` windows; a value of 0 makes those + // streams never poll, hanging callers (e.g. a metadata-only `count_rows`). It must clamp. + let store = ObjectStore::local(); + + // SAFETY: process-global env var, set and restored within this test. `io_parallelism()` + // only reads it, and a concurrent reader observes a valid clamped value, never 0. + unsafe { std::env::set_var("LANCE_IO_THREADS", "0") }; + assert_eq!( + store.io_parallelism(), + 1, + "LANCE_IO_THREADS=0 must clamp to 1" + ); + + unsafe { std::env::set_var("LANCE_IO_THREADS", "8") }; + assert_eq!( + store.io_parallelism(), + 8, + "a positive override must pass through unchanged" + ); + + unsafe { std::env::remove_var("LANCE_IO_THREADS") }; + assert!( + store.io_parallelism() >= 1, + "the configured default parallelism must be at least 1" + ); + } + + #[tokio::test] + async fn test_absolute_paths() { + let tmp_path = TempStrDir::default(); + write_to_file( + &format!("{tmp_path}/bar/foo.lance/test_file"), + "TEST_CONTENT", + ) + .unwrap(); + + // test a few variations of the same path + for uri in &[ + format!("{tmp_path}/bar/foo.lance"), + format!("{tmp_path}/./bar/foo.lance"), + format!("{tmp_path}/bar/foo.lance/../foo.lance"), + ] { + let (store, path) = ObjectStore::from_uri(uri).await.unwrap(); + let contents = read_from_store(store.as_ref(), &path.clone().join("test_file")) + .await + .unwrap(); + assert_eq!(contents, "TEST_CONTENT"); + } + } + + #[tokio::test] + async fn test_cloud_paths() { + let uri = "s3://bucket/foo.lance"; + let (store, path) = ObjectStore::from_uri(uri).await.unwrap(); + assert_eq!(store.scheme, "s3"); + assert_eq!(path.to_string(), "foo.lance"); + + let (store, path) = ObjectStore::from_uri("s3+ddb://bucket/foo.lance") + .await + .unwrap(); + assert_eq!(store.scheme, "s3"); + assert_eq!(path.to_string(), "foo.lance"); + + let (store, path) = ObjectStore::from_uri("gs://bucket/foo.lance") + .await + .unwrap(); + assert_eq!(store.scheme, "gs"); + assert_eq!(path.to_string(), "foo.lance"); + + let (store, path) = + ObjectStore::from_uri("abfss://filesystem@account.dfs.core.windows.net/foo.lance") + .await + .unwrap(); + assert_eq!(store.scheme, "abfss"); + assert_eq!(path.to_string(), "foo.lance"); + } + + async fn test_block_size_used_test_helper( + uri: &str, + storage_options: Option>, + default_expected_block_size: usize, + ) { + // Test the default + let registry = Arc::new(ObjectStoreRegistry::default()); + let accessor = storage_options + .clone() + .map(|opts| Arc::new(StorageOptionsAccessor::with_static_options(opts))); + let params = ObjectStoreParams { + storage_options_accessor: accessor.clone(), + ..ObjectStoreParams::default() + }; + let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms) + .await + .unwrap(); + assert_eq!(store.block_size, default_expected_block_size); + + // Ensure param is used + let registry = Arc::new(ObjectStoreRegistry::default()); + let params = ObjectStoreParams { + block_size: Some(1024), + storage_options_accessor: accessor, + ..ObjectStoreParams::default() + }; + let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms) + .await + .unwrap(); + assert_eq!(store.block_size, 1024); + } + + #[rstest] + #[case("s3://bucket/foo.lance", None)] + #[case("gs://bucket/foo.lance", None)] + #[case("az://account/bucket/foo.lance", + Some(HashMap::from([ + (String::from("account_name"), String::from("account")), + (String::from("container_name"), String::from("container")) + ])))] + #[case("abfss://filesystem@account.dfs.core.windows.net/foo.lance", + Some(HashMap::from([ + (String::from("account_name"), String::from("account")), + (String::from("container_name"), String::from("filesystem")) + ])))] + #[tokio::test] + async fn test_block_size_used_cloud( + #[case] uri: &str, + #[case] storage_options: Option>, + ) { + test_block_size_used_test_helper(uri, storage_options, 64 * 1024).await; + } + + #[rstest] + #[case("file")] + #[case("file-object-store")] + #[case("memory:///bucket/foo.lance")] + #[tokio::test] + async fn test_block_size_used_file(#[case] prefix: &str) { + let tmp_path = TempStrDir::default(); + let path = format!("{tmp_path}/bar/foo.lance/test_file"); + write_to_file(&path, "URL").unwrap(); + let uri = format!("{prefix}:///{path}"); + test_block_size_used_test_helper(&uri, None, 4 * 1024).await; + } + + #[tokio::test] + async fn test_relative_paths() { + let tmp_path = TempStrDir::default(); + write_to_file( + &format!("{tmp_path}/bar/foo.lance/test_file"), + "RELATIVE_URL", + ) + .unwrap(); + + set_current_dir(StdPath::new(tmp_path.as_ref())).expect("Error changing current dir"); + let (store, path) = ObjectStore::from_uri("./bar/foo.lance").await.unwrap(); + + let contents = read_from_store(store.as_ref(), &path.clone().join("test_file")) + .await + .unwrap(); + assert_eq!(contents, "RELATIVE_URL"); + } + + #[tokio::test] + async fn test_tilde_expansion() { + let uri = "~/foo.lance"; + write_to_file(&format!("{uri}/test_file"), "TILDE").unwrap(); + let (store, path) = ObjectStore::from_uri(uri).await.unwrap(); + let contents = read_from_store(store.as_ref(), &path.clone().join("test_file")) + .await + .unwrap(); + assert_eq!(contents, "TILDE"); + } + + #[tokio::test] + async fn test_read_directory() { + let path = TempStdDir::default(); + create_dir_all(path.join("foo").join("bar")).unwrap(); + create_dir_all(path.join("foo").join("zoo")).unwrap(); + create_dir_all(path.join("foo").join("zoo").join("abc")).unwrap(); + write_to_file( + path.join("foo").join("test_file").to_str().unwrap(), + "read_dir", + ) + .unwrap(); + let (store, base) = ObjectStore::from_uri(path.to_str().unwrap()).await.unwrap(); + + let sub_dirs = store.read_dir(base.clone().join("foo")).await.unwrap(); + assert_eq!(sub_dirs, vec!["bar", "zoo", "test_file"]); + } + + #[tokio::test] + async fn test_delete_directory_local_store() { + test_delete_directory("").await; + } + + #[tokio::test] + async fn test_delete_directory_file_object_store() { + test_delete_directory("file-object-store").await; + } + + async fn test_delete_directory(scheme: &str) { + let path = TempStdDir::default(); + create_dir_all(path.join("foo").join("bar")).unwrap(); + create_dir_all(path.join("foo").join("zoo")).unwrap(); + create_dir_all(path.join("foo").join("zoo").join("abc")).unwrap(); + write_to_file( + path.join("foo") + .join("bar") + .join("test_file") + .to_str() + .unwrap(), + "delete", + ) + .unwrap(); + let file_url = Url::from_directory_path(&path).unwrap(); + let url = if scheme.is_empty() { + file_url + } else { + let mut url = Url::parse(&format!("{scheme}:///")).unwrap(); + // Use the file:// URL's normalized path so this works on Windows too. + url.set_path(file_url.path()); + url + }; + let (store, base) = ObjectStore::from_uri(url.as_ref()).await.unwrap(); + store + .remove_dir_all(base.clone().join("foo")) + .await + .unwrap(); + + assert!(!path.join("foo").exists()); + } + + #[derive(Debug)] + struct TestWrapper { + called: AtomicBool, + + return_value: Arc, + } + + impl WrappingObjectStore for TestWrapper { + fn wrap( + &self, + _store_prefix: &str, + _original: Arc, + ) -> Arc { + self.called.store(true, Ordering::Relaxed); + + // return a mocked value so we can check if the final store is the one we expect + self.return_value.clone() + } + } + + impl TestWrapper { + fn called(&self) -> bool { + self.called.load(Ordering::Relaxed) + } + } + + #[tokio::test] + async fn test_wrapper_identity_is_stable_across_tasks() { + let wrapper = Arc::new(TestWrapper { + called: AtomicBool::new(false), + return_value: Arc::new(InMemory::new()), + }); + let initial_params = ObjectStoreParams { + object_store_wrapper: Some(wrapper.clone()), + ..ObjectStoreParams::default() + }; + let task_params = tokio::spawn(async move { + ObjectStoreParams { + object_store_wrapper: Some(wrapper), + ..ObjectStoreParams::default() + } + }) + .await + .unwrap(); + + assert_eq!(initial_params, task_params); + + let mut initial_hasher = std::hash::DefaultHasher::new(); + std::hash::Hash::hash(&initial_params, &mut initial_hasher); + let mut task_hasher = std::hash::DefaultHasher::new(); + std::hash::Hash::hash(&task_params, &mut task_hasher); + assert_eq!( + std::hash::Hasher::finish(&initial_hasher), + std::hash::Hasher::finish(&task_hasher) + ); + } + + #[tokio::test] + async fn test_wrapping_object_store_option_is_used() { + // Make a store for the inner store first + let mock_inner_store: Arc = Arc::new(InMemory::new()); + let registry = Arc::new(ObjectStoreRegistry::default()); + + assert_eq!(Arc::strong_count(&mock_inner_store), 1); + + let wrapper = Arc::new(TestWrapper { + called: AtomicBool::new(false), + return_value: mock_inner_store.clone(), + }); + + let params = ObjectStoreParams { + object_store_wrapper: Some(wrapper.clone()), + ..ObjectStoreParams::default() + }; + + // not called yet + assert!(!wrapper.called()); + + let _ = ObjectStore::from_uri_and_params(registry, "memory:///", ¶ms) + .await + .unwrap(); + + // called after construction + assert!(wrapper.called()); + + // hard to compare two trait pointers as the point to vtables + // using the ref count as a proxy to make sure that the store is correctly kept + assert_eq!(Arc::strong_count(&mock_inner_store), 2); + } + + #[tokio::test] + async fn test_local_paths() { + let file_path = TempStdFile::default(); + let mut writer = ObjectStore::create_local_writer(&file_path).await.unwrap(); + writer.write_all(b"LOCAL").await.unwrap(); + Writer::shutdown(&mut writer).await.unwrap(); + + let reader = ObjectStore::open_local(&file_path).await.unwrap(); + let buf = reader.get_range(0..5).await.unwrap(); + assert_eq!(buf.as_ref(), b"LOCAL"); + } + + #[tokio::test] + async fn test_read_one() { + let file_path = TempStdFile::default(); + let mut writer = ObjectStore::create_local_writer(&file_path).await.unwrap(); + writer.write_all(b"LOCAL").await.unwrap(); + Writer::shutdown(&mut writer).await.unwrap(); + + let file_path_os = object_store::path::Path::parse(file_path.to_str().unwrap()).unwrap(); + let obj_store = ObjectStore::local(); + let buf = obj_store.read_one_all(&file_path_os).await.unwrap(); + assert_eq!(buf.as_ref(), b"LOCAL"); + + let buf = obj_store.read_one_range(&file_path_os, 0..5).await.unwrap(); + assert_eq!(buf.as_ref(), b"LOCAL"); + } + + #[tokio::test] + #[cfg(windows)] + async fn test_windows_paths() { + use std::path::Component; + use std::path::Prefix; + use std::path::Prefix::*; + + fn get_path_prefix(path: &StdPath) -> Prefix<'_> { + match path.components().next().unwrap() { + Component::Prefix(prefix_component) => prefix_component.kind(), + _ => panic!(), + } + } + + fn get_drive_letter(prefix: Prefix) -> String { + match prefix { + Disk(bytes) => String::from_utf8(vec![bytes]).unwrap(), + _ => panic!(), + } + } + + let tmp_path = TempStdFile::default(); + let prefix = get_path_prefix(&tmp_path); + let drive_letter = get_drive_letter(prefix); + + write_to_file( + &(format!("{drive_letter}:/test_folder/test.lance") + "/test_file"), + "WINDOWS", + ) + .unwrap(); + + for uri in &[ + format!("{drive_letter}:/test_folder/test.lance"), + format!("{drive_letter}:\\test_folder\\test.lance"), + ] { + let (store, base) = ObjectStore::from_uri(uri).await.unwrap(); + let contents = read_from_store(store.as_ref(), &base.clone().join("test_file")) + .await + .unwrap(); + assert_eq!(contents, "WINDOWS"); + } + } + + #[tokio::test] + async fn test_cross_filesystem_copy() { + // Create two temporary directories that simulate different filesystems + let source_dir = TempStdDir::default(); + let dest_dir = TempStdDir::default(); + + // Create a test file in the source directory + let source_file_name = "test_file.txt"; + let source_file = source_dir.join(source_file_name); + std::fs::write(&source_file, b"test content").unwrap(); + + // Create ObjectStore for local filesystem + let (store, base_path) = ObjectStore::from_uri(source_dir.to_str().unwrap()) + .await + .unwrap(); + + // Create paths relative to the ObjectStore base + let from_path = base_path.clone().join(source_file_name); + + // Use object_store::Path::parse for the destination + let dest_file = dest_dir.join("copied_file.txt"); + let dest_str = dest_file.to_str().unwrap(); + let to_path = object_store::path::Path::parse(dest_str).unwrap(); + + // Perform the copy operation + store.copy(&from_path, &to_path).await.unwrap(); + + // Verify the file was copied correctly + assert!(dest_file.exists()); + let copied_content = std::fs::read(&dest_file).unwrap(); + assert_eq!(copied_content, b"test content"); + } + + #[tokio::test] + async fn test_copy_creates_parent_directories() { + let source_dir = TempStdDir::default(); + let dest_dir = TempStdDir::default(); + + // Create a test file in the source directory + let source_file_name = "test_file.txt"; + let source_file = source_dir.join(source_file_name); + std::fs::write(&source_file, b"test content").unwrap(); + + // Create ObjectStore for local filesystem + let (store, base_path) = ObjectStore::from_uri(source_dir.to_str().unwrap()) + .await + .unwrap(); + + // Create paths + let from_path = base_path.clone().join(source_file_name); + + // Create destination with nested directories that don't exist yet + let dest_file = dest_dir.join("nested").join("dirs").join("copied_file.txt"); + let dest_str = dest_file.to_str().unwrap(); + let to_path = object_store::path::Path::parse(dest_str).unwrap(); + + // Perform the copy operation - should create parent directories + store.copy(&from_path, &to_path).await.unwrap(); + + // Verify the file was copied correctly and directories were created + assert!(dest_file.exists()); + assert!(dest_file.parent().unwrap().exists()); + let copied_content = std::fs::read(&dest_file).unwrap(); + assert_eq!(copied_content, b"test content"); + } + + /// Inner store that forwards everything to `InMemory` except single-shot + /// server-side copy (`copy_opts`), which always fails. This lets a test + /// prove that `ObjectStore::copy` fell back to a streaming multipart copy + /// for an oversized source rather than issuing a single `CopyObject`. + #[derive(Debug)] + struct CopyFailingStore { + inner: InMemory, + } + + impl Display for CopyFailingStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CopyFailingStore") + } + } + + #[async_trait] + impl OSObjectStore for CopyFailingStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.inner.put_opts(location, bytes, opts).await + } + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.inner.get_opts(location, options).await + } + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.inner.get_ranges(location, ranges).await + } + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.inner.list(prefix) + } + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.inner.list_with_offset(prefix, offset) + } + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> { + Err(object_store::Error::Generic { + store: "CopyFailingStore", + source: "single-shot copy disabled in test".into(), + }) + } + } + + #[tokio::test] + async fn test_copy_streams_objects_larger_than_threshold() { + // memory:// is non-local but isn't an S3/GCS scheme, so copy() wouldn't + // enable the fallback on its own. Drive copy_impl directly with + // multipart_copy_fallback = true to exercise the streaming path. The + // inner store rejects any single-shot copy, so a successful copy can only + // have gone through the streaming branch. + let mut store = ObjectStore::memory(); + store.inner = Arc::new(CopyFailingStore { + inner: InMemory::new(), + }); + + let from = Path::from("source.bin"); + let contents = b"streaming multipart copy payload well past the tiny threshold"; + store.put(&from, contents).await.unwrap(); + + // Source size (61 bytes) exceeds the threshold -> must stream via a + // multipart write rather than a single-shot server-side copy. + let streamed = Path::from("streamed.bin"); + store.copy_impl(&from, &streamed, true, 8).await.unwrap(); + let copied = store.read_one_all(&streamed).await.unwrap(); + assert_eq!(copied.as_ref(), contents.as_slice()); + + // Source size below the threshold -> single-shot copy, which the inner + // store rejects, confirming that the streaming branch (not native copy) + // is what made the first copy succeed. + let native = Path::from("native.bin"); + assert!( + store + .copy_impl(&from, &native, true, u64::MAX) + .await + .is_err() + ); + } + + #[test] + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + fn test_client_options_extracts_headers() { + let opts = StorageOptions(HashMap::from([ + ("headers.x-custom-foo".to_string(), "bar".to_string()), + ("headers.x-ms-version".to_string(), "2023-11-03".to_string()), + ("region".to_string(), "us-west-2".to_string()), + ])); + let client_options = opts.client_options().unwrap(); + + // Verify non-header keys are not consumed as headers by creating + // another StorageOptions with no headers.* keys. + let opts_no_headers = StorageOptions(HashMap::from([( + "region".to_string(), + "us-west-2".to_string(), + )])); + opts_no_headers.client_options().unwrap(); + + // Smoke test: the client_options with headers should be usable + // in a builder (we can't inspect the headers directly, but building + // should not fail). + #[cfg(feature = "gcp")] + { + use object_store::gcp::GoogleCloudStorageBuilder; + let _builder = GoogleCloudStorageBuilder::new() + .with_client_options(client_options) + .with_url("gs://test-bucket"); + } + } + + #[test] + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + fn test_client_options_rejects_invalid_header_name() { + let opts = StorageOptions(HashMap::from([( + "headers.bad header".to_string(), + "value".to_string(), + )])); + let err = opts.client_options().unwrap_err(); + assert!(err.to_string().contains("invalid header name")); + } + + #[test] + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + fn test_client_options_rejects_invalid_header_value() { + let opts = StorageOptions(HashMap::from([( + "headers.x-good-name".to_string(), + "bad\x01value".to_string(), + )])); + let err = opts.client_options().unwrap_err(); + assert!(err.to_string().contains("invalid header value")); + } + + #[test] + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + fn test_client_options_empty_when_no_header_keys() { + let opts = StorageOptions(HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("access_key_id".to_string(), "AKID".to_string()), + ])); + opts.client_options().unwrap(); + } +} diff --git a/vendor/lance-io/src/object_store/dynamic_credentials.rs b/vendor/lance-io/src/object_store/dynamic_credentials.rs new file mode 100644 index 000000000..8a39e568d --- /dev/null +++ b/vendor/lance-io/src/object_store/dynamic_credentials.rs @@ -0,0 +1,411 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::fmt; +use std::marker::PhantomData; +use std::sync::Arc; + +use async_trait::async_trait; +use lance_core::error::{Error, Result}; +use object_store::{CredentialProvider, Result as ObjectStoreResult}; + +use crate::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; + +#[cfg(feature = "aws")] +use object_store::aws::AwsCredential as ObjectStoreAwsCredential; +#[cfg(feature = "azure")] +use object_store::azure::{AzureAccessKey, AzureCredential}; +#[cfg(feature = "gcp")] +use object_store::gcp::GcpCredential; + +/// Raw dynamic storage options fetched from a credential-vending source. +/// +/// Callers must convert this bag into a cloud-specific credential type via +/// `TryFrom`. +#[derive(Clone)] +pub struct DynamicCredentials(pub HashMap); + +impl fmt::Debug for DynamicCredentials { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("DynamicCredentials") + .field(&format_args!("[{} keys redacted]", self.0.len())) + .finish() + } +} + +#[derive(Clone)] +pub struct NamespaceCredentialsProvider { + accessor: Arc, + _credential: PhantomData, +} + +impl fmt::Debug for NamespaceCredentialsProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NamespaceCredentialsProvider") + .field("accessor", &self.accessor) + .field("credential_type", &std::any::type_name::()) + .finish() + } +} + +impl NamespaceCredentialsProvider { + pub fn new(accessor: Arc) -> Self { + Self { + accessor, + _credential: PhantomData, + } + } + + pub fn from_provider(provider: Arc) -> Self { + Self::new(Arc::new(StorageOptionsAccessor::with_provider(provider))) + } + + pub fn from_provider_with_initial( + provider: Arc, + initial_options: HashMap, + ) -> Self { + Self::new(Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial_options, + provider, + ))) + } +} + +/// Build a dynamic credential provider for any cloud type, returning `None` +/// if the accessor has no provider or the provider options are incompatible with `T`. +pub async fn build_dynamic_credential_provider( + accessor: Option>, +) -> Result>>> +where + T: TryFrom + fmt::Debug + Send + Sync + 'static, +{ + let Some(accessor) = accessor.filter(|a| a.has_provider()) else { + return Ok(None); + }; + + let compatible = if let Some(initial) = accessor.initial_storage_options() + && T::try_from(DynamicCredentials(initial.clone())).is_ok() + { + true + } else { + let fetched = accessor.refresh_storage_options().await?.0; + T::try_from(DynamicCredentials(fetched)).is_ok() + }; + + if !compatible { + return Ok(None); + } + + Ok(Some( + Arc::new(NamespaceCredentialsProvider::::new(accessor)) + as Arc>, + )) +} + +fn map_credential_error(error: Error) -> object_store::Error { + object_store::Error::Generic { + store: "NamespaceCredentialsProvider", + source: Box::new(error), + } +} + +#[async_trait] +impl CredentialProvider for NamespaceCredentialsProvider +where + T: TryFrom + fmt::Debug + Send + Sync + 'static, +{ + type Credential = T; + + async fn get_credential(&self) -> ObjectStoreResult> { + let storage_options = self + .accessor + .get_storage_options() + .await + .map_err(map_credential_error)?; + + let credential = match T::try_from(DynamicCredentials(storage_options.0)) { + Ok(credential) => credential, + Err(_) if self.accessor.has_provider() => { + let storage_options = self + .accessor + .refresh_storage_options() + .await + .map_err(map_credential_error)?; + T::try_from(DynamicCredentials(storage_options.0)).map_err(map_credential_error)? + } + Err(error) => return Err(map_credential_error(error)), + }; + + Ok(Arc::new(credential)) + } +} + +fn missing_dynamic_credential(kind: &str) -> Error { + Error::invalid_input(format!( + "Missing required {kind} credential fields in dynamic storage options" + )) +} + +#[cfg(feature = "azure")] +fn split_azure_sas(sas: &str) -> Result> { + let pairs = url::form_urlencoded::parse(sas.trim_start_matches('?').as_bytes()) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect::>(); + + if pairs.is_empty() { + return Err(Error::invalid_input( + "Azure SAS token is empty or invalid in dynamic storage options", + )); + } + + Ok(pairs) +} + +#[cfg(feature = "aws")] +impl TryFrom for ObjectStoreAwsCredential { + type Error = Error; + + fn try_from(credentials: DynamicCredentials) -> Result { + let key_id = credentials + .0 + .get("aws_access_key_id") + .or_else(|| credentials.0.get("access_key_id")) + .cloned(); + let secret_key = credentials + .0 + .get("aws_secret_access_key") + .or_else(|| credentials.0.get("secret_access_key")) + .cloned(); + let token = credentials + .0 + .get("aws_session_token") + .or_else(|| credentials.0.get("aws_token")) + .or_else(|| credentials.0.get("aws_security_token")) + .or_else(|| credentials.0.get("session_token")) + .or_else(|| credentials.0.get("token")) + .cloned(); + + match (key_id, secret_key) { + (Some(key_id), Some(secret_key)) => Ok(Self { + key_id, + secret_key, + token, + }), + _ => Err(missing_dynamic_credential("AWS")), + } + } +} + +#[cfg(feature = "azure")] +impl TryFrom for AzureCredential { + type Error = Error; + + fn try_from(credentials: DynamicCredentials) -> Result { + if let Some(sas) = credentials + .0 + .get("azure_storage_sas_token") + .or_else(|| credentials.0.get("azure_storage_sas_key")) + .or_else(|| credentials.0.get("sas_token")) + .or_else(|| credentials.0.get("sas_key")) + { + return Ok(Self::SASToken(split_azure_sas(sas)?)); + } + + if let Some(token) = credentials + .0 + .get("azure_storage_token") + .or_else(|| credentials.0.get("bearer_token")) + .or_else(|| credentials.0.get("token")) + { + return Ok(Self::BearerToken(token.clone())); + } + + if let Some(access_key) = credentials + .0 + .get("azure_storage_account_key") + .or_else(|| credentials.0.get("azure_storage_access_key")) + .or_else(|| credentials.0.get("azure_storage_master_key")) + .or_else(|| credentials.0.get("access_key")) + .or_else(|| credentials.0.get("master_key")) + .or_else(|| credentials.0.get("account_key")) + { + return Ok(Self::AccessKey( + AzureAccessKey::try_new(access_key).map_err(|source| { + Error::invalid_input(format!("Invalid Azure access key: {source}")) + })?, + )); + } + + Err(missing_dynamic_credential("Azure")) + } +} + +#[cfg(feature = "gcp")] +impl TryFrom for GcpCredential { + type Error = Error; + + fn try_from(credentials: DynamicCredentials) -> Result { + let bearer = credentials + .0 + .get("google_storage_token") + .cloned() + .ok_or_else(|| missing_dynamic_credential("GCP"))?; + + Ok(Self { bearer }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use super::*; + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + + #[cfg(feature = "aws")] + #[tokio::test] + async fn test_dynamic_aws_credentials() { + let provider = Arc::new(StaticMockStorageOptionsProvider { + options: HashMap::from([ + ("aws_access_key_id".to_string(), "AKID".to_string()), + ("aws_secret_access_key".to_string(), "SECRET".to_string()), + ("aws_session_token".to_string(), "TOKEN".to_string()), + ]), + }); + + let credentials = + NamespaceCredentialsProvider::::from_provider(provider) + .get_credential() + .await + .expect("aws credentials should convert"); + + assert_eq!(credentials.key_id, "AKID"); + assert_eq!(credentials.secret_key, "SECRET"); + assert_eq!(credentials.token.as_deref(), Some("TOKEN")); + } + + #[cfg(feature = "aws")] + #[tokio::test] + async fn test_dynamic_aws_credentials_aws_token_alias() { + let provider = Arc::new(StaticMockStorageOptionsProvider { + options: HashMap::from([ + ("aws_access_key_id".to_string(), "AKID".to_string()), + ("aws_secret_access_key".to_string(), "SECRET".to_string()), + ("aws_token".to_string(), "TOKEN".to_string()), + ]), + }); + + let credentials = + NamespaceCredentialsProvider::::from_provider(provider) + .get_credential() + .await + .expect("aws credentials should convert"); + + assert_eq!(credentials.token.as_deref(), Some("TOKEN")); + } + + #[cfg(feature = "aws")] + #[tokio::test] + async fn test_dynamic_credentials_fetch_provider_when_initial_has_metadata_only() { + let provider = Arc::new(StaticMockStorageOptionsProvider { + options: HashMap::from([ + ("aws_access_key_id".to_string(), "AKID".to_string()), + ("aws_secret_access_key".to_string(), "SECRET".to_string()), + ]), + }); + let accessor = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + HashMap::from([("region".to_string(), "us-west-2".to_string())]), + provider, + )); + + let credentials = + build_dynamic_credential_provider::(Some(accessor)) + .await + .expect("dynamic credential provider should build") + .expect("provider should be returned") + .get_credential() + .await + .expect("provider-vended aws credentials should convert"); + + assert_eq!(credentials.key_id, "AKID"); + assert_eq!(credentials.secret_key, "SECRET"); + } + + #[cfg(feature = "azure")] + #[tokio::test] + async fn test_dynamic_azure_credentials() { + let provider = Arc::new(StaticMockStorageOptionsProvider { + options: HashMap::from([( + "azure_storage_sas_token".to_string(), + "?sv=2022-11-02&sp=rl&sig=test".to_string(), + )]), + }); + + let credentials = NamespaceCredentialsProvider::::from_provider(provider) + .get_credential() + .await + .expect("azure credentials should convert"); + + match credentials.as_ref() { + AzureCredential::SASToken(pairs) => { + assert!( + pairs + .iter() + .any(|(key, value)| key == "sv" && value == "2022-11-02") + ); + assert!( + pairs + .iter() + .any(|(key, value)| key == "sig" && value == "test") + ); + } + other => panic!("expected SAS token, got {other:?}"), + } + } + + #[cfg(feature = "azure")] + #[tokio::test] + async fn test_dynamic_azure_credentials_short_sas_aliases() { + for key in ["sas_token", "sas_key"] { + let provider = Arc::new(StaticMockStorageOptionsProvider { + options: HashMap::from([( + key.to_string(), + "?sv=2022-11-02&sp=rl&sig=short".to_string(), + )]), + }); + + let credentials = + NamespaceCredentialsProvider::::from_provider(provider) + .get_credential() + .await + .unwrap_or_else(|_| panic!("azure credentials should convert for key '{key}'")); + + match credentials.as_ref() { + AzureCredential::SASToken(pairs) => { + assert!( + pairs.iter().any(|(k, v)| k == "sig" && v == "short"), + "SAS token from key '{key}' should contain sig=short" + ); + } + other => panic!("expected SAS token for key '{key}', got {other:?}"), + } + } + } + + #[cfg(feature = "gcp")] + #[tokio::test] + async fn test_dynamic_gcp_credentials() { + let provider = Arc::new(StaticMockStorageOptionsProvider { + options: HashMap::from([("google_storage_token".to_string(), "gcp-token".to_string())]), + }); + + let credentials = NamespaceCredentialsProvider::::from_provider(provider) + .get_credential() + .await + .expect("gcp credentials should convert"); + + assert_eq!(credentials.bearer, "gcp-token"); + } +} diff --git a/vendor/lance-io/src/object_store/dynamic_opendal.rs b/vendor/lance-io/src/object_store/dynamic_opendal.rs new file mode 100644 index 000000000..9eafa2569 --- /dev/null +++ b/vendor/lance-io/src/object_store/dynamic_opendal.rs @@ -0,0 +1,368 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::fmt; +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt, stream, stream::BoxStream}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore as OSObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + RenameOptions, +}; +use object_store_opendal::OpendalStore; +use tokio::sync::RwLock; + +use crate::object_store::StorageOptionsAccessor; +use lance_core::Result; + +type NormalizeConfigFn = fn(&HashMap) -> Result>; +type BuildStoreFn = fn(HashMap) -> Result; +type FilterDynamicOptionsFn = fn(&HashMap) -> HashMap; + +#[derive(Debug, Clone)] +struct CachedOpenDalStore { + config: HashMap, + store: Arc, +} + +#[derive(Clone)] +pub(in crate::object_store) struct DynamicOpenDalStore { + name: Arc, + base_options: Arc>, + accessor: Arc, + normalize_config: NormalizeConfigFn, + build_store: BuildStoreFn, + filter_dynamic_options: Option, + atomic_key_groups: Vec>, + protected_keys: Vec<&'static str>, + cache: Arc>>, +} + +impl fmt::Debug for DynamicOpenDalStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("DynamicOpenDalStore") + .field("name", &self.name) + .field("accessor", &self.accessor) + .finish() + } +} + +impl DynamicOpenDalStore { + pub(in crate::object_store) fn new( + name: impl Into>, + base_options: HashMap, + accessor: Arc, + normalize_config: NormalizeConfigFn, + build_store: BuildStoreFn, + ) -> Self { + Self { + name: name.into(), + base_options: Arc::new(base_options), + accessor, + normalize_config, + build_store, + filter_dynamic_options: None, + atomic_key_groups: Vec::new(), + protected_keys: Vec::new(), + cache: Arc::new(RwLock::new(None)), + } + } + + #[allow(dead_code)] + pub(in crate::object_store) fn with_protected_keys( + mut self, + keys: impl IntoIterator, + ) -> Self { + self.protected_keys = keys.into_iter().collect(); + self + } + + /// Restrict provider-vended updates before they are merged into the fixed store config. + pub(in crate::object_store) fn with_dynamic_options_filter( + mut self, + filter: FilterDynamicOptionsFn, + ) -> Self { + self.filter_dynamic_options = Some(filter); + self + } + + /// Treat a set of related options as one authority when provider values are present. + pub(in crate::object_store) fn with_atomic_key_group( + mut self, + keys: impl IntoIterator, + ) -> Self { + self.atomic_key_groups.push(keys.into_iter().collect()); + self + } + + fn merge_options( + &self, + mut dynamic_options: HashMap, + ) -> HashMap { + if let Some(filter) = self.filter_dynamic_options { + dynamic_options = filter(&dynamic_options); + } + for key in &self.protected_keys { + dynamic_options.remove(*key); + } + let mut merged = self.base_options.as_ref().clone(); + for group in &self.atomic_key_groups { + if group.iter().any(|key| dynamic_options.contains_key(*key)) { + for key in group { + merged.remove(*key); + } + } + } + merged.extend(dynamic_options); + merged + } + + pub(in crate::object_store) async fn current_store(&self) -> Result> { + let merged_options = self.merge_options(self.accessor.get_storage_options().await?.0); + let config = (self.normalize_config)(&merged_options)?; + + // Cache reuse depends on exact normalized config equality. Providers + // should return stable, canonicalized values for semantically identical + // configurations to avoid unnecessary store rebuilds. + { + let cache = self.cache.read().await; + if let Some(cached) = cache.as_ref() + && cached.config == config + { + return Ok(cached.store.clone()); + } + } + + let store = Arc::new((self.build_store)(config.clone())?); + let mut cache = self.cache.write().await; + if let Some(cached) = cache.as_ref() + && cached.config == config + { + return Ok(cached.store.clone()); + } + + *cache = Some(CachedOpenDalStore { + config, + store: store.clone(), + }); + Ok(store) + } + + fn map_store_error(&self, error: lance_core::Error) -> object_store::Error { + object_store::Error::Generic { + store: "DynamicOpenDalStore", + source: Box::new(error), + } + } +} + +impl fmt::Display for DynamicOpenDalStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "DynamicOpenDalStore({})", self.name) + } +} + +#[async_trait::async_trait] +impl OSObjectStore for DynamicOpenDalStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .put_opts(location, payload, opts) + .await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> object_store::Result> { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .put_multipart_opts(location, opts) + .await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .get_opts(location, options) + .await + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> object_store::Result> { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .get_ranges(location, ranges) + .await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + let this = self.clone(); + stream::once(async move { + let store = this + .current_store() + .await + .map_err(|e| this.map_store_error(e))?; + Ok::<_, object_store::Error>((store, locations)) + }) + .map_ok(|(store, locations)| store.delete_stream(locations)) + .try_flatten() + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result> { + let prefix = prefix.cloned(); + let this = self.clone(); + stream::once(async move { + this.current_store() + .await + .map_err(|e| this.map_store_error(e)) + }) + .map_ok(move |store| store.list(prefix.as_ref())) + .try_flatten() + .boxed() + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .list_with_delimiter(prefix) + .await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + opts: CopyOptions, + ) -> object_store::Result<()> { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .copy_opts(from, to, opts) + .await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + opts: RenameOptions, + ) -> object_store::Result<()> { + self.current_store() + .await + .map_err(|e| self.map_store_error(e))? + .rename_opts(from, to, opts) + .await + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use opendal::{Operator, services::Memory}; + + use super::*; + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + + #[tokio::test] + async fn test_dynamic_store_caches_by_normalized_config() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([("token".to_string(), "value".to_string())]), + }, + ))); + + let store = DynamicOpenDalStore::new( + "memory", + HashMap::new(), + accessor, + |options| Ok(options.clone()), + |_| { + let operator = Operator::new(Memory::default()).map_err(|e| { + lance_core::Error::invalid_input(format!( + "Failed to create memory operator: {e:?}" + )) + })?; + Ok(OpendalStore::new(operator)) + }, + ); + + let first = store + .current_store() + .await + .expect("first store should build"); + let second = store + .current_store() + .await + .expect("second store should reuse cache"); + + assert!(Arc::ptr_eq(&first, &second)); + } + + #[test] + fn test_merge_options_preserves_protected_base_keys() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::new(), + }, + ))); + let store = DynamicOpenDalStore::new( + "memory", + HashMap::from([ + ("bucket".to_string(), "url-bucket".to_string()), + ("root".to_string(), "/".to_string()), + ("token".to_string(), "base-token".to_string()), + ]), + accessor, + |options| Ok(options.clone()), + |_| { + let operator = Operator::new(Memory::default()).map_err(|e| { + lance_core::Error::invalid_input(format!( + "Failed to create memory operator: {e:?}" + )) + })?; + Ok(OpendalStore::new(operator)) + }, + ) + .with_protected_keys(["bucket", "root"]); + + let merged = store.merge_options(HashMap::from([ + ("bucket".to_string(), "provider-bucket".to_string()), + ("root".to_string(), "/provider-root".to_string()), + ("token".to_string(), "provider-token".to_string()), + ])); + + assert_eq!(merged.get("bucket").unwrap(), "url-bucket"); + assert_eq!(merged.get("root").unwrap(), "/"); + assert_eq!(merged.get("token").unwrap(), "provider-token"); + } +} diff --git a/vendor/lance-io/src/object_store/list_retry.rs b/vendor/lance-io/src/object_store/list_retry.rs new file mode 100644 index 000000000..f1ee65035 --- /dev/null +++ b/vendor/lance-io/src/object_store/list_retry.rs @@ -0,0 +1,403 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{future::Future, pin::Pin, sync::Arc, task::Poll, time::Duration}; + +use futures::stream::BoxStream; +use futures::{Stream, StreamExt}; +use object_store::{ObjectMeta, ObjectStore, path::Path}; +use rand::Rng; +use tokio::time::Sleep; + +const DEFAULT_BASE_RETRY_DELAY: Duration = Duration::from_millis(100); +const DEFAULT_MAX_RETRY_DELAY: Duration = Duration::from_secs(5); + +/// A stream that does outer retries on list operations. +/// +/// This is to handle request responses that ObjectStore doesn't handle, such as +/// the error `error decoding response body` from queries to GCS. +pub struct ListRetryStream { + object_store: Arc, + current_stream: BoxStream<'static, object_store::Result>, + prefix: Option, + last_successful_key: Option, + max_retries: usize, + current_retries: usize, + retry_sleep: Option>>, + base_retry_delay: Duration, + max_retry_delay: Duration, +} + +impl ListRetryStream { + pub fn new( + object_store: Arc, + prefix: Option, + max_retries: usize, + ) -> Self { + let current_stream = object_store.list(prefix.as_ref()); + Self { + object_store, + current_stream, + prefix, + last_successful_key: None, + max_retries, + current_retries: 0, + retry_sleep: None, + base_retry_delay: DEFAULT_BASE_RETRY_DELAY, + max_retry_delay: DEFAULT_MAX_RETRY_DELAY, + } + } + + #[cfg(test)] + fn new_with_backoff( + object_store: Arc, + prefix: Option, + max_retries: usize, + base_retry_delay: Duration, + max_retry_delay: Duration, + ) -> Self { + let current_stream = object_store.list(prefix.as_ref()); + Self { + object_store, + current_stream, + prefix, + last_successful_key: None, + max_retries, + current_retries: 0, + retry_sleep: None, + base_retry_delay, + max_retry_delay, + } + } + + fn is_retryable(error: &object_store::Error) -> bool { + !matches!( + error, + object_store::Error::NotFound { .. } + | object_store::Error::InvalidPath { .. } + | object_store::Error::NotSupported { .. } + | object_store::Error::NotImplemented { .. } + ) + } + + fn retry_delay(&self) -> Duration { + let exponent = self.current_retries.saturating_sub(1).min(16) as u32; + let base_ms = self.base_retry_delay.as_millis().max(1); + let max_ms = self.max_retry_delay.as_millis().max(base_ms); + let cap_ms = base_ms.saturating_mul(1_u128 << exponent).min(max_ms); + let min_ms = (cap_ms / 2).max(1); + let delay_ms = if cap_ms > min_ms { + rand::rng().random_range(min_ms..=cap_ms) + } else { + cap_ms + }; + Duration::from_millis(delay_ms.min(u64::MAX as u128) as u64) + } + + fn recreate_stream(&mut self) { + self.current_stream = if let Some(offset) = self.last_successful_key.clone() { + self.object_store + .list_with_offset(self.prefix.as_ref(), &offset) + } else { + self.object_store.list(self.prefix.as_ref()) + }; + } +} + +impl Stream for ListRetryStream { + type Item = Result; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + let this = self.get_mut(); + loop { + if let Some(sleep) = this.retry_sleep.as_mut() { + match sleep.as_mut().poll(cx) { + Poll::Ready(()) => { + this.retry_sleep = None; + this.recreate_stream(); + } + Poll::Pending => return Poll::Pending, + } + } + + match this.current_stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(meta))) => { + this.last_successful_key = Some(meta.location.clone()); + return Poll::Ready(Some(Ok(meta))); + } + Poll::Ready(None) => { + // If the stream is done, return None + return Poll::Ready(None); + } + Poll::Ready(Some(Err(error))) if Self::is_retryable(&error) => { + if this.current_retries < this.max_retries { + this.current_retries += 1; + this.retry_sleep = Some(Box::pin(tokio::time::sleep(this.retry_delay()))); + + continue; + } else { + return Poll::Ready(Some(Err(error))); + } + } + Poll::Ready(Some(Err(error))) => { + return Poll::Ready(Some(Err(error))); + } + Poll::Pending => { + return Poll::Pending; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::VecDeque; + use std::fmt::{Debug, Display, Formatter}; + use std::ops::Range; + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Instant; + + use async_trait::async_trait; + use bytes::Bytes; + use futures::stream; + use object_store::memory::InMemory; + use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, PutMultipartOptions, + PutOptions, PutPayload, PutResult, Result as OSResult, + }; + + fn assert_send() {} + + #[test] + fn test_list_retry_stream_send() { + // Ensure that ListRetryStream is Send + assert_send::(); + } + + fn object_meta(path: &str) -> ObjectMeta { + ObjectMeta { + location: Path::from(path), + last_modified: chrono::Utc::now(), + size: 1, + e_tag: None, + version: None, + } + } + + fn retryable_error() -> object_store::Error { + object_store::Error::Generic { + store: "scripted", + source: "retryable list error".into(), + } + } + + fn not_found_error() -> object_store::Error { + object_store::Error::NotFound { + path: "missing".to_string(), + source: "missing".into(), + } + } + + struct ScriptedListStore { + inner: InMemory, + list_streams: Mutex>>>, + offset_streams: Mutex>>>, + list_calls: AtomicUsize, + offset_calls: AtomicUsize, + last_offset: Mutex>, + } + + impl ScriptedListStore { + fn new( + list_streams: Vec>>, + offset_streams: Vec>>, + ) -> Self { + Self { + inner: InMemory::new(), + list_streams: Mutex::new(list_streams.into()), + offset_streams: Mutex::new(offset_streams.into()), + list_calls: AtomicUsize::new(0), + offset_calls: AtomicUsize::new(0), + last_offset: Mutex::new(None), + } + } + + fn list_calls(&self) -> usize { + self.list_calls.load(Ordering::SeqCst) + } + + fn offset_calls(&self) -> usize { + self.offset_calls.load(Ordering::SeqCst) + } + + fn last_offset(&self) -> Option { + self.last_offset.lock().unwrap().clone() + } + } + + impl Display for ScriptedListStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "ScriptedListStore") + } + } + + impl Debug for ScriptedListStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ScriptedListStore").finish() + } + } + + #[async_trait] + impl ObjectStore for ScriptedListStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.inner.put_opts(location, bytes, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.inner.get_opts(location, options).await + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + + fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.list_calls.fetch_add(1, Ordering::SeqCst); + let results = self + .list_streams + .lock() + .unwrap() + .pop_front() + .unwrap_or_default(); + stream::iter(results).boxed() + } + + fn list_with_offset( + &self, + _prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.offset_calls.fetch_add(1, Ordering::SeqCst); + *self.last_offset.lock().unwrap() = Some(offset.clone()); + let results = self + .offset_streams + .lock() + .unwrap() + .pop_front() + .unwrap_or_default(); + stream::iter(results).boxed() + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.inner.copy_opts(from, to, opts).await + } + } + + #[tokio::test] + async fn test_list_retry_stream_retries_after_backoff() { + let store = Arc::new(ScriptedListStore::new( + vec![ + vec![Err(retryable_error())], + vec![Ok(object_meta("prefix/file"))], + ], + vec![], + )); + let stream = ListRetryStream::new_with_backoff( + store.clone(), + Some(Path::from("prefix")), + 1, + Duration::from_millis(20), + Duration::from_millis(20), + ); + + let start = Instant::now(); + let items = stream.collect::>().await; + + assert_eq!(items.len(), 1); + assert!(items[0].is_ok()); + assert_eq!(store.list_calls(), 2); + assert!( + start.elapsed() >= Duration::from_millis(10), + "retry should wait before recreating the list stream" + ); + } + + #[tokio::test] + async fn test_list_retry_stream_resumes_after_last_successful_key() { + let store = Arc::new(ScriptedListStore::new( + vec![vec![Ok(object_meta("prefix/a")), Err(retryable_error())]], + vec![vec![Ok(object_meta("prefix/b"))]], + )); + let stream = ListRetryStream::new_with_backoff( + store.clone(), + Some(Path::from("prefix")), + 1, + Duration::from_millis(1), + Duration::from_millis(1), + ); + + let items = stream.collect::>().await; + + assert_eq!(items.len(), 2); + assert_eq!(items[0].as_ref().unwrap().location, Path::from("prefix/a")); + assert_eq!(items[1].as_ref().unwrap().location, Path::from("prefix/b")); + assert_eq!(store.list_calls(), 1); + assert_eq!(store.offset_calls(), 1); + assert_eq!(store.last_offset(), Some(Path::from("prefix/a"))); + } + + #[tokio::test] + async fn test_list_retry_stream_non_retryable_errors_return_immediately() { + let store = Arc::new(ScriptedListStore::new( + vec![vec![Err(not_found_error())]], + vec![], + )); + let stream = ListRetryStream::new_with_backoff( + store.clone(), + Some(Path::from("prefix")), + 5, + Duration::from_millis(1), + Duration::from_millis(1), + ); + + let items = stream.collect::>().await; + + assert_eq!(items.len(), 1); + assert!(matches!( + items.into_iter().next().unwrap(), + Err(object_store::Error::NotFound { .. }) + )); + assert_eq!(store.list_calls(), 1); + assert_eq!(store.offset_calls(), 0); + } +} diff --git a/vendor/lance-io/src/object_store/metrics.rs b/vendor/lance-io/src/object_store/metrics.rs new file mode 100644 index 000000000..de00c82ff --- /dev/null +++ b/vendor/lance-io/src/object_store/metrics.rs @@ -0,0 +1,1821 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Publishes object store metrics via the [`metrics`] crate. +//! +//! Two layers cooperate: +//! +//! * [`MeteredObjectStore`] wraps any [`object_store::ObjectStore`] and records +//! per-operation request counts, transferred bytes, latency, errors, and the +//! number of requests currently in flight. It works for every store +//! regardless of backend. +//! * [`MeteringHttpConnector`] wraps the HTTP client used by the native cloud +//! stores (S3 / GCS / Azure) and records throttle / retryable responses per +//! attempt. Because `object_store`'s retry loop re-issues each request +//! through the [`HttpService`](object_store::client::HttpService), this sees +//! every retried response, which a store-level wrapper cannot observe. +//! +//! The two layers have different coverage: every store gets the request-level +//! metrics from [`MeteredObjectStore`], but only the native cloud stores get +//! the HTTP-level throttle metrics. Opendal-backed stores (tos, oss, etc.) +//! bypass `object_store`'s HTTP client, so there is no place to install the +//! connector for them. +//! +//! Neither layer sees the optimized local reads and writes ([`LocalObjectReader`], +//! [`LocalWriter`], the io_uring readers, and the local `copy` / recursive delete +//! shortcuts), which go straight to the filesystem. Those publish the same +//! request-level metrics themselves through +//! [`IOTracker::begin_io`](crate::utils::tracking_store::IOTracker::begin_io). +//! The two are installed together, so a store either publishes for all of its +//! IO or for none of it. A store built by calling a provider's `new_store` +//! directly, bypassing both `ObjectStore` constructors — as +//! [`ObjectStore::local`](crate::object_store::ObjectStore::local) and +//! [`ObjectStore::memory`](crate::object_store::ObjectStore::memory) do — is in +//! the "none of it" case. +//! +//! [`LocalObjectReader`]: crate::local::LocalObjectReader +//! [`LocalWriter`]: crate::object_writer::LocalWriter +//! +//! Metrics carry a `base` label identifying the store. Its cardinality is +//! controlled by the `LANCE_OBJECT_STORE_METRICS_LABEL` environment variable +//! ([`BASE_LABEL_ENV_VAR`]): +//! +//! * `scheme` (default) — scheme only, e.g. `s3`; low, bounded cardinality. +//! * `full` — the full store prefix, e.g. `s3$bucket` or `az$container@account`, +//! so multiple buckets on the same cloud can be told apart. +//! * `off` — omit the `base` label entirely. +//! +//! The metric name constants ([`METRIC_REQUESTS`] etc.) and the recording +//! helpers ([`record_request`], [`record_count`], [`record_error`], +//! [`InFlightGuard`]) are public so custom object stores can emit the same +//! metrics. + +use std::ops::Range; +use std::pin::Pin; +use std::sync::{Arc, OnceLock}; +use std::task::{Context, Poll}; +use std::time::Instant; + +use bytes::Bytes; +use futures::stream::BoxStream; +use futures::{FutureExt, Stream, StreamExt}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, + PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult, + UploadPart, +}; + +/// Total number of object store requests, labelled by `operation` and `base`. +pub const METRIC_REQUESTS: &str = "lance_object_store_requests_total"; +/// Total bytes transferred by object store requests, labelled by `operation` and `base`. +pub const METRIC_BYTES: &str = "lance_object_store_request_bytes_total"; +/// Object store request latency in seconds, labelled by `operation` and `base`. +pub const METRIC_DURATION: &str = "lance_object_store_request_duration_seconds"; +/// Total number of failed object store requests, labelled by `operation` and `base`. +pub const METRIC_ERRORS: &str = "lance_object_store_errors_total"; +/// Total number of throttle responses (HTTP 429 / 503) seen at the HTTP layer, +/// labelled by `status` and `base`. Counts every attempt, including retries. +pub const METRIC_THROTTLE: &str = "lance_object_store_throttle_total"; +/// Total number of retryable responses (HTTP 5xx / 429 / 408) seen at the HTTP +/// layer, labelled by `status` and `base`. Counts every attempt, including +/// retries. This is a superset of [`METRIC_THROTTLE`]; 409 (conflict) is +/// deliberately excluded so commit conflicts are not counted as retries. +pub const METRIC_RETRYABLE: &str = "lance_object_store_retryable_responses_total"; +/// Number of object store requests currently in flight, labelled by `operation` +/// and `base`. +pub const METRIC_IN_FLIGHT: &str = "lance_object_store_in_flight_requests"; + +/// Environment variable controlling the cardinality of the `base` label. +pub const BASE_LABEL_ENV_VAR: &str = "LANCE_OBJECT_STORE_METRICS_LABEL"; + +/// Controls how much of a store's identity the `base` label carries, traded off +/// against metric cardinality. Selected via [`BASE_LABEL_ENV_VAR`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BaseLabelMode { + /// Full store prefix, e.g. `s3$bucket` or `az$container@account`. Highest + /// cardinality: one series family per bucket/container. + Full, + /// Scheme only, e.g. `s3`. The default: low, bounded cardinality. + Scheme, + /// Omit the `base` label entirely. + Off, +} + +fn parse_base_label_mode(value: Option<&str>) -> BaseLabelMode { + match value { + Some("full") => BaseLabelMode::Full, + Some("off") | Some("none") => BaseLabelMode::Off, + Some("scheme") | None => BaseLabelMode::Scheme, + Some(other) => { + tracing::warn!( + "Unrecognized {BASE_LABEL_ENV_VAR}={other:?}; \ + expected one of full, scheme, off. Defaulting to scheme." + ); + BaseLabelMode::Scheme + } + } +} + +/// The label mode is read once from the environment and cached for the process. +fn base_label_mode() -> BaseLabelMode { + static MODE: OnceLock = OnceLock::new(); + *MODE.get_or_init(|| parse_base_label_mode(std::env::var(BASE_LABEL_ENV_VAR).ok().as_deref())) +} + +/// Reduce a full store prefix (`scheme$authority`, or just `scheme` for stores +/// without buckets) to the configured `base` label value, or `None` when the +/// label should be omitted. +fn scoped_base(mode: BaseLabelMode, base: &str) -> Option { + match mode { + BaseLabelMode::Full => Some(base.to_owned()), + BaseLabelMode::Scheme => Some(base.split('$').next().unwrap_or(base).to_owned()), + BaseLabelMode::Off => None, + } +} + +/// Build the `operation` (+ optional `base`) label set shared by all +/// store-level metrics, honoring the configured label mode. +fn operation_labels(base: &str, operation: &'static str) -> Vec { + let mut labels = vec![metrics::Label::new("operation", operation)]; + if let Some(base) = scoped_base(base_label_mode(), base) { + labels.push(metrics::Label::new("base", base)); + } + labels +} + +/// Recommended histogram bucket boundaries for [`METRIC_DURATION`], in seconds. +/// +/// Object store requests can take anywhere from a few milliseconds to the +/// client timeout (commonly ~120s), so the boundaries are dense below 10s and +/// keep useful resolution through the timeout band out to 5 minutes. Exporters +/// that aggregate into fixed buckets (e.g. the OpenTelemetry bridge in the +/// Python bindings) use these. +pub const REQUEST_DURATION_BOUNDS: &[f64] = &[ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, // sub-10s + 10.0, 20.0, 30.0, 45.0, 60.0, 90.0, 120.0, 150.0, 180.0, 240.0, 300.0, // 10s–5min +]; + +/// Register descriptions (units and help text) for the object store metrics. +/// +/// This routes through whatever [`metrics::Recorder`] is currently installed, +/// so it must be called *after* the recorder is set. Exporters that build a +/// catalog of available metrics (such as the OpenTelemetry bridge) rely on +/// these descriptions to discover metric names, kinds, and units up front. +pub fn describe_metrics() { + metrics::describe_counter!( + METRIC_REQUESTS, + metrics::Unit::Count, + "Total number of object store requests, by operation and scheme." + ); + metrics::describe_counter!( + METRIC_BYTES, + metrics::Unit::Bytes, + "Total bytes transferred by object store requests, by operation and scheme." + ); + metrics::describe_histogram!( + METRIC_DURATION, + metrics::Unit::Seconds, + "Object store request latency in seconds, by operation and scheme." + ); + metrics::describe_counter!( + METRIC_ERRORS, + metrics::Unit::Count, + "Total number of failed object store requests, by operation and scheme." + ); + metrics::describe_counter!( + METRIC_THROTTLE, + metrics::Unit::Count, + "Total number of throttle responses (HTTP 429 / 503) seen at the HTTP layer, by status and scheme." + ); + metrics::describe_counter!( + METRIC_RETRYABLE, + metrics::Unit::Count, + "Total number of retryable responses (HTTP 5xx / 429 / 408) seen at the HTTP layer, by status and scheme." + ); + metrics::describe_gauge!( + METRIC_IN_FLIGHT, + metrics::Unit::Count, + "Number of object store requests currently in flight, by operation and scheme." + ); +} + +/// Recommended fixed bucket boundaries for the histogram metrics defined here, +/// as `(metric_name, boundaries)` pairs. Exporters that aggregate histograms +/// into fixed buckets read this to configure each histogram. +pub fn histogram_bounds() -> &'static [(&'static str, &'static [f64])] { + &[(METRIC_DURATION, REQUEST_DURATION_BOUNDS)] +} + +/// Record the outcome of a unary request: count, latency, bytes (on success), and errors. +pub fn record_request( + base: &str, + operation: &'static str, + start: Instant, + bytes: u64, + result: &OSResult, +) { + record_outcome(base, operation, start, bytes, result.is_err()); +} + +/// Record count, latency, and either transferred bytes or an error for a +/// completed request. Used both for unary requests and for streamed GETs whose +/// bytes are only known once the body finishes. +pub fn record_outcome( + base: &str, + operation: &'static str, + start: Instant, + bytes: u64, + is_error: bool, +) { + let elapsed = start.elapsed().as_secs_f64(); + let labels = operation_labels(base, operation); + metrics::counter!(METRIC_REQUESTS, labels.clone()).increment(1); + metrics::histogram!(METRIC_DURATION, labels.clone()).record(elapsed); + if is_error { + metrics::counter!(METRIC_ERRORS, labels).increment(1); + } else if bytes > 0 { + metrics::counter!(METRIC_BYTES, labels).increment(bytes); + } +} + +/// Record a single request count without latency, used for streaming operations +/// (list / delete) whose work happens lazily as the stream is polled. +pub fn record_count(base: &str, operation: &'static str) { + metrics::counter!(METRIC_REQUESTS, operation_labels(base, operation)).increment(1); +} + +/// Record a single error for an operation. +pub fn record_error(base: &str, operation: &'static str) { + metrics::counter!(METRIC_ERRORS, operation_labels(base, operation)).increment(1); +} + +/// Raises the in-flight gauge for an operation on creation and lowers it on +/// drop, so the count stays balanced even if the request future or stream is +/// cancelled or dropped before completing. +pub struct InFlightGuard { + labels: Vec, +} + +impl InFlightGuard { + pub fn new(base: &str, operation: &'static str) -> Self { + let labels = operation_labels(base, operation); + metrics::gauge!(METRIC_IN_FLIGHT, labels.clone()).increment(1.0); + Self { labels } + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + metrics::gauge!(METRIC_IN_FLIGHT, self.labels.clone()).decrement(1.0); + } +} + +#[derive(Debug)] +pub struct MeteredObjectStore { + target: Arc, + base: String, +} + +impl std::fmt::Display for MeteredObjectStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "MeteredObjectStore({})", self.target) + } +} + +#[async_trait::async_trait] +#[deny(clippy::missing_trait_methods)] +impl object_store::ObjectStore for MeteredObjectStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + let size = bytes.content_length() as u64; + let _in_flight = InFlightGuard::new(&self.base, "put"); + let start = Instant::now(); + let result = self.target.put_opts(location, bytes, opts).await; + record_request(&self.base, "put", start, size, &result); + result + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + let upload = self.target.put_multipart_opts(location, opts).await?; + Ok(Box::new(MeteredMultipartUpload { + target: upload, + base: self.base.clone(), + })) + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + // `head()` is implemented as a `get_opts` call with `head = true`, so we + // distinguish it here to keep HEAD and GET as separate operations. + let is_head = options.head; + let operation = if is_head { "head" } else { "get" }; + let in_flight = InFlightGuard::new(&self.base, operation); + let start = Instant::now(); + let result = self.target.get_opts(location, options).await; + + // A HEAD transfers only metadata, and errors carry no payload, so both + // are recorded immediately. `get_opts` only resolves once the response + // headers arrive; the body is streamed afterwards, so for a successful + // GET we defer recording until the body has been drained (see below). + if is_head || result.is_err() { + record_request(&self.base, operation, start, 0, &result); + return result; + } + + let result = result.expect("checked to be Ok above"); + Ok(meter_get_result( + result, + self.base.clone(), + start, + in_flight, + )) + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + let _in_flight = InFlightGuard::new(&self.base, "get"); + let start = Instant::now(); + let result = self.target.get_ranges(location, ranges).await; + let bytes = match &result { + Ok(parts) => parts.iter().map(|b| b.len() as u64).sum(), + Err(_) => 0, + }; + record_request(&self.base, "get", start, bytes, &result); + result + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + let base = self.base.clone(); + // Count one logical delete request per call, matching `list`: a single + // `delete_stream` maps to one batched request on stores that support it + // (e.g. S3's `DeleteObjects`), so counting per yielded path would + // over-count. Errors are still recorded per failing path. + record_count(&self.base, "delete"); + let in_flight = InFlightGuard::new(&self.base, "delete"); + self.target + .delete_stream(locations) + .map(move |result| { + // Reference `in_flight` so this `move` closure captures (owns) + // the guard, keeping the gauge raised until the stream is + // dropped (a move closure only captures the variables it uses). + let _in_flight = &in_flight; + if result.is_err() { + record_error(&base, "delete"); + } + result + }) + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + record_count(&self.base, "list"); + meter_list_stream( + self.target.list(prefix), + self.base.clone(), + InFlightGuard::new(&self.base, "list"), + ) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + record_count(&self.base, "list"); + meter_list_stream( + self.target.list_with_offset(prefix, offset), + self.base.clone(), + InFlightGuard::new(&self.base, "list"), + ) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + let _in_flight = InFlightGuard::new(&self.base, "list"); + let start = Instant::now(); + let result = self.target.list_with_delimiter(prefix).await; + record_request(&self.base, "list", start, 0, &result); + result + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + let _in_flight = InFlightGuard::new(&self.base, "copy"); + let start = Instant::now(); + let result = self.target.copy_opts(from, to, opts).await; + record_request(&self.base, "copy", start, 0, &result); + result + } + + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OSResult<()> { + let _in_flight = InFlightGuard::new(&self.base, "rename"); + let start = Instant::now(); + let result = self.target.rename_opts(from, to, opts).await; + record_request(&self.base, "rename", start, 0, &result); + result + } +} + +/// Count errors yielded while draining a list stream. The request itself is +/// counted once when the stream is created (a single LIST may return many items). +fn meter_list_stream( + stream: BoxStream<'static, OSResult>, + base: String, + in_flight: InFlightGuard, +) -> BoxStream<'static, OSResult> { + stream + .map(move |result| { + // Reference `in_flight` so this `move` closure captures (owns) the + // guard: a move closure only captures the variables it uses, and + // holding it here keeps the gauge raised until the stream is dropped. + let _in_flight = &in_flight; + if result.is_err() { + record_error(&base, "list"); + } + result + }) + .boxed() +} + +/// Wrap a successful GET so the request is recorded once its body has been +/// fully read, capturing the true transfer duration and byte count rather than +/// the time-to-first-byte and declared range. For payloads without a body +/// stream (e.g. a local file handle) the request is recorded immediately. +fn meter_get_result( + mut result: GetResult, + base: String, + start: Instant, + in_flight: InFlightGuard, +) -> GetResult { + match result.payload { + GetResultPayload::Stream(stream) => { + result.payload = GetResultPayload::Stream( + MeteredGetStream { + inner: stream, + base, + start, + bytes: 0, + errored: false, + recorded: false, + _in_flight: in_flight, + } + .boxed(), + ); + result + } + // No body stream to observe (e.g. a local file), so record now. + other => { + let bytes = result.range.end - result.range.start; + record_outcome(&base, "get", start, bytes, false); + result.payload = other; + result + } + } +} + +/// Stream wrapper over a GET body that records the request (count, duration, +/// bytes, errors) once the body is fully drained or the stream is dropped. +struct MeteredGetStream { + inner: BoxStream<'static, OSResult>, + base: String, + start: Instant, + bytes: u64, + errored: bool, + recorded: bool, + _in_flight: InFlightGuard, +} + +impl MeteredGetStream { + fn record(&mut self) { + if self.recorded { + return; + } + self.recorded = true; + record_outcome(&self.base, "get", self.start, self.bytes, self.errored); + } +} + +impl Stream for MeteredGetStream { + type Item = OSResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.inner.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(chunk))) => { + self.bytes += chunk.len() as u64; + Poll::Ready(Some(Ok(chunk))) + } + Poll::Ready(Some(Err(e))) => { + self.errored = true; + Poll::Ready(Some(Err(e))) + } + Poll::Ready(None) => { + self.record(); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for MeteredGetStream { + fn drop(&mut self) { + // Records the partial transfer if the body was dropped before it drained. + self.record(); + } +} + +#[derive(Debug)] +struct MeteredMultipartUpload { + target: Box, + base: String, +} + +#[async_trait::async_trait] +impl MultipartUpload for MeteredMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + // Each part upload is a distinct request, recorded under the `put_part` + // operation with the same count / bytes / latency / error set as a + // unary put. + let base = self.base.clone(); + let size = data.content_length() as u64; + let inner = self.target.put_part(data); + async move { + let _in_flight = InFlightGuard::new(&base, "put_part"); + let start = Instant::now(); + let result = inner.await; + record_request(&base, "put_part", start, size, &result); + result + } + .boxed() + } + + async fn complete(&mut self) -> OSResult { + // Completing a multipart upload issues its own request that can throttle + // or fail, so it is metered like any other operation. + let _in_flight = InFlightGuard::new(&self.base, "complete_multipart"); + let start = Instant::now(); + let result = self.target.complete().await; + record_request(&self.base, "complete_multipart", start, 0, &result); + result + } + + async fn abort(&mut self) -> OSResult<()> { + let _in_flight = InFlightGuard::new(&self.base, "abort_multipart"); + let start = Instant::now(); + let result = self.target.abort().await; + record_request(&self.base, "abort_multipart", start, 0, &result); + result + } +} + +pub trait ObjectStoreMetricsExt { + /// Wrap this store so its operations publish metrics under the given `base` label. + fn metered(self, base: String) -> Arc; +} + +impl ObjectStoreMetricsExt for Arc { + fn metered(self, base: String) -> Arc { + Arc::new(MeteredObjectStore { target: self, base }) + } +} + +// --- Layer 2: HTTP-level throttle metrics for native cloud stores --- + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +mod http { + use super::*; + use object_store::client::{ + ClientOptions, HttpClient, HttpConnector, HttpError, HttpRequest, HttpResponse, + HttpService, ReqwestConnector, + }; + + /// An [`HttpConnector`] that records throttle and retryable responses + /// observed by the underlying HTTP client. Install it on the S3 / GCS / + /// Azure builders via `with_http_connector`. + #[derive(Debug)] + pub struct MeteringHttpConnector { + base: String, + inner: ReqwestConnector, + } + + impl MeteringHttpConnector { + pub fn new(base: String) -> Self { + Self { + base, + inner: ReqwestConnector::default(), + } + } + } + + impl HttpConnector for MeteringHttpConnector { + fn connect(&self, options: &ClientOptions) -> object_store::Result { + let client = self.inner.connect(options)?; + Ok(HttpClient::new(MeteringHttpService { + base: self.base.clone(), + inner: client, + })) + } + } + + #[derive(Debug)] + struct MeteringHttpService { + base: String, + inner: HttpClient, + } + + #[async_trait::async_trait] + impl HttpService for MeteringHttpService { + async fn call(&self, req: HttpRequest) -> Result { + let response = self.inner.execute(req).await?; + let status = response.status().as_u16(); + // Each attempt that object_store may retry is recorded with its + // numeric status. Throttles (429 / 503) are a distinct, narrower + // signal than the broader set of retryable responses, so they get + // their own counter. 409 (conflict) is intentionally excluded from + // the retryable set so commit conflicts are not counted as retries. + let is_throttle = status == 429 || status == 503; + let is_retryable = status == 429 || status == 408 || (500..600).contains(&status); + if is_throttle { + metrics::counter!(METRIC_THROTTLE, status_labels(&self.base, status)).increment(1); + } + if is_retryable { + metrics::counter!(METRIC_RETRYABLE, status_labels(&self.base, status)).increment(1); + } + Ok(response) + } + } + + /// Build the `status` (+ optional `base`) label set for HTTP-layer metrics, + /// honoring the configured label mode. + fn status_labels(base: &str, status: u16) -> Vec { + let mut labels = vec![metrics::Label::new("status", status.to_string())]; + if let Some(base) = scoped_base(base_label_mode(), base) { + labels.push(metrics::Label::new("base", base)); + } + labels + } + + #[cfg(test)] + mod tests { + use super::*; + use metrics_util::debugging::{DebugValue, DebuggingRecorder}; + use object_store::client::{HttpRequestBody, HttpResponseBody}; + + /// A mock [`HttpService`] that always responds with a fixed status code. + #[derive(Debug)] + struct StaticStatusService { + status: u16, + } + + #[async_trait::async_trait] + impl HttpService for StaticStatusService { + async fn call(&self, _req: HttpRequest) -> Result { + Ok(::http::Response::builder() + .status(self.status) + .body(HttpResponseBody::from(Bytes::new())) + .unwrap()) + } + } + + fn request() -> HttpRequest { + ::http::Request::builder() + .method("GET") + .uri("http://example.com/obj") + .body(HttpRequestBody::empty()) + .unwrap() + } + + fn metric_count( + metrics: &[(metrics::Key, DebugValue)], + name: &str, + base: &str, + status: &str, + ) -> u64 { + for (key, value) in metrics { + if key.name() != name { + continue; + } + let labels: std::collections::HashMap<&str, &str> = + key.labels().map(|l| (l.key(), l.value())).collect(); + if labels.get("base") == Some(&base) + && labels.get("status") == Some(&status) + && let DebugValue::Counter(v) = value + { + return *v; + } + } + 0 + } + + #[test] + fn test_throttle_and_retryable_responses_counted_by_status() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + // Each attempt that object_store retries flows through `call` + // again; here we simulate that by issuing several responses. + // The base is baked into the connector, so it labels the + // metric. Bases here have no `$`, so they are unaffected by + // the label mode and this test isolates status handling. + for (base, status) in [ + ("s3", 429u16), + ("s3", 503), + ("s3", 503), + ("s3", 500), + ("s3", 408), + ("s3", 409), + ("s3", 200), + ("s3", 404), + ("gs", 429), + ] { + let service = MeteringHttpService { + base: base.into(), + inner: HttpClient::new(StaticStatusService { status }), + }; + service.call(request()).await.unwrap(); + } + }); + }); + + let recorded: Vec<_> = snapshotter + .snapshot() + .into_vec() + .into_iter() + .map(|(ck, _unit, _desc, value)| (ck.key().clone(), value)) + .collect(); + + let throttle = |base, status| metric_count(&recorded, METRIC_THROTTLE, base, status); + let retryable = |base, status| metric_count(&recorded, METRIC_RETRYABLE, base, status); + + // Throttles are only 429 and 503. + assert_eq!(throttle("s3", "429"), 1); + assert_eq!(throttle("s3", "503"), 2); + assert_eq!(throttle("s3", "500"), 0); + assert_eq!(throttle("s3", "408"), 0); + + // Retryable is the broader set: 5xx, 429, 408 (but not 409). + assert_eq!(retryable("s3", "429"), 1); + assert_eq!(retryable("s3", "503"), 2); + assert_eq!(retryable("s3", "500"), 1); + assert_eq!(retryable("s3", "408"), 1); + // 409 conflict is excluded so commit conflicts are not counted as retries. + assert_eq!(retryable("s3", "409"), 0); + + // Success and non-retryable client errors count as neither. + assert_eq!(throttle("s3", "200"), 0); + assert_eq!(retryable("s3", "404"), 0); + + // The base label is taken from the connector, not shared across stores. + assert_eq!(throttle("gs", "429"), 1); + assert_eq!(retryable("gs", "429"), 1); + assert_eq!(throttle("gs", "503"), 0); + } + } +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +pub use http::MeteringHttpConnector; + +#[cfg(test)] +mod tests { + use super::*; + + use lance_core::utils::tempfile::TempStdDir; + use metrics_util::debugging::{DebugValue, DebuggingRecorder, Snapshotter}; + use object_store::memory::InMemory; + use object_store::{ObjectStoreExt, PutPayload}; + use tokio::io::AsyncWriteExt; + use url::Url; + + use crate::object_store::ObjectStore as LanceObjectStore; + use crate::traits::Writer; + + fn payload(data: &[u8]) -> PutPayload { + PutPayload::from_bytes(Bytes::copy_from_slice(data)) + } + + fn metered_store() -> Arc { + (Arc::new(InMemory::new()) as Arc).metered("memory".into()) + } + + /// A single materialized snapshot of recorded metrics. It must be taken + /// only once: the snapshotter *drains* histogram samples on every + /// `snapshot()` call, so a second snapshot would see empty histograms. + type Metrics = Vec<(metrics::Key, DebugValue)>; + + /// Materialize the current recorder state. Histogram samples are *drained* + /// on each call, so a metric must be read from a single snapshot. + fn snapshot(snapshotter: &Snapshotter) -> Metrics { + snapshotter + .snapshot() + .into_vec() + .into_iter() + .map(|(ck, _unit, _desc, value)| (ck.key().clone(), value)) + .collect() + } + + /// Run an async closure with a thread-local metrics recorder installed and + /// return the resulting metrics. Uses a current-thread runtime so all polls + /// happen on the thread that holds the recorder guard. + fn capture_metrics(f: F) -> Metrics + where + F: FnOnce() -> Fut, + Fut: std::future::Future, + { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(f()); + }); + snapshot(&snapshotter) + } + + fn key_matches(key: &metrics::Key, name: &str, labels: &[(&str, &str)]) -> bool { + if key.name() != name { + return false; + } + let actual: std::collections::HashSet<(&str, &str)> = + key.labels().map(|l| (l.key(), l.value())).collect(); + labels.len() == actual.len() && labels.iter().all(|l| actual.contains(l)) + } + + fn counter_value(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> u64 { + for (key, value) in metrics { + if key_matches(key, name, labels) + && let DebugValue::Counter(v) = value + { + return *v; + } + } + 0 + } + + fn histogram_count(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> usize { + for (key, value) in metrics { + if key_matches(key, name, labels) + && let DebugValue::Histogram(samples) = value + { + return samples.len(); + } + } + 0 + } + + fn gauge_value(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> f64 { + for (key, value) in metrics { + if key_matches(key, name, labels) + && let DebugValue::Gauge(v) = value + { + return v.0; + } + } + 0.0 + } + + fn has_metric(metrics: &Metrics, name: &str, labels: &[(&str, &str)]) -> bool { + metrics + .iter() + .any(|(key, _)| key_matches(key, name, labels)) + } + + #[test] + fn test_parse_base_label_mode() { + assert_eq!(parse_base_label_mode(None), BaseLabelMode::Scheme); + assert_eq!(parse_base_label_mode(Some("scheme")), BaseLabelMode::Scheme); + assert_eq!(parse_base_label_mode(Some("full")), BaseLabelMode::Full); + assert_eq!(parse_base_label_mode(Some("off")), BaseLabelMode::Off); + assert_eq!(parse_base_label_mode(Some("none")), BaseLabelMode::Off); + // Unrecognized values fall back to the conservative default. + assert_eq!(parse_base_label_mode(Some("bogus")), BaseLabelMode::Scheme); + } + + #[test] + fn test_scoped_base() { + assert_eq!( + scoped_base(BaseLabelMode::Full, "s3$bucket").as_deref(), + Some("s3$bucket") + ); + assert_eq!( + scoped_base(BaseLabelMode::Scheme, "s3$bucket").as_deref(), + Some("s3") + ); + // Azure keeps only the scheme even though its prefix carries the account. + assert_eq!( + scoped_base(BaseLabelMode::Scheme, "az$container@account").as_deref(), + Some("az") + ); + // A prefix without `$` (e.g. memory/file) is unchanged by scheme mode. + assert_eq!( + scoped_base(BaseLabelMode::Scheme, "memory").as_deref(), + Some("memory") + ); + assert_eq!(scoped_base(BaseLabelMode::Off, "s3$bucket"), None); + } + + #[test] + fn test_base_label_defaults_to_scheme() { + // No env var is set in the test process, so the default `scheme` mode + // applies: the full prefix collapses to just the scheme. + let recorded = capture_metrics(|| async { + let store = (Arc::new(InMemory::new()) as Arc) + .metered("s3$my-bucket".into()); + store.put(&Path::from("a"), payload(b"x")).await.unwrap(); + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "put"), ("base", "s3")] + ), + 1 + ); + // The full prefix is not emitted as the label under the default mode. + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "put"), ("base", "s3$my-bucket")] + ), + 0 + ); + } + + #[test] + fn test_put_records_count_bytes_and_latency() { + let data = b"hello world"; + let recorded = capture_metrics(|| async { + let store = metered_store(); + store + .put(&Path::from("a/b.bin"), payload(data)) + .await + .unwrap(); + }); + + let labels = [("operation", "put"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &labels), + data.len() as u64 + ); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_get_records_count_and_bytes() { + let data = b"hello world"; + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(data)).await.unwrap(); + // The GET is only recorded once its body has been fully drained. + store.get(&path).await.unwrap().bytes().await.unwrap(); + }); + + let labels = [("operation", "get"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &labels), + data.len() as u64 + ); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_get_not_recorded_until_body_drained() { + let data = b"hello world"; + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "get"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(data)).await.unwrap(); + + // Holding the result without reading the body records nothing yet. + let result = store.get(&path).await.unwrap(); + assert_eq!( + counter_value(&snapshot(&snapshotter), METRIC_REQUESTS, &labels), + 0 + ); + + // Draining the body records the request with the true byte count. + let bytes = result.bytes().await.unwrap(); + assert_eq!(bytes.len(), data.len()); + let recorded = snapshot(&snapshotter); + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &labels), + data.len() as u64 + ); + }); + }); + } + + #[test] + fn test_head_is_a_separate_operation() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(b"hello world")).await.unwrap(); + store.head(&path).await.unwrap(); + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "head"), ("base", "memory")] + ), + 1 + ); + // The head call must not be counted as a get. + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "get"), ("base", "memory")] + ), + 0 + ); + // A HEAD transfers only metadata, so it records no payload bytes. + assert_eq!( + counter_value( + &recorded, + METRIC_BYTES, + &[("operation", "head"), ("base", "memory")] + ), + 0 + ); + } + + #[test] + fn test_delete_records_one_request_per_call() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + for i in 0..3 { + store + .put(&Path::from(format!("a/{i}.bin")), payload(b"x")) + .await + .unwrap(); + } + // `delete` drives `delete_stream`; deleting three paths is still one + // logical delete request (a single batched request on real stores). + let paths = + futures::stream::iter((0..3).map(|i| Ok(Path::from(format!("a/{i}.bin"))))).boxed(); + let _: Vec<_> = store.delete_stream(paths).collect().await; + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "delete"), ("base", "memory")] + ), + 1 + ); + } + + #[test] + fn test_list_counts_one_request_not_per_item() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + for i in 0..3 { + store + .put(&Path::from(format!("a/{i}.bin")), payload(b"x")) + .await + .unwrap(); + } + let _: Vec<_> = store.list(Some(&Path::from("a"))).collect().await; + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "list"), ("base", "memory")] + ), + 1 + ); + } + + #[test] + fn test_error_is_counted() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + // Getting a missing object errors. + let _ = store.get(&Path::from("does/not/exist")).await; + }); + + let labels = [("operation", "get"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &labels), 1); + // A failed request is still counted as a request, with latency recorded. + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + // No bytes are transferred on a failed get. + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + #[test] + fn test_get_ranges_sums_part_bytes_and_labels_get() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(b"hello world")).await.unwrap(); + // Two disjoint ranges of 3 bytes each. + store.get_ranges(&path, &[2..5, 6..9]).await.unwrap(); + }); + + let labels = [("operation", "get"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 6); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_copy_and_rename_record_zero_bytes() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + store + .put(&Path::from("a/src"), payload(b"x")) + .await + .unwrap(); + store + .copy(&Path::from("a/src"), &Path::from("a/copy")) + .await + .unwrap(); + store + .rename(&Path::from("a/copy"), &Path::from("a/moved")) + .await + .unwrap(); + }); + + for operation in ["copy", "rename"] { + let labels = [("operation", operation), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + } + + #[test] + fn test_list_with_delimiter_records_latency() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + store.put(&Path::from("a/b"), payload(b"x")).await.unwrap(); + store + .list_with_delimiter(Some(&Path::from("a"))) + .await + .unwrap(); + }); + + let labels = [("operation", "list"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + #[test] + fn test_list_with_offset_counts_one_request() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + for i in 0..3 { + store + .put(&Path::from(format!("a/{i}")), payload(b"x")) + .await + .unwrap(); + } + let _: Vec<_> = store + .list_with_offset(Some(&Path::from("a")), &Path::from("a/0")) + .collect() + .await; + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "list"), ("base", "memory")] + ), + 1 + ); + } + + #[test] + fn test_multipart_records_each_part_and_complete() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let mut upload = store.put_multipart(&Path::from("a/big")).await.unwrap(); + upload.put_part(payload(b"hello")).await.unwrap(); // 5 bytes + upload.put_part(payload(b"world!!")).await.unwrap(); // 7 bytes + upload.complete().await.unwrap(); + }); + + let part_labels = [("operation", "put_part"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &part_labels), 2); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &part_labels), 12); + // Each part records its own latency sample, like a unary put. + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &part_labels), 2); + // A successful part upload records no error. + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &part_labels), 0); + + // Completing the upload is its own metered request. + let complete_labels = [("operation", "complete_multipart"), ("base", "memory")]; + assert_eq!( + counter_value(&recorded, METRIC_REQUESTS, &complete_labels), + 1 + ); + assert_eq!( + histogram_count(&recorded, METRIC_DURATION, &complete_labels), + 1 + ); + } + + #[test] + fn test_multipart_abort_is_recorded() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let mut upload = store.put_multipart(&Path::from("a/big")).await.unwrap(); + upload.put_part(payload(b"hello")).await.unwrap(); + upload.abort().await.unwrap(); + }); + + let labels = [("operation", "abort_multipart"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + } + + #[test] + fn test_multipart_part_error_is_counted() { + let recorded = capture_metrics(|| async { + let store = (Arc::new(FailingStreamStore) as Arc) + .metered("memory".into()); + let mut upload = store.put_multipart(&Path::from("a/big")).await.unwrap(); + let _ = upload.put_part(payload(b"data")).await; + }); + + let labels = [("operation", "put_part"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + // A failed part transfers no counted bytes. + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + #[test] + fn test_in_flight_guard_tracks_and_releases() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "get"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let g1 = InFlightGuard::new("memory", "get"); + let g2 = InFlightGuard::new("memory", "get"); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 2.0 + ); + drop(g1); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(g2); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + } + + #[test] + fn test_in_flight_gauge_is_wired_and_balances() { + let recorded = capture_metrics(|| async { + let store = metered_store(); + let path = Path::from("a/b.bin"); + store.put(&path, payload(b"hello")).await.unwrap(); + store.get(&path).await.unwrap(); + }); + + // The gauge is emitted for each operation (guard is wired in) and, once + // the operation completes, balances back to zero. + for operation in ["put", "get"] { + let labels = [("operation", operation), ("base", "memory")]; + assert!(has_metric(&recorded, METRIC_IN_FLIGHT, &labels)); + assert_eq!(gauge_value(&recorded, METRIC_IN_FLIGHT, &labels), 0.0); + } + } + + #[test] + fn test_list_stream_holds_in_flight_until_dropped() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "list"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let store = metered_store(); + store.put(&Path::from("a/x"), payload(b"x")).await.unwrap(); + + // Creating the stream raises the gauge; it stays raised until the + // stream is dropped, even before any items are drained. + let stream = store.list(Some(&Path::from("a"))); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(stream); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + }); + } + + #[test] + fn test_delete_stream_holds_in_flight_until_dropped() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "delete"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let store = metered_store(); + let locations = futures::stream::iter(vec![Ok(Path::from("a/b"))]).boxed(); + + // Like list, creating the delete stream raises the gauge and holds + // it until the stream is dropped, before any items are drained. + let stream = store.delete_stream(locations); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(stream); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + }); + } + + #[test] + fn test_in_flight_released_when_operation_future_dropped() { + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "get"), ("base", "memory")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let started = Arc::new(tokio::sync::Notify::new()); + // Never signalled: the request stays blocked mid-flight. + let release = Arc::new(tokio::sync::Notify::new()); + let store = (Arc::new(BlockingStore { + started: started.clone(), + release, + }) as Arc) + .metered("memory".into()); + + let path = Path::from("a/b"); + let mut fut = Box::pin(store.get(&path)); + // Drive the request until it is blocked inside the inner store. + tokio::select! { + _ = &mut fut => unreachable!("the blocking store never returns"), + _ = started.notified() => {} + } + + // The gauge is raised while the request is outstanding, and + // dropping the future before it completes releases it. + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 1.0 + ); + drop(fut); + assert_eq!( + gauge_value(&snapshot(&snapshotter), METRIC_IN_FLIGHT, &labels), + 0.0 + ); + }); + }); + } + + #[test] + fn test_streaming_errors_are_counted() { + let recorded = capture_metrics(|| async { + let delete_store = (Arc::new(FailingStreamStore) as Arc) + .metered("memory".into()); + let _ = delete_store.delete(&Path::from("a/b")).await; + + let list_store = (Arc::new(FailingStreamStore) as Arc) + .metered("memory".into()); + let _: Vec<_> = list_store.list(None).collect().await; + }); + + // delete_stream counts the item and records an error when it fails. + let delete_labels = [("operation", "delete"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &delete_labels), 1); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &delete_labels), 1); + + // A list request is counted once; a failure while draining records an error. + let list_labels = [("operation", "list"), ("base", "memory")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &list_labels), 1); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &list_labels), 1); + } + + /// The optimized local reads and writes talk to the filesystem directly, so + /// they never reach [`MeteredObjectStore`] and publish these metrics + /// themselves. They must land under the same `base` label as the store's + /// metered operations, which for a local store is its scheme. + #[test] + fn test_local_filesystem_io_is_metered() { + let tmp = TempStdDir::default(); + let dir = tmp.join("sub"); + let data = b"hello world"; + let recorded = capture_metrics(|| async { + // Built through the registry, like any store opened from a URI. + let (store, path) = LanceObjectStore::from_uri(dir.join("a.bin").to_str().unwrap()) + .await + .unwrap(); + // Writes go through LocalWriter. + store.put(&path, data).await.unwrap(); + + // Reads go through LocalObjectReader. + let reader = store.open(&path).await.unwrap(); + assert_eq!(reader.size().await.unwrap(), data.len()); + assert_eq!(reader.get_range(0..5).await.unwrap().len(), 5); + assert_eq!(reader.get_all().await.unwrap().len(), data.len()); + // The file is smaller than the block size, so it streams as one chunk. + let chunks: Vec<_> = reader.get_stream().await.unwrap().collect().await; + assert_eq!(chunks.len(), 1); + + // Copy and recursive delete both shortcut to the filesystem too. + store + .copy(&path, &Path::from_absolute_path(dir.join("b.bin")).unwrap()) + .await + .unwrap(); + store + .remove_dir_all(Path::from_absolute_path(&dir).unwrap()) + .await + .unwrap(); + }); + + let put_labels = [("operation", "put"), ("base", "file")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &put_labels), 1); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &put_labels), + data.len() as u64 + ); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &put_labels), 1); + assert_eq!(gauge_value(&recorded, METRIC_IN_FLIGHT, &put_labels), 0.0); + + // One request each for the range read, the full read and the single + // streamed chunk. + let get_labels = [("operation", "get"), ("base", "file")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &get_labels), 3); + assert_eq!( + counter_value(&recorded, METRIC_BYTES, &get_labels), + (5 + 2 * data.len()) as u64 + ); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &get_labels), 3); + assert_eq!(gauge_value(&recorded, METRIC_IN_FLIGHT, &get_labels), 0.0); + + // The size lookup is the local equivalent of a HEAD, and transfers no + // payload bytes. + let head_labels = [("operation", "head"), ("base", "file")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &head_labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &head_labels), 0); + + for operation in ["copy", "delete"] { + let labels = [("operation", operation), ("base", "file")]; + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &get_labels), 0); + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &put_labels), 0); + } + + #[test] + fn test_local_read_error_is_counted() { + let tmp = TempStdDir::default(); + let recorded = capture_metrics(|| async { + let (store, path) = LanceObjectStore::from_uri(tmp.join("a.bin").to_str().unwrap()) + .await + .unwrap(); + store.put(&path, b"hello").await.unwrap(); + + let reader = store.open(&path).await.unwrap(); + // Reading past the end of the file fails. + assert!(reader.get_range(0..100).await.is_err()); + }); + + let labels = [("operation", "get"), ("base", "file")]; + assert_eq!(counter_value(&recorded, METRIC_ERRORS, &labels), 1); + // A failed read is still counted as a request, with latency recorded. + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(histogram_count(&recorded, METRIC_DURATION, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 0); + } + + /// A local write is reported as a single `put` covering the whole file, so it + /// stays in flight until the file is persisted under its final path. + #[test] + fn test_local_write_is_in_flight_until_persisted() { + let tmp = TempStdDir::default(); + let recorder = DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + let labels = [("operation", "put"), ("base", "file")]; + metrics::with_local_recorder(&recorder, || { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + let (store, path) = LanceObjectStore::from_uri(tmp.join("a.bin").to_str().unwrap()) + .await + .unwrap(); + let mut writer = store.create(&path).await.unwrap(); + writer.write_all(b"hello").await.unwrap(); + + let recorded = snapshot(&snapshotter); + assert_eq!(gauge_value(&recorded, METRIC_IN_FLIGHT, &labels), 1.0); + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 0); + + Writer::shutdown(writer.as_mut()).await.unwrap(); + let recorded = snapshot(&snapshotter); + assert_eq!(gauge_value(&recorded, METRIC_IN_FLIGHT, &labels), 0.0); + assert_eq!(counter_value(&recorded, METRIC_REQUESTS, &labels), 1); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), 5); + }); + }); + } + + /// A store handed in by the caller is metered like one built by the registry. + #[test] + fn test_caller_supplied_store_is_metered() { + let recorded = capture_metrics(|| async { + #[allow(deprecated)] + let params = crate::object_store::ObjectStoreParams { + object_store: Some(( + Arc::new(InMemory::new()) as Arc, + Url::parse("memory:///").unwrap(), + )), + ..Default::default() + }; + let (store, _) = LanceObjectStore::from_uri_and_params( + Arc::new(crate::object_store::ObjectStoreRegistry::default()), + "memory:///", + ¶ms, + ) + .await + .unwrap(); + store.put(&Path::from("a"), b"hello").await.unwrap(); + }); + + assert_eq!( + counter_value( + &recorded, + METRIC_REQUESTS, + &[("operation", "put"), ("base", "memory")] + ), + 1 + ); + } + + /// `ObjectStore::new` is how `DatasetBuilder` wraps a caller-supplied store, + /// so it must meter both halves of the store: the operations that go through + /// `inner`, and the local ones that bypass it. Metering only one half would + /// report a partial picture that reads like a complete one. + #[test] + fn test_store_built_from_new_is_metered() { + let tmp = TempStdDir::default(); + let recorded = capture_metrics(|| async { + let store = LanceObjectStore::new( + Arc::new(object_store::local::LocalFileSystem::new()), + Url::parse("file:///").unwrap(), + None, + None, + false, + false, + 1, + 3, + None, + ); + let path = Path::from_absolute_path(tmp.join("a.bin")).unwrap(); + // put and open bypass `inner` and publish for themselves. + store.put(&path, b"hello").await.unwrap(); + let reader = store.open(&path).await.unwrap(); + assert_eq!(reader.get_all().await.unwrap().len(), 5); + // delete goes through `inner`, so only MeteredObjectStore can count it. + store.delete(&path).await.unwrap(); + }); + + for (operation, bytes) in [("put", 5), ("get", 5), ("delete", 0)] { + let labels = [("operation", operation), ("base", "file")]; + assert_eq!( + counter_value(&recorded, METRIC_REQUESTS, &labels), + 1, + "expected one {operation} request" + ); + assert_eq!(counter_value(&recorded, METRIC_BYTES, &labels), bytes); + } + } + + /// A store whose stream-producing operations always yield an error, used to + /// exercise the error branches of the streaming wrappers. + #[derive(Debug)] + struct FailingStreamStore; + + impl std::fmt::Display for FailingStreamStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "FailingStreamStore") + } + } + + fn test_error() -> object_store::Error { + object_store::Error::Generic { + store: "FailingStreamStore", + source: "injected failure".into(), + } + } + + #[async_trait::async_trait] + impl object_store::ObjectStore for FailingStreamStore { + async fn put_opts( + &self, + _location: &Path, + _bytes: PutPayload, + _opts: PutOptions, + ) -> OSResult { + unimplemented!() + } + + async fn put_multipart_opts( + &self, + _location: &Path, + _opts: PutMultipartOptions, + ) -> OSResult> { + Ok(Box::new(FailingUpload)) + } + + async fn get_opts(&self, _location: &Path, _options: GetOptions) -> OSResult { + unimplemented!() + } + + fn delete_stream( + &self, + _locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + futures::stream::once(async { Err(test_error()) }).boxed() + } + + fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + futures::stream::once(async { Err(test_error()) }).boxed() + } + + fn list_with_offset( + &self, + _prefix: Option<&Path>, + _offset: &Path, + ) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> OSResult { + unimplemented!() + } + + async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> { + unimplemented!() + } + + async fn rename_opts( + &self, + _from: &Path, + _to: &Path, + _opts: RenameOptions, + ) -> OSResult<()> { + unimplemented!() + } + } + + /// A [`MultipartUpload`] whose part uploads always fail, used to exercise the + /// error branch of the metered `put_part`. + #[derive(Debug)] + struct FailingUpload; + + #[async_trait::async_trait] + impl MultipartUpload for FailingUpload { + fn put_part(&mut self, _data: PutPayload) -> UploadPart { + async { Err(test_error()) }.boxed() + } + + async fn complete(&mut self) -> OSResult { + unimplemented!() + } + + async fn abort(&mut self) -> OSResult<()> { + unimplemented!() + } + } + + /// A store whose `get_opts` blocks after signalling `started`, so a request + /// can be observed mid-flight and then dropped before it completes. + #[derive(Debug)] + struct BlockingStore { + started: Arc, + release: Arc, + } + + impl std::fmt::Display for BlockingStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "BlockingStore") + } + } + + #[async_trait::async_trait] + impl object_store::ObjectStore for BlockingStore { + async fn put_opts( + &self, + _location: &Path, + _bytes: PutPayload, + _opts: PutOptions, + ) -> OSResult { + unimplemented!() + } + + async fn put_multipart_opts( + &self, + _location: &Path, + _opts: PutMultipartOptions, + ) -> OSResult> { + unimplemented!() + } + + async fn get_opts(&self, _location: &Path, _options: GetOptions) -> OSResult { + self.started.notify_one(); + self.release.notified().await; + unreachable!("release is never signalled in the test") + } + + fn delete_stream( + &self, + _locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + fn list_with_offset( + &self, + _prefix: Option<&Path>, + _offset: &Path, + ) -> BoxStream<'static, OSResult> { + unimplemented!() + } + + async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> OSResult { + unimplemented!() + } + + async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> { + unimplemented!() + } + + async fn rename_opts( + &self, + _from: &Path, + _to: &Path, + _opts: RenameOptions, + ) -> OSResult<()> { + unimplemented!() + } + } +} diff --git a/vendor/lance-io/src/object_store/providers.rs b/vendor/lance-io/src/object_store/providers.rs new file mode 100644 index 000000000..b84dc7362 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers.rs @@ -0,0 +1,520 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + collections::HashMap, + sync::{ + Arc, RwLock, Weak, + atomic::{AtomicU64, Ordering}, + }, +}; + +use object_store::path::Path; +use url::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}; + +#[cfg(feature = "aws")] +pub mod aws; +#[cfg(feature = "azure")] +pub mod azure; +#[cfg(feature = "gcp")] +pub mod gcp; +#[cfg(feature = "goosefs")] +pub mod goosefs; +#[cfg(feature = "huggingface")] +pub mod huggingface; +pub mod local; +pub mod memory; +#[cfg(feature = "oss")] +pub mod oss; +pub mod shared_memory; +#[cfg(feature = "tencent")] +pub mod tencent; +#[cfg(feature = "tos")] +pub mod tos; + +#[async_trait::async_trait] +pub trait ObjectStoreProvider: std::fmt::Debug + Sync + Send { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result; + + /// Extract the path relative to the base of the store. + /// + /// For example, in S3 the path is relative to the bucket. So a URL of + /// `s3://bucket/path/to/file` would return `path/to/file`. + /// + /// Meanwhile, for a file store, the path is relative to the filesystem root. + /// So a URL of `file:///path/to/file` would return `/path/to/file`. + fn extract_path(&self, url: &Url) -> Result { + // url.path() returns a percent-encoded string (per the WHATWG URL spec). + // Path::from_url_path decodes it first so the Path internal representation + // holds the raw UTF-8 string. This prevents double-encoding when the + // object store client later percent-encodes the path for HTTP requests. + Path::from_url_path(url.path()).map_err(|e| { + Error::invalid_input(format!("Invalid path in URL '{}': {}", url.path(), e)) + }) + } + + /// Calculate the unique prefix that should be used for this object store. + /// + /// For object stores that don't have the concept of buckets, this will just be something like + /// 'file' or 'memory'. + /// + /// In object stores where all bucket names are unique, like s3, this will be + /// simply 's3$my_bucket_name' or similar. + /// + /// In Azure, only the combination of (account name, container name) is unique, so + /// this will be something like 'az$account_name@container' + /// + /// Providers should override this if they have special requirements like Azure's. + fn calculate_object_store_prefix( + &self, + url: &Url, + _storage_options: Option<&HashMap>, + ) -> Result { + Ok(format!("{}${}", url.scheme(), url.authority())) + } +} + +/// Statistics for the object store registry cache. +#[derive(Debug, Clone, Default)] +pub struct ObjectStoreRegistryStats { + /// Number of cache hits (store was already cached and reused). + pub hits: u64, + /// Number of cache misses (new store had to be created). + pub misses: u64, + /// Number of currently active object stores in the cache. + pub active_stores: usize, +} + +/// A registry of object store providers. +/// +/// Use [`Self::default()`] to create one with the available default providers. +/// This includes (depending on features enabled): +/// - `memory`: An in-memory object store. +/// - `file`: A local file object store, with optimized code paths. +/// - `file-object-store`: A local file object store that uses the ObjectStore API, +/// for all operations. Used for testing with ObjectStore wrappers. +/// - `file+uring`: A local file object store using io_uring (Linux only). +/// - `s3`: An S3 object store. +/// - `s3+ddb`: An S3 object store with DynamoDB for metadata. +/// - `az`: An Azure Blob Storage object store. +/// - `gs`: A Google Cloud Storage object store. +/// - `tos`: A Volcengine TOS object store. +/// +/// 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()`]. +#[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>>, + // Cache statistics + hits: AtomicU64, + misses: AtomicU64, +} + +impl ObjectStoreRegistry { + /// Create a new registry with no providers registered. + /// + /// Typically, you want to use [`Self::default()`] instead, so you get the + /// default providers. + pub fn empty() -> Self { + Self { + providers: RwLock::new(HashMap::new()), + active_stores: RwLock::new(HashMap::new()), + hits: AtomicU64::new(0), + misses: AtomicU64::new(0), + } + } + + /// Get the object store provider for a given scheme. + pub fn get_provider(&self, scheme: &str) -> Option> { + self.providers + .read() + .expect("ObjectStoreRegistry lock poisoned") + .get(scheme) + .cloned() + } + + /// Get a list of all active object stores. + /// + /// Calling this will also clean up any weak references to object stores that + /// are no longer valid. + pub fn active_stores(&self) -> Vec> { + let mut found_inactive = false; + let output = self + .active_stores + .read() + .expect("ObjectStoreRegistry lock poisoned") + .values() + .filter_map(|weak| match weak.upgrade() { + Some(store) => Some(store), + None => { + found_inactive = true; + None + } + }) + .collect(); + + if found_inactive { + // Clean up the cache by removing any weak references that are no longer valid + let mut cache_lock = self + .active_stores + .write() + .expect("ObjectStoreRegistry lock poisoned"); + cache_lock.retain(|_, weak| weak.upgrade().is_some()); + } + output + } + + /// 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). + 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); + ObjectStoreRegistryStats { + hits: self.hits.load(Ordering::Relaxed), + misses: self.misses.load(Ordering::Relaxed), + active_stores, + } + } + + fn scheme_not_found_error(&self, scheme: &str) -> Error { + let mut message = format!("No object store provider found for scheme: '{}'", scheme); + if let Ok(providers) = self.providers.read() { + let valid_schemes = providers.keys().cloned().collect::>().join(", "); + message.push_str(&format!("\nValid schemes: {}", valid_schemes)); + } + Error::invalid_input(message) + } + + /// Get an object store for a given base path and parameters. + /// + /// 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. + pub async fn get_store( + &self, + base_path: Url, + params: &ObjectStoreParams, + ) -> Result> { + // Base-scoped storage options (`base_.`) are directives for + // other registered base paths; resolve them away before building or + // caching a store for this location. Params already resolved for a + // base contain no scoped entries, so this is a no-op for them. + let params = params.scoped_to_base(None); + let params = params.as_ref(); + let scheme = base_path.scheme(); + let Some(provider) = self.get_provider(scheme) else { + return Err(self.scheme_not_found_error(scheme)); + }; + + let cache_path = + 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); + } + } + } + } + + self.misses.fetch_add(1, Ordering::Relaxed); + + let mut store = provider.new_store(base_path, params).await?; + + store.inner = store.inner.traced(); + + // Label metrics by the store's unique prefix (e.g. `s3$bucket`, + // `az$container@account`) so multiple stores on one cloud differ. + crate::object_store::meter_store(&mut store.inner, &mut store.io_tracker, &cache_path); + + 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 store = Arc::new(store); + + { + // 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)); + } + + Ok(store) + } + + /// Calculate the datastore prefix based on the URI and the storage options. + /// The data store prefix should uniquely identify the datastore. + pub fn calculate_object_store_prefix( + &self, + uri: &str, + storage_options: Option<&HashMap>, + ) -> Result { + let url = uri_to_url(uri)?; + match self.get_provider(url.scheme()) { + None => { + if url.scheme() == "file" || url.scheme().len() == 1 { + Ok("file".to_string()) + } else { + Err(self.scheme_not_found_error(url.scheme())) + } + } + Some(provider) => provider.calculate_object_store_prefix(&url, storage_options), + } + } +} + +impl Default for ObjectStoreRegistry { + fn default() -> Self { + let mut providers: HashMap> = HashMap::new(); + + providers.insert("memory".into(), Arc::new(memory::MemoryStoreProvider)); + providers.insert( + "shared-memory".into(), + Arc::new(shared_memory::SharedMemoryStoreProvider::default()), + ); + providers.insert("file".into(), Arc::new(local::FileStoreProvider)); + // The "file" scheme has special optimized code paths that bypass + // the ObjectStore API for better performance. However, this can make it + // hard to test when using ObjectStore wrappers, such as IOTrackingStore. + // So we provide a "file-object-store" scheme that uses the ObjectStore API. + // The specialized code paths are differentiated by the scheme name. + providers.insert( + "file-object-store".into(), + Arc::new(local::FileStoreProvider), + ); + #[cfg(target_os = "linux")] + providers.insert("file+uring".into(), Arc::new(local::FileStoreProvider)); + + #[cfg(feature = "aws")] + { + let aws = Arc::new(aws::AwsStoreProvider); + providers.insert("s3".into(), aws.clone()); + providers.insert("s3+ddb".into(), aws); + } + #[cfg(feature = "azure")] + { + let azure = Arc::new(azure::AzureBlobStoreProvider); + providers.insert("az".into(), azure.clone()); + providers.insert("abfss".into(), azure); + } + #[cfg(feature = "gcp")] + providers.insert("gs".into(), Arc::new(gcp::GcsStoreProvider)); + #[cfg(feature = "goosefs")] + providers.insert("goosefs".into(), Arc::new(goosefs::GooseFsStoreProvider)); + #[cfg(feature = "oss")] + providers.insert("oss".into(), Arc::new(oss::OssStoreProvider)); + #[cfg(feature = "tencent")] + providers.insert("cos".into(), Arc::new(tencent::TencentStoreProvider)); + #[cfg(feature = "huggingface")] + providers.insert("hf".into(), Arc::new(huggingface::HuggingfaceStoreProvider)); + #[cfg(feature = "tos")] + providers.insert("tos".into(), Arc::new(tos::TosStoreProvider)); + Self { + providers: RwLock::new(providers), + active_stores: RwLock::new(HashMap::new()), + hits: AtomicU64::new(0), + misses: AtomicU64::new(0), + } + } +} + +impl ObjectStoreRegistry { + /// Add a new object store provider to the registry. The provider will be used + /// in [`Self::get_store()`] when a URL is passed with a matching scheme. + pub fn insert(&self, scheme: &str, provider: Arc) { + self.providers + .write() + .expect("ObjectStoreRegistry lock poisoned") + .insert(scheme.into(), provider); + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + + #[derive(Debug)] + struct DummyProvider; + + #[async_trait::async_trait] + impl ObjectStoreProvider for DummyProvider { + async fn new_store( + &self, + _base_path: Url, + _params: &ObjectStoreParams, + ) -> Result { + unreachable!("This test doesn't create stores") + } + } + + #[test] + fn test_calculate_object_store_prefix() { + let provider = DummyProvider; + let url = Url::parse("dummy://blah/path").unwrap(); + assert_eq!( + "dummy$blah", + provider.calculate_object_store_prefix(&url, None).unwrap() + ); + } + + #[tokio::test] + async fn test_get_store_resolves_base_scoped_options() { + use crate::object_store::StorageOptionsAccessor; + + let registry = ObjectStoreRegistry::default(); + let url = Url::parse("memory://test").unwrap(); + + let with_scoped = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("shared".to_string(), "value".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ]), + ))), + ..Default::default() + }; + let without_scoped = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([("shared".to_string(), "value".to_string())]), + ))), + ..Default::default() + }; + + // Base-scoped entries are resolved away before the store is built and + // cached, so params with and without them yield the same cached store. + let store_scoped = registry.get_store(url.clone(), &with_scoped).await.unwrap(); + let store_plain = registry.get_store(url, &without_scoped).await.unwrap(); + assert!(Arc::ptr_eq(&store_scoped, &store_plain)); + } + + #[test] + fn test_calculate_object_store_scheme_not_found() { + let registry = ObjectStoreRegistry::empty(); + registry.insert("dummy", Arc::new(DummyProvider)); + let s = "Invalid user input: No object store provider found for scheme: 'dummy2'\nValid schemes: dummy"; + let result = registry + .calculate_object_store_prefix("dummy2://mybucket/my/long/path", None) + .expect_err("expected error") + .to_string(); + assert_eq!(s, &result[..s.len()]); + } + + // Test that paths without a scheme get treated as local paths. + #[test] + fn test_calculate_object_store_prefix_for_local() { + let registry = ObjectStoreRegistry::empty(); + assert_eq!( + "file", + registry + .calculate_object_store_prefix("/tmp/foobar", None) + .unwrap() + ); + } + + // Test that paths with a single-letter scheme that is not registered for anything get treated as local paths. + #[test] + fn test_calculate_object_store_prefix_for_local_windows_path() { + let registry = ObjectStoreRegistry::empty(); + assert_eq!( + "file", + registry + .calculate_object_store_prefix("c://dos/path", None) + .unwrap() + ); + } + + // Test that paths with a given scheme get mapped to that storage provider. + #[test] + fn test_calculate_object_store_prefix_for_dummy_path() { + let registry = ObjectStoreRegistry::empty(); + registry.insert("dummy", Arc::new(DummyProvider)); + assert_eq!( + "dummy$mybucket", + registry + .calculate_object_store_prefix("dummy://mybucket/my/long/path", None) + .unwrap() + ); + } + + #[tokio::test] + async fn test_stats_hit_miss_tracking() { + use crate::object_store::StorageOptionsAccessor; + let registry = ObjectStoreRegistry::default(); + let url = Url::parse("memory://test").unwrap(); + + let params1 = ObjectStoreParams::default(); + let params2 = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([("k".into(), "v".into())]), + ))), + ..Default::default() + }; + + // (hits, misses, active) + let cases: &[(&ObjectStoreParams, (u64, u64, usize))] = &[ + (¶ms1, (0, 1, 1)), // miss: new params + (¶ms1, (1, 1, 1)), // hit: same params + (¶ms2, (1, 2, 2)), // miss: different storage_options + ]; + + let mut stores = vec![]; // retain the stores + for (params, (hits, misses, active)) in cases { + stores.push(registry.get_store(url.clone(), params).await.unwrap()); + let s = registry.stats(); + assert_eq!( + (s.hits, s.misses, s.active_stores), + (*hits, *misses, *active) + ); + } + + // Same params returns same instance + assert!(Arc::ptr_eq(&stores[0], &stores[1])); + } +} diff --git a/vendor/lance-io/src/object_store/providers/aws.rs b/vendor/lance-io/src/object_store/providers/aws.rs new file mode 100644 index 000000000..d9567f2c1 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/aws.rs @@ -0,0 +1,1531 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, str::FromStr, sync::Arc, time::Duration}; + +#[cfg(test)] +use mock_instant::thread_local::{SystemTime, UNIX_EPOCH}; + +#[cfg(not(test))] +use std::time::{SystemTime, UNIX_EPOCH}; + +use object_store::ObjectStore as OSObjectStore; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::S3}; + +use aws_config::Region; +use aws_config::default_provider::credentials::DefaultCredentialsChain; +use aws_config::ecs::EcsCredentialsProvider; +use aws_config::provider_config::ProviderConfig; +use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider; +use aws_credential_types::provider::ProvideCredentials; +use object_store::{ + ClientOptions, CredentialProvider, Result as ObjectStoreResult, RetryConfig, + StaticCredentialProvider, + aws::{ + AmazonS3Builder, AmazonS3ConfigKey, AwsCredential as ObjectStoreAwsCredential, + AwsCredentialProvider, + }, +}; +use tokio::sync::RwLock; +use url::Url; + +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, + StorageOptionsProvider, + dynamic_credentials::{NamespaceCredentialsProvider, build_dynamic_credential_provider}, + dynamic_opendal::DynamicOpenDalStore, + throttle::{AimdThrottleConfig, AimdThrottleState, AimdThrottledStore, cloud_http_connector}, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct AwsStoreProvider; + +const AWS_ACCESS_KEY_ID: &str = "aws_access_key_id"; +const AWS_SECRET_ACCESS_KEY: &str = "aws_secret_access_key"; +const AWS_SESSION_TOKEN: &str = "aws_session_token"; + +fn is_aws_credential_key(key: &AmazonS3ConfigKey) -> bool { + matches!( + key, + AmazonS3ConfigKey::AccessKeyId + | AmazonS3ConfigKey::SecretAccessKey + | AmazonS3ConfigKey::Token + ) +} + +fn has_aws_credential_member(options: &HashMap) -> bool { + options.keys().any(|key| { + AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) + .is_ok_and(|key| is_aws_credential_key(&key)) + }) +} + +/// Merge AWS environment options without combining two credential authorities. +fn with_atomic_env_s3(options: &mut StorageOptions) { + merge_atomic_aws_environment( + options, + std::env::vars_os() + .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))), + ); +} + +/// Merge a supplied environment into S3 options without mixing credential families. +/// +/// This is public only so downstream regression tests can provide a deterministic environment. +#[doc(hidden)] +pub fn merge_atomic_aws_environment( + options: &mut StorageOptions, + environment: impl IntoIterator, +) { + let has_explicit_credential = has_aws_credential_member(&options.0); + for (key, value) in environment { + if let Ok(config_key) = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) + && !(has_explicit_credential && is_aws_credential_key(&config_key)) + { + options + .0 + .entry(config_key.as_ref().to_string()) + .or_insert(value); + } + } +} + +fn canonical_opendal_s3_config(options: &HashMap) -> HashMap { + options + .iter() + .map(|(key, value)| { + let key = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) + .map(|key| key.as_ref().to_string()) + .unwrap_or_else(|_| key.clone()); + (key, value.clone()) + }) + .collect() +} + +fn dynamic_aws_credential_options(options: &HashMap) -> HashMap { + options + .iter() + .filter_map(|(key, value)| { + let key = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()).ok()?; + is_aws_credential_key(&key).then(|| (key.as_ref().to_string(), value.clone())) + }) + .collect() +} + +fn normalize_opendal_s3_config( + options: &HashMap, +) -> Result> { + let mut options = canonical_opendal_s3_config(options); + let key_id = options.get(AWS_ACCESS_KEY_ID); + let secret_key = options.get(AWS_SECRET_ACCESS_KEY); + let token = options.get(AWS_SESSION_TOKEN); + if key_id.is_some() || secret_key.is_some() || token.is_some() { + if key_id.is_none() || secret_key.is_none() { + return Err(Error::invalid_input( + "Explicit AWS credentials require both aws_access_key_id and aws_secret_access_key", + )); + } + // Once one explicit family wins, OpenDAL must not consult a second ambient authority. + options.insert("disable_config_load".to_string(), "true".to_string()); + } + Ok(options) +} + +fn build_opendal_s3_store(config: HashMap) -> Result { + let operator = Operator::from_iter::(config).map_err(|error| { + Error::invalid_input(format!("Failed to create S3 operator: {error:?}")) + })?; + Ok(OpendalStore::new(operator)) +} + +#[derive(Debug)] +struct AwsCredentialStorageOptionsProvider { + provider: AwsCredentialProvider, +} + +#[async_trait::async_trait] +impl StorageOptionsProvider for AwsCredentialStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let credential = self + .provider + .get_credential() + .await + .map_err(|error| Error::io_source(Box::new(error)))?; + let mut options = HashMap::from([ + (AWS_ACCESS_KEY_ID.to_string(), credential.key_id.clone()), + ( + AWS_SECRET_ACCESS_KEY.to_string(), + credential.secret_key.clone(), + ), + // The object_store provider owns its own cache, so ask it on every store access. + ("expires_at_millis".to_string(), "0".to_string()), + ]); + if let Some(token) = &credential.token { + options.insert(AWS_SESSION_TOKEN.to_string(), token.clone()); + } + Ok(Some(options)) + } + + fn provider_id(&self) -> String { + format!( + "aws-credential-provider[{:p}]", + Arc::as_ptr(&self.provider) as *const () + ) + } +} + +impl AwsStoreProvider { + async fn build_amazon_s3_store( + &self, + base_path: &mut Url, + params: &ObjectStoreParams, + storage_options: &StorageOptions, + is_s3_express: bool, + throttle_state: Option<&AimdThrottleState>, + ) -> Result> { + // Use a low retry count since the AIMD throttle layer handles + // throttle recovery with its own retry loop. + let retry_config = RetryConfig { + backoff: Default::default(), + max_retries: storage_options.client_max_retries(), + retry_timeout: Duration::from_secs(storage_options.client_retry_timeout()), + }; + + let mut s3_storage_options = storage_options.as_s3_options(); + let region = resolve_s3_region(base_path, &s3_storage_options).await?; + + // Get accessor from params + let accessor = params.get_accessor(); + + let provider_scheme = storage_options.aws_provider_scheme()?; + + let (aws_creds, region) = build_aws_credential( + params.s3_credentials_refresh_offset, + params.aws_credentials.clone(), + Some(&s3_storage_options), + region, + accessor, + provider_scheme, + ) + .await?; + + // Set S3Express flag if detected + if is_s3_express { + s3_storage_options.insert(AmazonS3ConfigKey::S3Express, true.to_string()); + } + + // Compute the metrics label before rewriting the URL below so it + // matches the prefix the registry uses to key this store. + let store_prefix = + self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?; + + // before creating the OSObjectStore we need to rewrite the url to drop ddb related parts + base_path.set_scheme("s3").unwrap(); + base_path.set_query(None); + + // we can't use parse_url_opts here because we need to manually set the credentials provider + let mut builder = + AmazonS3Builder::new().with_client_options(storage_options.client_options()?); + for (key, value) in s3_storage_options { + builder = builder.with_config(key, value); + } + builder = builder + .with_url(base_path.as_ref()) + .with_credentials(aws_creds) + .with_retry(retry_config) + .with_region(region); + + builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix)); + + Ok(Arc::new(builder.build()?) as Arc) + } + + async fn build_opendal_s3_store( + &self, + base_path: &Url, + params: &ObjectStoreParams, + storage_options: &StorageOptions, + ) -> Result> { + let bucket = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("S3 URL must contain bucket name"))? + .to_string(); + + let prefix = base_path.path().trim_start_matches('/').to_string(); + + if let Some(provider_scheme) = storage_options.aws_provider_scheme()? { + return Result::Err(Error::not_supported(format!( + "OpendalStore does not currently support an explicit provider_scheme (currently set to {:?})", + provider_scheme + ))); + } + + let mut config_map = canonical_opendal_s3_config(&storage_options.0); + config_map.insert("bucket".to_string(), bucket); + + if !prefix.is_empty() { + config_map.insert("root".to_string(), "/".to_string()); + } + + let dynamic_accessor = if let Some(provider) = params.aws_credentials.clone() { + Some(Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + AwsCredentialStorageOptionsProvider { provider }, + )))) + } else { + params + .get_accessor() + .filter(|accessor| accessor.has_provider()) + }; + + if let Some(accessor) = dynamic_accessor { + // Dynamic OpenDAL refresh is deliberately credential-only. Noncredential changes can + // alter the outer ObjectStore contract and require a new registry entry instead. + let store = DynamicOpenDalStore::new( + "s3", + config_map, + accessor, + normalize_opendal_s3_config, + build_opendal_s3_store, + ) + .with_dynamic_options_filter(dynamic_aws_credential_options) + .with_atomic_key_group([ + AWS_ACCESS_KEY_ID, + AWS_SECRET_ACCESS_KEY, + AWS_SESSION_TOKEN, + ]); + // Validate the currently active family and prime the normalized-config cache. + store.current_store().await?; + Ok(Arc::new(store) as Arc) + } else { + let config_map = normalize_opendal_s3_config(&config_map)?; + Ok(Arc::new(build_opendal_s3_store(config_map)?) as Arc) + } + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for AwsStoreProvider { + async fn new_store( + &self, + mut base_path: Url, + params: &ObjectStoreParams, + ) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let mut storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); + with_atomic_env_s3(&mut storage_options); + let download_retry_count = storage_options.download_retry_count(); + + let use_opendal = storage_options + .0 + .get("use_opendal") + .map(|v| v == "true") + .unwrap_or(false); + + // Determine S3 Express and constant size upload parts before building the store + let is_s3_express = check_s3_express(&base_path, &storage_options); + + let use_constant_size_upload_parts = storage_options + .0 + .get("aws_endpoint") + .map(|endpoint| endpoint.contains("r2.cloudflarestorage.com")) + .unwrap_or(false); + + let throttle_config = AimdThrottleConfig::from_storage_options(params.storage_options())?; + let throttle_state = if throttle_config.is_disabled() { + None + } else { + Some(AimdThrottleState::new(throttle_config)?) + }; + + let inner = if use_opendal { + // Use OpenDAL implementation + self.build_opendal_s3_store(&base_path, params, &storage_options) + .await? + } else { + // Use default Amazon S3 implementation + self.build_amazon_s3_store( + &mut base_path, + params, + &storage_options, + is_s3_express, + throttle_state.as_ref(), + ) + .await? + }; + let inner = if let Some(throttle_state) = throttle_state { + Arc::new(AimdThrottledStore::new_with_state( + inner, + throttle_state, + !use_opendal, + )) as Arc + } else { + inner + }; + + Ok(ObjectStore { + inner, + scheme: String::from(base_path.scheme()), + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts, + list_is_lexically_ordered: !is_s3_express, + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count, + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } +} + +/// Check if the storage is S3 Express +fn check_s3_express(url: &Url, storage_options: &StorageOptions) -> bool { + storage_options + .0 + .get("s3_express") + .map(|v| v == "true") + .unwrap_or(false) + || url.authority().ends_with("--x-s3") +} + +/// Figure out the S3 region of the bucket. +/// +/// This resolves in order of precedence: +/// 1. The region provided in the storage options +/// 2. (If endpoint is not set), the region returned by the S3 API for the bucket +/// +/// It can return None if no region is provided and the endpoint is set. +async fn resolve_s3_region( + url: &Url, + storage_options: &HashMap, +) -> Result> { + if let Some(region) = storage_options.get(&AmazonS3ConfigKey::Region) { + Ok(Some(region.clone())) + } else if storage_options.get(&AmazonS3ConfigKey::Endpoint).is_none() { + // If no endpoint is set, we can assume this is AWS S3 and the region + // can be resolved from the bucket. + let bucket = url.host_str().ok_or_else(|| { + Error::invalid_input(format!("Could not parse bucket from url: {}", url)) + })?; + + let mut client_options = ClientOptions::default(); + for (key, value) in storage_options { + if let AmazonS3ConfigKey::Client(client_key) = key { + client_options = client_options.with_config(*client_key, value.clone()); + } + } + + let bucket_region = + object_store::aws::resolve_bucket_region(bucket, &client_options).await?; + Ok(Some(bucket_region)) + } else { + Ok(None) + } +} + +/// Selects which AWS credential provider to use for a dataset. +/// +/// When set, overrides automatic credential resolution for everything except an +/// explicitly-supplied `credentials` provider or `storage_options_accessor`. +#[derive(Debug, Clone, PartialEq)] +pub enum AwsProviderScheme { + /// Require static access-key credentials (`aws_access_key_id` + + /// `aws_secret_access_key`). Returns an error if they are absent. + Token, + /// Use the ECS/Pod Identity container credential endpoint. + /// The endpoint URI is read from the `AWS_CONTAINER_CREDENTIALS_FULL_URI` + /// or `AWS_CONTAINER_CREDENTIALS_RELATIVE_URI` environment variables. + Ecs, + /// Use IRSA (IAM Roles for Service Accounts) web identity token credentials. + /// The token file and role ARN are read from the `AWS_WEB_IDENTITY_TOKEN_FILE` + /// and `AWS_ROLE_ARN` environment variables. + Irsa, +} + +/// Build AWS credentials +/// +/// This resolves credentials from the following sources in order: +/// 1. An explicit `credentials` provider +/// 2. An explicit `storage_options_accessor` with a provider +/// 3. If `provider_scheme` is set: +/// - [`AwsProviderScheme::Token`]: static access-key credentials (error if absent) +/// - [`AwsProviderScheme::Ecs`]: ECS container credential provider +/// - [`AwsProviderScheme::Irsa`]: web identity token (IRSA) provider +/// 4. Static access-key credentials from `storage_options`, if present +/// 5. The default AWS credential provider chain +/// +/// # Storage Options Accessor +/// +/// When `storage_options_accessor` is provided and has a dynamic provider, +/// credentials are fetched and cached by the accessor with automatic refresh +/// before expiration. +/// +/// `credentials_refresh_offset` is the amount of time before expiry to refresh credentials. +pub async fn build_aws_credential( + credentials_refresh_offset: Duration, + credentials: Option, + storage_options: Option<&HashMap>, + region: Option, + storage_options_accessor: Option>, + provider_scheme: Option, +) -> Result<(AwsCredentialProvider, String)> { + use aws_config::meta::region::RegionProviderChain; + const DEFAULT_REGION: &str = "us-west-2"; + + let region = if let Some(region) = region { + region + } else { + RegionProviderChain::default_provider() + .or_else(DEFAULT_REGION) + .region() + .await + .map(|r| r.as_ref().to_string()) + .unwrap_or(DEFAULT_REGION.to_string()) + }; + + // If the user supplied their own credential provider that takes top priority + if let Some(creds) = credentials { + return Ok((creds, region)); + } + + // Otherwise, if the user provided a storage_options_accessor, try and use that + if let Some(dynamic_creds) = build_dynamic_credential_provider::( + storage_options_accessor.clone(), + ) + .await? + { + return Ok((dynamic_creds, region)); + } + + // If the user provided a storage_options_accessor, then it must not have matched AWS. + // Log a message and ignore it. + if storage_options_accessor + .as_ref() + .is_some_and(|a| a.has_provider()) + { + log::debug!( + "Storage options from provider do not contain explicit AWS credentials, \ + falling back to default AWS credentials chain." + ); + } + + // If the caller specified an explicit provider scheme, use only that provider. + if let Some(scheme) = provider_scheme { + return match scheme { + AwsProviderScheme::Token => { + let creds = storage_options + .and_then(extract_static_s3_credentials) + .ok_or_else(|| { + Error::invalid_input( + "aws_provider_scheme=token requires aws_access_key_id \ + and aws_secret_access_key to be set", + ) + })?; + Ok((Arc::new(creds), region)) + } + AwsProviderScheme::Ecs => { + let provider = EcsCredentialsProvider::builder().build(); + Ok(( + Arc::new(AwsCredentialAdapter::new( + Arc::new(provider), + credentials_refresh_offset, + )), + region, + )) + } + AwsProviderScheme::Irsa => { + let conf = ProviderConfig::default().with_region(Some(Region::new(region.clone()))); + let provider = WebIdentityTokenCredentialsProvider::builder() + .configure(&conf) + .build(); + Ok(( + Arc::new(AwsCredentialAdapter::new( + Arc::new(provider), + credentials_refresh_offset, + )), + region, + )) + } + }; + } + + if let Some(opts) = storage_options { + // Check for static credentials (access key & secret) + if let Some(creds) = extract_static_s3_credentials(opts) { + return Ok((Arc::new(creds), region)); + } + if opts.keys().any(is_aws_credential_key) { + return Err(Error::invalid_input( + "Explicit AWS credentials require both aws_access_key_id and aws_secret_access_key", + )); + } + } + + let credentials_provider = DefaultCredentialsChain::builder().build().await; + Ok(( + Arc::new(AwsCredentialAdapter::new( + Arc::new(credentials_provider), + credentials_refresh_offset, + )), + region, + )) +} + +fn extract_static_s3_credentials( + options: &HashMap, +) -> Option> { + let key_id = options.get(&AmazonS3ConfigKey::AccessKeyId).cloned(); + let secret_key = options.get(&AmazonS3ConfigKey::SecretAccessKey).cloned(); + let token = options.get(&AmazonS3ConfigKey::Token).cloned(); + match (key_id, secret_key, token) { + (Some(key_id), Some(secret_key), token) => { + Some(StaticCredentialProvider::new(ObjectStoreAwsCredential { + key_id, + secret_key, + token, + })) + } + _ => None, + } +} + +/// Adapt an AWS SDK cred into object_store credentials +#[derive(Debug)] +pub struct AwsCredentialAdapter { + pub inner: Arc, + + // RefCell can't be shared across threads, so we use HashMap + cache: Arc>>>, + + // The amount of time before expiry to refresh credentials + credentials_refresh_offset: Duration, +} + +impl AwsCredentialAdapter { + pub fn new( + provider: Arc, + credentials_refresh_offset: Duration, + ) -> Self { + Self { + inner: provider, + cache: Arc::new(RwLock::new(HashMap::new())), + credentials_refresh_offset, + } + } +} + +const AWS_CREDS_CACHE_KEY: &str = "aws_credentials"; + +/// Convert std::time::SystemTime from AWS SDK to our mockable SystemTime +fn to_system_time(time: std::time::SystemTime) -> SystemTime { + let duration_since_epoch = time + .duration_since(std::time::UNIX_EPOCH) + .expect("time should be after UNIX_EPOCH"); + UNIX_EPOCH + duration_since_epoch +} + +#[async_trait::async_trait] +impl CredentialProvider for AwsCredentialAdapter { + type Credential = ObjectStoreAwsCredential; + + async fn get_credential(&self) -> ObjectStoreResult> { + let cached_creds = { + let cache_value = self.cache.read().await.get(AWS_CREDS_CACHE_KEY).cloned(); + let expired = cache_value + .clone() + .map(|cred| { + cred.expiry() + .map(|exp| { + to_system_time(exp) + .checked_sub(self.credentials_refresh_offset) + .expect("this time should always be valid") + < SystemTime::now() + }) + // no expiry is never expire + .unwrap_or(false) + }) + .unwrap_or(true); // no cred is the same as expired; + if expired { None } else { cache_value.clone() } + }; + + if let Some(creds) = cached_creds { + Ok(Arc::new(Self::Credential { + key_id: creds.access_key_id().to_string(), + secret_key: creds.secret_access_key().to_string(), + token: creds.session_token().map(|s| s.to_string()), + })) + } else { + let refreshed_creds = + Arc::new(self.inner.provide_credentials().await.map_err(|e| { + Error::internal(format!("Failed to get AWS credentials: {:?}", e)) + })?); + + self.cache + .write() + .await + .insert(AWS_CREDS_CACHE_KEY.to_string(), refreshed_creds.clone()); + + Ok(Arc::new(Self::Credential { + key_id: refreshed_creds.access_key_id().to_string(), + secret_key: refreshed_creds.secret_access_key().to_string(), + token: refreshed_creds.session_token().map(|s| s.to_string()), + })) + } + } +} + +impl StorageOptions { + /// Add values from the environment to storage options. + /// + /// Only adds keys that are not already present, so explicitly-set options + /// (including empty-string sentinels) always take precedence over env vars. + pub fn with_env_s3(&mut self) { + for (os_key, os_value) in std::env::vars_os() { + if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) + && let Ok(config_key) = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) + && !self.0.contains_key(config_key.as_ref()) + { + self.0 + .insert(config_key.as_ref().to_string(), value.to_string()); + } + } + } + + /// Subset of options relevant for s3 storage + pub fn as_s3_options(&self) -> HashMap { + self.0 + .iter() + .filter_map(|(key, value)| { + let s3_key = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()).ok()?; + Some((s3_key, value.clone())) + }) + .collect() + } + + /// Parse the `aws_provider_scheme` storage option, if set. + pub fn aws_provider_scheme(&self) -> Result> { + match self.0.get("aws_provider_scheme").map(|s| s.as_str()) { + None | Some("") => Ok(None), + Some("token") => Ok(Some(AwsProviderScheme::Token)), + Some("ecs") => Ok(Some(AwsProviderScheme::Ecs)), + Some("irsa") => Ok(Some(AwsProviderScheme::Irsa)), + Some(other) => Err(Error::invalid_input(format!( + "Invalid aws_provider_scheme '{}'. Valid values are: token, ecs, irsa", + other + ))), + } + } +} + +impl ObjectStoreParams { + /// Create a new instance of [`ObjectStoreParams`] based on the AWS credentials. + pub fn with_aws_credentials( + aws_credentials: Option, + region: Option, + ) -> Self { + let storage_options_accessor = region.map(|region| { + let opts: HashMap = + [("region".into(), region)].iter().cloned().collect(); + Arc::new(StorageOptionsAccessor::with_static_options(opts)) + }); + Self { + aws_credentials, + storage_options_accessor, + ..Default::default() + } + } +} + +pub type DynamicStorageOptionsCredentialProvider = + NamespaceCredentialsProvider; + +#[cfg(test)] +mod tests { + use crate::object_store::ObjectStoreRegistry; + use crate::object_store::StorageOptionsProvider; + use mock_instant::thread_local::MockClock; + use object_store::path::Path; + use std::sync::atomic::{AtomicBool, Ordering}; + + use super::*; + + #[derive(Debug, Default)] + struct MockAwsCredentialsProvider { + called: AtomicBool, + } + + #[async_trait::async_trait] + impl CredentialProvider for MockAwsCredentialsProvider { + type Credential = ObjectStoreAwsCredential; + + async fn get_credential(&self) -> ObjectStoreResult> { + self.called.store(true, Ordering::Relaxed); + Ok(Arc::new(Self::Credential { + key_id: "".to_string(), + secret_key: "".to_string(), + token: None, + })) + } + } + + #[tokio::test] + async fn test_injected_aws_creds_option_is_used() { + let mock_provider = Arc::new(MockAwsCredentialsProvider::default()); + let registry = Arc::new(ObjectStoreRegistry::default()); + + let params = ObjectStoreParams { + aws_credentials: Some(mock_provider.clone() as AwsCredentialProvider), + ..ObjectStoreParams::default() + }; + + // Not called yet + assert!(!mock_provider.called.load(Ordering::Relaxed)); + + let (store, _) = ObjectStore::from_uri_and_params(registry, "s3://not-a-bucket", ¶ms) + .await + .unwrap(); + + // fails, but we don't care + let _ = store + .open(&Path::parse("/").unwrap()) + .await + .unwrap() + .get_range(0..1) + .await; + + // Not called yet + assert!(mock_provider.called.load(Ordering::Relaxed)); + } + + #[test] + fn test_s3_path_parsing() { + let provider = AwsStoreProvider; + + let cases = [ + ("s3://bucket/path/to/file", "path/to/file"), + // for non ASCII string tests: the URL encodes them, extract_path must decode back + ("s3://bucket/测试path/to/file", "测试path/to/file"), + ("s3://bucket/path/&to/file", "path/&to/file"), + ("s3://bucket/path/=to/file", "path/=to/file"), + ( + "s3+ddb://bucket/path/to/file?ddbTableName=test", + "path/to/file", + ), + ]; + + for (uri, expected_path) in cases { + let url = Url::parse(uri).unwrap(); + let path = provider.extract_path(&url).unwrap(); + // extract_path decodes url.path(), so the Path stores the raw (decoded) + // string. Path::parse keeps its input verbatim, matching that, whereas + // Path::from would percent-encode non-ASCII bytes and not match. + let expected_path = Path::parse(expected_path).unwrap(); + assert_eq!(path, expected_path) + } + } + + // Regression test for https://github.com/lance-format/lance/issues/6643 + // extract_path must NOT double-encode paths that contain non-ASCII characters. + // url.path() returns a percent-encoded string; we must decode it back to raw + // UTF-8 before storing it in a Path, so the object store HTTP client can apply + // a single, correct percent-encoding when building the request URL. + #[test] + fn test_s3_non_ascii_path_no_double_encoding() { + let provider = AwsStoreProvider; + + // "s3://bucket/中文路径" → url.path() == "/%E4%B8%AD%E6%96%87%E8%B7%AF%E5%BE%84". + // The buggy Path::parse(url.path()) stored "%E4%B8%AD..." verbatim; the S3 + // client then percent-encodes the '%' again, yielding "%25E4%25B8%25AD...". + // With Path::from_url_path the Path stores the decoded UTF-8 instead. + let url = Url::parse("s3://bucket/中文路径").unwrap(); + let path = provider.extract_path(&url).unwrap(); + + // The Path must hold the decoded UTF-8, not the percent-encoded form. + assert_eq!(path.as_ref(), "中文路径"); + } + + #[test] + fn test_is_s3_express() { + let cases = [ + ( + "s3://bucket/path/to/file", + HashMap::from([("s3_express".to_string(), "true".to_string())]), + true, + ), + ( + "s3://bucket/path/to/file", + HashMap::from([("s3_express".to_string(), "false".to_string())]), + false, + ), + ("s3://bucket/path/to/file", HashMap::from([]), false), + ( + "s3://bucket--x-s3/path/to/file", + HashMap::from([("s3_express".to_string(), "true".to_string())]), + true, + ), + ( + "s3://bucket--x-s3/path/to/file", + HashMap::from([("s3_express".to_string(), "false".to_string())]), + true, // URL takes precedence + ), + ("s3://bucket--x-s3/path/to/file", HashMap::from([]), true), + ]; + + for (uri, storage_map, expected) in cases { + let url = Url::parse(uri).unwrap(); + let storage_options = StorageOptions(storage_map); + let is_s3_express = check_s3_express(&url, &storage_options); + assert_eq!(is_s3_express, expected); + } + } + + #[tokio::test] + async fn test_use_opendal_flag() { + use crate::object_store::StorageOptionsAccessor; + let provider = AwsStoreProvider; + let url = Url::parse("s3://test-bucket/path").unwrap(); + let params_with_flag = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("use_opendal".to_string(), "true".to_string()), + ("region".to_string(), "us-west-2".to_string()), + ]), + ))), + ..Default::default() + }; + + let store = provider + .new_store(url.clone(), ¶ms_with_flag) + .await + .unwrap(); + assert_eq!(store.scheme, "s3"); + } + + #[derive(Debug)] + struct MockStorageOptionsProvider { + call_count: Arc>, + expires_in_millis: Option, + } + + impl MockStorageOptionsProvider { + fn new(expires_in_millis: Option) -> Self { + Self { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis, + } + } + + async fn get_call_count(&self) -> usize { + *self.call_count.read().await + } + } + + #[async_trait::async_trait] + impl StorageOptionsProvider for MockStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let count = { + let mut c = self.call_count.write().await; + *c += 1; + *c + }; + + let mut options = HashMap::from([ + ("aws_access_key_id".to_string(), format!("AKID_{}", count)), + ( + "aws_secret_access_key".to_string(), + format!("SECRET_{}", count), + ), + ("aws_session_token".to_string(), format!("TOKEN_{}", count)), + ]); + + if let Some(expires_in) = self.expires_in_millis { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + let expires_at = now_ms + expires_in; + options.insert("expires_at_millis".to_string(), expires_at.to_string()); + } + + Ok(Some(options)) + } + + fn provider_id(&self) -> String { + let ptr = Arc::as_ptr(&self.call_count) as usize; + format!("MockStorageOptionsProvider {{ id: {} }}", ptr) + } + } + + #[tokio::test] + async fn test_dynamic_credential_provider_with_initial_cache() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + + // Create a mock provider that returns credentials expiring in 10 minutes + let mock = Arc::new(MockStorageOptionsProvider::new(Some( + 600_000, // Expires in 10 minutes + ))); + + // Create initial options with cached credentials that expire in 10 minutes + let expires_at = now_ms + 600_000; // 10 minutes from now + let initial_options = HashMap::from([ + ("aws_access_key_id".to_string(), "AKID_CACHED".to_string()), + ( + "aws_secret_access_key".to_string(), + "SECRET_CACHED".to_string(), + ), + ("aws_session_token".to_string(), "TOKEN_CACHED".to_string()), + ("expires_at_millis".to_string(), expires_at.to_string()), + ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset + ]); + + let provider = DynamicStorageOptionsCredentialProvider::from_provider_with_initial( + mock.clone(), + initial_options, + ); + + // First call should use cached credentials (not expired yet) + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_CACHED"); + assert_eq!(cred.secret_key, "SECRET_CACHED"); + assert_eq!(cred.token, Some("TOKEN_CACHED".to_string())); + + // Should not have called the provider yet + assert_eq!(mock.get_call_count().await, 0); + } + + #[tokio::test] + async fn test_dynamic_credential_provider_with_expired_cache() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + + // Create a mock provider that returns credentials expiring in 10 minutes + let mock = Arc::new(MockStorageOptionsProvider::new(Some( + 600_000, // Expires in 10 minutes + ))); + + // Create initial options with credentials that expired 1 second ago + let expired_time = now_ms - 1_000; // 1 second ago + let initial_options = HashMap::from([ + ("aws_access_key_id".to_string(), "AKID_EXPIRED".to_string()), + ( + "aws_secret_access_key".to_string(), + "SECRET_EXPIRED".to_string(), + ), + ("expires_at_millis".to_string(), expired_time.to_string()), + ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset + ]); + + let provider = DynamicStorageOptionsCredentialProvider::from_provider_with_initial( + mock.clone(), + initial_options, + ); + + // First call should fetch new credentials because cached ones are expired + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(cred.secret_key, "SECRET_1"); + assert_eq!(cred.token, Some("TOKEN_1".to_string())); + + // Should have called the provider once + assert_eq!(mock.get_call_count().await, 1); + } + + #[tokio::test] + async fn test_dynamic_credential_provider_refresh_lead_time() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + // Create a mock provider that returns credentials expiring in 30 seconds + let mock = Arc::new(MockStorageOptionsProvider::new(Some( + 30_000, // Expires in 30 seconds + ))); + + // Create credential provider with default 60 second refresh offset + // This means credentials should be refreshed when they have less than 60 seconds left + let provider = DynamicStorageOptionsCredentialProvider::from_provider(mock.clone()); + + // First call should fetch credentials from provider (no initial cache) + // Credentials expire in 30 seconds, which is less than our 60 second refresh offset, + // so they should be considered "needs refresh" immediately + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(mock.get_call_count().await, 1); + + // Second call should trigger refresh because credentials expire in 30 seconds + // but our refresh lead time is 60 seconds (now + 60sec > expires_at) + // The mock will return new credentials (AKID_2) with the same expiration + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_2"); + assert_eq!(mock.get_call_count().await, 2); + } + + #[tokio::test] + async fn test_dynamic_credential_provider_no_initial_cache() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + // Create a mock provider that returns credentials expiring in 2 minutes + let mock = Arc::new(MockStorageOptionsProvider::new(Some( + 120_000, // Expires in 2 minutes + ))); + + // Create credential provider without initial cache, using default 60 second refresh offset + let provider = DynamicStorageOptionsCredentialProvider::from_provider(mock.clone()); + + // First call should fetch from provider (call count = 1) + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(cred.secret_key, "SECRET_1"); + assert_eq!(cred.token, Some("TOKEN_1".to_string())); + assert_eq!(mock.get_call_count().await, 1); + + // Second call should use cached credentials (not expired yet, still > 60 seconds remaining) + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(mock.get_call_count().await, 1); // Still 1, didn't fetch again + + // Advance time to 90 seconds - should trigger refresh (within 60 sec refresh offset) + // At this point, credentials expire in 30 seconds (< 60 sec offset) + MockClock::set_system_time(Duration::from_secs(100_000 + 90)); + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_2"); + assert_eq!(cred.secret_key, "SECRET_2"); + assert_eq!(cred.token, Some("TOKEN_2".to_string())); + assert_eq!(mock.get_call_count().await, 2); + + // Advance time to 210 seconds total (90 + 120) - should trigger another refresh + MockClock::set_system_time(Duration::from_secs(100_000 + 210)); + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_3"); + assert_eq!(cred.secret_key, "SECRET_3"); + assert_eq!(mock.get_call_count().await, 3); + } + + #[tokio::test] + async fn test_dynamic_credential_provider_with_initial_options() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + + // Create a mock provider that returns credentials expiring in 10 minutes + let mock = Arc::new(MockStorageOptionsProvider::new(Some( + 600_000, // Expires in 10 minutes + ))); + + // Create initial options with expiration in 10 minutes + let expires_at = now_ms + 600_000; // 10 minutes from now + let initial_options = HashMap::from([ + ("aws_access_key_id".to_string(), "AKID_INITIAL".to_string()), + ( + "aws_secret_access_key".to_string(), + "SECRET_INITIAL".to_string(), + ), + ("aws_session_token".to_string(), "TOKEN_INITIAL".to_string()), + ("expires_at_millis".to_string(), expires_at.to_string()), + ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset + ]); + + // Create credential provider with initial options + let provider = DynamicStorageOptionsCredentialProvider::from_provider_with_initial( + mock.clone(), + initial_options, + ); + + // First call should use the initial credential (not expired yet) + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_INITIAL"); + assert_eq!(cred.secret_key, "SECRET_INITIAL"); + assert_eq!(cred.token, Some("TOKEN_INITIAL".to_string())); + + // Should not have called the provider yet + assert_eq!(mock.get_call_count().await, 0); + + // Advance time to 6 minutes - this should trigger a refresh + // (5 minute refresh offset means we refresh 5 minutes before expiration) + MockClock::set_system_time(Duration::from_secs(100_000 + 360)); + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(cred.secret_key, "SECRET_1"); + assert_eq!(cred.token, Some("TOKEN_1".to_string())); + + // Should have called the provider once + assert_eq!(mock.get_call_count().await, 1); + + // Advance time to 11 minutes total - this should trigger another refresh + MockClock::set_system_time(Duration::from_secs(100_000 + 660)); + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_2"); + assert_eq!(cred.secret_key, "SECRET_2"); + assert_eq!(cred.token, Some("TOKEN_2".to_string())); + + // Should have called the provider twice + assert_eq!(mock.get_call_count().await, 2); + + // Advance time to 16 minutes total - this should trigger yet another refresh + MockClock::set_system_time(Duration::from_secs(100_000 + 960)); + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_3"); + assert_eq!(cred.secret_key, "SECRET_3"); + assert_eq!(cred.token, Some("TOKEN_3".to_string())); + + // Should have called the provider three times + assert_eq!(mock.get_call_count().await, 3); + } + + #[tokio::test] + async fn test_dynamic_credential_provider_concurrent_access() { + // Create a mock provider with far future expiration + let mock = Arc::new(MockStorageOptionsProvider::new(Some(9999999999999))); + + let provider = Arc::new(DynamicStorageOptionsCredentialProvider::from_provider( + mock.clone(), + )); + + // Spawn 10 concurrent tasks that all try to get credentials at the same time + let mut handles = vec![]; + for i in 0..10 { + let provider = provider.clone(); + let handle = tokio::spawn(async move { + let cred = provider.get_credential().await.unwrap(); + // Verify we got the correct credentials (should all be AKID_1 from first fetch) + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(cred.secret_key, "SECRET_1"); + assert_eq!(cred.token, Some("TOKEN_1".to_string())); + i // Return task number for verification + }); + handles.push(handle); + } + + // Wait for all tasks to complete + let results: Vec<_> = futures::future::join_all(handles) + .await + .into_iter() + .map(|r| r.unwrap()) + .collect(); + + // Verify all 10 tasks completed successfully + assert_eq!(results.len(), 10); + for i in 0..10 { + assert!(results.contains(&i)); + } + + // The provider should have been called exactly once (first request triggers fetch, + // subsequent requests use cache) + let call_count = mock.get_call_count().await; + assert_eq!( + call_count, 1, + "Provider should be called exactly once despite concurrent access" + ); + } + + #[tokio::test] + async fn test_dynamic_credential_provider_concurrent_refresh() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + + // Create initial options with credentials that expired in the past (1000 seconds ago) + let expires_at = now_ms - 1_000_000; + let initial_options = HashMap::from([ + ("aws_access_key_id".to_string(), "AKID_OLD".to_string()), + ( + "aws_secret_access_key".to_string(), + "SECRET_OLD".to_string(), + ), + ("aws_session_token".to_string(), "TOKEN_OLD".to_string()), + ("expires_at_millis".to_string(), expires_at.to_string()), + ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset + ]); + + // Mock will return credentials expiring in 1 hour + let mock = Arc::new(MockStorageOptionsProvider::new(Some( + 3_600_000, // Expires in 1 hour + ))); + + let provider = Arc::new( + DynamicStorageOptionsCredentialProvider::from_provider_with_initial( + mock.clone(), + initial_options, + ), + ); + + // Spawn 20 concurrent tasks that all try to get credentials at the same time + // Since the initial credential is expired, they'll all try to refresh + let mut handles = vec![]; + for i in 0..20 { + let provider = provider.clone(); + let handle = tokio::spawn(async move { + let cred = provider.get_credential().await.unwrap(); + // All should get the new credentials (AKID_1 from first fetch) + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(cred.secret_key, "SECRET_1"); + assert_eq!(cred.token, Some("TOKEN_1".to_string())); + i + }); + handles.push(handle); + } + + // Wait for all tasks to complete + let results: Vec<_> = futures::future::join_all(handles) + .await + .into_iter() + .map(|r| r.unwrap()) + .collect(); + + // Verify all 20 tasks completed successfully + assert_eq!(results.len(), 20); + + // The provider should have been called at least once, but possibly more times + // due to the try_write mechanism and race conditions + let call_count = mock.get_call_count().await; + assert!( + call_count >= 1, + "Provider should be called at least once, was called {} times", + call_count + ); + + // It shouldn't be called 20 times though - the lock should prevent most concurrent fetches + assert!( + call_count < 10, + "Provider should not be called too many times due to lock contention, was called {} times", + call_count + ); + } + + #[tokio::test] + async fn test_explicit_aws_credentials_takes_precedence_over_accessor() { + // Create a mock storage options provider that should NOT be called + let mock_storage_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000))); + + // Create an accessor with the mock provider + let accessor = Arc::new(StorageOptionsAccessor::with_provider( + mock_storage_provider.clone(), + )); + + // Create an explicit AWS credentials provider + let explicit_cred_provider = Arc::new(MockAwsCredentialsProvider::default()); + + // Build credentials with both aws_credentials AND accessor + // The explicit aws_credentials should take precedence + let (result, _region) = build_aws_credential( + Duration::from_secs(300), + Some(explicit_cred_provider.clone() as AwsCredentialProvider), + None, // no storage_options + Some("us-west-2".to_string()), + Some(accessor), + None, + ) + .await + .unwrap(); + + // Get credential from the result + let cred = result.get_credential().await.unwrap(); + + // The explicit provider should have been called (it returns empty strings) + assert!(explicit_cred_provider.called.load(Ordering::Relaxed)); + + // The storage options provider should NOT have been called + assert_eq!( + mock_storage_provider.get_call_count().await, + 0, + "Storage options provider should not be called when explicit aws_credentials is provided" + ); + + // Verify we got credentials from the explicit provider (empty strings) + assert_eq!(cred.key_id, ""); + assert_eq!(cred.secret_key, ""); + } + + #[tokio::test] + async fn test_accessor_used_when_no_explicit_aws_credentials() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + + // Create a mock storage options provider + let mock_storage_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000))); + + // Create initial options + let expires_at = now_ms + 600_000; // 10 minutes from now + let initial_options = HashMap::from([ + ( + "aws_access_key_id".to_string(), + "AKID_FROM_ACCESSOR".to_string(), + ), + ( + "aws_secret_access_key".to_string(), + "SECRET_FROM_ACCESSOR".to_string(), + ), + ( + "aws_session_token".to_string(), + "TOKEN_FROM_ACCESSOR".to_string(), + ), + ("expires_at_millis".to_string(), expires_at.to_string()), + ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset + ]); + + // Create an accessor with initial options and provider + let accessor = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial_options, + mock_storage_provider.clone(), + )); + + // Build credentials with accessor but NO explicit aws_credentials + let (result, _region) = build_aws_credential( + Duration::from_secs(300), + None, // no explicit aws_credentials + None, // no storage_options + Some("us-west-2".to_string()), + Some(accessor), + None, + ) + .await + .unwrap(); + + // Get credential - should use the initial accessor credentials + let cred = result.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_FROM_ACCESSOR"); + assert_eq!(cred.secret_key, "SECRET_FROM_ACCESSOR"); + + // Storage options provider should NOT have been called yet (using cached initial creds) + assert_eq!(mock_storage_provider.get_call_count().await, 0); + + // Advance time to trigger refresh (past the 5 minute refresh offset) + MockClock::set_system_time(Duration::from_secs(100_000 + 360)); + + // Get credential again - should now fetch from provider + let cred = result.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID_1"); + assert_eq!(cred.secret_key, "SECRET_1"); + + // Storage options provider should have been called once + assert_eq!(mock_storage_provider.get_call_count().await, 1); + } + + // Test that aws_provider_scheme=token selects static credentials. + #[tokio::test] + async fn test_provider_scheme_token() { + let opts = HashMap::from([ + (AmazonS3ConfigKey::AccessKeyId, "AKID".to_string()), + (AmazonS3ConfigKey::SecretAccessKey, "SECRET".to_string()), + ]); + + let (provider, _) = build_aws_credential( + Duration::from_secs(300), + None, + Some(&opts), + Some("us-east-1".to_string()), + None, + Some(AwsProviderScheme::Token), + ) + .await + .unwrap(); + + let cred = provider.get_credential().await.unwrap(); + assert_eq!(cred.key_id, "AKID"); + assert_eq!(cred.secret_key, "SECRET"); + } + + // Test that aws_provider_scheme=token errors when no static credentials are present. + #[tokio::test] + async fn test_provider_scheme_token_errors_without_credentials() { + let opts: HashMap = HashMap::new(); + + let result = build_aws_credential( + Duration::from_secs(300), + None, + Some(&opts), + Some("us-east-1".to_string()), + None, + Some(AwsProviderScheme::Token), + ) + .await; + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("aws_provider_scheme=token"), + "error should mention aws_provider_scheme=token" + ); + } + + // Test that aws_provider_scheme=ecs builds a provider without error. + // The ECS provider itself reads from env vars lazily; construction always succeeds. + #[tokio::test] + async fn test_provider_scheme_ecs() { + let opts: HashMap = HashMap::new(); + + let result = build_aws_credential( + Duration::from_secs(300), + None, + Some(&opts), + Some("us-east-1".to_string()), + None, + Some(AwsProviderScheme::Ecs), + ) + .await; + assert!(result.is_ok(), "ECS provider should build without error"); + } + + // Test that aws_provider_scheme=irsa builds a provider and attempts credential + // retrieval (which fails with a provider error, not a config error like + // "Missing Region" — confirming the region is wired through to the STS client). + #[tokio::test] + async fn test_provider_scheme_irsa() { + let opts: HashMap = HashMap::new(); + + let (provider, _) = build_aws_credential( + Duration::from_secs(300), + None, + Some(&opts), + Some("us-east-1".to_string()), + None, + Some(AwsProviderScheme::Irsa), + ) + .await + .unwrap(); + + // Credential retrieval must fail with a provider error (missing env vars or + // network), NOT a configuration error like "Invalid Configuration: Missing Region". + let err = provider.get_credential().await.unwrap_err(); + assert!( + !err.to_string().contains("Missing Region"), + "should not fail with Missing Region; region was provided. got: {err}" + ); + } + + // Test that an invalid aws_provider_scheme value produces a clear error. + #[test] + fn test_provider_scheme_invalid_value() { + let opts = StorageOptions::new(HashMap::from([( + "aws_provider_scheme".to_string(), + "magic".to_string(), + )])); + let result = opts.aws_provider_scheme(); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("magic")); + } + + // Test that no aws_provider_scheme falls through to DefaultCredentialsChain without error. + #[tokio::test] + async fn test_no_provider_scheme_uses_default_chain() { + let opts: HashMap = HashMap::new(); + + let result = build_aws_credential( + Duration::from_secs(300), + None, + Some(&opts), + Some("us-east-1".to_string()), + None, + None, + ) + .await; + assert!(result.is_ok()); + } +} diff --git a/vendor/lance-io/src/object_store/providers/azure.rs b/vendor/lance-io/src/object_store/providers/azure.rs new file mode 100644 index 000000000..2d407cd6d --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/azure.rs @@ -0,0 +1,674 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + collections::HashMap, + str::FromStr, + sync::{Arc, LazyLock}, + time::Duration, +}; + +use object_store::ObjectStore as OSObjectStore; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Azblob, services::Azdls}; + +use object_store::{ + RetryConfig, + azure::{AzureConfigKey, AzureCredential, MicrosoftAzureBuilder}, +}; +use url::Url; + +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, + dynamic_credentials::build_dynamic_credential_provider, + throttle::{AimdThrottleConfig, AimdThrottleState, AimdThrottledStore, cloud_http_connector}, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct AzureBlobStoreProvider; + +impl AzureBlobStoreProvider { + /// Normalize Azure storage options for OpenDAL, resolving aliases for + /// well-known keys while passing through all other options (e.g. + /// `client_id`, `tenant_id`, `encryption_key`, etc.) so that OpenDAL + /// can use them directly. + fn normalize_opendal_azure_options( + options: &HashMap, + ) -> HashMap { + // Start with all options so unknown keys are forwarded to OpenDAL. + let mut config_map = options.clone(); + + // Normalize well-known aliases into canonical OpenDAL key names. + // Remove the alias after resolving to avoid duplicate/conflicting entries. + let alias_groups: &[(&str, &[&str])] = &[ + ("account_name", &["azure_storage_account_name"]), + ("endpoint", &["azure_storage_endpoint", "azure_endpoint"]), + ( + "account_key", + &[ + "azure_storage_account_key", + "azure_storage_access_key", + "azure_storage_master_key", + "access_key", + "master_key", + ], + ), + ( + "sas_token", + &[ + "azure_storage_sas_token", + "azure_storage_sas_key", + "sas_key", + ], + ), + ]; + + for (canonical, aliases) in alias_groups { + if !config_map.contains_key(*canonical) { + for alias in *aliases { + if let Some(value) = config_map.remove(*alias) { + config_map.insert(canonical.to_string(), value); + break; + } + } + } else { + // Canonical key exists; remove aliases to avoid conflicts. + for alias in *aliases { + config_map.remove(*alias); + } + } + } + + config_map + } + + fn build_opendal_operator( + base_path: &Url, + storage_options: &StorageOptions, + ) -> Result { + // Start with all storage options as the config map + // OpenDAL will handle environment variables through its default credentials chain + let mut config_map = Self::normalize_opendal_azure_options(&storage_options.0); + + match base_path.scheme() { + "az" => { + let container = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("Azure URL must contain container name"))? + .to_string(); + + config_map.insert("container".to_string(), container); + + let prefix = base_path.path().trim_start_matches('/'); + if !prefix.is_empty() { + config_map.insert("root".to_string(), format!("/{}", prefix)); + } + + Operator::from_iter::(config_map).map_err(|e| { + Error::invalid_input(format!("Failed to create Azure Blob operator: {:?}", e)) + }) + } + "abfss" => { + let filesystem = base_path.username(); + if filesystem.is_empty() { + return Err(Error::invalid_input( + "abfss:// URL must include account: abfss://@.dfs.core.windows.net/path", + )); + } + let host = base_path.host_str().ok_or_else(|| { + Error::invalid_input( + "abfss:// URL must include account: abfss://@.dfs.core.windows.net/path" + ) + })?; + + config_map.insert("filesystem".to_string(), filesystem.to_string()); + config_map.insert("endpoint".to_string(), format!("https://{}", host)); + config_map + .entry("account_name".to_string()) + .or_insert_with(|| host.split('.').next().unwrap_or(host).to_string()); + + let root_path = base_path.path().trim_start_matches('/'); + if !root_path.is_empty() { + config_map.insert("root".to_string(), format!("/{}", root_path)); + } + + Operator::from_iter::(config_map).map_err(|e| { + Error::invalid_input(format!( + "Failed to create Azure DFS (ADLS Gen2) operator: {:?}", + e + )) + }) + } + _ => Err(Error::invalid_input(format!( + "Unsupported Azure scheme: {}", + base_path.scheme() + ))), + } + } + + async fn build_opendal_azure_store( + &self, + base_path: &Url, + storage_options: &StorageOptions, + ) -> Result> { + let operator = Self::build_opendal_operator(base_path, storage_options)?; + Ok(Arc::new(OpendalStore::new(operator))) + } + + async fn build_microsoft_azure_store( + &self, + base_path: &Url, + storage_options: &StorageOptions, + accessor: Option>, + throttle_state: Option<&AimdThrottleState>, + ) -> Result> { + // Use a low retry count since the AIMD throttle layer handles + // throttle recovery with its own retry loop. + let retry_config = RetryConfig { + backoff: Default::default(), + max_retries: storage_options.client_max_retries(), + retry_timeout: Duration::from_secs(storage_options.client_retry_timeout()), + }; + + let mut builder = MicrosoftAzureBuilder::new() + .with_url(base_path.as_ref()) + .with_retry(retry_config) + .with_client_options(storage_options.client_options()?); + for (key, value) in storage_options.as_azure_options() { + builder = builder.with_config(key, value); + } + + if let Some(credentials) = + build_dynamic_credential_provider::(accessor).await? + { + builder = builder.with_credentials(credentials); + } + + let store_prefix = + self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?; + builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix)); + + Ok(Arc::new(builder.build()?) as Arc) + } + + fn calculate_object_store_prefix_with_env( + url: &Url, + storage_options: Option<&HashMap>, + env_options: &HashMap, + ) -> Result { + let authority = url.authority(); + let (container, account) = match authority.find("@") { + Some(at_index) => { + // The URI has an: + // - az:// schema type and is similar to 'az://container@account.dfs.core.windows.net/path-part/file + // or possibly 'az://container@account/path-part/file' (the short version). + // - abfss:// schema type and is similar to 'abfss://filesystem@account.dfs.core.windows.net/path-part/file'. + let container = &authority[..at_index]; + let account = &authority[at_index + 1..]; + ( + container, + account.split(".").next().unwrap_or_default().to_string(), + ) + } + None => { + // The URI looks like 'az://container/path-part/file'. + // We must look at the storage options to find the account. + let mut account = match storage_options { + Some(opts) => StorageOptions::find_configured_storage_account(opts), + None => None, + }; + if account.is_none() { + account = StorageOptions::find_configured_storage_account(env_options); + } + let account = account.ok_or(Error::invalid_input("Unable to find object store prefix: no Azure account name in URI, and no storage account configured."))?; + (authority, account) + } + }; + Ok(format!("{}${}@{}", url.scheme(), container, account)) + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for AzureBlobStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let scheme = base_path.scheme().to_string(); + if scheme != "az" && scheme != "abfss" { + return Err(Error::invalid_input(format!( + "Unsupported Azure scheme '{}', expected 'az' or 'abfss'", + scheme + ))); + } + + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let mut storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); + storage_options.with_env_azure(); + let download_retry_count = storage_options.download_retry_count(); + + let use_opendal = storage_options + .0 + .get("use_opendal") + .map(|v| v.as_str() == "true") + .unwrap_or(false); + + let accessor = params.get_accessor(); + + let throttle_config = AimdThrottleConfig::from_storage_options(params.storage_options())?; + let throttle_state = if throttle_config.is_disabled() { + None + } else { + Some(AimdThrottleState::new(throttle_config)?) + }; + + let inner: Arc = if use_opendal { + // OpenDAL Azure intentionally uses static/environment-backed configuration only. + // Namespace-vended dynamic credentials are supported on the native object_store path. + self.build_opendal_azure_store(&base_path, &storage_options) + .await? + } else { + self.build_microsoft_azure_store( + &base_path, + &storage_options, + accessor, + throttle_state.as_ref(), + ) + .await? + }; + let inner = if let Some(throttle_state) = throttle_state { + Arc::new(AimdThrottledStore::new_with_state( + inner, + throttle_state, + !use_opendal, + )) as Arc + } else { + inner + }; + + Ok(ObjectStore { + inner, + scheme, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: false, + list_is_lexically_ordered: true, + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count, + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } + + fn calculate_object_store_prefix( + &self, + url: &Url, + storage_options: Option<&HashMap>, + ) -> Result { + Self::calculate_object_store_prefix_with_env(url, storage_options, &ENV_OPTIONS.0) + } +} + +static ENV_OPTIONS: LazyLock = LazyLock::new(StorageOptions::from_env); + +impl StorageOptions { + /// Iterate over all environment variables, looking for anything related to Azure. + fn from_env() -> Self { + let mut opts = HashMap::::new(); + for (os_key, os_value) in std::env::vars_os() { + if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) + && let Ok(config_key) = AzureConfigKey::from_str(&key.to_ascii_lowercase()) + { + opts.insert(config_key.as_ref().to_string(), value.to_string()); + } + } + Self(opts) + } + + /// Add values from the environment to storage options + pub fn with_env_azure(&mut self) { + for (os_key, os_value) in &ENV_OPTIONS.0 { + if !self.0.contains_key(os_key) { + self.0.insert(os_key.clone(), os_value.clone()); + } + } + } + + /// Subset of options relevant for azure storage + pub fn as_azure_options(&self) -> HashMap { + self.0 + .iter() + .filter_map(|(key, value)| { + let az_key = AzureConfigKey::from_str(&key.to_ascii_lowercase()).ok()?; + Some((az_key, value.clone())) + }) + .collect() + } + + #[allow(clippy::manual_map)] + fn find_configured_storage_account(map: &HashMap) -> Option { + if let Some(account) = map.get("azure_storage_account_name") { + Some(account.clone()) + } else if let Some(account) = map.get("account_name") { + Some(account.clone()) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + use crate::object_store::{ObjectStoreParams, StorageOptionsAccessor}; + use std::collections::HashMap; + + #[test] + fn test_azure_store_path() { + let provider = AzureBlobStoreProvider; + + let url = Url::parse("az://bucket/path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = object_store::path::Path::from("path/to/file"); + assert_eq!(path, expected_path); + } + + #[tokio::test] + async fn test_use_opendal_flag() { + let provider = AzureBlobStoreProvider; + let url = Url::parse("az://test-container/path").unwrap(); + let params_with_flag = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("use_opendal".to_string(), "true".to_string()), + ("account_name".to_string(), "test_account".to_string()), + ( + "endpoint".to_string(), + "https://test_account.blob.core.windows.net".to_string(), + ), + ( + "account_key".to_string(), + "dGVzdF9hY2NvdW50X2tleQ==".to_string(), + ), + ]), + ))), + ..Default::default() + }; + + let store = provider + .new_store(url.clone(), ¶ms_with_flag) + .await + .unwrap(); + assert_eq!(store.scheme, "az"); + let inner_desc = store.inner.to_string(); + assert!( + inner_desc.contains("Opendal") && inner_desc.contains("azblob"), + "az:// with use_opendal=true should use OpenDAL Azblob, got: {}", + inner_desc + ); + } + + #[tokio::test] + async fn test_dynamic_azure_credentials_provider() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([( + "azure_storage_sas_token".to_string(), + "?sv=2022-11-02&sp=rl&sig=test".to_string(), + )]), + }, + ))); + + let credentials = build_dynamic_credential_provider::(Some(accessor)) + .await + .expect("dynamic azure credentials should build") + .expect("expected credential provider") + .get_credential() + .await + .expect("expected azure credential"); + + match credentials.as_ref() { + AzureCredential::SASToken(pairs) => { + assert!( + pairs + .iter() + .any(|(key, value)| key == "sig" && value == "test") + ); + } + other => panic!("expected SAS token, got {other:?}"), + } + } + + #[test] + fn test_find_configured_storage_account() { + assert_eq!( + Some("myaccount".to_string()), + StorageOptions::find_configured_storage_account(&HashMap::from_iter( + [ + ("access_key".to_string(), "myaccesskey".to_string()), + ( + "azure_storage_account_name".to_string(), + "myaccount".to_string() + ) + ] + .into_iter() + )) + ); + } + + #[test] + fn test_calculate_object_store_prefix_from_url_and_options() { + let provider = AzureBlobStoreProvider; + let options = HashMap::from_iter([("account_name".to_string(), "bob".to_string())]); + assert_eq!( + "az$container@bob", + provider + .calculate_object_store_prefix( + &Url::parse("az://container/path").unwrap(), + Some(&options) + ) + .unwrap() + ); + } + + #[test] + fn test_calculate_object_store_prefix_from_url_and_ignored_options() { + let provider = AzureBlobStoreProvider; + let options = HashMap::from_iter([("account_name".to_string(), "bob".to_string())]); + assert_eq!( + "az$container@account", + provider + .calculate_object_store_prefix( + &Url::parse("az://container@account.dfs.core.windows.net/path").unwrap(), + Some(&options) + ) + .unwrap() + ); + } + + #[test] + fn test_calculate_object_store_prefix_from_url_short_account() { + let provider = AzureBlobStoreProvider; + let options = HashMap::from_iter([("account_name".to_string(), "bob".to_string())]); + assert_eq!( + "az$container@account", + provider + .calculate_object_store_prefix( + &Url::parse("az://container@account/path").unwrap(), + Some(&options) + ) + .unwrap() + ); + } + + #[test] + fn test_fail_to_calculate_object_store_prefix_from_url() { + let options = HashMap::from_iter([("access_key".to_string(), "myaccesskey".to_string())]); + let expected = "Invalid user input: Unable to find object store prefix: no Azure account name in URI, and no storage account configured."; + let result = AzureBlobStoreProvider::calculate_object_store_prefix_with_env( + &Url::parse("az://container/path").unwrap(), + Some(&options), + &HashMap::new(), + ) + .expect_err("expected error") + .to_string(); + assert_eq!(expected, &result[..expected.len()]); + } + + // --- abfss:// tests --- + + #[test] + fn test_abfss_extract_path() { + let provider = AzureBlobStoreProvider; + let url = Url::parse("abfss://myfs@myaccount.dfs.core.windows.net/path/to/dataset.lance") + .unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!( + path, + object_store::path::Path::from("path/to/dataset.lance") + ); + } + + #[test] + fn test_calculate_abfss_prefix() { + let provider = AzureBlobStoreProvider; + let url = Url::parse("abfss://myfs@myaccount.dfs.core.windows.net/path/to/data").unwrap(); + let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + assert_eq!(prefix, "abfss$myfs@myaccount"); + } + + #[test] + fn test_calculate_abfss_prefix_ignores_storage_options() { + let provider = AzureBlobStoreProvider; + let options = + HashMap::from_iter([("account_name".to_string(), "other_account".to_string())]); + let url = Url::parse("abfss://myfs@myaccount.dfs.core.windows.net/path").unwrap(); + let prefix = provider + .calculate_object_store_prefix(&url, Some(&options)) + .unwrap(); + assert_eq!(prefix, "abfss$myfs@myaccount"); + } + + #[tokio::test] + async fn test_abfss_default_uses_microsoft_builder() { + use crate::object_store::StorageOptionsAccessor; + let provider = AzureBlobStoreProvider; + let url = Url::parse("abfss://testfs@testaccount.dfs.core.windows.net/data").unwrap(); + let params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("account_name".to_string(), "testaccount".to_string()), + ("account_key".to_string(), "dGVzdA==".to_string()), + ]), + ))), + ..Default::default() + }; + + let store = provider.new_store(url, ¶ms).await.unwrap(); + assert_eq!(store.scheme, "abfss"); + assert!(!store.is_local()); + assert!(store.is_cloud()); + let inner_desc = store.inner.to_string(); + assert!( + inner_desc.contains("MicrosoftAzure"), + "abfss:// without use_opendal should use MicrosoftAzureBuilder, got: {}", + inner_desc + ); + } + + #[tokio::test] + async fn test_unsupported_scheme_rejected() { + use crate::object_store::StorageOptionsAccessor; + let provider = AzureBlobStoreProvider; + let url = Url::parse("wasbs://container@myaccount.blob.core.windows.net/path").unwrap(); + let params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("account_name".to_string(), "myaccount".to_string()), + ("account_key".to_string(), "dGVzdA==".to_string()), + ]), + ))), + ..Default::default() + }; + + let err = provider + .new_store(url, ¶ms) + .await + .expect_err("expected error for unsupported scheme"); + assert!( + err.to_string().contains("Unsupported Azure scheme"), + "unexpected error: {}", + err + ); + } + + #[tokio::test] + async fn test_abfss_with_opendal_uses_azdls() { + use crate::object_store::StorageOptionsAccessor; + let provider = AzureBlobStoreProvider; + let url = Url::parse("abfss://testfs@testaccount.dfs.core.windows.net/data").unwrap(); + let params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("use_opendal".to_string(), "true".to_string()), + ("account_name".to_string(), "testaccount".to_string()), + ("account_key".to_string(), "dGVzdA==".to_string()), + ]), + ))), + ..Default::default() + }; + + let store = provider.new_store(url, ¶ms).await.unwrap(); + assert_eq!(store.scheme, "abfss"); + assert!(!store.is_local()); + assert!(store.is_cloud()); + let inner_desc = store.inner.to_string(); + assert!( + inner_desc.contains("Opendal") && inner_desc.contains("azdls"), + "abfss:// with use_opendal=true should use OpenDAL Azdls, got: {}", + inner_desc + ); + } + + #[test] + fn test_azdls_capabilities_differ_from_azblob() { + let common_opts = StorageOptions(HashMap::from([ + ("account_name".to_string(), "testaccount".to_string()), + ("account_key".to_string(), "dGVzdA==".to_string()), + ( + "endpoint".to_string(), + "https://testaccount.blob.core.windows.net".to_string(), + ), + ])); + + // Build az:// operator (uses Azblob backend) + let az_url = Url::parse("az://test-container/path").unwrap(); + let az_operator = + AzureBlobStoreProvider::build_opendal_operator(&az_url, &common_opts).unwrap(); + + // Build abfss:// operator (uses Azdls backend) + let abfss_url = Url::parse("abfss://testfs@testaccount.dfs.core.windows.net/data").unwrap(); + let abfss_operator = + AzureBlobStoreProvider::build_opendal_operator(&abfss_url, &common_opts).unwrap(); + + let azblob_cap = az_operator.info().capability(); + let azdls_cap = abfss_operator.info().capability(); + + // Both support basic operations + assert!(azblob_cap.read); + assert!(azdls_cap.read); + assert!(azblob_cap.write); + assert!(azdls_cap.write); + assert!(azblob_cap.list); + assert!(azdls_cap.list); + + // Azdls supports rename and create_dir (HNS features); Azblob does not + assert!(azdls_cap.rename, "Azdls should support rename"); + assert!(azdls_cap.create_dir, "Azdls should support create_dir"); + assert!(!azblob_cap.rename, "Azblob should not support rename"); + } +} diff --git a/vendor/lance-io/src/object_store/providers/gcp.rs b/vendor/lance-io/src/object_store/providers/gcp.rs new file mode 100644 index 000000000..1a93c3ac9 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/gcp.rs @@ -0,0 +1,262 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, str::FromStr, sync::Arc, time::Duration}; + +use object_store::ObjectStore as OSObjectStore; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Gcs}; + +use object_store::{ + RetryConfig, StaticCredentialProvider, + gcp::{GcpCredential, GoogleCloudStorageBuilder, GoogleConfigKey}, +}; +use url::Url; + +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, + dynamic_credentials::build_dynamic_credential_provider, + throttle::{AimdThrottleConfig, AimdThrottleState, AimdThrottledStore, cloud_http_connector}, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct GcsStoreProvider; + +impl GcsStoreProvider { + async fn build_opendal_gcs_store( + &self, + base_path: &Url, + storage_options: &StorageOptions, + ) -> Result> { + let bucket = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("GCS URL must contain bucket name"))? + .to_string(); + + let prefix = base_path.path().trim_start_matches('/').to_string(); + + // Start with all storage options as the config map + // OpenDAL will handle environment variables through its default credentials chain + let mut config_map: HashMap = storage_options.0.clone(); + + // Set required OpenDAL configuration + config_map.insert("bucket".to_string(), bucket); + + if !prefix.is_empty() { + config_map.insert("root".to_string(), format!("/{}", prefix)); + } + + let operator = Operator::from_iter::(config_map) + .map_err(|e| Error::invalid_input(format!("Failed to create GCS operator: {:?}", e)))?; + + Ok(Arc::new(OpendalStore::new(operator)) as Arc) + } + + async fn build_google_cloud_store( + &self, + base_path: &Url, + storage_options: &StorageOptions, + accessor: Option>, + throttle_state: Option<&AimdThrottleState>, + ) -> Result> { + // Use a low retry count since the AIMD throttle layer handles + // throttle recovery with its own retry loop. + let retry_config = RetryConfig { + backoff: Default::default(), + max_retries: storage_options.client_max_retries(), + retry_timeout: Duration::from_secs(storage_options.client_retry_timeout()), + }; + + let mut builder = GoogleCloudStorageBuilder::new() + .with_url(base_path.as_ref()) + .with_retry(retry_config) + .with_client_options(storage_options.client_options()?); + for (key, value) in storage_options.as_gcs_options() { + builder = builder.with_config(key, value); + } + + if let Some(credentials) = + build_dynamic_credential_provider::(accessor).await? + { + builder = builder.with_credentials(credentials); + } else if let Some(storage_token) = storage_options.get("google_storage_token") { + let credential = GcpCredential { + bearer: storage_token.clone(), + }; + let credential_provider = Arc::new(StaticCredentialProvider::new(credential)) as _; + builder = builder.with_credentials(credential_provider); + } + + let store_prefix = + self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?; + builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix)); + + Ok(Arc::new(builder.build()?) as Arc) + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for GcsStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let mut storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); + storage_options.with_env_gcs(); + let download_retry_count = storage_options.download_retry_count(); + + let use_opendal = storage_options + .0 + .get("use_opendal") + .map(|v| v.as_str() == "true") + .unwrap_or(false); + + let accessor = params.get_accessor(); + + let throttle_config = AimdThrottleConfig::from_storage_options(params.storage_options())?; + let throttle_state = if throttle_config.is_disabled() { + None + } else { + Some(AimdThrottleState::new(throttle_config)?) + }; + + let inner = if use_opendal { + // OpenDAL GCS intentionally uses static/environment-backed configuration only. + // Namespace-vended dynamic credentials are supported on the native object_store path. + self.build_opendal_gcs_store(&base_path, &storage_options) + .await? + } else { + self.build_google_cloud_store( + &base_path, + &storage_options, + accessor, + throttle_state.as_ref(), + ) + .await? + }; + let inner = if let Some(throttle_state) = throttle_state { + Arc::new(AimdThrottledStore::new_with_state( + inner, + throttle_state, + !use_opendal, + )) as Arc + } else { + inner + }; + + Ok(ObjectStore { + inner, + scheme: String::from("gs"), + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: false, + list_is_lexically_ordered: true, + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count, + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } +} + +impl StorageOptions { + /// Add values from the environment to storage options + pub fn with_env_gcs(&mut self) { + for (os_key, os_value) in std::env::vars_os() { + if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) { + let lowercase_key = key.to_ascii_lowercase(); + let token_key = "google_storage_token"; + + if let Ok(config_key) = GoogleConfigKey::from_str(&lowercase_key) { + if !self.0.contains_key(config_key.as_ref()) { + self.0 + .insert(config_key.as_ref().to_string(), value.to_string()); + } + } + // Check for GOOGLE_STORAGE_TOKEN until GoogleConfigKey supports storage token + else if lowercase_key == token_key && !self.0.contains_key(token_key) { + self.0.insert(token_key.to_string(), value.to_string()); + } + } + } + } + + /// Subset of options relevant for gcs storage + pub fn as_gcs_options(&self) -> HashMap { + self.0 + .iter() + .filter_map(|(key, value)| { + let gcs_key = GoogleConfigKey::from_str(&key.to_ascii_lowercase()).ok()?; + Some((gcs_key, value.clone())) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + use crate::object_store::{ObjectStoreParams, StorageOptionsAccessor}; + use std::collections::HashMap; + + #[test] + fn test_gcs_store_path() { + let provider = GcsStoreProvider; + + let url = Url::parse("gs://bucket/path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = object_store::path::Path::from("path/to/file"); + assert_eq!(path, expected_path); + } + + #[tokio::test] + async fn test_use_opendal_flag() { + let provider = GcsStoreProvider; + let url = Url::parse("gs://test-bucket/path").unwrap(); + let params_with_flag = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("use_opendal".to_string(), "true".to_string()), + ( + "service_account".to_string(), + "test@example.iam.gserviceaccount.com".to_string(), + ), + ]), + ))), + ..Default::default() + }; + + let store = provider + .new_store(url.clone(), ¶ms_with_flag) + .await + .unwrap(); + assert_eq!(store.scheme, "gs"); + } + + #[tokio::test] + async fn test_dynamic_gcp_credentials_provider() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([( + "google_storage_token".to_string(), + "gcp-token".to_string(), + )]), + }, + ))); + + let credentials = build_dynamic_credential_provider::(Some(accessor)) + .await + .expect("dynamic gcp credentials should build") + .expect("expected credential provider") + .get_credential() + .await + .expect("expected gcp credential"); + + assert_eq!(credentials.bearer, "gcp-token"); + } +} diff --git a/vendor/lance-io/src/object_store/providers/goosefs.rs b/vendor/lance-io/src/object_store/providers/goosefs.rs new file mode 100644 index 000000000..fe3002bcc --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/goosefs.rs @@ -0,0 +1,379 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::GooseFs}; +use url::Url; + +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::{Error, Result}; + +/// Default GooseFS Master gRPC port. +const DEFAULT_GOOSEFS_PORT: u16 = 9200; + +/// GooseFS object store provider. +/// +/// Uses OpenDAL's GooseFs service to access GooseFS via gRPC. +/// URL format: `goosefs://host:port/path` +/// +/// Where: +/// - `host:port` is the GooseFS Master address (default port: 9200) +/// - `/path` is the filesystem path within GooseFS +/// +/// Path handling model (S3-style): +/// - The OpenDAL `root` is fixed to `/` (or a user-supplied cluster-wide base) +/// so that a single `Operator` can serve every dataset under the same +/// master. This keeps the `ObjectStoreRegistry` cache correct: two URLs +/// like `goosefs://host:9200/a.lance` and `goosefs://host:9200/b.lance` +/// share one store and each request carries its own object key. +/// - Path extraction relies on the default [`ObjectStoreProvider::extract_path`] +/// implementation, which returns the URL path (percent-decoded) as the key +/// passed to `ObjectStore::get`, `put`, etc. — mirroring how `s3://bucket/k` +/// yields key `k`. +/// +/// Supported configuration keys (via `storage_options` or environment variables, +/// resolved with priority: `storage_options` > env var > URL authority > default): +/// +/// | storage_options key | env var | purpose | +/// |---------------------------|-------------------------|-----------------------------------------------------------------------------------------------| +/// | `goosefs_master_addr` | `GOOSEFS_MASTER_ADDR` | Master gRPC address, e.g. `host:9200`. Supports HA: `addr1:port,addr2:port`. | +/// | `goosefs_root` | `GOOSEFS_ROOT` | Cluster-wide OpenDAL root shared by all datasets under the same master. Defaults to `/`. | +/// | `goosefs_write_type` | `GOOSEFS_WRITE_TYPE` | GooseFS write type (e.g. `MUST_CACHE`, `CACHE_THROUGH`, `THROUGH`, `ASYNC_THROUGH`). | +/// | `goosefs_block_size` | `GOOSEFS_BLOCK_SIZE` | GooseFS block size (bytes). Distinct from Lance's own `block_size`. | +/// | `goosefs_chunk_size` | `GOOSEFS_CHUNK_SIZE` | GooseFS chunk size (bytes) used by the client. | +/// | `goosefs_auth_type` | `GOOSEFS_AUTH_TYPE` | Authentication mode: `nosasl` or `simple`. | +/// | `goosefs_auth_username` | `GOOSEFS_AUTH_USERNAME` | Username for `simple` auth mode. | +/// +/// Note on `goosefs_root`: it is deliberately cluster-wide (not per-URL) so +/// that many datasets under the same master share a single cached `Operator`. +/// A custom root also participates in the `ObjectStoreRegistry` cache prefix, +/// so stores rooted at different subtrees do not collide. +#[derive(Default, Debug)] +pub struct GooseFsStoreProvider; + +impl GooseFsStoreProvider { + /// Resolve the GooseFS Master address from storage_options, environment, or URL. + /// + /// Priority: + /// 1. `storage_options["goosefs_master_addr"]` (supports HA: "addr1:port,addr2:port") + /// 2. `GOOSEFS_MASTER_ADDR` environment variable + /// 3. URL authority (host:port from the URL) + fn resolve_master_addr(url: &Url, storage_options: &StorageOptions) -> Result { + // 1. storage_options + if let Some(addr) = storage_options + .0 + .get("goosefs_master_addr") + .filter(|v| !v.is_empty()) + { + return Ok(addr.clone()); + } + + // 2. Environment variable + if let Ok(addr) = std::env::var("GOOSEFS_MASTER_ADDR") + && !addr.is_empty() + { + return Ok(addr); + } + + // 3. URL authority + let host = url.host_str().ok_or_else(|| { + Error::invalid_input( + "GooseFS URL must contain a master address (host), e.g. goosefs://host:port/path", + ) + })?; + + let port = url.port().unwrap_or(DEFAULT_GOOSEFS_PORT); + Ok(format!("{}:{}", host, port)) + } + + /// Resolve a storage option from storage_options or environment variable. + fn resolve_option( + storage_options: &StorageOptions, + option_key: &str, + env_key: &str, + ) -> Option { + storage_options + .0 + .get(option_key) + .cloned() + .or_else(|| std::env::var(env_key).ok()) + .filter(|v| !v.is_empty()) + } + + /// Resolve the OpenDAL `root` for this Operator. See the file-level docs on + /// [`GooseFsStoreProvider`] for the semantics of `goosefs_root`. + fn resolve_root(storage_options: &StorageOptions) -> String { + Self::resolve_option(storage_options, "goosefs_root", "GOOSEFS_ROOT") + .unwrap_or_else(|| "/".to_string()) + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for GooseFsStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + + // Resolve master address + let master_addr = Self::resolve_master_addr(&base_path, &storage_options)?; + + // Resolve a stable cluster-wide root. The URL path is *not* used here + // because it varies per dataset; per-request keys are supplied by + // `extract_path` instead. + let root = Self::resolve_root(&storage_options); + + // Build OpenDAL config map + let mut config_map: HashMap = HashMap::new(); + config_map.insert("master_addr".to_string(), master_addr); + config_map.insert("root".to_string(), root); + + // Optional: write_type + if let Some(wt) = + Self::resolve_option(&storage_options, "goosefs_write_type", "GOOSEFS_WRITE_TYPE") + { + config_map.insert("write_type".to_string(), wt); + } + + // Optional: block_size (for GooseFS, not Lance block_size) + if let Some(bs) = + Self::resolve_option(&storage_options, "goosefs_block_size", "GOOSEFS_BLOCK_SIZE") + { + config_map.insert("block_size".to_string(), bs); + } + + // Optional: chunk_size + if let Some(cs) = + Self::resolve_option(&storage_options, "goosefs_chunk_size", "GOOSEFS_CHUNK_SIZE") + { + config_map.insert("chunk_size".to_string(), cs); + } + + // Optional: auth_type (nosasl / simple) + if let Some(at) = + Self::resolve_option(&storage_options, "goosefs_auth_type", "GOOSEFS_AUTH_TYPE") + { + config_map.insert("auth_type".to_string(), at); + } + + // Optional: auth_username (used in SIMPLE auth mode) + if let Some(au) = Self::resolve_option( + &storage_options, + "goosefs_auth_username", + "GOOSEFS_AUTH_USERNAME", + ) { + config_map.insert("auth_username".to_string(), au); + } + + // Create OpenDAL Operator with GooseFS service + let operator = Operator::from_iter::(config_map).map_err(|e| { + Error::invalid_input(format!("Failed to create GooseFS operator: {:?}", e)) + })?; + + // Wrap as object_store::ObjectStore via OpendalStore bridge + let opendal_store = Arc::new(OpendalStore::new(operator)); + + Ok(ObjectStore { + scheme: "goosefs".to_string(), + inner: opendal_store, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(false), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count: storage_options.download_retry_count(), + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } + + // `extract_path` uses the default `ObjectStoreProvider` trait implementation: + // it percent-decodes the URL path and returns it as the object key, exactly + // like S3 does for `s3://bucket/key`. Overriding it here would only + // duplicate that behavior. See the file-level doc comment above for the + // full path-handling model. + + /// Calculate the object store prefix used as the registry cache key. + /// + /// Format: `goosefs$host:port`. Because the OpenDAL root is now cluster- + /// wide (not per-URL), all datasets under the same master intentionally + /// share the same cached [`ObjectStore`]; the URL path is disambiguated + /// by [`Self::extract_path`] on each request. This is analogous to how + /// two `s3://bucket/a` and `s3://bucket/b` URLs share one store. + fn calculate_object_store_prefix( + &self, + url: &Url, + storage_options: Option<&HashMap>, + ) -> Result { + // If a custom `goosefs_root` is provided, include it in the prefix so + // that stores built with different roots don't accidentally collide. + let opts = StorageOptions(storage_options.cloned().unwrap_or_default()); + let root = Self::resolve_root(&opts); + if root == "/" { + Ok(format!("{}${}", url.scheme(), url.authority())) + } else { + Ok(format!("{}${}#{}", url.scheme(), url.authority(), root)) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_goosefs_extract_path_basic() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://10.0.0.1:9200/data/embeddings.lance").unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.to_string(), "data/embeddings.lance"); + } + + #[test] + fn test_goosefs_extract_path_root() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://10.0.0.1:9200/").unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.to_string(), ""); + } + + #[test] + fn test_goosefs_extract_path_deep() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://master:9200/a/b/c/d.lance").unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.to_string(), "a/b/c/d.lance"); + } + + #[test] + fn test_goosefs_extract_path_percent_decoded() { + // The URL contains a percent-encoded space; extract_path must decode + // it once so the ObjectStore layer does not double-encode later. + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://master:9200/dir/with%20space/f.lance").unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.to_string(), "dir/with space/f.lance"); + } + + #[test] + fn test_calculate_object_store_prefix_default_root() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap(); + let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + assert_eq!(prefix, "goosefs$10.0.0.1:9200"); + } + + #[test] + fn test_calculate_object_store_prefix_with_hostname() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://myhost:9200/data").unwrap(); + let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + assert_eq!(prefix, "goosefs$myhost:9200"); + } + + /// Regression test: two URLs pointing at different datasets under the + /// same master must produce the *same* cache prefix so they share one + /// Operator, and correctness must come from `extract_path` returning + /// distinct keys — never from a per-URL root baked into the prefix. + #[test] + fn test_prefix_shared_across_datasets_same_master() { + let provider = GooseFsStoreProvider; + let url_a = Url::parse("goosefs://10.0.0.1:9200/repro/a.lance").unwrap(); + let url_b = Url::parse("goosefs://10.0.0.1:9200/repro/b.lance").unwrap(); + + let pa = provider + .calculate_object_store_prefix(&url_a, None) + .unwrap(); + let pb = provider + .calculate_object_store_prefix(&url_b, None) + .unwrap(); + assert_eq!(pa, pb, "same master must share one cache prefix"); + + // Extracted keys must differ so the shared Operator can route + // requests to the correct dataset. + assert_ne!( + provider.extract_path(&url_a).unwrap(), + provider.extract_path(&url_b).unwrap(), + "distinct URLs must yield distinct object keys", + ); + } + + /// Different masters must never share a cache entry. + #[test] + fn test_prefix_isolated_across_masters() { + let provider = GooseFsStoreProvider; + let u1 = Url::parse("goosefs://host-a:9200/x.lance").unwrap(); + let u2 = Url::parse("goosefs://host-b:9200/x.lance").unwrap(); + assert_ne!( + provider.calculate_object_store_prefix(&u1, None).unwrap(), + provider.calculate_object_store_prefix(&u2, None).unwrap(), + ); + } + + /// A user-supplied `goosefs_root` participates in the cache prefix so + /// stores rooted at different subtrees don't collide. + #[test] + fn test_prefix_includes_custom_root() { + let provider = GooseFsStoreProvider; + let url = Url::parse("goosefs://host:9200/x.lance").unwrap(); + + let default_prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + let custom_opts: HashMap = + HashMap::from([("goosefs_root".to_string(), "/tenant-a".to_string())]); + let custom_prefix = provider + .calculate_object_store_prefix(&url, Some(&custom_opts)) + .unwrap(); + + assert_eq!(default_prefix, "goosefs$host:9200"); + assert_eq!(custom_prefix, "goosefs$host:9200#/tenant-a"); + assert_ne!(default_prefix, custom_prefix); + } + + #[test] + fn test_resolve_master_addr_from_url() { + let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap(); + let storage_options = StorageOptions(HashMap::new()); + let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap(); + assert_eq!(addr, "10.0.0.1:9200"); + } + + #[test] + fn test_resolve_master_addr_default_port() { + let url = Url::parse("goosefs://10.0.0.1/data").unwrap(); + let storage_options = StorageOptions(HashMap::new()); + let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap(); + assert_eq!(addr, "10.0.0.1:9200"); + } + + #[test] + fn test_resolve_master_addr_from_storage_options() { + let url = Url::parse("goosefs://10.0.0.1:9200/data").unwrap(); + let storage_options = StorageOptions(HashMap::from([( + "goosefs_master_addr".to_string(), + "10.0.0.2:9200,10.0.0.3:9200".to_string(), + )])); + let addr = GooseFsStoreProvider::resolve_master_addr(&url, &storage_options).unwrap(); + assert_eq!(addr, "10.0.0.2:9200,10.0.0.3:9200"); + } + + #[test] + fn test_resolve_root_defaults_to_slash() { + let opts = StorageOptions(HashMap::new()); + assert_eq!(GooseFsStoreProvider::resolve_root(&opts), "/"); + } + + #[test] + fn test_resolve_root_from_storage_options() { + let opts = StorageOptions(HashMap::from([( + "goosefs_root".to_string(), + "/tenant-a".to_string(), + )])); + assert_eq!(GooseFsStoreProvider::resolve_root(&opts), "/tenant-a"); + } +} diff --git a/vendor/lance-io/src/object_store/providers/huggingface.rs b/vendor/lance-io/src/object_store/providers/huggingface.rs new file mode 100644 index 000000000..cda56e36f --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/huggingface.rs @@ -0,0 +1,440 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::ObjectStore as OSObjectStore; +use object_store::path::Path; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Huggingface}; +use url::Url; + +use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::parse_hf_repo_id; +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::{Error, Result}; + +/// Hugging Face object store provider backed by OpenDAL. +#[derive(Default, Debug)] +pub struct HuggingfaceStoreProvider; + +/// Parsed components from a Hugging Face URL. +#[derive(Debug, PartialEq, Eq)] +struct ParsedHfUrl { + repo_type: String, + repo_id: String, + relative_path: String, +} + +fn parse_hf_url(url: &Url) -> Result { + let mut repo_type = url + .host_str() + .ok_or_else(|| Error::invalid_input("Huggingface URL must contain repo type"))? + .to_string(); + // OpenDAL expects `dataset` instead of `datasets`; keep the workaround here and adapt tests. + if repo_type == "datasets" { + repo_type = "dataset".to_string(); + } + + let mut segments = url.path().trim_start_matches('/').split('/'); + let owner = segments + .next() + .ok_or_else(|| Error::invalid_input("Huggingface URL must contain owner"))?; + let repo_name = segments + .next() + .ok_or_else(|| Error::invalid_input("Huggingface URL must contain repository name"))?; + + let relative_path = segments.collect::>().join("/"); + + Ok(ParsedHfUrl { + repo_type, + repo_id: format!("{owner}/{repo_name}"), + relative_path, + }) +} + +fn build_hf_base_options( + repo_type: &str, + repo_id: &str, + storage_options: &StorageOptions, +) -> HashMap { + let mut options = storage_options.0.clone(); + options.insert("repo_type".to_string(), repo_type.to_string()); + options.insert("repo_id".to_string(), repo_id.to_string()); + options +} + +fn normalize_download_mode(download_mode: String) -> Result { + match download_mode.to_lowercase().as_str() { + "xet" => Ok("xet".to_string()), + "http" => Ok("http".to_string()), + _ => Err(Error::invalid_input(format!( + "Invalid Huggingface download_mode: {download_mode}. Expected one of: xet, http" + ))), + } +} + +fn normalize_hf_config(options: &HashMap) -> Result> { + let mut config_map = HashMap::new(); + + let repo_type = options + .get("repo_type") + .cloned() + .ok_or_else(|| Error::invalid_input("Huggingface repo_type is required"))?; + let repo_id = options + .get("repo_id") + .cloned() + .ok_or_else(|| Error::invalid_input("Huggingface repo_id is required"))?; + + config_map.insert("repo_type".to_string(), repo_type); + config_map.insert("repo_id".to_string(), repo_id); + + if let Some(revision) = options + .get("hf_revision") + .cloned() + .or_else(|| options.get("revision").cloned()) + { + config_map.insert("revision".to_string(), revision); + } + + if let Some(root) = options + .get("hf_root") + .cloned() + .or_else(|| options.get("root").cloned()) + && !root.is_empty() + { + config_map.insert("root".to_string(), root); + } + + if let Some(token) = options + .get("hf_token") + .cloned() + .or_else(|| options.get("token").cloned()) + && !token.is_empty() + { + config_map.insert("token".to_string(), token); + } + + let download_mode = options + .get("hf_download_mode") + .filter(|download_mode| !download_mode.is_empty()) + .cloned() + .or_else(|| { + options + .get("download_mode") + .filter(|download_mode| !download_mode.is_empty()) + .cloned() + }) + .unwrap_or_else(|| "http".to_string()); + config_map.insert( + "download_mode".to_string(), + normalize_download_mode(download_mode)?, + ); + + Ok(config_map) +} + +fn build_hf_store(config_map: HashMap) -> Result { + let repo_type = config_map + .get("repo_type") + .ok_or_else(|| Error::invalid_input("Huggingface repo_type is required"))?; + let repo_id = config_map + .get("repo_id") + .ok_or_else(|| Error::invalid_input("Huggingface repo_id is required"))?; + + let mut builder = Huggingface::default().repo_type(repo_type).repo_id(repo_id); + if let Some(revision) = config_map.get("revision") { + builder = builder.revision(revision); + } + if let Some(root) = config_map.get("root") { + builder = builder.root(root); + } + if let Some(token) = config_map.get("token") { + builder = builder.token(token); + } + if let Some(download_mode) = config_map.get("download_mode") { + builder = builder.download_mode(download_mode); + } + + let operator = Operator::new(builder).map_err(|e| { + Error::invalid_input(format!("Failed to create Huggingface operator: {:?}", e)) + })?; + + Ok(OpendalStore::new(operator)) +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for HuggingfaceStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let ParsedHfUrl { + repo_type, repo_id, .. + } = parse_hf_url(&base_path)?; + + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let download_retry_count = storage_options.download_retry_count(); + + let mut base_options = build_hf_base_options(&repo_type, &repo_id, &storage_options); + if !base_options.contains_key("hf_token") && !base_options.contains_key("token") { + if let Ok(token) = std::env::var("HF_TOKEN") { + base_options.insert("hf_token".to_string(), token); + } else if let Ok(token) = std::env::var("HUGGINGFACE_TOKEN") { + base_options.insert("hf_token".to_string(), token); + } + } + + let accessor = params.get_accessor(); + let inner: Arc = + if let Some(accessor) = accessor.filter(|a| a.has_provider()) { + Arc::new( + DynamicOpenDalStore::new( + format!("hf:{}", base_path), + base_options, + accessor, + normalize_hf_config, + build_hf_store, + ) + .with_protected_keys(["repo_type", "repo_id"]), + ) + } else { + Arc::new(build_hf_store(normalize_hf_config(&base_options)?)?) + }; + + Ok(ObjectStore { + scheme: "hf".to_string(), + inner, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count, + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } + + fn extract_path(&self, url: &Url) -> Result { + let parsed = parse_hf_url(url)?; + Path::from_url_path(&parsed.relative_path).map_err(|_| { + Error::invalid_input(format!("Invalid path in Huggingface URL: {}", url.path())) + }) + } + + fn calculate_object_store_prefix( + &self, + url: &Url, + _storage_options: Option<&HashMap>, + ) -> Result { + let repo_id = parse_hf_repo_id(url)?; + Ok(format!("{}${}@{}", url.scheme(), url.authority(), repo_id)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use crate::object_store::StorageOptionsAccessor; + use crate::object_store::dynamic_opendal::DynamicOpenDalStore; + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + + #[test] + fn parse_basic_url() { + let url = Url::parse("hf://datasets/acme/repo/path/to/table.lance").unwrap(); + let parsed = parse_hf_url(&url).unwrap(); + assert_eq!( + parsed, + ParsedHfUrl { + repo_type: "dataset".to_string(), + repo_id: "acme/repo".to_string(), + relative_path: "path/to/table.lance".to_string(), + } + ); + } + + #[test] + fn storage_option_revision_takes_precedence() { + use crate::object_store::StorageOptionsAccessor; + use std::sync::Arc; + let url = Url::parse("hf://datasets/acme/repo/data/file").unwrap(); + let params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([(String::from("hf_revision"), String::from("stable"))]), + ))), + ..Default::default() + }; + // new_store should accept without creating operator; test precedence via builder config + let ParsedHfUrl { + repo_type, repo_id, .. + } = parse_hf_url(&url).unwrap(); + + // Build config map the same way new_store would to assert precedence logic. + let mut config_map: HashMap = HashMap::new(); + config_map.insert("repo_type".to_string(), repo_type); + config_map.insert("repo".to_string(), repo_id); + if let Some(rev) = params + .storage_options() + .unwrap() + .get("hf_revision") + .cloned() + { + config_map.insert("revision".to_string(), rev); + } + assert_eq!(config_map.get("revision").unwrap(), "stable"); + } + + #[test] + fn storage_options_cannot_override_url_repo_identity() { + let config = normalize_hf_config(&build_hf_base_options( + "dataset", + "acme/repo", + &crate::object_store::StorageOptions(HashMap::from([ + ("repo_type".to_string(), "model".to_string()), + ("repo_id".to_string(), "other/repo".to_string()), + ("hf_revision".to_string(), "stable".to_string()), + ])), + )) + .unwrap(); + + assert_eq!(config.get("repo_type").unwrap(), "dataset"); + assert_eq!(config.get("repo_id").unwrap(), "acme/repo"); + assert_eq!(config.get("revision").unwrap(), "stable"); + } + + #[test] + fn storage_option_download_mode_takes_hf_prefix_precedence() { + let config = normalize_hf_config(&build_hf_base_options( + "dataset", + "acme/repo", + &crate::object_store::StorageOptions(HashMap::from([ + ("download_mode".to_string(), "xet".to_string()), + ("hf_download_mode".to_string(), "http".to_string()), + ])), + )) + .unwrap(); + + assert_eq!(config.get("download_mode").unwrap(), "http"); + } + + #[test] + fn storage_option_download_mode_defaults_to_http() { + let config = normalize_hf_config(&build_hf_base_options( + "dataset", + "acme/repo", + &crate::object_store::StorageOptions(HashMap::new()), + )) + .unwrap(); + + assert_eq!(config.get("download_mode").unwrap(), "http"); + } + + #[test] + fn storage_option_download_mode_rejects_invalid_value() { + let err = normalize_hf_config(&build_hf_base_options( + "dataset", + "acme/repo", + &crate::object_store::StorageOptions(HashMap::from([( + "hf_download_mode".to_string(), + "invalid".to_string(), + )])), + )) + .unwrap_err(); + + assert!( + err.to_string().contains("download_mode"), + "unexpected error: {}", + err + ); + } + + #[test] + fn parse_hf_repo_id_with_type_and_owner_repo() { + let url = Url::parse("hf://models/owner/repo/path/to/file").unwrap(); + let repo = crate::object_store::parse_hf_repo_id(&url).unwrap(); + assert_eq!(repo, "owner/repo"); + } + + #[test] + fn parse_hf_repo_id_legacy_without_type() { + let url = Url::parse("hf://owner/repo/path/to/file").unwrap(); + let repo = crate::object_store::parse_hf_repo_id(&url).unwrap(); + assert_eq!(repo, "owner/repo"); + } + + #[test] + fn parse_hf_repo_id_strips_revision() { + let url = Url::parse("hf://datasets/owner/repo@main/data").unwrap(); + let repo = crate::object_store::parse_hf_repo_id(&url).unwrap(); + assert_eq!(repo, "owner/repo"); + } + + #[test] + fn parse_hf_repo_id_missing_segments_errors() { + let url = Url::parse("hf://datasets/only-owner").unwrap(); + let err = crate::object_store::parse_hf_repo_id(&url).unwrap_err(); + assert!( + err.to_string().contains("owner/repo"), + "unexpected error: {}", + err + ); + } + + #[test] + fn extract_path_returns_relative() { + let url = Url::parse("hf://datasets/acme/repo/sub/dir/table.lance").unwrap(); + let provider = HuggingfaceStoreProvider; + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.to_string(), "sub/dir/table.lance"); + } + + #[test] + fn calculate_prefix_uses_repo_id() { + let provider = HuggingfaceStoreProvider; + let url = Url::parse("hf://datasets/acme/repo/path").unwrap(); + let prefix = provider.calculate_object_store_prefix(&url, None).unwrap(); + assert_eq!(prefix, "hf$datasets@acme/repo"); + } + + #[test] + fn parse_invalid_url_errors() { + let url = Url::parse("hf://datasets/only-owner").unwrap(); + let err = parse_hf_url(&url).unwrap_err(); + assert!(err.to_string().contains("repository name")); + } + + #[tokio::test] + async fn test_dynamic_opendal_hf_store_uses_provider_token() { + let parsed = parse_hf_url(&Url::parse("hf://datasets/acme/repo/path").unwrap()).unwrap(); + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([("hf_token".to_string(), "dynamic-token".to_string())]), + }, + ))); + + let store = DynamicOpenDalStore::new( + "hf", + build_hf_base_options( + &parsed.repo_type, + &parsed.repo_id, + &crate::object_store::StorageOptions(HashMap::new()), + ), + accessor, + normalize_hf_config, + build_hf_store, + ); + + let current_store = store + .current_store() + .await + .expect("dynamic OpenDAL HuggingFace store should build"); + + assert!(current_store.to_string().contains("Opendal")); + } +} diff --git a/vendor/lance-io/src/object_store/providers/local.rs b/vendor/lance-io/src/object_store/providers/local.rs new file mode 100644 index 000000000..9f0762916 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/local.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, sync::Arc}; + +use crate::object_store::{ + DEFAULT_LOCAL_BLOCK_SIZE, DEFAULT_LOCAL_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::Error; +use lance_core::error::Result; +use object_store::{local::LocalFileSystem, path::Path}; +use url::Url; + +#[derive(Default, Debug)] +pub struct FileStoreProvider; + +#[async_trait::async_trait] +impl ObjectStoreProvider for FileStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_LOCAL_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let download_retry_count = storage_options.download_retry_count(); + Ok(ObjectStore { + inner: Arc::new(LocalFileSystem::new()), + scheme: base_path.scheme().to_owned(), + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: false, + list_is_lexically_ordered: false, + io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM, + download_retry_count, + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } + + fn extract_path(&self, url: &Url) -> Result { + if let Ok(file_path) = url.to_file_path() + && let Ok(path) = Path::from_absolute_path(&file_path) + { + return Ok(path); + } + + Path::from_url_path(url.path()).map_err(|e| { + Error::invalid_input(format!("Failed to parse path '{}': {}", url.path(), e)) + }) + } + + fn calculate_object_store_prefix( + &self, + url: &Url, + _storage_options: Option<&HashMap>, + ) -> Result { + Ok(url.scheme().to_string()) + } +} + +#[cfg(test)] +mod tests { + use crate::object_store::uri_to_url; + + use super::*; + + #[test] + fn test_file_store_path() { + let provider = FileStoreProvider; + + let cases = [ + ("file:///", ""), + ("file:///usr/local/bin", "usr/local/bin"), + ("file-object-store:///path/to/file", "path/to/file"), + ("file:///path/to/foo/../bar", "path/to/bar"), + ]; + + for (uri, expected_path) in cases { + let url = uri_to_url(uri).unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.as_ref(), expected_path, "uri: '{}'", uri); + } + } + + #[test] + fn test_calculate_object_store_prefix() { + let provider = FileStoreProvider; + assert_eq!( + "file", + provider + .calculate_object_store_prefix(&Url::parse("file:///etc").unwrap(), None) + .unwrap() + ); + } + + #[test] + fn test_calculate_object_store_prefix_for_file_object_store() { + let provider = FileStoreProvider; + assert_eq!( + "file-object-store", + provider + .calculate_object_store_prefix( + &Url::parse("file-object-store:///etc").unwrap(), + None + ) + .unwrap() + ); + } + + #[test] + #[cfg(windows)] + fn test_file_store_path_windows() { + let provider = FileStoreProvider; + + let cases = [ + ( + "C:\\Users\\ADMINI~1\\AppData\\Local\\", + "C:/Users/ADMINI~1/AppData/Local", + ), + ( + "C:\\Users\\ADMINI~1\\AppData\\Local\\..\\", + "C:/Users/ADMINI~1/AppData", + ), + ( + "file-object-store:///C:/Users/ADMINI~1/AppData/Local", + "C:/Users/ADMINI~1/AppData/Local", + ), + ( + "file:///C:/Users/RUNNER~1/AppData/Local/Temp/tmpm49j_w0f", + "C:/Users/RUNNER~1/AppData/Local/Temp/tmpm49j_w0f", + ), + ]; + + for (uri, expected_path) in cases { + let url = uri_to_url(uri).unwrap(); + let path = provider.extract_path(&url).unwrap(); + assert_eq!(path.as_ref(), expected_path); + } + } +} diff --git a/vendor/lance-io/src/object_store/providers/memory.rs b/vendor/lance-io/src/object_store/providers/memory.rs new file mode 100644 index 000000000..dd72edc46 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/memory.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{collections::HashMap, sync::Arc}; + +use crate::object_store::{ + DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_LOCAL_BLOCK_SIZE, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::Result; +use object_store::{memory::InMemory, path::Path}; +use url::Url; + +/// Provides a fresh in-memory object store for each call to `new_store`. +#[derive(Default, Debug)] +pub struct MemoryStoreProvider; + +#[async_trait::async_trait] +impl ObjectStoreProvider for MemoryStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_LOCAL_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let download_retry_count = storage_options.download_retry_count(); + Ok(ObjectStore { + inner: Arc::new(InMemory::new()), + scheme: String::from("memory"), + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: false, + list_is_lexically_ordered: true, + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count, + io_tracker: Default::default(), + store_prefix: self + .calculate_object_store_prefix(&base_path, params.storage_options())?, + }) + } + + fn extract_path(&self, url: &Url) -> Result { + let mut output = String::new(); + if let Some(domain) = url.domain() { + output.push_str(domain); + } + output.push_str(url.path()); + // The in-memory store uses the Path directly as a key with no HTTP layer, + // so there is no re-encoding step and thus no double-encoding to avoid. + // Path::from also tolerates the empty segments that local temp paths embed. + Ok(Path::from(output)) + } + + fn calculate_object_store_prefix( + &self, + _url: &Url, + _storage_options: Option<&HashMap>, + ) -> Result { + Ok("memory".to_string()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_memory_store_path() { + let provider = MemoryStoreProvider; + + let url = Url::parse("memory://path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = Path::from("path/to/file"); + assert_eq!(path, expected_path); + } + + #[test] + fn test_calculate_object_store_prefix() { + let provider = MemoryStoreProvider; + assert_eq!( + "memory", + provider + .calculate_object_store_prefix(&Url::parse("memory://etc").unwrap(), None) + .unwrap() + ); + } +} diff --git a/vendor/lance-io/src/object_store/providers/oss.rs b/vendor/lance-io/src/object_store/providers/oss.rs new file mode 100644 index 000000000..3d116e2e3 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/oss.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::ObjectStore as OSObjectStore; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Oss}; +use url::Url; + +use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct OssStoreProvider; + +impl OssStoreProvider { + fn base_oss_options( + base_path: &Url, + storage_options: &StorageOptions, + ) -> Result> { + let bucket = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("OSS URL must contain bucket name"))? + .to_string(); + + let prefix = base_path.path().trim_start_matches('/').to_string(); + + // Snapshot env-backed OSS defaults at store construction time. Dynamic provider + // options can still override these values during per-request config merging. + let mut config_map: HashMap = std::env::vars() + .filter(|(key, _)| { + key.starts_with("OSS_") + || key.starts_with("AWS_") + || key.starts_with("ALIBABA_CLOUD_") + }) + .map(|(key, value)| { + let normalized_key = key + .to_lowercase() + .replace("oss_", "") + .replace("aws_", "") + .replace("alibaba_cloud_", ""); + (normalized_key, value) + }) + .collect(); + + config_map.extend(storage_options.0.clone()); + + config_map.insert("bucket".to_string(), bucket); + if prefix.is_empty() { + config_map.remove("root"); + } else { + config_map.insert("root".to_string(), "/".to_string()); + } + + Ok(config_map) + } + + /// Normalize OSS storage options, resolving aliases for well-known keys + /// while passing through all other options (e.g. `role_arn`, + /// `sts_endpoint`, `allow_anonymous`, etc.) so that OpenDAL can use them. + fn normalize_oss_config(options: &HashMap) -> Result> { + let mut config_map = options.clone(); + + let alias_groups: &[(&str, &[&str])] = &[ + ("endpoint", &["oss_endpoint"]), + ("access_key_id", &["oss_access_key_id"]), + ("access_key_secret", &["oss_secret_access_key"]), + ("region", &["oss_region"]), + ("security_token", &["oss_security_token"]), + ]; + + for (canonical, aliases) in alias_groups { + for alias in *aliases { + if let Some(value) = config_map.remove(*alias) { + config_map.insert(canonical.to_string(), value); + break; + } + } + } + + if !config_map.contains_key("endpoint") { + return Err(Error::invalid_input( + "OSS endpoint is required. Please provide 'oss_endpoint' in storage options or set OSS_ENDPOINT environment variable", + )); + } + + Ok(config_map) + } + + fn build_oss_store(config_map: HashMap) -> Result { + let operator = Operator::from_iter::(config_map) + .map_err(|e| Error::invalid_input(format!("Failed to create OSS operator: {:?}", e)))?; + + Ok(OpendalStore::new(operator)) + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for OssStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + + let base_options = Self::base_oss_options(&base_path, &storage_options)?; + let accessor = params.get_accessor(); + + let inner: Arc = + if let Some(accessor) = accessor.filter(|a| a.has_provider()) { + Arc::new( + DynamicOpenDalStore::new( + format!("oss:{}", base_path), + base_options, + accessor, + Self::normalize_oss_config, + Self::build_oss_store, + ) + .with_protected_keys(["bucket", "root"]), + ) + } else { + Arc::new(Self::build_oss_store(Self::normalize_oss_config( + &base_options, + )?)?) + }; + + let mut url = base_path; + if !url.path().ends_with('/') { + url.set_path(&format!("{}/", url.path())); + } + + Ok(ObjectStore { + scheme: "oss".to_string(), + inner, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count: storage_options.download_retry_count(), + io_tracker: Default::default(), + store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use super::OssStoreProvider; + use crate::object_store::dynamic_opendal::DynamicOpenDalStore; + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + use crate::object_store::{ObjectStoreProvider, StorageOptionsAccessor}; + use url::Url; + + #[test] + fn test_oss_store_path() { + let provider = OssStoreProvider; + + let url = Url::parse("oss://bucket/path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = object_store::path::Path::from("path/to/file"); + assert_eq!(path, expected_path); + } + + #[test] + fn test_oss_alias_options_override_canonical_env_options() { + let config = OssStoreProvider::normalize_oss_config(&HashMap::from([ + ( + "endpoint".to_string(), + "https://env.example.com".to_string(), + ), + ( + "oss_endpoint".to_string(), + "https://user.example.com".to_string(), + ), + ("access_key_id".to_string(), "env-akid".to_string()), + ("oss_access_key_id".to_string(), "user-akid".to_string()), + ("access_key_secret".to_string(), "env-secret".to_string()), + ( + "oss_secret_access_key".to_string(), + "user-secret".to_string(), + ), + ("region".to_string(), "env-region".to_string()), + ("oss_region".to_string(), "user-region".to_string()), + ("security_token".to_string(), "env-token".to_string()), + ("oss_security_token".to_string(), "user-token".to_string()), + ("bucket".to_string(), "bucket".to_string()), + ])) + .unwrap(); + + assert_eq!(config.get("endpoint").unwrap(), "https://user.example.com"); + assert_eq!(config.get("access_key_id").unwrap(), "user-akid"); + assert_eq!(config.get("access_key_secret").unwrap(), "user-secret"); + assert_eq!(config.get("region").unwrap(), "user-region"); + assert_eq!(config.get("security_token").unwrap(), "user-token"); + assert!(!config.contains_key("oss_endpoint")); + assert!(!config.contains_key("oss_security_token")); + } + + #[test] + fn test_oss_url_bucket_and_root_are_authoritative() { + let storage_options = crate::object_store::StorageOptions(HashMap::from([ + ( + "oss_endpoint".to_string(), + "https://oss-cn-hangzhou.aliyuncs.com".to_string(), + ), + ("bucket".to_string(), "storage-options-bucket".to_string()), + ("root".to_string(), "/storage-options-root".to_string()), + ])); + let base_options = OssStoreProvider::base_oss_options( + &Url::parse("oss://url-bucket/path").unwrap(), + &storage_options, + ) + .unwrap(); + let config = OssStoreProvider::normalize_oss_config(&base_options).unwrap(); + + assert_eq!(config.get("bucket").unwrap(), "url-bucket"); + assert_eq!(config.get("root").unwrap(), "/"); + } + + #[test] + fn test_oss_empty_url_path_removes_storage_option_root() { + let storage_options = crate::object_store::StorageOptions(HashMap::from([ + ( + "oss_endpoint".to_string(), + "https://oss-cn-hangzhou.aliyuncs.com".to_string(), + ), + ("root".to_string(), "/storage-options-root".to_string()), + ])); + let base_options = OssStoreProvider::base_oss_options( + &Url::parse("oss://url-bucket").unwrap(), + &storage_options, + ) + .unwrap(); + let config = OssStoreProvider::normalize_oss_config(&base_options).unwrap(); + + assert_eq!(config.get("bucket").unwrap(), "url-bucket"); + assert!(!config.contains_key("root")); + } + + #[tokio::test] + async fn test_dynamic_opendal_oss_store_uses_provider_credentials() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([ + ( + "oss_endpoint".to_string(), + "https://oss-cn-hangzhou.aliyuncs.com".to_string(), + ), + ("oss_access_key_id".to_string(), "akid".to_string()), + ("oss_secret_access_key".to_string(), "secret".to_string()), + ("oss_security_token".to_string(), "token".to_string()), + ]), + }, + ))); + + let base_options = OssStoreProvider::base_oss_options( + &Url::parse("oss://bucket/path").unwrap(), + &crate::object_store::StorageOptions(HashMap::new()), + ) + .unwrap(); + + let store = DynamicOpenDalStore::new( + "oss", + base_options, + accessor, + OssStoreProvider::normalize_oss_config, + OssStoreProvider::build_oss_store, + ); + + let current_store = store + .current_store() + .await + .expect("dynamic OpenDAL OSS store should build"); + + assert!(current_store.to_string().contains("Opendal")); + } +} diff --git a/vendor/lance-io/src/object_store/providers/shared_memory.rs b/vendor/lance-io/src/object_store/providers/shared_memory.rs new file mode 100644 index 000000000..5e7ef4e8d --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/shared_memory.rs @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{ + collections::HashMap, + sync::{Arc, LazyLock, Mutex}, +}; + +use crate::object_store::{ + ObjectStore, ObjectStoreParams, ObjectStoreProvider, providers::memory::MemoryStoreProvider, +}; +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. +/// +/// 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. +static SHARED_BACKENDS: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +fn shared_backend_for(url: &Url) -> Arc { + SHARED_BACKENDS + .lock() + .expect("SHARED_BACKENDS mutex poisoned") + .entry(url.authority().to_string()) + .or_insert_with(|| Arc::new(InMemory::new())) + .clone() +} + +/// Like [`MemoryStoreProvider`], but every caller resolving a `shared-memory:///...` +/// URL with the same `` sees the same backing bytes — across `ObjectStoreRegistry` +/// instances, threads, and unrelated components in the same process. +/// +/// Intended for tests and harnesses that need multiple actors to coordinate through a +/// common in-memory object store (e.g. a writer and an independent reader, multi-pod +/// fence simulations). Choose distinct authorities for isolation +/// (`shared-memory://test-a` vs `shared-memory://test-b`). +/// +/// Unlike `memory://` — which mints a fresh `InMemory` per `new_store` call — this +/// provider is opt-in precisely so existing tests relying on per-call isolation are +/// unaffected. +#[derive(Default, Debug)] +pub struct SharedMemoryStoreProvider { + inner: MemoryStoreProvider, +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for SharedMemoryStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let mut store = self.inner.new_store(base_path.clone(), params).await?; + store.inner = shared_backend_for(&base_path); + store.scheme = String::from("shared-memory"); + store.store_prefix = self.calculate_object_store_prefix(&base_path, None)?; + Ok(store) + } + + fn extract_path(&self, url: &Url) -> Result { + // The authority is the bucket; the URL path is the object path within it. + Ok(Path::from(url.path().trim_start_matches('/'))) + } + + fn calculate_object_store_prefix( + &self, + url: &Url, + _storage_options: Option<&HashMap>, + ) -> Result { + Ok(format!("shared-memory${}", url.authority())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::object_store::ObjectStoreRegistry; + use bytes::Bytes; + use object_store::{ObjectStoreExt as _, PutPayload}; + + async fn store_for(uri: &str) -> (Arc, Path) { + let registry = Arc::new(ObjectStoreRegistry::default()); + let (store, path) = ObjectStore::from_uri_and_params(registry, uri, &Default::default()) + .await + .unwrap(); + (store, path) + } + + #[tokio::test] + async fn same_authority_shares_bytes_across_registries() { + let (writer, _) = store_for("shared-memory://bucket-a/").await; + writer + .inner + .put(&Path::from("file"), PutPayload::from_static(b"hello")) + .await + .unwrap(); + + // Build a *separate* registry — no shared state at the registry layer. + let (reader, _) = store_for("shared-memory://bucket-a/").await; + let bytes = reader.inner.get(&Path::from("file")).await.unwrap(); + assert_eq!(bytes.bytes().await.unwrap(), Bytes::from_static(b"hello")); + } + + #[tokio::test] + async fn different_authorities_are_isolated() { + let (a, _) = store_for("shared-memory://iso-a/").await; + let (b, _) = store_for("shared-memory://iso-b/").await; + a.inner + .put(&Path::from("k"), PutPayload::from_static(b"in-a")) + .await + .unwrap(); + assert!(b.inner.get(&Path::from("k")).await.is_err()); + } + + #[tokio::test] + async fn extract_path_strips_authority() { + let provider = SharedMemoryStoreProvider::default(); + let url = Url::parse("shared-memory://bucket/foo/bar").unwrap(); + assert_eq!(provider.extract_path(&url).unwrap(), Path::from("foo/bar")); + } + + #[tokio::test] + async fn from_uri_and_params_resolves_path_correctly() { + let (store, path) = store_for("shared-memory://path-test/sub/dir/obj").await; + assert_eq!(path, Path::from("sub/dir/obj")); + store + .inner + .put(&path, PutPayload::from_static(b"payload")) + .await + .unwrap(); + + let (peer, peer_path) = store_for("shared-memory://path-test/sub/dir/obj").await; + let bytes = peer.inner.get(&peer_path).await.unwrap(); + assert_eq!(bytes.bytes().await.unwrap(), Bytes::from_static(b"payload")); + } + + #[test] + fn calculate_prefix_is_per_authority() { + let provider = SharedMemoryStoreProvider::default(); + let a = provider + .calculate_object_store_prefix(&Url::parse("shared-memory://x/p").unwrap(), None) + .unwrap(); + let b = provider + .calculate_object_store_prefix(&Url::parse("shared-memory://y/p").unwrap(), None) + .unwrap(); + assert_ne!(a, b); + assert_eq!(a, "shared-memory$x"); + } +} diff --git a/vendor/lance-io/src/object_store/providers/tencent.rs b/vendor/lance-io/src/object_store/providers/tencent.rs new file mode 100644 index 000000000..d29d5a6ad --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/tencent.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Cos}; +use url::Url; + +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct TencentStoreProvider; + +#[async_trait::async_trait] +impl ObjectStoreProvider for TencentStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + + let bucket = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("Tencent Cos URL must contain bucket name"))? + .to_string(); + + let prefix = base_path.path().trim_start_matches('/').to_string(); + + // Start with environment variables as base configuration + let mut config_map: HashMap = std::env::vars() + .filter(|(k, _)| k.starts_with("COS_") || k.starts_with("TENCENTCLOUD_")) + .map(|(k, v)| { + // Convert env var names to opendal config keys + let key = k + .to_lowercase() + .replace("cos_", "") + .replace("tencentcloud_", ""); + (key, v) + }) + .collect(); + + config_map.insert("bucket".to_string(), bucket); + + if !prefix.is_empty() { + config_map.insert("root".to_string(), "/".to_string()); + } + + // Override with storage options if provided + if let Some(endpoint) = storage_options.0.get("cos_endpoint") { + config_map.insert("endpoint".to_string(), endpoint.clone()); + } + + if let Some(secret_id) = storage_options.0.get("cos_secret_id") { + config_map.insert("secret_id".to_string(), secret_id.clone()); + } + + if let Some(secret_key) = storage_options.0.get("cos_secret_key") { + config_map.insert("secret_key".to_string(), secret_key.clone()); + } + + if let Some(enable_versioning) = storage_options.0.get("cos_enable_versioning") { + config_map.insert("enable_versioning".to_string(), enable_versioning.clone()); + } + + // Currently, the configuration options for CosConfig in OpenDAL are very limited. + // Most configurations need to be entered via environment variables, such as TENCENTCLOUD_SECURITY_TOKEN, TENCENTCLOUD_REGION, etc. + // (more env config details: https://github.com/apache/opendal-reqsign/blob/v0.16.5/src/tencent/config.rs) + // Therefore, we need to keep `disable_config_load` always false to allow configurations to be loaded from environment variables. + // TODO: improve CosConfig in opendal and add more storage_option here + config_map.insert("disable_config_load".to_string(), "false".to_string()); + + if !config_map.contains_key("endpoint") { + return Err(Error::invalid_input( + "COS endpoint is required. Please provide 'cos_endpoint' in storage options or set COS_ENDPOINT environment variable", + )); + } + + let operator = Operator::from_iter::(config_map) + .map_err(|e| Error::invalid_input(format!("Failed to create COS operator: {:?}", e)))?; + + let opendal_store = Arc::new(OpendalStore::new(operator)); + + let mut url = base_path; + if !url.path().ends_with('/') { + url.set_path(&format!("{}/", url.path())); + } + + Ok(ObjectStore { + scheme: "cos".to_string(), + inner: opendal_store, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count: storage_options.download_retry_count(), + io_tracker: Default::default(), + store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + }) + } +} + +#[cfg(test)] +mod tests { + use super::TencentStoreProvider; + use crate::object_store::ObjectStoreProvider; + use url::Url; + + #[test] + fn test_cos_store_path() { + let provider = TencentStoreProvider; + + let url = Url::parse("cos://bucket/path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = object_store::path::Path::from("path/to/file"); + assert_eq!(path, expected_path); + } +} diff --git a/vendor/lance-io/src/object_store/providers/tos.rs b/vendor/lance-io/src/object_store/providers/tos.rs new file mode 100644 index 000000000..7dee659f5 --- /dev/null +++ b/vendor/lance-io/src/object_store/providers/tos.rs @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::ObjectStore as OSObjectStore; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Tos}; +use url::Url; + +use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct TosStoreProvider; + +impl TosStoreProvider { + fn tos_env_options_from_iter(vars: I) -> HashMap + where + I: IntoIterator, + K: Into, + V: Into, + { + let vars = vars + .into_iter() + .map(|(key, value)| (key.into(), value.into())) + .collect::>(); + let mut config_map = HashMap::new(); + + for prefix in ["VOLCENGINE_", "TOS_"] { + for (key, value) in &vars { + if let Some(stripped_key) = key.strip_prefix(prefix) { + config_map.insert(stripped_key.to_ascii_lowercase(), value.clone()); + } + } + } + + config_map + } + + fn base_tos_options( + base_path: &Url, + storage_options: &StorageOptions, + ) -> Result> { + let bucket = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("TOS URL must contain bucket name"))? + .to_string(); + + let prefix = base_path.path().trim_start_matches('/').to_string(); + + let mut config_map = Self::tos_env_options_from_iter(std::env::vars()); + + config_map.extend(storage_options.0.clone()); + + config_map.insert("bucket".to_string(), bucket); + if prefix.is_empty() { + config_map.remove("root"); + } else { + config_map.insert("root".to_string(), "/".to_string()); + } + + Ok(config_map) + } + + /// Normalize TOS storage options, resolving aliases for well-known keys + /// while passing through all other options so that OpenDAL can use them. + fn normalize_tos_config(options: &HashMap) -> Result> { + let mut config_map = options.clone(); + + let alias_groups: &[(&str, &[&str])] = &[ + ("endpoint", &["tos_endpoint"]), + ("region", &["tos_region"]), + ("access_key_id", &["tos_access_key_id"]), + ("secret_access_key", &["tos_secret_access_key"]), + ("security_token", &["tos_security_token"]), + ]; + + for (canonical, aliases) in alias_groups { + for alias in *aliases { + if let Some(value) = config_map.remove(*alias) { + config_map.insert(canonical.to_string(), value); + break; + } + } + } + + if !config_map.contains_key("endpoint") { + return Err(Error::invalid_input( + "TOS endpoint is required. Please provide 'tos_endpoint' in storage options or set TOS_ENDPOINT environment variable", + )); + } + + Ok(config_map) + } + + fn build_tos_store(config_map: HashMap) -> Result { + let operator = Operator::from_iter::(config_map) + .map_err(|e| Error::invalid_input(format!("Failed to create TOS operator: {:?}", e)))?; + + Ok(OpendalStore::new(operator)) + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for TosStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + + let base_options = Self::base_tos_options(&base_path, &storage_options)?; + let accessor = params.get_accessor(); + + let inner: Arc = + if let Some(accessor) = accessor.filter(|a| a.has_provider()) { + Arc::new( + DynamicOpenDalStore::new( + format!("tos:{}", base_path), + base_options, + accessor, + Self::normalize_tos_config, + Self::build_tos_store, + ) + .with_protected_keys(["bucket", "root"]), + ) + } else { + Arc::new(Self::build_tos_store(Self::normalize_tos_config( + &base_options, + )?)?) + }; + + let mut url = base_path; + if !url.path().ends_with('/') { + url.set_path(&format!("{}/", url.path())); + } + + Ok(ObjectStore { + scheme: "tos".to_string(), + inner, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count: storage_options.download_retry_count(), + io_tracker: Default::default(), + store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use super::TosStoreProvider; + use crate::object_store::dynamic_opendal::DynamicOpenDalStore; + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + use crate::object_store::{ObjectStoreProvider, StorageOptionsAccessor}; + use url::Url; + + #[test] + fn test_tos_store_path() { + let provider = TosStoreProvider; + + let url = Url::parse("tos://bucket/path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = object_store::path::Path::from("path/to/file"); + assert_eq!(path, expected_path); + } + + #[test] + fn test_tos_env_options_normalize_supported_prefixes() { + let config = TosStoreProvider::tos_env_options_from_iter([ + ("VOLCENGINE_ENDPOINT", "https://tos-cn-beijing.volces.com"), + ("TOS_ACCESS_KEY_ID", "tos-akid"), + ("TOS_SECRET_ACCESS_KEY", "tos-secret"), + ]); + + assert_eq!( + config.get("endpoint").unwrap(), + "https://tos-cn-beijing.volces.com" + ); + assert_eq!(config.get("access_key_id").unwrap(), "tos-akid"); + assert_eq!(config.get("secret_access_key").unwrap(), "tos-secret"); + } + + #[test] + fn test_tos_alias_options_override_canonical_env_options() { + let config = TosStoreProvider::normalize_tos_config(&HashMap::from([ + ( + "endpoint".to_string(), + "https://env.example.com".to_string(), + ), + ( + "tos_endpoint".to_string(), + "https://user.example.com".to_string(), + ), + ("region".to_string(), "env-region".to_string()), + ("tos_region".to_string(), "user-region".to_string()), + ("access_key_id".to_string(), "env-akid".to_string()), + ("tos_access_key_id".to_string(), "user-akid".to_string()), + ("secret_access_key".to_string(), "env-secret".to_string()), + ( + "tos_secret_access_key".to_string(), + "user-secret".to_string(), + ), + ("security_token".to_string(), "env-token".to_string()), + ("tos_security_token".to_string(), "user-token".to_string()), + ("bucket".to_string(), "bucket".to_string()), + ])) + .unwrap(); + + assert_eq!(config.get("endpoint").unwrap(), "https://user.example.com"); + assert_eq!(config.get("region").unwrap(), "user-region"); + assert_eq!(config.get("access_key_id").unwrap(), "user-akid"); + assert_eq!(config.get("secret_access_key").unwrap(), "user-secret"); + assert_eq!(config.get("security_token").unwrap(), "user-token"); + assert!(!config.contains_key("tos_endpoint")); + assert!(!config.contains_key("tos_secret_access_key")); + assert!(!config.contains_key("tos_security_token")); + } + + #[test] + fn test_tos_url_bucket_and_root_are_authoritative() { + let storage_options = crate::object_store::StorageOptions(HashMap::from([ + ( + "tos_endpoint".to_string(), + "https://tos-cn-beijing.volces.com".to_string(), + ), + ("bucket".to_string(), "storage-options-bucket".to_string()), + ("root".to_string(), "/storage-options-root".to_string()), + ])); + let base_options = TosStoreProvider::base_tos_options( + &Url::parse("tos://url-bucket/path").unwrap(), + &storage_options, + ) + .unwrap(); + let config = TosStoreProvider::normalize_tos_config(&base_options).unwrap(); + + assert_eq!(config.get("bucket").unwrap(), "url-bucket"); + assert_eq!(config.get("root").unwrap(), "/"); + + let base_options = TosStoreProvider::base_tos_options( + &Url::parse("tos://url-bucket").unwrap(), + &storage_options, + ) + .unwrap(); + let config = TosStoreProvider::normalize_tos_config(&base_options).unwrap(); + + assert_eq!(config.get("bucket").unwrap(), "url-bucket"); + assert!(!config.contains_key("root")); + } + + #[tokio::test] + async fn test_dynamic_opendal_tos_store_uses_provider_credentials() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([ + ( + "tos_endpoint".to_string(), + "https://tos-cn-beijing.volces.com".to_string(), + ), + ("tos_region".to_string(), "cn-beijing".to_string()), + ("tos_access_key_id".to_string(), "akid".to_string()), + ("tos_secret_access_key".to_string(), "secret".to_string()), + ("tos_security_token".to_string(), "token".to_string()), + ]), + }, + ))); + + let base_options = TosStoreProvider::base_tos_options( + &Url::parse("tos://url-bucket/path").unwrap(), + &crate::object_store::StorageOptions(HashMap::new()), + ) + .unwrap(); + + let store = DynamicOpenDalStore::new( + "tos", + base_options, + accessor, + TosStoreProvider::normalize_tos_config, + TosStoreProvider::build_tos_store, + ) + .with_protected_keys(["bucket", "root"]); + + let current_store = store + .current_store() + .await + .expect("dynamic OpenDAL TOS store should build"); + + assert!(current_store.to_string().contains("Opendal")); + } +} diff --git a/vendor/lance-io/src/object_store/storage_options.rs b/vendor/lance-io/src/object_store/storage_options.rs new file mode 100644 index 000000000..355845fd9 --- /dev/null +++ b/vendor/lance-io/src/object_store/storage_options.rs @@ -0,0 +1,1244 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Storage options provider and accessor for dynamic credential fetching +//! +//! This module provides: +//! - [`StorageOptionsProvider`] trait for fetching storage options from various sources +//! (namespace servers, secret managers, etc.) with support for expiration tracking +//! - [`StorageOptionsAccessor`] for unified access to storage options with automatic +//! caching and refresh + +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +#[cfg(test)] +use mock_instant::thread_local::{SystemTime, UNIX_EPOCH}; + +#[cfg(not(test))] +use std::time::{SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use lance_namespace::LanceNamespace; +use lance_namespace::models::DescribeTableRequest; +use tokio::sync::RwLock; + +use crate::{Error, Result}; + +/// Key for the expiration timestamp in storage options HashMap +pub const EXPIRES_AT_MILLIS_KEY: &str = "expires_at_millis"; + +/// Key for the refresh offset in storage options HashMap (milliseconds before expiry to refresh) +pub const REFRESH_OFFSET_MILLIS_KEY: &str = "refresh_offset_millis"; + +/// Default refresh offset: 60 seconds before expiration +const DEFAULT_REFRESH_OFFSET_MILLIS: u64 = 60_000; + +/// Trait for providing storage options with expiration tracking +/// +/// Implementations can fetch storage options from various sources (namespace servers, +/// secret managers, etc.) and are usable from Python/Java. +/// +/// # Current Use Cases +/// +/// - **Temporary Credentials**: Fetch short-lived AWS temporary credentials that expire +/// after a set time period, with automatic refresh before expiration +/// +/// # Future Possible Use Cases +/// +/// - **Dynamic Storage Location Resolution**: Resolve logical names to actual storage +/// locations (bucket aliases, S3 Access Points, region-specific endpoints) that may +/// change based on region, tier, data migration, or failover scenarios +/// - **Runtime S3 Tags Assignment**: Inject cost allocation tags, security labels, or +/// compliance metadata into S3 requests based on the current execution context (user, +/// application, workspace, etc.) +/// - **Dynamic Endpoint Configuration**: Update storage endpoints for disaster recovery, +/// A/B testing, or gradual migration scenarios +/// - **Just-in-time Permission Elevation**: Request elevated permissions only when needed +/// for sensitive operations, then immediately revoke them +/// - **Secret Manager Integration**: Fetch encryption keys from AWS Secrets Manager, +/// Azure Key Vault, or Google Secret Manager with automatic rotation +/// - **OIDC/SAML Federation**: Integrate with identity providers to obtain storage +/// credentials based on user identity and group membership +/// +/// # Equality and Hashing +/// +/// Implementations must provide `provider_id()` which returns a unique identifier for +/// equality and hashing purposes. Two providers with the same ID are considered equal +/// and will share the same cached ObjectStore in the registry. +#[async_trait] +pub trait StorageOptionsProvider: Send + Sync + fmt::Debug { + /// Fetch fresh storage options + /// + /// Returns None if no storage options are available, or Some(HashMap) with the options. + /// If the [`EXPIRES_AT_MILLIS_KEY`] key is present in the HashMap, it should contain the + /// epoch time in milliseconds when the options expire, and credentials will automatically + /// refresh before expiration. + /// If [`EXPIRES_AT_MILLIS_KEY`] is not provided, the options are considered to never expire. + async fn fetch_storage_options(&self) -> Result>>; + + /// Fetch fresh storage options, bypassing caches along the chain. + /// + /// Providers that serve from an upstream cache (e.g. base-scoped wrappers + /// reading through a parent accessor) override this to force the upstream + /// to re-fetch. Defaults to [`Self::fetch_storage_options`]. + async fn force_fetch_storage_options(&self) -> Result>> { + self.fetch_storage_options().await + } + + /// Return a human-readable unique identifier for this provider instance + /// + /// This is used for equality comparison and hashing in the object store registry. + /// Two providers with the same ID will be treated as equal and share the same cached + /// ObjectStore. + /// + /// The ID should be human-readable for debugging and logging purposes. + /// For example: `"namespace[dir(root=/data)],table[db$schema$table1]"` + /// + /// The ID should uniquely identify the provider's configuration. + fn provider_id(&self) -> String; +} + +/// StorageOptionsProvider implementation that fetches options from a LanceNamespace +pub struct LanceNamespaceStorageOptionsProvider { + namespace_client: Arc, + table_id: Vec, +} + +impl fmt::Debug for LanceNamespaceStorageOptionsProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.provider_id()) + } +} + +impl fmt::Display for LanceNamespaceStorageOptionsProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.provider_id()) + } +} + +impl LanceNamespaceStorageOptionsProvider { + /// Create a new LanceNamespaceStorageOptionsProvider + /// + /// # Arguments + /// * `namespace_client` - The namespace implementation to fetch storage options from + /// * `table_id` - The table identifier + pub fn new(namespace_client: Arc, table_id: Vec) -> Self { + Self { + namespace_client, + table_id, + } + } +} + +#[async_trait] +impl StorageOptionsProvider for LanceNamespaceStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let request = DescribeTableRequest { + id: Some(self.table_id.clone()), + // Some server implementations may not return credentials unless explicitly requested + vend_credentials: Some(true), + ..Default::default() + }; + + let response = self + .namespace_client + .describe_table(request) + .await + .map_err(|e| { + Error::io_source(Box::new(std::io::Error::other(format!( + "Failed to fetch storage options: {}", + e + )))) + })?; + + Ok(response.storage_options) + } + + fn provider_id(&self) -> String { + format!( + "LanceNamespaceStorageOptionsProvider {{ namespace_client: {}, table_id: {:?} }}", + self.namespace_client.namespace_id(), + self.table_id + ) + } +} + +/// Prefix marking a storage option as scoped to a single registered base path. +/// +/// A key of the form `base_.` applies `` only to the base path +/// with manifest id ``, overriding the shared (unscoped) options for that +/// base. For example `base_1.account_key = abc` gives the base with id 1 the +/// option `account_key = abc` while it inherits every unscoped option. +pub const BASE_SCOPED_OPTION_PREFIX: &str = "base_"; + +/// Parse a base-scoped storage option key of the form `base_.`. +/// +/// Returns `Some((base_id, key))` only for keys that match the convention +/// exactly: the `base_` prefix, an all-digit u32 base id, a `.` separator, and +/// a non-empty remainder. Any other key (e.g. `base_url`, `base_x.key`, +/// `base_1.`) is not base-scoped. +pub fn parse_base_scoped_key(key: &str) -> Option<(u32, &str)> { + let rest = key.strip_prefix(BASE_SCOPED_OPTION_PREFIX)?; + let (id_str, scoped_key) = rest.split_once('.')?; + if scoped_key.is_empty() || id_str.is_empty() || !id_str.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + let id = id_str.parse::().ok()?; + Some((id, scoped_key)) +} + +/// Returns true if any key in `options` is base-scoped (`base_.`). +pub fn has_base_scoped_options(options: &HashMap) -> bool { + options + .keys() + .any(|key| parse_base_scoped_key(key).is_some()) +} + +/// Resolve the effective storage options for one base path scope. +/// +/// All base-scoped keys are removed from the result. When `base_id` is +/// `Some(id)`, entries scoped to that id are overlaid on the unscoped options, +/// adding or overriding keys. `None` resolves the default scope (the primary +/// dataset base), which simply drops every base-scoped entry. +pub fn resolve_base_scoped_options( + options: &HashMap, + base_id: Option, +) -> HashMap { + let mut resolved = HashMap::with_capacity(options.len()); + let mut overrides = Vec::new(); + for (key, value) in options { + match parse_base_scoped_key(key) { + Some((id, scoped_key)) => { + if Some(id) == base_id { + overrides.push((scoped_key.to_string(), value.clone())); + } + } + None => { + resolved.insert(key.clone(), value.clone()); + } + } + } + resolved.extend(overrides); + resolved +} + +/// A [`StorageOptionsProvider`] that resolves another accessor's options for a +/// single base path scope. +/// +/// Fetching through this provider first refreshes the parent accessor when its +/// options have expired, then resolves the refreshed options for the scope. As +/// a result, dynamically vended per-base credentials (e.g. a namespace server +/// returning `base_.` entries in one flat map) stay fresh per base. +#[derive(Debug)] +pub struct BaseScopedStorageOptionsProvider { + inner: Arc, + base_id: Option, +} + +impl BaseScopedStorageOptionsProvider { + pub fn new(inner: Arc, base_id: Option) -> Self { + Self { inner, base_id } + } +} + +#[async_trait] +impl StorageOptionsProvider for BaseScopedStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let options = self.inner.get_storage_options().await?; + Ok(Some(resolve_base_scoped_options(&options.0, self.base_id))) + } + + async fn force_fetch_storage_options(&self) -> Result>> { + let options = self.inner.refresh_storage_options().await?; + Ok(Some(resolve_base_scoped_options(&options.0, self.base_id))) + } + + fn provider_id(&self) -> String { + match self.base_id { + Some(id) => format!("base-scoped[base_id={}]({})", id, self.inner.accessor_id()), + None => format!("base-scoped[default]({})", self.inner.accessor_id()), + } + } +} + +/// Unified access to storage options with automatic caching and refresh +/// +/// This struct bundles static storage options with an optional dynamic provider, +/// handling all caching and refresh logic internally. It provides a single entry point +/// for accessing storage options regardless of whether they're static or dynamic. +/// +/// # Behavior +/// +/// - If only static options are provided, returns those options +/// - If a provider is configured, fetches from provider and caches results +/// - Automatically refreshes cached options before expiration (based on refresh_offset) +/// - Uses `expires_at_millis` key to track expiration +/// +/// # Thread Safety +/// +/// The accessor is thread-safe and can be shared across multiple tasks. +/// Concurrent refresh attempts are deduplicated using a try-lock mechanism. +pub struct StorageOptionsAccessor { + /// Initial/fallback static storage options + initial_options: Option>, + + /// Optional dynamic provider for refreshing options + provider: Option>, + + /// Cached storage options with expiration tracking + cache: Arc>>, + + /// Duration before expiry to trigger refresh + refresh_offset: Duration, + + /// True when this accessor was produced by [`Self::scoped_to_base`]; its + /// options are already resolved for one base path scope, so scoping again + /// is a no-op. + scope_resolved: bool, +} + +impl fmt::Debug for StorageOptionsAccessor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("StorageOptionsAccessor") + .field("has_initial_options", &self.initial_options.is_some()) + .field("has_provider", &self.provider.is_some()) + .field("refresh_offset", &self.refresh_offset) + .finish() + } +} + +#[derive(Debug, Clone)] +struct CachedStorageOptions { + options: HashMap, + expires_at_millis: Option, +} + +impl StorageOptionsAccessor { + /// Extract refresh offset from storage options, or use default + fn extract_refresh_offset(options: &HashMap) -> Duration { + options + .get(REFRESH_OFFSET_MILLIS_KEY) + .and_then(|s| s.parse::().ok()) + .map(Duration::from_millis) + .unwrap_or(Duration::from_millis(DEFAULT_REFRESH_OFFSET_MILLIS)) + } + + /// Effective expiration of a raw options map: the minimum of the unscoped + /// `expires_at_millis` and every `base_.expires_at_millis` entry. + /// + /// A flat map may vend per-base credentials that expire before the shared + /// ones. Refreshing when the earliest credential is due keeps base-scoped + /// accessors (which refresh through this accessor) from re-resolving stale + /// per-base credentials out of a still-"valid" cache. + fn effective_expires_at_millis(options: &HashMap) -> Option { + options + .iter() + .filter(|(key, _)| { + key.as_str() == EXPIRES_AT_MILLIS_KEY + || matches!( + parse_base_scoped_key(key), + Some((_, scoped_key)) if scoped_key == EXPIRES_AT_MILLIS_KEY + ) + }) + .filter_map(|(_, value)| value.parse::().ok()) + .min() + } + + /// Create an accessor with only static options (no refresh capability) + /// + /// The returned accessor will always return the provided options. + /// This is useful when credentials don't expire or are managed externally. + pub fn with_static_options(options: HashMap) -> Self { + let expires_at_millis = Self::effective_expires_at_millis(&options); + let refresh_offset = Self::extract_refresh_offset(&options); + + Self { + initial_options: Some(options.clone()), + provider: None, + cache: Arc::new(RwLock::new(Some(CachedStorageOptions { + options, + expires_at_millis, + }))), + refresh_offset, + scope_resolved: false, + } + } + + /// Create an accessor with a dynamic provider (no initial options) + /// + /// The accessor will fetch from the provider on first access and cache + /// the results. Refresh happens automatically before expiration. + /// Uses the default refresh offset (60 seconds) until options are fetched. + /// + /// # Arguments + /// * `provider` - The storage options provider for fetching fresh options + pub fn with_provider(provider: Arc) -> Self { + Self { + initial_options: None, + provider: Some(provider), + cache: Arc::new(RwLock::new(None)), + refresh_offset: Duration::from_millis(DEFAULT_REFRESH_OFFSET_MILLIS), + scope_resolved: false, + } + } + + /// Create an accessor with initial options and a dynamic provider + /// + /// Initial options are used until they expire, then the provider is called. + /// This avoids an immediate fetch when initial credentials are still valid. + /// The `refresh_offset_millis` key in initial_options controls refresh timing. + /// + /// # Arguments + /// * `initial_options` - Initial storage options to cache + /// * `provider` - The storage options provider for refreshing + pub fn with_initial_and_provider( + initial_options: HashMap, + provider: Arc, + ) -> Self { + let expires_at_millis = Self::effective_expires_at_millis(&initial_options); + let refresh_offset = Self::extract_refresh_offset(&initial_options); + + Self { + initial_options: Some(initial_options.clone()), + provider: Some(provider), + cache: Arc::new(RwLock::new(Some(CachedStorageOptions { + options: initial_options, + expires_at_millis, + }))), + refresh_offset, + scope_resolved: false, + } + } + + /// Get current valid storage options + /// + /// - Returns cached options if not expired + /// - Fetches from provider if expired or not cached + /// - Falls back to initial_options if provider returns None + /// + /// # Errors + /// + /// Returns an error if: + /// - The provider fails to fetch options + /// - No options are available (no cache, no provider, no initial options) + pub async fn get_storage_options(&self) -> Result { + loop { + match self.do_get_storage_options().await? { + Some(options) => return Ok(options), + None => { + // Lock was busy, wait 10ms before retrying + tokio::time::sleep(Duration::from_millis(10)).await; + continue; + } + } + } + } + + /// Fetch fresh options from the provider and update the cache. + /// + /// This bypasses the cache for callers that need to validate provider-vended + /// credentials even when initial metadata has no expiration. The force + /// propagates through provider chains (e.g. base-scoped wrappers), so the + /// origin provider is re-queried even when intermediate caches are valid. + pub(crate) async fn refresh_storage_options(&self) -> Result { + let Some(provider) = &self.provider else { + return self.get_storage_options().await; + }; + + log::debug!( + "Refreshing storage options from provider: {}", + provider.provider_id() + ); + + let storage_options_map = provider.force_fetch_storage_options().await.map_err(|e| { + Error::io_source(Box::new(std::io::Error::other(format!( + "Failed to fetch storage options: {}", + e + )))) + })?; + + let Some(options) = storage_options_map else { + if let Some(initial) = &self.initial_options { + return Ok(super::StorageOptions(initial.clone())); + } + log::debug!( + "Provider {} returned no storage options, using default credentials", + provider.provider_id() + ); + return Ok(super::StorageOptions(HashMap::new())); + }; + + let expires_at_millis = Self::effective_expires_at_millis(&options); + + let mut cache = self.cache.write().await; + *cache = Some(CachedStorageOptions { + options: options.clone(), + expires_at_millis, + }); + + Ok(super::StorageOptions(options)) + } + + async fn do_get_storage_options(&self) -> Result> { + // Check if we have valid cached options with read lock + { + let cached = self.cache.read().await; + if !self.needs_refresh(&cached) + && let Some(cached_opts) = &*cached + { + return Ok(Some(super::StorageOptions(cached_opts.options.clone()))); + } + } + + // If no provider, return initial options or use defaults + let Some(provider) = &self.provider else { + return if let Some(initial) = &self.initial_options { + Ok(Some(super::StorageOptions(initial.clone()))) + } else { + // No provider and no initial options - use default credentials + Ok(Some(super::StorageOptions(HashMap::new()))) + }; + }; + + // Try to acquire write lock - if it fails, return None and let caller retry + let Ok(mut cache) = self.cache.try_write() else { + return Ok(None); + }; + + // Double-check if options are still stale after acquiring write lock + // (another thread might have refreshed them) + if !self.needs_refresh(&cache) + && let Some(cached_opts) = &*cache + { + return Ok(Some(super::StorageOptions(cached_opts.options.clone()))); + } + log::debug!( + "Refreshing storage options from provider: {}", + provider.provider_id() + ); + + let storage_options_map = provider.fetch_storage_options().await.map_err(|e| { + Error::io_source(Box::new(std::io::Error::other(format!( + "Failed to fetch storage options: {}", + e + )))) + })?; + + let Some(options) = storage_options_map else { + // Provider returned None, fall back to initial options or use defaults + if let Some(initial) = &self.initial_options { + return Ok(Some(super::StorageOptions(initial.clone()))); + } + // Provider returned None and no initial options - use default credentials + // This is valid when namespace doesn't vend credentials (e.g., directory namespace + // where environment credentials are used) + log::debug!( + "Provider {} returned no storage options, using default credentials", + provider.provider_id() + ); + return Ok(Some(super::StorageOptions(HashMap::new()))); + }; + + let expires_at_millis = Self::effective_expires_at_millis(&options); + + if let Some(expires_at) = expires_at_millis { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::from_secs(0)) + .as_millis() as u64; + let expires_in_secs = (expires_at.saturating_sub(now_ms)) / 1000; + log::debug!( + "Successfully refreshed storage options from provider: {}, options expire in {} seconds", + provider.provider_id(), + expires_in_secs + ); + } else { + log::debug!( + "Successfully refreshed storage options from provider: {} (no expiration)", + provider.provider_id() + ); + } + + *cache = Some(CachedStorageOptions { + options: options.clone(), + expires_at_millis, + }); + + Ok(Some(super::StorageOptions(options))) + } + + fn needs_refresh(&self, cached: &Option) -> bool { + match cached { + None => true, + Some(cached_opts) => { + if let Some(expires_at_millis) = cached_opts.expires_at_millis { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::from_secs(0)) + .as_millis() as u64; + + // Refresh if we're within the refresh offset of expiration + let refresh_offset_millis = self.refresh_offset.as_millis() as u64; + now_ms + refresh_offset_millis >= expires_at_millis + } else { + // No expiration means options never expire + false + } + } + } + } + + /// Get the initial storage options without refresh + /// + /// Returns the initial options that were provided when creating the accessor. + /// This does not trigger any refresh, even if the options have expired. + pub fn initial_storage_options(&self) -> Option<&HashMap> { + self.initial_options.as_ref() + } + + /// Get the accessor ID for equality/hashing + /// + /// Returns the provider_id if a provider exists, otherwise generates + /// a stable ID from the initial options hash. + pub fn accessor_id(&self) -> String { + if let Some(provider) = &self.provider { + provider.provider_id() + } else if let Some(initial) = &self.initial_options { + // Generate a stable ID from initial options + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + + let mut hasher = DefaultHasher::new(); + let mut keys: Vec<_> = initial.keys().collect(); + keys.sort(); + for key in keys { + key.hash(&mut hasher); + initial.get(key).hash(&mut hasher); + } + format!("static_options_{:x}", hasher.finish()) + } else { + "empty_accessor".to_string() + } + } + + /// Resolve this accessor for a single base path scope. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to one registered base path. The returned accessor resolves + /// options for `base_id`: entries scoped to that base overlay the unscoped + /// defaults, and all other scoped entries are dropped. `None` resolves the + /// default scope used for the primary dataset base. + /// + /// A static accessor whose options contain no base-scoped entries is + /// returned unchanged, preserving accessor identity (and thus object store + /// registry cache keys). A provider-backed accessor is always wrapped + /// through [`BaseScopedStorageOptionsProvider`] — fetched options may vend + /// base-scoped entries at any refresh, even when the initial options carry + /// none — so refreshed options are re-resolved for the scope on every + /// fetch. Accessors already produced by this method are returned unchanged. + pub fn scoped_to_base(self: &Arc, base_id: Option) -> Arc { + if self.scope_resolved { + return self.clone(); + } + if self.has_provider() { + let provider = Arc::new(BaseScopedStorageOptionsProvider::new(self.clone(), base_id)); + let mut scoped = match self.initial_storage_options() { + Some(initial) => Self::with_initial_and_provider( + resolve_base_scoped_options(initial, base_id), + provider, + ), + None => Self::with_provider(provider), + }; + scoped.scope_resolved = true; + Arc::new(scoped) + } else { + match self.initial_storage_options() { + Some(initial) if has_base_scoped_options(initial) => { + let mut scoped = + Self::with_static_options(resolve_base_scoped_options(initial, base_id)); + scoped.scope_resolved = true; + Arc::new(scoped) + } + // Static options never change, so there is nothing to scope. + _ => self.clone(), + } + } + } + + /// Check if this accessor has a dynamic provider + pub fn has_provider(&self) -> bool { + self.provider.is_some() + } + + /// Get the refresh offset duration + pub fn refresh_offset(&self) -> Duration { + self.refresh_offset + } + + /// Get the storage options provider, if any + pub fn provider(&self) -> Option<&Arc> { + self.provider.as_ref() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mock_instant::thread_local::MockClock; + + #[derive(Debug)] + struct MockStorageOptionsProvider { + call_count: Arc>, + expires_in_millis: Option, + } + + impl MockStorageOptionsProvider { + fn new(expires_in_millis: Option) -> Self { + Self { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis, + } + } + + async fn get_call_count(&self) -> usize { + *self.call_count.read().await + } + } + + #[async_trait] + impl StorageOptionsProvider for MockStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let count = { + let mut c = self.call_count.write().await; + *c += 1; + *c + }; + + let mut options = HashMap::from([ + ("aws_access_key_id".to_string(), format!("AKID_{}", count)), + ( + "aws_secret_access_key".to_string(), + format!("SECRET_{}", count), + ), + ("aws_session_token".to_string(), format!("TOKEN_{}", count)), + ]); + + if let Some(expires_in) = self.expires_in_millis { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + let expires_at = now_ms + expires_in; + options.insert(EXPIRES_AT_MILLIS_KEY.to_string(), expires_at.to_string()); + } + + Ok(Some(options)) + } + + fn provider_id(&self) -> String { + let ptr = Arc::as_ptr(&self.call_count) as usize; + format!("MockStorageOptionsProvider {{ id: {} }}", ptr) + } + } + + #[tokio::test] + async fn test_static_options_only() { + let options = HashMap::from([ + ("key1".to_string(), "value1".to_string()), + ("key2".to_string(), "value2".to_string()), + ]); + let accessor = StorageOptionsAccessor::with_static_options(options.clone()); + + let result = accessor.get_storage_options().await.unwrap(); + assert_eq!(result.0, options); + assert!(!accessor.has_provider()); + assert_eq!(accessor.initial_storage_options(), Some(&options)); + } + + #[tokio::test] + async fn test_provider_only() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let mock_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000))); + let accessor = StorageOptionsAccessor::with_provider(mock_provider.clone()); + + let result = accessor.get_storage_options().await.unwrap(); + assert!(result.0.contains_key("aws_access_key_id")); + assert_eq!(result.0.get("aws_access_key_id").unwrap(), "AKID_1"); + assert!(accessor.has_provider()); + assert_eq!(accessor.initial_storage_options(), None); + assert_eq!(mock_provider.get_call_count().await, 1); + } + + #[tokio::test] + async fn test_initial_and_provider_uses_initial_first() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + let expires_at = now_ms + 600_000; // 10 minutes from now + + let initial = HashMap::from([ + ("aws_access_key_id".to_string(), "INITIAL_KEY".to_string()), + ( + "aws_secret_access_key".to_string(), + "INITIAL_SECRET".to_string(), + ), + (EXPIRES_AT_MILLIS_KEY.to_string(), expires_at.to_string()), + ]); + let mock_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000))); + + let accessor = StorageOptionsAccessor::with_initial_and_provider( + initial.clone(), + mock_provider.clone(), + ); + + // First call uses initial + let result = accessor.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("aws_access_key_id").unwrap(), "INITIAL_KEY"); + assert_eq!(mock_provider.get_call_count().await, 0); // Provider not called yet + } + + #[tokio::test] + async fn test_caching_and_refresh() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let mock_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000))); // 10 min expiry + // Use with_initial_and_provider to set custom refresh_offset_millis (5 min = 300000ms) + let now_ms = MockClock::system_time().as_millis() as u64; + let expires_at = now_ms + 600_000; // 10 minutes from now + let initial = HashMap::from([ + (EXPIRES_AT_MILLIS_KEY.to_string(), expires_at.to_string()), + (REFRESH_OFFSET_MILLIS_KEY.to_string(), "300000".to_string()), // 5 min refresh offset + ]); + let accessor = + StorageOptionsAccessor::with_initial_and_provider(initial, mock_provider.clone()); + + // First call uses initial cached options + let result = accessor.get_storage_options().await.unwrap(); + assert!(result.0.contains_key(EXPIRES_AT_MILLIS_KEY)); + assert_eq!(mock_provider.get_call_count().await, 0); + + // Advance time to 6 minutes - should trigger refresh (within 5 min refresh offset) + MockClock::set_system_time(Duration::from_secs(100_000 + 360)); + let result = accessor.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("aws_access_key_id").unwrap(), "AKID_1"); + assert_eq!(mock_provider.get_call_count().await, 1); + } + + #[tokio::test] + async fn test_expired_initial_triggers_refresh() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let now_ms = MockClock::system_time().as_millis() as u64; + let expired_time = now_ms - 1_000; // Expired 1 second ago + + let initial = HashMap::from([ + ("aws_access_key_id".to_string(), "EXPIRED_KEY".to_string()), + (EXPIRES_AT_MILLIS_KEY.to_string(), expired_time.to_string()), + ]); + let mock_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000))); + + let accessor = + StorageOptionsAccessor::with_initial_and_provider(initial, mock_provider.clone()); + + // Should fetch from provider since initial is expired + let result = accessor.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("aws_access_key_id").unwrap(), "AKID_1"); + assert_eq!(mock_provider.get_call_count().await, 1); + } + + #[tokio::test] + async fn test_accessor_id_with_provider() { + let mock_provider = Arc::new(MockStorageOptionsProvider::new(None)); + let accessor = StorageOptionsAccessor::with_provider(mock_provider); + + let id = accessor.accessor_id(); + assert!(id.starts_with("MockStorageOptionsProvider")); + } + + #[tokio::test] + async fn test_accessor_id_static() { + let options = HashMap::from([("key".to_string(), "value".to_string())]); + let accessor = StorageOptionsAccessor::with_static_options(options); + + let id = accessor.accessor_id(); + assert!(id.starts_with("static_options_")); + } + + #[tokio::test] + async fn test_concurrent_access() { + // Create a mock provider with far future expiration + let mock_provider = Arc::new(MockStorageOptionsProvider::new(Some(9999999999999))); + + let accessor = Arc::new(StorageOptionsAccessor::with_provider(mock_provider.clone())); + + // Spawn 10 concurrent tasks that all try to get options at the same time + let mut handles = vec![]; + for i in 0..10 { + let acc = accessor.clone(); + let handle = tokio::spawn(async move { + let result = acc.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("aws_access_key_id").unwrap(), "AKID_1"); + i + }); + handles.push(handle); + } + + // Wait for all tasks to complete + let results: Vec<_> = futures::future::join_all(handles) + .await + .into_iter() + .map(|r| r.unwrap()) + .collect(); + + // Verify all 10 tasks completed successfully + assert_eq!(results.len(), 10); + + // The provider should have been called exactly once + let call_count = mock_provider.get_call_count().await; + assert_eq!( + call_count, 1, + "Provider should be called exactly once despite concurrent access" + ); + } + + #[tokio::test] + async fn test_no_expiration_never_refreshes() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + let mock_provider = Arc::new(MockStorageOptionsProvider::new(None)); // No expiration + let accessor = StorageOptionsAccessor::with_provider(mock_provider.clone()); + + // First call fetches + accessor.get_storage_options().await.unwrap(); + assert_eq!(mock_provider.get_call_count().await, 1); + + // Advance time significantly + MockClock::set_system_time(Duration::from_secs(200_000)); + + // Should still use cached options + accessor.get_storage_options().await.unwrap(); + assert_eq!(mock_provider.get_call_count().await, 1); + } + + #[test] + fn test_parse_base_scoped_key() { + assert_eq!( + parse_base_scoped_key("base_1.account_key"), + Some((1, "account_key")) + ); + assert_eq!( + parse_base_scoped_key("base_12.headers.x-ms-version"), + Some((12, "headers.x-ms-version")) + ); + assert_eq!(parse_base_scoped_key("base_0.region"), Some((0, "region"))); + + // Not base-scoped keys + assert_eq!(parse_base_scoped_key("account_key"), None); + assert_eq!(parse_base_scoped_key("base_url"), None); + assert_eq!(parse_base_scoped_key("base_hot.account_key"), None); + assert_eq!(parse_base_scoped_key("base_1x.account_key"), None); + assert_eq!(parse_base_scoped_key("base_+1.account_key"), None); + assert_eq!(parse_base_scoped_key("base_.account_key"), None); + assert_eq!(parse_base_scoped_key("base_1."), None); + assert_eq!(parse_base_scoped_key("base_1"), None); + // Overflows u32 + assert_eq!(parse_base_scoped_key("base_4294967296.key"), None); + } + + #[test] + fn test_resolve_base_scoped_options() { + let options = HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ("base_2.account_key".to_string(), "base2-key".to_string()), + ("base_2.endpoint".to_string(), "http://b2".to_string()), + ]); + assert!(has_base_scoped_options(&options)); + + let base1 = resolve_base_scoped_options(&options, Some(1)); + assert_eq!( + base1, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "base1-key".to_string()), + ]) + ); + + let base2 = resolve_base_scoped_options(&options, Some(2)); + assert_eq!( + base2, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "base2-key".to_string()), + ("endpoint".to_string(), "http://b2".to_string()), + ]) + ); + + // A base without scoped entries inherits only the unscoped options + let base3 = resolve_base_scoped_options(&options, Some(3)); + assert_eq!( + base3, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "shared-key".to_string()), + ]) + ); + + // The default scope drops every scoped entry + let default = resolve_base_scoped_options(&options, None); + assert_eq!(default, base3); + + assert!(!has_base_scoped_options(&HashMap::from([( + "account_key".to_string(), + "shared-key".to_string() + )]))); + } + + #[tokio::test] + async fn test_scoped_to_base_identity_and_idempotency() { + // Static accessors without scoped keys are returned unchanged. + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [("account_key".to_string(), "shared-key".to_string())], + ))); + assert!(Arc::ptr_eq(&accessor.scoped_to_base(Some(1)), &accessor)); + assert!(Arc::ptr_eq(&accessor.scoped_to_base(None), &accessor)); + + // Scoping an already-scoped accessor is a no-op (the registry choke + // point re-applies the default scope to every params it sees). + let scoped = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [ + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ], + ))) + .scoped_to_base(Some(1)); + assert!(Arc::ptr_eq(&scoped.scoped_to_base(None), &scoped)); + + let provider_scoped = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + MockStorageOptionsProvider::new(None), + ))) + .scoped_to_base(Some(1)); + assert!(Arc::ptr_eq( + &provider_scoped.scoped_to_base(None), + &provider_scoped + )); + } + + #[tokio::test] + async fn test_scoped_to_base_provider_only_resolves_vended_options() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + // No initial options: scoped entries arrive only through the provider. + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let parent = Arc::new(StorageOptionsAccessor::with_provider(provider.clone())); + + let base1 = parent.scoped_to_base(Some(1)); + assert!(!Arc::ptr_eq(&base1, &parent)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert!(!result.0.contains_key("base_1.account_key")); + + let default = parent.scoped_to_base(None); + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_1"); + assert!(!result.0.contains_key("base_1.account_key")); + + // Both scopes were served from one parent fetch. + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_earlier_base_expiry_refreshes_parent() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + // Base 1 credentials expire before the shared ones; the parent must + // refresh when the earliest credential is due, or the scoped accessor + // would keep re-resolving stale base-1 credentials from a still- + // "valid" parent cache. + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ( + format!("base_1.{}", EXPIRES_AT_MILLIS_KEY), + (now_ms + 120_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + + let base1 = parent.scoped_to_base(Some(1)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_0"); + assert_eq!(*provider.call_count.read().await, 0); + + // Past the base-1 expiry but before the shared expiry: the parent's + // effective expiry is the earlier one, so the refresh chain fetches + // fresh credentials instead of re-serving BASE1_0. + MockClock::set_system_time(Duration::from_secs(100_000 + 121)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_to_base_static() { + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [ + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ], + ))); + + let base1 = accessor.scoped_to_base(Some(1)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!( + result.0, + HashMap::from([("account_key".to_string(), "base1-key".to_string())]) + ); + assert!(!base1.has_provider()); + + let default = accessor.scoped_to_base(None); + let result = default.get_storage_options().await.unwrap(); + assert_eq!( + result.0, + HashMap::from([("account_key".to_string(), "shared-key".to_string())]) + ); + + // Scoped accessor ids are stable across derivations and distinct per scope + assert_eq!( + accessor.scoped_to_base(Some(1)).accessor_id(), + base1.accessor_id() + ); + assert_ne!(base1.accessor_id(), default.accessor_id()); + assert_ne!(base1.accessor_id(), accessor.accessor_id()); + } + + #[derive(Debug)] + struct MockBaseScopedVendingProvider { + call_count: Arc>, + expires_in_millis: u64, + } + + #[async_trait] + impl StorageOptionsProvider for MockBaseScopedVendingProvider { + async fn fetch_storage_options(&self) -> Result>> { + let count = { + let mut c = self.call_count.write().await; + *c += 1; + *c + }; + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + Ok(Some(HashMap::from([ + ("account_key".to_string(), format!("SHARED_{}", count)), + ("base_1.account_key".to_string(), format!("BASE1_{}", count)), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + self.expires_in_millis).to_string(), + ), + ]))) + } + + fn provider_id(&self) -> String { + "MockBaseScopedVendingProvider".to_string() + } + } + + #[tokio::test] + async fn test_scoped_to_base_refreshes_through_parent() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + + let base1 = parent.scoped_to_base(Some(1)); + let default = parent.scoped_to_base(None); + assert!(base1.has_provider()); + + // Initial options are resolved per scope without fetching + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_0"); + assert!(!result.0.contains_key("base_1.account_key")); + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_0"); + assert_eq!(*provider.call_count.read().await, 0); + + // Expire the vended options; the scoped accessor refreshes through the + // parent and re-resolves the refreshed options for its scope. + MockClock::set_system_time(Duration::from_secs(100_000 + 601)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + + // The parent refresh is shared: other scopes see it without refetching + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_1"); + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_forced_refresh_reaches_origin_provider() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + let base1 = parent.scoped_to_base(Some(1)); + + // A forced refresh must reach the origin provider even though both the + // scoped and the parent caches are still valid. + let result = base1.refresh_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + } +} diff --git a/vendor/lance-io/src/object_store/test_utils.rs b/vendor/lance-io/src/object_store/test_utils.rs new file mode 100644 index 000000000..b22ff4912 --- /dev/null +++ b/vendor/lance-io/src/object_store/test_utils.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; + +use async_trait::async_trait; + +use super::StorageOptionsProvider; +use lance_core::Result; + +#[derive(Debug)] +pub struct StaticMockStorageOptionsProvider { + pub options: HashMap, +} + +#[async_trait] +impl StorageOptionsProvider for StaticMockStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + Ok(Some(self.options.clone())) + } + + fn provider_id(&self) -> String { + "StaticMockStorageOptionsProvider".to_string() + } +} diff --git a/vendor/lance-io/src/object_store/throttle.rs b/vendor/lance-io/src/object_store/throttle.rs new file mode 100644 index 000000000..4876de01b --- /dev/null +++ b/vendor/lance-io/src/object_store/throttle.rs @@ -0,0 +1,2097 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! AIMD-controlled token bucket rate limiter for ObjectStore operations. +//! +//! Wraps any [`object_store::ObjectStore`] with per-category token buckets +//! whose fill rates are dynamically adjusted by AIMD controllers. When cloud +//! stores return HTTP 429/503, the fill rate decreases multiplicatively. During +//! sustained success windows, it increases additively. +//! +//! Operations are split into four independent categories — **read**, **write**, +//! **delete**, **list** — each with its own AIMD controller and token bucket. +//! This prevents a burst of reads from starving writes, and vice versa. +//! +//! # Example +//! +//! ```ignore +//! use lance_io::object_store::throttle::{AimdThrottleConfig, AimdThrottledStore}; +//! +//! let throttled = AimdThrottledStore::new(target, AimdThrottleConfig::default()).unwrap(); +//! ``` + +use std::collections::HashMap; +use std::fmt::{Debug, Display, Formatter}; +use std::ops::Range; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; +use lance_core::utils::aimd::{AimdConfig, AimdController, RequestOutcome}; +use lance_core::utils::tracing::TRACE_OBJECT_STORE_THROTTLE; +#[cfg(test)] +use object_store::ObjectStoreExt; +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +use object_store::client::{ + ClientOptions, HttpClient, HttpConnector, HttpError, HttpErrorKind, HttpRequest, HttpResponse, + HttpResponseBody, HttpService, +}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, + PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult, + UploadPart, +}; +use rand::Rng; +use tokio::sync::Mutex; +use tracing::{debug, warn}; + +/// Check whether an `object_store::Error` represents a throttle response +/// (HTTP 429 / 503) from a cloud object store. +/// +/// Regrettably, this information is not fully exposed by the `object_store` crate. +/// There is no generic mechanism for a custom object store to return a throttle error. +/// +/// However, the builtin object stores all use RetryError when retries are configured and +/// throttle errors are returned. Sadly, RetryError is not a public type, so we have to +/// infer it from the error message. This is potentially dangerous because these errors +/// often include the URI itself and that URI could have any characters in it (e.g. if we +/// look for 429 then we might match a 429 in a UUID).These error messages currently look like: +/// +/// ", after ... retries, max_retries: ..., retry_timeout: ..." +/// +/// So, as a crude heuristic, which should work for the builtin object stores, but won't +/// work for custom object stores, we simply look for the string "retries, max_retries" +/// in the error message. +pub fn is_throttle_error(err: &object_store::Error) -> bool { + // Only Generic errors can carry throttle responses + if let object_store::Error::Generic { source, .. } = err { + let message = source.to_string(); + let lowercase = message.to_ascii_lowercase(); + lowercase.contains("retries, max_retries") + || lowercase.contains("serverbusy") + || lowercase.contains("server busy") + || lowercase.contains("egress is over the account limit") + || lowercase.contains("http 429") + || lowercase.contains("status code: 429") + || lowercase.contains("429 too many requests") + || lowercase.contains("too many requests") + || lowercase.contains("slowdown") + || lowercase.contains("please reduce your request rate") + || lowercase.contains("rate limit") + || lowercase.contains("throttling") + || lowercase.contains("throttled") + } else { + false + } +} + +/// Configuration for the AIMD-throttled ObjectStore wrapper. +/// +/// Each operation category (read, write, delete, list) has its own AIMD config. +/// Use [`with_aimd`](AimdThrottleConfig::with_aimd) to set all categories at +/// once, or per-category methods like [`with_read_aimd`](AimdThrottleConfig::with_read_aimd) +/// for fine-grained control. +#[derive(Debug, Clone)] +pub struct AimdThrottleConfig { + /// AIMD configuration for read operations (get, get_opts, get_range, get_ranges, head). + pub read: AimdConfig, + /// AIMD configuration for write operations (put, put_opts, put_multipart, copy, rename, etc.). + pub write: AimdConfig, + /// AIMD configuration for delete operations. + pub delete: AimdConfig, + /// AIMD configuration for list operations. + pub list: AimdConfig, + /// Maximum tokens that can accumulate for bursts (shared across all categories). + pub burst_capacity: u32, + /// Maximum number of retries for throttle errors within the AIMD layer. + pub max_retries: usize, + /// Minimum backoff in milliseconds between retry attempts. + pub min_backoff_ms: u64, + /// Maximum backoff in milliseconds between retry attempts. + pub max_backoff_ms: u64, +} + +impl Default for AimdThrottleConfig { + fn default() -> Self { + let aimd = AimdConfig::default(); + Self { + read: aimd.clone(), + write: aimd.clone(), + delete: aimd.clone(), + list: aimd, + burst_capacity: 100, + max_retries: 3, + min_backoff_ms: 100, + max_backoff_ms: 300, + } + } +} + +impl AimdThrottleConfig { + /// Set the AIMD configuration for all four operation categories at once. + pub fn with_aimd(self, aimd: AimdConfig) -> Self { + Self { + read: aimd.clone(), + write: aimd.clone(), + delete: aimd.clone(), + list: aimd, + ..self + } + } + + /// Set the AIMD configuration for read operations. + pub fn with_read_aimd(self, aimd: AimdConfig) -> Self { + Self { read: aimd, ..self } + } + + /// Set the AIMD configuration for write operations. + pub fn with_write_aimd(self, aimd: AimdConfig) -> Self { + Self { + write: aimd, + ..self + } + } + + /// Set the AIMD configuration for delete operations. + pub fn with_delete_aimd(self, aimd: AimdConfig) -> Self { + Self { + delete: aimd, + ..self + } + } + + /// Set the AIMD configuration for list operations. + pub fn with_list_aimd(self, aimd: AimdConfig) -> Self { + Self { list: aimd, ..self } + } + + /// Returns `true` when the AIMD throttle layer should be bypassed entirely. + pub fn is_disabled(&self) -> bool { + self.max_retries == 0 + } + + pub fn with_burst_capacity(self, burst_capacity: u32) -> Self { + Self { + burst_capacity, + ..self + } + } + + /// Build an `AimdThrottleConfig` from storage options and environment variables. + /// + /// Storage options take precedence over environment variables, which take + /// precedence over defaults. A single AIMD config is applied to all four + /// operation categories (read/write/delete/list). + /// + /// | Setting | Storage Option Key | Env Var | Default | + /// |----------------------|----------------------------------|----------------------------------|---------| + /// | Initial rate | `lance_aimd_initial_rate` | `LANCE_AIMD_INITIAL_RATE` | 2000 | + /// | Min rate | `lance_aimd_min_rate` | `LANCE_AIMD_MIN_RATE` | 1 | + /// | Max rate | `lance_aimd_max_rate` | `LANCE_AIMD_MAX_RATE` | 5000 | + /// | Decrease factor | `lance_aimd_decrease_factor` | `LANCE_AIMD_DECREASE_FACTOR` | 0.5 | + /// | Additive increment | `lance_aimd_additive_increment` | `LANCE_AIMD_ADDITIVE_INCREMENT` | 300 | + /// | Burst capacity | `lance_aimd_burst_capacity` | `LANCE_AIMD_BURST_CAPACITY` | 100 | + /// | Max retries | `lance_aimd_max_retries` | `LANCE_AIMD_MAX_RETRIES` | 3 | + /// | Min backoff ms | `lance_aimd_min_backoff_ms` | `LANCE_AIMD_MIN_BACKOFF_MS` | 100 | + /// | Max backoff ms | `lance_aimd_max_backoff_ms` | `LANCE_AIMD_MAX_BACKOFF_MS` | 300 | + pub fn from_storage_options( + storage_options: Option<&HashMap>, + ) -> lance_core::Result { + fn resolve_f64( + key: &str, + storage_options: Option<&HashMap>, + default: f64, + ) -> lance_core::Result { + let env_key = key.to_ascii_uppercase(); + if let Some(val) = storage_options.and_then(|opts| opts.get(key)) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for storage option '{key}': '{val}'" + )) + }) + } else if let Ok(val) = std::env::var(&env_key) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for env var '{env_key}': '{val}'" + )) + }) + } else { + Ok(default) + } + } + + fn resolve_u32( + key: &str, + storage_options: Option<&HashMap>, + default: u32, + ) -> lance_core::Result { + let env_key = key.to_ascii_uppercase(); + if let Some(val) = storage_options.and_then(|opts| opts.get(key)) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for storage option '{key}': '{val}'" + )) + }) + } else if let Ok(val) = std::env::var(&env_key) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for env var '{env_key}': '{val}'" + )) + }) + } else { + Ok(default) + } + } + + fn resolve_usize( + key: &str, + storage_options: Option<&HashMap>, + default: usize, + ) -> lance_core::Result { + let env_key = key.to_ascii_uppercase(); + if let Some(val) = storage_options.and_then(|opts| opts.get(key)) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for storage option '{key}': '{val}'" + )) + }) + } else if let Ok(val) = std::env::var(&env_key) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for env var '{env_key}': '{val}'" + )) + }) + } else { + Ok(default) + } + } + + fn resolve_u64( + key: &str, + storage_options: Option<&HashMap>, + default: u64, + ) -> lance_core::Result { + let env_key = key.to_ascii_uppercase(); + if let Some(val) = storage_options.and_then(|opts| opts.get(key)) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for storage option '{key}': '{val}'" + )) + }) + } else if let Ok(val) = std::env::var(&env_key) { + val.parse::().map_err(|_| { + lance_core::Error::invalid_input(format!( + "Invalid value for env var '{env_key}': '{val}'" + )) + }) + } else { + Ok(default) + } + } + + let initial_rate = resolve_f64("lance_aimd_initial_rate", storage_options, 2000.0)?; + let min_rate = resolve_f64("lance_aimd_min_rate", storage_options, 1.0)?; + let max_rate = resolve_f64("lance_aimd_max_rate", storage_options, 5000.0)?; + let decrease_factor = resolve_f64("lance_aimd_decrease_factor", storage_options, 0.5)?; + let additive_increment = + resolve_f64("lance_aimd_additive_increment", storage_options, 300.0)?; + let burst_capacity = resolve_u32("lance_aimd_burst_capacity", storage_options, 100)?; + let max_retries = resolve_usize("lance_aimd_max_retries", storage_options, 3)?; + let min_backoff_ms = resolve_u64("lance_aimd_min_backoff_ms", storage_options, 100)?; + let max_backoff_ms = resolve_u64("lance_aimd_max_backoff_ms", storage_options, 300)?; + + let aimd = AimdConfig::default() + .with_initial_rate(initial_rate) + .with_min_rate(min_rate) + .with_max_rate(max_rate) + .with_decrease_factor(decrease_factor) + .with_additive_increment(additive_increment); + + Ok(Self { + max_retries, + min_backoff_ms, + max_backoff_ms, + ..Self::default() + .with_aimd(aimd) + .with_burst_capacity(burst_capacity) + }) + } +} + +struct TokenBucketState { + tokens: f64, + last_refill: tokio::time::Instant, + rate: f64, +} + +/// Per-category throttle state: an AIMD controller paired with a token bucket. +struct OperationThrottle { + controller: AimdController, + bucket: Mutex, + burst_capacity: f64, + max_retries: usize, + min_backoff_ms: u64, + max_backoff_ms: u64, +} + +impl OperationThrottle { + fn new( + aimd_config: AimdConfig, + burst_capacity: f64, + max_retries: usize, + min_backoff_ms: u64, + max_backoff_ms: u64, + ) -> lance_core::Result { + let initial_rate = aimd_config.initial_rate; + let controller = AimdController::new(aimd_config)?; + Ok(Self { + controller, + bucket: Mutex::new(TokenBucketState { + tokens: burst_capacity, + last_refill: tokio::time::Instant::now(), + rate: initial_rate, + }), + burst_capacity, + max_retries, + min_backoff_ms, + max_backoff_ms, + }) + } + + /// Acquire a token from the bucket, sleeping if none are available. + /// + /// Each caller reserves a token immediately (allowing `tokens` to go + /// negative) so that concurrent waiters queue behind each other instead + /// of all waking at the same instant (thundering herd). + async fn acquire_token(&self) { + let sleep_duration = { + let mut bucket = self.bucket.lock().await; + let now = tokio::time::Instant::now(); + let elapsed = now.duration_since(bucket.last_refill).as_secs_f64(); + bucket.tokens = (bucket.tokens + elapsed * bucket.rate).min(self.burst_capacity); + bucket.last_refill = now; + + // Reserve a token (may go negative to queue behind other waiters) + bucket.tokens -= 1.0; + + if bucket.tokens >= 0.0 { + // Had a token available, no need to sleep + return; + } + + // Sleep proportional to our position in the queue + std::time::Duration::from_secs_f64(-bucket.tokens / bucket.rate) + }; + + tokio::time::sleep(sleep_duration).await; + } + + /// Update the bucket's fill rate from the controller. + async fn update_bucket_rate(&self, new_rate: f64) { + let mut bucket = self.bucket.lock().await; + bucket.rate = new_rate; + } + + /// Classify a result and feed it back to the AIMD controller without + /// acquiring a token. Uses `try_lock` for the bucket update so that if the + /// bucket lock is contended the rate update is deferred to the next + /// `throttled()` call. + fn observe_outcome(&self, result: &OSResult) { + let outcome = match result { + Ok(_) => RequestOutcome::Success, + Err(err) if is_throttle_error(err) => { + debug!( + target: TRACE_OBJECT_STORE_THROTTLE, + error = %err, + "Throttle error detected in stream" + ); + RequestOutcome::Throttled + } + Err(_) => RequestOutcome::Success, + }; + let error = result + .as_ref() + .err() + .map(|error| error as &dyn std::fmt::Display); + let new_rate = self.record_outcome(outcome, error); + if let Ok(mut bucket) = self.bucket.try_lock() { + bucket.rate = new_rate; + } + } + + fn record_outcome( + &self, + outcome: RequestOutcome, + error: Option<&dyn std::fmt::Display>, + ) -> f64 { + let prev_rate = self.controller.current_rate(); + let new_rate = self.controller.record_outcome(outcome); + if new_rate < prev_rate { + if let Some(error) = error { + warn!( + target: TRACE_OBJECT_STORE_THROTTLE, + previous_rate = format!("{prev_rate:.1}"), + new_rate = format!("{new_rate:.1}"), + error = %error, + "AIMD throttle: rate reduced due to throttle errors" + ); + } else { + warn!( + target: TRACE_OBJECT_STORE_THROTTLE, + previous_rate = format!("{prev_rate:.1}"), + new_rate = format!("{new_rate:.1}"), + "AIMD throttle: rate reduced due to throttle errors" + ); + } + } + new_rate + } + + /// Execute an operation with throttling: acquire token, run, classify result. + /// On throttle errors, retries up to `max_retries` times with a random + /// backoff between `min_backoff_ms` and `max_backoff_ms` between attempts. + async fn throttled(&self, f: F) -> OSResult + where + F: Fn() -> Fut, + Fut: std::future::Future>, + { + for attempt in 0..=self.max_retries { + self.acquire_token().await; + let result = f().await; + let outcome = match &result { + Ok(_) => RequestOutcome::Success, + Err(err) if is_throttle_error(err) => { + debug!( + target: TRACE_OBJECT_STORE_THROTTLE, + error = %err, + "Throttle error detected" + ); + RequestOutcome::Throttled + } + Err(_) => RequestOutcome::Success, // Non-throttle errors don't indicate capacity problems + }; + let error = result + .as_ref() + .err() + .map(|error| error as &dyn std::fmt::Display); + let new_rate = self.record_outcome(outcome, error); + self.update_bucket_rate(new_rate).await; + + match &result { + Err(err) if is_throttle_error(err) && attempt < self.max_retries => { + let backoff_ms = + rand::rng().random_range(self.min_backoff_ms..=self.max_backoff_ms); + debug!( + target: TRACE_OBJECT_STORE_THROTTLE, + attempt = attempt + 1, + max_retries = self.max_retries, + backoff_ms, + error = %err, + "Retrying after throttle error" + ); + tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; + continue; + } + _ => return result, + } + } + unreachable!() + } +} + +impl Debug for OperationThrottle { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OperationThrottle") + .field("controller", &self.controller) + .field("burst_capacity", &self.burst_capacity) + .finish() + } +} + +#[derive(Clone)] +pub(crate) struct AimdThrottleState { + read: Arc, + write: Arc, + delete: Arc, + list: Arc, +} + +impl AimdThrottleState { + pub(crate) fn new(config: AimdThrottleConfig) -> lance_core::Result { + let burst_capacity = config.burst_capacity as f64; + let max_retries = config.max_retries; + let min_backoff_ms = config.min_backoff_ms; + let max_backoff_ms = config.max_backoff_ms; + Ok(Self { + read: Arc::new(OperationThrottle::new( + config.read, + burst_capacity, + max_retries, + min_backoff_ms, + max_backoff_ms, + )?), + write: Arc::new(OperationThrottle::new( + config.write, + burst_capacity, + max_retries, + min_backoff_ms, + max_backoff_ms, + )?), + delete: Arc::new(OperationThrottle::new( + config.delete, + burst_capacity, + max_retries, + min_backoff_ms, + max_backoff_ms, + )?), + list: Arc::new(OperationThrottle::new( + config.list, + burst_capacity, + max_retries, + min_backoff_ms, + max_backoff_ms, + )?), + }) + } +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +#[derive(Debug)] +pub(crate) struct AimdMultipartUploadConnector { + inner: C, + write: Option>, +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +impl AimdMultipartUploadConnector { + fn new(inner: C, state: Option<&AimdThrottleState>) -> Self { + Self { + inner, + write: state.map(|state| Arc::clone(&state.write)), + } + } +} + +#[cfg(all( + any(feature = "aws", feature = "azure", feature = "gcp"), + feature = "metrics" +))] +pub(crate) fn cloud_http_connector( + state: Option<&AimdThrottleState>, + metrics_base: String, +) -> AimdMultipartUploadConnector { + AimdMultipartUploadConnector::new( + crate::object_store::metrics::MeteringHttpConnector::new(metrics_base), + state, + ) +} + +#[cfg(all( + any(feature = "aws", feature = "azure", feature = "gcp"), + not(feature = "metrics") +))] +pub(crate) fn cloud_http_connector( + state: Option<&AimdThrottleState>, + _metrics_base: String, +) -> AimdMultipartUploadConnector { + AimdMultipartUploadConnector::new(object_store::client::ReqwestConnector::default(), state) +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +impl HttpConnector for AimdMultipartUploadConnector { + fn connect(&self, options: &ClientOptions) -> object_store::Result { + Ok(HttpClient::new(AimdMultipartUploadService { + inner: self.inner.connect(options)?, + write: self.write.clone(), + })) + } +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +#[derive(Debug)] +struct AimdMultipartUploadService { + inner: HttpClient, + write: Option>, +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +fn is_multipart_part_request(request: &HttpRequest) -> bool { + if request.method() != ::http::Method::PUT { + return false; + } + request.uri().query().is_some_and(|query| { + url::form_urlencoded::parse(query.as_bytes()).any(|(key, value)| { + key.eq_ignore_ascii_case("partNumber") + || (key.eq_ignore_ascii_case("comp") && value.eq_ignore_ascii_case("block")) + }) + }) +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +fn is_retryable_http_error(error: &HttpError) -> bool { + matches!( + error.kind(), + HttpErrorKind::Connect + | HttpErrorKind::Request + | HttpErrorKind::Timeout + | HttpErrorKind::Interrupted + ) +} + +#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] +#[async_trait] +impl HttpService for AimdMultipartUploadService { + async fn call(&self, request: HttpRequest) -> Result { + let Some(write) = self.write.as_ref() else { + return self.inner.execute(request).await; + }; + if !is_multipart_part_request(&request) { + return self.inner.execute(request).await; + } + + for attempt in 0..=write.max_retries { + write.acquire_token().await; + let mut result = self.inner.execute(request.clone()).await; + let mut is_retryable = result.as_ref().err().is_some_and(is_retryable_http_error); + let mut is_throttle = false; + let mut response_status = None; + + if let Ok(response) = result { + let status = response.status(); + response_status = Some(status); + is_retryable = status == ::http::StatusCode::REQUEST_TIMEOUT + || status == ::http::StatusCode::TOO_MANY_REQUESTS + || status.is_server_error(); + is_throttle = status == ::http::StatusCode::TOO_MANY_REQUESTS + || status == ::http::StatusCode::SERVICE_UNAVAILABLE; + + let (parts, body) = response.into_parts(); + result = match body.bytes().await { + Ok(bytes) => { + let body = String::from_utf8_lossy(&bytes).to_ascii_lowercase(); + let is_throttle_body = body.contains("requesttimeout") + || body.contains("slowdown") + || body.contains("serverbusy") + || body.contains("throttl"); + is_retryable |= is_throttle_body; + is_throttle |= is_throttle_body; + Ok(HttpResponse::from_parts( + parts, + HttpResponseBody::from(bytes), + )) + } + Err(error) => { + is_retryable = is_retryable_http_error(&error); + Err(error) + } + }; + } + + let detail = response_status + .filter(|status| !status.is_success()) + .map(|status| format!("HTTP status {status}")); + let error = result + .as_ref() + .err() + .map(|error| error as &dyn std::fmt::Display) + .or_else(|| { + detail + .as_ref() + .map(|detail| detail as &dyn std::fmt::Display) + }); + let outcome = if is_throttle { + RequestOutcome::Throttled + } else { + RequestOutcome::Success + }; + let new_rate = write.record_outcome(outcome, error); + write.update_bucket_rate(new_rate).await; + + if is_retryable && attempt < write.max_retries { + let backoff_ms = + rand::rng().random_range(write.min_backoff_ms..=write.max_backoff_ms); + debug!( + target: TRACE_OBJECT_STORE_THROTTLE, + attempt = attempt + 1, + max_retries = write.max_retries, + backoff_ms, + "Retrying multipart upload part after retryable HTTP response" + ); + tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; + continue; + } + return result; + } + unreachable!() + } +} + +/// A [`MultipartUpload`] wrapper that applies the write AIMD controller. +struct ThrottledMultipartUpload { + target: Box, + write: Arc, + parts_throttled_at_http: bool, +} + +impl Debug for ThrottledMultipartUpload { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ThrottledMultipartUpload").finish() + } +} + +#[async_trait] +impl MultipartUpload for ThrottledMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + // Call put_part synchronously to preserve part ordering regardless + // of which futures are awaited first. + let fut = self.target.put_part(data); + if self.parts_throttled_at_http { + return fut; + } + let write = Arc::clone(&self.write); + Box::pin(async move { + write.acquire_token().await; + let result = fut.await; + write.observe_outcome(&result); + result + }) + } + + async fn complete(&mut self) -> OSResult { + let target = &mut self.target; + for attempt in 0..=self.write.max_retries { + self.write.acquire_token().await; + let result = target.complete().await; + self.write.observe_outcome(&result); + + match &result { + Err(err) if is_throttle_error(err) && attempt < self.write.max_retries => { + let backoff_ms = rand::rng() + .random_range(self.write.min_backoff_ms..=self.write.max_backoff_ms); + tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; + continue; + } + _ => return result, + } + } + unreachable!() + } + + async fn abort(&mut self) -> OSResult<()> { + let target = &mut self.target; + for attempt in 0..=self.write.max_retries { + self.write.acquire_token().await; + let result = target.abort().await; + self.write.observe_outcome(&result); + + match &result { + Err(err) if is_throttle_error(err) && attempt < self.write.max_retries => { + let backoff_ms = rand::rng() + .random_range(self.write.min_backoff_ms..=self.write.max_backoff_ms); + tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; + continue; + } + _ => return result, + } + } + unreachable!() + } +} + +/// An ObjectStore wrapper that rate-limits operations using per-category token +/// buckets whose fill rates are controlled by AIMD algorithms. +/// +/// Operations are split into four independent categories: +/// - **read**: `get`, `get_opts`, `get_range`, `get_ranges`, `head` +/// - **write**: `put`, `put_opts`, `put_multipart`, `put_multipart_opts`, `copy`, `copy_if_not_exists`, `rename`, `rename_if_not_exists` +/// - **delete**: `delete` +/// - **list**: `list`, `list_with_offset`, `list_with_delimiter` +/// +/// Streaming list operations acquire a token before starting the underlying list stream. +/// Streaming operations also observe each yielded item and feed the result back to the +/// AIMD controller so it can adjust the rate for other operations in the same category. +/// +/// This is not perfect but probably as close as we can get without moving the throttle into +/// the object_store crate itself. +pub struct AimdThrottledStore { + target: Arc, + read: Arc, + write: Arc, + delete: Arc, + list: Arc, + multipart_parts_throttled_at_http: bool, +} + +impl Debug for AimdThrottledStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AimdThrottledStore") + .field("target", &self.target) + .field("read", &self.read) + .field("write", &self.write) + .field("delete", &self.delete) + .field("list", &self.list) + .field( + "multipart_parts_throttled_at_http", + &self.multipart_parts_throttled_at_http, + ) + .finish() + } +} + +impl Display for AimdThrottledStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "AimdThrottledStore({})", self.target) + } +} + +impl AimdThrottledStore { + pub fn new( + target: Arc, + config: AimdThrottleConfig, + ) -> lance_core::Result { + Ok(Self::new_with_state( + target, + AimdThrottleState::new(config)?, + false, + )) + } + + pub(crate) fn new_with_state( + target: Arc, + state: AimdThrottleState, + multipart_parts_throttled_at_http: bool, + ) -> Self { + Self { + target, + read: state.read, + write: state.write, + delete: state.delete, + list: state.list, + multipart_parts_throttled_at_http, + } + } +} + +#[async_trait] +#[deny(clippy::missing_trait_methods)] +impl ObjectStore for AimdThrottledStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.write + .throttled(|| self.target.put_opts(location, bytes.clone(), opts.clone())) + .await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + let target = self + .write + .throttled(|| self.target.put_multipart_opts(location, opts.clone())) + .await?; + Ok(Box::new(ThrottledMultipartUpload { + target, + write: Arc::clone(&self.write), + parts_throttled_at_http: self.multipart_parts_throttled_at_http, + })) + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.read + .throttled(|| self.target.get_opts(location, options.clone())) + .await + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.read + .throttled(|| self.target.get_ranges(location, ranges)) + .await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + let delete = Arc::clone(&self.delete); + self.target + .delete_stream(locations) + .map(move |item| { + delete.observe_outcome(&item); + item + }) + .boxed() + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + let throttle = Arc::clone(&self.list); + let throttle_for_start = Arc::clone(&throttle); + let target = Arc::clone(&self.target); + let prefix = prefix.cloned(); + futures::stream::once(async move { + throttle_for_start.acquire_token().await; + target.list(prefix.as_ref()) + }) + .flatten() + .map(move |item| { + throttle.observe_outcome(&item); + item + }) + .boxed() + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + let throttle = Arc::clone(&self.list); + let throttle_for_start = Arc::clone(&throttle); + let target = Arc::clone(&self.target); + let prefix = prefix.cloned(); + let offset = offset.clone(); + futures::stream::once(async move { + throttle_for_start.acquire_token().await; + target.list_with_offset(prefix.as_ref(), &offset) + }) + .flatten() + .map(move |item| { + throttle.observe_outcome(&item); + item + }) + .boxed() + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.list + .throttled(|| self.target.list_with_delimiter(prefix)) + .await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.write + .throttled(|| self.target.copy_opts(from, to, opts.clone())) + .await + } + + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OSResult<()> { + self.write + .throttled(|| self.target.rename_opts(from, to, opts.clone())) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use object_store::memory::InMemory; + use rstest::rstest; + use std::collections::VecDeque; + use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; + + const THROTTLE_ERROR_RESPONSE: &str = "request failed, after 3 retries, max_retries: 3, retry_timeout: 30s - Server returned non-2xx status code: 503: x-ms-request-id: azure-request-id"; + + fn make_generic_error(msg: &str) -> object_store::Error { + object_store::Error::Generic { + store: "test", + source: msg.into(), + } + } + + #[rstest] + #[case::retry_error("Error after 10 retries, max_retries: 10, retry_timeout: 180s", true)] + #[case::retries_in_message( + "request failed, after 3 retries, max_retries: 5, retry_timeout: 60s", + true + )] + #[case::not_found("Object not found", false)] + #[case::permission_denied("Access denied", false)] + #[case::timeout("Connection timed out", false)] + #[case::http_429_without_retries("HTTP 429 Too Many Requests", true)] + #[case::slowdown_without_retries("SlowDown: Please reduce your request rate", true)] + #[case::azure_server_busy("Code: ServerBusy", true)] + #[case::azure_egress_limit("Message: Egress is over the account limit", true)] + fn test_is_throttle_error(#[case] msg: &str, #[case] expected: bool) { + let err = make_generic_error(msg); + assert_eq!( + is_throttle_error(&err), + expected, + "is_throttle_error for '{}' should be {}", + msg, + expected + ); + } + + #[test] + fn test_non_generic_errors_are_not_throttle() { + let err = object_store::Error::NotFound { + path: "test".to_string(), + source: "not found".into(), + }; + assert!(!is_throttle_error(&err)); + } + + #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + #[rstest] + #[case::s3("https://bucket/object?partNumber=1&uploadId=id", true)] + #[case::azure_block("https://account/object?comp=block&blockid=id", true)] + #[case::azure_block_list("https://account/object?comp=blocklist", false)] + #[case::ordinary_put("https://bucket/object", false)] + fn test_is_multipart_part_request(#[case] uri: &str, #[case] expected: bool) { + let request = ::http::Request::builder() + .method(::http::Method::PUT) + .uri(uri) + .body(object_store::client::HttpRequestBody::empty()) + .unwrap(); + assert_eq!(is_multipart_part_request(&request), expected); + } + + #[tokio::test] + async fn test_basic_put_get_through_wrapper() { + let store = Arc::new(InMemory::new()); + let config = AimdThrottleConfig::default(); + let throttled = AimdThrottledStore::new(store, config).unwrap(); + + let path = Path::from("test/file.txt"); + let data = PutPayload::from_static(b"hello world"); + throttled.put(&path, data).await.unwrap(); + + let result = throttled.get(&path).await.unwrap(); + let bytes = result.bytes().await.unwrap(); + assert_eq!(bytes.as_ref(), b"hello world"); + } + + #[tokio::test] + async fn test_rate_decreases_on_throttle() { + let store = Arc::new(InMemory::new()); + let config = AimdThrottleConfig::default().with_aimd( + AimdConfig::default() + .with_initial_rate(100.0) + .with_decrease_factor(0.5) + .with_window_duration(std::time::Duration::from_millis(10)), + ); + let throttled = AimdThrottledStore::new(store, config).unwrap(); + + let initial_rate = throttled.read.controller.current_rate(); + assert_eq!(initial_rate, 100.0); + + // Simulate a throttle outcome directly + throttled + .read + .controller + .record_outcome(RequestOutcome::Throttled); + + // Wait for window to expire and trigger evaluation + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + throttled + .read + .controller + .record_outcome(RequestOutcome::Success); + + let new_rate = throttled.read.controller.current_rate(); + assert!( + new_rate < initial_rate, + "Rate should decrease after throttle: {} < {}", + new_rate, + initial_rate + ); + } + + #[tokio::test] + async fn test_rate_recovers_on_success() { + let store = Arc::new(InMemory::new()); + let config = AimdThrottleConfig::default().with_aimd( + AimdConfig::default() + .with_initial_rate(100.0) + .with_decrease_factor(0.5) + .with_additive_increment(10.0) + .with_window_duration(std::time::Duration::from_millis(10)), + ); + let throttled = AimdThrottledStore::new(store, config).unwrap(); + + // First decrease via throttle + throttled + .read + .controller + .record_outcome(RequestOutcome::Throttled); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + throttled + .read + .controller + .record_outcome(RequestOutcome::Success); + let decreased_rate = throttled.read.controller.current_rate(); + assert_eq!(decreased_rate, 50.0); + + // Now recover via success + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + throttled + .read + .controller + .record_outcome(RequestOutcome::Success); + let recovered_rate = throttled.read.controller.current_rate(); + assert_eq!(recovered_rate, 60.0); + } + + #[tokio::test] + async fn test_as_dyn_object_store() { + let store: Arc = Arc::new(InMemory::new()); + let throttled: Arc = + Arc::new(AimdThrottledStore::new(store, AimdThrottleConfig::default()).unwrap()); + + let path = Path::from("test/data.bin"); + let data = PutPayload::from_static(b"test data"); + throttled.put(&path, data).await.unwrap(); + + let result = throttled.get(&path).await.unwrap(); + let bytes = result.bytes().await.unwrap(); + assert_eq!(bytes.as_ref(), b"test data"); + } + + #[tokio::test] + async fn test_token_bucket_delays_when_exhausted() { + let store = Arc::new(InMemory::new()); + // Very low rate and burst capacity to force waiting + let config = AimdThrottleConfig::default() + .with_burst_capacity(1) + .with_aimd(AimdConfig::default().with_initial_rate(10.0)); + let throttled = Arc::new(AimdThrottledStore::new(store, config).unwrap()); + + let path = Path::from("test/file.txt"); + let data = PutPayload::from_static(b"data"); + throttled.put(&path, data).await.unwrap(); + + // After consuming the burst token, the next request should take ~100ms + // (1 token / 10 tokens-per-sec). We verify it takes at least 50ms. + let start = std::time::Instant::now(); + let data2 = PutPayload::from_static(b"data2"); + throttled.put(&path, data2).await.unwrap(); + let elapsed = start.elapsed(); + + assert!( + elapsed >= std::time::Duration::from_millis(50), + "Expected delay for token refill, but elapsed was {:?}", + elapsed + ); + } + + #[tokio::test] + async fn test_list_observes_outcomes() { + let store = Arc::new(InMemory::new()); + let config = AimdThrottleConfig::default(); + let throttled = AimdThrottledStore::new(store.clone(), config).unwrap(); + + let path = Path::from("prefix/file.txt"); + let data = PutPayload::from_static(b"data"); + store.put(&path, data).await.unwrap(); + + let items: Vec<_> = throttled.list(Some(&Path::from("prefix"))).collect().await; + assert_eq!(items.len(), 1); + assert!(items[0].is_ok()); + } + + /// A mock store whose `list` stream yields a configurable sequence of + /// Ok / throttle-error items. Used to verify that the AIMD wrapper + /// observes errors surfaced inside list streams. + struct ThrottlingListMockStore { + inner: InMemory, + /// Number of throttle errors to inject at the start of each list call. + throttle_count: usize, + } + + impl Display for ThrottlingListMockStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "ThrottlingListMockStore") + } + } + + impl Debug for ThrottlingListMockStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ThrottlingListMockStore").finish() + } + } + + #[async_trait] + impl ObjectStore for ThrottlingListMockStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.inner.put_opts(location, bytes, opts).await + } + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.inner.get_opts(location, options).await + } + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.inner.get_ranges(location, ranges).await + } + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + let n = self.throttle_count; + let inner_stream = self.inner.list(prefix); + let errors = futures::stream::iter((0..n).map(|_| { + Err(object_store::Error::Generic { + store: "ThrottlingListMock", + source: "request failed, after 3 retries, max_retries: 5, retry_timeout: 60s" + .into(), + }) + })); + errors.chain(inner_stream).boxed() + } + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.inner.list_with_offset(prefix, offset) + } + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.inner.copy_opts(from, to, opts).await + } + } + + #[tokio::test] + async fn test_list_stream_throttle_errors_decrease_rate() { + let mock = Arc::new(ThrottlingListMockStore { + inner: InMemory::new(), + throttle_count: 5, + }); + + // Seed a file so the real items come through after the errors. + mock.put( + &Path::from("prefix/file.txt"), + PutPayload::from_static(b"data"), + ) + .await + .unwrap(); + + let config = AimdThrottleConfig::default().with_list_aimd( + AimdConfig::default() + .with_initial_rate(100.0) + .with_decrease_factor(0.5) + .with_window_duration(std::time::Duration::from_millis(10)), + ); + let throttled = AimdThrottledStore::new(mock as Arc, config).unwrap(); + + let initial_rate = throttled.list.controller.current_rate(); + assert_eq!(initial_rate, 100.0); + + let items: Vec<_> = throttled.list(Some(&Path::from("prefix"))).collect().await; + + // 5 errors + 1 real item + assert_eq!(items.len(), 6); + assert!(items[0].is_err()); + assert!(items[5].is_ok()); + + // Wait for the AIMD window to expire and trigger evaluation. + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + throttled + .list + .controller + .record_outcome(RequestOutcome::Success); + + let new_rate = throttled.list.controller.current_rate(); + assert!( + new_rate < initial_rate, + "List rate should decrease after stream throttle errors: {} < {}", + new_rate, + initial_rate + ); + } + + struct CountingListStartStore { + inner: InMemory, + list_calls: AtomicUsize, + offset_calls: AtomicUsize, + } + + impl Default for CountingListStartStore { + fn default() -> Self { + Self { + inner: InMemory::new(), + list_calls: AtomicUsize::new(0), + offset_calls: AtomicUsize::new(0), + } + } + } + + impl CountingListStartStore { + fn list_calls(&self) -> usize { + self.list_calls.load(Ordering::SeqCst) + } + + fn offset_calls(&self) -> usize { + self.offset_calls.load(Ordering::SeqCst) + } + } + + impl Display for CountingListStartStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CountingListStartStore") + } + } + + impl Debug for CountingListStartStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CountingListStartStore").finish() + } + } + + #[async_trait] + impl ObjectStore for CountingListStartStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.inner.put_opts(location, bytes, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.inner.get_opts(location, options).await + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.list_calls.fetch_add(1, Ordering::SeqCst); + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.offset_calls.fetch_add(1, Ordering::SeqCst); + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.inner.copy_opts(from, to, opts).await + } + } + + fn list_start_throttle_config() -> AimdThrottleConfig { + // Use a low rate (10 tokens/s) so that the token-acquisition sleep is + // 1/10 = 100 ms — well above the 50 ms timeout used in assertions, + // avoiding flakiness from coarse OS timer resolution (e.g. Windows ~16 ms). + AimdThrottleConfig::default() + .with_burst_capacity(0) + .with_list_aimd(AimdConfig::default().with_initial_rate(10.0)) + } + + #[tokio::test(start_paused = true)] + async fn test_list_acquires_token_before_starting_underlying_stream() { + let store = Arc::new(CountingListStartStore::default()); + store + .put( + &Path::from("prefix/file.txt"), + PutPayload::from_static(b"data"), + ) + .await + .unwrap(); + let throttled = AimdThrottledStore::new( + store.clone() as Arc, + list_start_throttle_config(), + ) + .unwrap(); + + let mut stream = throttled.list(Some(&Path::from("prefix"))); + assert_eq!(store.list_calls(), 0); + // With rate=10 tokens/s and burst_capacity=0, the token acquisition + // sleeps for 100 ms. A 50 ms timeout must expire before that. + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), stream.next()) + .await + .is_err() + ); + assert_eq!(store.list_calls(), 0); + + let item = tokio::time::timeout(std::time::Duration::from_millis(300), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(item.location, Path::from("prefix/file.txt")); + assert_eq!(store.list_calls(), 1); + } + + #[tokio::test(start_paused = true)] + async fn test_list_with_offset_acquires_token_before_starting_underlying_stream() { + let store = Arc::new(CountingListStartStore::default()); + store + .put(&Path::from("prefix/b"), PutPayload::from_static(b"data")) + .await + .unwrap(); + let throttled = AimdThrottledStore::new( + store.clone() as Arc, + list_start_throttle_config(), + ) + .unwrap(); + + let mut stream = + throttled.list_with_offset(Some(&Path::from("prefix")), &Path::from("prefix/a")); + assert_eq!(store.offset_calls(), 0); + // With rate=10 tokens/s and burst_capacity=0, the token acquisition + // sleeps for 100 ms. A 50 ms timeout must expire before that. + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), stream.next()) + .await + .is_err() + ); + assert_eq!(store.offset_calls(), 0); + + let item = tokio::time::timeout(std::time::Duration::from_millis(300), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(item.location, Path::from("prefix/b")); + assert_eq!(store.offset_calls(), 1); + } + + #[tokio::test] + async fn test_per_category_independence() { + let store = Arc::new(InMemory::new()); + let config = AimdThrottleConfig::default().with_aimd( + AimdConfig::default() + .with_initial_rate(100.0) + .with_decrease_factor(0.5) + .with_window_duration(std::time::Duration::from_millis(10)), + ); + let throttled = AimdThrottledStore::new(store, config).unwrap(); + + // Push the read controller into a throttled state + throttled + .read + .controller + .record_outcome(RequestOutcome::Throttled); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + throttled + .read + .controller + .record_outcome(RequestOutcome::Success); + + let read_rate = throttled.read.controller.current_rate(); + let write_rate = throttled.write.controller.current_rate(); + let delete_rate = throttled.delete.controller.current_rate(); + let list_rate = throttled.list.controller.current_rate(); + + assert_eq!(read_rate, 50.0, "Read rate should have decreased"); + assert_eq!(write_rate, 100.0, "Write rate should be unaffected"); + assert_eq!(delete_rate, 100.0, "Delete rate should be unaffected"); + assert_eq!(list_rate, 100.0, "List rate should be unaffected"); + } + + #[tokio::test] + async fn test_per_category_config() { + let store = Arc::new(InMemory::new()); + let config = AimdThrottleConfig::default() + .with_read_aimd(AimdConfig::default().with_initial_rate(200.0)) + .with_write_aimd(AimdConfig::default().with_initial_rate(100.0)) + .with_delete_aimd(AimdConfig::default().with_initial_rate(50.0)) + .with_list_aimd(AimdConfig::default().with_initial_rate(25.0)); + let throttled = AimdThrottledStore::new(store, config).unwrap(); + + assert_eq!(throttled.read.controller.current_rate(), 200.0); + assert_eq!(throttled.write.controller.current_rate(), 100.0); + assert_eq!(throttled.delete.controller.current_rate(), 50.0); + assert_eq!(throttled.list.controller.current_rate(), 25.0); + } + + /// A mock [`ObjectStore`] that measures request rate over a sliding window + /// and returns 503 errors when the rate exceeds a configurable threshold. + /// Write and metadata-only operations are not rate-limited. + struct RateLimitingMockStore { + inner: InMemory, + /// Timestamps of recent successful (admitted) requests. + timestamps: std::sync::Mutex>, + /// Maximum requests allowed within `window`. + max_per_window: usize, + /// Sliding window duration. + window: std::time::Duration, + success_count: AtomicU64, + throttle_count: AtomicU64, + } + + impl RateLimitingMockStore { + fn new(max_per_window: usize, window: std::time::Duration) -> Self { + Self { + inner: InMemory::new(), + timestamps: std::sync::Mutex::new(VecDeque::new()), + max_per_window, + window, + success_count: AtomicU64::new(0), + throttle_count: AtomicU64::new(0), + } + } + + /// Returns `true` if the request is admitted, `false` if throttled. + fn check_rate(&self) -> bool { + let mut ts = self.timestamps.lock().unwrap(); + let now = std::time::Instant::now(); + while let Some(&front) = ts.front() { + if now.duration_since(front) > self.window { + ts.pop_front(); + } else { + break; + } + } + if ts.len() >= self.max_per_window { + self.throttle_count.fetch_add(1, Ordering::Relaxed); + false + } else { + ts.push_back(now); + self.success_count.fetch_add(1, Ordering::Relaxed); + true + } + } + + fn throttle_error() -> object_store::Error { + object_store::Error::Generic { + store: "RateLimitingMock", + source: THROTTLE_ERROR_RESPONSE.into(), + } + } + } + + impl Display for RateLimitingMockStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "RateLimitingMockStore") + } + } + + impl Debug for RateLimitingMockStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RateLimitingMockStore").finish() + } + } + + #[async_trait] + impl ObjectStore for RateLimitingMockStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.inner.put_opts(location, bytes, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + if self.check_rate() { + self.inner.get_opts(location, options).await + } else { + Err(Self::throttle_error()) + } + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + if self.check_rate() { + self.inner.get_ranges(location, ranges).await + } else { + Err(Self::throttle_error()) + } + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.inner.copy_opts(from, to, opts).await + } + } + + /// Verify that multiple concurrent readers sharing an AIMD-throttled store + /// converge to the backend's actual capacity. + /// + /// Setup: + /// - Mock backend allows 30 requests per 100ms (= 300 req/s). + /// - 5 reader tasks, each with their own [`AimdThrottledStore`] wrapping + /// the shared mock. + /// - AIMD: 100ms window, initial rate 100 req/s, decrease 0.5, increase 2. + /// - Readers issue `head()` requests as fast as the throttle allows for 2s. + /// + /// Expected behaviour: + /// - Initial burst (100 burst tokens × 5 readers) overshoots the mock + /// capacity, causing many 503s. Each reader's AIMD halves its rate. + /// - After the transient, each reader converges to ~60 req/s (300/5). + /// - Over 2 seconds, total successful requests should be in the range + /// [300, 900] (theoretical max ≈ 600). + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn test_aimd_throttle_under_concurrent_load() { + let mock = Arc::new(RateLimitingMockStore::new( + 30, + std::time::Duration::from_millis(100), + )); + + // Seed a test file so head() succeeds when admitted. + let path = Path::from("test/data.bin"); + mock.put(&path, PutPayload::from_static(b"test data")) + .await + .unwrap(); + + let aimd = AimdConfig::default() + .with_initial_rate(100.0) + .with_decrease_factor(0.5) + .with_additive_increment(2.0) + .with_window_duration(std::time::Duration::from_millis(100)); + let throttle_config = AimdThrottleConfig::default() + .with_aimd(aimd) + .with_burst_capacity(100); + + let num_readers = 5; + let test_duration = std::time::Duration::from_secs(2); + let mut handles = Vec::new(); + + for _ in 0..num_readers { + let store = Arc::new( + AimdThrottledStore::new( + mock.clone() as Arc, + throttle_config.clone(), + ) + .unwrap(), + ); + let p = path.clone(); + handles.push(tokio::spawn(async move { + let deadline = std::time::Instant::now() + test_duration; + let mut count = 0u64; + while std::time::Instant::now() < deadline { + let _ = store.head(&p).await; + count += 1; + } + count + })); + } + + let mut total_reader_requests = 0u64; + for handle in handles { + total_reader_requests += handle.await.unwrap(); + } + + let successes = mock.success_count.load(Ordering::Relaxed); + let throttled = mock.throttle_count.load(Ordering::Relaxed); + let total_mock = successes + throttled; + + // Mock-side count >= reader-side count because the AIMD layer retries + // throttle errors internally, causing multiple mock calls per reader call. + assert!( + total_mock >= total_reader_requests, + "Mock-side count ({total_mock}) should be >= reader-side count ({total_reader_requests})" + ); + + // Mock capacity is 30/100ms = 300 req/s. Over 2s the theoretical max is + // ~600 successful requests. With AIMD ramp-up, expect somewhat fewer. + assert!( + successes >= 300, + "Expected >= 300 successful requests over 2s, got {successes}" + ); + assert!( + successes <= 900, + "Expected <= 900 successful requests, got {successes}" + ); + + // The initial burst exceeds mock capacity, so throttling must occur. + assert!(throttled > 0, "Expected some throttled requests but got 0"); + + // Without AIMD, raw tokio tasks against InMemory would fire 100k+ req/s. + // AIMD should keep the total well under 5000 over 2s. + assert!( + total_mock <= 5000, + "AIMD should limit total requests, got {total_mock}" + ); + } + + /// A mock store that returns a configurable number of throttle errors + /// before succeeding on `get` operations. Used to test the retry logic + /// inside `OperationThrottle::throttled()`. + struct RetryTestMockStore { + inner: InMemory, + /// Number of throttle errors remaining before success. + errors_remaining: std::sync::Mutex, + /// Total number of `get` calls observed. + get_call_count: AtomicU64, + } + + impl RetryTestMockStore { + fn new(errors_before_success: usize) -> Self { + Self { + inner: InMemory::new(), + errors_remaining: std::sync::Mutex::new(errors_before_success), + get_call_count: AtomicU64::new(0), + } + } + } + + impl Display for RetryTestMockStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "RetryTestMockStore") + } + } + + impl Debug for RetryTestMockStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RetryTestMockStore").finish() + } + } + + #[async_trait] + impl ObjectStore for RetryTestMockStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.inner.put_opts(location, bytes, opts).await + } + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + self.inner.put_multipart_opts(location, opts).await + } + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + self.get_call_count.fetch_add(1, Ordering::Relaxed); + let should_error = { + let mut remaining = self.errors_remaining.lock().unwrap(); + if *remaining > 0 { + *remaining -= 1; + true + } else { + false + } + }; + if should_error { + Err(object_store::Error::Generic { + store: "RetryTestMock", + source: THROTTLE_ERROR_RESPONSE.into(), + }) + } else { + self.inner.get_opts(location, options).await + } + } + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.inner.get_ranges(location, ranges).await + } + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.inner.delete_stream(locations) + } + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.inner.list(prefix) + } + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.inner.list_with_offset(prefix, offset) + } + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.inner.list_with_delimiter(prefix).await + } + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.inner.copy_opts(from, to, opts).await + } + } + + #[tokio::test] + async fn test_throttled_retries_on_throttle_error_then_succeeds() { + // Mock returns 2 throttle errors then succeeds (within MAX_RETRIES=3) + let mock = Arc::new(RetryTestMockStore::new(2)); + let path = Path::from("test/retry.txt"); + mock.put(&path, PutPayload::from_static(b"retry data")) + .await + .unwrap(); + + let config = AimdThrottleConfig::default(); + let throttled = + AimdThrottledStore::new(mock.clone() as Arc, config).unwrap(); + + let result = throttled.get(&path).await; + assert!(result.is_ok(), "Expected success after retries"); + + let bytes = result.unwrap().bytes().await.unwrap(); + assert_eq!(bytes.as_ref(), b"retry data"); + + // Should have called get 3 times total: 2 failures + 1 success + assert_eq!(mock.get_call_count.load(Ordering::Relaxed), 3); + } + + #[tokio::test] + async fn test_throttled_fails_after_max_retries_exceeded() { + // Mock returns 4 throttle errors (more than MAX_RETRIES=3), + // so all 4 attempts (initial + 3 retries) will fail. + let mock = Arc::new(RetryTestMockStore::new(10)); + let path = Path::from("test/fail.txt"); + mock.put(&path, PutPayload::from_static(b"fail data")) + .await + .unwrap(); + + let config = AimdThrottleConfig::default(); + let throttled = + AimdThrottledStore::new(mock.clone() as Arc, config).unwrap(); + + let result = throttled.get(&path).await; + assert!(result.is_err(), "Expected error after max retries"); + let err = result.unwrap_err(); + assert!(is_throttle_error(&err)); + + let lance_error = lance_core::Error::from(err); + let error_message = lance_error.to_string(); + assert!(error_message.contains("x-ms-request-id")); + assert!(error_message.contains("azure-request-id")); + + // Should have called get 4 times: initial attempt + 3 retries + assert_eq!(mock.get_call_count.load(Ordering::Relaxed), 4); + } + + #[cfg(feature = "aws")] + #[derive(Debug)] + struct MultipartRetryState { + failures_remaining: AtomicUsize, + part_uris: std::sync::Mutex>, + } + + #[cfg(feature = "aws")] + #[derive(Debug)] + struct MultipartRetryConnector { + state: Arc, + } + + #[cfg(feature = "aws")] + impl HttpConnector for MultipartRetryConnector { + fn connect(&self, _options: &ClientOptions) -> object_store::Result { + Ok(HttpClient::new(MultipartRetryService { + state: Arc::clone(&self.state), + })) + } + } + + #[cfg(feature = "aws")] + #[derive(Debug)] + struct MultipartRetryService { + state: Arc, + } + + #[cfg(feature = "aws")] + #[async_trait] + impl HttpService for MultipartRetryService { + async fn call(&self, request: HttpRequest) -> Result { + let method = request.method().clone(); + let query = request.uri().query().unwrap_or_default(); + let (status, body, e_tag) = if method == ::http::Method::POST + && query + .split('&') + .any(|part| part == "uploads" || part == "uploads=") + { + ( + ::http::StatusCode::OK, + "bucketobjectupload-id", + None, + ) + } else if method == ::http::Method::PUT && query.contains("partNumber=") { + self.state + .part_uris + .lock() + .unwrap() + .push(request.uri().to_string()); + let mut remaining = self.state.failures_remaining.load(Ordering::SeqCst); + let should_fail = loop { + let Some(next) = remaining.checked_sub(1) else { + break false; + }; + match self.state.failures_remaining.compare_exchange_weak( + remaining, + next, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => break true, + Err(actual) => remaining = actual, + } + }; + if should_fail { + ( + ::http::StatusCode::SERVICE_UNAVAILABLE, + "SlowDownPlease reduce your request rate.", + None, + ) + } else { + (::http::StatusCode::OK, "", Some("\"part-etag\"")) + } + } else if method == ::http::Method::POST && query.contains("uploadId=") { + ( + ::http::StatusCode::OK, + "https://bucket/objectbucketobject\"object-etag\"", + None, + ) + } else { + (::http::StatusCode::BAD_REQUEST, "unexpected request", None) + }; + + let mut response = ::http::Response::builder().status(status); + if let Some(e_tag) = e_tag { + response = response.header(::http::header::ETAG, e_tag); + } + Ok(response + .body(HttpResponseBody::from(body.to_string())) + .unwrap()) + } + } + + /// Retries must remain inside the original S3 `put_part` call. Re-entering + /// `MultipartUpload::put_part` would allocate a new part number and leave a + /// gap that makes `complete` fail with "Missing part". + #[cfg(feature = "aws")] + #[tokio::test(start_paused = true)] + async fn test_multipart_http_retry_reuses_part_number() { + use object_store::RetryConfig; + use object_store::aws::AmazonS3Builder; + + let retry_state = Arc::new(MultipartRetryState { + failures_remaining: AtomicUsize::new(3), + part_uris: std::sync::Mutex::new(Vec::new()), + }); + let throttle_state = AimdThrottleState::new(AimdThrottleConfig::default()).unwrap(); + let connector = AimdMultipartUploadConnector::new( + MultipartRetryConnector { + state: Arc::clone(&retry_state), + }, + Some(&throttle_state), + ); + let store = AmazonS3Builder::new() + .with_bucket_name("bucket") + .with_region("us-east-1") + .with_skip_signature(true) + .with_retry(RetryConfig { + max_retries: 0, + ..Default::default() + }) + .with_http_connector(connector) + .build() + .unwrap(); + + let mut upload = store.put_multipart(&Path::from("object")).await.unwrap(); + upload + .put_part(PutPayload::from_static(b"payload")) + .await + .unwrap(); + upload.complete().await.unwrap(); + + let part_uris = retry_state.part_uris.lock().unwrap(); + assert_eq!(part_uris.len(), 4); + assert!(part_uris.iter().all(|uri| uri == &part_uris[0])); + assert!(part_uris[0].contains("partNumber=1")); + } + + #[tokio::test] + async fn test_throttled_multipart_reorders_parts() { + let store = Arc::new(InMemory::new()) as Arc; + let config = AimdThrottleConfig::default(); + let throttled = AimdThrottledStore::new(store.clone(), config).unwrap(); + + let path = Path::from("test/multipart_ordering.bin"); + let mut upload = throttled.put_multipart(&path).await.unwrap(); + + // Create futures for two parts in order: A then B. + let fut_a = upload.put_part(PutPayload::from_static(b"AAAA")); + let fut_b = upload.put_part(PutPayload::from_static(b"BBBB")); + + // Await in REVERSE order. Part ordering should be determined by + // creation order (put_part call order), not by await order. + fut_b.await.unwrap(); + fut_a.await.unwrap(); + + upload.complete().await.unwrap(); + + let result = store.get(&path).await.unwrap(); + let bytes = result.bytes().await.unwrap(); + + assert_eq!( + bytes.as_ref(), + b"AAAABBBB", + "Parts were reordered! Got {:?} instead of AAAABBBB.", + std::str::from_utf8(&bytes).unwrap_or(""), + ); + } +} diff --git a/vendor/lance-io/src/object_store/tracing.rs b/vendor/lance-io/src/object_store/tracing.rs new file mode 100644 index 000000000..8c31e72e5 --- /dev/null +++ b/vendor/lance-io/src/object_store/tracing.rs @@ -0,0 +1,401 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Wrappers around object_store that apply tracing + +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; +use lance_core::utils::tracing::StreamTracingExt; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult, + UploadPart, +}; +use tracing::{Instrument, Span, instrument}; + +#[derive(Debug)] +pub struct TracedMultipartUpload { + write_span: Span, + target: Box, + write_size: usize, +} + +#[async_trait::async_trait] +impl MultipartUpload for TracedMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + let write_span = self.write_span.clone(); + self.write_size += data.content_length(); + let fut = self.target.put_part(data); + Box::pin(fut.instrument(write_span)) + } + + #[instrument(level = "debug", skip_all)] + async fn complete(&mut self) -> OSResult { + let res = self.target.complete().await?; + self.write_span.record("size", self.write_size); + Ok(res) + } + + #[instrument(level = "debug", skip_all)] + async fn abort(&mut self) -> OSResult<()> { + self.target.abort().await + } +} + +#[derive(Debug)] +pub struct TracedObjectStore { + target: Arc, +} + +impl std::fmt::Display for TracedObjectStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!("TracedObjectStore({})", self.target)) + } +} + +#[async_trait::async_trait] +#[deny(clippy::missing_trait_methods)] +impl object_store::ObjectStore for TracedObjectStore { + #[instrument(level = "debug", skip(self, bytes, location, opts), fields(path = location.as_ref(), size = bytes.content_length()))] + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + self.target.put_opts(location, bytes, opts).await + } + + #[instrument(level = "debug", skip(self, location, opts), fields(path = location.as_ref(), size = tracing::field::Empty))] + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + let upload = self.target.put_multipart_opts(location, opts).await?; + Ok(Box::new(TracedMultipartUpload { + target: upload, + write_span: tracing::Span::current(), + write_size: 0, + })) + } + + #[instrument(level = "debug", skip(self, options, location), fields(path = location.as_ref(), size = tracing::field::Empty))] + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + let res = self.target.get_opts(location, options).await?; + + let span = tracing::Span::current(); + span.record("size", res.range.end - res.range.start); + + Ok(res) + } + + #[instrument(level = "debug", skip(self, location), fields(path = location.as_ref(), size = ranges.iter().map(|r| r.end - r.start).sum::()))] + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + self.target.get_ranges(location, ranges).await + } + + #[instrument(level = "debug", skip_all)] + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + self.target + .delete_stream(locations) + .stream_in_current_span() + .boxed() + } + + #[instrument(level = "debug", skip(self, prefix), fields(prefix = prefix.map(|p| p.as_ref())))] + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + self.target.list(prefix).stream_in_current_span().boxed() + } + + #[instrument(level = "debug", skip(self, prefix, offset), fields(prefix = prefix.map(|p| p.as_ref()), offset = offset.as_ref()))] + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.target + .list_with_offset(prefix, offset) + .stream_in_current_span() + .boxed() + } + + #[instrument(level = "debug", skip(self, prefix), fields(prefix = prefix.map(|p| p.as_ref())))] + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + self.target.list_with_delimiter(prefix).await + } + + #[instrument(level = "debug", skip(self, from, to, opts), fields(from = from.as_ref(), to = to.as_ref()))] + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + self.target.copy_opts(from, to, opts).await + } + + #[instrument(level = "debug", skip(self, from, to, opts), fields(from = from.as_ref(), to = to.as_ref()))] + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OSResult<()> { + self.target.rename_opts(from, to, opts).await + } +} + +pub trait ObjectStoreTracingExt { + fn traced(self) -> Arc; +} + +impl ObjectStoreTracingExt for Arc { + fn traced(self) -> Arc { + Arc::new(TracedObjectStore { target: self }) + } +} + +impl ObjectStoreTracingExt for Arc { + fn traced(self) -> Arc { + Arc::new(TracedObjectStore { target: self }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use bytes::Bytes; + use object_store::memory::InMemory; + use object_store::path::Path; + use object_store::{ObjectStoreExt, PutPayload}; + use tracing_mock::{expect, subscriber}; + + fn payload(data: &[u8]) -> PutPayload { + PutPayload::from_bytes(Bytes::copy_from_slice(data)) + } + + fn make_store() -> Arc { + Arc::new(InMemory::new()).traced() + } + + #[tokio::test(flavor = "current_thread")] + async fn test_put_records_path_and_size() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + + let span = expect::span().named("put_opts"); + let (sub, handle) = subscriber::mock() + .new_span( + span.clone().with_fields( + expect::field("path") + .with_value(&"a/b.bin") + .and(expect::field("size").with_value(&data.len())) + .only(), + ), + ) + .enter(span.clone()) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + make_store().put(&path, payload(data)).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_get_records_path_and_size() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + let size = data.len() as u64; // meta.size is u64 + + // Seed without an active mock subscriber. + let store = make_store(); + store.put(&path, payload(data)).await.unwrap(); + + let span = expect::span().named("get_opts"); + let (sub, handle) = subscriber::mock() + .new_span( + // size = Empty at span creation, so only path is visited. + span.clone() + .with_fields(expect::field("path").with_value(&"a/b.bin").only()), + ) + .enter(span.clone()) + .record(span.clone(), expect::field("size").with_value(&size)) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + store.get(&path).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_get_range_records_path_and_size() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + + let store = make_store(); + store.put(&path, payload(data)).await.unwrap(); + + let range = 2u64..7u64; + let size = range.end - range.start; + + let span = expect::span().named("get_opts"); + let (sub, handle) = subscriber::mock() + .new_span( + span.clone() + .with_fields(expect::field("path").with_value(&"a/b.bin").only()), + ) + .enter(span.clone()) + .record(span.clone(), expect::field("size").with_value(&size)) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + store.get_range(&path, range).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_get_ranges_records_path_and_total_size() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + + let store = make_store(); + store.put(&path, payload(data)).await.unwrap(); + + let ranges = [2u64..5u64, 6u64..9u64]; + let size: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + + let span = expect::span().named("get_ranges"); + let (sub, handle) = subscriber::mock() + .new_span( + // `ranges` is also captured automatically as a debug field since + // it is not in the skip list, so we don't use `.only()` here. + span.clone().with_fields( + expect::field("path") + .with_value(&"a/b.bin") + .and(expect::field("size").with_value(&size)), + ), + ) + .enter(span.clone()) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + store.get_ranges(&path, &ranges).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_head_records_path() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + let size = data.len() as u64; + + let store = make_store(); + store.put(&path, payload(data)).await.unwrap(); + + let span = expect::span().named("get_opts"); + let (sub, handle) = subscriber::mock() + .new_span( + span.clone() + .with_fields(expect::field("path").with_value(&"a/b.bin").only()), + ) + .enter(span.clone()) + .record(span.clone(), expect::field("size").with_value(&size)) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + store.head(&path).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_delete_records_path() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + + let store = make_store(); + store.put(&path, payload(data)).await.unwrap(); + + let span = expect::span().named("delete_stream"); + let (sub, handle) = subscriber::mock() + .new_span(span.clone()) + .enter(span.clone()) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + store.delete(&path).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_copy_records_from_and_to() { + let from = Path::from("a/src.bin"); + let to = Path::from("a/dst.bin"); + let data = b"hello world"; + + let store = make_store(); + store.put(&from, payload(data)).await.unwrap(); + + let span = expect::span().named("copy_opts"); + let (sub, handle) = subscriber::mock() + .new_span( + span.clone().with_fields( + expect::field("from") + .with_value(&"a/src.bin") + .and(expect::field("to").with_value(&"a/dst.bin")) + .only(), + ), + ) + .enter(span.clone()) + .exit(span.clone()) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + store.copy(&from, &to).await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } + + #[tokio::test(flavor = "current_thread")] + async fn test_put_multipart_records_path() { + let path = Path::from("a/b.bin"); + let data = b"hello world"; + + let put_mp_span = expect::span().named("put_multipart_opts"); + // Expect only the span creation; any subsequent enter/exit/record + // events are not in the queue so they are silently ignored. + let (sub, handle) = subscriber::mock() + .new_span( + // size = Empty at span creation, so only path is visited. + put_mp_span.with_fields(expect::field("path").with_value(&"a/b.bin").only()), + ) + .run_with_handle(); + + let _guard = tracing::subscriber::set_default(sub); + let store = make_store(); + let mut upload = store.put_multipart(&path).await.unwrap(); + upload.put_part(payload(data)).await.unwrap(); + upload.complete().await.unwrap(); + drop(_guard); + + handle.assert_finished(); + } +} diff --git a/vendor/lance-io/src/object_writer.rs b/vendor/lance-io/src/object_writer.rs new file mode 100644 index 000000000..73167112e --- /dev/null +++ b/vendor/lance-io/src/object_writer.rs @@ -0,0 +1,806 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::io; +use std::pin::Pin; +use std::sync::{Arc, OnceLock}; +use std::task::Poll; + +use crate::object_store::ObjectStore as LanceObjectStore; +use async_trait::async_trait; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; +use object_store::{MultipartUpload, ObjectStoreExt}; +use object_store::{ObjectStore, Result as OSResult, path::Path}; +use tokio::io::{AsyncWrite, AsyncWriteExt}; +use tokio::task::JoinSet; + +use lance_core::{Error, Result}; +use tracing::Instrument; + +use crate::traits::Writer; +use crate::utils::tracking_store::{IOTracker, IoMetricsGuard}; +use tokio::runtime::Handle; + +/// Start at 5MB. +const INITIAL_UPLOAD_STEP: usize = 1024 * 1024 * 5; + +fn max_upload_parallelism() -> usize { + static MAX_UPLOAD_PARALLELISM: OnceLock = OnceLock::new(); + *MAX_UPLOAD_PARALLELISM.get_or_init(|| { + std::env::var("LANCE_UPLOAD_CONCURRENCY") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(10) + }) +} + +/// Maximum body size for a single S3 PUT: strictly less than 5 GiB. +/// AWS rejects single-PUT bodies of exactly 5 GiB (= 5 * 1024^3) with +/// `EntityTooLarge`, so we clamp `LANCE_INITIAL_UPLOAD_SIZE` one byte +/// below that threshold to keep the buffer-fills-to-clamp single-PUT +/// path safe. See lance#6750 for the related txn-file write fix. +const MAX_UPLOAD_PART_SIZE: usize = 1024 * 1024 * 1024 * 5 - 1; + +/// Clamps a requested upload part size to the valid [5MB, 5GB] range. +/// Returns the clamped value and whether clamping was necessary. +fn clamp_initial_upload_size(raw: usize) -> (usize, bool) { + let clamped = raw.clamp(INITIAL_UPLOAD_STEP, MAX_UPLOAD_PART_SIZE); + (clamped, clamped != raw) +} + +fn initial_upload_size() -> usize { + static LANCE_INITIAL_UPLOAD_SIZE: OnceLock = OnceLock::new(); + *LANCE_INITIAL_UPLOAD_SIZE.get_or_init(|| { + let Some(raw) = std::env::var("LANCE_INITIAL_UPLOAD_SIZE") + .ok() + .and_then(|s| s.parse::().ok()) + else { + return INITIAL_UPLOAD_STEP; + }; + let (clamped, was_clamped) = clamp_initial_upload_size(raw); + if was_clamped { + // OnceLock caches the result, so this warning fires at most once per process. + tracing::warn!( + requested = raw, + clamped, + "LANCE_INITIAL_UPLOAD_SIZE must be between 5MB and 5GB; clamping to valid range" + ); + } + clamped + }) +} + +/// Writer to an object in an object store. +/// +/// If the object is small enough, the writer will upload the object in a single +/// PUT request. If the object is larger, the writer will create a multipart +/// upload and upload parts in parallel. +/// +/// This implements the `AsyncWrite` trait. +pub struct ObjectWriter { + state: UploadState, + path: Arc, + cursor: usize, + buffer: Vec, + // TODO: use constant size to support R2 + use_constant_size_upload_parts: bool, +} + +#[derive(Debug, Clone, Default)] +pub struct WriteResult { + pub size: usize, + pub e_tag: Option, +} + +enum UploadState { + /// The writer has been opened but no data has been written yet. Will be in + /// this state until the buffer is full or the writer is shut down. + Started(Arc), + /// The writer is in the process of creating a multipart upload. + CreatingUpload(BoxFuture<'static, OSResult>>), + /// The writer is in the process of uploading parts. + InProgress { + part_idx: u16, + upload: Box, + futures: JoinSet>, + }, + /// The writer is in the process of uploading data in a single PUT request. + /// This happens when shutdown is called before the buffer is full. + PuttingSingle(BoxFuture<'static, OSResult>), + /// The writer is in the process of completing the multipart upload. + Completing(BoxFuture<'static, OSResult>), + /// The writer has been shut down and all data has been written. + Done(WriteResult), +} + +/// Methods for state transitions. +impl UploadState { + fn started_to_putting_single(&mut self, path: Arc, buffer: Vec) { + // To get owned self, we temporarily swap with Done. + let this = std::mem::replace(self, Self::Done(WriteResult::default())); + *self = match this { + Self::Started(store) => { + let fut = async move { + let size = buffer.len(); + let res = store.put(&path, buffer.into()).await?; + Ok(WriteResult { + size, + e_tag: res.e_tag, + }) + }; + Self::PuttingSingle(Box::pin(fut)) + } + _ => unreachable!(), + } + } + + fn in_progress_to_completing(&mut self) { + // To get owned self, we temporarily swap with Done. + let this = std::mem::replace(self, Self::Done(WriteResult::default())); + *self = match this { + Self::InProgress { + mut upload, + futures, + .. + } => { + debug_assert!(futures.is_empty()); + let fut = async move { + let res = upload.complete().await?; + Ok(WriteResult { + size: 0, // This will be set properly later. + e_tag: res.e_tag, + }) + }; + Self::Completing(Box::pin(fut)) + } + _ => unreachable!(), + }; + } +} + +impl ObjectWriter { + pub async fn new(object_store: &LanceObjectStore, path: &Path) -> Result { + Ok(Self { + state: UploadState::Started(object_store.inner.clone()), + cursor: 0, + path: Arc::new(path.clone()), + buffer: Vec::with_capacity(initial_upload_size()), + use_constant_size_upload_parts: object_store.use_constant_size_upload_parts, + }) + } + + /// Returns the contents of `buffer` as a `Bytes` object and resets `buffer`. + /// The new capacity of `buffer` is determined by the current part index. + fn next_part_buffer(buffer: &mut Vec, part_idx: u16, constant_upload_size: bool) -> Bytes { + let new_capacity = if constant_upload_size { + // The store does not support variable part sizes, so use the initial size. + initial_upload_size() + } else { + // Increase the upload size every 100 parts. This gives maximum part size of 2.5TB. + initial_upload_size().max(((part_idx / 100) as usize + 1) * INITIAL_UPLOAD_STEP) + }; + let new_buffer = Vec::with_capacity(new_capacity); + let part = std::mem::replace(buffer, new_buffer); + Bytes::from(part) + } + + fn put_part( + upload: &mut dyn MultipartUpload, + buffer: Bytes, + ) -> BoxFuture<'static, OSResult<()>> { + log::debug!( + "MultipartUpload submitting part with {} bytes", + buffer.len() + ); + upload.put_part(buffer.into()) + } + + fn poll_tasks( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::result::Result<(), io::Error> { + let mut_self = &mut *self; + loop { + match &mut mut_self.state { + UploadState::Started(_) | UploadState::Done(_) => break, + UploadState::CreatingUpload(fut) => match fut.poll_unpin(cx) { + Poll::Ready(Ok(mut upload)) => { + let mut futures = JoinSet::new(); + + let data = Self::next_part_buffer( + &mut mut_self.buffer, + 0, + mut_self.use_constant_size_upload_parts, + ); + futures.spawn(Self::put_part(upload.as_mut(), data)); + + mut_self.state = UploadState::InProgress { + part_idx: 1, // We just used 0 + futures, + upload, + }; + } + Poll::Ready(Err(e)) => return Err(std::io::Error::other(e)), + Poll::Pending => break, + }, + UploadState::InProgress { futures, .. } => { + while let Poll::Ready(Some(res)) = futures.poll_join_next(cx) { + match res { + Ok(Ok(())) => {} + Err(err) => return Err(std::io::Error::other(err)), + Ok(Err(err)) => return Err(err.into()), + } + } + break; + } + UploadState::PuttingSingle(fut) | UploadState::Completing(fut) => { + match fut.poll_unpin(cx) { + Poll::Ready(Ok(mut res)) => { + res.size = mut_self.cursor; + mut_self.state = UploadState::Done(res) + } + Poll::Ready(Err(e)) => return Err(std::io::Error::other(e)), + Poll::Pending => break, + } + } + } + } + Ok(()) + } + + pub async fn abort(&mut self) { + let state = std::mem::replace(&mut self.state, UploadState::Done(WriteResult::default())); + if let UploadState::InProgress { mut upload, .. } = state { + let _ = upload.abort().await; + } + } +} + +impl Drop for ObjectWriter { + fn drop(&mut self) { + // If there is a multipart upload started but not finished, we should abort it. + if matches!(self.state, UploadState::InProgress { .. }) { + // Take ownership of the state. + let state = + std::mem::replace(&mut self.state, UploadState::Done(WriteResult::default())); + if let UploadState::InProgress { mut upload, .. } = state + && let Ok(handle) = Handle::try_current() + { + handle.spawn(async move { + let _ = upload.abort().await; + }); + } + } + } +} + +impl AsyncWrite for ObjectWriter { + fn poll_write( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &[u8], + ) -> std::task::Poll> { + self.as_mut().poll_tasks(cx)?; + + // Fill buffer up to remaining capacity. + let remaining_capacity = self.buffer.capacity() - self.buffer.len(); + let bytes_to_write = std::cmp::min(remaining_capacity, buf.len()); + self.buffer.extend_from_slice(&buf[..bytes_to_write]); + self.cursor += bytes_to_write; + + // Rust needs a little help to borrow self mutably and immutably at the same time + // through a Pin. + let mut_self = &mut *self; + + // Instantiate next request, if available. + if mut_self.buffer.capacity() == mut_self.buffer.len() { + match &mut mut_self.state { + UploadState::Started(store) => { + let path = mut_self.path.clone(); + let store = store.clone(); + let fut = Box::pin(async move { store.put_multipart(path.as_ref()).await }); + self.state = UploadState::CreatingUpload(fut); + } + // TODO: Make max concurrency configurable from storage options. + UploadState::InProgress { + upload, + part_idx, + futures, + .. + } if futures.len() < max_upload_parallelism() => { + let data = Self::next_part_buffer( + &mut mut_self.buffer, + *part_idx, + mut_self.use_constant_size_upload_parts, + ); + futures.spawn( + Self::put_part(upload.as_mut(), data).instrument(tracing::Span::current()), + ); + *part_idx += 1; + } + _ => {} + } + } + + self.poll_tasks(cx)?; + + match bytes_to_write { + 0 => Poll::Pending, + _ => Poll::Ready(Ok(bytes_to_write)), + } + } + + fn poll_flush( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.as_mut().poll_tasks(cx)?; + + match &self.state { + UploadState::Started(_) | UploadState::Done(_) => Poll::Ready(Ok(())), + UploadState::CreatingUpload(_) + | UploadState::Completing(_) + | UploadState::PuttingSingle(_) => Poll::Pending, + UploadState::InProgress { futures, .. } => { + if futures.is_empty() { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + } + } + + fn poll_shutdown( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + loop { + self.as_mut().poll_tasks(cx)?; + + // Rust needs a little help to borrow self mutably and immutably at the same time + // through a Pin. + let mut_self = &mut *self; + match &mut mut_self.state { + UploadState::Done(_) => return Poll::Ready(Ok(())), + UploadState::CreatingUpload(_) + | UploadState::PuttingSingle(_) + | UploadState::Completing(_) => return Poll::Pending, + UploadState::Started(_) => { + // If we didn't start a multipart upload, we can just do a single put. + let part = std::mem::take(&mut mut_self.buffer); + let path = mut_self.path.clone(); + self.state.started_to_putting_single(path, part); + } + UploadState::InProgress { + upload, futures, .. + } => { + // Flush final batch + if !mut_self.buffer.is_empty() && futures.len() < max_upload_parallelism() { + // We can just use `take` since we don't need the buffer anymore. + let data = Bytes::from(std::mem::take(&mut mut_self.buffer)); + futures.spawn( + Self::put_part(upload.as_mut(), data) + .instrument(tracing::Span::current()), + ); + // We need to go back to beginning of loop to poll the + // new feature and get the waker registered on the ctx. + continue; + } + + // We handle the transition from in progress to completing here. + if futures.is_empty() { + self.state.in_progress_to_completing(); + } else { + return Poll::Pending; + } + } + } + } + } +} + +#[async_trait] +impl Writer for ObjectWriter { + async fn tell(&mut self) -> Result { + Ok(self.cursor) + } + + async fn shutdown(&mut self) -> Result { + AsyncWriteExt::shutdown(self).await.map_err(|e| { + Error::io(format!( + "failed to shutdown object writer for {}: {}", + self.path, e + )) + })?; + if let UploadState::Done(result) = &self.state { + Ok(result.clone()) + } else { + unreachable!() + } + } +} + +pub struct LocalWriter { + path: Path, + state: LocalWriteState, +} + +#[derive(Default)] +enum LocalWriteState { + Writing(Box), + Finishing { + size: usize, + future: BoxFuture<'static, Result>, + }, + Done(WriteResult), + #[default] + Poisoned, +} + +struct WritingState { + writer: tokio::io::BufWriter, + cursor: usize, + /// Temp path that auto-deletes on drop. Set to `None` after `persist()`. + temp_path: tempfile::TempPath, + io_tracker: Arc, + /// The whole file is reported as a single `put`, so this covers everything + /// from opening the file to it being durable under its final path. A writer + /// dropped before `persist()` records nothing, like an aborted upload. + metrics: IoMetricsGuard, +} + +impl LocalWriter { + pub fn new( + file: tokio::fs::File, + path: Path, + temp_path: tempfile::TempPath, + io_tracker: Arc, + ) -> Self { + Self { + path, + state: LocalWriteState::Writing(Box::new(WritingState { + writer: tokio::io::BufWriter::new(file), + cursor: 0, + temp_path, + metrics: io_tracker.begin_io("put"), + io_tracker, + })), + } + } + + fn already_closed_err(path: &Path) -> io::Error { + io::Error::other(format!( + "cannot write to LocalWriter for {} after shutdown", + path + )) + } + + fn poisoned_err(path: &Path) -> io::Error { + io::Error::other(format!("LocalWriter for {} is in poisoned state", path)) + } + + async fn persist( + temp_path: tempfile::TempPath, + final_path: Path, + size: usize, + io_tracker: Arc, + metrics: IoMetricsGuard, + ) -> Result { + let local_path = crate::local::to_local_path(&final_path); + let persisted = tokio::task::spawn_blocking(move || -> Result { + temp_path.persist(&local_path).map_err(|e| { + Error::io(format!( + "failed to persist temp file to {}: {}", + local_path, e.error + )) + })?; + + let metadata = std::fs::metadata(&local_path).map_err(|e| { + Error::io(format!("failed to read metadata for {}: {}", local_path, e)) + })?; + Ok(get_etag(&metadata)) + }) + .await + .map_err(|e| Error::io(format!("spawn_blocking failed: {}", e))) + .and_then(|e_tag| e_tag); + + metrics.record(&persisted, size as u64); + let e_tag = persisted?; + + io_tracker.record_write("put", final_path, size as u64); + + Ok(WriteResult { + size, + e_tag: Some(e_tag), + }) + } +} + +impl AsyncWrite for LocalWriter { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &[u8], + ) -> Poll> { + if let LocalWriteState::Writing(state) = &mut self.state { + let poll = Pin::new(&mut state.writer).poll_write(cx, buf); + if let Poll::Ready(Ok(n)) = &poll { + state.cursor += *n; + } + poll + } else { + Poll::Ready(Err(Self::already_closed_err(&self.path))) + } + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + if let LocalWriteState::Writing(state) = &mut self.state { + Pin::new(&mut state.writer).poll_flush(cx) + } else { + Poll::Ready(Err(Self::already_closed_err(&self.path))) + } + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll> { + let mut_self = &mut *self; + loop { + match &mut mut_self.state { + LocalWriteState::Writing(state) => { + if Pin::new(&mut state.writer).poll_shutdown(cx).is_pending() { + return Poll::Pending; + } + + // Write is complete, we can transition to persisting. + let LocalWriteState::Writing(state) = + std::mem::replace(&mut mut_self.state, LocalWriteState::Poisoned) + else { + unreachable!() + }; + let size = state.cursor; + mut_self.state = LocalWriteState::Finishing { + size, + future: Box::pin(Self::persist( + state.temp_path, + mut_self.path.clone(), + size, + state.io_tracker, + state.metrics, + )), + }; + } + LocalWriteState::Finishing { future, .. } => match future.poll_unpin(cx) { + Poll::Ready(Ok(result)) => mut_self.state = LocalWriteState::Done(result), + Poll::Ready(Err(e)) => { + return Poll::Ready(Err(io::Error::other(e))); + } + Poll::Pending => return Poll::Pending, + }, + LocalWriteState::Done(_) => return Poll::Ready(Ok(())), + LocalWriteState::Poisoned => { + return Poll::Ready(Err(Self::poisoned_err(&self.path))); + } + } + } + } +} + +#[async_trait] +impl Writer for LocalWriter { + async fn tell(&mut self) -> Result { + match &mut self.state { + LocalWriteState::Writing(state) => Ok(state.cursor), + LocalWriteState::Finishing { size, .. } => Ok(*size), + LocalWriteState::Done(result) => Ok(result.size), + LocalWriteState::Poisoned => Err(Self::poisoned_err(&self.path).into()), + } + } + + async fn shutdown(&mut self) -> Result { + AsyncWriteExt::shutdown(self).await.map_err(|e| { + Error::io(format!( + "failed to shutdown local writer for {}: {}", + self.path, e + )) + })?; + + match &self.state { + LocalWriteState::Done(result) => Ok(result.clone()), + _ => unreachable!(), + } + } +} + +// Based on object store's implementation. +pub fn get_etag(metadata: &std::fs::Metadata) -> String { + let inode = get_inode(metadata); + let size = metadata.len(); + let mtime = metadata + .modified() + .ok() + .and_then(|mtime| mtime.duration_since(std::time::SystemTime::UNIX_EPOCH).ok()) + .unwrap_or_default() + .as_micros(); + + // Use an ETag scheme based on that used by many popular HTTP servers + // + format!("{inode:x}-{mtime:x}-{size:x}") +} + +#[cfg(unix)] +fn get_inode(metadata: &std::fs::Metadata) -> u64 { + std::os::unix::fs::MetadataExt::ino(metadata) +} + +#[cfg(not(unix))] +fn get_inode(_metadata: &std::fs::Metadata) -> u64 { + 0 +} + +#[cfg(test)] +mod tests { + use tokio::io::AsyncWriteExt; + + use super::*; + + #[tokio::test] + async fn test_write() { + let store = LanceObjectStore::memory(); + + let mut object_writer = ObjectWriter::new(&store, &Path::from("/foo")) + .await + .unwrap(); + assert_eq!(object_writer.tell().await.unwrap(), 0); + + let buf = vec![0; 256]; + assert_eq!(object_writer.write(buf.as_slice()).await.unwrap(), 256); + assert_eq!(object_writer.tell().await.unwrap(), 256); + + assert_eq!(object_writer.write(buf.as_slice()).await.unwrap(), 256); + assert_eq!(object_writer.tell().await.unwrap(), 512); + + assert_eq!(object_writer.write(buf.as_slice()).await.unwrap(), 256); + assert_eq!(object_writer.tell().await.unwrap(), 256 * 3); + + let res = Writer::shutdown(&mut object_writer).await.unwrap(); + assert_eq!(res.size, 256 * 3); + + // Trigger multi part upload + let mut object_writer = ObjectWriter::new(&store, &Path::from("/bar")) + .await + .unwrap(); + let buf = vec![0; INITIAL_UPLOAD_STEP / 3 * 2]; + for i in 0..5 { + // Write more data to trigger the multipart upload + // This should be enough to trigger a multipart upload + object_writer.write_all(buf.as_slice()).await.unwrap(); + // Check the cursor + assert_eq!(object_writer.tell().await.unwrap(), (i + 1) * buf.len()); + } + let res = Writer::shutdown(&mut object_writer).await.unwrap(); + assert_eq!(res.size, buf.len() * 5); + } + + #[tokio::test] + async fn test_abort_write() { + let store = LanceObjectStore::memory(); + + let mut object_writer = ObjectWriter::new(&store, &Path::from("/foo")) + .await + .unwrap(); + object_writer.abort().await; + } + + #[tokio::test] + async fn test_local_writer_shutdown() { + let tmp = lance_core::utils::tempfile::TempStdDir::default(); + let file_path = tmp.join("test_local_writer.bin"); + let os_path = Path::from_absolute_path(&file_path).unwrap(); + let io_tracker = Arc::new(IOTracker::default()); + + let named_temp = tempfile::NamedTempFile::new_in(&*tmp).unwrap(); + let temp_file_path = named_temp.path().to_owned(); + let (std_file, temp_path) = named_temp.into_parts(); + let file = tokio::fs::File::from_std(std_file); + let mut writer = LocalWriter::new(file, os_path, temp_path, io_tracker.clone()); + + let data = b"hello local writer"; + writer.write_all(data).await.unwrap(); + + // Before shutdown, the final path should not exist + assert!(!file_path.exists()); + // But the temp file should exist + assert!(temp_file_path.exists()); + + let result = Writer::shutdown(&mut writer).await.unwrap(); + assert_eq!(result.size, data.len()); + assert!(result.e_tag.is_some()); + assert!(!result.e_tag.as_ref().unwrap().is_empty()); + + // After shutdown, the final path should exist and temp should be gone + assert!(file_path.exists()); + assert!(!temp_file_path.exists()); + + let stats = io_tracker.stats(); + assert_eq!(stats.write_iops, 1); + assert_eq!(stats.written_bytes, data.len() as u64); + } + + #[tokio::test] + async fn test_local_writer_drop_cleans_up() { + let tmp = lance_core::utils::tempfile::TempStdDir::default(); + let file_path = tmp.join("test_drop.bin"); + let os_path = Path::from_absolute_path(&file_path).unwrap(); + let io_tracker = Arc::new(IOTracker::default()); + + let named_temp = tempfile::NamedTempFile::new_in(&*tmp).unwrap(); + let temp_file_path = named_temp.path().to_owned(); + let (std_file, temp_path) = named_temp.into_parts(); + let file = tokio::fs::File::from_std(std_file); + let mut writer = LocalWriter::new(file, os_path, temp_path, io_tracker); + + writer.write_all(b"some data").await.unwrap(); + assert!(temp_file_path.exists()); + + // Drop without shutdown should clean up the temp file + drop(writer); + assert!(!temp_file_path.exists()); + assert!(!file_path.exists()); + } + + #[test] + fn clamp_initial_upload_size_below_min_is_clamped_up() { + assert_eq!(clamp_initial_upload_size(0), (INITIAL_UPLOAD_STEP, true)); + assert_eq!( + clamp_initial_upload_size(INITIAL_UPLOAD_STEP - 1), + (INITIAL_UPLOAD_STEP, true) + ); + } + + #[test] + fn clamp_initial_upload_size_within_range_is_unchanged() { + assert_eq!( + clamp_initial_upload_size(INITIAL_UPLOAD_STEP), + (INITIAL_UPLOAD_STEP, false) + ); + assert_eq!( + clamp_initial_upload_size(MAX_UPLOAD_PART_SIZE), + (MAX_UPLOAD_PART_SIZE, false) + ); + let mid = INITIAL_UPLOAD_STEP * 8; // 40MB, in range + assert_eq!(clamp_initial_upload_size(mid), (mid, false)); + } + + #[test] + fn clamp_initial_upload_size_above_max_is_clamped_down() { + assert_eq!( + clamp_initial_upload_size(MAX_UPLOAD_PART_SIZE + 1), + (MAX_UPLOAD_PART_SIZE, true) + ); + assert_eq!( + clamp_initial_upload_size(usize::MAX), + (MAX_UPLOAD_PART_SIZE, true) + ); + } + + /// Regression for the foot-gun where `LANCE_INITIAL_UPLOAD_SIZE=5368709120` + /// (exactly 5 GiB, Pucheng's setting) caused a single-PUT of 5 GiB on + /// shutdown — which S3 rejects with `EntityTooLarge`. After tightening + /// `MAX_UPLOAD_PART_SIZE` to 5 GiB - 1, raw 5 GiB must clamp DOWN. + #[test] + fn clamp_initial_upload_size_at_5gib_clamps_down() { + let exactly_5_gib: usize = 5 * 1024 * 1024 * 1024; + assert_eq!( + clamp_initial_upload_size(exactly_5_gib), + (MAX_UPLOAD_PART_SIZE, true) + ); + } +} diff --git a/vendor/lance-io/src/scheduler.rs b/vendor/lance-io/src/scheduler.rs new file mode 100644 index 000000000..60f9910e5 --- /dev/null +++ b/vendor/lance-io/src/scheduler.rs @@ -0,0 +1,2585 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use bytes::Bytes; +use futures::channel::oneshot; +use futures::{FutureExt, TryFutureExt}; +use object_store::path::Path; +use std::collections::BinaryHeap; +use std::fmt::Debug; +use std::future::Future; +use std::num::NonZero; +use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; +use tokio::sync::Notify; + +use lance_core::utils::io_stats::IoStatsRecorder; +use lance_core::utils::parse::str_is_truthy; +use lance_core::{Error, Result}; + +use crate::object_store::ObjectStore; +use crate::traits::Reader; +use crate::utils::CachedFileSize; + +mod lite; + +// Don't log backpressure warnings until at least this many seconds have passed +const BACKPRESSURE_MIN: u64 = 5; +// Don't log backpressure warnings more than once / minute +const BACKPRESSURE_DEBOUNCE: u64 = 60; +const SCHEDULER_STATE_EVENT_TARGET: &str = "lance_io::scheduler::state"; + +// Global counter of how many IOPS we have issued +static IOPS_COUNTER: AtomicU64 = AtomicU64::new(0); +// Global counter of how many bytes were read by the scheduler +static BYTES_READ_COUNTER: AtomicU64 = AtomicU64::new(0); + +pub fn iops_counter() -> u64 { + IOPS_COUNTER.load(Ordering::Acquire) +} + +pub fn bytes_read_counter() -> u64 { + BYTES_READ_COUNTER.load(Ordering::Acquire) +} + +// We want to allow requests that have a lower priority than any +// currently in-flight request. This helps avoid potential deadlocks +// related to backpressure. Unfortunately, it is quite expensive to +// keep track of which priorities are in-flight. +// +// TODO: At some point it would be nice if we can optimize this away but +// in_flight should remain relatively small (generally less than 256 items) +// and has not shown itself to be a bottleneck yet. +struct PrioritiesInFlight { + in_flight: Vec, +} + +impl PrioritiesInFlight { + fn new(capacity: u32) -> Self { + Self { + in_flight: Vec::with_capacity(capacity as usize * 2), + } + } + + fn min_in_flight(&self) -> u128 { + self.in_flight.first().copied().unwrap_or(u128::MAX) + } + + fn contains(&self, prio: u128) -> bool { + self.in_flight.binary_search(&prio).is_ok() + } + + fn push(&mut self, prio: u128) { + let pos = match self.in_flight.binary_search(&prio) { + Ok(pos) => pos, + Err(pos) => pos, + }; + self.in_flight.insert(pos, prio); + } + + fn remove(&mut self, prio: u128) { + if let Ok(pos) = self.in_flight.binary_search(&prio) { + self.in_flight.remove(pos); + } + } + + fn len(&self) -> usize { + self.in_flight.len() + } + + fn is_empty(&self) -> bool { + self.in_flight.is_empty() + } +} + +struct IoQueueState { + // The configured number of IOPS that can be issued concurrently. + io_capacity: u32, + // Number of IOPS we can issue concurrently before pausing I/O + iops_avail: u32, + // The configured byte budget for unread I/O. + io_buffer_size: u64, + // Number of bytes we are allowed to buffer in memory before pausing I/O + // + // This can dip below 0 due to I/O prioritization + bytes_avail: i64, + // Pending I/O requests + pending_requests: BinaryHeap, + // Priorities of in-flight requests + priorities_in_flight: PrioritiesInFlight, + // Set when the scheduler is finished to notify the I/O loop to shut down + // once all outstanding requests have been completed. + done_scheduling: bool, + // Time when the scheduler started + start: Instant, + // Last time we warned about backpressure + last_warn: AtomicU64, + // When true, skip all byte-based backpressure checks (set when io_buffer_size == 0) + no_backpressure: bool, +} + +impl IoQueueState { + fn new(io_capacity: u32, io_buffer_size: u64) -> Self { + Self { + io_capacity, + iops_avail: io_capacity, + io_buffer_size, + bytes_avail: io_buffer_size as i64, + pending_requests: BinaryHeap::new(), + priorities_in_flight: PrioritiesInFlight::new(io_capacity), + done_scheduling: false, + start: Instant::now(), + last_warn: AtomicU64::from(0), + no_backpressure: io_buffer_size == 0, + } + } + + fn scheduler_state_event(&self) -> Option { + if !tracing::enabled!(target: SCHEDULER_STATE_EVENT_TARGET, tracing::Level::TRACE) { + return None; + } + + let pending_bytes = self + .pending_requests + .iter() + .map(IoTask::num_bytes) + .sum::(); + let head_task = self.pending_requests.peek(); + let min_in_flight_priority = if self.priorities_in_flight.is_empty() { + None + } else { + Some(self.priorities_in_flight.min_in_flight()) + }; + let head_task_priority_bypass = head_task.map(|task| { + self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight() + }); + let head_task_blocked_by_iops = head_task.map(|_| self.iops_avail == 0); + let head_task_blocked_by_bytes = head_task.map(|task| { + let bypasses_bytes = self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight(); + !bypasses_bytes && task.num_bytes() as i64 > self.bytes_avail + }); + let head_task_can_deliver = head_task.map(|task| self.can_deliver_without_warning(task)); + let head_task_bytes = head_task.map(IoTask::num_bytes); + let (head_task_priority_high, head_task_priority_low) = + split_priority(head_task.map(|task| task.priority)); + let (min_in_flight_priority_high, min_in_flight_priority_low) = + split_priority(min_in_flight_priority); + + Some(SchedulerStateEvent { + queue_kind: "standard", + io_capacity: u64::from(self.io_capacity), + iops_available: u64::from(self.iops_avail), + active_iops: u64::from(self.io_capacity.saturating_sub(self.iops_avail)), + pending_iops: self.pending_requests.len() as u64, + pending_bytes, + bytes_available: self.bytes_avail, + bytes_reserved: self.io_buffer_size as i64 - self.bytes_avail, + io_buffer_size_bytes: self.io_buffer_size, + priorities_in_flight: self.priorities_in_flight.len() as u64, + no_backpressure: self.no_backpressure, + head_task_bytes, + head_task_priority_high, + head_task_priority_low, + min_in_flight_priority_high, + min_in_flight_priority_low, + head_task_can_deliver, + head_task_priority_bypass, + head_task_blocked_by_iops, + head_task_blocked_by_bytes, + }) + } + + fn warn_if_needed(&self) { + let seconds_elapsed = self.start.elapsed().as_secs(); + let last_warn = self.last_warn.load(Ordering::Acquire); + let since_last_warn = seconds_elapsed - last_warn; + if (last_warn == 0 + && seconds_elapsed > BACKPRESSURE_MIN + && seconds_elapsed < BACKPRESSURE_DEBOUNCE) + || since_last_warn > BACKPRESSURE_DEBOUNCE + { + tracing::event!(tracing::Level::DEBUG, "Backpressure throttle exceeded"); + log::debug!( + "Backpressure throttle is full, I/O will pause until buffer is drained. Max I/O bandwidth will not be achieved because CPU is falling behind" + ); + self.last_warn + .store(seconds_elapsed.max(1), Ordering::Release); + } + } + + fn can_deliver(&self, task: &IoTask) -> bool { + let can_deliver = self.can_deliver_without_warning(task); + if !can_deliver + && self.iops_avail > 0 + && !(self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight()) + && task.num_bytes() as i64 > self.bytes_avail + { + self.warn_if_needed(); + } + can_deliver + } + + fn can_deliver_without_warning(&self, task: &IoTask) -> bool { + if self.iops_avail == 0 { + false + } else if self.no_backpressure + || task.bypass_backpressure + || task.priority <= self.priorities_in_flight.min_in_flight() + // Chunks from an admitted logical request must keep moving. A + // higher-priority request may be scheduled later and remain + // unconsumed while the caller awaits this request. + || self.priorities_in_flight.contains(task.priority) + { + true + } else { + task.num_bytes() as i64 <= self.bytes_avail + } + } + + fn next_task(&mut self) -> Option { + let task = self.pending_requests.peek()?; + if self.can_deliver(task) { + let skip_bytes_accounting = self.no_backpressure || task.bypass_backpressure; + self.priorities_in_flight.push(task.priority); + self.iops_avail -= 1; + if !skip_bytes_accounting { + self.bytes_avail -= task.num_bytes() as i64; + if self.bytes_avail < 0 { + // This can happen when we admit special priority requests + log::debug!( + "Backpressure throttle temporarily exceeded by {} bytes due to priority I/O", + -self.bytes_avail + ); + } + } + Some(self.pending_requests.pop().unwrap()) + } else { + None + } + } +} + +// This is modeled after the MPSC queue described here: https://docs.rs/tokio/latest/tokio/sync/struct.Notify.html +// +// However, it only needs to be SPSC since there is only one "scheduler thread" +// and one I/O loop. +struct IoQueue { + // Queue state + state: Mutex, + // Used to signal new I/O requests have arrived that might potentially be runnable + notify: Notify, + stats: IoStats, +} + +impl IoQueue { + fn new(io_capacity: u32, io_buffer_size: u64, stats: IoStats) -> Self { + Self { + state: Mutex::new(IoQueueState::new(io_capacity, io_buffer_size)), + notify: Notify::new(), + stats, + } + } + + fn push(&self, task: IoTask) { + log::trace!( + "Inserting I/O request for {} bytes with priority ({},{}) into I/O queue", + task.num_bytes(), + task.priority >> 64, + task.priority & 0xFFFFFFFFFFFFFFFF + ); + let event = { + let mut state = self.state.lock().unwrap(); + state.pending_requests.push(task); + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); + + self.notify.notify_one(); + } + + async fn pop(&self) -> Option { + loop { + { + let mut state = self.state.lock().unwrap(); + if let Some(task) = state.next_task() { + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); + return Some(task); + } + + if state.done_scheduling { + return None; + } + } + + self.notify.notified().await; + } + } + + fn on_iop_complete(&self) { + let event = { + let mut state = self.state.lock().unwrap(); + state.iops_avail += 1; + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); + + self.notify.notify_one(); + } + + fn on_bytes_consumed(&self, bytes: u64, priority: u128, num_reqs: usize) { + let event = { + let mut state = self.state.lock().unwrap(); + state.bytes_avail += bytes as i64; + for _ in 0..num_reqs { + state.priorities_in_flight.remove(priority); + } + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); + + self.notify.notify_one(); + } + + fn close(&self) { + let (pending_requests, event) = { + let mut state = self.state.lock().unwrap(); + state.done_scheduling = true; + let pending_requests = std::mem::take(&mut state.pending_requests); + let event = state.scheduler_state_event(); + (pending_requests, event) + }; + emit_scheduler_state_event(event, &self.stats); + for request in pending_requests { + request.cancel(); + } + + self.notify.notify_one(); + } +} + +// There is one instance of MutableBatch shared by all the I/O operations +// that make up a single request. When all the I/O operations complete +// then the MutableBatch goes out of scope and the batch request is considered +// complete +struct MutableBatch { + when_done: Option, + data_buffers: Vec, + num_bytes: u64, + priority: u128, + num_reqs: usize, + num_delivered: usize, + err: Option, + // When true, report 0 bytes consumed so the backpressure budget is unaffected + bypass_backpressure: bool, + // Queue the batch's backpressure reservation is refunded to once its response + // is delivered or discarded (see `Response`'s `Drop`). + io_queue: Arc, +} + +impl MutableBatch { + fn new( + when_done: F, + num_data_buffers: u32, + priority: u128, + num_reqs: usize, + bypass_backpressure: bool, + io_queue: Arc, + ) -> Self { + Self { + when_done: Some(when_done), + data_buffers: vec![Bytes::default(); num_data_buffers as usize], + num_bytes: 0, + priority, + num_reqs, + num_delivered: 0, + err: None, + bypass_backpressure, + io_queue, + } + } +} + +// Rather than keep track of when all the I/O requests are finished so that we +// can deliver the batch of data we let Rust do that for us. When all I/O's are +// done then the MutableBatch will go out of scope and we know we have all the +// data. +impl Drop for MutableBatch { + fn drop(&mut self) { + // If we have an error, return that. Otherwise return the data, as long as the I/O requests have been processed. + let result = if let Some(err) = self.err.take() { + Err(err) + } else if self.num_delivered < self.data_buffers.len() { + // This usually happens on tokio runtime shutdown + Err(Error::io(format!( + "I/O request was dropped before completion ({} of {} reads delivered)", + self.num_delivered, + self.data_buffers.len() + ))) + } else { + let mut data = Vec::new(); + std::mem::swap(&mut data, &mut self.data_buffers); + Ok(data) + }; + // We don't really care if no one is around to receive it, just let + // the result go out of scope and get cleaned up + let response = Response { + data: Some(result), + io_queue: self.io_queue.clone(), + // Report 0 bytes for bypass tasks so the backpressure budget is unaffected + num_bytes: if self.bypass_backpressure { + 0 + } else { + self.num_bytes + }, + priority: self.priority, + num_reqs: self.num_reqs, + }; + (self.when_done.take().unwrap())(response); + } +} + +struct DataChunk { + task_idx: usize, + num_bytes: u64, + data: Result, +} + +trait DataSink: Send { + fn deliver_data(&mut self, data: DataChunk); +} + +impl DataSink for MutableBatch { + // Called by worker tasks to add data to the MutableBatch + fn deliver_data(&mut self, data: DataChunk) { + self.num_bytes += data.num_bytes; + self.num_delivered += 1; + match data.data { + Ok(data_bytes) => { + self.data_buffers[data.task_idx] = data_bytes; + } + Err(err) => { + // This keeps the original error, if present + self.err.get_or_insert(err); + } + } + } +} + +struct IoTask { + reader: Arc, + to_read: Range, + when_done: Box) + Send>, + priority: u128, + bypass_backpressure: bool, +} + +fn validate_read_length( + file_path: &Path, + requested_range: &Range, + bytes: Bytes, +) -> Result { + let expected_len = requested_range.end - requested_range.start; + if bytes.len() as u64 != expected_len { + return Err(Error::io(format!( + "I/O request for file {file_path} and range {}..{} returned {} bytes, expected {expected_len} bytes", + requested_range.start, + requested_range.end, + bytes.len() + ))); + } + Ok(bytes) +} + +impl Eq for IoTask {} + +impl PartialEq for IoTask { + fn eq(&self, other: &Self) -> bool { + self.bypass_backpressure == other.bypass_backpressure && self.priority == other.priority + } +} + +impl PartialOrd for IoTask { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for IoTask { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Bypass tasks are always delivered before normal tasks. + // Within the same bypass class, this is a min-heap on priority. + self.bypass_backpressure + .cmp(&other.bypass_backpressure) + .then(other.priority.cmp(&self.priority)) + } +} + +impl IoTask { + fn num_bytes(&self) -> u64 { + self.to_read.end - self.to_read.start + } + fn cancel(self) { + (self.when_done)(Err(Error::internal( + "Scheduler closed before I/O was completed".to_string(), + ))); + } + + async fn run(self) { + let file_path = self.reader.path().as_ref(); + let num_bytes = self.num_bytes(); + let bytes = if self.to_read.start == self.to_read.end { + Ok(Bytes::new()) + } else { + let bytes_fut = self + .reader + .get_range(self.to_read.start as usize..self.to_read.end as usize); + IOPS_COUNTER.fetch_add(1, Ordering::Release); + let num_bytes = self.num_bytes(); + bytes_fut + .inspect(move |_| { + BYTES_READ_COUNTER.fetch_add(num_bytes, Ordering::Release); + }) + .await + .map_err(Error::from) + .and_then(|bytes| validate_read_length(self.reader.path(), &self.to_read, bytes)) + }; + // Emit per-file I/O trace event only when tracing is enabled + tracing::trace!( + file = file_path, + bytes_read = num_bytes, + requests = 1, + range_start = self.to_read.start, + range_end = self.to_read.end, + "File I/O completed" + ); + (self.when_done)(bytes); + } +} + +// Every time a scheduler starts up it launches a task to run the I/O loop. This loop +// repeats endlessly until the scheduler is destroyed. +async fn run_io_loop(tasks: Arc) { + // Pop the first finished task off the queue and submit another until + // we are done + loop { + let next_task = tasks.pop().await; + match next_task { + Some(task) => { + tokio::spawn(task.run()); + } + None => { + // The sender has been dropped, we are done + return; + } + } + } +} + +#[derive(Debug)] +struct StatsCollector { + iops: AtomicU64, + requests: AtomicU64, + bytes_read: AtomicU64, +} + +impl StatsCollector { + fn new() -> Self { + Self { + iops: AtomicU64::new(0), + requests: AtomicU64::new(0), + bytes_read: AtomicU64::new(0), + } + } + + fn iops(&self) -> u64 { + self.iops.load(Ordering::Relaxed) + } + + fn bytes_read(&self) -> u64 { + self.bytes_read.load(Ordering::Relaxed) + } + + fn requests(&self) -> u64 { + self.requests.load(Ordering::Relaxed) + } + + fn record_request(&self, request: &[Range]) { + self.requests.fetch_add(1, Ordering::Relaxed); + self.iops.fetch_add(request.len() as u64, Ordering::Relaxed); + self.bytes_read.fetch_add( + request.iter().map(|r| r.end - r.start).sum::(), + Ordering::Relaxed, + ); + } + + /// Add already-aggregated counts (e.g. a snapshot captured from another + /// scheduler) into these counters. + fn add(&self, iops: u64, requests: u64, bytes_read: u64) { + self.iops.fetch_add(iops, Ordering::Relaxed); + self.requests.fetch_add(requests, Ordering::Relaxed); + self.bytes_read.fetch_add(bytes_read, Ordering::Relaxed); + } +} + +impl IoStatsRecorder for StatsCollector { + fn record_request(&self, request: &[Range]) { + // Inherent methods take precedence in resolution, so this delegates to + // the inherent `record_request` above rather than recursing. + Self::record_request(self, request) + } +} + +#[derive(Debug, Clone, Copy, Default)] +pub struct ScanStats { + pub iops: u64, + pub requests: u64, + pub bytes_read: u64, +} + +impl ScanStats { + fn new(stats: &StatsCollector) -> Self { + Self { + iops: stats.iops(), + requests: stats.requests(), + bytes_read: stats.bytes_read(), + } + } +} + +fn split_priority(priority: Option) -> (Option, Option) { + priority + .map(|priority| ((priority >> 64) as u64, priority as u64)) + .unzip() +} + +#[derive(Debug, Clone, Copy)] +pub(super) struct SchedulerStateEvent { + pub(super) queue_kind: &'static str, + pub(super) io_capacity: u64, + pub(super) iops_available: u64, + pub(super) active_iops: u64, + pub(super) pending_iops: u64, + pub(super) pending_bytes: u64, + pub(super) bytes_available: i64, + pub(super) bytes_reserved: i64, + pub(super) io_buffer_size_bytes: u64, + pub(super) priorities_in_flight: u64, + pub(super) no_backpressure: bool, + pub(super) head_task_bytes: Option, + pub(super) head_task_priority_high: Option, + pub(super) head_task_priority_low: Option, + pub(super) min_in_flight_priority_high: Option, + pub(super) min_in_flight_priority_low: Option, + pub(super) head_task_can_deliver: Option, + pub(super) head_task_priority_bypass: Option, + pub(super) head_task_blocked_by_iops: Option, + pub(super) head_task_blocked_by_bytes: Option, +} + +impl SchedulerStateEvent { + fn trace(self, stats: ScanStats) { + tracing::event!( + target: SCHEDULER_STATE_EVENT_TARGET, + tracing::Level::TRACE, + queue_kind = self.queue_kind, + scheduler_iops = stats.iops, + scheduler_requests = stats.requests, + scheduler_bytes_read = stats.bytes_read, + io_capacity = self.io_capacity, + iops_available = self.iops_available, + active_iops = self.active_iops, + pending_iops = self.pending_iops, + pending_bytes = self.pending_bytes, + bytes_available = self.bytes_available, + bytes_reserved = self.bytes_reserved, + io_buffer_size_bytes = self.io_buffer_size_bytes, + priorities_in_flight = self.priorities_in_flight, + no_backpressure = self.no_backpressure, + head_task_bytes_present = self.head_task_bytes.is_some(), + head_task_bytes = self.head_task_bytes.unwrap_or_default(), + head_task_priority_high_present = self.head_task_priority_high.is_some(), + head_task_priority_high = self.head_task_priority_high.unwrap_or_default(), + head_task_priority_low_present = self.head_task_priority_low.is_some(), + head_task_priority_low = self.head_task_priority_low.unwrap_or_default(), + min_in_flight_priority_high_present = self.min_in_flight_priority_high.is_some(), + min_in_flight_priority_high = self.min_in_flight_priority_high.unwrap_or_default(), + min_in_flight_priority_low_present = self.min_in_flight_priority_low.is_some(), + min_in_flight_priority_low = self.min_in_flight_priority_low.unwrap_or_default(), + head_task_can_deliver_present = self.head_task_can_deliver.is_some(), + head_task_can_deliver = self.head_task_can_deliver.unwrap_or(false), + head_task_priority_bypass_present = self.head_task_priority_bypass.is_some(), + head_task_priority_bypass = self.head_task_priority_bypass.unwrap_or(false), + head_task_blocked_by_iops_present = self.head_task_blocked_by_iops.is_some(), + head_task_blocked_by_iops = self.head_task_blocked_by_iops.unwrap_or(false), + head_task_blocked_by_bytes_present = self.head_task_blocked_by_bytes.is_some(), + head_task_blocked_by_bytes = self.head_task_blocked_by_bytes.unwrap_or(false), + "Scheduler state" + ); + } +} + +pub(super) fn emit_scheduler_state_event(event: Option, stats: &IoStats) { + if let Some(event) = event { + event.trace(stats.snapshot()); + } +} + +/// A shareable, cloneable handle to a set of cumulative I/O counters. +/// +/// All clones share the same underlying counters. This serves two purposes: +/// +/// 1. It backs each [`ScanScheduler`]'s own running totals. +/// 2. It can be attached to an individual [`FileScheduler`] (via +/// [`FileScheduler::with_io_stats`]) as a *secondary* sink, so a caller can +/// measure the exact bytes/IOPS performed through that file handle for a +/// bounded scope (e.g. a single query) without disturbing the scheduler's +/// global totals. Read the result back with [`IoStats::snapshot`]. +#[derive(Debug, Clone)] +pub struct IoStats(Arc); + +impl IoStats { + pub fn new() -> Self { + Self(Arc::new(StatsCollector::new())) + } + + /// Record a single completed request. `request` holds the byte ranges as + /// actually submitted to storage (post coalescing/splitting), so the counts + /// reflect physical I/O. + pub fn record_request(&self, request: &[Range]) { + self.0.record_request(request); + } + + /// Take an immutable snapshot of the current cumulative counters. + pub fn snapshot(&self) -> ScanStats { + ScanStats::new(self.0.as_ref()) + } + + /// Return this handle as a type-erased [`IoStatsRecorder`], suitable for + /// attaching to a file reader (e.g. `FileReader::with_io_stats`). The + /// returned recorder shares the same underlying counters as `self`. + pub fn recorder(&self) -> Arc { + self.0.clone() + } + + /// Add a snapshot of already-aggregated statistics into this sink. Used to + /// fold in I/O measured on a separate scheduler (e.g. the one-time reads + /// performed while opening an index). + pub fn add_scan_stats(&self, stats: &ScanStats) { + self.0.add(stats.iops, stats.requests, stats.bytes_read); + } +} + +impl Default for IoStats { + fn default() -> Self { + Self::new() + } +} + +enum IoQueueType { + Standard(Arc), + Lite(Arc), +} + +/// An I/O scheduler which wraps an ObjectStore and throttles the amount of +/// parallel I/O that can be run. +/// +/// The ScanScheduler will cancel any outstanding I/O requests when it is dropped. +/// For this reason it should be kept alive until all I/O has finished. +/// +/// Note: The 2.X file readers already do this so this is only a concern if you are +/// using the ScanScheduler directly. +pub struct ScanScheduler { + object_store: Arc, + io_queue: IoQueueType, + stats: IoStats, +} + +impl Debug for ScanScheduler { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ScanScheduler") + .field("object_store", &self.object_store) + .finish() + } +} + +struct Response { + // `Option` so the caller can take the data out while the response (and its + // backpressure refund on drop) stays intact. + data: Option>>, + io_queue: Arc, + priority: u128, + num_reqs: usize, + num_bytes: u64, +} + +// Refund the batch's backpressure reservation when the response is dropped, be +// that on delivery or when a cancelled request's undelivered response is +// discarded. This releases the budget even if the caller drops the future early. +impl Drop for Response { + fn drop(&mut self) { + self.io_queue + .on_bytes_consumed(self.num_bytes, self.priority, self.num_reqs); + } +} + +#[derive(Debug, Clone, Copy)] +pub struct SchedulerConfig { + /// the # of bytes that can be buffered but not yet requested. + /// This controls back pressure. If data is not processed quickly enough then this + /// buffer will fill up and the I/O loop will pause until the buffer is drained. + pub io_buffer_size_bytes: u64, + /// Whether to use the lite scheduler. + /// + /// - `Some(true)` forces the lite scheduler (e.g. from env var or programmatic). + /// - `Some(false)` forces the standard scheduler. + /// - `None` defers to the object store's preference (see [`ObjectStore::prefers_lite_scheduler`]). + pub use_lite_scheduler: Option, +} + +impl SchedulerConfig { + pub fn new(io_buffer_size_bytes: u64) -> Self { + Self { + io_buffer_size_bytes, + use_lite_scheduler: std::env::var("LANCE_USE_LITE_SCHEDULER") + .ok() + .map(|v| str_is_truthy(v.trim())), + } + } + + /// Big enough for unit testing + pub fn default_for_testing() -> Self { + Self { + io_buffer_size_bytes: 256 * 1024 * 1024, + use_lite_scheduler: None, + } + } + + /// Configuration that should generally maximize bandwidth (not trying to save RAM + /// at all). We assume a max page size of 32MiB and then allow 32MiB per I/O thread + pub fn max_bandwidth(store: &ObjectStore) -> Self { + Self::new(32 * 1024 * 1024 * store.io_parallelism() as u64) + } + + pub fn with_lite_scheduler(self) -> Self { + Self { + use_lite_scheduler: Some(true), + ..self + } + } +} + +impl ScanScheduler { + /// Create a new scheduler with the given I/O capacity + /// + /// # Arguments + /// + /// * object_store - the store to wrap + /// * config - configuration settings for the scheduler + pub fn new(object_store: Arc, config: SchedulerConfig) -> Arc { + let io_capacity = object_store.io_parallelism(); + let stats = IoStats::new(); + let use_lite = config + .use_lite_scheduler + .unwrap_or_else(|| object_store.prefers_lite_scheduler()); + let io_queue = if use_lite { + let io_queue = Arc::new(lite::IoQueue::new( + io_capacity as u64, + config.io_buffer_size_bytes, + stats.clone(), + )); + IoQueueType::Lite(io_queue) + } else { + let io_queue = Arc::new(IoQueue::new( + io_capacity as u32, + config.io_buffer_size_bytes, + stats.clone(), + )); + let io_queue_clone = io_queue.clone(); + // Best we can do here is fire and forget. If the I/O loop is still running when the scheduler is + // dropped we can't wait for it to finish or we'd block a tokio thread. We could spawn a blocking task + // to wait for it to finish but that doesn't seem helpful. + tokio::task::spawn(async move { run_io_loop(io_queue_clone).await }); + IoQueueType::Standard(io_queue) + }; + Arc::new(Self { + object_store, + io_queue, + stats, + }) + } + + /// Open a file for reading + /// + /// # Arguments + /// + /// * path - the path to the file to open + /// * base_priority - the base priority for I/O requests submitted to this file scheduler + /// this will determine the upper 64 bits of priority (the lower 64 bits + /// come from `submit_request` and `submit_single`) + pub async fn open_file_with_priority( + self: &Arc, + path: &Path, + base_priority: u64, + file_size_bytes: &CachedFileSize, + ) -> Result { + let file_size_bytes = if let Some(size) = file_size_bytes.get() { + u64::from(size) + } else { + let size = self.object_store.size(path).await?; + if let Some(size) = NonZero::new(size) { + file_size_bytes.set(size); + } + size + }; + let reader = self + .object_store + .open_with_size(path, file_size_bytes as usize) + .await?; + let block_size = self.object_store.block_size() as u64; + let max_iop_size = self.object_store.max_iop_size(); + Ok(FileScheduler { + reader: reader.into(), + block_size, + root: self.clone(), + base_priority, + max_iop_size, + bypass_backpressure: false, + extra_stats: None, + }) + } + + /// Open a file with a default priority of 0 + /// + /// See [`Self::open_file_with_priority`] for more information on the priority + pub async fn open_file( + self: &Arc, + path: &Path, + file_size_bytes: &CachedFileSize, + ) -> Result { + self.open_file_with_priority(path, 0, file_size_bytes).await + } + + /// Open a [`FileScheduler`] over an already-open [`Reader`]. + /// + /// Unlike [`Self::open_file`], this skips the path lookup and size probe and + /// schedules I/O against `reader` directly. This is useful when the reader + /// was produced outside the scheduler's object store (e.g. a spill file + /// opened via [`crate::spill::Spill::reader`]), since a bare `Reader` + /// cannot otherwise drive a v2 `FileReader` (which needs a scheduler). + /// + /// Uses a base priority of 0; chain [`FileScheduler::with_priority`] to set + /// a different one. + pub fn open_reader(self: &Arc, reader: Arc) -> FileScheduler { + FileScheduler { + reader, + block_size: self.object_store.block_size() as u64, + root: self.clone(), + base_priority: 0, + max_iop_size: self.object_store.max_iop_size(), + bypass_backpressure: false, + extra_stats: None, + } + } + + fn do_submit_request( + &self, + reader: Arc, + request: Vec>, + tx: oneshot::Sender, + priority: u128, + io_queue: &Arc, + bypass_backpressure: bool, + ) { + let num_iops = request.len() as u32; + + let when_all_io_done = move |bytes_and_permits| { + // We don't care if the receiver has given up so discard the result + let _ = tx.send(bytes_and_permits); + }; + + let dest = Arc::new(Mutex::new(Box::new(MutableBatch::new( + when_all_io_done, + num_iops, + priority, + request.len(), + bypass_backpressure, + io_queue.clone(), + )))); + + for (task_idx, iop) in request.into_iter().enumerate() { + let dest = dest.clone(); + let io_queue_clone = io_queue.clone(); + let num_bytes = iop.end - iop.start; + let task = IoTask { + reader: reader.clone(), + to_read: iop, + priority, + bypass_backpressure, + when_done: Box::new(move |data| { + io_queue_clone.on_iop_complete(); + let mut dest = dest.lock().unwrap(); + let chunk = DataChunk { + data, + task_idx, + num_bytes, + }; + dest.deliver_data(chunk); + }), + }; + io_queue.push(task); + } + } + + fn submit_request_standard( + &self, + reader: Arc, + request: Vec>, + priority: u128, + io_queue: &Arc, + bypass_backpressure: bool, + ) -> impl Future>> + Send + use<> { + let (tx, rx) = oneshot::channel::(); + + self.do_submit_request(reader, request, tx, priority, io_queue, bypass_backpressure); + + rx.map(|wrapped_rsp| { + // A cancel error can't occur: the sender always sends before dropping. + // The reservation is refunded on `Response` drop, so just take the data. + let mut rsp = wrapped_rsp.unwrap(); + rsp.data.take().unwrap() + }) + } + + fn submit_request_lite( + &self, + reader: Arc, + request: Vec>, + priority: u128, + io_queue: &Arc, + bypass_backpressure: bool, + ) -> impl Future>> + Send + use<> { + // It's important that we submit all requests _before_ we await anything + let maybe_tasks = request + .into_iter() + .map(|task| { + let reader = reader.clone(); + let queue = io_queue.clone(); + let requested_range = task.clone(); + let run_fn = Box::new(move || { + let bytes_fut = reader + .get_range(requested_range.start as usize..requested_range.end as usize); + async move { + let bytes = bytes_fut.await.map_err(Error::from)?; + validate_read_length(reader.path(), &requested_range, bytes) + } + .boxed() + }); + queue.submit(task, priority, run_fn, bypass_backpressure) + }) + .collect::>>(); + match maybe_tasks { + Ok(tasks) => async move { + let mut results = Vec::with_capacity(tasks.len()); + for task in tasks { + results.push(task.await?); + } + Ok(results) + } + .boxed(), + Err(e) => async move { Err(e) }.boxed(), + } + } + + pub fn submit_request( + &self, + reader: Arc, + request: Vec>, + priority: u128, + bypass_backpressure: bool, + ) -> impl Future>> + Send + use<> { + match &self.io_queue { + IoQueueType::Standard(io_queue) => { + futures::future::Either::Left(self.submit_request_standard( + reader, + request, + priority, + io_queue, + bypass_backpressure, + )) + } + IoQueueType::Lite(io_queue) => futures::future::Either::Right( + self.submit_request_lite(reader, request, priority, io_queue, bypass_backpressure), + ), + } + } + + pub fn stats(&self) -> ScanStats { + self.stats.snapshot() + } + + #[cfg(test)] + fn uses_lite_scheduler(&self) -> bool { + matches!(self.io_queue, IoQueueType::Lite(_)) + } +} + +impl Drop for ScanScheduler { + fn drop(&mut self) { + // If the user is dropping the ScanScheduler then they _should_ be done with I/O. This can happen + // even when I/O is in progress if, for example, the user is dropping a scan mid-read because they found + // the data they wanted (limit after filter or some other example). + // + // Closing the I/O queue will cancel any requests that have not yet been sent to the I/O loop. However, + // it will not terminate the I/O loop itself. This is to help prevent deadlock and ensure that all I/O + // requests that are submitted will terminate. + // + // In theory, this isn't strictly necessary, as callers should drop any task expecting I/O before they + // drop the scheduler. In practice, this can be difficult to do, and it is better to spend a little bit + // of time letting the I/O loop drain so that we can avoid any potential deadlocks. + match &self.io_queue { + IoQueueType::Standard(io_queue) => io_queue.close(), + IoQueueType::Lite(io_queue) => io_queue.close(), + } + } +} + +/// A throttled file reader +#[derive(Clone, Debug)] +pub struct FileScheduler { + reader: Arc, + root: Arc, + block_size: u64, + base_priority: u64, + max_iop_size: u64, + bypass_backpressure: bool, + /// Optional secondary statistics sink. When set, every request submitted + /// through this handle is also recorded here, in addition to the + /// scheduler's global totals. Used to measure per-scope I/O. + extra_stats: Option>, +} + +fn is_close_together(range1: &Range, range2: &Range, block_size: u64) -> bool { + // Note that range1.end <= range2.start is possible (e.g. when decoding string arrays) + range2.start <= (range1.end + block_size) +} + +fn is_overlapping(range1: &Range, range2: &Range) -> bool { + range1.start < range2.end && range2.start < range1.end +} + +impl FileScheduler { + /// Submit a batch of I/O requests to the reader + /// + /// The requests will be queued in a FIFO manner and, when all requests + /// have been fulfilled, the returned future will be completed. + /// + /// Each request has a given priority. If the I/O loop is full then requests + /// will be buffered and requests with the *lowest* priority will be released + /// from the buffer first. + /// + /// Each request has a backpressure ID which controls which backpressure throttle + /// is applied to the request. Requests made to the same backpressure throttle + /// will be throttled together. + pub fn submit_request( + &self, + request: Vec>, + priority: u64, + ) -> impl Future>> + Send + use<> { + // The final priority is a combination of the row offset and the file number + let priority = ((self.base_priority as u128) << 64) + priority as u128; + + let mut merged_requests = Vec::with_capacity(request.len()); + + if !request.is_empty() { + let mut curr_interval = request[0].clone(); + + for req in request.iter().skip(1) { + if is_close_together(&curr_interval, req, self.block_size) { + curr_interval.end = curr_interval.end.max(req.end); + } else { + merged_requests.push(curr_interval); + curr_interval = req.clone(); + } + } + + merged_requests.push(curr_interval); + } + + let mut updated_requests = Vec::with_capacity(merged_requests.len()); + for req in merged_requests { + if req.is_empty() { + updated_requests.push(req); + } else { + let num_requests = (req.end - req.start).div_ceil(self.max_iop_size); + let bytes_per_request = (req.end - req.start) / num_requests; + for i in 0..num_requests { + let start = req.start + i * bytes_per_request; + let end = if i == num_requests - 1 { + // Last request is a bit bigger due to rounding + req.end + } else { + start + bytes_per_request + }; + updated_requests.push(start..end); + } + } + } + + self.root.stats.record_request(&updated_requests); + if let Some(extra_stats) = &self.extra_stats { + extra_stats.record_request(&updated_requests); + } + + let bytes_vec_fut = self.root.submit_request( + self.reader.clone(), + updated_requests.clone(), + priority, + self.bypass_backpressure, + ); + + let mut updated_index = 0; + let mut final_bytes = Vec::with_capacity(request.len()); + + async move { + let bytes_vec = bytes_vec_fut.await?; + + let mut orig_index = 0; + while (updated_index < updated_requests.len()) && (orig_index < request.len()) { + let updated_range = &updated_requests[updated_index]; + let orig_range = &request[orig_index]; + let byte_offset = updated_range.start as usize; + + if is_overlapping(updated_range, orig_range) { + // We need to undo the coalescing and splitting done earlier + let start = orig_range.start as usize - byte_offset; + if orig_range.end <= updated_range.end { + // The original range is fully contained in the updated range, can do + // zero-copy slice + let end = orig_range.end as usize - byte_offset; + final_bytes.push(bytes_vec[updated_index].slice(start..end)); + } else { + // The original read was split into multiple requests, need to copy + // back into a single buffer + let orig_size = orig_range.end - orig_range.start; + let mut merged_bytes = Vec::with_capacity(orig_size as usize); + merged_bytes.extend_from_slice(&bytes_vec[updated_index].slice(start..)); + let mut copy_offset = merged_bytes.len() as u64; + while copy_offset < orig_size { + updated_index += 1; + let next_range = &updated_requests[updated_index]; + let bytes_to_take = + (orig_size - copy_offset).min(next_range.end - next_range.start); + merged_bytes.extend_from_slice( + &bytes_vec[updated_index].slice(0..bytes_to_take as usize), + ); + copy_offset += bytes_to_take; + } + final_bytes.push(Bytes::from(merged_bytes)); + } + orig_index += 1; + } else { + updated_index += 1; + } + } + + Ok(final_bytes) + } + } + + pub fn with_priority(&self, priority: u64) -> Self { + Self { + reader: self.reader.clone(), + root: self.root.clone(), + block_size: self.block_size, + max_iop_size: self.max_iop_size, + base_priority: priority, + bypass_backpressure: self.bypass_backpressure, + extra_stats: self.extra_stats.clone(), + } + } + + /// Returns a copy of this scheduler that additionally records the I/O it + /// performs into `stats`, on top of the scheduler's global statistics. + /// + /// This is the mechanism for measuring exact per-scope (e.g. per-query) I/O: + /// attach a recorder here (e.g. via [`IoStats::recorder`]), perform the reads + /// through the returned handle, then read the totals back with + /// [`IoStats::snapshot`]. The returned handle is cheap to create (a few + /// `Arc` clones) and reuses the same underlying reader, so it does not + /// re-open the file. + pub fn with_io_stats(&self, stats: Arc) -> Self { + Self { + extra_stats: Some(stats), + ..self.clone() + } + } + + /// Returns a copy of this scheduler that bypasses backpressure for all requests. + /// + /// This should be used for indirect I/O (e.g. fetching items after decoding offsets) where + /// blocking on backpressure could cause a deadlock or excessive latency. + pub fn with_bypass_backpressure(&self) -> Self { + Self { + bypass_backpressure: true, + ..self.clone() + } + } + + /// Submit a single IOP to the reader + /// + /// If you have multiple IOPS to perform then [`Self::submit_request`] is going + /// to be more efficient. + /// + /// See [`Self::submit_request`] for more information on the priority and backpressure. + pub fn submit_single( + &self, + range: Range, + priority: u64, + ) -> impl Future> + Send { + self.submit_request(vec![range], priority) + .map_ok(|vec_bytes| vec_bytes.into_iter().next().unwrap()) + } + + /// Provides access to the underlying reader + /// + /// Do not use this for reading data as it will bypass any I/O scheduling! + /// This is mainly exposed to allow metadata operations (e.g size, block_size,) + /// which either aren't IOPS or we don't throttle + pub fn reader(&self) -> &Arc { + &self.reader + } +} + +#[cfg(test)] +mod tests { + use std::{collections::VecDeque, time::Duration}; + + use futures::poll; + use lance_core::utils::tempfile::TempObjFile; + use rand::RngCore; + use rstest::rstest; + + use object_store::{GetRange, ObjectStore as OSObjectStore, ObjectStoreExt, memory::InMemory}; + use tokio::{runtime::Handle, time::timeout}; + use url::Url; + + use crate::{ + object_store::{DEFAULT_DOWNLOAD_RETRY_COUNT, DEFAULT_MAX_IOP_SIZE}, + testing::MockObjectStore, + }; + + use super::*; + + fn make_task(priority: u128, bypass_backpressure: bool) -> IoTask { + IoTask { + reader: Arc::new(TrackingReader { + get_range_count: Arc::new(AtomicU64::new(0)), + path: Path::parse("test").unwrap(), + }), + to_read: 0..1, + when_done: Box::new(|_| {}), + priority, + bypass_backpressure, + } + } + + #[test] + fn test_scheduler_state_event_fields() { + use tracing_mock::{expect, subscriber}; + + let event = expect::event() + .with_target(SCHEDULER_STATE_EVENT_TARGET) + .at_level(tracing::Level::TRACE) + .with_fields( + expect::field("queue_kind") + .with_value(&"standard") + .and(expect::field("scheduler_iops").with_value(&7u64)) + .and(expect::field("scheduler_requests").with_value(&3u64)) + .and(expect::field("scheduler_bytes_read").with_value(&4096u64)) + .and(expect::field("io_capacity").with_value(&4u64)) + .and(expect::field("pending_iops").with_value(&1u64)) + .and(expect::field("bytes_available").with_value(&128i64)) + .and(expect::field("head_task_bytes_present").with_value(&true)) + .and(expect::field("head_task_bytes").with_value(&1u64)) + .and(expect::field("head_task_can_deliver_present").with_value(&true)) + .and(expect::field("head_task_can_deliver").with_value(&true)), + ); + let (subscriber, handle) = subscriber::mock().event(event).run_with_handle(); + + let stats = IoStats::new(); + stats.add_scan_stats(&ScanStats { + iops: 7, + requests: 3, + bytes_read: 4096, + }); + let mut state = IoQueueState::new(4, 192); + state.iops_avail = 2; + state.bytes_avail = 128; + state.pending_requests.push(make_task(1, false)); + + tracing::subscriber::with_default(subscriber, || { + emit_scheduler_state_event(state.scheduler_state_event(), &stats); + }); + + handle.assert_finished(); + } + + #[test] + fn test_iotask_ordering() { + // Bypass tasks must come out of the heap before non-bypass tasks. + // Within each group, lower priority number (= higher priority) comes first. + let mut heap = BinaryHeap::new(); + heap.push(make_task(10, false)); // non-bypass, low priority + heap.push(make_task(1, false)); // non-bypass, high priority + heap.push(make_task(20, true)); // bypass, low priority + heap.push(make_task(5, true)); // bypass, high priority + + let order: Vec<(u128, bool)> = std::iter::from_fn(|| heap.pop()) + .map(|t| (t.priority, t.bypass_backpressure)) + .collect(); + + assert_eq!(order, vec![(5, true), (20, true), (1, false), (10, false)]); + } + + #[test] + fn test_batch_with_undelivered_slot_is_error() { + let response = Arc::new(Mutex::new(None)); + let response_clone = response.clone(); + let io_queue = Arc::new(IoQueue::new(1, 1024, IoStats::new())); + let batch = MutableBatch::new( + move |rsp| *response_clone.lock().unwrap() = Some(rsp), + 2, // num_data_buffers + 0, // priority + 2, // num_reqs + false, + io_queue, + ); + drop(batch); + + let mut rsp = response.lock().unwrap().take().unwrap(); + let data = rsp.data.take().unwrap(); + assert!( + data.is_err(), + "undelivered slot must yield an error, got {data:?}", + ); + } + + #[tokio::test] + async fn test_full_seq_read() { + let tmp_file = TempObjFile::default(); + + let obj_store = Arc::new(ObjectStore::local()); + + // Write 1MiB of data + const DATA_SIZE: u64 = 1024 * 1024; + let mut some_data = vec![0; DATA_SIZE as usize]; + rand::rng().fill_bytes(&mut some_data); + obj_store.put(&tmp_file, &some_data).await.unwrap(); + + let config = SchedulerConfig::default_for_testing(); + + let scheduler = ScanScheduler::new(obj_store, config); + + let file_scheduler = scheduler + .open_file(&tmp_file, &CachedFileSize::unknown()) + .await + .unwrap(); + + // Read it back 4KiB at a time + const READ_SIZE: u64 = 4 * 1024; + let mut reqs = VecDeque::new(); + let mut offset = 0; + while offset < DATA_SIZE { + reqs.push_back( + #[allow(clippy::single_range_in_vec_init)] + file_scheduler + .submit_request(vec![offset..offset + READ_SIZE], 0) + .await + .unwrap(), + ); + offset += READ_SIZE; + } + + offset = 0; + // Note: we should get parallel I/O even though we are consuming serially + while offset < DATA_SIZE { + let data = reqs.pop_front().unwrap(); + let actual = &data[0]; + let expected = &some_data[offset as usize..(offset + READ_SIZE) as usize]; + assert_eq!(expected, actual); + offset += READ_SIZE; + } + } + + #[tokio::test] + async fn test_open_reader_bridge() { + let tmp_file = TempObjFile::default(); + + let obj_store = Arc::new(ObjectStore::local()); + + const DATA_SIZE: u64 = 64 * 1024; + let mut some_data = vec![0; DATA_SIZE as usize]; + rand::rng().fill_bytes(&mut some_data); + obj_store.put(&tmp_file, &some_data).await.unwrap(); + + let config = SchedulerConfig::default_for_testing(); + let scheduler = ScanScheduler::new(obj_store.clone(), config); + + // Open a bare Reader ourselves, then bridge it into a FileScheduler. + let reader: Arc = obj_store.open(&tmp_file).await.unwrap().into(); + let file_scheduler = scheduler.open_reader(reader); + + let bytes = file_scheduler + .submit_request(vec![0..DATA_SIZE], 0) + .await + .unwrap(); + assert_eq!(bytes[0], some_data); + } + + #[derive(Debug)] + struct ShortReader { + path: Path, + } + + impl lance_core::deepsize::DeepSizeOf for ShortReader { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + impl Reader for ShortReader { + fn path(&self) -> &Path { + &self.path + } + + fn block_size(&self) -> usize { + 4096 + } + + fn io_parallelism(&self) -> usize { + 1 + } + + fn size(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(0) }) + } + + fn get_range( + &self, + _range: Range, + ) -> futures::future::BoxFuture<'static, object_store::Result> { + Box::pin(async { Ok(Bytes::new()) }) + } + + fn get_all(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(Bytes::new()) }) + } + } + + #[rstest] + #[case::standard(false)] + #[case::lite(true)] + #[tokio::test] + async fn test_short_read_returns_io_error(#[case] use_lite_scheduler: bool) { + let config = SchedulerConfig { + use_lite_scheduler: Some(use_lite_scheduler), + ..SchedulerConfig::default_for_testing() + }; + let scheduler = ScanScheduler::new(Arc::new(ObjectStore::memory()), config); + let reader = Arc::new(ShortReader { + path: Path::parse("short-file").unwrap(), + }); + let file_scheduler = scheduler.open_reader(reader); + + let error = file_scheduler + .submit_request(vec![0..8], 0) + .await + .unwrap_err(); + + assert!(matches!(error, Error::IO { .. }), "{error:?}"); + assert!( + error.to_string().contains( + "I/O request for file short-file and range 0..8 returned 0 bytes, expected 8 bytes" + ), + "{error}" + ); + } + + #[tokio::test] + async fn test_split_coalesce() { + let tmp_file = TempObjFile::default(); + + let obj_store = Arc::new(ObjectStore::local()); + + // Write 75MiB of data + const DATA_SIZE: u64 = 75 * 1024 * 1024; + let mut some_data = vec![0; DATA_SIZE as usize]; + rand::rng().fill_bytes(&mut some_data); + obj_store.put(&tmp_file, &some_data).await.unwrap(); + + let config = SchedulerConfig::default_for_testing(); + + let scheduler = ScanScheduler::new(obj_store, config); + + let file_scheduler = scheduler + .open_file(&tmp_file, &CachedFileSize::unknown()) + .await + .unwrap(); + + // These 3 requests should be coalesced into a single I/O because they are within 4KiB + // of each other + let req = + file_scheduler.submit_request(vec![50_000..51_000, 52_000..53_000, 54_000..55_000], 0); + + let bytes = req.await.unwrap(); + + assert_eq!(bytes[0], &some_data[50_000..51_000]); + assert_eq!(bytes[1], &some_data[52_000..53_000]); + assert_eq!(bytes[2], &some_data[54_000..55_000]); + + assert_eq!(1, scheduler.stats().iops); + + // This should be split into 5 requests because it is so large + let req = file_scheduler.submit_request(vec![0..DATA_SIZE], 0); + let bytes = req.await.unwrap(); + assert!(bytes[0] == some_data, "data is not the same"); + + assert_eq!(6, scheduler.stats().iops); + + // None of these requests are bigger than the max IOP size but they will be coalesced into + // one IOP that is bigger and then split back into 2 requests that don't quite align with the original + // ranges. + let chunk_size = *DEFAULT_MAX_IOP_SIZE; + let req = file_scheduler.submit_request( + vec![ + 10..chunk_size, + chunk_size + 10..(chunk_size * 2) - 20, + chunk_size * 2..(chunk_size * 2) + 10, + ], + 0, + ); + + let bytes = req.await.unwrap(); + let chunk_size = chunk_size as usize; + assert!( + bytes[0] == some_data[10..chunk_size], + "data is not the same" + ); + assert!( + bytes[1] == some_data[chunk_size + 10..(chunk_size * 2) - 20], + "data is not the same" + ); + assert!( + bytes[2] == some_data[chunk_size * 2..(chunk_size * 2) + 10], + "data is not the same" + ); + assert_eq!(8, scheduler.stats().iops); + + let reads = (0..44) + .map(|i| i * 1_000_000..(i + 1) * 1_000_000) + .collect::>(); + let req = file_scheduler.submit_request(reads, 0); + let bytes = req.await.unwrap(); + for (i, bytes) in bytes.iter().enumerate() { + assert!( + bytes == &some_data[i * 1_000_000..(i + 1) * 1_000_000], + "data is not the same" + ); + } + assert_eq!(11, scheduler.stats().iops); + } + + #[tokio::test] + async fn test_io_stats_sink() { + let tmp_file = TempObjFile::default(); + let obj_store = Arc::new(ObjectStore::local()); + + const DATA_SIZE: u64 = 1024 * 1024; + let mut some_data = vec![0; DATA_SIZE as usize]; + rand::rng().fill_bytes(&mut some_data); + obj_store.put(&tmp_file, &some_data).await.unwrap(); + + let scheduler = ScanScheduler::new(obj_store, SchedulerConfig::default_for_testing()); + + // Attach a per-scope sink to one file handle. + let sink = IoStats::new(); + let file_scheduler = scheduler + .open_file(&tmp_file, &CachedFileSize::unknown()) + .await + .unwrap() + .with_io_stats(sink.recorder()); + + // Three reads within 4KiB coalesce into a single physical IOP. The sink + // and the scheduler's global totals must agree exactly, because both are + // recorded from the same post-coalescing request. + file_scheduler + .submit_request(vec![50_000..51_000, 52_000..53_000, 54_000..55_000], 0) + .await + .unwrap(); + + let global = scheduler.stats(); + let scoped = sink.snapshot(); + assert_eq!(1, scoped.iops); + assert_eq!(1, scoped.requests); + // Coalesced range 50_000..55_000 => 5000 physical bytes. + assert_eq!(5000, scoped.bytes_read); + assert_eq!(global.iops, scoped.iops); + assert_eq!(global.requests, scoped.requests); + assert_eq!(global.bytes_read, scoped.bytes_read); + + // A sibling handle without the sink: the global totals advance but the + // sink stays put, proving per-scope isolation. + let other = scheduler + .open_file(&tmp_file, &CachedFileSize::unknown()) + .await + .unwrap(); + other.submit_request(vec![0..1000], 0).await.unwrap(); + + let global_after = scheduler.stats(); + let scoped_after = sink.snapshot(); + assert_eq!(global.bytes_read + 1000, global_after.bytes_read); + assert_eq!(scoped.bytes_read, scoped_after.bytes_read); + assert_eq!(scoped.iops, scoped_after.iops); + } + + #[tokio::test] + async fn test_priority() { + let some_path = Path::parse("foo").unwrap(); + let base_store = Arc::new(InMemory::new()); + base_store + .put(&some_path, vec![0; 1000].into()) + .await + .unwrap(); + + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let mut obj_store = MockObjectStore::default(); + let semaphore_copy = semaphore.clone(); + obj_store + .expect_get_opts() + .returning(move |location, options| { + let semaphore = semaphore.clone(); + let base_store = base_store.clone(); + let location = location.clone(); + async move { + semaphore.acquire().await.unwrap().forget(); + base_store.get_opts(&location, options).await + } + .boxed() + }); + let obj_store = Arc::new(ObjectStore::new( + Arc::new(obj_store), + Url::parse("mem://").unwrap(), + Some(500), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + + let config = SchedulerConfig { + io_buffer_size_bytes: 1024 * 1024, + use_lite_scheduler: None, + }; + + let scan_scheduler = ScanScheduler::new(obj_store, config); + + let file_scheduler = scan_scheduler + .open_file(&Path::parse("foo").unwrap(), &CachedFileSize::new(1000)) + .await + .unwrap(); + + // Issue a request, priority doesn't matter, it will be submitted + // immediately (it will go pending) + // Note: the timeout is to prevent a deadlock if the test fails. + let first_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..10, 0), + ) + .boxed(); + + // Issue another low priority request (it will go in queue) + let mut second_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..20, 100), + ) + .boxed(); + + // Issue a high priority request (it will go in queue and should bump + // the other queued request down) + let mut third_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..30, 0), + ) + .boxed(); + + // Finish one file, should be the in-flight first request + semaphore_copy.add_permits(1); + assert!(first_fut.await.unwrap().unwrap().len() == 10); + // Other requests should not be finished + assert!(poll!(&mut second_fut).is_pending()); + assert!(poll!(&mut third_fut).is_pending()); + + // Next should be high priority request + semaphore_copy.add_permits(1); + assert!(third_fut.await.unwrap().unwrap().len() == 30); + assert!(poll!(&mut second_fut).is_pending()); + + // Finally, the low priority request + semaphore_copy.add_permits(1); + assert!(second_fut.await.unwrap().unwrap().len() == 20); + } + + #[tokio::test] + async fn test_standard_scheduler_state_tracks_queue_state() { + let some_path = Path::parse("foo").unwrap(); + let base_store = Arc::new(InMemory::new()); + base_store + .put(&some_path, vec![0; 1000].into()) + .await + .unwrap(); + + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let mut obj_store = MockObjectStore::default(); + let semaphore_copy = semaphore.clone(); + obj_store + .expect_get_opts() + .returning(move |location, options| { + let semaphore = semaphore.clone(); + let base_store = base_store.clone(); + let location = location.clone(); + async move { + semaphore.acquire().await.unwrap().forget(); + base_store.get_opts(&location, options).await + } + .boxed() + }); + let obj_store = Arc::new(ObjectStore::new( + Arc::new(obj_store), + Url::parse("mem://").unwrap(), + Some(500), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + + let scheduler = ScanScheduler::new( + obj_store, + SchedulerConfig { + io_buffer_size_bytes: 1024 * 1024, + use_lite_scheduler: Some(false), + }, + ); + let file_scheduler = scheduler + .open_file(&Path::parse("foo").unwrap(), &CachedFileSize::new(1000)) + .await + .unwrap(); + + let first_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..10, 0), + ) + .boxed(); + let second_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..20, 100), + ) + .boxed(); + let third_fut = timeout( + Duration::from_secs(10), + file_scheduler.submit_single(0..30, 0), + ) + .boxed(); + + let io_queue = match &scheduler.io_queue { + IoQueueType::Standard(io_queue) => io_queue.clone(), + IoQueueType::Lite(_) => unreachable!("test forces the standard scheduler"), + }; + let ( + io_capacity, + iops_available, + pending_bytes, + bytes_reserved, + priorities_in_flight, + head_task_bytes, + head_task_blocked_by_iops, + head_task_blocked_by_bytes, + ) = timeout(Duration::from_secs(5), async { + loop { + let observed = { + let state = io_queue.state.lock().unwrap(); + let active_iops = state.io_capacity.saturating_sub(state.iops_avail); + if active_iops == 1 && state.pending_requests.len() == 2 { + let pending_bytes = state + .pending_requests + .iter() + .map(IoTask::num_bytes) + .sum::(); + let head_task = state.pending_requests.peek().unwrap(); + let bypasses_bytes = state.no_backpressure + || head_task.bypass_backpressure + || head_task.priority <= state.priorities_in_flight.min_in_flight(); + Some(( + state.io_capacity, + state.iops_avail, + pending_bytes, + state.io_buffer_size as i64 - state.bytes_avail, + state.priorities_in_flight.len(), + head_task.num_bytes(), + state.iops_avail == 0, + !bypasses_bytes && head_task.num_bytes() as i64 > state.bytes_avail, + )) + } else { + None + } + }; + if let Some(observed) = observed { + break observed; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + assert_eq!(io_capacity, 1); + assert_eq!(iops_available, 0); + assert_eq!(pending_bytes, 50); + assert_eq!(bytes_reserved, 10); + assert_eq!(priorities_in_flight, 1); + assert_eq!(head_task_bytes, 30); + assert!(head_task_blocked_by_iops); + assert!(!head_task_blocked_by_bytes); + + semaphore_copy.add_permits(3); + assert_eq!(first_fut.await.unwrap().unwrap().len(), 10); + assert_eq!(third_fut.await.unwrap().unwrap().len(), 30); + assert_eq!(second_fut.await.unwrap().unwrap().len(), 20); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_backpressure() { + let some_path = Path::parse("foo").unwrap(); + let base_store = Arc::new(InMemory::new()); + base_store + .put(&some_path, vec![0; 100000].into()) + .await + .unwrap(); + + let bytes_read = Arc::new(AtomicU64::from(0)); + let mut obj_store = MockObjectStore::default(); + let bytes_read_copy = bytes_read.clone(); + // Wraps the obj_store to keep track of how many bytes have been read + obj_store + .expect_get_opts() + .returning(move |location, options| { + let range = options.range.as_ref().unwrap(); + let num_bytes = match range { + GetRange::Bounded(bounded) => bounded.end - bounded.start, + _ => panic!(), + }; + bytes_read_copy.fetch_add(num_bytes, Ordering::Release); + let location = location.clone(); + let base_store = base_store.clone(); + async move { base_store.get_opts(&location, options).await }.boxed() + }); + let obj_store = Arc::new(ObjectStore::new( + Arc::new(obj_store), + Url::parse("mem://").unwrap(), + Some(500), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + + let config = SchedulerConfig { + io_buffer_size_bytes: 10, + use_lite_scheduler: None, + }; + + let scan_scheduler = ScanScheduler::new(obj_store.clone(), config); + + let file_scheduler = scan_scheduler + .open_file(&Path::parse("foo").unwrap(), &CachedFileSize::new(100000)) + .await + .unwrap(); + + let wait_for_idle = || async move { + let handle = Handle::current(); + while handle.metrics().num_alive_tasks() != 1 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }; + let wait_for_bytes_read_and_idle = |target_bytes: u64| { + // We need to move `target` but don't want to move `bytes_read` + let bytes_read = &bytes_read; + async move { + let bytes_read_copy = bytes_read.clone(); + while bytes_read_copy.load(Ordering::Acquire) < target_bytes { + tokio::time::sleep(Duration::from_millis(10)).await; + } + wait_for_idle().await; + } + }; + + // This read will begin immediately + let first_fut = file_scheduler.submit_single(0..5, 0); + // This read should also begin immediately + let second_fut = file_scheduler.submit_single(0..5, 0); + // This read will be throttled + let third_fut = file_scheduler.submit_single(0..3, 0); + // Two tasks (third_fut and unit test) + wait_for_bytes_read_and_idle(10).await; + + assert_eq!(first_fut.await.unwrap().len(), 5); + // One task (unit test) + wait_for_bytes_read_and_idle(13).await; + + // 2 bytes are ready but 5 bytes requested, read will be blocked + let fourth_fut = file_scheduler.submit_single(0..5, 0); + wait_for_bytes_read_and_idle(13).await; + + // Out of order completion is ok, will unblock backpressure + assert_eq!(third_fut.await.unwrap().len(), 3); + wait_for_bytes_read_and_idle(18).await; + + assert_eq!(second_fut.await.unwrap().len(), 5); + // At this point there are 5 bytes available in backpressure queue + // Now we issue multi-read that can be partially fulfilled, it will read some bytes but + // not all of them. (using large range gap to ensure request not coalesced) + // + // I'm actually not sure this behavior is great. It's possible that we should just + // block until we can fulfill the entire request. + let fifth_fut = file_scheduler.submit_request(vec![0..3, 90000..90007], 0); + wait_for_bytes_read_and_idle(21).await; + + // Fifth future should eventually finish due to deadlock prevention + let fifth_bytes = tokio::time::timeout(Duration::from_secs(10), fifth_fut) + .await + .unwrap(); + assert_eq!( + fifth_bytes.unwrap().iter().map(|b| b.len()).sum::(), + 10 + ); + + // And now let's just make sure that we can read the rest of the data + assert_eq!(fourth_fut.await.unwrap().len(), 5); + wait_for_bytes_read_and_idle(28).await; + + // Ensure deadlock prevention timeout can be disabled + let config = SchedulerConfig { + io_buffer_size_bytes: 10, + use_lite_scheduler: None, + }; + + let scan_scheduler = ScanScheduler::new(obj_store, config); + let file_scheduler = scan_scheduler + .open_file(&Path::parse("foo").unwrap(), &CachedFileSize::new(100000)) + .await + .unwrap(); + + let first_fut = file_scheduler.submit_single(0..10, 0); + let second_fut = file_scheduler.submit_single(0..10, 0); + + std::thread::sleep(Duration::from_millis(100)); + assert_eq!(first_fut.await.unwrap().len(), 10); + assert_eq!(second_fut.await.unwrap().len(), 10); + } + + #[derive(Debug)] + struct BlockingReader { + semaphore: Arc, + get_range_count: Arc, + path: Path, + } + + impl lance_core::deepsize::DeepSizeOf for BlockingReader { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + impl Reader for BlockingReader { + fn path(&self) -> &Path { + &self.path + } + + fn block_size(&self) -> usize { + 4096 + } + + fn io_parallelism(&self) -> usize { + 1 + } + + fn size(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(1_000_000) }) + } + + fn get_range( + &self, + range: Range, + ) -> futures::future::BoxFuture<'static, object_store::Result> { + self.get_range_count.fetch_add(1, Ordering::Release); + let semaphore = self.semaphore.clone(); + let num_bytes = range.end - range.start; + Box::pin(async move { + semaphore.acquire().await.unwrap().forget(); + Ok(Bytes::from(vec![0u8; num_bytes])) + }) + } + + fn get_all(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(Bytes::from(vec![0u8; 1_000_000])) }) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_same_priority_chunks_continue_after_higher_priority_request() { + let obj_store = Arc::new(ObjectStore::new( + Arc::new(InMemory::new()), + Url::parse("mem://").unwrap(), + Some(4096), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + let scheduler = ScanScheduler::new( + obj_store, + SchedulerConfig { + io_buffer_size_bytes: 10, + use_lite_scheduler: Some(false), + }, + ); + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let reader: Arc = Arc::new(BlockingReader { + semaphore: semaphore.clone(), + get_range_count: Arc::new(AtomicU64::new(0)), + path: Path::parse("test").unwrap(), + }); + + let low_priority = + scheduler.submit_request(reader.clone(), vec![0..6, 100..106], 10, false); + let high_priority = scheduler.submit_request(reader, vec![200..204], 0, false); + + semaphore.add_permits(3); + let low_priority = timeout(Duration::from_secs(5), low_priority) + .await + .unwrap() + .unwrap(); + assert_eq!( + low_priority.iter().map(|bytes| bytes.len()).sum::(), + 12 + ); + + let high_priority = timeout(Duration::from_secs(5), high_priority) + .await + .unwrap() + .unwrap(); + assert_eq!(high_priority[0].len(), 4); + } + + /// A Reader that tracks how many times get_range has been called. + #[derive(Debug)] + struct TrackingReader { + get_range_count: Arc, + path: Path, + } + + impl lance_core::deepsize::DeepSizeOf for TrackingReader { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + 0 + } + } + + impl Reader for TrackingReader { + fn path(&self) -> &Path { + &self.path + } + + fn block_size(&self) -> usize { + 4096 + } + + fn io_parallelism(&self) -> usize { + 1 + } + + fn size(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(1_000_000) }) + } + + fn get_range( + &self, + range: Range, + ) -> futures::future::BoxFuture<'static, object_store::Result> { + self.get_range_count.fetch_add(1, Ordering::Release); + let num_bytes = range.end - range.start; + Box::pin(async move { Ok(Bytes::from(vec![0u8; num_bytes])) }) + } + + fn get_all(&self) -> futures::future::BoxFuture<'_, object_store::Result> { + Box::pin(async { Ok(Bytes::from(vec![0u8; 1_000_000])) }) + } + } + + #[tokio::test] + async fn test_lite_scheduler_submits_eagerly() { + let obj_store = Arc::new(ObjectStore::memory()); + let config = SchedulerConfig::default_for_testing().with_lite_scheduler(); + let scheduler = ScanScheduler::new(obj_store, config); + + let get_range_count = Arc::new(AtomicU64::new(0)); + let reader: Arc = Arc::new(TrackingReader { + get_range_count: get_range_count.clone(), + path: Path::parse("test").unwrap(), + }); + + // Submit several requests. The lite scheduler should call get_range + // eagerly during submit (before the returned future is polled). + let fut1 = scheduler.submit_request(reader.clone(), vec![0..100], 0, false); + let fut2 = scheduler.submit_request(reader.clone(), vec![100..200], 10, false); + let fut3 = scheduler.submit_request(reader.clone(), vec![200..300], 20, false); + + // get_range must have been called for all 3 requests already. + assert_eq!(get_range_count.load(Ordering::Acquire), 3); + + // The futures should still resolve with the correct data. + assert_eq!(fut1.await.unwrap()[0].len(), 100); + assert_eq!(fut2.await.unwrap()[0].len(), 100); + assert_eq!(fut3.await.unwrap()[0].len(), 100); + } + + #[tokio::test] + async fn test_object_store_selects_scheduler() { + // A memory:// store should use the standard scheduler when config is None + let memory_store = Arc::new(ObjectStore::memory()); + assert!(!memory_store.prefers_lite_scheduler()); + let config = SchedulerConfig { + io_buffer_size_bytes: 256 * 1024 * 1024, + use_lite_scheduler: None, + }; + let scheduler = ScanScheduler::new(memory_store.clone(), config); + assert!(!scheduler.uses_lite_scheduler()); + + // A file+uring:// store should use the lite scheduler when config is None + let uring_store = Arc::new(ObjectStore::new( + Arc::new(InMemory::new()), + Url::parse("file+uring:///tmp").unwrap(), + None, + None, + false, + false, + 8, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + assert!(uring_store.prefers_lite_scheduler()); + let config = SchedulerConfig { + io_buffer_size_bytes: 256 * 1024 * 1024, + use_lite_scheduler: None, + }; + let scheduler = ScanScheduler::new(uring_store.clone(), config); + assert!(scheduler.uses_lite_scheduler()); + + // Explicit Some(false) overrides a file+uring:// store's preference + let config = SchedulerConfig { + io_buffer_size_bytes: 256 * 1024 * 1024, + use_lite_scheduler: Some(false), + }; + let scheduler = ScanScheduler::new(uring_store, config); + assert!(!scheduler.uses_lite_scheduler()); + + // Explicit Some(true) overrides a memory:// store's preference + let config = SchedulerConfig { + io_buffer_size_bytes: 256 * 1024 * 1024, + use_lite_scheduler: Some(true), + }; + let scheduler = ScanScheduler::new(memory_store, config); + assert!(scheduler.uses_lite_scheduler()); + } + + #[test_log::test(tokio::test(flavor = "multi_thread"))] + async fn stress_backpressure() { + // This test ensures that the backpressure mechanism works correctly with + // regards to priority. In other words, as long as all requests are consumed + // in priority order then the backpressure mechanism should not deadlock + let some_path = Path::parse("foo").unwrap(); + let obj_store = Arc::new(ObjectStore::memory()); + obj_store + .put(&some_path, vec![0; 100000].as_slice()) + .await + .unwrap(); + + // Only one request will be allowed in + let config = SchedulerConfig { + io_buffer_size_bytes: 1, + use_lite_scheduler: None, + }; + let scan_scheduler = ScanScheduler::new(obj_store.clone(), config); + let file_scheduler = scan_scheduler + .open_file(&some_path, &CachedFileSize::unknown()) + .await + .unwrap(); + + let mut futs = Vec::with_capacity(10000); + for idx in 0..10000 { + futs.push(file_scheduler.submit_single(idx..idx + 1, idx)); + } + + for fut in futs { + fut.await.unwrap(); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_zero_buffer_size_no_backpressure() { + // With io_buffer_size_bytes=0 (no_backpressure=true), reads at any priority go + // through without blocking, even though a zero budget would normally halt all I/O. + let obj_store = Arc::new(ObjectStore::memory()); + let config = SchedulerConfig { + io_buffer_size_bytes: 0, + use_lite_scheduler: Some(false), + }; + let scheduler = ScanScheduler::new(obj_store, config); + + let get_range_count = Arc::new(AtomicU64::new(0)); + let reader: Arc = Arc::new(TrackingReader { + get_range_count: get_range_count.clone(), + path: Path::parse("test").unwrap(), + }); + + // Submit three reads at increasing priorities without awaiting any first. + // Priority 1 and 2 would deadlock under a real 0-byte budget without no_backpressure. + let fut1 = scheduler.submit_request(reader.clone(), vec![0..1000], 0, false); + let fut2 = scheduler.submit_request(reader.clone(), vec![1000..2000], 1, false); + let fut3 = scheduler.submit_request(reader.clone(), vec![2000..3000], 2, false); + + let bytes1 = timeout(Duration::from_secs(5), fut1) + .await + .unwrap() + .unwrap(); + let bytes2 = timeout(Duration::from_secs(5), fut2) + .await + .unwrap() + .unwrap(); + let bytes3 = timeout(Duration::from_secs(5), fut3) + .await + .unwrap() + .unwrap(); + assert_eq!(bytes1[0].len(), 1000); + assert_eq!(bytes2[0].len(), 1000); + assert_eq!(bytes3[0].len(), 1000); + assert_eq!(get_range_count.load(Ordering::Acquire), 3); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_file_scheduler_bypass_backpressure() { + // A FileScheduler obtained via with_bypass_backpressure() submits reads that bypass + // the byte budget, allowing them to proceed even when the budget is exhausted. + let some_path = Path::parse("foo").unwrap(); + let base_store = Arc::new(InMemory::new()); + base_store + .put(&some_path, vec![0u8; 1000].into()) + .await + .unwrap(); + + let bytes_dispatched = Arc::new(AtomicU64::from(0)); + let mut obj_store = MockObjectStore::default(); + let bytes_dispatched_copy = bytes_dispatched.clone(); + obj_store + .expect_get_opts() + .returning(move |location, options| { + let range = options.range.as_ref().unwrap(); + let num_bytes = match range { + GetRange::Bounded(bounded) => bounded.end - bounded.start, + _ => panic!(), + }; + bytes_dispatched_copy.fetch_add(num_bytes, Ordering::Release); + let location = location.clone(); + let base_store = base_store.clone(); + async move { base_store.get_opts(&location, options).await }.boxed() + }); + let obj_store = Arc::new(ObjectStore::new( + Arc::new(obj_store), + Url::parse("mem://").unwrap(), + Some(500), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + + // Budget = 10 bytes. + let config = SchedulerConfig { + io_buffer_size_bytes: 10, + use_lite_scheduler: Some(false), + }; + let scan_scheduler = ScanScheduler::new(obj_store, config); + let file_scheduler = scan_scheduler + .open_file(&Path::parse("foo").unwrap(), &CachedFileSize::new(1000)) + .await + .unwrap(); + let bypass_scheduler = file_scheduler.with_bypass_backpressure(); + + // Fill the 10-byte budget with a priority-0 read. + let blocker_fut = file_scheduler.submit_single(0..10, 0); + while bytes_dispatched.load(Ordering::Acquire) < 10 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + + // A normal read at priority 2 is blocked: budget = 0, priority 2 > min-in-flight 0. + // A bypass read at priority 1 (higher priority in the queue) bypasses the budget check. + let normal_fut = file_scheduler.submit_single(0..10, 2); + let bypass_fut = bypass_scheduler.submit_single(0..10, 1); + + // Bypass read is dispatched; normal read is still blocked. + while bytes_dispatched.load(Ordering::Acquire) < 20 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!( + bytes_dispatched.load(Ordering::Acquire), + 20, + "normal read should still be blocked while budget is exhausted" + ); + + // Consuming the blocker releases its 10-byte budget → normal read can proceed. + timeout(Duration::from_secs(5), blocker_fut) + .await + .unwrap() + .unwrap(); + timeout(Duration::from_secs(5), bypass_fut) + .await + .unwrap() + .unwrap(); + timeout(Duration::from_secs(5), normal_fut) + .await + .unwrap() + .unwrap(); + assert_eq!(bytes_dispatched.load(Ordering::Acquire), 30); + } + + // Against a 100-byte budget: submit fut1 (50 bytes, priority 0), drop it while + // its read is still blocked in get_range, then submit fut2 (60 bytes, priority 1). + // fut2's priority can't win the priority-bypass, so it needs 60 of the budget -- + // available only if fut1's dropped reservation was refunded. Returns whether fut2 + // completed within 2s (false = the reservation leaked and fut2 deadlocked). + async fn run_caller_drop_scenario(use_lite_scheduler: bool) -> (bool, Duration) { + let obj_store = Arc::new(ObjectStore::new( + Arc::new(InMemory::new()), + Url::parse("mem://").unwrap(), + Some(4096), + None, + false, + false, + 1, + DEFAULT_DOWNLOAD_RETRY_COUNT, + None, + )); + let scheduler = ScanScheduler::new( + obj_store, + SchedulerConfig { + io_buffer_size_bytes: 100, + use_lite_scheduler: Some(use_lite_scheduler), + }, + ); + + let semaphore = Arc::new(tokio::sync::Semaphore::new(0)); + let get_range_count = Arc::new(AtomicU64::new(0)); + let reader: Arc = Arc::new(BlockingReader { + semaphore: semaphore.clone(), + get_range_count: get_range_count.clone(), + path: Path::parse("test").unwrap(), + }); + + // Step 1: reserve 50 of the 100 budget bytes with a read we never consume. + // Spawn it so we can cancel the caller-side future while it is still parked + // waiting for the (blocked) read to finish. + let fut1 = scheduler.submit_request(reader.clone(), vec![0..50], 0, false); + let handle = tokio::spawn(async move { + let _ = fut1.await; + }); + + // Wait until the read is genuinely in flight (blocked on the semaphore). + // This guarantees the 50-byte reservation has been taken before we drop + // the caller, closing the race between the I/O loop and the abort. + while get_range_count.load(Ordering::Acquire) == 0 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + + // Step 2: drop the caller-side future while its `rx` is still pending. + handle.abort(); + let _ = handle.await; + + // Step 3: let the in-flight read finish. The reservation should be refunded + // now that the request is done, whether or not the caller is still around. + semaphore.add_permits(1); + // Give the read time to run to completion so the refund would already have + // happened. + tokio::time::sleep(Duration::from_millis(50)).await; + + // Step 4: submit the follow-up. Add a permit up front so that, if it *is* + // admitted, its own read can complete rather than block on the semaphore. + semaphore.add_permits(1); + let fut2 = scheduler.submit_request(reader, vec![100..160], 1, false); + + let start = std::time::Instant::now(); + let outcome = timeout(Duration::from_secs(2), fut2).await; + let elapsed = start.elapsed(); + match outcome { + Ok(res) => { + assert_eq!(res.unwrap().iter().map(|b| b.len()).sum::(), 60); + (true, elapsed) + } + Err(_) => (false, elapsed), + } + } + + /// Dropping a standard-scheduler request future while its read is in flight must + /// still refund the backpressure reservation, so a later request that needs the + /// budget does not deadlock. + #[tokio::test(flavor = "multi_thread")] + async fn standard_scheduler_refunds_reservation_on_caller_drop() { + let (completed, elapsed) = run_caller_drop_scenario(false).await; + assert!( + completed, + "standard scheduler deadlocked the follow-up request (elapsed {elapsed:?}); \ + the dropped request's reservation was not refunded" + ); + } + + /// Same guarantee for the lite scheduler: dropping a request future mid-read + /// releases its reservation via the `TaskHandle` drop path. + #[tokio::test(flavor = "multi_thread")] + async fn lite_scheduler_refunds_reservation_on_caller_drop() { + let (completed, elapsed) = run_caller_drop_scenario(true).await; + assert!( + completed, + "lite scheduler deadlocked the follow-up request (elapsed {elapsed:?}); \ + the dropped request's reservation was not refunded" + ); + } +} diff --git a/vendor/lance-io/src/scheduler/lite.rs b/vendor/lance-io/src/scheduler/lite.rs new file mode 100644 index 000000000..b7a06e11b --- /dev/null +++ b/vendor/lance-io/src/scheduler/lite.rs @@ -0,0 +1,955 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! A lightweight I/O scheduler primarily intended for use with I/O uring. +//! +//! This scheduler attempts to avoid any kind of task switching whenever possible +//! to minimize context switching overhead. +//! +//! There are a few limitations compared to the standard scheduler: +//! +//! * There is no concurrency limit. The scheduler will allow as many IOPS to run +//! as possible as long as the backpressure throttle is not exceeded. +//! * There is no "babysitting" of IOPS. An I/O task will only be polled when its +//! future is polled. The standard scheduler will `spawn` I/O tasks and so they +//! are always polled by tokio's runtime. This is important for operations like +//! cloud requests where intermittent polling is required to clear out network +//! buffers and keep the TCP connection moving. + +use std::{ + collections::{BinaryHeap, HashMap}, + fmt::Debug, + future::Future, + ops::Range, + pin::Pin, + sync::{ + Arc, Mutex, MutexGuard, + atomic::{AtomicU64, Ordering}, + }, + task::{Context, Poll, Waker}, + time::Instant, +}; + +use bytes::Bytes; +use lance_core::{Error, Result}; + +use super::{ + BACKPRESSURE_DEBOUNCE, BACKPRESSURE_MIN, IoStats, SCHEDULER_STATE_EVENT_TARGET, + SchedulerStateEvent, emit_scheduler_state_event, +}; + +type RunFn = Box Pin> + Send>> + Send>; + +/// The state of an I/O task +/// +/// The state machine is as follows: +/// +/// * `Broken` - The task is in an error state and cannot be run, should never happen +/// * `Initial` - The task has been submitted but does not have a backpressure reservation +/// * `Reserved` - The task has a backpressure reservation +/// * `Running` - The task is running and has a future to poll +/// * `Finished` - The task has finished and has a result +enum TaskState { + Broken, + Initial { + idle_waker: Option, + run_fn: RunFn, + }, + Reserved { + idle_waker: Option, + backpressure_reservation: BackpressureReservation, + run_fn: RunFn, + }, + Running { + backpressure_reservation: BackpressureReservation, + inner: Pin> + Send>>, + }, + Finished { + backpressure_reservation: BackpressureReservation, + data: Result, + }, +} + +impl TaskState { + fn backpressure_reservation(&self) -> Option { + match self { + Self::Reserved { + backpressure_reservation, + .. + } + | Self::Running { + backpressure_reservation, + .. + } + | Self::Finished { + backpressure_reservation, + .. + } => Some(*backpressure_reservation), + Self::Initial { .. } | Self::Broken => None, + } + } +} + +/// A custom error type that might have a backpressure reservation +/// +/// This is used instead of Lance's standard error type so we can ensure +/// we release the reservation before returning the error. +struct BrokenTaskError { + message: String, + backpressure_reservation: Option, +} + +/// The result type corresponding to BrokenTaskError +type TaskResult = std::result::Result<(), BrokenTaskError>; + +impl BrokenTaskError { + // Create a BrokenTaskError from a task state + // + // This will capture any backpressure reservation the task has and put it into the + // error so we make sure to release it when returning the error. + fn new(task_state: TaskState, message: String) -> Self { + match task_state.backpressure_reservation() { + None => Self { + message, + backpressure_reservation: None, + }, + Some(reservation) => Self { + message, + backpressure_reservation: Some(reservation), + }, + } + } +} + +/// An I/O task represents a single read operation +struct IoTask { + /// The unique identifier of the task (only used for debugging) + id: u64, + /// The number of bytes to read + num_bytes: u64, + /// The priority of the task, lower values are higher priority + priority: u128, + /// The current state of the task + state: TaskState, + /// When true, the task bypasses backpressure + bypass_backpressure: bool, +} + +impl IoTask { + fn is_reserved(&self) -> bool { + !matches!(self.state, TaskState::Initial { .. }) + } + + fn cancel(&mut self) -> bool { + let was_running = matches!(self.state, TaskState::Running { .. }); + self.state = TaskState::Finished { + backpressure_reservation: BackpressureReservation { + num_bytes: 0, + priority: 0, + }, + data: Err(Error::io_source(Box::new(Error::io_source( + "I/O Task cancelled".to_string().into(), + )))), + }; + was_running + } + + fn reserve(&mut self, backpressure_reservation: BackpressureReservation) -> TaskResult { + let state = std::mem::replace(&mut self.state, TaskState::Broken); + let TaskState::Initial { idle_waker, run_fn } = state else { + return Err(BrokenTaskError::new( + state, + format!("Task with id {} not in initial state", self.id), + )); + }; + self.state = TaskState::Reserved { + idle_waker, + backpressure_reservation, + run_fn, + }; + Ok(()) + } + + fn start(&mut self) -> TaskResult { + let state = std::mem::replace(&mut self.state, TaskState::Broken); + let TaskState::Reserved { + backpressure_reservation, + idle_waker, + run_fn, + } = state + else { + return Err(BrokenTaskError::new( + state, + format!("Task with id {} not in reserved state", self.id), + )); + }; + let inner = run_fn(); + self.state = TaskState::Running { + backpressure_reservation, + inner, + }; + + // If someone is already waiting for this task let them know it is now running + // so they can poll it + if let Some(idle_waker) = idle_waker { + idle_waker.wake(); + } + Ok(()) + } + + fn poll(&mut self, cx: &mut Context<'_>) -> Poll<()> { + match &mut self.state { + TaskState::Broken => Poll::Ready(()), + TaskState::Initial { idle_waker, .. } | TaskState::Reserved { idle_waker, .. } => { + idle_waker.replace(cx.waker().clone()); + Poll::Pending + } + TaskState::Running { + inner, + backpressure_reservation, + } => match inner.as_mut().poll(cx) { + Poll::Ready(data) => { + self.state = TaskState::Finished { + data, + backpressure_reservation: *backpressure_reservation, + }; + Poll::Ready(()) + } + Poll::Pending => Poll::Pending, + }, + TaskState::Finished { .. } => Poll::Ready(()), + } + } + + fn consume(self) -> Result<(Result, BackpressureReservation)> { + let TaskState::Finished { + data, + backpressure_reservation, + } = self.state + else { + return Err(Error::internal(format!( + "Task with id {} not in finished state", + self.id + ))); + }; + Ok((data, backpressure_reservation)) + } +} + +#[derive(Debug, Clone, Copy)] +struct BackpressureReservation { + num_bytes: u64, + priority: u128, +} + +/// A throttle to control how many bytes can be read before we pause to let compute catch up +trait BackpressureThrottle: Send { + fn try_acquire(&mut self, num_bytes: u64, priority: u128) -> Option; + fn release(&mut self, reservation: BackpressureReservation); + /// Unconditionally acquire a zero-cost reservation, tracking only the priority. + /// Used for bypass tasks that must never be blocked by backpressure. + fn force_acquire(&mut self, priority: u128) -> BackpressureReservation; + fn state(&self) -> BackpressureState; +} + +// We want to allow requests that have a lower priority than any +// currently in-flight request. This helps avoid potential deadlocks +// related to backpressure. Unfortunately, it is quite expensive to +// keep track of which priorities are in-flight. +// +// TODO: At some point it would be nice if we can optimize this away but +// in_flight should remain relatively small (generally less than 256 items) +// and has not shown itself to be a bottleneck yet. +struct PrioritiesInFlight { + in_flight: Vec, +} + +impl PrioritiesInFlight { + fn new(capacity: u64) -> Self { + Self { + in_flight: Vec::with_capacity(capacity as usize * 2), + } + } + + fn min_in_flight(&self) -> u128 { + self.in_flight.first().copied().unwrap_or(u128::MAX) + } + + fn contains(&self, prio: u128) -> bool { + self.in_flight.binary_search(&prio).is_ok() + } + + fn push(&mut self, prio: u128) { + let pos = match self.in_flight.binary_search(&prio) { + Ok(pos) => pos, + Err(pos) => pos, + }; + self.in_flight.insert(pos, prio); + } + + fn remove(&mut self, prio: u128) { + if let Ok(pos) = self.in_flight.binary_search(&prio) { + self.in_flight.remove(pos); + } + } + + fn len(&self) -> usize { + self.in_flight.len() + } +} + +#[derive(Debug, Clone, Copy)] +struct BackpressureState { + max_bytes: u64, + bytes_available: i64, + priorities_in_flight: u64, + no_backpressure: bool, +} + +struct SimpleBackpressureThrottle { + max_bytes: u64, + start: Instant, + last_warn: AtomicU64, + bytes_available: i64, + priorities_in_flight: PrioritiesInFlight, + // When true, skip all byte-based backpressure checks (set when max_bytes == 0) + no_backpressure: bool, +} + +impl SimpleBackpressureThrottle { + fn new(max_bytes: u64, max_concurrency: u64) -> Self { + if max_bytes > i64::MAX as u64 { + // This is unlikely to ever be an issue + panic!("Max bytes must be less than {}", i64::MAX); + } + Self { + max_bytes, + start: Instant::now(), + last_warn: AtomicU64::new(0), + bytes_available: max_bytes as i64, + priorities_in_flight: PrioritiesInFlight::new(max_concurrency), + no_backpressure: max_bytes == 0, + } + } + + fn warn_if_needed(&self) { + let seconds_elapsed = self.start.elapsed().as_secs(); + let last_warn = self.last_warn.load(Ordering::Acquire); + let since_last_warn = seconds_elapsed - last_warn; + if (last_warn == 0 + && seconds_elapsed > BACKPRESSURE_MIN + && seconds_elapsed < BACKPRESSURE_DEBOUNCE) + || since_last_warn > BACKPRESSURE_DEBOUNCE + { + tracing::event!(tracing::Level::DEBUG, "Backpressure throttle exceeded"); + log::debug!( + "Backpressure throttle is full, I/O will pause until buffer is drained. Max I/O bandwidth will not be achieved because CPU is falling behind" + ); + self.last_warn + .store(seconds_elapsed.max(1), Ordering::Release); + } + } +} + +impl BackpressureThrottle for SimpleBackpressureThrottle { + fn try_acquire(&mut self, num_bytes: u64, priority: u128) -> Option { + if self.no_backpressure + || self.bytes_available >= num_bytes as i64 + || self.priorities_in_flight.min_in_flight() >= priority + // Chunks from an admitted logical request must keep moving. A + // higher-priority request may be scheduled later and remain + // unconsumed while the caller awaits this request. + || self.priorities_in_flight.contains(priority) + { + self.bytes_available -= num_bytes as i64; + self.priorities_in_flight.push(priority); + Some(BackpressureReservation { + num_bytes, + priority, + }) + } else { + self.warn_if_needed(); + None + } + } + + fn release(&mut self, reservation: BackpressureReservation) { + self.bytes_available += reservation.num_bytes as i64; + self.priorities_in_flight.remove(reservation.priority); + } + + fn force_acquire(&mut self, priority: u128) -> BackpressureReservation { + self.priorities_in_flight.push(priority); + BackpressureReservation { + num_bytes: 0, + priority, + } + } + + fn state(&self) -> BackpressureState { + BackpressureState { + max_bytes: self.max_bytes, + bytes_available: self.bytes_available, + priorities_in_flight: self.priorities_in_flight.len() as u64, + no_backpressure: self.no_backpressure, + } + } +} + +struct TaskEntry { + task_id: u64, + priority: u128, + reserved: bool, +} + +impl Ord for TaskEntry { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + // Prefer reserved tasks over unreserved tasks and then highest priority tasks over lowest + // priority tasks. + // + // This is a max-heap so we sort by reserved in normal order (true > false) and priority + // in reverse order (lowest priority first) + self.reserved + .cmp(&other.reserved) + .then(other.priority.cmp(&self.priority)) + } +} + +impl PartialOrd for TaskEntry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl PartialEq for TaskEntry { + fn eq(&self, other: &Self) -> bool { + self.priority == other.priority + } +} + +impl Eq for TaskEntry {} + +struct IoQueueState { + backpressure_throttle: Box, + pending_tasks: BinaryHeap, + tasks: HashMap, + next_task_id: u64, +} + +impl IoQueueState { + fn new(max_concurrency: u64, max_bytes: u64) -> Self { + Self { + backpressure_throttle: Box::new(SimpleBackpressureThrottle::new( + max_bytes, + max_concurrency, + )), + pending_tasks: BinaryHeap::new(), + tasks: HashMap::new(), + next_task_id: 0, + } + } + + // If a task is in an unexpected state then we need to release any reservations that were made + // before we return an error. + // + // Note: this is perhaps a bit paranoid as a task should never be in an unexpected state. + fn handle_result(&mut self, result: TaskResult) -> Result<()> { + if let Err(error) = result { + if let Some(reservation) = error.backpressure_reservation { + self.backpressure_throttle.release(reservation); + } + Err(Error::internal(error.message)) + } else { + Ok(()) + } + } + + fn scheduler_state_event(&self) -> Option { + if !tracing::enabled!(target: SCHEDULER_STATE_EVENT_TARGET, tracing::Level::TRACE) { + return None; + } + + let backpressure = self.backpressure_throttle.state(); + let pending_bytes = self + .pending_tasks + .iter() + .filter_map(|entry| self.tasks.get(&entry.task_id)) + .map(|task| task.num_bytes) + .sum::(); + let active_iops = self + .tasks + .values() + .filter(|task| matches!(task.state, TaskState::Running { .. })) + .count() as u64; + + Some(SchedulerStateEvent { + queue_kind: "lite", + io_capacity: 0, + iops_available: 0, + active_iops, + pending_iops: self.pending_tasks.len() as u64, + pending_bytes, + bytes_available: backpressure.bytes_available, + bytes_reserved: backpressure.max_bytes as i64 - backpressure.bytes_available, + io_buffer_size_bytes: backpressure.max_bytes, + priorities_in_flight: backpressure.priorities_in_flight, + no_backpressure: backpressure.no_backpressure, + head_task_bytes: None, + head_task_priority_high: None, + head_task_priority_low: None, + min_in_flight_priority_high: None, + min_in_flight_priority_low: None, + head_task_can_deliver: None, + head_task_priority_bypass: None, + head_task_blocked_by_iops: None, + head_task_blocked_by_bytes: None, + }) + } +} + +/// A queue of I/O tasks to be shared between the I/O scheduler and the I/O decoder. +/// +/// The queue is protected by two different throttles. The first controls memory backpressure, and +/// will only allow a certain number of bytes to be allocated for reads. This throttle is released +/// as soon as the decoder consumes the bytes (not when the bytes have been fully processed). This +/// throttle is currently scoped to the scheduler and not shared across the process. This will likely +/// change in the future. +/// +/// The second throttle controls how many IOPS can be issued concurrently. This throttle is released +/// as soon as the IOP is finished. This throttle has both a local per-scheduler limit and also a +/// process-wide limit. +/// +/// Note: unlike the standard scheduler, there is no dedicated I/O loop thread. If the decoder is not +/// polling the I/O tasks then nothing else will. This scheduler is currently intended for use with I/O +/// uring where I/O tasks are bunched together and polling one task advances all outstanding I/O. It +/// would not be suitable for cloud storage where each task is an independent HTTP request and needs to +/// be polled individually (though presumably one could use I/O uring for networked cloud storage some +/// day as well) +pub(super) struct IoQueue { + state: Arc>, + stats: IoStats, +} + +impl IoQueue { + pub fn new(max_concurrency: u64, max_bytes: u64, stats: IoStats) -> Self { + Self { + state: Arc::new(Mutex::new(IoQueueState::new(max_concurrency, max_bytes))), + stats, + } + } + + fn push(&self, mut task: IoTask, mut state: MutexGuard) -> Result<()> { + let task_id = task.id; + let maybe_reservation = if task.bypass_backpressure { + Some(state.backpressure_throttle.force_acquire(task.priority)) + } else { + state + .backpressure_throttle + .try_acquire(task.num_bytes, task.priority) + }; + if let Some(reservation) = maybe_reservation { + state.handle_result(task.reserve(reservation))?; + state.handle_result(task.start())?; + state.tasks.insert(task_id, task); + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); + return Ok(()); + } + + state.pending_tasks.push(TaskEntry { + task_id, + priority: task.priority, + reserved: task.is_reserved(), + }); + state.tasks.insert(task_id, task); + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); + Ok(()) + } + + pub(super) fn submit( + self: Arc, + range: Range, + priority: u128, + run_fn: RunFn, + bypass_backpressure: bool, + ) -> Result { + log::trace!( + "Submitting I/O task with range {:?}, priority {:?}", + range, + priority + ); + let mut state = self.state.lock().unwrap(); + let task_id = state.next_task_id; + state.next_task_id += 1; + + let task = IoTask { + id: task_id, + num_bytes: range.end - range.start, + priority, + bypass_backpressure, + state: TaskState::Initial { + idle_waker: None, + run_fn, + }, + }; + self.push(task, state)?; + Ok(TaskHandle { + task_id, + queue: self, + }) + } + + // When a task completes we should check to see if any other tasks are now runnable + fn on_task_complete(&self, mut state: MutexGuard) -> Result<()> { + let result = { + let state_ref = &mut *state; + let mut task_result = TaskResult::Ok(()); + while !state_ref.pending_tasks.is_empty() { + // Unwrap safe here since we just checked the queue is not empty + let task_id = state_ref.pending_tasks.peek().unwrap().task_id; + let Some(task) = state_ref.tasks.get_mut(&task_id) else { + // The caller dropped this task's handle (see `abandon`); discard the + // stale queue entry instead of spinning on it. + state_ref.pending_tasks.pop(); + continue; + }; + if !task.is_reserved() { + let Some(reservation) = state_ref + .backpressure_throttle + .try_acquire(task.num_bytes, task.priority) + else { + break; + }; + if let Err(e) = task.reserve(reservation) { + task_result = Err(e); + break; + } + } + state_ref.pending_tasks.pop(); + if let Err(e) = task.start() { + task_result = Err(e); + break; + } + } + state_ref.handle_result(task_result) + }; + let event = state.scheduler_state_event(); + drop(state); + emit_scheduler_state_event(event, &self.stats); + result + } + + fn poll(&self, task_id: u64, cx: &mut Context<'_>) -> Poll> { + let mut state = self.state.lock().unwrap(); + let Some(task) = state.tasks.get_mut(&task_id) else { + // This should never happen and indicates a bug + return Poll::Ready(Err(Error::internal(format!( + "Task with id {} was lost", + task_id + )))); + }; + match task.poll(cx) { + Poll::Ready(_) => { + let task = state.tasks.remove(&task_id).unwrap(); + let (bytes, reservation) = task.consume()?; + state.backpressure_throttle.release(reservation); + // We run on_task_complete even if not newly finished because we released the backpressure reservation + match self.on_task_complete(state) { + Ok(_) => Poll::Ready(bytes), + Err(e) => Poll::Ready(Err(e)), + } + } + Poll::Pending => Poll::Pending, + } + } + + pub(super) fn close(&self) { + let event = { + let mut state = self.state.lock().unwrap(); + for task in std::mem::take(&mut state.tasks).values_mut() { + task.cancel(); + } + state.scheduler_state_event() + }; + emit_scheduler_state_event(event, &self.stats); + } + + // Called when a caller drops a task's handle before the task finishes. Removes + // the task and returns any backpressure reservation it holds to the budget, then + // re-checks the queue so newly-affordable tasks can start. Unlike the standard + // release path (`poll`), this runs without the task being polled to completion, + // so a cancelled read does not leak its reservation. + fn abandon(&self, task_id: u64) { + let mut state = self.state.lock().unwrap(); + let Some(task) = state.tasks.remove(&task_id) else { + // Already consumed by `poll`; nothing to release. + return; + }; + + if let Some(reservation) = task.state.backpressure_reservation() { + state.backpressure_throttle.release(reservation); + } + // Freed budget may make queued tasks runnable; there is no caller to surface + // an error to here. + let _ = self.on_task_complete(state); + } +} + +pub(super) struct TaskHandle { + task_id: u64, + queue: Arc, +} + +impl Future for TaskHandle { + type Output = Result; + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.queue.poll(self.task_id, cx) + } +} + +impl Drop for TaskHandle { + fn drop(&mut self) { + self.queue.abandon(self.task_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::oneshot; + + #[tokio::test] + async fn test_priority_ordering() { + // Backpressure budget of 10 bytes: only one 10-byte task runs at a time. + let queue = Arc::new(IoQueue::new(128, 10, IoStats::default())); + + // Records the priority of each task when its run_fn is invoked (i.e. when + // the task transitions to Running). + let start_order: Arc>> = Arc::new(Mutex::new(Vec::new())); + + // Helper: builds a RunFn that records `prio` in start_order and then + // waits on the oneshot receiver for its result bytes. + let make_run_fn = + |prio: u128, rx: oneshot::Receiver, order: Arc>>| -> RunFn { + Box::new(move || { + order.lock().unwrap().push(prio); + Box::pin(async move { Ok(rx.await.unwrap()) }) + }) + }; + + // Submit a blocker task (priority 0, 10 bytes). + // It starts immediately because there is enough backpressure budget. + let (blocker_tx, blocker_rx) = oneshot::channel(); + let blocker = queue + .clone() + .submit( + 0..10, + 0, + make_run_fn(0, blocker_rx, start_order.clone()), + false, + ) + .unwrap(); + + // Submit four tasks with out-of-order priorities. + // All are queued because the blocker consumed the full budget. + let (tx_30, rx_30) = oneshot::channel(); + let h30 = queue + .clone() + .submit( + 0..10, + 30, + make_run_fn(30, rx_30, start_order.clone()), + false, + ) + .unwrap(); + + let (tx_10, rx_10) = oneshot::channel(); + let h10 = queue + .clone() + .submit( + 0..10, + 10, + make_run_fn(10, rx_10, start_order.clone()), + false, + ) + .unwrap(); + + let (tx_50, rx_50) = oneshot::channel(); + let h50 = queue + .clone() + .submit( + 0..10, + 50, + make_run_fn(50, rx_50, start_order.clone()), + false, + ) + .unwrap(); + + let (tx_20, rx_20) = oneshot::channel(); + let h20 = queue + .clone() + .submit( + 0..10, + 20, + make_run_fn(20, rx_20, start_order.clone()), + false, + ) + .unwrap(); + + // Only the blocker has started so far. + assert_eq!(*start_order.lock().unwrap(), vec![0]); + + // Complete the blocker -> frees budget -> starts priority 10 (lowest value = highest priority). + blocker_tx.send(Bytes::from_static(b"x")).unwrap(); + blocker.await.unwrap(); + assert_eq!(*start_order.lock().unwrap(), vec![0, 10]); + + // Complete priority 10 -> starts priority 20. + tx_10.send(Bytes::from_static(b"x")).unwrap(); + h10.await.unwrap(); + assert_eq!(*start_order.lock().unwrap(), vec![0, 10, 20]); + + // Complete priority 20 -> starts priority 30. + tx_20.send(Bytes::from_static(b"x")).unwrap(); + h20.await.unwrap(); + assert_eq!(*start_order.lock().unwrap(), vec![0, 10, 20, 30]); + + // Complete priority 30 -> starts priority 50. + tx_30.send(Bytes::from_static(b"x")).unwrap(); + h30.await.unwrap(); + assert_eq!(*start_order.lock().unwrap(), vec![0, 10, 20, 30, 50]); + + // Complete priority 50 -> no more pending tasks. + tx_50.send(Bytes::from_static(b"x")).unwrap(); + h50.await.unwrap(); + assert_eq!(*start_order.lock().unwrap(), vec![0, 10, 20, 30, 50]); + } + + #[tokio::test] + async fn test_zero_buffer_bypasses_backpressure() { + // Budget = 0 sets no_backpressure = true, so all tasks start immediately + // regardless of how many bytes are "outstanding". + let queue = Arc::new(IoQueue::new(128, 0, IoStats::default())); + let start_order: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let make_run_fn = + |prio: u128, rx: oneshot::Receiver, order: Arc>>| -> RunFn { + Box::new(move || { + order.lock().unwrap().push(prio); + Box::pin(async move { Ok(rx.await.unwrap()) }) + }) + }; + + let (tx0, rx0) = oneshot::channel(); + let h0 = queue + .clone() + .submit(0..10, 0, make_run_fn(0, rx0, start_order.clone()), false) + .unwrap(); + let (tx1, rx1) = oneshot::channel(); + let h1 = queue + .clone() + .submit(0..10, 1, make_run_fn(1, rx1, start_order.clone()), false) + .unwrap(); + let (tx2, rx2) = oneshot::channel(); + let h2 = queue + .clone() + .submit(0..10, 2, make_run_fn(2, rx2, start_order.clone()), false) + .unwrap(); + + // All three tasks start immediately — no backpressure budget check when max_bytes=0. + assert_eq!(*start_order.lock().unwrap(), vec![0, 1, 2]); + + tx0.send(Bytes::from_static(b"done")).unwrap(); + tx1.send(Bytes::from_static(b"done")).unwrap(); + tx2.send(Bytes::from_static(b"done")).unwrap(); + h0.await.unwrap(); + h1.await.unwrap(); + h2.await.unwrap(); + } + + #[tokio::test] + async fn test_bypass_flag_proceeds_past_exhausted_budget() { + // Budget of 10 bytes. A blocker task fills it. A task with bypass=true starts + // immediately despite the exhausted budget; a normal task stays queued. + let queue = Arc::new(IoQueue::new(128, 10, IoStats::default())); + let start_order: Arc>> = Arc::new(Mutex::new(Vec::new())); + + let make_run_fn = + |prio: u128, rx: oneshot::Receiver, order: Arc>>| -> RunFn { + Box::new(move || { + order.lock().unwrap().push(prio); + Box::pin(async move { Ok(rx.await.unwrap()) }) + }) + }; + + // Blocker (priority 0, 10 bytes): fills the budget. + let (blocker_tx, blocker_rx) = oneshot::channel(); + let blocker = queue + .clone() + .submit( + 0..10, + 0, + make_run_fn(0, blocker_rx, start_order.clone()), + false, + ) + .unwrap(); + + // Normal (priority 1, 10 bytes): blocked — budget exhausted, no priority bypass. + let (normal_tx, normal_rx) = oneshot::channel(); + let normal = queue + .clone() + .submit( + 0..10, + 1, + make_run_fn(1, normal_rx, start_order.clone()), + false, + ) + .unwrap(); + + // Bypass (priority 2, 10 bytes): starts immediately via force_acquire. + let (bypass_tx, bypass_rx) = oneshot::channel(); + let bypass = queue + .clone() + .submit( + 0..10, + 2, + make_run_fn(2, bypass_rx, start_order.clone()), + true, + ) + .unwrap(); + + // Blocker (0) and bypass (2) have started; normal (1) is still queued. + assert_eq!(*start_order.lock().unwrap(), vec![0, 2]); + + // Completing the blocker frees the budget and unblocks the normal task. + blocker_tx.send(Bytes::from_static(b"done")).unwrap(); + blocker.await.unwrap(); + assert_eq!(*start_order.lock().unwrap(), vec![0, 2, 1]); + + bypass_tx.send(Bytes::from_static(b"done")).unwrap(); + bypass.await.unwrap(); + normal_tx.send(Bytes::from_static(b"done")).unwrap(); + normal.await.unwrap(); + } + + #[test] + fn test_same_priority_reservation_continues_after_higher_priority() { + let mut throttle = SimpleBackpressureThrottle::new(10, 128); + + let low_priority_first = throttle.try_acquire(6, 10).unwrap(); + let high_priority = throttle.try_acquire(4, 0).unwrap(); + let low_priority_next = throttle.try_acquire(6, 10); + + assert!( + low_priority_next.is_some(), + "chunks from an already admitted logical request should continue" + ); + + throttle.release(low_priority_first); + throttle.release(high_priority); + throttle.release(low_priority_next.unwrap()); + } +} diff --git a/vendor/lance-io/src/spill.rs b/vendor/lance-io/src/spill.rs new file mode 100644 index 000000000..16b4c10f1 --- /dev/null +++ b/vendor/lance-io/src/spill.rs @@ -0,0 +1,543 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Reclaimable scratch storage. +//! +//! A [`SpillStore`] hands out scratch space for temporary state that is too +//! large to keep in memory and is read back later in the same process (for +//! example, posting lists or shuffle runs accumulated while building an index). +//! The backing storage is reclaimed automatically when the handle is dropped. +//! +//! [`SpillStore::new_spill`] returns a [`Writer`] paired with a [`Spill`] +//! handle: the writer is the byte sink (feed it to `FileWriter::try_new`, or +//! write to it directly); the [`Spill`] reads the bytes back (via +//! [`crate::scheduler::ScanScheduler::open_reader`] for a v2 `FileReader`) and +//! owns the file's lifetime. +//! +//! # Lifecycle +//! +//! - **Write-once.** The only way to obtain a writer is `new_spill`, and each +//! call allocates a fresh unit of storage, so a single spill cannot be +//! written twice — there is no second-writer path to guard against. +//! - **Write-before-read.** [`Spill::reader`] fails until the writer has been +//! shut down, so partially written bytes are never read back. +//! - **RAII.** Dropping the [`Spill`] deletes the file and releases its bytes +//! back to the store's disk budget. The store's temp directory is the +//! backstop for anything leaked if a handle is forgotten. +//! +//! # Disk cap +//! +//! [`LocalSpillStore::with_cap`] enforces a byte budget shared across all live +//! handles, returning a typed [`lance_core::Error::DiskCapExceeded`] rather than +//! silently filling the disk. Accounting is reserve-on-write + release-on-drop +//! (by stat), which is exact for the write-once contract. Two minor +//! inexactnesses are not engineered around: a write aborted at the cap leaks its +//! reservation until the store is dropped, and a file whose size cannot be +//! stat-ed on drop is not released. + +use std::io; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; + +use async_trait::async_trait; +use object_store::path::Path; +use tokio::io::AsyncWrite; + +use lance_core::{Error, Result}; + +use crate::object_store::ObjectStore; +use crate::object_writer::WriteResult; +use crate::traits::{Reader, Writer}; + +/// A factory for scratch storage. +/// +/// The trait is object-safe and `Send + Sync` so it can be held behind an +/// `Arc` (e.g. inside a `Session`). Implementations need not be +/// backed by local files (e.g. in-memory buffers, remote object stores). +#[async_trait] +pub trait SpillStore: Send + Sync + 'static { + /// Allocate a unit of scratch storage. + /// + /// Returns the byte sink to write it with and a [`Spill`] handle to read it + /// back. For a capped store, writes that would exceed the cap fail with + /// [`lance_core::Error::DiskCapExceeded`]. The storage is reclaimed when the + /// [`Spill`] is dropped. + async fn new_spill(&self) -> Result<(Box, Box)>; +} + +/// The readable half of a spill, and the owner of its backing storage. +/// +/// Dropping it reclaims the storage. The trait is object-safe so it can be +/// returned as `Box` from [`SpillStore::new_spill`]. +#[async_trait] +pub trait Spill: Send + Sync { + /// Open a reader over the spilled bytes. + /// + /// Fails until the paired writer has been shut down, since the bytes are not + /// complete before then. + async fn reader(&self) -> Result>; +} + +/// A shared, cloneable byte budget. +/// +/// Cloning produces another handle to the *same* underlying counter, so a quota +/// shared across many writers enforces a single combined cap. +#[derive(Debug, Clone)] +struct DiskQuota { + cap_bytes: u64, + used: Arc>, +} + +impl DiskQuota { + fn new(cap_bytes: u64) -> Self { + Self { + cap_bytes, + used: Arc::new(Mutex::new(0)), + } + } + + /// Try to reserve `n` bytes, failing with [`Error::DiskCapExceeded`] if the + /// reservation would push total usage past the cap. + fn try_reserve(&self, n: u64) -> Result<()> { + // The lock is held only for a couple of arithmetic ops and never across + // an `.await`, so a std `Mutex` is the simplest correct choice. + let mut used = self.used.lock().unwrap(); + let next = used.saturating_add(n); + if next > self.cap_bytes { + return Err(Error::disk_cap_exceeded(self.cap_bytes, *used)); + } + *used = next; + Ok(()) + } + + /// Release `n` previously reserved bytes back to the budget. + fn release(&self, n: u64) { + // Saturating sub keeps a stray double-release from underflowing. + let mut used = self.used.lock().unwrap(); + *used = used.saturating_sub(n); + } +} + +/// The byte sink handed out by [`SpillStore::new_spill`]. +/// +/// It optionally reserves a [`DiskQuota`] as bytes are written (keeping cap +/// enforcement inside the spill store rather than in [`ObjectStore`], and +/// working for any backend the store opens), and flips a shared `finished` flag +/// on shutdown so the paired [`Spill`] knows the bytes are complete. +struct SpillWriter { + inner: Box, + quota: Option, + finished: Arc, +} + +impl AsyncWrite for SpillWriter { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + let Some(quota) = &this.quota else { + return Pin::new(this.inner.as_mut()).poll_write(cx, buf); + }; + // Reserve up-front for the bytes we intend to write, then release the + // remainder the inner writer did not accept so the reservation tracks + // bytes actually buffered (and, for a write-once file, the file size). + if let Err(e) = quota.try_reserve(buf.len() as u64) { + return Poll::Ready(Err(io::Error::other(e))); + } + let poll = Pin::new(this.inner.as_mut()).poll_write(cx, buf); + match &poll { + Poll::Ready(Ok(n)) => quota.release((buf.len() - *n) as u64), + _ => quota.release(buf.len() as u64), + } + poll + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(self.get_mut().inner.as_mut()).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let poll = Pin::new(this.inner.as_mut()).poll_shutdown(cx); + if matches!(poll, Poll::Ready(Ok(()))) { + // Mirrors `Writer::shutdown` so the flag is set whichever shutdown + // surface the consumer drives (`AsyncWrite` vs the `Writer` trait). + this.finished.store(true, Ordering::Relaxed); + } + poll + } +} + +#[async_trait] +impl Writer for SpillWriter { + async fn tell(&mut self) -> Result { + self.inner.tell().await + } + + async fn shutdown(&mut self) -> Result { + let result = self.inner.shutdown().await?; + // Signal the paired `Spill` that the bytes are now complete. `Relaxed` + // is sufficient: this only flags that shutdown happened; the file + // contents are synchronized through the filesystem, not this flag. + self.finished.store(true, Ordering::Relaxed); + Ok(result) + } +} + +/// A [`SpillStore`] that writes temporary files to a local temp directory. +/// +/// By default there is no disk cap. Use [`LocalSpillStore::with_cap`] to +/// configure one shared across every handle this store produces. +/// +/// The temp directory is deleted when the store is dropped, cleaning up any +/// files whose handles have already been dropped. +pub struct LocalSpillStore { + store: Arc, + /// Backstop cleanup: removes the whole scratch directory on drop. + temp_dir: Arc, + file_counter: Arc, + /// Byte budget shared across every handle, enforced while writing. + quota: Option, +} + +impl LocalSpillStore { + /// Create a store with no disk cap. + pub fn new() -> Result { + Ok(Self { + store: Arc::new(ObjectStore::local()), + temp_dir: Arc::new(tempfile::tempdir()?), + file_counter: Arc::new(AtomicU64::new(0)), + quota: None, + }) + } + + /// Create a store that returns [`lance_core::Error::DiskCapExceeded`] once + /// total bytes written across all live handles would exceed `cap_bytes`. + pub fn with_cap(cap_bytes: u64) -> Result { + Ok(Self { + store: Arc::new(ObjectStore::local()), + temp_dir: Arc::new(tempfile::tempdir()?), + file_counter: Arc::new(AtomicU64::new(0)), + quota: Some(DiskQuota::new(cap_bytes)), + }) + } +} + +impl Default for LocalSpillStore { + fn default() -> Self { + Self::new().expect("failed to create temp directory for LocalSpillStore") + } +} + +#[async_trait] +impl SpillStore for LocalSpillStore { + async fn new_spill(&self) -> Result<(Box, Box)> { + let idx = self.file_counter.fetch_add(1, Ordering::Relaxed); + let fs_path = self.temp_dir.path().join(format!("spill_{idx:06}.bin")); + let os_path = Path::from_absolute_path(&fs_path)?; + let finished = Arc::new(AtomicBool::new(false)); + + let writer = Box::new(SpillWriter { + inner: self.store.create(&os_path).await?, + quota: self.quota.clone(), + finished: finished.clone(), + }); + let spill = Box::new(LocalSpill { + store: self.store.clone(), + os_path, + fs_path, + quota: self.quota.clone(), + finished, + _temp_dir: self.temp_dir.clone(), + }); + Ok((writer, spill)) + } +} + +/// The readable half of a [`LocalSpillStore`] spill; reclaims the file on drop. +struct LocalSpill { + store: Arc, + os_path: Path, + fs_path: PathBuf, + quota: Option, + /// Set by the paired [`SpillWriter`] once it has been shut down. + finished: Arc, + /// Keep the store's temp directory alive for at least this file's lifetime. + _temp_dir: Arc, +} + +#[async_trait] +impl Spill for LocalSpill { + async fn reader(&self) -> Result> { + // `Relaxed` is sufficient: the flag only gates "has the writer shut + // down"; the bytes themselves are synchronized through the filesystem, + // not this load. + if !self.finished.load(Ordering::Relaxed) { + return Err(Error::invalid_input( + "spill reader requested before the writer was shut down", + )); + } + self.store.open(&self.os_path).await + } +} + +impl Drop for LocalSpill { + fn drop(&mut self) { + // Release the bytes this file occupied back to the budget. We stat the + // persisted file rather than tracking writes, which is exact for the + // write-once contract. + if let Some(quota) = &self.quota + && let Ok(metadata) = std::fs::metadata(&self.fs_path) + { + quota.release(metadata.len()); + } + // Best-effort removal; the temp dir is the backstop. + let _ = std::fs::remove_file(&self.fs_path); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncWriteExt; + + /// Write `data` to a fresh writer and shut it down. + async fn finish_writer(mut writer: Box, data: &[u8]) -> Result<()> { + writer.write_all(data).await?; + Writer::shutdown(writer.as_mut()).await?; + Ok(()) + } + + #[test] + fn test_disk_quota_reserve_release() { + let quota = DiskQuota::new(100); + quota.try_reserve(60).unwrap(); + assert!(quota.try_reserve(60).is_err()); + quota.release(60); + quota.try_reserve(60).unwrap(); + // Reserving exactly up to the cap succeeds; one byte past it fails. + quota.try_reserve(40).unwrap(); + assert!(quota.try_reserve(1).is_err()); + } + + #[tokio::test] + async fn test_write_then_read() { + let store = LocalSpillStore::new().unwrap(); + let (writer, spill) = store.new_spill().await.unwrap(); + + let data = b"hello spill world"; + finish_writer(writer, data).await.unwrap(); + + let reader = spill.reader().await.unwrap(); + let read_back = reader.get_all().await.unwrap(); + assert_eq!(read_back.as_ref(), data); + } + + #[tokio::test] + async fn test_reader_requires_finished_writer() { + let store = LocalSpillStore::new().unwrap(); + let (mut writer, spill) = store.new_spill().await.unwrap(); + writer.write_all(b"partial").await.unwrap(); + + // Reading before the writer is shut down is rejected. + let Err(err) = spill.reader().await else { + panic!("reader before shutdown should be rejected"); + }; + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput, got {err:?}" + ); + + // After shutdown the reader sees the bytes. + Writer::shutdown(writer.as_mut()).await.unwrap(); + let reader = spill.reader().await.unwrap(); + assert_eq!(reader.get_all().await.unwrap().as_ref(), b"partial"); + } + + #[tokio::test] + async fn test_reader_ready_after_async_shutdown() { + // Shutting down through the `AsyncWrite` surface (not the `Writer` + // trait) must also mark the spill readable — covers poll_shutdown's + // flag set, the path the `Writer::shutdown` tests don't reach. + let store = LocalSpillStore::new().unwrap(); + let (mut writer, spill) = store.new_spill().await.unwrap(); + writer.write_all(b"async").await.unwrap(); + AsyncWriteExt::shutdown(&mut writer).await.unwrap(); + + let reader = spill.reader().await.unwrap(); + assert_eq!(reader.get_all().await.unwrap().as_ref(), b"async"); + } + + #[tokio::test] + async fn test_empty_spill() { + // A spill written with no bytes round-trips empty, and the capped path + // handles the zero-byte reserve/stat without error. + let store = LocalSpillStore::with_cap(100).unwrap(); + let (writer, spill) = store.new_spill().await.unwrap(); + finish_writer(writer, b"").await.unwrap(); + + let reader = spill.reader().await.unwrap(); + assert!(reader.get_all().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn test_raii_cleanup() { + let store = LocalSpillStore::new().unwrap(); + let (writer, spill) = store.new_spill().await.unwrap(); + finish_writer(writer, b"some bytes").await.unwrap(); + + // The first spill gets a deterministic name under the store's temp dir. + let path = store.temp_dir.path().join("spill_000000.bin"); + assert!(path.exists()); + drop(spill); + assert!(!path.exists(), "spill file should be deleted on drop"); + } + + #[tokio::test] + async fn test_cap_exceeded() { + let store = LocalSpillStore::with_cap(100).unwrap(); + let (writer, _spill) = store.new_spill().await.unwrap(); + let err = finish_writer(writer, &[0u8; 101]).await.unwrap_err(); + assert!( + matches!(err, Error::DiskCapExceeded { cap_bytes: 100, .. }), + "expected DiskCapExceeded, got {err:?}" + ); + } + + #[tokio::test] + async fn test_cap_shared_across_files() { + let store = LocalSpillStore::with_cap(100).unwrap(); + let (writer_a, _spill_a) = store.new_spill().await.unwrap(); + let (writer_b, _spill_b) = store.new_spill().await.unwrap(); + + finish_writer(writer_a, &[0u8; 60]).await.unwrap(); + // 60 already reserved by `a`; writing 60 more would reach 120 > 100. + let err = finish_writer(writer_b, &[0u8; 60]).await.unwrap_err(); + assert!( + matches!(err, Error::DiskCapExceeded { cap_bytes: 100, .. }), + "expected DiskCapExceeded, got {err:?}" + ); + } + + #[tokio::test] + async fn test_cap_freed_on_drop() { + let store = LocalSpillStore::with_cap(100).unwrap(); + + { + let (writer, spill) = store.new_spill().await.unwrap(); + finish_writer(writer, &[0u8; 80]).await.unwrap(); + // `spill` drops at the end of this block, releasing its 80 bytes. + drop(spill); + } + + let (writer, _spill) = store.new_spill().await.unwrap(); + // Succeeds because the cap is no longer under pressure. + finish_writer(writer, &[0u8; 80]).await.unwrap(); + } + + #[tokio::test] + async fn test_custom_implementation() { + // A custom store can satisfy the traits without a local file. + struct MemStore; + struct MemSpill; + + #[async_trait] + impl Spill for MemSpill { + async fn reader(&self) -> Result> { + ObjectStore::memory().open(&Path::from("/mem")).await + } + } + + #[async_trait] + impl SpillStore for MemStore { + async fn new_spill(&self) -> Result<(Box, Box)> { + let writer = ObjectStore::memory().create(&Path::from("/mem")).await?; + Ok((writer, Box::new(MemSpill))) + } + } + + let store = MemStore; + // Exercise the factory + trait objects; the in-memory store is a fresh + // instance per call so we don't round-trip data here. + let (_writer, _spill) = store.new_spill().await.unwrap(); + } + + /// A [`Writer`] whose `poll_write` accepts a fixed number of bytes per call, + /// or fails, so we can drive the [`SpillWriter`] release arms that the local + /// backend (which accepts every write in full) never hits. + struct ControlledWriter { + outcome: Poll>, + } + + impl AsyncWrite for ControlledWriter { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + match &self.outcome { + Poll::Ready(Ok(n)) => Poll::Ready(Ok((*n).min(buf.len()))), + Poll::Ready(Err(e)) => Poll::Ready(Err(io::Error::new(e.kind(), e.to_string()))), + Poll::Pending => Poll::Pending, + } + } + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + #[async_trait] + impl Writer for ControlledWriter { + async fn tell(&mut self) -> Result { + Ok(0) + } + async fn shutdown(&mut self) -> Result { + Ok(WriteResult::default()) + } + } + + #[tokio::test] + async fn test_spill_writer_releases_unaccepted_bytes() { + // Short write: the inner writer accepts only 10 of the 40 reserved bytes, + // so the 30-byte remainder must be returned to the budget. + let quota = DiskQuota::new(100); + let mut writer = SpillWriter { + inner: Box::new(ControlledWriter { + outcome: Poll::Ready(Ok(10)), + }), + quota: Some(quota.clone()), + finished: Arc::new(AtomicBool::new(false)), + }; + let n = writer.write(&[0u8; 40]).await.unwrap(); + assert_eq!(n, 10); + assert_eq!( + *quota.used.lock().unwrap(), + 10, + "only the accepted bytes should remain reserved" + ); + + // Failed write: the full reservation must be released. + let quota = DiskQuota::new(100); + let mut writer = SpillWriter { + inner: Box::new(ControlledWriter { + outcome: Poll::Ready(Err(io::Error::other("boom"))), + }), + quota: Some(quota.clone()), + finished: Arc::new(AtomicBool::new(false)), + }; + writer.write(&[0u8; 40]).await.unwrap_err(); + assert_eq!( + *quota.used.lock().unwrap(), + 0, + "a failed write should release its entire reservation" + ); + } +} diff --git a/vendor/lance-io/src/stream.rs b/vendor/lance-io/src/stream.rs new file mode 100644 index 000000000..a1c2bb393 --- /dev/null +++ b/vendor/lance-io/src/stream.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::pin::Pin; +use std::task::{Context, Poll}; + +use arrow_array::RecordBatch; +use arrow_schema::{ArrowError, SchemaRef}; +use futures::stream::BoxStream; +use futures::{Stream, StreamExt}; +use pin_project::pin_project; + +use lance_core::Result; + +pub type BatchStream = BoxStream<'static, Result>; + +pub fn arrow_stream_to_lance_stream( + arrow_stream: BoxStream<'static, std::result::Result>, +) -> BatchStream { + arrow_stream.map(|r| r.map_err(Into::into)).boxed() +} + +/// RecordBatch Stream trait. +pub trait RecordBatchStream: Stream> + Send { + /// Returns the schema of the stream. + fn schema(&self) -> SchemaRef; +} + +/// Combines a [`Stream`] with a [`SchemaRef`] implementing +/// [`RecordBatchStream`] for the combination +#[pin_project] +pub struct RecordBatchStreamAdapter { + schema: SchemaRef, + + #[pin] + stream: S, +} + +impl RecordBatchStreamAdapter { + /// Creates a new [`RecordBatchStreamAdapter`] from the provided schema and stream + pub fn new(schema: SchemaRef, stream: S) -> Self { + Self { schema, stream } + } +} + +impl std::fmt::Debug for RecordBatchStreamAdapter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RecordBatchStreamAdapter") + .field("schema", &self.schema) + .finish() + } +} + +impl RecordBatchStream for RecordBatchStreamAdapter +where + S: Stream> + Send + 'static, +{ + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +impl Stream for RecordBatchStreamAdapter +where + S: Stream>, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().stream.poll_next(cx) + } + + fn size_hint(&self) -> (usize, Option) { + self.stream.size_hint() + } +} diff --git a/vendor/lance-io/src/testing.rs b/vendor/lance-io/src/testing.rs new file mode 100644 index 000000000..aca9b925a --- /dev/null +++ b/vendor/lance-io/src/testing.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors +use std::fmt::{self, Display, Formatter}; + +use async_trait::async_trait; +use futures::stream::BoxStream; +use mockall::mock; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore as OSObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + Result as OSResult, path::Path, +}; +use std::future::Future; + +mock! { + pub ObjectStore {} + + #[async_trait] + impl OSObjectStore for ObjectStore { + async fn put_opts(&self, location: &Path, bytes: PutPayload, opts: PutOptions) -> OSResult; + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult>; + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, + options: GetOptions + ) -> std::pin::Pin > +Send+'async_trait> > where + Self: 'async_trait, + 'life0: 'async_trait, + 'life1: 'async_trait; + fn delete_stream(&self, locations: BoxStream<'static, OSResult>) -> BoxStream<'static, OSResult>; + fn list<'a>(&'a self, prefix: Option<&'a Path>) -> BoxStream<'_, OSResult>; + async fn list_with_delimiter<'a, 'b>(&'a self, prefix: Option<&'b Path>) -> OSResult; + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()>; + } +} + +impl std::fmt::Debug for MockObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "MockObjectStore") + } +} + +impl Display for MockObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "MockObjectStore") + } +} diff --git a/vendor/lance-io/src/traits.rs b/vendor/lance-io/src/traits.rs new file mode 100644 index 000000000..6a40171b6 --- /dev/null +++ b/vendor/lance-io/src/traits.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::ops::Range; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::{StreamExt, future::BoxFuture, stream::BoxStream}; +use lance_core::deepsize::DeepSizeOf; +use object_store::path::Path; +use prost::Message; +use tokio::io::{AsyncWrite, AsyncWriteExt}; + +use lance_core::Result; + +use crate::object_writer::WriteResult; + +pub trait ProtoStruct { + type Proto: Message; +} + +pub type ByteStream = BoxStream<'static, object_store::Result>; + +/// A trait for writing to a file on local file system or object store. +#[async_trait] +pub trait Writer: AsyncWrite + Unpin + Send { + /// Tell the current offset. + async fn tell(&mut self) -> Result; + + /// Flush all buffered data and finalize the write, returning metadata about + /// the written object. + async fn shutdown(&mut self) -> Result; +} + +#[async_trait] +impl Writer for Box { + async fn tell(&mut self) -> Result { + self.as_mut().tell().await + } + + async fn shutdown(&mut self) -> Result { + self.as_mut().shutdown().await + } +} + +/// Lance Write Extension. +#[async_trait] +pub trait WriteExt { + /// Write a Protobuf message to the [Writer], and returns the file position + /// where the protobuf is written. + async fn write_protobuf(&mut self, msg: &impl Message) -> Result; + + async fn write_struct< + 'b, + M: Message + From<&'b T>, + T: ProtoStruct + Send + Sync + 'b, + >( + &mut self, + obj: &'b T, + ) -> Result { + let msg: M = M::from(obj); + self.write_protobuf(&msg).await + } + /// Write magics to the tail of a file before closing the file. + async fn write_magics( + &mut self, + pos: usize, + major_version: i16, + minor_version: i16, + magic: &[u8], + ) -> Result<()>; + + async fn copy_from_reader(&mut self, reader: &dyn Reader) -> Result; + + async fn copy_range_from_reader( + &mut self, + reader: &dyn Reader, + range: Range, + ) -> Result; +} + +#[async_trait] +impl WriteExt for W { + async fn write_protobuf(&mut self, msg: &impl Message) -> Result { + let offset = self.tell().await?; + + let len = msg.encoded_len(); + + self.write_u32_le(len as u32).await?; + self.write_all(&msg.encode_to_vec()).await?; + + Ok(offset) + } + + async fn write_magics( + &mut self, + pos: usize, + major_version: i16, + minor_version: i16, + magic: &[u8], + ) -> Result<()> { + self.write_i64_le(pos as i64).await?; + self.write_i16_le(major_version).await?; + self.write_i16_le(minor_version).await?; + self.write_all(magic).await?; + Ok(()) + } + + async fn copy_from_reader(&mut self, reader: &dyn Reader) -> Result { + let mut stream = reader.get_stream().await?; + let mut copied = 0usize; + while let Some(chunk) = stream.next().await { + let bytes = chunk?; + copied += bytes.len(); + self.write_all(&bytes).await?; + } + Ok(copied) + } + + async fn copy_range_from_reader( + &mut self, + reader: &dyn Reader, + range: Range, + ) -> Result { + let mut stream = reader.get_range_stream(range).await?; + let mut copied = 0usize; + while let Some(chunk) = stream.next().await { + let bytes = chunk?; + copied += bytes.len(); + self.write_all(&bytes).await?; + } + Ok(copied) + } +} + +pub trait Reader: std::fmt::Debug + Send + Sync + DeepSizeOf { + fn path(&self) -> &Path; + + /// Suggest optimal I/O size per storage device. + fn block_size(&self) -> usize; + + /// Suggest optimal I/O parallelism per storage device. + fn io_parallelism(&self) -> usize; + + /// Object/File Size. + fn size(&self) -> BoxFuture<'_, object_store::Result>; + + /// Read a range of bytes from the object. + /// + /// TODO: change to read_at()? + fn get_range(&self, range: Range) -> BoxFuture<'static, object_store::Result>; + + /// Read all bytes from the object. + /// + /// By default this reads the size in a separate IOP but some implementations + /// may not need the size beforehand. + fn get_all(&self) -> BoxFuture<'_, object_store::Result>; + + /// Read the entire object as a byte stream. + fn get_stream(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { + let bytes = self.get_all().await?; + Ok(futures::stream::once(async move { Ok(bytes) }).boxed()) + }) + } + + /// Read a byte range as a byte stream. + fn get_range_stream( + &self, + range: Range, + ) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { + let bytes = self.get_range(range).await?; + Ok(futures::stream::once(async move { Ok(bytes) }).boxed()) + }) + } +} diff --git a/vendor/lance-io/src/uring.rs b/vendor/lance-io/src/uring.rs new file mode 100644 index 000000000..067861781 --- /dev/null +++ b/vendor/lance-io/src/uring.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! io_uring-based I/O for disks with high IOPS capacity (e.g. NVMe) +//! +//! This module provides two implementations of the [`Reader`](crate::traits::Reader) trait +//! using Linux's io_uring interface for asynchronous I/O. +//! +//! One of these uses a pool of dedicated background threads which each own an io_uring instance. +//! Read requests are submitted to a background thread's pool. +//! +//! The other implementation uses a thread-local io_uring instance. This only works if the future +//! is polled by the same thread that submitted the request. This means that the runtime must be +//! a single-threaded runtime. +//! +//! # Configuration +//! +//! The io_uring reader is enabled by using the `file+uring://` URI scheme instead of `file://`. +//! Additional tuning parameters are controlled by environment variables: +//! +//! - `LANCE_URING_CURRENT_THREAD` - Use thread-local io_uring (default: false) +//! - `LANCE_URING_BLOCK_SIZE` - Block size in bytes (default: 4KB) +//! - `LANCE_URING_IO_PARALLELISM` - Max concurrent operations (default: 128) +//! - `LANCE_URING_QUEUE_DEPTH` - io_uring queue depth (default: 16K) +//! - `LANCE_URING_THREAD_COUNT` - Number of io_uring threads to use (default: 2) +//! - `LANCE_URING_SUBMIT_BATCH_SIZE` - Number of requests to batch before submitting (default: 128) +//! - `LANCE_URING_POLL_TIMEOUT_MS` - Thread poll timeout in milliseconds (default: 10) +//! +//! Note: the block size and io parallelism are not actually used by the io_uring implementation. These +//! variables just control what the filesystem reports up to Lance. +//! +//! # Platform Support +//! +//! This module is only available on Linux and requires kernel 5.1 or newer. +//! On other platforms, the code falls back to [`LocalObjectReader`](crate::local::LocalObjectReader). +//! +//! # Example +//! +//! ```no_run +//! # use lance_io::object_store::ObjectStore; +//! # async fn example() -> lance_core::Result<()> { +//! // Enable io_uring by using the file+uring:// scheme +//! let uri = "file+uring:///path/to/file.dat"; +//! let (store, path) = ObjectStore::from_uri(uri).await?; +//! let reader = store.open(&path).await?; +//! +//! // Reader will use io_uring +//! let data = reader.get_range(0..1024).await?; +//! # Ok(()) +//! # } +//! ``` + +mod future; +mod reader; +mod requests; +mod thread; + +// Thread-local io_uring implementation for current-thread runtimes +pub(crate) mod current_thread; +pub(crate) mod current_thread_future; + +#[cfg(test)] +mod tests; + +use std::sync::LazyLock; + +pub(crate) use current_thread::UringCurrentThreadReader; +pub use reader::UringReader; + +/// Default block size for io_uring reads (4KB) +pub const DEFAULT_URING_BLOCK_SIZE: usize = 4 * 1024; + +/// Default I/O parallelism for io_uring (128 concurrent operations) +pub const DEFAULT_URING_IO_PARALLELISM: usize = 128; + +/// Default io_uring queue depth (16K entries) +pub const DEFAULT_URING_QUEUE_DEPTH: usize = 16 * 1024; + +/// Cached `LANCE_URING_BLOCK_SIZE` env var, read once at first access. +pub(crate) static URING_BLOCK_SIZE: LazyLock> = LazyLock::new(|| { + std::env::var("LANCE_URING_BLOCK_SIZE") + .ok() + .and_then(|s| s.parse().ok()) +}); diff --git a/vendor/lance-io/src/uring/current_thread.rs b/vendor/lance-io/src/uring/current_thread.rs new file mode 100644 index 000000000..ae2e3414e --- /dev/null +++ b/vendor/lance-io/src/uring/current_thread.rs @@ -0,0 +1,430 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Thread-local io_uring implementation for current-thread runtimes. +//! +//! This implementation creates a thread-local IoUring instance per thread +//! and directly processes completions during future polling, eliminating +//! the need for background threads and MPSC channels. + +use super::requests::{IoRequest, RequestState}; +use super::{DEFAULT_URING_BLOCK_SIZE, DEFAULT_URING_IO_PARALLELISM, URING_BLOCK_SIZE}; +use crate::local::to_local_path; +use crate::traits::Reader; +use crate::uring::DEFAULT_URING_QUEUE_DEPTH; +use crate::utils::tracking_store::IOTracker; +use bytes::{Bytes, BytesMut}; +use futures::FutureExt; +use futures::future::BoxFuture; +use io_uring::{IoUring, opcode, types}; +use lance_core::deepsize::DeepSizeOf; +use lance_core::{Error, Result}; +use object_store::path::Path; + +use std::cell::{LazyCell, RefCell}; +use std::collections::HashMap; +use std::fs::File; +use std::future::Future; +use std::io::{self, ErrorKind}; +use std::ops::Range; +use std::pin::Pin; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use tracing::instrument; + +// Re-use file handle types from reader.rs +use super::reader::{CacheKey, CachedReaderData, HANDLE_CACHE, UringFileHandle}; + +/// Global counter for generating unique user_data values +static USER_DATA_COUNTER: AtomicU64 = AtomicU64::new(1); + +/// Thread-local io_uring instance with pending requests +struct ThreadLocalUring { + ring: IoUring, + pending: HashMap>, +} + +thread_local! { + static URING: LazyCell> = LazyCell::new(|| { + let queue_depth = std::env::var("LANCE_URING_QUEUE_DEPTH") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_URING_QUEUE_DEPTH); + + let ring = IoUring::builder() + // Ensures work is only done in submit_and_wait + .setup_defer_taskrun() + // Enable perf. optimization when there is only one issuer thread + .setup_single_issuer() + .build(queue_depth as u32) + .expect("Failed to create io_uring"); + + log::debug!( + "Created thread-local io_uring with queue depth {}", + queue_depth + ); + + RefCell::new(ThreadLocalUring { + ring, + pending: HashMap::new(), + }) + }); +} + +/// Push request to thread-local submission queue +pub(super) fn push_request(request: Arc) -> io::Result<()> { + URING.with(|cell| { + let mut uring = cell.borrow_mut(); + + // Generate unique user_data + let user_data = USER_DATA_COUNTER.fetch_add(1, Ordering::Relaxed); + + // Get buffer pointer, adjusting for any bytes already read (short read retry) + let (buffer_ptr, read_offset, read_length) = { + let state = request.state.lock().unwrap(); + let br = state.bytes_read; + ( + unsafe { state.buffer.as_ptr().add(br) as *mut u8 }, + request.offset + br as u64, + (request.length - br) as u32, + ) + }; + + // Prepare read operation + let read_op = + opcode::Read::new(types::Fd(request.fd), buffer_ptr, read_length).offset(read_offset); + + // Get submission queue + let mut sq = uring.ring.submission(); + + // Check if SQ has space + if sq.is_full() { + drop(sq); + return Err(io::Error::new( + io::ErrorKind::WouldBlock, + "io_uring submission queue full", + )); + } + + // Push to SQ + unsafe { + sq.push(&read_op.build().user_data(user_data)) + .map_err(|_| io::Error::other("Failed to push to SQ"))?; + } + drop(sq); + + // Track request in pending map + uring.pending.insert(user_data, request); + + // Don't submit here - let the future handle submission + + Ok(()) + }) +} + +/// Process completions from thread-local IoUring +pub(super) fn process_thread_local_completions() -> io::Result { + URING.with(|cell| { + let mut uring = cell.borrow_mut(); + let mut completed = 0; + let mut retries: Vec> = Vec::new(); + + // Collect completions first to avoid borrowing ring and pending simultaneously + let cqes: Vec<_> = uring + .ring + .completion() + .map(|cqe| (cqe.user_data(), cqe.result())) + .collect(); + + for (user_data, result) in cqes { + if let Some(request) = uring.pending.remove(&user_data) { + let mut state = request.state.lock().unwrap(); + + if result < 0 { + // Kernel error + state.err = Some(io::Error::from_raw_os_error(-result)); + state.completed = true; + } else if result == 0 { + // EOF before full read completed + let br = state.bytes_read; + state.err = Some(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("unexpected EOF: read {} of {} bytes", br, request.length), + )); + state.buffer.truncate(br); + state.completed = true; + } else { + // Positive result: n bytes read + let n = result as usize; + state.bytes_read += n; + let br = state.bytes_read; + + if br >= request.length { + // Full read complete + state.buffer.truncate(br); + state.completed = true; + } else { + // Short read — need retry; don't mark completed or wake + drop(state); + retries.push(request); + + continue; + } + } + + // Wake waiting future + if let Some(waker) = state.waker.take() { + drop(state); + waker.wake(); + } + + completed += 1; + } else { + log::warn!("Received completion for unknown user_data: {}", user_data); + } + } + + // Resubmit short-read retries + for request in retries { + // Generate unique user_data + let user_data = USER_DATA_COUNTER.fetch_add(1, Ordering::Relaxed); + + let (buffer_ptr, read_offset, read_length) = { + let state = request.state.lock().unwrap(); + let br = state.bytes_read; + ( + unsafe { state.buffer.as_ptr().add(br) as *mut u8 }, + request.offset + br as u64, + (request.length - br) as u32, + ) + }; + + let read_op = opcode::Read::new(types::Fd(request.fd), buffer_ptr, read_length) + .offset(read_offset); + + let mut sq = uring.ring.submission(); + if sq.is_full() { + drop(sq); + request.fail(io::Error::new( + io::ErrorKind::WouldBlock, + "io_uring submission queue full during retry", + )); + continue; + } + + unsafe { + if sq.push(&read_op.build().user_data(user_data)).is_err() { + request.fail(io::Error::other("Failed to push short-read retry to SQ")); + continue; + } + } + drop(sq); + + uring.pending.insert(user_data, request); + } + + if completed > 0 { + log::trace!("Processed {} completions", completed); + } + + Ok(completed) + }) +} + +/// Submit all pending requests and wait with timeout 0 (non-blocking) +pub(super) fn submit_and_wait_thread_local() -> io::Result<()> { + URING.with(|cell| { + let uring = cell.borrow_mut(); + // Submit with wait=1 (do at least some work) + uring.ring.submit_and_wait(1)?; + Ok(()) + }) +} + +/// Thread-local io_uring-based reader for current-thread runtimes +#[derive(Debug)] +pub struct UringCurrentThreadReader { + /// File handle + handle: Arc, + + /// Block size for I/O operations + block_size: usize, + + /// File size (determined at open time) + size: usize, + + /// I/O tracker for monitoring operations + io_tracker: Arc, +} + +impl DeepSizeOf for UringCurrentThreadReader { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // Skip file handle (just a system resource) + // Only count the path's deep size + self.handle.path.as_ref().deep_size_of_children(context) + } +} + +impl UringCurrentThreadReader { + /// Open a file with thread-local io_uring + /// + /// This reuses the file handle caching infrastructure from UringReader + #[instrument(level = "debug")] + pub(crate) async fn open( + path: &Path, + block_size: usize, + known_size: Option, + io_tracker: Arc, + ) -> Result> { + // Determine block size with environment variable override + let block_size = URING_BLOCK_SIZE.unwrap_or(block_size.max(DEFAULT_URING_BLOCK_SIZE)); + + let cache_key = CacheKey::new(path, block_size); + + // Try to get from cache first + if let Some(data) = HANDLE_CACHE.get(&cache_key).await { + // Use known_size if provided, otherwise use cached size + let size = known_size.unwrap_or(data.size); + return Ok(Box::new(Self { + handle: data.handle, + block_size, + size, + io_tracker, + }) as Box); + } + + // Cache miss - open file and get size + let path_clone = path.clone(); + let local_path = to_local_path(path); + + let data = tokio::task::spawn_blocking(move || { + let file = File::open(&local_path).map_err(|e| match e.kind() { + ErrorKind::NotFound => Error::not_found(path_clone.to_string()), + _ => e.into(), + })?; + + // Get size from known_size or file metadata + let size = match known_size { + Some(s) => s, + None => file.metadata()?.len() as usize, + }; + + Ok::<_, Error>(CachedReaderData { + handle: Arc::new(UringFileHandle::new(file, path_clone)), + size, + }) + }) + .await??; + + // Insert into cache + HANDLE_CACHE.insert(cache_key, data.clone()).await; + + // Return new reader instance + Ok(Box::new(Self { + handle: data.handle.clone(), + block_size, + size: data.size, + io_tracker, + }) as Box) + } + + /// Submit a read request and return a future + fn submit_read( + &self, + offset: u64, + length: usize, + ) -> Pin> + Send>> { + let mut buffer = BytesMut::with_capacity(length); + unsafe { + buffer.set_len(length); + } + + let request = Arc::new(IoRequest { + fd: self.handle.fd, + offset, + length, + thread_id: std::thread::current().id(), + state: Mutex::new(RequestState { + completed: false, + waker: None, + err: None, + buffer, + bytes_read: 0, + }), + }); + + match push_request(request.clone()) { + Ok(()) => Box::pin(super::current_thread_future::UringCurrentThreadFuture::new( + request, + )), + Err(e) => Box::pin(async move { + Err(object_store::Error::Generic { + store: "io_uring_ct", + source: Box::new(e), + }) + }), + } + } +} + +impl Reader for UringCurrentThreadReader { + fn path(&self) -> &Path { + &self.handle.path + } + + fn block_size(&self) -> usize { + self.block_size + } + + fn io_parallelism(&self) -> usize { + std::env::var("LANCE_URING_IO_PARALLELISM") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_URING_IO_PARALLELISM) + } + + /// Returns the file size + fn size(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { Ok(self.size) }) + } + + /// Read a range of bytes using thread-local io_uring + #[instrument(level = "debug", skip(self))] + fn get_range(&self, range: Range) -> BoxFuture<'static, object_store::Result> { + let io_tracker = self.io_tracker.clone(); + let path = self.handle.path.clone(); + let num_bytes = range.len() as u64; + let range_u64 = (range.start as u64)..(range.end as u64); + + let metrics = self.io_tracker.begin_io("get"); + self.submit_read(range.start as u64, range.len()) + .map(move |result| { + metrics.record(&result, num_bytes); + if result.is_ok() { + io_tracker.record_read("get_range", path, num_bytes, Some(range_u64)); + } + result + }) + .boxed() + } + + /// Read the entire file using thread-local io_uring + #[instrument(level = "debug", skip(self))] + fn get_all(&self) -> BoxFuture<'static, object_store::Result> { + let size = self.size; + let io_tracker = self.io_tracker.clone(); + let path = self.handle.path.clone(); + + let metrics = self.io_tracker.begin_io("get"); + self.submit_read(0, size) + .map(move |result| { + let num_bytes = result.as_ref().map_or(0, |bytes| bytes.len() as u64); + metrics.record(&result, num_bytes); + if result.is_ok() { + io_tracker.record_read("get_all", path, num_bytes, None); + } + result + }) + .boxed() + } +} diff --git a/vendor/lance-io/src/uring/current_thread_future.rs b/vendor/lance-io/src/uring/current_thread_future.rs new file mode 100644 index 000000000..dbcf22426 --- /dev/null +++ b/vendor/lance-io/src/uring/current_thread_future.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Future implementation for thread-local io_uring operations. +//! +//! This future actively processes completions during polling instead of +//! relying on background tasks. + +use super::current_thread::{process_thread_local_completions, submit_and_wait_thread_local}; +use super::requests::IoRequest; +use bytes::Bytes; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +/// Future that awaits completion of a thread-local io_uring read operation +pub struct UringCurrentThreadFuture { + request: Arc, +} + +impl UringCurrentThreadFuture { + pub(super) fn new(request: Arc) -> Self { + Self { request } + } +} + +impl Future for UringCurrentThreadFuture { + type Output = object_store::Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // Check thread safety + if self.request.thread_id != std::thread::current().id() { + panic!("Request thread ID does not match current thread ID"); + } + + // First, check if we've been completed by some other future polling for completions. + let mut state = self.request.state.lock().unwrap(); + + if state.completed { + // Take result and return Ready + match state.err.take() { + Some(err) => { + return Poll::Ready(Err(object_store::Error::Generic { + store: "io_uring_ct", + source: Box::new(err), + })); + } + None => { + let br = state.bytes_read; + state.buffer.truncate(br); + let bytes = std::mem::take(&mut state.buffer).freeze(); + return Poll::Ready(Ok(bytes)); + } + } + } + + drop(state); + + // If not, then we should do any available work and then process completions. + if let Err(e) = submit_and_wait_thread_local() { + log::debug!("Submit and wait error: {:?}", e); + } + + if let Err(e) = process_thread_local_completions() { + log::warn!("Error processing completions: {:?}", e); + } + + // Check if our request completed + let mut state = self.request.state.lock().unwrap(); + + if state.completed { + // Take result and return Ready + match state.err.take() { + Some(err) => { + return Poll::Ready(Err(object_store::Error::Generic { + store: "io_uring_ct", + source: Box::new(err), + })); + } + None => { + let br = state.bytes_read; + state.buffer.truncate(br); + let bytes = std::mem::take(&mut state.buffer).freeze(); + return Poll::Ready(Ok(bytes)); + } + } + } + + // Not done yet - immediately wake and return Pending (don't store waker) + // which will force the future to be polled again. This is intentionally + // a busy loop. io_uring is intended for fast disks where read latency is + // so small that the cost of a true context switch (parking and unparking) + // would be too high. + // + // We are effectively doing a "yield" here while we wait for + // the io_uring thread to complete the request. + drop(state); + cx.waker().wake_by_ref(); + Poll::Pending + } +} diff --git a/vendor/lance-io/src/uring/future.rs b/vendor/lance-io/src/uring/future.rs new file mode 100644 index 000000000..64b389936 --- /dev/null +++ b/vendor/lance-io/src/uring/future.rs @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Future implementation for io_uring read operations. + +use super::requests::IoRequest; +use bytes::Bytes; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +/// Future that awaits completion of an io_uring read operation. +/// +/// This future is woken by the io_uring thread when the operation completes. +pub(super) struct UringReadFuture { + pub(super) request: Arc, +} + +impl Future for UringReadFuture { + type Output = object_store::Result; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut state = self.request.state.lock().unwrap(); + + if state.completed { + // Operation completed - take the result + match state.err.take() { + Some(err) => Poll::Ready(Err(object_store::Error::Generic { + store: "io_uring", + source: Box::new(err), + })), + None => { + let br = state.bytes_read; + state.buffer.truncate(br); + let bytes = std::mem::take(&mut state.buffer).freeze(); + Poll::Ready(Ok(bytes)) + } + } + } else { + // Operation not yet complete - store waker and return Pending + state.waker = Some(cx.waker().clone()); + Poll::Pending + } + } +} diff --git a/vendor/lance-io/src/uring/reader.rs b/vendor/lance-io/src/uring/reader.rs new file mode 100644 index 000000000..2d6bb11c4 --- /dev/null +++ b/vendor/lance-io/src/uring/reader.rs @@ -0,0 +1,301 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! UringReader implementation. + +use super::future::UringReadFuture; +use super::requests::IoRequest; +use super::thread::{SUBMITTED_COUNTER, THREAD_SELECTOR, URING_THREADS}; +use super::{DEFAULT_URING_BLOCK_SIZE, DEFAULT_URING_IO_PARALLELISM, URING_BLOCK_SIZE}; +use crate::local::to_local_path; +use crate::traits::Reader; +use crate::uring::requests::RequestState; +use crate::utils::tracking_store::IOTracker; +use bytes::{Bytes, BytesMut}; +use futures::FutureExt; +use futures::future::BoxFuture; +use lance_core::deepsize::DeepSizeOf; +use lance_core::{Error, Result}; +use object_store::path::Path; +use std::fs::File; +use std::future::Future; +use std::io::{self, ErrorKind}; +use std::ops::Range; +use std::os::unix::io::{AsRawFd, RawFd}; +use std::pin::Pin; +use std::sync::atomic::Ordering; +use std::sync::{Arc, LazyLock, Mutex}; +use std::time::Duration; +use tracing::instrument; + +/// Cache key for UringReader instances. +/// We cache by (path, block_size) because block_size affects reader behavior. +#[derive(Clone, Debug, Hash, Eq, PartialEq)] +pub(super) struct CacheKey { + path: String, + block_size: usize, +} + +impl CacheKey { + pub(super) fn new(path: &Path, block_size: usize) -> Self { + Self { + path: path.to_string(), + block_size, + } + } +} + +/// Data stored in the cache for each opened file. +#[derive(Clone)] +pub(super) struct CachedReaderData { + pub(super) handle: Arc, + pub(super) size: usize, +} + +/// Global cache of open file handles. +/// Entries expire after 60 seconds to ensure files are eventually closed. +pub(super) static HANDLE_CACHE: LazyLock> = + LazyLock::new(|| { + moka::future::Cache::builder() + .time_to_live(Duration::from_secs(60)) + .max_capacity(10_000) + .build() + }); + +/// File handle for io_uring operations. +/// +/// Keeps the file alive and provides the raw file descriptor. +#[derive(Debug)] +pub(super) struct UringFileHandle { + /// The file (kept alive via Arc) + #[allow(unused)] + file: Arc, + + /// Raw file descriptor for io_uring + pub(super) fd: RawFd, + + /// Object store path + pub(super) path: Path, +} + +impl UringFileHandle { + pub(super) fn new(file: File, path: Path) -> Self { + let fd = file.as_raw_fd(); + Self { + file: Arc::new(file), + fd, + path, + } + } +} + +/// io_uring-based reader for local files. +/// +/// This reader uses a dedicated process-wide thread running an io_uring event loop +/// for high-performance asynchronous I/O. +#[derive(Debug)] +pub struct UringReader { + /// File handle + handle: Arc, + + /// Block size for I/O operations + block_size: usize, + + /// File size (determined at open time) + size: usize, + + /// I/O tracker for monitoring operations + io_tracker: Arc, +} + +impl DeepSizeOf for UringReader { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // Skip file handle (just a system resource) + // Only count the path's deep size + self.handle.path.as_ref().deep_size_of_children(context) + } +} + +impl UringReader { + /// Open a file with io_uring. + /// + /// This is the internal constructor used by ObjectStore. + #[instrument(level = "debug")] + pub(crate) async fn open( + path: &Path, + block_size: usize, + known_size: Option, + io_tracker: Arc, + ) -> Result> { + // Determine block size with environment variable override + let block_size = URING_BLOCK_SIZE.unwrap_or(block_size.max(DEFAULT_URING_BLOCK_SIZE)); + + let cache_key = CacheKey::new(path, block_size); + + // Try to get from cache first + if let Some(data) = HANDLE_CACHE.get(&cache_key).await { + // Use known_size if provided, otherwise use cached size + let size = known_size.unwrap_or(data.size); + return Ok(Box::new(Self { + handle: data.handle, + block_size, + size, + io_tracker, + }) as Box); + } + + // Cache miss - open file and get size + let path_clone = path.clone(); + let local_path = to_local_path(path); + + let data = tokio::task::spawn_blocking(move || { + let file = File::open(&local_path).map_err(|e| match e.kind() { + ErrorKind::NotFound => Error::not_found(path_clone.to_string()), + _ => e.into(), + })?; + + // Get size from known_size or file metadata + let size = match known_size { + Some(s) => s, + None => file.metadata()?.len() as usize, + }; + + Ok::<_, Error>(CachedReaderData { + handle: Arc::new(UringFileHandle::new(file, path_clone)), + size, + }) + }) + .await??; + + // Insert into cache + HANDLE_CACHE.insert(cache_key, data.clone()).await; + + // Return new reader instance + Ok(Box::new(Self { + handle: data.handle.clone(), + block_size, + size: data.size, + io_tracker, + }) as Box) + } + + /// Submit a read request to the io_uring thread via channel and return a future. + fn submit_read( + &self, + offset: u64, + length: usize, + ) -> Pin> + Send>> { + let mut buffer = BytesMut::with_capacity(length); + unsafe { + buffer.set_len(length); + } + + // Create IoRequest with all data + let request = Arc::new(IoRequest { + fd: self.handle.fd, + offset, + length, + thread_id: std::thread::current().id(), + state: Mutex::new(RequestState { + completed: false, + waker: None, + err: None, + buffer, + bytes_read: 0, + }), + }); + + // Increment submitted counter before sending to channel + SUBMITTED_COUNTER.fetch_add(1, Ordering::Relaxed); + + // Select thread in round-robin fashion + let thread_idx = + (THREAD_SELECTOR.fetch_add(1, Ordering::Relaxed) as usize) % URING_THREADS.len(); + + // Send to selected thread via channel + match URING_THREADS[thread_idx] + .request_tx + .send(Arc::clone(&request)) + { + Ok(()) => { + // Return future that will be woken when operation completes + Box::pin(UringReadFuture { request }) + } + Err(_) => { + // Thread died - decrement counter and return error future + SUBMITTED_COUNTER.fetch_sub(1, Ordering::Relaxed); + Box::pin(async move { + Err(object_store::Error::Generic { + store: "UringReader", + source: Box::new(io::Error::new( + io::ErrorKind::BrokenPipe, + "io_uring thread died", + )), + }) + }) + } + } + } +} + +impl Reader for UringReader { + fn path(&self) -> &Path { + &self.handle.path + } + + fn block_size(&self) -> usize { + self.block_size + } + + fn io_parallelism(&self) -> usize { + std::env::var("LANCE_URING_IO_PARALLELISM") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_URING_IO_PARALLELISM) + } + + /// Returns the file size. + fn size(&self) -> BoxFuture<'_, object_store::Result> { + Box::pin(async move { Ok(self.size) }) + } + + /// Read a range of bytes using io_uring. + #[instrument(level = "debug", skip(self))] + fn get_range(&self, range: Range) -> BoxFuture<'static, object_store::Result> { + let io_tracker = self.io_tracker.clone(); + let path = self.handle.path.clone(); + let num_bytes = range.len() as u64; + let range_u64 = (range.start as u64)..(range.end as u64); + + let metrics = self.io_tracker.begin_io("get"); + self.submit_read(range.start as u64, range.len()) + .map(move |result| { + metrics.record(&result, num_bytes); + if result.is_ok() { + io_tracker.record_read("get_range", path, num_bytes, Some(range_u64)); + } + result + }) + .boxed() + } + + /// Read the entire file using io_uring. + #[instrument(level = "debug", skip(self))] + fn get_all(&self) -> BoxFuture<'static, object_store::Result> { + let size = self.size; + let io_tracker = self.io_tracker.clone(); + let path = self.handle.path.clone(); + + let metrics = self.io_tracker.begin_io("get"); + self.submit_read(0, size) + .map(move |result| { + let num_bytes = result.as_ref().map_or(0, |bytes| bytes.len() as u64); + metrics.record(&result, num_bytes); + if result.is_ok() { + io_tracker.record_read("get_all", path, num_bytes, None); + } + result + }) + .boxed() + } +} diff --git a/vendor/lance-io/src/uring/requests.rs b/vendor/lance-io/src/uring/requests.rs new file mode 100644 index 000000000..fa257507d --- /dev/null +++ b/vendor/lance-io/src/uring/requests.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Protocol types for communication between UringReader and the io_uring thread. + +use bytes::BytesMut; +use std::io; +use std::os::unix::io::RawFd; +use std::sync::Mutex; +use std::task::Waker; +use std::thread::ThreadId; + +pub(super) struct RequestState { + pub completed: bool, + pub waker: Option, + pub err: Option, + pub buffer: BytesMut, + /// Accumulated bytes read across retries (for handling short reads). + pub bytes_read: usize, +} + +/// I/O request object that contains all state for a single read operation. +/// This is shared between the submitter, uring thread, and future via Arc. +pub(super) struct IoRequest { + /// File descriptor to read from. + pub fd: RawFd, + + /// Byte offset to start reading from. + pub offset: u64, + + /// Number of bytes to read. + pub length: usize, + + pub thread_id: ThreadId, + + /// Completion flag - set to true when operation completes. + pub state: Mutex, +} + +impl IoRequest { + /// Mark this request as failed with the given error. + /// + /// Sets the error, marks completed, and wakes any waiting future. + /// Used when a request cannot be submitted (e.g. SQ full). + pub(super) fn fail(&self, err: io::Error) { + let mut state = self.state.lock().unwrap(); + state.err = Some(err); + state.completed = true; + if let Some(waker) = state.waker.take() { + drop(state); + waker.wake(); + } + } +} diff --git a/vendor/lance-io/src/uring/tests.rs b/vendor/lance-io/src/uring/tests.rs new file mode 100644 index 000000000..19d931da6 --- /dev/null +++ b/vendor/lance-io/src/uring/tests.rs @@ -0,0 +1,392 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Tests for io_uring reader implementation. + +use crate::object_store::ObjectStore; +use lance_core::Result; +use std::io::Write; +use std::time::Duration; +use tempfile::NamedTempFile; + +/// Helper to create a temporary file with test data +fn create_test_file(size: usize) -> Result<(NamedTempFile, Vec)> { + let mut file = NamedTempFile::new()?; + let data: Vec = (0..size).map(|i| (i % 256) as u8).collect(); + file.write_all(&data)?; + file.flush()?; + Ok((file, data)) +} + +#[tokio::test] +async fn test_read_small_file() -> Result<()> { + let (file, expected_data) = create_test_file(1024)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Read entire file + let data = reader.get_all().await.unwrap(); + assert_eq!(data.as_ref(), expected_data.as_slice()); + + Ok(()) +} + +#[tokio::test] +async fn test_read_range() -> Result<()> { + let (file, expected_data) = create_test_file(4096)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Read a range in the middle + let range = 1000..2000; + let data = reader.get_range(range.clone()).await.unwrap(); + assert_eq!(data.as_ref(), &expected_data[range]); + + Ok(()) +} + +#[tokio::test] +async fn test_read_multiple_ranges() -> Result<()> { + let (file, expected_data) = create_test_file(8192)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Read multiple ranges + let ranges = vec![0..100, 500..600, 2000..3000]; + for range in ranges { + let data = reader.get_range(range.clone()).await.unwrap(); + assert_eq!(data.as_ref(), &expected_data[range]); + } + + Ok(()) +} + +#[tokio::test] +async fn test_file_size() -> Result<()> { + let size = 5000; + let (file, _) = create_test_file(size)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + assert_eq!(reader.size().await.unwrap(), size); + + Ok(()) +} + +#[tokio::test] +async fn test_concurrent_reads() -> Result<()> { + let (file, expected_data) = create_test_file(16384)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + + // Perform multiple concurrent reads + let mut tasks = vec![]; + for i in 0..10 { + let reader_clone = store.open(&path).await?; + let expected = expected_data.clone(); + tasks.push(tokio::spawn(async move { + let range = (i * 1000)..((i + 1) * 1000); + let data = reader_clone.get_range(range.clone()).await.unwrap(); + assert_eq!(data.as_ref(), &expected[range]); + })); + } + + // Wait for all tasks + for task in tasks { + task.await.unwrap(); + } + + Ok(()) +} + +#[tokio::test] +async fn test_large_file_read() -> Result<()> { + // Test with a larger file (1MB) + let size = 1024 * 1024; + let (file, expected_data) = create_test_file(size)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Read entire file + let data = reader.get_all().await.unwrap(); + assert_eq!(data.len(), size); + assert_eq!(data.as_ref(), expected_data.as_slice()); + + Ok(()) +} + +#[tokio::test] +async fn test_read_edge_cases() -> Result<()> { + let (file, expected_data) = create_test_file(4096)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Read from start + let data = reader.get_range(0..100).await.unwrap(); + assert_eq!(data.as_ref(), &expected_data[0..100]); + + // Read to end + let data = reader.get_range(4000..4096).await.unwrap(); + assert_eq!(data.as_ref(), &expected_data[4000..4096]); + + // Read single byte + let data = reader.get_range(2000..2001).await.unwrap(); + assert_eq!(data.as_ref(), &expected_data[2000..2001]); + + Ok(()) +} + +#[tokio::test] +async fn test_file_not_found() { + let uri = "file+uring:///nonexistent/file.dat"; + let (store, path) = ObjectStore::from_uri(uri).await.unwrap(); + + // Should fail to open non-existent file + let result = store.open(&path).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_block_size_and_parallelism() -> Result<()> { + let (file, _) = create_test_file(1024)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Check default values (or configured values) + assert!(reader.block_size() > 0); + assert!(reader.io_parallelism() > 0); + + Ok(()) +} + +#[tokio::test] +async fn test_path() -> Result<()> { + let (file, _) = create_test_file(1024)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Verify path is preserved + assert_eq!(reader.path(), &path); + + Ok(()) +} + +/// Test that reading past EOF returns an error. +/// +/// This exercises the case where `known_size` passed to `open_with_size` is larger +/// than the actual file, causing io_uring to hit EOF before the full read completes. +#[tokio::test] +async fn test_short_read_get_all() -> Result<()> { + let actual_size: usize = 8192; + let (file, _expected_data) = create_test_file(actual_size)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + + // Open with inflated known_size — the reader will think the file is 2x its real size + let inflated_size = actual_size * 2; + let reader = store.open_with_size(&path, inflated_size).await?; + + // get_all() will submit a read for inflated_size bytes from an actual_size file. + // The kernel reads actual_size bytes then returns 0 (EOF) — this should be an error. + let result = reader.get_all().await; + assert!(result.is_err(), "reading past EOF should return an error"); + + Ok(()) +} + +/// Test that a range read extending past EOF returns an error. +#[tokio::test] +async fn test_short_read_get_range_past_eof() -> Result<()> { + let actual_size: usize = 8192; + let (file, _expected_data) = create_test_file(actual_size)?; + let file_path = file.path().to_str().unwrap(); + let uri = format!("file+uring://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Request a range that starts inside the file but extends past EOF. + // File is 8192 bytes; reading 4096..16384 hits EOF — this should be an error. + let range_start = 4096; + let range_end = actual_size * 2; // 16384, well past EOF + let result = reader.get_range(range_start..range_end).await; + assert!( + result.is_err(), + "range extending past EOF should return an error" + ); + + Ok(()) +} + +/// Test that when push_to_sq fails (SQ full), the request's future returns +/// an error instead of hanging forever. +/// +/// This directly tests the thread-path scenario: create an IoUring with +/// queue_depth=2, fill the SQ, then try to push a 3rd request. The 3rd +/// request's future should return an error within the timeout. +/// +/// BUG: currently the failed push silently drops the request, so the +/// future hangs and the timeout fires. +#[tokio::test] +async fn test_retry_sq_full_thread() -> Result<()> { + use super::future::UringReadFuture; + use super::requests::{IoRequest, RequestState}; + use super::thread::push_to_sq; + use bytes::BytesMut; + use io_uring::IoUring; + use std::collections::HashMap; + use std::os::unix::io::AsRawFd; + use std::sync::{Arc, Mutex}; + + let (file, _) = create_test_file(4096)?; + let fd = file.as_file().as_raw_fd(); + + // Create a tiny ring with queue_depth=2 + let mut ring = IoUring::new(2).unwrap(); + let mut pending: HashMap> = HashMap::new(); + + // Helper to create a request + let make_request = || { + Arc::new(IoRequest { + fd, + offset: 0, + length: 4096, + thread_id: std::thread::current().id(), + state: Mutex::new(RequestState { + completed: false, + waker: None, + err: None, + buffer: BytesMut::zeroed(4096), + bytes_read: 0, + }), + }) + }; + + // Fill the SQ (capacity=2) + let _r1 = make_request(); + let _r2 = make_request(); + push_to_sq(&mut ring, &mut pending, _r1).unwrap(); + push_to_sq(&mut ring, &mut pending, _r2).unwrap(); + + // 3rd push should fail — SQ is full + let r3 = make_request(); + let push_result = push_to_sq(&mut ring, &mut pending, r3.clone()); + assert!(push_result.is_err(), "3rd push should fail (SQ full)"); + + // r3's future should return an error, not hang forever. + // BUG: currently nobody sets completed=true or err on r3, so the future hangs. + let future = UringReadFuture { request: r3 }; + let result = tokio::time::timeout(Duration::from_secs(2), future).await; + assert!( + result.is_ok(), + "future timed out — request was dropped without error on SQ-full push failure" + ); + + Ok(()) +} + +/// Test that when push_to_sq fails (SQ full) on the current-thread path, +/// the request's future returns an error instead of hanging forever. +/// +/// Uses UringCurrentThreadFuture (which will be a no-op poller since the +/// thread-local URING has no knowledge of this request) after push_to_sq +/// has already completed the request with an error. +#[tokio::test(flavor = "current_thread")] +async fn test_retry_sq_full_current_thread() -> Result<()> { + use super::current_thread_future::UringCurrentThreadFuture; + use super::requests::{IoRequest, RequestState}; + use super::thread::push_to_sq; + use bytes::BytesMut; + use io_uring::IoUring; + use std::collections::HashMap; + use std::os::unix::io::AsRawFd; + use std::sync::{Arc, Mutex}; + + let (file, _) = create_test_file(4096)?; + let fd = file.as_file().as_raw_fd(); + + // Create a tiny ring with queue_depth=2 + let mut ring = IoUring::new(2).unwrap(); + let mut pending: HashMap> = HashMap::new(); + + let make_request = || { + Arc::new(IoRequest { + fd, + offset: 0, + length: 4096, + thread_id: std::thread::current().id(), + state: Mutex::new(RequestState { + completed: false, + waker: None, + err: None, + buffer: BytesMut::zeroed(4096), + bytes_read: 0, + }), + }) + }; + + // Fill the SQ (capacity=2) + push_to_sq(&mut ring, &mut pending, make_request()).unwrap(); + push_to_sq(&mut ring, &mut pending, make_request()).unwrap(); + + // 3rd push should fail — SQ is full + let r3 = make_request(); + let push_result = push_to_sq(&mut ring, &mut pending, r3.clone()); + assert!(push_result.is_err(), "3rd push should fail (SQ full)"); + + // r3's future should return an error, not hang forever. + let future = UringCurrentThreadFuture::new(r3); + let result = tokio::time::timeout(Duration::from_secs(2), future).await; + assert!( + result.is_ok(), + "future timed out — request was dropped without error on SQ-full push failure" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_uring_not_enabled_with_file_scheme() -> Result<()> { + // Verify that files opened with file:// don't use uring + let (file, expected_data) = create_test_file(1024)?; + let file_path = file.path().to_str().unwrap(); + // Use regular file:// scheme, should NOT use uring + let uri = format!("file://{}", file_path); + + let (store, path) = ObjectStore::from_uri(&uri).await?; + let reader = store.open(&path).await?; + + // Should still be able to read, just won't use uring + let data = reader.get_all().await.unwrap(); + assert_eq!(data.as_ref(), expected_data.as_slice()); + + Ok(()) +} diff --git a/vendor/lance-io/src/uring/thread.rs b/vendor/lance-io/src/uring/thread.rs new file mode 100644 index 000000000..d2ef19794 --- /dev/null +++ b/vendor/lance-io/src/uring/thread.rs @@ -0,0 +1,396 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Dedicated thread for io_uring operations. +//! +//! This module provides a background thread that owns an io_uring instance +//! and processes read requests from a channel. Readers send requests via +//! an MPSC channel, and the thread handles submission and completion processing. + +use super::DEFAULT_URING_QUEUE_DEPTH; +use super::requests::IoRequest; +use io_uring::{IoUring, opcode, types}; +use std::collections::HashMap; +use std::io; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::mpsc::{Receiver, RecvTimeoutError, SyncSender, sync_channel}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +/// Handle to the io_uring background thread. +/// +/// This provides a channel sender for submitting read requests to the thread. +pub(super) struct UringThreadHandle { + pub request_tx: SyncSender>, +} + +/// Lazy-initialized io_uring thread pool. +/// +/// Multiple threads are spawned on first access and run until process exit. +pub(super) static URING_THREADS: LazyLock> = LazyLock::new(|| { + let queue_depth = get_queue_depth(); + let thread_count = get_thread_count(); + + let mut threads = Vec::with_capacity(thread_count); + + for i in 0..thread_count { + let (tx, rx) = sync_channel(queue_depth); + + std::thread::Builder::new() + .name(format!("lance-uring-{}", i)) + .spawn(move || run_uring_thread(rx, queue_depth, i)) + .expect("Failed to spawn io_uring thread"); + + threads.push(UringThreadHandle { request_tx: tx }); + } + + log::info!( + "io_uring thread pool spawned ({} threads, queue_depth={})", + thread_count, + queue_depth + ); + + threads +}); + +/// Atomic counter for round-robin thread selection. +pub(super) static THREAD_SELECTOR: AtomicU64 = AtomicU64::new(0); + +/// Counter for generating unique user_data values. +/// +/// Each io_uring operation needs a unique user_data ID to match completions +/// with their corresponding requests. +static USER_DATA_COUNTER: AtomicU64 = AtomicU64::new(1); + +/// Counter for requests that have been submitted to the thread but not yet received. +/// +/// This tracks requests sitting in the channel queue waiting to be received by the thread. +pub(super) static SUBMITTED_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// Default batch size for submission - how many requests to batch before calling submit(). +const DEFAULT_SUBMIT_BATCH_SIZE: usize = 128; + +/// Default number of io_uring threads. +const DEFAULT_URING_THREAD_COUNT: usize = 2; + +/// Get the configured queue depth from environment variable. +fn get_queue_depth() -> usize { + std::env::var("LANCE_URING_QUEUE_DEPTH") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_URING_QUEUE_DEPTH) +} + +/// Get the configured poll timeout from environment variable. +fn get_poll_timeout() -> Duration { + let timeout_ms = std::env::var("LANCE_URING_POLL_TIMEOUT_MS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(10); + + Duration::from_millis(timeout_ms) +} + +/// Get the configured submit batch size from environment variable. +fn get_submit_batch_size() -> usize { + std::env::var("LANCE_URING_SUBMIT_BATCH_SIZE") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_SUBMIT_BATCH_SIZE) +} + +/// Get the configured number of uring threads from environment variable. +fn get_thread_count() -> usize { + std::env::var("LANCE_URING_THREAD_COUNT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_URING_THREAD_COUNT) +} + +/// Main loop for the io_uring thread. +/// +/// This thread: +/// 1. Receives requests from the channel +/// 2. Submits them to io_uring +/// 3. Processes completions +/// 4. Wakes futures via their wakers +fn run_uring_thread(request_rx: Receiver>, queue_depth: usize, thread_id: usize) { + // Create local io_uring instance + let mut ring = IoUring::builder() + // .setup_sqpoll(100) + .build(queue_depth as u32) + .expect("Failed to create io_uring"); + + let mut pending: HashMap> = HashMap::with_capacity(queue_depth); + let poll_timeout = get_poll_timeout(); + let submit_batch_size = get_submit_batch_size(); + let mut last_log = Instant::now(); + let log_interval = Duration::from_millis(100); + let mut completed_iops = 0usize; + let mut completed_sectors = 0usize; + let mut min_in_flight = usize::MAX; + + loop { + // Track minimum in-flight count + let in_flight = pending.len(); + min_in_flight = min_in_flight.min(in_flight); + + // Log in-flight requests every 100ms + let now = Instant::now(); + if now.duration_since(last_log) >= log_interval { + let submitted = SUBMITTED_COUNTER.load(Ordering::Relaxed); + log::info!( + "io_uring[{}]: {} submitted, {} in flight (min {}), {} iops completed, {} sectors completed", + thread_id, + submitted, + in_flight, + min_in_flight, + completed_iops, + completed_sectors + ); + last_log = now; + completed_iops = 0; // Reset counter after logging + completed_sectors = 0; // Reset counter after logging + min_in_flight = usize::MAX; // Reset min tracker + } + + // Process all available completions first + let mut needs_submit = false; + let completions = process_completions(&mut ring, &mut pending); + match completions { + Ok(result) => { + completed_iops += result.iops; + completed_sectors += result.sectors; + + // Resubmit any short-read retries + for request in result.retries { + if let Err(e) = push_to_sq(&mut ring, &mut pending, request) { + log::error!("Failed to resubmit short read: {}", e); + } else { + needs_submit = true; + } + } + } + Err(e) => { + log::error!("Error processing io_uring completions: {}", e); + } + } + + min_in_flight = min_in_flight.min(pending.len()); + + // Batch submit requests - keep pulling from channel and pushing to SQ + // until we hit batch size or channel is empty + let mut batch_count = 0; + loop { + // Try to receive new request + // Use recv_timeout only when pending is empty, otherwise use try_recv + let recv_result = if pending.is_empty() && batch_count == 0 { + // No operations in flight and no batch started - we can afford to wait with timeout + request_rx.recv_timeout(poll_timeout).map_err(|e| match e { + RecvTimeoutError::Timeout => std::sync::mpsc::TryRecvError::Empty, + RecvTimeoutError::Disconnected => std::sync::mpsc::TryRecvError::Disconnected, + }) + } else { + // Operations in flight or batch in progress - busy loop with try_recv + request_rx.try_recv() + }; + + match recv_result { + Ok(request) => { + // Decrement submitted counter when we receive the request from channel + SUBMITTED_COUNTER.fetch_sub(1, Ordering::Relaxed); + + // Push to submission queue (but don't submit yet) + if let Err(e) = push_to_sq(&mut ring, &mut pending, request) { + log::error!("Failed to push to io_uring SQ: {}", e); + } else { + batch_count += 1; + } + + // Break if we've hit the batch size limit + if batch_count >= submit_batch_size { + break; + } + } + Err(std::sync::mpsc::TryRecvError::Empty) => { + // No more requests in channel - break to submit the batch + break; + } + Err(std::sync::mpsc::TryRecvError::Disconnected) => { + // All senders dropped - submit batch and shutdown + if batch_count > 0 + && let Err(e) = ring.submit() + { + log::error!( + "io_uring[{}]: Failed to submit io_uring batch: {}", + thread_id, + e + ); + } + log::info!( + "io_uring thread {} shutting down (channel disconnected)", + thread_id + ); + return; + } + } + } + + // Submit if we have any requests (from channel or retries) + if (batch_count > 0 || needs_submit) + && let Err(e) = ring.submit() + { + log::error!( + "Failed to submit io_uring batch of {} requests: {}", + batch_count, + e + ); + } + } +} + +/// Push a read request to the io_uring submission queue (without submitting). +/// +/// This generates a unique user_data ID, prepares the read operation, +/// and pushes it to the SQ. The caller is responsible for calling ring.submit(). +pub(super) fn push_to_sq( + ring: &mut IoUring, + pending: &mut HashMap>, + request: Arc, +) -> io::Result<()> { + // Generate unique user_data + let user_data = USER_DATA_COUNTER.fetch_add(1, Ordering::Relaxed); + + // Get buffer pointer, adjusting for any bytes already read (short read retry) + let (buffer_ptr, read_offset, read_length) = { + let state = request.state.lock().unwrap(); + let br = state.bytes_read; + ( + unsafe { state.buffer.as_ptr().add(br) as *mut u8 }, + request.offset + br as u64, + (request.length - br) as u32, + ) + }; + + // Prepare read operation + let read_op = + opcode::Read::new(types::Fd(request.fd), buffer_ptr, read_length).offset(read_offset); + + // Get submission queue + let mut sq = ring.submission(); + + // Check if SQ has space + if sq.is_full() { + drop(sq); + request.fail(io::Error::new( + io::ErrorKind::WouldBlock, + "io_uring submission queue full", + )); + return Err(io::Error::new( + io::ErrorKind::WouldBlock, + "io_uring submission queue full", + )); + } + + // Push to SQ + unsafe { + if sq.push(&read_op.build().user_data(user_data)).is_err() { + drop(sq); + request.fail(io::Error::other("Failed to push to SQ")); + return Err(io::Error::other("Failed to push to SQ")); + } + } + drop(sq); + + // Track request in pending map + pending.insert(user_data, request); + + Ok(()) +} + +struct CompletionResult { + iops: usize, + sectors: usize, + retries: Vec>, +} + +/// Process all available completions from the io_uring. +/// +/// This iterates through the completion queue, matches completions to requests, +/// updates their state, and wakes any waiting futures. Short reads are collected +/// into `retries` for resubmission; EOF before a full read is an error. +/// +/// Returns completion stats and a list of requests needing resubmission. +fn process_completions( + ring: &mut IoUring, + pending: &mut HashMap>, +) -> io::Result { + let mut iops = 0; + let mut sectors = 0; + let mut retries = Vec::new(); + + // Process all available completions + for cqe in ring.completion() { + let user_data = cqe.user_data(); + let result = cqe.result(); + + // Look up request + if let Some(request) = pending.remove(&user_data) { + let mut state = request.state.lock().unwrap(); + + if result < 0 { + // Kernel error + state.err = Some(io::Error::from_raw_os_error(-result)); + state.completed = true; + } else if result == 0 { + // EOF before full read completed + let br = state.bytes_read; + state.err = Some(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("unexpected EOF: read {} of {} bytes", br, request.length), + )); + state.buffer.truncate(br); + state.completed = true; + } else { + // Positive result: n bytes read + let n = result as usize; + state.bytes_read += n; + let br = state.bytes_read; + + if br >= request.length { + // Full read complete + state.buffer.truncate(br); + state.completed = true; + + if request.length > 0 { + let first_sector = request.offset / 4096; + let last_sector = (request.offset + request.length as u64 - 1) / 4096; + let num_sectors = (last_sector - first_sector + 1) as usize; + sectors += num_sectors; + } + } else { + // Short read — need retry; don't mark completed or wake + drop(state); + retries.push(request); + continue; + } + } + + // Wake the future if it's waiting + if let Some(waker) = state.waker.take() { + drop(state); // Release lock before waking + waker.wake(); + } + + iops += 1; + } else { + log::warn!("Received completion for unknown user_data: {}", user_data); + } + } + + Ok(CompletionResult { + iops, + sectors, + retries, + }) +} diff --git a/vendor/lance-io/src/utils.rs b/vendor/lance-io/src/utils.rs new file mode 100644 index 000000000..4b471d8b0 --- /dev/null +++ b/vendor/lance-io/src/utils.rs @@ -0,0 +1,291 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::{cmp::min, num::NonZero, sync::atomic::AtomicU64}; + +use byteorder::{ByteOrder, LittleEndian}; +use bytes::Bytes; +use lance_core::deepsize::DeepSizeOf; +use prost::Message; +use serde::{Deserialize, Serialize}; + +use crate::traits::{ProtoStruct, Reader}; +use lance_core::{Error, Result}; + +pub mod tracking_store; + +/// Read a protobuf message at file position 'pos'. +/// +/// We write protobuf by first writing the length of the message as a u32, +/// followed by the message itself. +pub async fn read_message(reader: &dyn Reader, pos: usize) -> Result { + let file_size = reader.size().await?; + // A message is a u32 length prefix followed by its body; both must lie before + // the end. A `pos` too close to the end means the reader size is too small + // (e.g. a stale cached size). Reject it rather than slice a short buffer and + // panic. + if pos + 4 > file_size { + return Err(Error::io("file size is too small".to_string())); + } + + let range = pos..min(pos + reader.block_size(), file_size); + let buf = reader.get_range(range.clone()).await?; + let msg_len = LittleEndian::read_u32(&buf) as usize; + + if msg_len + 4 > buf.len() { + let remaining_range = range.end..min(4 + pos + msg_len, file_size); + let remaining_bytes = reader.get_range(remaining_range).await?; + let buf = [buf, remaining_bytes].concat(); + if buf.len() < msg_len + 4 { + return Err(Error::io("file size is too small".to_string())); + } + Ok(M::decode(&buf[4..4 + msg_len])?) + } else { + Ok(M::decode(&buf[4..4 + msg_len])?) + } +} + +/// Read a Protobuf-backed struct at file position: `pos`. +// TODO: pub(crate) +pub async fn read_struct< + M: Message + Default + 'static, + T: ProtoStruct + TryFrom, +>( + reader: &dyn Reader, + pos: usize, +) -> Result { + let msg = read_message::(reader, pos).await?; + T::try_from(msg) +} + +pub async fn read_last_block(reader: &dyn Reader) -> object_store::Result { + let file_size = reader.size().await?; + let block_size = reader.block_size(); + let begin = file_size.saturating_sub(block_size); + reader.get_range(begin..file_size).await +} + +pub fn read_metadata_offset(bytes: &Bytes) -> Result { + let len = bytes.len(); + if len < 16 { + return Err(Error::io(format!( + "does not have sufficient data, len: {}, bytes: {:?}", + len, bytes + ))); + } + let offset_bytes = bytes.slice(len - 16..len - 8); + Ok(LittleEndian::read_u64(offset_bytes.as_ref()) as usize) +} + +/// Read the version from the footer bytes +pub fn read_version(bytes: &Bytes) -> Result<(u16, u16)> { + let len = bytes.len(); + if len < 8 { + return Err(Error::io(format!( + "does not have sufficient data, len: {}, bytes: {:?}", + len, bytes + ))); + } + + let major_version = LittleEndian::read_u16(bytes.slice(len - 8..len - 6).as_ref()); + let minor_version = LittleEndian::read_u16(bytes.slice(len - 6..len - 4).as_ref()); + Ok((major_version, minor_version)) +} + +/// Read protobuf from a buffer. +pub fn read_message_from_buf(buf: &Bytes) -> Result { + let msg_len = LittleEndian::read_u32(buf) as usize; + Ok(M::decode(&buf[4..4 + msg_len])?) +} + +/// Read a Protobuf-backed struct from a buffer. +pub fn read_struct_from_buf< + M: Message + Default, + T: ProtoStruct + TryFrom, +>( + buf: &Bytes, +) -> Result { + let msg: M = read_message_from_buf(buf)?; + T::try_from(msg) +} + +/// A cached file size. +/// +/// This wraps an atomic u64 to allow setting the cached file size without +/// needed a mutable reference. +/// +/// Zero is interpreted as unknown. +#[derive(Debug, DeepSizeOf)] +pub struct CachedFileSize(AtomicU64); + +impl<'de> Deserialize<'de> for CachedFileSize { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + let size = Option::::deserialize(deserializer)?.unwrap_or(0); + Ok(Self::new(size)) + } +} + +impl Serialize for CachedFileSize { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let size = self.0.load(std::sync::atomic::Ordering::Relaxed); + if size == 0 { + serializer.serialize_none() + } else { + serializer.serialize_u64(size) + } + } +} + +impl From>> for CachedFileSize { + fn from(size: Option>) -> Self { + match size { + Some(size) => Self(AtomicU64::new(size.into())), + None => Self(AtomicU64::new(0)), + } + } +} + +impl Default for CachedFileSize { + fn default() -> Self { + Self(AtomicU64::new(0)) + } +} + +impl Clone for CachedFileSize { + fn clone(&self) -> Self { + Self(AtomicU64::new( + self.0.load(std::sync::atomic::Ordering::Relaxed), + )) + } +} + +impl PartialEq for CachedFileSize { + fn eq(&self, other: &Self) -> bool { + self.0.load(std::sync::atomic::Ordering::Relaxed) + == other.0.load(std::sync::atomic::Ordering::Relaxed) + } +} + +impl Eq for CachedFileSize {} + +impl CachedFileSize { + /// Create a `CachedFileSize` from a raw byte count. + /// + /// Passing `0` is equivalent to calling [`unknown`](Self::unknown): the + /// type interprets zero as "size not yet known". + pub fn new(size: u64) -> Self { + Self(AtomicU64::new(size)) + } + + pub fn unknown() -> Self { + Self(AtomicU64::new(0)) + } + + pub fn get(&self) -> Option> { + NonZero::new(self.0.load(std::sync::atomic::Ordering::Relaxed)) + } + + pub fn set(&self, size: NonZero) { + self.0 + .store(size.into(), std::sync::atomic::Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + use object_store::path::Path; + + use crate::{ + Error, Result, + object_reader::CloudObjectReader, + object_store::{DEFAULT_DOWNLOAD_RETRY_COUNT, ObjectStore}, + object_writer::ObjectWriter, + traits::{ProtoStruct, WriteExt, Writer}, + utils::read_struct, + }; + + // Bytes is a prost::Message, since we don't have any .proto files in this crate we + // can use it to simulate a real message object. + #[derive(Debug, PartialEq)] + struct BytesWrapper(Bytes); + + impl ProtoStruct for BytesWrapper { + type Proto = Bytes; + } + + impl From<&BytesWrapper> for Bytes { + fn from(value: &BytesWrapper) -> Self { + value.0.clone() + } + } + + impl TryFrom for BytesWrapper { + type Error = Error; + fn try_from(value: Bytes) -> Result { + Ok(Self(value)) + } + } + + #[tokio::test] + async fn test_write_proto_structs() { + let store = ObjectStore::memory(); + let path = Path::from("/foo"); + + let mut object_writer = ObjectWriter::new(&store, &path).await.unwrap(); + assert_eq!(object_writer.tell().await.unwrap(), 0); + + let some_message = BytesWrapper(Bytes::from(vec![10, 20, 30])); + + let pos = object_writer.write_struct(&some_message).await.unwrap(); + assert_eq!(pos, 0); + object_writer.shutdown().await.unwrap(); + + let object_reader = + CloudObjectReader::new(store.inner, path, 1024, None, DEFAULT_DOWNLOAD_RETRY_COUNT) + .unwrap(); + let actual: BytesWrapper = read_struct(&object_reader, pos).await.unwrap(); + assert_eq!(some_message, actual); + } + + #[tokio::test] + async fn test_copy_reader_to_writer() { + let store = ObjectStore::memory(); + let src = Path::from("/src"); + let dst = Path::from("/dst"); + store.put(&src, b"abcdef").await.unwrap(); + + let reader = store.open(&src).await.unwrap(); + let mut writer = store.create(&dst).await.unwrap(); + let copied = writer.copy_from_reader(reader.as_ref()).await.unwrap(); + writer.shutdown().await.unwrap(); + + assert_eq!(copied, 6); + assert_eq!(store.read_one_all(&dst).await.unwrap().as_ref(), b"abcdef"); + } + + #[tokio::test] + async fn test_copy_reader_range_to_writer() { + let store = ObjectStore::memory(); + let src = Path::from("/src-range"); + let dst = Path::from("/dst-range"); + store.put(&src, b"abcdef").await.unwrap(); + + let reader = store.open(&src).await.unwrap(); + let mut writer = store.create(&dst).await.unwrap(); + let copied = writer + .copy_range_from_reader(reader.as_ref(), 2..5) + .await + .unwrap(); + writer.shutdown().await.unwrap(); + + assert_eq!(copied, 3); + assert_eq!(store.read_one_all(&dst).await.unwrap().as_ref(), b"cde"); + } +} diff --git a/vendor/lance-io/src/utils/tracking_store.rs b/vendor/lance-io/src/utils/tracking_store.rs new file mode 100644 index 000000000..703061920 --- /dev/null +++ b/vendor/lance-io/src/utils/tracking_store.rs @@ -0,0 +1,583 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Make assertions about IO operations to an [ObjectStore]. +//! +//! When testing code that performs IO, you will often want to make assertions +//! about the number of reads and writes performed, the amount of data read or +//! written, and the number of disjoint periods where at least one IO is in-flight. +//! +//! This modules provides [`IOTracker`] which can be used to wrap any object store. +use std::fmt::{Display, Formatter}; +use std::ops::Range; +#[cfg(feature = "test-util")] +use std::sync::atomic::AtomicU16; +use std::sync::{Arc, Mutex}; +#[cfg(feature = "metrics")] +use std::time::Instant; + +use bytes::Bytes; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream::BoxStream; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + Result as OSResult, UploadPart, +}; + +use crate::object_store::WrappingObjectStore; +#[cfg(feature = "metrics")] +use crate::object_store::metrics::{InFlightGuard, record_outcome}; + +#[derive(Debug, Default, Clone)] +pub struct IOTracker { + stats: Arc>, + /// The `base` label for the object store metrics published by IO that + /// bypasses the `object_store` layer (see [`Self::begin_io`]). `None` when + /// the IO cannot be attributed to a store, in which case no metrics are + /// published. + #[cfg(feature = "metrics")] + metrics_base: Option>, +} + +impl IOTracker { + /// Get IO statistics and reset the counters (incremental pattern). + /// + /// This returns the accumulated statistics since the last call and resets + /// the internal counters to zero. + pub fn incremental_stats(&self) -> IoStats { + std::mem::take(&mut *self.stats.lock().unwrap()) + } + + /// Get a snapshot of current IO statistics without resetting counters. + /// + /// This returns a clone of the current statistics without modifying the + /// internal state. Use this when you need to check stats without resetting. + pub fn stats(&self) -> IoStats { + self.stats.lock().unwrap().clone() + } + + /// Record a read operation for tracking. + /// + /// This is used by readers that bypass the ObjectStore layer (like LocalObjectReader) + /// to ensure their IO operations are still tracked. + pub fn record_read( + &self, + #[allow(unused_variables)] method: &'static str, + #[allow(unused_variables)] path: Path, + num_bytes: u64, + #[allow(unused_variables)] range: Option>, + ) { + let mut stats = self.stats.lock().unwrap(); + stats.read_iops += 1; + stats.read_bytes += num_bytes; + #[cfg(feature = "test-util")] + stats.requests.push(IoRequestRecord { + method, + path, + range, + }); + } + + /// Record a write operation for tracking. + /// + /// This is used by writers that bypass the ObjectStore layer (like LocalWriter) + /// to ensure their IO operations are still tracked. + pub fn record_write( + &self, + #[allow(unused_variables)] method: &'static str, + #[allow(unused_variables)] path: Path, + num_bytes: u64, + ) { + let mut stats = self.stats.lock().unwrap(); + stats.write_iops += 1; + stats.written_bytes += num_bytes; + #[cfg(feature = "test-util")] + stats.requests.push(IoRequestRecord { + method, + path, + range: None, + }); + } + + /// Label the metrics published through [`Self::begin_io`] with the prefix of + /// the store this tracker belongs to, so IO that bypasses the `object_store` + /// layer carries the same `base` label as the store's metered operations. + /// + /// Only `meter_store` should call this, so that labelling the tracker and + /// wrapping the store stay inseparable — see the rationale there. + #[cfg(feature = "metrics")] + pub(crate) fn set_metrics_base(&mut self, base: &str) { + self.metrics_base = Some(base.into()); + } + + /// Begin an operation that talks to storage without going through the + /// `object_store` layer, and so is invisible to the `MeteredObjectStore` + /// wrapper: the optimized local reads and writes go straight to the + /// filesystem. `operation` must be one of the labels that wrapper uses + /// (`get`, `put`, `head`, ...) so this IO aggregates with the rest. + /// + /// The returned guard keeps the in-flight gauge raised until it is dropped. + #[cfg(feature = "metrics")] + pub fn begin_io(&self, operation: &'static str) -> IoMetricsGuard { + IoMetricsGuard { + state: self.metrics_base.as_ref().map(|base| IoMetricsState { + _in_flight: InFlightGuard::new(base, operation), + base: base.clone(), + operation, + start: Instant::now(), + }), + } + } + + /// Without the `metrics` feature there is nothing to publish. + #[cfg(not(feature = "metrics"))] + pub fn begin_io(&self, _operation: &'static str) -> IoMetricsGuard { + IoMetricsGuard {} + } +} + +/// Publishes the object store metrics for a single operation that bypassed the +/// `object_store` layer (see [`IOTracker::begin_io`]). +/// +/// The operation is only counted by [`Self::record`]; one dropped before that — +/// a cancelled read, an abandoned write — counts as neither a success nor a +/// failure, and only lowers the in-flight gauge. +#[must_use = "the operation is not recorded until `record` is called"] +pub struct IoMetricsGuard { + #[cfg(feature = "metrics")] + state: Option, +} + +#[cfg(feature = "metrics")] +struct IoMetricsState { + base: Arc, + operation: &'static str, + start: Instant, + /// Lowers the in-flight gauge when the guard is dropped. + _in_flight: InFlightGuard, +} + +impl IoMetricsGuard { + /// Record the operation's count and latency, along with `num_bytes` + /// transferred if `result` is `Ok` or an error if it is not. + pub fn record(self, result: &std::result::Result, num_bytes: u64) { + #[cfg(feature = "metrics")] + if let Some(state) = self.state { + record_outcome( + &state.base, + state.operation, + state.start, + num_bytes, + result.is_err(), + ); + } + #[cfg(not(feature = "metrics"))] + let _ = (result, num_bytes); + } +} + +impl WrappingObjectStore for IOTracker { + fn wrap(&self, _store_prefix: &str, target: Arc) -> Arc { + Arc::new(IoTrackingStore::new(target, self.stats.clone())) + } +} + +#[derive(Debug, Default, Clone)] +pub struct IoStats { + pub read_iops: u64, + pub read_bytes: u64, + pub write_iops: u64, + pub written_bytes: u64, + // This is only really meaningful in tests where there isn't any concurrent IO. + #[cfg(feature = "test-util")] + /// Number of disjoint periods where at least one IO is in-flight. + pub num_stages: u64, + #[cfg(feature = "test-util")] + pub requests: Vec, +} + +/// Assertions on IO statistics. +/// assert_io_eq!(io_stats, read_iops, 1); +/// assert_io_eq!(io_stats, write_iops, 0, "should be no writes"); +/// assert_io_eq!(io_stats, num_hops, 1, "should be just {}", "one hop"); +#[cfg(feature = "test-util")] +#[macro_export] +macro_rules! assert_io_eq { + ($io_stats:expr, $field:ident, $expected:expr) => { + assert_eq!( + $io_stats.$field, $expected, + "Expected {} to be {}, got {}. Requests: {:#?}", + stringify!($field), + $expected, + $io_stats.$field, + $io_stats.requests + ); + }; + ($io_stats:expr, $field:ident, $expected:expr, $($arg:tt)+) => { + assert_eq!( + $io_stats.$field, $expected, + "Expected {} to be {}, got {}. Requests: {:#?} {}", + stringify!($field), + $expected, + $io_stats.$field, + $io_stats.requests, + format_args!($($arg)+) + ); + }; +} + +#[cfg(feature = "test-util")] +#[macro_export] +macro_rules! assert_io_gt { + ($io_stats:expr, $field:ident, $expected:expr) => { + assert!( + $io_stats.$field > $expected, + "Expected {} to be > {}, got {}. Requests: {:#?}", + stringify!($field), + $expected, + $io_stats.$field, + $io_stats.requests + ); + }; + ($io_stats:expr, $field:ident, $expected:expr, $($arg:tt)+) => { + assert!( + $io_stats.$field > $expected, + "Expected {} to be > {}, got {}. Requests: {:#?} {}", + stringify!($field), + $expected, + $io_stats.$field, + $io_stats.requests, + format_args!($($arg)+) + ); + }; +} + +#[cfg(feature = "test-util")] +#[macro_export] +macro_rules! assert_io_lt { + ($io_stats:expr, $field:ident, $expected:expr) => { + assert!( + $io_stats.$field < $expected, + "Expected {} to be < {}, got {}. Requests: {:#?}", + stringify!($field), + $expected, + $io_stats.$field, + $io_stats.requests + ); + }; + ($io_stats:expr, $field:ident, $expected:expr, $($arg:tt)+) => { + assert!( + $io_stats.$field < $expected, + "Expected {} to be < {}, got {}. Requests: {:#?} {}", + stringify!($field), + $expected, + $io_stats.$field, + $io_stats.requests, + format_args!($($arg)+) + ); + }; +} + +// These request records only exist for test-only diagnostics. +#[cfg(feature = "test-util")] +#[derive(Clone)] +pub struct IoRequestRecord { + pub method: &'static str, + pub path: Path, + pub range: Option>, +} + +#[cfg(feature = "test-util")] +impl std::fmt::Debug for IoRequestRecord { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + // For example: "put /path/to/file range: 0-100" + write!( + f, + "IORequest(method={}, path=\"{}\"", + self.method, self.path + )?; + if let Some(range) = &self.range { + write!(f, ", range={:?}", range)?; + } + write!(f, ")")?; + Ok(()) + } +} + +impl Display for IoStats { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{:#?}", self) + } +} + +#[derive(Debug)] +pub struct IoTrackingStore { + target: Arc, + stats: Arc>, + #[cfg(feature = "test-util")] + active_requests: Arc, +} + +impl Display for IoTrackingStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{:#?}", self) + } +} + +impl IoTrackingStore { + pub fn new(target: Arc, stats: Arc>) -> Self { + Self { + target, + stats, + #[cfg(feature = "test-util")] + active_requests: Arc::new(AtomicU16::new(0)), + } + } + + fn record_read( + &self, + method: &'static str, + path: Path, + num_bytes: u64, + range: Option>, + ) { + let mut stats = self.stats.lock().unwrap(); + stats.read_iops += 1; + stats.read_bytes += num_bytes; + #[cfg(feature = "test-util")] + stats.requests.push(IoRequestRecord { + method, + path, + range, + }); + #[cfg(not(feature = "test-util"))] + let _ = (method, path, range); // Suppress unused variable warnings + } + + fn record_write(&self, method: &'static str, path: Path, num_bytes: u64) { + let mut stats = self.stats.lock().unwrap(); + stats.write_iops += 1; + stats.written_bytes += num_bytes; + #[cfg(feature = "test-util")] + stats.requests.push(IoRequestRecord { + method, + path, + range: None, + }); + #[cfg(not(feature = "test-util"))] + let _ = (method, path); // Suppress unused variable warnings + } + + #[cfg(feature = "test-util")] + fn stage_guard(&self) -> StageGuard { + StageGuard::new(self.active_requests.clone(), self.stats.clone()) + } + + #[cfg(not(feature = "test-util"))] + fn stage_guard(&self) -> StageGuard { + StageGuard + } +} + +#[async_trait::async_trait] +#[deny(clippy::missing_trait_methods)] +impl ObjectStore for IoTrackingStore { + async fn put_opts( + &self, + location: &Path, + bytes: PutPayload, + opts: PutOptions, + ) -> OSResult { + let _guard = self.stage_guard(); + self.record_write( + "put_opts", + location.to_owned(), + bytes.content_length() as u64, + ); + self.target.put_opts(location, bytes, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OSResult> { + let _guard = self.stage_guard(); + let target = self.target.put_multipart_opts(location, opts).await?; + Ok(Box::new(IoTrackingMultipartUpload { + target, + stats: self.stats.clone(), + #[cfg(feature = "test-util")] + path: location.to_owned(), + #[cfg(feature = "test-util")] + _guard, + })) + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { + let _guard = self.stage_guard(); + let range = match &options.range { + Some(GetRange::Bounded(range)) => Some(range.clone()), + _ => None, // TODO: fill in other options. + }; + let result = self.target.get_opts(location, options).await; + if let Ok(result) = &result { + let num_bytes = result.range.end - result.range.start; + + self.record_read("get_opts", location.to_owned(), num_bytes, range); + } + result + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OSResult> { + let _guard = self.stage_guard(); + let result = self.target.get_ranges(location, ranges).await; + if let Ok(result) = &result { + self.record_read( + "get_ranges", + location.to_owned(), + result.iter().map(|b| b.len() as u64).sum(), + None, + ); + } + result + } + + fn delete_stream( + &self, + locations: BoxStream<'static, OSResult>, + ) -> BoxStream<'static, OSResult> { + let stats = Arc::clone(&self.stats); + let tracked = locations + .map_ok(move |path| { + let mut stats = stats.lock().unwrap(); + stats.write_iops += 1; + #[cfg(feature = "test-util")] + stats.requests.push(IoRequestRecord { + method: "delete", + path: path.clone(), + range: None, + }); + path + }) + .boxed(); + self.target.delete_stream(tracked) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { + let _guard = self.stage_guard(); + self.record_read("list", prefix.cloned().unwrap_or_default(), 0, None); + self.target.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OSResult> { + self.record_read( + "list_with_offset", + prefix.cloned().unwrap_or_default(), + 0, + None, + ); + self.target.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { + let _guard = self.stage_guard(); + self.record_read( + "list_with_delimiter", + prefix.cloned().unwrap_or_default(), + 0, + None, + ); + self.target.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> { + let _guard = self.stage_guard(); + self.record_write("copy", from.to_owned(), 0); + self.target.copy_opts(from, to, opts).await + } + + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> OSResult<()> { + let _guard = self.stage_guard(); + self.record_write("rename", from.to_owned(), 0); + self.target.rename_opts(from, to, opts).await + } +} + +#[derive(Debug)] +struct IoTrackingMultipartUpload { + target: Box, + #[cfg(feature = "test-util")] + path: Path, + stats: Arc>, + #[cfg(feature = "test-util")] + _guard: StageGuard, +} + +#[async_trait::async_trait] +impl MultipartUpload for IoTrackingMultipartUpload { + async fn abort(&mut self) -> OSResult<()> { + self.target.abort().await + } + + async fn complete(&mut self) -> OSResult { + self.target.complete().await + } + + fn put_part(&mut self, payload: PutPayload) -> UploadPart { + { + let mut stats = self.stats.lock().unwrap(); + stats.write_iops += 1; + stats.written_bytes += payload.content_length() as u64; + #[cfg(feature = "test-util")] + stats.requests.push(IoRequestRecord { + method: "put_part", + path: self.path.to_owned(), + range: None, + }); + } + self.target.put_part(payload) + } +} + +#[cfg(feature = "test-util")] +#[derive(Debug)] +struct StageGuard { + active_requests: Arc, + stats: Arc>, +} + +#[cfg(not(feature = "test-util"))] +struct StageGuard; + +#[cfg(feature = "test-util")] +impl StageGuard { + fn new(active_requests: Arc, stats: Arc>) -> Self { + active_requests.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Self { + active_requests, + stats, + } + } +} + +#[cfg(feature = "test-util")] +impl Drop for StageGuard { + fn drop(&mut self) { + if self + .active_requests + .fetch_sub(1, std::sync::atomic::Ordering::SeqCst) + == 1 + { + let mut stats = self.stats.lock().unwrap(); + stats.num_stages += 1; + } + } +}