diff --git a/Cargo.lock b/Cargo.lock index 6cf44e4..2f6fbd4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -444,6 +444,12 @@ dependencies = [ "slab", ] +[[package]] +name = "asyncband" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e52766975a4f080528a898235c51e82e65df9db713419067b5368040eeb5659" + [[package]] name = "atoi" version = "2.0.0" @@ -1293,6 +1299,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "countio" version = "0.3.0" @@ -2148,6 +2165,12 @@ dependencies = [ "url", ] +[[package]] +name = "datasketches" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46c4cf71a36b46dcfc00e5014c0c20ccad2b1b6a008304d7d57d2749b2d41b3d" + [[package]] name = "defmt" version = "1.1.1" @@ -2366,6 +2389,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8eb564c5c7423d25c886fb561d1e4ee69f72354d16918afa32c08811f6b6a55" +[[package]] +name = "fastant" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e825441bfb2d831c47c97d05821552db8832479f44c571b97fededbf0099c07" +dependencies = [ + "small_ctor", + "web-time", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -2437,12 +2470,133 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "foyer" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f911e6f0b4909f23d65a95c5d27bcf2f92855b8a0f629b47b743d53d14f2828b" +dependencies = [ + "anyhow", + "asyncband 0.7.1", + "equivalent", + "foyer-common", + "foyer-memory", + "foyer-storage", + "foyer-tokio", + "futures-util", + "mixtrics", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-common" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05cdcae6cedec72c28e97ada0b453e33b1df57fbc209708091107be2a283d577" +dependencies = [ + "anyhow", + "bytes", + "cfg-if 1.0.4", + "foyer-tokio", + "mixtrics", + "parking_lot", + "pin-project", + "twox-hash", +] + +[[package]] +name = "foyer-intrusive-collections" +version = "0.10.0-dev" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4fee46bea69e0596130e3210e65d3424e0ac1e6df3bde6636304bdf1ca4a3b" +dependencies = [ + "memoffset", +] + +[[package]] +name = "foyer-memory" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a51c8ce8e1e323a1e087ac45bf629d953676fc5e0037d14a799fadd272b42db" +dependencies = [ + "anyhow", + "asyncband 0.7.1", + "bitflags 2.11.0", + "datasketches", + "equivalent", + "foyer-common", + "foyer-intrusive-collections", + "foyer-tokio", + "futures-util", + "hashbrown 0.17.1", + "itertools 0.15.0", + "mixtrics", + "parking_lot", + "paste", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-storage" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "127a64057f63e361123cf62b3fd50a36147783687d4cd36e7087c27f9909265b" +dependencies = [ + "allocator-api2", + "anyhow", + "asyncband 0.7.1", + "bytes", + "core_affinity", + "equivalent", + "fastant", + "foyer-common", + "foyer-memory", + "foyer-tokio", + "fs4", + "futures-core", + "futures-util", + "hashbrown 0.17.1", + "io-uring", + "itertools 0.15.0", + "libc", + "lz4", + "parking_lot", + "pin-project", + "rand 0.10.1", + "tracing", + "twox-hash", + "zstd", +] + +[[package]] +name = "foyer-tokio" +version = "0.22.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbdbb9f39443cb348a069baa1a0ec73bcea848a4a383eb2da1a5ea7a0ca05941" +dependencies = [ + "tokio", +] + [[package]] name = "frostem" version = "1.20260821.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "36a80a7406da302e04bfd2ca987907590d3a1f3c69958947c43890abd7426b2f" +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -3485,6 +3639,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -3789,8 +3952,10 @@ dependencies = [ "arrow-array", "arrow-schema", "async-trait", + "bytes", "chrono", "datafusion", + "foyer", "futures", "half", "lance", @@ -3804,6 +3969,7 @@ dependencies = [ "lance-table", "libc", "log", + "object_store", "opendal", "pin-project", "prost", @@ -4494,6 +4660,15 @@ dependencies = [ "libc", ] +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "mime" version = "0.3.17" @@ -4531,6 +4706,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "mixtrics" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c46b5adfb7a3ae4996d327a5bdc90e78fec025806dd312bdbe6f07a755e0ec9" +dependencies = [ + "itertools 0.15.0", + "parking_lot", +] + [[package]] name = "moka" version = "0.12.15" @@ -4842,7 +5027,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48dbcef97d3eb7591db2c18d5cae95c836bcce07359b98d98dd6f4e861eb77b7" dependencies = [ "anyhow", - "asyncband", + "asyncband 0.6.7", "base64 0.23.1", "bytes", "futures", @@ -4881,7 +5066,7 @@ version = "0.58.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03f9e144b5228d741c3763ade8711d9b72e5fb6d998e779f2d7a09da0b5a3eba" dependencies = [ - "asyncband", + "asyncband 0.6.7", "futures", "http 1.4.0", "opendal-core", @@ -4945,7 +5130,7 @@ version = "0.58.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2d564484a8f7d091827e825cfc91ed45bd48e64d262451ee041fe843db81bd8a" dependencies = [ - "asyncband", + "asyncband 0.6.7", "base64 0.23.1", "bytes", "http 1.4.0", @@ -6670,6 +6855,12 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "small_ctor" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" + [[package]] name = "smallvec" version = "1.15.1" @@ -7932,6 +8123,15 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + [[package]] name = "windows-sys" version = "0.60.2" diff --git a/Cargo.toml b/Cargo.toml index 489a609..c73208f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ datafusion = { version = "54.0.0", default-features = false } arrow = { version = "58.0.0", features = ["prettyprint", "ffi"] } arrow-array = "58.0.0" arrow-schema = "58.0.0" +bytes = "1" # Direct to name `chrono::TimeDelta` (the field type of lance's public # `AutoCleanupParams`) and `chrono::DateTime`/`Utc` (index metadata # timestamps); already in the graph transitively via lance. @@ -41,10 +42,12 @@ chrono = { version = "0.4", default-features = false } half = "2" tokio = { version = "1", features = ["rt-multi-thread", "sync"] } futures = "0.3" +foyer = "=0.22.5" log = "0.4" libc = "0.2" # Explicitly install the HTTP transport when embedded in a static C/C++ executable. opendal = { version = "=0.58.2", default-features = false, features = ["http-transport-reqwest"] } +object_store = "0.13.2" pin-project = "1.0" prost = "0.14" snafu = "0.9" diff --git a/README.md b/README.md index 9ceccdc..b1097f0 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ Based on the [liblance RFC](https://github.com/lance-format/lance/discussions/60 | [x] | Async scan | Callback-based `lance_scanner_scan_async()` for non-blocking scans | | [x] | Dataset metadata | `lance_dataset_version()`, `lance_dataset_count_rows()`, `lance_dataset_latest_version()` | | [x] | Filter pushdown | `lance_scanner_set_substrait_filter()` accepts a serialized Substrait `ExtendedExpression`; `lance_scanner_additional_sql_filter()` adds SQL predicates with AND before scanning starts | +| [x] | Data-file cache | Optional Foyer disk cache shared with serializable index entries | ## Multi-vector search @@ -231,6 +232,42 @@ auto ds = lance::Dataset::open_with_session(session, "data.lance"); auto stats = session.cache_stats(); ``` +To cache remote Lance data and serializable index entries on disk, configure +one Foyer directory and total capacity. Data-file blocks and serialized index +entries share that capacity and eviction pool. L1 is the Session's index and +metadata memory caches, with their existing separate budgets. L2 is the shared +Foyer disk cache; Foyer's memory-cache capacity is zero. Index entries follow +L1 -> L2 -> origin, data-file blocks follow L2 -> origin, and metadata keeps its +existing L1 -> origin path. Data reads cover direct `data/*.lance` children; index +entries use Lance's `CacheCodec` rather than caching raw index files. Manifests, +deletion files, and conditional/versioned data reads keep their normal paths. +Share one session per cache directory; concurrent owners are rejected. A change +to capacity or disk layout requires a new directory. Capacity must be a multiple +of 4096 and at least 64 MiB. Data reads use 1 MiB blocks; disk storage uses 16 MiB +blocks, so a trailing partial storage block is unused. Serialized entries that +exceed the storage block's payload limit remain eligible for the memory cache +but cannot be retained on disk. Cache writes are +asynchronous; recovery can reuse fully persisted entries but is not a durability +guarantee. + +```cpp +lance::FoyerCacheOptions disk_cache{ + "/var/cache/my-service/lance", + 100ULL * 1024 * 1024 * 1024, // total data + index disk capacity +}; +lance::Session session( + 6ULL * 1024 * 1024 * 1024, // L1 index capacity + 1ULL * 1024 * 1024 * 1024, // L1 metadata capacity + disk_cache); +auto ds = lance::Dataset::open_with_session(session, "s3://bucket/data.lance"); +``` + +Index cache statistics expose L1 `memory_hits`/`memory_misses` and L2 +`disk_hits`/`disk_misses` separately. An L2 hit counts as an L1 miss and an L2 hit. +The Session's aggregate index hit counter includes both tiers: its backend's +`was_cached` result means the origin loader was skipped. Requests coalesced by +L1 also count as hits. Data-file byte statistics remain scoped to each dataset. + ### Open at a specific version `lance_dataset_open` takes a `version` argument — `0` means the latest, any diff --git a/include/lance/lance.h b/include/lance/lance.h index 7630913..818b9da 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -214,6 +214,53 @@ typedef struct LanceSessionCacheStats { uint64_t metadata_cache_size_bytes; } LanceSessionCacheStats; +/** + * One shared L2 disk cache for immutable data-file blocks and serializable index + * entries. Both use the same directory and total capacity, with no per-type quota. + * Share one session per directory. Use a new directory when changing capacity. + * Block sizes are managed internally (1 MiB data reads, 16 MiB storage blocks). + */ +typedef struct LanceFoyerCacheOptions { + const char* directory; + /** Total disk capacity in bytes; a multiple of 4096 and at least 64 MiB. */ + uint64_t disk_capacity_bytes; +} LanceFoyerCacheOptions; + +/** + * Cumulative L1 index and L2 serialized-index statistics for one shared session. + * + * `memory_*` counts the Session's L1 index cache (including coalesced requests). + * `disk_*` counts lookups issued to Foyer L2 after an L1 miss. An L2 hit + * increments memory_misses and disk_hits; it is a hit in LanceSessionCacheStats. + * L2 hits/read bytes may be served from pending-write buffers and do not + * measure physical disk I/O. + * disk_write_bytes counts serialized bytes submitted, not confirmed disk writes. + * disk_write_errors counts serialization failures, not background I/O failures. + */ +typedef struct LanceIndexDiskCacheStats { + uint64_t memory_hits; + uint64_t memory_misses; + uint64_t disk_hits; + uint64_t disk_misses; + uint64_t disk_read_bytes; + uint64_t disk_write_bytes; + uint64_t decode_failures; + uint64_t disk_read_errors; + uint64_t disk_write_errors; +} LanceIndexDiskCacheStats; + +/** + * Cumulative Foyer data-cache statistics for one opened dataset handle. + * + * Successful reads are accumulated. Both fields measure bytes returned to the + * dataset reader. Their sum is the logical data-file range bytes observed by + * the Foyer wrapper; block-aligned origin read amplification is not included. + */ +typedef struct LanceDataCacheStatistics { + uint64_t bytes_read_from_cache; + uint64_t bytes_read_from_remote; +} LanceDataCacheStatistics; + /** * Create a session that can share metadata and index caches across datasets. * @@ -225,6 +272,22 @@ LanceSession* lance_session_new( uint64_t metadata_cache_size_bytes ); +/** + * Create a session with one shared Foyer disk cache for data files and index + * entries. NULL options disable the disk cache. A non-NULL options pointer + * requires a non-empty UTF-8 directory and a valid positive disk capacity. + * Values are copied before return. Index and metadata memory budgets remain + * controlled by the first two arguments (L1). Foyer L2 adds no memory-cache + * capacity. Metadata does not enter L2. + * + * @return Session handle, or NULL on error + */ +LanceSession* lance_session_new_with_foyer_cache( + uint64_t index_cache_size_bytes, + uint64_t metadata_cache_size_bytes, + const LanceFoyerCacheOptions* foyer_cache_options +); + /** * Close a session handle. Safe to call with NULL. Datasets previously opened * with the session remain valid and retain the shared cache state. @@ -241,6 +304,12 @@ int32_t lance_session_get_cache_stats( LanceSessionCacheStats* out_stats ); +/** Copy current index disk-tier statistics to `out_stats`. */ +int32_t lance_session_get_index_disk_cache_stats( + const LanceSession* session, + LanceIndexDiskCacheStats* out_stats +); + /* ─── Dataset lifecycle ─── */ /** @@ -281,6 +350,20 @@ LanceDataset* lance_dataset_open_with_session( const LanceSession* session ); +/** + * Copy this dataset handle's cumulative data-cache statistics. + * + * A dataset not opened with a data cache reports all-zero statistics. The + * snapshot belongs only to this dataset handle; the underlying cache may + * still be shared by other datasets through a session. + * + * @return 0 on success, -1 on error + */ +int32_t lance_dataset_get_data_cache_statistics( + const LanceDataset* dataset, + LanceDataCacheStatistics* out_statistics +); + /** Close and free a dataset handle. Safe to call with NULL. */ void lance_dataset_close(LanceDataset* dataset); diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 286724e..b5dac7c 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -176,6 +176,11 @@ struct SqlColumn { // ─── Shared Session ────────────────────────────────────────────────────────── +struct FoyerCacheOptions { + std::string directory; + uint64_t disk_capacity_bytes; +}; + class Session { Handle handle_; @@ -185,6 +190,19 @@ class Session { if (!handle_) check_error(); } + Session(uint64_t index_cache_size_bytes, + uint64_t metadata_cache_size_bytes, + const FoyerCacheOptions& foyer_cache_options) { + LanceFoyerCacheOptions options{ + foyer_cache_options.directory.c_str(), + foyer_cache_options.disk_capacity_bytes, + }; + handle_ = Handle( + lance_session_new_with_foyer_cache( + index_cache_size_bytes, metadata_cache_size_bytes, &options)); + if (!handle_) check_error(); + } + LanceSessionCacheStats cache_stats() const { LanceSessionCacheStats stats{}; if (lance_session_get_cache_stats(handle_.get(), &stats) != 0) @@ -192,6 +210,13 @@ class Session { return stats; } + LanceIndexDiskCacheStats index_disk_cache_stats() const { + LanceIndexDiskCacheStats stats{}; + if (lance_session_get_index_disk_cache_stats(handle_.get(), &stats) != 0) + check_error(); + return stats; + } + const LanceSession* c_handle() const { return handle_.get(); } }; @@ -350,6 +375,13 @@ class Dataset { return Dataset(ds); } + LanceDataCacheStatistics data_cache_statistics() const { + LanceDataCacheStatistics statistics{}; + if (lance_dataset_get_data_cache_statistics(handle_.get(), &statistics) != 0) + check_error(); + return statistics; + } + /// Write an Arrow record batch stream to a Lance dataset and return the /// open dataset at the committed version. /// diff --git a/src/data_cache.rs b/src/data_cache.rs new file mode 100644 index 0000000..430b4c8 --- /dev/null +++ b/src/data_cache.rs @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Private bridge between a shared data-cache backend and dataset handles. + +use std::fmt::Debug; +use std::sync::Arc; + +use lance::Dataset; +use lance_core::Result; + +use crate::dataset::LanceDataset; +use crate::error::ffi_try; + +/// Data-cache statistics owned by one opened dataset. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct LanceDataCacheStatistics { + /// Requested bytes returned from usable data-cache entries. + pub bytes_read_from_cache: u64, + /// Requested bytes returned after a data-cache miss or fallback. + pub bytes_read_from_remote: u64, +} + +pub(crate) trait DatasetDataCache: Debug + Send + Sync { + fn snapshot(&self) -> LanceDataCacheStatistics; + + fn attach_fresh(&self, dataset: Dataset) -> (Dataset, Arc); +} + +pub(crate) trait DataCacheFactory: Debug + Send + Sync { + fn attach(&self, dataset: Dataset) -> (Dataset, Arc); +} + +/// Copy this dataset handle's cumulative data-cache statistics into +/// `out_statistics`. +/// +/// A dataset not opened with a data cache reports all-zero statistics. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_dataset_get_data_cache_statistics( + dataset: *const LanceDataset, + out_statistics: *mut LanceDataCacheStatistics, +) -> i32 { + ffi_try!( + unsafe { dataset_get_data_cache_statistics_inner(dataset, out_statistics) }, + neg + ) +} + +unsafe fn dataset_get_data_cache_statistics_inner( + dataset: *const LanceDataset, + out_statistics: *mut LanceDataCacheStatistics, +) -> Result { + if dataset.is_null() || out_statistics.is_null() { + return Err(lance_core::Error::invalid_input_source( + "dataset and out_statistics must not be NULL".into(), + )); + } + let dataset = unsafe { &*dataset }; + let statistics = dataset + .data_cache + .as_ref() + .map_or_else(LanceDataCacheStatistics::default, |cache| cache.snapshot()); + unsafe { + std::ptr::write_unaligned(out_statistics, statistics); + } + Ok(0) +} diff --git a/src/dataset.rs b/src/dataset.rs index cc1f87c..76fd39e 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -14,6 +14,7 @@ use lance::Dataset; use lance::dataset::builder::DatasetBuilder; use lance_core::Result; +use crate::data_cache::DatasetDataCache; use crate::error::{ffi_try, swallow_unwind}; use crate::helpers; use crate::runtime::block_on; @@ -23,6 +24,7 @@ use crate::stream_guard::guarded_ffi_stream_from_reader; /// Opaque handle representing an opened Lance dataset. pub struct LanceDataset { pub(crate) inner: RwLock>, + pub(crate) data_cache: Option>, } impl LanceDataset { @@ -182,8 +184,16 @@ unsafe fn open_dataset_inner( } let dataset = block_on(builder.load())?; + let (dataset, data_cache) = + if let Some(factory) = session.and_then(|session| session.data_cache_factory.clone()) { + let (dataset, data_cache) = factory.attach(dataset); + (dataset, Some(data_cache)) + } else { + (dataset, None) + }; let handle = LanceDataset { inner: RwLock::new(Arc::new(dataset)), + data_cache, }; Ok(Box::into_raw(Box::new(handle))) } @@ -519,6 +529,7 @@ mod tests { .unwrap(); let handle = LanceDataset { inner: RwLock::new(Arc::new(dataset)), + data_cache: None, }; (tmp, handle) } diff --git a/src/foyer_cache.rs b/src/foyer_cache.rs new file mode 100644 index 0000000..4d691b2 --- /dev/null +++ b/src/foyer_cache.rs @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Shared L2 Foyer disk cache for data-file blocks and serialized index entries. +//! L1 is the Session's index and metadata memory cache; Foyer adds no L1 capacity. + +use std::ffi::c_char; +use std::fs::{self, File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::Path; +use std::sync::Arc; + +use bytes::Bytes; +use foyer::{ + BlockEngineConfig, DeviceBuilder, EventListener, FsDeviceBuilder, HybridCache, + HybridCacheBuilder, HybridCachePolicy, PsyncIoEngineConfig, +}; + +use crate::error::ffi_try; +use crate::foyer_data_cache::FoyerDataCache; +use crate::foyer_index_cache::FoyerIndexCache; +use crate::helpers; +use crate::runtime::block_on; +use crate::session::{LanceSession, session_new_with_factories, u64_to_usize}; + +pub(crate) const PAGE_SIZE: usize = 4096; +pub(crate) const DEFAULT_READ_BLOCK_SIZE: usize = 1024 * 1024; +pub(crate) const DEFAULT_STORAGE_BLOCK_SIZE: usize = 16 * 1024 * 1024; + +/// One disk path and total capacity shared by data blocks and index entries. +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct LanceFoyerCacheOptions { + pub directory: *const c_char, + pub disk_capacity_bytes: u64, +} + +/// Create a session with a shared Foyer data/index disk cache. +/// NULL options disable the disk cache. Block sizes are managed internally. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_session_new_with_foyer_cache( + index_cache_size_bytes: u64, + metadata_cache_size_bytes: u64, + foyer_cache_options: *const LanceFoyerCacheOptions, +) -> *mut LanceSession { + ffi_try!( + unsafe { + session_new_with_foyer_cache_inner( + index_cache_size_bytes, + metadata_cache_size_bytes, + foyer_cache_options, + ) + }, + null + ) +} + +unsafe fn session_new_with_foyer_cache_inner( + index_cache_size_bytes: u64, + metadata_cache_size_bytes: u64, + foyer_cache_options: *const LanceFoyerCacheOptions, +) -> lance_core::Result<*mut LanceSession> { + let memory_capacity = u64_to_usize(index_cache_size_bytes, "index_cache_size_bytes")?; + u64_to_usize(metadata_cache_size_bytes, "metadata_cache_size_bytes")?; + let Some(options) = (unsafe { foyer_cache_options.as_ref() }) else { + return session_new_with_factories( + index_cache_size_bytes, + metadata_cache_size_bytes, + None, + None, + None, + ); + }; + let directory = unsafe { helpers::parse_c_string(options.directory)? }.ok_or_else(|| { + lance_core::Error::invalid_input_source( + "foyer_cache_options.directory must not be NULL".into(), + ) + })?; + if directory.is_empty() { + return Err(lance_core::Error::invalid_input_source( + "foyer_cache_options.directory must not be empty".into(), + )); + } + let disk_capacity = u64_to_usize(options.disk_capacity_bytes, "disk_capacity_bytes")?; + validate_disk_cache_sizes(disk_capacity, DEFAULT_STORAGE_BLOCK_SIZE)?; + let disk = block_on(build_disk_cache( + Path::new(directory), + disk_capacity, + DEFAULT_STORAGE_BLOCK_SIZE, + )) + .map_err(|error| { + lance_core::Error::io(format!( + "failed to initialize Foyer cache at {directory:?}: {error}", + )) + })?; + let data_cache = FoyerDataCache::from_cache(disk.clone(), DEFAULT_READ_BLOCK_SIZE); + let index_cache = FoyerIndexCache::from_cache(disk, memory_capacity, Path::new(directory)) + .map_err(|error| { + lance_core::Error::io(format!( + "failed to initialize index cache generation at {directory:?}: {error}", + )) + })?; + let index_stats = index_cache.stats(); + session_new_with_factories( + index_cache_size_bytes, + metadata_cache_size_bytes, + Some(Arc::new(index_cache)), + Some(index_stats), + Some(Arc::new(data_cache)), + ) +} + +pub(crate) fn validate_disk_cache_sizes( + disk_capacity: usize, + storage_block_size: usize, +) -> lance_core::Result<()> { + // A storage block includes a 4 KiB blob index. Reserve blocks for the + // flusher and reclaimer as well as retained entries. + if storage_block_size <= PAGE_SIZE || !storage_block_size.is_multiple_of(PAGE_SIZE) { + return Err(lance_core::Error::invalid_input_source(format!( + "storage_block_size_bytes={storage_block_size} must be a multiple of {PAGE_SIZE} and greater than {PAGE_SIZE}" + ).into())); + } + let minimum_capacity = storage_block_size.checked_mul(4).ok_or_else(|| { + lance_core::Error::invalid_input_source( + format!("storage_block_size_bytes={storage_block_size} is too large").into(), + ) + })?; + if disk_capacity < minimum_capacity || !disk_capacity.is_multiple_of(PAGE_SIZE) { + return Err(lance_core::Error::invalid_input_source(format!( + "disk_capacity_bytes={disk_capacity} must be a multiple of {PAGE_SIZE} and at least {minimum_capacity} (four storage blocks of {storage_block_size} bytes)" + ).into())); + } + Ok(()) +} + +// Foyer does not lock its FsDevice directory. Its event listener also stays +// alive while HybridCache's asynchronous drop drains storage, keeping the +// directory locked until pending background I/O completes. +struct CacheDirectory { + lock: File, +} + +impl Drop for CacheDirectory { + fn drop(&mut self) { + // Closing our descriptor alone can leave the lock held by a copy + // inherited during a concurrent process spawn. Release the lock when + // the last cache owner finishes, including initialization failures. + if let Err(error) = self.lock.unlock() { + log::warn!("failed to release Foyer cache directory lock: {error}"); + } + } +} + +impl EventListener for CacheDirectory { + type Key = String; + type Value = Bytes; +} + +fn lock_directory( + directory: &Path, + disk_capacity: usize, + storage_block_size: usize, +) -> std::result::Result, foyer::Error> { + fs::create_dir_all(directory).map_err(foyer::Error::io_error)?; + let lock = OpenOptions::new() + .create(true) + .truncate(false) + .read(true) + .write(true) + .open(directory.join("lance-cache-layout")) + .map_err(foyer::Error::io_error)?; + lock.try_lock().map_err(|error| { + foyer::Error::new(foyer::ErrorKind::Config, format!( + "cannot exclusively lock cache directory {directory:?}: {error}; share one LanceSession per directory" + )) + })?; + let mut directory_guard = CacheDirectory { lock }; + let lock = &mut directory_guard.lock; + let expected = format!("lance-disk-v1 {disk_capacity} {storage_block_size}\n"); + let mut layout = String::new(); + lock.read_to_string(&mut layout) + .map_err(foyer::Error::io_error)?; + if layout.is_empty() { + // Adopt an older cache only if opening it will not truncate existing + // partitions or leave files outside the requested capacity. + let mut partitions = 0; + for entry in fs::read_dir(directory).map_err(foyer::Error::io_error)? { + let entry = entry.map_err(foyer::Error::io_error)?; + if entry + .file_name() + .to_string_lossy() + .starts_with("foyer-storage-direct-fs-") + { + let size = entry.metadata().map_err(foyer::Error::io_error)?.len(); + if size != storage_block_size as u64 { + return Err(foyer::Error::new( + foyer::ErrorKind::Config, + format!( + "cache partition {:?} has size {size}, expected {storage_block_size}; use a new cache directory", + entry.path() + ), + )); + } + partitions += 1; + } + } + if partitions > disk_capacity / storage_block_size { + return Err(foyer::Error::new( + foyer::ErrorKind::Config, + format!( + "cache directory {directory:?} has {partitions} partitions exceeding disk_capacity_bytes={disk_capacity}; use a new cache directory" + ), + )); + } + lock.write_all(expected.as_bytes()) + .map_err(foyer::Error::io_error)?; + lock.sync_all().map_err(foyer::Error::io_error)?; + } else if layout != expected { + return Err(foyer::Error::new( + foyer::ErrorKind::Config, + format!( + "cache directory {directory:?} has layout {layout:?}, requested {expected:?}; use a new cache directory when changing disk capacity or storage block size" + ), + )); + } + Ok(Arc::new(directory_guard)) +} + +/// Build a disk-only Foyer cache. The caller supplies logical namespaces in +/// its keys, allowing data-file and index entries to share one capacity pool +/// without colliding. +pub(crate) async fn build_disk_cache( + directory: &Path, + disk_capacity: usize, + storage_block_size: usize, +) -> std::result::Result, foyer::Error> { + validate_disk_cache_sizes(disk_capacity, storage_block_size) + .map_err(|error| foyer::Error::new(foyer::ErrorKind::Config, error.to_string()))?; + let directory_lock = lock_directory(directory, disk_capacity, storage_block_size)?; + let device = FsDeviceBuilder::new(directory) + .with_capacity(disk_capacity) + .build()?; + let engine = BlockEngineConfig::new(device).with_block_size(storage_block_size); + HybridCacheBuilder::new() + .with_name("lance_disk") + .with_event_listener(directory_lock) + .with_policy(HybridCachePolicy::WriteOnInsertion) + .with_flush_on_close(false) + .memory(0) + .with_shards(1) + .with_weighter(|_key: &String, value: &Bytes| value.len().max(1)) + .storage() + .with_io_engine_config(PsyncIoEngineConfig::new()) + .with_engine_config(engine) + .build() + .await +} + +#[cfg(test)] +pub(crate) async fn wait_for_directory_release(directory: &Path) { + let file = OpenOptions::new() + .read(true) + .write(true) + .open(directory.join("lance-cache-layout")) + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + match file.try_lock() { + Ok(()) => { + file.unlock().unwrap(); + return; + } + Err(std::fs::TryLockError::WouldBlock) => { + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + Err(error) => panic!("cannot lock closed cache directory: {error}"), + } + } + }) + .await + .expect("cache directory remained locked after all owners closed"); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_unusable_storage_geometry() { + assert!(validate_disk_cache_sizes(1024 * 1024, 4096).is_err()); + assert!(validate_disk_cache_sizes(3 * 65536, 65536).is_err()); + assert!(validate_disk_cache_sizes(4 * 65536, 65536).is_ok()); + assert!(validate_disk_cache_sizes(4 * 65536 + 1, 65536).is_err()); + } + + #[tokio::test] + async fn directory_lock_follows_shared_cache_lifetime() { + let directory = tempfile::tempdir().unwrap(); + let cache = build_disk_cache(directory.path(), 1024 * 1024, 65536) + .await + .unwrap(); + let shared = cache.clone(); + drop(cache); + assert!( + build_disk_cache(directory.path(), 1024 * 1024, 65536) + .await + .is_err() + ); + shared.insert("test".to_owned(), Bytes::from_static(b"cached")); + shared.close().await.unwrap(); + drop(shared); + wait_for_directory_release(directory.path()).await; + let reopened = build_disk_cache(directory.path(), 1024 * 1024, 65536) + .await + .unwrap(); + assert_eq!( + reopened + .get("test") + .await + .unwrap() + .unwrap() + .value() + .as_ref(), + b"cached" + ); + reopened.close().await.unwrap(); + } + + #[test] + fn directory_guard_releases_lock_even_if_descriptor_was_duplicated() { + let directory = tempfile::tempdir().unwrap(); + let guard = lock_directory(directory.path(), 1024 * 1024, 65536).unwrap(); + // A concurrent process spawn can temporarily inherit the same open + // file description. Closing only our descriptor would retain its lock. + let inherited = guard.lock.try_clone().unwrap(); + drop(guard); + let reopened = lock_directory(directory.path(), 1024 * 1024, 65536).unwrap(); + drop(inherited); + drop(reopened); + } + + #[test] + fn rejects_layout_changes_and_legacy_partition_truncation() { + let directory = tempfile::tempdir().unwrap(); + let guard = lock_directory(directory.path(), 1024 * 1024, 65536).unwrap(); + drop(guard); + assert!(lock_directory(directory.path(), 2 * 1024 * 1024, 65536).is_err()); + assert!(lock_directory(directory.path(), 1024 * 1024, 131072).is_err()); + lock_directory(directory.path(), 1024 * 1024, 65536).unwrap(); + + let legacy = tempfile::tempdir().unwrap(); + let partition = legacy.path().join("foyer-storage-direct-fs-00000000"); + File::create(&partition).unwrap().set_len(65536).unwrap(); + assert!(lock_directory(legacy.path(), 1024 * 1024, 131072).is_err()); + assert_eq!(fs::metadata(&partition).unwrap().len(), 65536); + lock_directory(legacy.path(), 1024 * 1024, 65536).unwrap(); + } +} diff --git a/src/foyer_data_cache.rs b/src/foyer_data_cache.rs new file mode 100644 index 0000000..a9756dc --- /dev/null +++ b/src/foyer_data_cache.rs @@ -0,0 +1,1249 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Foyer-backed cache for immutable Lance data-file reads. + +use std::collections::{BTreeMap, HashMap}; +use std::fmt::{Debug, Display, Formatter}; +use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, Weak}; + +use async_trait::async_trait; +use bytes::{Bytes, BytesMut}; +use foyer::HybridCache; +use futures::StreamExt; +use futures::stream::BoxStream; +use lance_io::object_store::WrappingObjectStore; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, + RenameOptions, Result, +}; + +use crate::data_cache::{DataCacheFactory, DatasetDataCache, LanceDataCacheStatistics}; +const CACHE_KEY_VERSION: &str = "lance-data-v1"; +// Bound concurrent disk-tier lookups so one large range request cannot flood +// Foyer's storage executor or starve other Lance queries. +const DATA_CACHE_LOOKUP_CONCURRENCY: usize = 8; + +/// Process-local owner of a Foyer hybrid cache. +#[derive(Clone)] +pub(crate) struct FoyerDataCache { + cache: HybridCache, + read_block_size: usize, + wrapped_stores: Arc>>, +} + +struct WrappedStore { + wrapper: Weak, + origin: Weak, +} + +impl Debug for FoyerDataCache { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FoyerDataCache") + .field("read_block_size", &self.read_block_size) + .finish_non_exhaustive() + } +} + +impl FoyerDataCache { + #[cfg(test)] + async fn try_new( + directory: &std::path::Path, + disk_capacity: usize, + read_block_size: usize, + ) -> std::result::Result { + let cache = + crate::foyer_cache::build_disk_cache(directory, disk_capacity, 2 * read_block_size) + .await?; + Ok(Self::from_cache(cache, read_block_size)) + } + + pub(crate) fn from_cache(cache: HybridCache, read_block_size: usize) -> Self { + Self { + cache, + read_block_size, + wrapped_stores: Arc::new(Mutex::new(HashMap::new())), + } + } + + fn is_cacheable_data_file(location: &Path) -> bool { + let mut parts = location.as_ref().rsplit('/'); + matches!( + (parts.next(), parts.next()), + (Some(file), Some("data")) if file.ends_with(".lance") + ) + } + + fn key(&self, store_prefix: &str, location: &Path, block_index: u64) -> String { + format!( + "{CACHE_KEY_VERSION}\0{}\0{store_prefix}\0{}\0{block_index}", + self.read_block_size, + location.as_ref() + ) + } + + fn size_key(&self, store_prefix: &str, location: &Path) -> String { + format!( + "{CACHE_KEY_VERSION}\0{}\0{store_prefix}\0{}\0size", + self.read_block_size, + location.as_ref() + ) + } + + fn create_scope(&self) -> Arc { + Arc::new(DatasetFoyerDataCache { + cache: self.clone(), + statistics: Arc::new(FoyerDataCacheStatistics::default()), + }) + } + + fn unwrap_store(&self, store: Arc) -> Arc { + let identity = Arc::as_ptr(&store) as *const () as usize; + let origin = { + let mut wrapped_stores = self + .wrapped_stores + .lock() + .unwrap_or_else(|error| error.into_inner()); + let origin = wrapped_stores.get(&identity).and_then(|entry| { + let wrapper = entry.wrapper.upgrade()?; + if Arc::ptr_eq(&wrapper, &store) { + entry.origin.upgrade() + } else { + None + } + }); + if origin.is_none() { + wrapped_stores.remove(&identity); + } + origin + }; + match origin { + Some(origin) => origin, + None => store, + } + } + + fn remember_wrapper(&self, wrapper: &Arc, origin: &Arc) { + let identity = Arc::as_ptr(wrapper) as *const () as usize; + self.wrapped_stores + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert( + identity, + WrappedStore { + wrapper: Arc::downgrade(wrapper), + origin: Arc::downgrade(origin), + }, + ); + } + + fn forget_wrapper(&self, identity: usize) { + self.wrapped_stores + .lock() + .unwrap_or_else(|error| error.into_inner()) + .remove(&identity); + } +} + +#[derive(Debug, Default)] +struct FoyerDataCacheStatistics { + bytes_read_from_cache: AtomicU64, + bytes_read_from_remote: AtomicU64, +} + +impl FoyerDataCacheStatistics { + fn record(&self, bytes_read_from_cache: u64, bytes_read_from_remote: u64) { + self.bytes_read_from_cache + .fetch_add(bytes_read_from_cache, Ordering::Relaxed); + self.bytes_read_from_remote + .fetch_add(bytes_read_from_remote, Ordering::Relaxed); + } + + fn snapshot(&self) -> LanceDataCacheStatistics { + LanceDataCacheStatistics { + bytes_read_from_cache: self.bytes_read_from_cache.load(Ordering::Relaxed), + bytes_read_from_remote: self.bytes_read_from_remote.load(Ordering::Relaxed), + } + } +} + +impl DataCacheFactory for FoyerDataCache { + fn attach(&self, dataset: lance::Dataset) -> (lance::Dataset, Arc) { + let scope = self.create_scope(); + let wrapper: Arc = scope.clone(); + let dataset = dataset.with_object_store_wrappers([wrapper]); + (dataset, scope) + } +} + +#[derive(Debug)] +struct DatasetFoyerDataCache { + cache: FoyerDataCache, + statistics: Arc, +} + +impl DatasetDataCache for DatasetFoyerDataCache { + fn snapshot(&self) -> LanceDataCacheStatistics { + self.statistics.snapshot() + } + + fn attach_fresh(&self, dataset: lance::Dataset) -> (lance::Dataset, Arc) { + self.cache.attach(dataset) + } +} + +impl WrappingObjectStore for DatasetFoyerDataCache { + fn wrap(&self, store_prefix: &str, original: Arc) -> Arc { + // A derived Dataset can already contain this cache wrapper. Resolve + // that exact wrapper back to its origin before attaching fresh + // dataset-scoped counters. + let original = self.cache.unwrap_store(original); + let reader = DataCacheReader { + cache: self.cache.clone(), + store_prefix: store_prefix.to_owned(), + original: original.clone(), + statistics: self.statistics.clone(), + }; + let cached_store = + Arc::new_cyclic(|weak: &Weak| DataCacheObjectStore { + reader, + identity: weak.as_ptr() as usize, + }); + let wrapped: Arc = cached_store.clone(); + self.cache.remember_wrapper(&wrapped, &original); + wrapped + } +} + +#[derive(Debug)] +struct DataCacheObjectStore { + reader: DataCacheReader, + identity: usize, +} + +#[derive(Clone, Debug)] +struct DataCacheReader { + cache: FoyerDataCache, + store_prefix: String, + original: Arc, + statistics: Arc, +} + +impl Drop for DataCacheObjectStore { + fn drop(&mut self) { + self.reader.cache.forget_wrapper(self.identity); + } +} + +impl Display for DataCacheObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FoyerDataCache({})", self.reader.original) + } +} + +impl DataCacheObjectStore { + fn is_cache_safe_get(options: &GetOptions) -> bool { + !options.head + && options.if_match.is_none() + && options.if_none_match.is_none() + && options.if_modified_since.is_none() + && options.if_unmodified_since.is_none() + && options.version.is_none() + && options.extensions.is_empty() + } + + async fn cached_get(&self, location: &Path, options: GetOptions) -> Result { + // Fetch metadata separately so the returned GetResult retains the origin's identity while + // its payload uses the same block cache as get_ranges(). This also provides the object size + // needed to resolve bounded, offset, and suffix ranges. + let GetResult { + meta: metadata, + attributes, + .. + } = self + .reader + .original + .get_opts( + location, + GetOptions { + head: true, + ..Default::default() + }, + ) + .await?; + let object_size = metadata.size; + self.reader.cache.cache.insert( + self.reader + .cache + .size_key(&self.reader.store_prefix, location), + Bytes::copy_from_slice(&object_size.to_le_bytes()), + ); + + let range = match options.range.clone() { + Some(requested) => match requested.as_range(object_size) { + Ok(range) if !range.is_empty() => range, + // Preserve the origin's exact error for invalid or empty ranges. + _ => return self.reader.original.get_opts(location, options).await, + }, + None => 0..object_size, + }; + + let reader = self.reader.clone(); + let stream_location = location.clone(); + let stream_range = range.clone(); + let stream = futures::stream::try_unfold( + (reader, stream_location, stream_range), + |(reader, location, remaining)| async move { + if remaining.is_empty() { + return Ok(None); + } + + // Yield no more than the remainder of one cache block. The next block is not + // requested until the consumer polls again, so cancellation drops the pending + // range without downloading or retaining the rest of the object. + let block_size = reader.cache.read_block_size as u64; + let bytes_to_boundary = block_size - remaining.start % block_size; + let end = remaining + .start + .saturating_add(bytes_to_boundary) + .min(remaining.end); + let chunk_range = remaining.start..end; + let chunk = reader + .cached_ranges(&location, std::slice::from_ref(&chunk_range)) + .await? + .into_iter() + .next() + .ok_or_else(|| cache_error(format!("missing get result for {location}")))?; + Ok(Some((chunk, (reader, location, end..remaining.end)))) + }, + ); + let payload = GetResultPayload::Stream(Box::pin(stream)); + Ok(GetResult { + payload, + meta: metadata, + range, + attributes, + }) + } +} + +impl DataCacheReader { + async fn read_origin_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> Result> { + let bytes = self.original.get_ranges(location, ranges).await?; + self.statistics.record(0, total_bytes(&bytes)); + Ok(bytes) + } + + async fn object_size(&self, location: &Path) -> Result { + let key = self.cache.size_key(&self.store_prefix, location); + match self.cache.cache.get(&key).await { + Ok(Some(entry)) => match entry.value().as_ref().try_into() { + Ok(bytes) => return Ok(u64::from_le_bytes(bytes)), + Err(_) => log::warn!( + "Foyer data-cache size entry was malformed for {location}; refreshing it" + ), + }, + Ok(None) => {} + Err(error) => { + log::warn!("Foyer data-cache size lookup failed for {location}: {error}"); + } + } + + let size = self.original.head(location).await?.size; + self.cache + .cache + .insert(key, Bytes::copy_from_slice(&size.to_le_bytes())); + Ok(size) + } + + async fn cached_ranges(&self, location: &Path, ranges: &[Range]) -> Result> { + if ranges.is_empty() { + return Ok(Vec::new()); + } + if ranges.iter().any(|range| range.start >= range.end) { + return self.read_origin_ranges(location, ranges).await; + } + + let object_size = self.object_size(location).await?; + if ranges.iter().any(|range| range.start >= object_size) { + // Preserve the origin's exact error for ranges that start at or + // beyond EOF. + return self.read_origin_ranges(location, ranges).await; + } + let readable_ranges = ranges + .iter() + .map(|range| range.start..range.end.min(object_size)) + .collect::>(); + + let block_size = self.cache.read_block_size as u64; + let mut blocks = BTreeMap::>::new(); + for range in &readable_ranges { + let first = range.start / block_size; + let last = (range.end - 1) / block_size; + for block_index in first..=last { + blocks.entry(block_index).or_default(); + } + } + + // Data entries are disk-only. Look up independent blocks concurrently, + // but keep a fixed bound to avoid flooding the storage executor. + let cache = self.cache.cache.clone(); + let lookup_items = blocks + .keys() + .copied() + .map(|block_index| { + ( + block_index, + self.cache.key(&self.store_prefix, location, block_index), + ) + }) + .collect::>(); + let lookup_concurrency = lookup_items.len().clamp(1, DATA_CACHE_LOOKUP_CONCURRENCY); + let mut lookup_stream = + futures::stream::iter(lookup_items.into_iter().map(|(block_index, key)| { + let cache = cache.clone(); + async move { (block_index, cache.get(&key).await) } + })) + .buffer_unordered(lookup_concurrency); + while let Some((block_index, result)) = lookup_stream.next().await { + match result { + Ok(Some(entry)) => { + let expected_size = (object_size - block_index * block_size).min(block_size); + if entry.value().len() as u64 != expected_size { + // Reject the whole malformed block, even if its short + // contents would cover this particular requested slice. + self.cache.cache.remove(&self.cache.key( + &self.store_prefix, + location, + block_index, + )); + continue; + } + if let Some(block) = blocks.get_mut(&block_index) { + *block = Some(entry.value().clone()); + } + } + Ok(None) => {} + Err(error) => { + // Cache availability must not affect query correctness. + log::warn!("Foyer data-cache lookup failed for {location}: {error}"); + } + } + } + + let (bytes_read_from_cache, bytes_read_from_remote) = + requested_bytes_by_cache_status(&readable_ranges, block_size, &blocks); + let missing: Vec = blocks + .iter() + .filter_map(|(block_index, value)| value.is_none().then_some(*block_index)) + .collect(); + let mut runs = Vec::>::new(); + for block_index in missing { + let start = block_index + .checked_mul(block_size) + .ok_or_else(|| cache_error("data-cache block offset overflow"))?; + let end = start + .checked_add(block_size) + .ok_or_else(|| cache_error("data-cache block end overflow"))? + .min(object_size); + match runs.last_mut() { + Some(run) if run.end == start => run.end = end, + _ => runs.push(start..end), + } + } + + if !runs.is_empty() { + // Fetch all contiguous miss runs together so a large Lance read is + // not expanded into one remote request per cache block. + let fetched = self.original.get_ranges(location, &runs).await?; + for (run, bytes) in runs.into_iter().zip(fetched) { + let first_block = run.start / block_size; + for (offset, chunk) in bytes.chunks(self.cache.read_block_size).enumerate() { + let block_index = first_block + offset as u64; + let value = Bytes::copy_from_slice(chunk); + let key = self.cache.key(&self.store_prefix, location, block_index); + self.cache.cache.insert(key, value.clone()); + if let Some(block) = blocks.get_mut(&block_index) { + *block = Some(value); + } + } + } + } + + let assembled = readable_ranges + .iter() + .map(|range| assemble_range(location, range, block_size, &blocks)) + .collect::>>(); + match assembled { + Ok(bytes) => { + self.statistics + .record(bytes_read_from_cache, bytes_read_from_remote); + Ok(bytes) + } + Err(error) => { + // A malformed or incomplete cached entry must never turn a + // valid source read into a query failure. + log::warn!( + "Foyer data-cache entry was unusable for {location}; bypassing cache: {error}" + ); + let bytes = self.original.get_ranges(location, ranges).await?; + self.statistics.record(0, total_bytes(&bytes)); + Ok(bytes) + } + } + } +} + +fn total_bytes(ranges: &[Bytes]) -> u64 { + ranges.iter().fold(0_u64, |total, bytes| { + total.saturating_add(bytes.len() as u64) + }) +} + +fn requested_bytes_by_cache_status( + ranges: &[Range], + block_size: u64, + blocks: &BTreeMap>, +) -> (u64, u64) { + let mut hit_bytes = 0_u64; + let mut miss_bytes = 0_u64; + for range in ranges { + let mut start = range.start; + while start < range.end { + let block_index = start / block_size; + let block_start = block_index * block_size; + let end = range.end.min(block_start.saturating_add(block_size)); + let bytes = end - start; + if blocks + .get(&block_index) + .is_some_and(|block| block.is_some()) + { + hit_bytes = hit_bytes.saturating_add(bytes); + } else { + miss_bytes = miss_bytes.saturating_add(bytes); + } + start = end; + } + } + (hit_bytes, miss_bytes) +} + +fn assemble_range( + location: &Path, + range: &Range, + block_size: u64, + blocks: &BTreeMap>, +) -> Result { + if range.is_empty() { + return Ok(Bytes::new()); + } + // The caller already clamps the range to the object's EOF. A short cached value + // must not shorten it again: reject incomplete coverage so cached_ranges can + // retry the original request against the source store. + let first = range.start / block_size; + let last = (range.end - 1) / block_size; + if first == last { + let block = blocks + .get(&first) + .and_then(Option::as_ref) + .ok_or_else(|| cache_error(format!("missing block {first} for {location}")))?; + let block_start = first * block_size; + let start = usize::try_from(range.start - block_start) + .map_err(|_| cache_error("data-cache slice start exceeds usize::MAX"))?; + let end = usize::try_from((range.end - block_start).min(block_size)) + .map_err(|_| cache_error("data-cache slice end exceeds usize::MAX"))?; + if end > block.len() { + return Err(cache_error(format!( + "short data-cache block {first} for {location}: need {start}..{end}, got {} bytes", + block.len() + ))); + } + return Ok(block.slice(start..end)); + } + + let requested_len = usize::try_from(range.end - range.start) + .map_err(|_| cache_error(format!("range {range:?} for {location} exceeds usize::MAX")))?; + let mut output = BytesMut::with_capacity(requested_len); + for block_index in first..=last { + let block = blocks + .get(&block_index) + .and_then(Option::as_ref) + .ok_or_else(|| cache_error(format!("missing block {block_index} for {location}")))?; + let block_start = block_index * block_size; + let start = usize::try_from(range.start.saturating_sub(block_start)) + .map_err(|_| cache_error("data-cache slice start exceeds usize::MAX"))?; + let end_in_block = range.end.saturating_sub(block_start).min(block_size); + let end = usize::try_from(end_in_block) + .map_err(|_| cache_error("data-cache slice end exceeds usize::MAX"))?; + if end > block.len() { + return Err(cache_error(format!( + "short data-cache block {block_index} for {location}: need {start}..{end}, got {} bytes", + block.len() + ))); + } + output.extend_from_slice(&block[start..end]); + } + if output.len() != requested_len { + return Err(cache_error(format!( + "incomplete data-cache range {range:?} for {location}: expected {requested_len} bytes, got {}", + output.len() + ))); + } + Ok(output.freeze()) +} + +fn cache_error(message: impl Into) -> object_store::Error { + object_store::Error::Generic { + store: "foyer_data_cache", + source: Box::new(std::io::Error::other(message.into())), + } +} + +#[async_trait] +#[deny(clippy::missing_trait_methods)] +impl ObjectStore for DataCacheObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.reader.original.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + self.reader + .original + .put_multipart_opts(location, opts) + .await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + if FoyerDataCache::is_cacheable_data_file(location) && Self::is_cache_safe_get(&options) { + self.cached_get(location, options).await + } else { + self.reader.original.get_opts(location, options).await + } + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> Result> { + if FoyerDataCache::is_cacheable_data_file(location) { + self.reader.cached_ranges(location, ranges).await + } else { + self.reader.original.get_ranges(location, ranges).await + } + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.reader.original.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.reader.original.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + self.reader.original.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + self.reader.original.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { + self.reader.original.copy_opts(from, to, options).await + } + + async fn rename_opts(&self, from: &Path, to: &Path, options: RenameOptions) -> Result<()> { + self.reader.original.rename_opts(from, to, options).await + } +} + +#[cfg(test)] +mod tests { + use std::sync::mpsc; + use std::time::Duration; + + use futures::StreamExt; + use lance_io::object_store::ChainedWrappingObjectStore; + use object_store::GetRange; + use object_store::memory::InMemory; + + use super::*; + + fn wrap_for_test( + cache: &FoyerDataCache, + original: Arc, + ) -> (Arc, Arc) { + let scope = cache.create_scope(); + (scope.wrap("memory://test", original), scope) + } + + #[test] + fn assemble_range_rejects_incomplete_cached_coverage() { + let location = Path::from("table.lance/data/part-0.lance"); + let full_block = Bytes::from_static(b"abcdefgh"); + // Exercise a single block (including a start beyond the cached bytes), + // every position in a multi-block range, and an empty block after output. + for (range, short_index, short_len) in [ + (0..8, 0, 4), + (2..8, 0, 4), + (4..8, 0, 4), + (6..8, 0, 4), + (0..8, 0, 0), + (0..24, 0, 4), + (0..24, 1, 4), + (0..24, 2, 4), + (0..24, 1, 0), + ] { + let mut blocks = BTreeMap::from([ + (0, Some(full_block.clone())), + (1, Some(full_block.clone())), + (2, Some(full_block.clone())), + ]); + blocks.insert(short_index, Some(full_block.slice(..short_len))); + let error = assemble_range(&location, &range, 8, &blocks).unwrap_err(); + assert!( + error.to_string().contains("short data-cache block"), + "range={range:?}, short_index={short_index}, short_len={short_len}: {error}" + ); + } + } + + #[test] + fn assemble_range_preserves_exact_slices_and_eof_tail() { + let location = Path::from("table.lance/data/part-0.lance"); + let data = Bytes::from_static(b"abcdefghijk"); + let blocks = BTreeMap::from([(0, Some(data.slice(..8))), (1, Some(data.slice(8..)))]); + // The last block is shorter than the cache block size but fully covers + // each requested slice. Rejecting every short value would break EOF reads. + for range in [0..0, 0..8, 2..6, 3..11, 8..11, 9..11] { + assert_eq!( + assemble_range(&location, &range, 8, &blocks).unwrap(), + data.slice(range.start as usize..range.end as usize), + "range={range:?}" + ); + } + assert!( + assemble_range(&location, &(0..0), 8, &BTreeMap::new()) + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn incomplete_cache_hits_fall_back_to_complete_source_ranges() { + const BLOCK_SIZE: usize = 64 * 1024; + let original = Arc::new(InMemory::new()); + let data = Bytes::from( + (0..3 * BLOCK_SIZE) + .map(|value| value as u8) + .collect::>(), + ); + + for (case, range, short_index, short_len) in [ + ("single", 0..BLOCK_SIZE as u64, 0, BLOCK_SIZE / 2), + ("first", 0..data.len() as u64, 0, BLOCK_SIZE / 2), + ("middle", 0..data.len() as u64, 1, BLOCK_SIZE / 2), + ("last", 0..data.len() as u64, 2, BLOCK_SIZE / 2), + ("empty", 0..data.len() as u64, 1, 0), + ("partial", 0..16, 0, BLOCK_SIZE / 2), + ] { + let directory = tempfile::tempdir().unwrap(); + // Keep every injected entry available to exercise malformed hits, + // rather than ordinary misses caused by eviction. + let cache = FoyerDataCache::try_new(directory.path(), 64 * BLOCK_SIZE, BLOCK_SIZE) + .await + .unwrap(); + let location = Path::from(format!("table.lance/data/{case}.lance")); + original.put(&location, data.clone().into()).await.unwrap(); + cache.cache.insert( + cache.size_key("memory://test", &location), + Bytes::copy_from_slice(&(data.len() as u64).to_le_bytes()), + ); + for block_index in 0..3 { + let value = if block_index == short_index { + Bytes::from(vec![255; short_len]) + } else { + let start = block_index * BLOCK_SIZE; + data.slice(start..start + BLOCK_SIZE) + }; + cache.cache.insert( + cache.key("memory://test", &location, block_index as u64), + value, + ); + } + // Recover fully persisted corruption. This also avoids racing the + // fixture's pending writes against the cache's repair writes. + cache.cache.close().await.unwrap(); + drop(cache); + crate::foyer_cache::wait_for_directory_release(directory.path()).await; + let cache = FoyerDataCache::try_new(directory.path(), 64 * BLOCK_SIZE, BLOCK_SIZE) + .await + .unwrap(); + let short_key = cache.key("memory://test", &location, short_index as u64); + assert_eq!( + cache + .cache + .get(&short_key) + .await + .unwrap() + .unwrap() + .value() + .len(), + short_len + ); + let (wrapped, statistics) = wrap_for_test(&cache, original.clone()); + assert_eq!( + wrapped + .get_ranges(&location, std::slice::from_ref(&range)) + .await + .unwrap(), + vec![data.slice(range.start as usize..range.end as usize)], + "case={case}" + ); + let bad_start = (short_index * BLOCK_SIZE) as u64; + let bad_end = bad_start + BLOCK_SIZE as u64; + let remote_bytes = range.end.min(bad_end) - range.start.max(bad_start); + assert_eq!( + statistics.snapshot().bytes_read_from_cache, + range.end - range.start - remote_bytes, + "case={case}" + ); + assert_eq!( + statistics.snapshot().bytes_read_from_remote, + remote_bytes, + "case={case}" + ); + // The malformed block is repaired, so the next read is fully cached. + let before = statistics.snapshot(); + wrapped + .get_ranges(&location, std::slice::from_ref(&range)) + .await + .unwrap(); + assert_eq!( + statistics.snapshot().bytes_read_from_remote, + before.bytes_read_from_remote + ); + } + } + + #[tokio::test] + async fn caches_only_immutable_data_file_ranges() { + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let original = Arc::new(InMemory::new()); + let data_path = Path::from("table.lance/data/part-0.lance"); + let second_data_path = Path::from("table.lance/data/part-1.lance"); + let manifest_path = Path::from("table.lance/_versions/1.manifest"); + let data = Bytes::from((0..200_000).map(|value| value as u8).collect::>()); + let second_data = Bytes::from_static(b"second fragment"); + original.put(&data_path, data.clone().into()).await.unwrap(); + original + .put(&second_data_path, second_data.clone().into()) + .await + .unwrap(); + original + .put(&manifest_path, Bytes::from_static(b"manifest").into()) + .await + .unwrap(); + + let (wrapped, statistics) = wrap_for_test(&cache, original.clone()); + let ranges = vec![10..90_000, 65_000..140_000, 190_000..220_000]; + let second_data_range = 0..15; + let first = wrapped.get_ranges(&data_path, &ranges).await.unwrap(); + assert_eq!(first[0], data.slice(10..90_000)); + assert_eq!(first[1], data.slice(65_000..140_000)); + assert_eq!(first[2], data.slice(190_000..200_000)); + assert_eq!( + wrapped + .get_ranges(&second_data_path, std::slice::from_ref(&second_data_range)) + .await + .unwrap(), + vec![second_data.clone()] + ); + + assert_eq!( + statistics.snapshot(), + LanceDataCacheStatistics { + bytes_read_from_cache: 0, + bytes_read_from_remote: 175_005, + } + ); + + original.delete(&data_path).await.unwrap(); + original.delete(&second_data_path).await.unwrap(); + let second = wrapped.get_ranges(&data_path, &ranges).await.unwrap(); + assert_eq!(second, first); + assert_eq!( + wrapped + .get_ranges(&second_data_path, std::slice::from_ref(&second_data_range)) + .await + .unwrap(), + vec![second_data] + ); + + assert_eq!( + statistics.snapshot(), + LanceDataCacheStatistics { + bytes_read_from_cache: 175_005, + bytes_read_from_remote: 175_005, + } + ); + + assert_eq!( + wrapped.get_range(&manifest_path, 0..8).await.unwrap(), + Bytes::from_static(b"manifest") + ); + original.delete(&manifest_path).await.unwrap(); + assert!(wrapped.get_range(&manifest_path, 0..8).await.is_err()); + } + + #[tokio::test] + async fn caches_small_data_file_whole_object_reads() { + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let original = Arc::new(InMemory::new()); + let data_path = Path::from("table.lance/data/small.lance"); + let data = Bytes::from(vec![7; 42_000]); + original.put(&data_path, data.clone().into()).await.unwrap(); + + let (wrapped, statistics) = wrap_for_test(&cache, original); + let first = wrapped + .get(&data_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert_eq!(first, data); + assert_eq!( + statistics.snapshot(), + LanceDataCacheStatistics { + bytes_read_from_cache: 0, + bytes_read_from_remote: 42_000, + } + ); + + let second = wrapped + .get(&data_path) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert_eq!(second, data); + assert_eq!( + statistics.snapshot(), + LanceDataCacheStatistics { + bytes_read_from_cache: 42_000, + bytes_read_from_remote: 42_000, + } + ); + } + + #[tokio::test] + async fn streams_large_data_file_gets_with_bounded_read_ahead() { + let directory = tempfile::tempdir().unwrap(); + let block_size = 64 * 1024; + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, block_size) + .await + .unwrap(); + let original = Arc::new(InMemory::new()); + let data_path = Path::from("table.lance/data/large.lance"); + let data = Bytes::from(vec![7; 4 * block_size]); + original.put(&data_path, data.clone().into()).await.unwrap(); + + let (wrapped, statistics) = wrap_for_test(&cache, original); + let result = wrapped.get(&data_path).await.unwrap(); + assert_eq!(statistics.snapshot(), LanceDataCacheStatistics::default()); + + let mut stream = result.into_stream(); + let first = stream.next().await.unwrap().unwrap(); + assert_eq!(first, data.slice(..block_size)); + assert_eq!( + statistics.snapshot(), + LanceDataCacheStatistics { + bytes_read_from_cache: 0, + bytes_read_from_remote: block_size as u64, + } + ); + + drop(stream); + assert_eq!( + statistics.snapshot().bytes_read_from_remote, + block_size as u64 + ); + } + + #[tokio::test] + async fn caches_single_data_file_range_reads() { + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let original = Arc::new(InMemory::new()); + let data = Bytes::from((0..100_000).map(|value| value as u8).collect::>()); + let cases = [ + ( + Path::from("table.lance/data/bounded.lance"), + GetRange::Bounded(1_000..2_000), + 1_000..2_000, + ), + ( + Path::from("table.lance/data/offset.lance"), + GetRange::Offset(90_000), + 90_000..100_000, + ), + ( + Path::from("table.lance/data/suffix.lance"), + GetRange::Suffix(500), + 99_500..100_000, + ), + ]; + for (path, _, _) in &cases { + original.put(path, data.clone().into()).await.unwrap(); + } + + let (wrapped, statistics) = wrap_for_test(&cache, original); + for (path, requested, expected_range) in cases { + let expected = data.slice(expected_range.start as usize..expected_range.end as usize); + let before = statistics.snapshot(); + let first = wrapped + .get_opts(&path, GetOptions::new().with_range(Some(requested.clone()))) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert_eq!(first, expected); + assert_eq!( + statistics.snapshot().bytes_read_from_remote, + before.bytes_read_from_remote + expected.len() as u64 + ); + + let second = wrapped + .get_opts(&path, GetOptions::new().with_range(Some(requested))) + .await + .unwrap() + .bytes() + .await + .unwrap(); + assert_eq!(second, expected); + assert_eq!( + statistics.snapshot().bytes_read_from_cache, + before.bytes_read_from_cache + expected.len() as u64 + ); + } + + let before = statistics.snapshot(); + let conditional = GetOptions::new().with_if_match(Some("wrong-etag")); + assert!( + wrapped + .get_opts(&Path::from("table.lance/data/bounded.lance"), conditional) + .await + .is_err() + ); + assert_eq!(statistics.snapshot(), before); + + assert!( + wrapped + .get_range( + &Path::from("table.lance/data/bounded.lance"), + 100_000..100_001 + ) + .await + .is_err() + ); + assert_eq!(statistics.snapshot(), before); + } + + #[tokio::test] + async fn recovers_cached_data_from_disk() { + let directory = tempfile::tempdir().unwrap(); + let original = Arc::new(InMemory::new()); + let data_path = Path::from("table.lance/data/part-0.lance"); + let data = Bytes::from((0..100_000).map(|value| value as u8).collect::>()); + original.put(&data_path, data.clone().into()).await.unwrap(); + + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let (wrapped, _) = wrap_for_test(&cache, original.clone()); + let requested_range = 10..90_000; + assert_eq!( + wrapped + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(), + vec![data.slice(10..90_000)] + ); + drop(wrapped); + cache.cache.close().await.unwrap(); + drop(cache); + crate::foyer_cache::wait_for_directory_release(directory.path()).await; + + original.delete(&data_path).await.unwrap(); + let recovered = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let (wrapped, statistics) = wrap_for_test(&recovered, original); + assert_eq!( + wrapped + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(), + vec![data.slice(10..90_000)] + ); + assert_eq!(statistics.snapshot().bytes_read_from_cache, 89_990); + assert_eq!(statistics.snapshot().bytes_read_from_remote, 0); + } + + #[tokio::test] + async fn dataset_scopes_share_cache_without_sharing_statistics() { + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let original = Arc::new(InMemory::new()); + let data_path = Path::from("table.lance/data/part-0.lance"); + let data = Bytes::from((0..100_000).map(|value| value as u8).collect::>()); + original.put(&data_path, data.clone().into()).await.unwrap(); + + let source_scope = cache.create_scope(); + let source_store = source_scope.wrap("memory://test", original.clone()); + + // A restored Dataset is derived from an already-wrapped source + // Dataset. The fresh scope must unwrap to the registered origin rather + // than nesting over the source scope. + let restored_scope = cache.create_scope(); + let restored_store = restored_scope.wrap("memory://test", source_store.clone()); + let requested_range = 10..90_000; + assert_eq!( + restored_store + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(), + vec![data.slice(10..90_000)] + ); + assert_eq!(source_scope.snapshot(), LanceDataCacheStatistics::default()); + assert_eq!(restored_scope.snapshot().bytes_read_from_remote, 89_990); + + original.delete(&data_path).await.unwrap(); + assert_eq!( + source_store + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(), + vec![data.slice(10..90_000)] + ); + assert_eq!(source_scope.snapshot().bytes_read_from_cache, 89_990); + assert_eq!(restored_scope.snapshot().bytes_read_from_remote, 89_990); + + drop(restored_store); + drop(source_store); + assert!(cache.wrapped_stores.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn chained_scopes_drop_intermediate_store_without_deadlocking() { + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let first: Arc = cache.create_scope(); + let second: Arc = cache.create_scope(); + let chained = ChainedWrappingObjectStore::new(vec![first, second]); + let original: Arc = Arc::new(InMemory::new()); + let (sender, receiver) = mpsc::channel(); + + let thread = std::thread::spawn(move || { + sender + .send(chained.wrap("memory://test", original)) + .unwrap(); + }); + let wrapped = receiver + .recv_timeout(Duration::from_secs(2)) + .expect("chained cache wrappers deadlocked while dropping the intermediate store"); + thread.join().unwrap(); + + assert_eq!(cache.wrapped_stores.lock().unwrap().len(), 1); + drop(wrapped); + assert!(cache.wrapped_stores.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn truncates_ranges_at_eof_before_enumerating_cache_blocks() { + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerDataCache::try_new(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + let original = Arc::new(InMemory::new()); + let data_path = Path::from("table.lance/data/part-0.lance"); + let data = Bytes::from((0..100_000).map(|value| value as u8).collect::>()); + original.put(&data_path, data.into()).await.unwrap(); + + let requested_range = 99_990..300_000; + let expected = original + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(); + let (wrapped, statistics) = wrap_for_test(&cache, original.clone()); + let actual = wrapped + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(); + assert_eq!(actual, expected); + assert_eq!(statistics.snapshot().bytes_read_from_remote, 10); + + original.delete(&data_path).await.unwrap(); + assert_eq!( + wrapped + .get_ranges(&data_path, std::slice::from_ref(&requested_range)) + .await + .unwrap(), + expected + ); + assert_eq!(statistics.snapshot().bytes_read_from_cache, 10); + assert_eq!(statistics.snapshot().bytes_read_from_remote, 10); + } + + #[test] + fn recognizes_only_direct_data_children() { + assert!(FoyerDataCache::is_cacheable_data_file(&Path::from( + "dataset/data/part.lance" + ))); + assert!(!FoyerDataCache::is_cacheable_data_file(&Path::from( + "dataset/data/nested/part.lance" + ))); + assert!(!FoyerDataCache::is_cacheable_data_file(&Path::from( + "dataset/indices/index.lance" + ))); + assert!(!FoyerDataCache::is_cacheable_data_file(&Path::from( + "dataset/_versions/1.manifest" + ))); + } +} diff --git a/src/foyer_index_cache.rs b/src/foyer_index_cache.rs new file mode 100644 index 0000000..ab02181 --- /dev/null +++ b/src/foyer_index_cache.rs @@ -0,0 +1,965 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Two-tier cache backend for serializable Lance index entries. +//! +//! L1 is the Session's index memory cache, backed by `QuickCacheBackend`. +//! Metadata uses the Session's separate L1 cache and does not enter this backend. +//! L2 is the shared Foyer disk cache. It stores the Lance `CacheCodec` envelope, +//! so an entry written +//! by an older process is either decoded safely or treated as a miss. + +use std::fmt::{Debug, Formatter}; +use std::fs; +use std::future::Future; +use std::io::Write; +use std::path::{Path as FsPath, PathBuf}; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::RwLock; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +use async_trait::async_trait; +use bytes::Bytes; +use foyer::HybridCache; +use lance_core::Result; +use lance_core::cache::{ + CACHE_KEY_FORMAT, CacheBackend, CacheCodec, CacheDecode, CacheEntry, InternalCacheKey, + QuickCacheBackend, +}; + +use crate::error::ffi_try; +use crate::session::LanceSession; + +const CACHE_KEY_VERSION: &str = "lance-index-v1"; +const LOGICAL_SIZE_BYTES: usize = std::mem::size_of::(); +const INDEX_GENERATION_FILE: &str = "lance-index-generation"; + +/// Cumulative L1 index and L2 serialized-index counters for one shared session. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct LanceIndexDiskCacheStats { + pub memory_hits: u64, + pub memory_misses: u64, + pub disk_hits: u64, + pub disk_misses: u64, + pub disk_read_bytes: u64, + pub disk_write_bytes: u64, + pub decode_failures: u64, + pub disk_read_errors: u64, + pub disk_write_errors: u64, +} + +#[derive(Debug, Default)] +pub(crate) struct IndexDiskCacheStats { + memory_hits: AtomicU64, + memory_misses: AtomicU64, + disk_hits: AtomicU64, + disk_misses: AtomicU64, + disk_read_bytes: AtomicU64, + disk_write_bytes: AtomicU64, + decode_failures: AtomicU64, + disk_read_errors: AtomicU64, + disk_write_errors: AtomicU64, +} + +impl IndexDiskCacheStats { + fn snapshot(&self) -> LanceIndexDiskCacheStats { + LanceIndexDiskCacheStats { + memory_hits: self.memory_hits.load(Ordering::Relaxed), + memory_misses: self.memory_misses.load(Ordering::Relaxed), + disk_hits: self.disk_hits.load(Ordering::Relaxed), + disk_misses: self.disk_misses.load(Ordering::Relaxed), + disk_read_bytes: self.disk_read_bytes.load(Ordering::Relaxed), + disk_write_bytes: self.disk_write_bytes.load(Ordering::Relaxed), + decode_failures: self.decode_failures.load(Ordering::Relaxed), + disk_read_errors: self.disk_read_errors.load(Ordering::Relaxed), + disk_write_errors: self.disk_write_errors.load(Ordering::Relaxed), + } + } +} + +/// Lance index cache with decoded values in L1 and serialized values in L2. +pub(crate) struct FoyerIndexCache { + state: RwLock>, + memory_capacity: usize, + disk: HybridCache, + generation_path: PathBuf, + stats: Arc, +} + +#[derive(Debug)] +struct IndexGeneration { + memory: QuickCacheBackend, + generation: u64, +} + +impl Debug for FoyerIndexCache { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FoyerIndexCache") + .field("state", &self.state) + .field("disk", &self.disk) + .finish_non_exhaustive() + } +} + +impl FoyerIndexCache { + #[cfg(test)] + async fn try_new( + directory: &FsPath, + memory_capacity: usize, + disk_capacity: usize, + storage_block_size: usize, + ) -> std::result::Result { + let disk = + crate::foyer_cache::build_disk_cache(directory, disk_capacity, storage_block_size) + .await?; + Self::from_cache(disk, memory_capacity, directory) + } + + pub(crate) fn from_cache( + disk: HybridCache, + memory_capacity: usize, + directory: &FsPath, + ) -> std::result::Result { + let generation = load_generation(directory)?; + Ok(Self { + state: RwLock::new(Arc::new(IndexGeneration { + memory: QuickCacheBackend::with_capacity(memory_capacity), + generation, + })), + memory_capacity, + disk, + generation_path: generation_path(directory), + stats: Arc::new(IndexDiskCacheStats::default()), + }) + } + + pub(crate) fn stats(&self) -> Arc { + self.stats.clone() + } + + #[cfg(test)] + fn disk_key(key: &InternalCacheKey, codec: CacheCodec) -> String { + Self::disk_key_for_generation(key, codec, 0) + } + + fn disk_key_for_generation( + key: &InternalCacheKey, + codec: CacheCodec, + generation: u64, + ) -> String { + let mut encoded = String::with_capacity(32); + for byte in key.as_bytes() { + use std::fmt::Write; + let _ = write!(encoded, "{byte:02x}"); + } + format!( + "{CACHE_KEY_VERSION}:{CACHE_KEY_FORMAT}:{generation}:{}:{encoded}", + codec.type_id(), + ) + } + + fn state(&self) -> Arc { + self.state + .read() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } + + fn encode(codec: CacheCodec, entry: &CacheEntry, logical_size: usize) -> Result { + let mut bytes = Vec::new(); + bytes.extend_from_slice(&(logical_size as u64).to_le_bytes()); + // Vec's Write implementation appends after the size prefix. A fresh + // Cursor starts at zero and would overwrite that prefix with the codec. + codec.serialize(entry, &mut bytes)?; + Ok(Bytes::from(bytes)) + } + + fn decode(codec: CacheCodec, bytes: &Bytes) -> Option<(CacheEntry, usize)> { + if bytes.len() < LOGICAL_SIZE_BYTES { + return None; + } + let size = u64::from_le_bytes(bytes[..LOGICAL_SIZE_BYTES].try_into().ok()?); + let size = usize::try_from(size).ok()?; + let payload = bytes.slice(LOGICAL_SIZE_BYTES..); + match codec.deserialize(&payload) { + CacheDecode::Hit(entry) => Some((entry, size)), + CacheDecode::Miss(_) => None, + } + } + + async fn load_disk( + disk: &HybridCache, + stats: &IndexDiskCacheStats, + key: &str, + codec: CacheCodec, + ) -> Option<(CacheEntry, usize)> { + let entry = match disk.get(key).await { + Ok(Some(entry)) => entry, + Ok(None) => { + stats.disk_misses.fetch_add(1, Ordering::Relaxed); + return None; + } + Err(error) => { + stats.disk_misses.fetch_add(1, Ordering::Relaxed); + stats.disk_read_errors.fetch_add(1, Ordering::Relaxed); + log::warn!("Foyer index cache lookup failed for {key}: {error}"); + return None; + } + }; + let bytes = entry.value(); + stats + .disk_read_bytes + .fetch_add(bytes.len() as u64, Ordering::Relaxed); + match Self::decode(codec, bytes) { + Some(decoded) => { + stats.disk_hits.fetch_add(1, Ordering::Relaxed); + Some(decoded) + } + None => { + stats.disk_misses.fetch_add(1, Ordering::Relaxed); + stats.decode_failures.fetch_add(1, Ordering::Relaxed); + disk.remove(key); + None + } + } + } + + fn store_disk( + disk: &HybridCache, + stats: &IndexDiskCacheStats, + key: &str, + codec: CacheCodec, + entry: &CacheEntry, + logical_size: usize, + ) { + match Self::encode(codec, entry, logical_size) { + Ok(bytes) => { + stats + .disk_write_bytes + .fetch_add(bytes.len() as u64, Ordering::Relaxed); + disk.insert(key.to_owned(), bytes); + } + Err(error) => { + stats.disk_write_errors.fetch_add(1, Ordering::Relaxed); + log::warn!("failed to serialize Lance index cache entry {key}: {error}"); + } + } + } +} + +#[async_trait] +impl CacheBackend for FoyerIndexCache { + async fn get(&self, key: &InternalCacheKey, codec: Option) -> Option { + let state = self.state(); + if let Some(entry) = state.memory.get(key, None).await { + self.stats.memory_hits.fetch_add(1, Ordering::Relaxed); + return Some(entry); + } + self.stats.memory_misses.fetch_add(1, Ordering::Relaxed); + let codec = codec?; + let disk_key = Self::disk_key_for_generation(key, codec, state.generation); + let (entry, size) = Self::load_disk(&self.disk, &self.stats, &disk_key, codec).await?; + state.memory.insert(key, entry.clone(), size, None).await; + Some(entry) + } + + async fn insert( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + size_bytes: usize, + codec: Option, + ) { + let state = self.state(); + state + .memory + .insert(key, entry.clone(), size_bytes, None) + .await; + if let Some(codec) = codec { + let disk_key = Self::disk_key_for_generation(key, codec, state.generation); + Self::store_disk( + &self.disk, + &self.stats, + &disk_key, + codec, + &entry, + size_bytes, + ); + } + } + + async fn get_or_insert<'a>( + &self, + key: &InternalCacheKey, + loader: Pin> + Send + 'a>>, + codec: Option, + ) -> Result<(CacheEntry, bool)> { + let key = *key; + let disk = self.disk.clone(); + let stats = self.stats.clone(); + let disk_codec = codec; + let state = self.state(); + let generation = state.generation; + let disk_hit = AtomicBool::new(false); + let disk_hit_ref = &disk_hit; + let loader = async move { + if let Some(codec) = disk_codec { + let disk_key = Self::disk_key_for_generation(&key, codec, generation); + if let Some((entry, size)) = Self::load_disk(&disk, &stats, &disk_key, codec).await + { + disk_hit_ref.store(true, Ordering::Relaxed); + return Ok((entry, size)); + } + let result = loader.await?; + Self::store_disk(&disk, &stats, &disk_key, codec, &result.0, result.1); + Ok(result) + } else { + loader.await + } + }; + let result = state + .memory + .get_or_insert(&key, Box::pin(loader), None) + .await; + match &result { + Ok((_, true)) => { + self.stats.memory_hits.fetch_add(1, Ordering::Relaxed); + } + Ok((_, false)) | Err(_) => { + self.stats.memory_misses.fetch_add(1, Ordering::Relaxed); + } + } + // The inner result describes L1 only; the CacheBackend contract + // describes whether the original loader was skipped across both tiers. + result.map(|(entry, memory_hit)| (entry, memory_hit || disk_hit.load(Ordering::Relaxed))) + } + + async fn clear(&self) { + let mut state = self + .state + .write() + .unwrap_or_else(|error| error.into_inner()); + let next_generation = new_generation(); + if let Err(error) = persist_generation_file(&self.generation_path, next_generation) { + log::warn!( + "failed to persist Foyer index cache generation; using a fresh in-process namespace: {error}" + ); + } + // Replace L1 with the namespace. In-flight loads retain the old state + // and cannot repopulate the new memory cache or its disk namespace. + *state = Arc::new(IndexGeneration { + memory: QuickCacheBackend::with_capacity(self.memory_capacity), + generation: next_generation, + }); + } + + async fn num_entries(&self) -> usize { + self.state().memory.num_entries().await + } + + async fn size_bytes(&self) -> usize { + self.state().memory.size_bytes().await + } + + fn approx_num_entries(&self) -> usize { + self.state().memory.approx_num_entries() + } + + fn approx_size_bytes(&self) -> usize { + self.state().memory.approx_size_bytes() + } + + fn deep_size_of_entries( + &self, + context: &mut lance_core::deepsize::Context, + size_of_entry: &dyn Fn(&CacheEntry, &mut lance_core::deepsize::Context) -> Option, + ) -> Option { + self.state() + .memory + .deep_size_of_entries(context, size_of_entry) + } +} + +fn generation_path(directory: &FsPath) -> PathBuf { + directory.join(INDEX_GENERATION_FILE) +} + +fn new_generation() -> u64 { + uuid::Uuid::new_v4().as_u64_pair().0 +} + +fn load_generation(directory: &FsPath) -> std::result::Result { + let path = generation_path(directory); + match fs::read_to_string(&path) { + Ok(value) => match value.trim().parse::() { + Ok(generation) => Ok(generation), + Err(error) => { + log::warn!("invalid index cache generation in {path:?}; starting cold: {error}"); + let generation = new_generation(); + persist_generation_file(&path, generation)?; + Ok(generation) + } + }, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + let generation = new_generation(); + persist_generation_file(&path, generation)?; + Ok(generation) + } + Err(error) => Err(foyer::Error::io_error(error)), + } +} + +fn persist_generation_file( + path: &FsPath, + generation: u64, +) -> std::result::Result<(), foyer::Error> { + let temporary = path.with_file_name(format!( + "{INDEX_GENERATION_FILE}.{}.tmp", + uuid::Uuid::new_v4() + )); + let result = (|| { + let mut file = fs::OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary)?; + file.write_all(generation.to_string().as_bytes())?; + file.sync_all()?; + fs::rename(&temporary, path) + })(); + if result.is_err() { + let _ = fs::remove_file(&temporary); + } + result.map_err(foyer::Error::io_error) +} + +/// Copy the disk-tier counters into `out_stats`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_session_get_index_disk_cache_stats( + session: *const LanceSession, + out_stats: *mut LanceIndexDiskCacheStats, +) -> i32 { + ffi_try!( + unsafe { index_disk_cache_stats_inner(session, out_stats) }, + neg + ) +} + +unsafe fn index_disk_cache_stats_inner( + session: *const LanceSession, + out_stats: *mut LanceIndexDiskCacheStats, +) -> Result { + if session.is_null() || out_stats.is_null() { + return Err(lance_core::Error::invalid_input_source( + "session and out_stats must not be NULL".into(), + )); + } + let session = unsafe { &*session }; + let stats = session + .index_disk_cache_stats + .as_ref() + .map(|stats| stats.snapshot()) + .unwrap_or_default(); + unsafe { + std::ptr::write_unaligned(out_stats, stats); + } + Ok(0) +} + +#[cfg(test)] +mod tests { + use tokio::sync::oneshot; + + use super::*; + + fn test_codec() -> CacheCodec { + CacheCodec::new( + "lance.test.U64", + 1, + |value, writer| writer.write_raw(&value.downcast_ref::().unwrap().to_le_bytes()), + |reader| { + let raw = reader.read_raw()?; + let bytes = raw.as_ref().try_into().map_err(|_| { + lance_core::Error::io(format!("expected 8 test entry bytes, got {}", raw.len())) + })?; + Ok(Arc::new(u64::from_le_bytes(bytes))) + }, + ) + } + + async fn test_cache(directory: &FsPath) -> FoyerIndexCache { + FoyerIndexCache::try_new(directory, 1024 * 1024, 1024 * 1024, 65536) + .await + .unwrap() + } + + #[tokio::test] + async fn disk_hits_preserve_layer_statistics_and_report_cached() { + let directory = tempfile::tempdir().unwrap(); + let cache = test_cache(directory.path()).await; + let key = InternalCacheKey::from_bytes([1; 16]); + let codec = Some(test_codec()); + cache.insert(&key, Arc::new(42_u64), 8, codec).await; + cache.state().memory.clear().await; + + for _ in 0..2 { + let (value, was_cached) = cache + .get_or_insert( + &key, + Box::pin(async { panic!("cache hit must not execute source loader") }), + codec, + ) + .await + .unwrap(); + assert_eq!(*value.downcast_ref::().unwrap(), 42); + assert!(was_cached); + } + let stats = cache.stats.snapshot(); + assert_eq!((stats.memory_hits, stats.memory_misses), (1, 1)); + assert_eq!((stats.disk_hits, stats.disk_misses), (1, 0)); + + let missing = InternalCacheKey::from_bytes([2; 16]); + let (_, was_cached) = cache + .get_or_insert( + &missing, + Box::pin(async { Ok((Arc::new(7_u64) as CacheEntry, 8)) }), + codec, + ) + .await + .unwrap(); + assert!(!was_cached); + let stats = cache.stats.snapshot(); + assert_eq!((stats.memory_hits, stats.memory_misses), (1, 2)); + assert_eq!((stats.disk_hits, stats.disk_misses), (1, 1)); + } + + #[tokio::test] + async fn persisted_l2_hits_promote_to_l1_before_any_further_disk_lookup() { + let directory = tempfile::tempdir().unwrap(); + let cache = test_cache(directory.path()).await; + let keys = [ + InternalCacheKey::from_bytes([21; 16]), + InternalCacheKey::from_bytes([22; 16]), + ]; + for key in &keys { + cache + .insert(key, Arc::new(42_u64), 8, Some(test_codec())) + .await; + } + cache.disk.close().await.unwrap(); + drop(cache); + crate::foyer_cache::wait_for_directory_release(directory.path()).await; + + // Both capacities are nonzero, but neither L1 nor Foyer's pending-write + // buffers survive reopening. A hit must read a persisted disk entry. + let cache = test_cache(directory.path()).await; + assert_eq!(cache.num_entries().await, 0); + for (index, key) in keys.iter().enumerate() { + let reads = cache.disk.storage().statistics().disk_read_ios(); + let value = if index == 0 { + cache.get(key, Some(test_codec())).await.unwrap() + } else { + let (value, cached) = cache + .get_or_insert( + key, + Box::pin(async { panic!("persisted L2 hit must skip the source") }), + Some(test_codec()), + ) + .await + .unwrap(); + assert!(cached); + value + }; + assert_eq!(*value.downcast_ref::().unwrap(), 42); + assert!(cache.disk.storage().statistics().disk_read_ios() > reads); + assert!(cache.state().memory.get(key, None).await.is_some()); + + // Remove the lower-tier copy: both read APIs must now use L1, + // without inspecting L2 or polling their source loader. + let physical = FoyerIndexCache::disk_key_for_generation( + key, + test_codec(), + cache.state().generation, + ); + cache.disk.remove(&physical); + cache.disk.storage().wait().await; + let before = cache.stats.snapshot(); + let reads = cache.disk.storage().statistics().disk_read_ios(); + assert_eq!( + *cache + .get(key, Some(test_codec())) + .await + .unwrap() + .downcast_ref::() + .unwrap(), + 42 + ); + let (value, cached) = cache + .get_or_insert( + key, + Box::pin(async { panic!("L1 hit must skip the source") }), + Some(test_codec()), + ) + .await + .unwrap(); + assert!(cached); + assert_eq!(*value.downcast_ref::().unwrap(), 42); + let after = cache.stats.snapshot(); + assert_eq!(after.memory_hits, before.memory_hits + 2); + assert_eq!(after.memory_misses, before.memory_misses); + assert_eq!(after.disk_hits, before.disk_hits); + assert_eq!(after.disk_misses, before.disk_misses); + assert_eq!(cache.disk.storage().statistics().disk_read_ios(), reads); + } + } + + #[tokio::test] + async fn l1_capacity_eviction_preserves_l2_entries() { + const L1_CAPACITY: usize = 1024; + let directory = tempfile::tempdir().unwrap(); + let cache = FoyerIndexCache::try_new(directory.path(), L1_CAPACITY, 4 * 1024 * 1024, 65536) + .await + .unwrap(); + let keys = (0..128_u8) + .map(|value| InternalCacheKey::from_bytes([value; 16])) + .collect::>(); + for (value, key) in keys.iter().enumerate() { + cache + .insert(key, Arc::new(value as u64), 8, Some(test_codec())) + .await; + } + cache.disk.storage().wait().await; + assert!(cache.num_entries().await > 0); + assert!(cache.num_entries().await < keys.len()); + assert!(cache.size_bytes().await <= L1_CAPACITY); + + // Find a real capacity eviction; don't manually clear or disable L1. + let mut evicted = None; + for (value, key) in keys.iter().enumerate() { + if cache.state().memory.get(key, None).await.is_none() { + evicted = Some((value as u64, key)); + break; + } + } + let (expected, key) = evicted.expect("L1 budget must evict some entries"); + let (value, cached) = cache + .get_or_insert( + key, + Box::pin(async { panic!("L1 eviction must not discard the L2 copy") }), + Some(test_codec()), + ) + .await + .unwrap(); + assert!(cached); + assert_eq!(*value.downcast_ref::().unwrap(), expected); + assert_eq!(cache.stats.snapshot().disk_hits, 1); + assert!(cache.state().memory.get(key, None).await.is_some()); + assert!(cache.size_bytes().await <= L1_CAPACITY); + } + + #[tokio::test] + async fn entries_larger_than_l1_capacity_remain_readable_from_disk() { + let directory = tempfile::tempdir().unwrap(); + let key = InternalCacheKey::from_bytes([23; 16]); + let codec = CacheCodec::new( + "lance.test.Bytes", + 1, + |value, writer| writer.write_raw(value.downcast_ref::>().unwrap()), + |reader| Ok(Arc::new(reader.read_raw()?.to_vec())), + ); + let value = vec![23_u8; 2048]; + let cache = FoyerIndexCache::try_new(directory.path(), 1024, 1024 * 1024, 65536) + .await + .unwrap(); + cache + .insert(&key, Arc::new(value.clone()), value.len(), Some(codec)) + .await; + assert_eq!(cache.num_entries().await, 0); + cache.disk.close().await.unwrap(); + drop(cache); + crate::foyer_cache::wait_for_directory_release(directory.path()).await; + + let cache = FoyerIndexCache::try_new(directory.path(), 1024, 1024 * 1024, 65536) + .await + .unwrap(); + for _ in 0..2 { + let (entry, cached) = cache + .get_or_insert( + &key, + Box::pin(async { panic!("an entry refused by L1 can still hit L2") }), + Some(codec), + ) + .await + .unwrap(); + assert!(cached); + assert_eq!(entry.downcast_ref::>().unwrap(), &value); + assert_eq!(cache.size_bytes().await, 0); + } + let stats = cache.stats.snapshot(); + assert_eq!((stats.memory_hits, stats.memory_misses), (0, 2)); + assert_eq!((stats.disk_hits, stats.disk_misses), (2, 0)); + } + + #[tokio::test] + async fn shared_l2_capacity_eviction_preserves_l1_and_allows_source_refill() { + let directory = tempfile::tempdir().unwrap(); + let cache = test_cache(directory.path()).await; + let key = InternalCacheKey::from_bytes([24; 16]); + let physical = + FoyerIndexCache::disk_key_for_generation(&key, test_codec(), cache.state().generation); + cache + .insert(&key, Arc::new(42_u64), 8, Some(test_codec())) + .await; + cache.disk.storage().wait().await; + assert!(cache.disk.get(&physical).await.unwrap().is_some()); + + // Data and index entries compete for the same 1 MiB disk budget. + // Drain each write so admission buffering cannot mask disk eviction. + for block in 0..128 { + cache.disk.insert( + format!("lance-data-v1\0pressure\0{block}"), + Bytes::from(vec![0_u8; 32768]), + ); + cache.disk.storage().wait().await; + } + assert!(cache.disk.get(&physical).await.unwrap().is_none()); + let before = cache.stats.snapshot(); + let (entry, cached) = cache + .get_or_insert( + &key, + Box::pin(async { panic!("L2 eviction must not invalidate L1") }), + Some(test_codec()), + ) + .await + .unwrap(); + assert!(cached); + assert_eq!(*entry.downcast_ref::().unwrap(), 42); + assert_eq!(cache.stats.snapshot().disk_misses, before.disk_misses); + + cache.state().memory.clear().await; + let loads = AtomicU64::new(0); + for expected_cached in [false, true] { + let (entry, cached) = cache + .get_or_insert( + &key, + Box::pin(async { + loads.fetch_add(1, Ordering::Relaxed); + Ok((Arc::new(42_u64) as CacheEntry, 8)) + }), + Some(test_codec()), + ) + .await + .unwrap(); + assert_eq!(cached, expected_cached); + assert_eq!(*entry.downcast_ref::().unwrap(), 42); + } + assert_eq!(loads.load(Ordering::Relaxed), 1); + assert_eq!(cache.stats.snapshot().disk_misses, before.disk_misses + 1); + let allocated = fs::read_dir(directory.path()) + .unwrap() + .map(|entry| entry.unwrap()) + .filter(|entry| { + entry + .file_name() + .to_string_lossy() + .starts_with("foyer-storage-direct-fs-") + }) + .map(|entry| entry.metadata().unwrap().len()) + .sum::(); + assert!(allocated > 0 && allocated <= 1024 * 1024); + } + + #[tokio::test] + async fn concurrent_requests_share_one_source_load() { + let directory = tempfile::tempdir().unwrap(); + let cache = test_cache(directory.path()).await; + let key = InternalCacheKey::from_bytes([7; 16]); + let loads = AtomicU64::new(0); + let results = futures::future::join_all((0..8).map(|_| { + cache.get_or_insert( + &key, + Box::pin(async { + loads.fetch_add(1, Ordering::Relaxed); + tokio::task::yield_now().await; + Ok((Arc::new(42_u64) as CacheEntry, 8)) + }), + Some(test_codec()), + ) + })) + .await; + assert_eq!(loads.load(Ordering::Relaxed), 1); + let hits = results + .into_iter() + .map(|result| result.unwrap()) + .filter(|(_, hit)| *hit) + .count(); + assert_eq!(hits, 7); + let stats = cache.stats.snapshot(); + assert_eq!((stats.memory_hits, stats.memory_misses), (7, 1)); + assert_eq!((stats.disk_hits, stats.disk_misses), (0, 1)); + } + + #[tokio::test] + async fn malformed_entries_miss_and_entries_without_codecs_stay_in_memory() { + let directory = tempfile::tempdir().unwrap(); + let cache = test_cache(directory.path()).await; + let key = InternalCacheKey::from_bytes([3; 16]); + let physical = + FoyerIndexCache::disk_key_for_generation(&key, test_codec(), cache.state().generation); + cache.disk.insert(physical, Bytes::from_static(b"bad")); + assert!(cache.get(&key, Some(test_codec())).await.is_none()); + assert_eq!(cache.stats.snapshot().decode_failures, 1); + + let before = cache.stats.snapshot(); + cache.insert(&key, Arc::new(9_u64), 8, None).await; + assert!(cache.get(&key, None).await.is_some()); + cache.state().memory.clear().await; + assert!(cache.get(&key, None).await.is_none()); + let after = cache.stats.snapshot(); + assert_eq!(after.disk_write_bytes, before.disk_write_bytes); + assert_eq!(after.disk_misses, before.disk_misses); + } + + #[tokio::test] + async fn clear_isolates_in_flight_loads() { + let directory = tempfile::tempdir().unwrap(); + let cache = Arc::new(test_cache(directory.path()).await); + let key = InternalCacheKey::from_bytes([4; 16]); + let (started_tx, started_rx) = oneshot::channel(); + let (resume_tx, resume_rx) = oneshot::channel(); + let loading = cache.clone(); + let task = tokio::spawn(async move { + loading + .get_or_insert( + &key, + Box::pin(async move { + started_tx.send(()).unwrap(); + resume_rx.await.unwrap(); + Ok((Arc::new(17_u64) as CacheEntry, 8)) + }), + Some(test_codec()), + ) + .await + .unwrap() + }); + started_rx.await.unwrap(); + cache.clear().await; + cache + .insert(&key, Arc::new(42_u64), 8, Some(test_codec())) + .await; + resume_tx.send(()).unwrap(); + assert_eq!(*task.await.unwrap().0.downcast_ref::().unwrap(), 17); + assert_eq!( + *cache + .get(&key, Some(test_codec())) + .await + .unwrap() + .downcast_ref::() + .unwrap(), + 42 + ); + cache.state().memory.clear().await; + assert_eq!( + *cache + .get(&key, Some(test_codec())) + .await + .unwrap() + .downcast_ref::() + .unwrap(), + 42 + ); + } + + #[tokio::test] + async fn clear_invalidates_even_if_generation_cannot_be_persisted() { + let directory = tempfile::tempdir().unwrap(); + let mut cache = test_cache(directory.path()).await; + let key = InternalCacheKey::from_bytes([5; 16]); + cache + .insert(&key, Arc::new(7_u64), 8, Some(test_codec())) + .await; + cache.generation_path = directory.path().join("missing-parent/generation"); + cache.clear().await; + assert!(cache.get(&key, Some(test_codec())).await.is_none()); + } + + #[tokio::test] + async fn recovers_index_entries_and_persisted_invalidation() { + let directory = tempfile::tempdir().unwrap(); + let key = InternalCacheKey::from_bytes([6; 16]); + let cache = test_cache(directory.path()).await; + cache + .insert(&key, Arc::new(99_u64), 8, Some(test_codec())) + .await; + cache.disk.close().await.unwrap(); + drop(cache); + crate::foyer_cache::wait_for_directory_release(directory.path()).await; + + let recovered = test_cache(directory.path()).await; + assert_eq!( + *recovered + .get(&key, Some(test_codec())) + .await + .unwrap() + .downcast_ref::() + .unwrap(), + 99 + ); + assert_eq!(recovered.stats.snapshot().disk_hits, 1); + recovered.clear().await; + recovered.disk.close().await.unwrap(); + drop(recovered); + crate::foyer_cache::wait_for_directory_release(directory.path()).await; + + let cleared = test_cache(directory.path()).await; + assert!(cleared.get(&key, Some(test_codec())).await.is_none()); + } + + #[test] + fn corrupt_generation_starts_a_new_namespace() { + let directory = tempfile::tempdir().unwrap(); + let path = generation_path(directory.path()); + fs::write(&path, "truncated").unwrap(); + let generation = load_generation(directory.path()).unwrap(); + assert_eq!(load_generation(directory.path()).unwrap(), generation); + assert_eq!(fs::read_to_string(path).unwrap(), generation.to_string()); + } + + #[test] + fn codec_type_is_part_of_the_physical_key() { + let key = + InternalCacheKey::from_bytes([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]); + let codec = CacheCodec::new( + "lance.test.IndexEntry", + 1, + |_value, _writer| Ok(()), + |_reader| unreachable!(), + ); + let physical_key = FoyerIndexCache::disk_key(&key, codec); + assert!(physical_key.starts_with("lance-index-v1:blake3-128-v1:0:lance.test.IndexEntry:")); + assert!(!physical_key.contains('\0')); + assert_ne!( + physical_key, + FoyerIndexCache::disk_key_for_generation(&key, codec, 1) + ); + } + + #[tokio::test] + async fn clear_does_not_clear_shared_data_entries() { + let directory = tempfile::tempdir().unwrap(); + let disk = crate::foyer_cache::build_disk_cache(directory.path(), 1024 * 1024, 64 * 1024) + .await + .unwrap(); + disk.insert( + "lance-data-v1\0test".to_owned(), + Bytes::from_static(b"data-entry"), + ); + let index = + FoyerIndexCache::from_cache(disk.clone(), 1024 * 1024, directory.path()).unwrap(); + + index.clear().await; + + assert!(disk.get("lance-data-v1\0test").await.unwrap().is_some()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 3481760..958e822 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,11 +26,15 @@ mod async_dispatcher; mod batch; mod blob; mod compact; +mod data_cache; mod data_statistics; mod dataset; mod delete; mod drop_columns; mod error; +mod foyer_cache; +mod foyer_data_cache; +mod foyer_index_cache; mod fragment_writer; mod fts_query; mod helpers; @@ -55,6 +59,7 @@ pub use alter_columns::*; pub use batch::*; pub use blob::*; pub use compact::*; +pub use data_cache::{LanceDataCacheStatistics, lance_dataset_get_data_cache_statistics}; pub use data_statistics::*; pub use dataset::*; pub use delete::*; @@ -62,6 +67,8 @@ pub use drop_columns::*; pub use error::{ LanceErrorCode, lance_free_string, lance_last_error_code, lance_last_error_message, }; +pub use foyer_cache::{LanceFoyerCacheOptions, lance_session_new_with_foyer_cache}; +pub use foyer_index_cache::{LanceIndexDiskCacheStats, lance_session_get_index_disk_cache_stats}; pub use fragment_writer::*; pub use fts_query::*; pub use index::*; diff --git a/src/restore.rs b/src/restore.rs index 7804b55..fa2d26d 100644 --- a/src/restore.rs +++ b/src/restore.rs @@ -65,8 +65,16 @@ unsafe fn restore_inner(dataset: *const LanceDataset, version: u64) -> Result<*m Ok::<_, lance_core::Error>(checked_out) })?; + let (restored, data_cache) = if let Some(data_cache) = &ds.data_cache { + let (restored, data_cache) = data_cache.attach_fresh(restored); + (restored, Some(data_cache)) + } else { + (restored, None) + }; + let handle = LanceDataset { inner: RwLock::new(Arc::new(restored)), + data_cache, }; Ok(Box::into_raw(Box::new(handle))) } diff --git a/src/session.rs b/src/session.rs index 60a1623..b1376d9 100644 --- a/src/session.rs +++ b/src/session.rs @@ -7,13 +7,18 @@ use std::sync::Arc; use lance::session::Session; use lance_core::Result; +use lance_core::cache::CacheBackend; +use crate::data_cache::DataCacheFactory; use crate::error::{ffi_try, swallow_unwind}; +use crate::foyer_index_cache::IndexDiskCacheStats; use crate::runtime::block_on; -/// Opaque handle for sharing Lance metadata and index caches across datasets. +/// Opaque handle for shared Lance caches across datasets. pub struct LanceSession { pub(crate) inner: Arc, + pub(crate) data_cache_factory: Option>, + pub(crate) index_disk_cache_stats: Option>, } /// Snapshot of a session's metadata and index cache statistics. @@ -47,17 +52,42 @@ pub extern "C" fn lance_session_new( fn session_new_inner( index_cache_size_bytes: u64, metadata_cache_size_bytes: u64, +) -> Result<*mut LanceSession> { + session_new_with_factories( + index_cache_size_bytes, + metadata_cache_size_bytes, + None, + None, + None, + ) +} + +pub(crate) fn session_new_with_factories( + index_cache_size_bytes: u64, + metadata_cache_size_bytes: u64, + index_cache_backend: Option>, + index_disk_cache_stats: Option>, + data_cache_factory: Option>, ) -> Result<*mut LanceSession> { let index_cache_size_bytes = u64_to_usize(index_cache_size_bytes, "index_cache_size_bytes")?; let metadata_cache_size_bytes = u64_to_usize(metadata_cache_size_bytes, "metadata_cache_size_bytes")?; - let session = Session::new( - index_cache_size_bytes, - metadata_cache_size_bytes, - Default::default(), - ); + let session = match index_cache_backend { + Some(backend) => Session::with_index_cache_backend( + backend, + metadata_cache_size_bytes, + Default::default(), + ), + None => Session::new( + index_cache_size_bytes, + metadata_cache_size_bytes, + Default::default(), + ), + }; Ok(Box::into_raw(Box::new(LanceSession { inner: Arc::new(session), + data_cache_factory, + index_disk_cache_stats, }))) } @@ -119,7 +149,7 @@ unsafe fn session_get_cache_stats_inner( Ok(0) } -fn u64_to_usize(value: u64, field: &'static str) -> Result { +pub(crate) fn u64_to_usize(value: u64, field: &'static str) -> Result { usize::try_from(value).map_err(|_| { lance_core::Error::invalid_input_source( format!("{field}={value} exceeds usize::MAX on this target").into(), diff --git a/src/writer.rs b/src/writer.rs index 1971510..ba51c87 100644 --- a/src/writer.rs +++ b/src/writer.rs @@ -282,6 +282,7 @@ unsafe fn write_dataset_inner( if !out_dataset.is_null() { let handle = LanceDataset { inner: RwLock::new(Arc::new(dataset)), + data_cache: None, }; // SAFETY: `out_dataset` is non-NULL (checked above) and the caller // guarantees it points to caller-owned, writable storage of size diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 0524210..23d6552 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -100,10 +100,81 @@ fn create_large_dataset(num_rows: i32) -> (tempfile::TempDir, String) { (tmp, uri) } +/// Helper: create two fragments large enough for Lance's batched range-read +/// path, which is the path wrapped by the Foyer data cache. +fn create_large_multi_fragment_dataset(num_rows_per_fragment: i32) -> (tempfile::TempDir, String) { + let (tmp, uri) = create_large_dataset(num_rows_per_fragment); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Float32, true), + Field::new("label", DataType::Utf8, true), + ])); + let ids: Vec = (num_rows_per_fragment..2 * num_rows_per_fragment).collect(); + let values: Vec = ids.iter().map(|id| *id as f32 * 0.5).collect(); + let labels: Vec = ids.iter().map(|id| format!("row_{id}")).collect(); + let label_refs: Vec<&str> = labels.iter().map(String::as_str).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(Float32Array::from(values)), + Arc::new(StringArray::from(label_refs)), + ], + ) + .unwrap(); + + lance_c::runtime::block_on(async { + let mut dataset = Dataset::open(&uri).await.unwrap(); + dataset + .append( + arrow::record_batch::RecordBatchIterator::new(vec![Ok(batch)], schema), + None, + ) + .await + .unwrap(); + }); + + (tmp, uri) +} + fn c_str(s: &str) -> CString { CString::new(s).unwrap() } +fn file_object_store_uri(path: &str) -> CString { + let path = path.replace('\\', "/"); + let leading_slash = if path.starts_with('/') { "" } else { "/" }; + c_str(&format!("file-object-store://{leading_slash}{path}")) +} + +fn create_data_cache_session() -> (tempfile::TempDir, *mut LanceSession) { + let directory = tempfile::tempdir().unwrap(); + let c_directory = c_str(directory.path().to_str().unwrap()); + let options = LanceFoyerCacheOptions { + directory: c_directory.as_ptr(), + disk_capacity_bytes: 64 * 1024 * 1024, + }; + let session = unsafe { lance_session_new_with_foyer_cache(0, 16 * 1024 * 1024, &options) }; + assert!(!session.is_null(), "data-cache session should be created"); + (directory, session) +} + +fn data_cache_statistics(dataset: *const LanceDataset) -> LanceDataCacheStatistics { + let mut statistics = LanceDataCacheStatistics::default(); + assert_eq!( + unsafe { lance_dataset_get_data_cache_statistics(dataset, &mut statistics) }, + 0 + ); + statistics +} + +fn scanned_row_count(dataset: *const LanceDataset) -> usize { + scan_all_rows(dataset) + .iter() + .map(RecordBatch::num_rows) + .sum() +} + #[derive(Default)] struct CapturedScanStatistics { calls: usize, @@ -445,6 +516,275 @@ fn test_shared_session_rejects_null_inputs() { } } +#[test] +fn test_session_with_foyer_cache_validates_options() { + let session = unsafe { lance_session_new_with_foyer_cache(0, 0, ptr::null()) }; + assert!(!session.is_null()); + let mut stats = LanceIndexDiskCacheStats::default(); + assert_eq!( + unsafe { lance_session_get_index_disk_cache_stats(session, &mut stats) }, + 0 + ); + assert_eq!(stats, LanceIndexDiskCacheStats::default()); + assert_eq!( + unsafe { lance_session_get_index_disk_cache_stats(session, ptr::null_mut()) }, + -1 + ); + assert_eq!( + unsafe { lance_session_get_index_disk_cache_stats(ptr::null(), &mut stats) }, + -1 + ); + unsafe { lance_session_close(session) }; + + let directory = tempfile::tempdir().unwrap(); + let c_directory = c_str(directory.path().to_str().unwrap()); + let empty = c_str(""); + for (path, capacity) in [ + (ptr::null(), 64 * 1024 * 1024), + (empty.as_ptr(), 64 * 1024 * 1024), + (c_directory.as_ptr(), 0), + (c_directory.as_ptr(), 64 * 1024 * 1024 - 4096), + (c_directory.as_ptr(), 64 * 1024 * 1024 + 1), + ] { + let options = LanceFoyerCacheOptions { + directory: path, + disk_capacity_bytes: capacity, + }; + assert!(unsafe { lance_session_new_with_foyer_cache(0, 0, &options) }.is_null()); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + } + + let options = LanceFoyerCacheOptions { + directory: c_directory.as_ptr(), + disk_capacity_bytes: 64 * 1024 * 1024, + }; + let session = unsafe { lance_session_new_with_foyer_cache(0, 0, &options) }; + assert!(!session.is_null()); + assert!(unsafe { lance_session_new_with_foyer_cache(0, 0, &options) }.is_null()); + unsafe { lance_session_close(session) }; +} + +#[test] +fn test_session_with_foyer_cache_shares_data_and_index_across_fragments() { + let (_tmp, uri) = create_large_multi_fragment_dataset(10_000); + let c_uri = file_object_store_uri(&uri); + let source = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!source.is_null()); + let column = c_str("id"); + assert_eq!( + unsafe { + lance_dataset_create_scalar_index( + source, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::BTree as i32, + ptr::null(), + false, + ) + }, + 0 + ); + unsafe { lance_dataset_close(source) }; + + let query = |dataset: *const LanceDataset| { + let filter = c_str("id = 3 OR id = 10003"); + let scanner = unsafe { lance_scanner_new(dataset, ptr::null(), filter.as_ptr()) }; + assert!(!scanner.is_null()); + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0 + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.unwrap(); + let mut ids = reader + .flat_map(|batch| { + batch + .unwrap() + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![3, 10003]); + unsafe { lance_scanner_close(scanner) }; + }; + + let disk_stats = |session: *const LanceSession| { + let mut stats = LanceIndexDiskCacheStats::default(); + assert_eq!( + unsafe { lance_session_get_index_disk_cache_stats(session, &mut stats) }, + 0 + ); + stats + }; + let session_stats = |session: *const LanceSession| { + let mut stats = LanceSessionCacheStats::default(); + assert_eq!( + unsafe { lance_session_get_cache_stats(session, &mut stats) }, + 0 + ); + stats + }; + + for l1_capacity in [0, 16 * 1024 * 1024] { + let directory = tempfile::tempdir().unwrap(); + let c_directory = c_str(directory.path().to_str().unwrap()); + let options = LanceFoyerCacheOptions { + directory: c_directory.as_ptr(), + disk_capacity_bytes: 128 * 1024 * 1024, + }; + // The second session has cold L1 and no pending-write buffers from the + // first session. Its index and data hits must come from persisted L2. + for reopened in [false, true] { + let session = unsafe { + lance_session_new_with_foyer_cache(l1_capacity, 16 * 1024 * 1024, &options) + }; + assert!(!session.is_null()); + let dataset = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!(!dataset.is_null()); + assert_eq!(scanned_row_count(dataset), 20_000); + let data = data_cache_statistics(dataset); + if reopened { + assert!(data.bytes_read_from_cache > 0); + assert_eq!(data.bytes_read_from_remote, 0); + } else { + assert!(data.bytes_read_from_remote > 0); + } + + query(dataset); + let before = disk_stats(session); + let overall_before = session_stats(session); + if reopened { + assert!(before.disk_hits > 0); + assert!(before.disk_read_bytes > 0); + assert_eq!(before.decode_failures, 0); + } else { + assert!(before.disk_write_bytes > 0); + } + + query(dataset); + let after = disk_stats(session); + let overall_after = session_stats(session); + if l1_capacity == 0 { + assert!(after.disk_hits > before.disk_hits); + assert!(after.memory_misses > before.memory_misses); + assert_eq!(overall_after.index_cache_size_bytes, 0); + } else { + assert!(after.memory_hits > before.memory_hits); + assert_eq!(after.memory_misses, before.memory_misses); + assert_eq!(after.disk_hits, before.disk_hits); + assert_eq!(after.disk_misses, before.disk_misses); + assert_eq!(after.disk_read_bytes, before.disk_read_bytes); + assert!(overall_after.index_cache_size_bytes > 0); + assert!(overall_after.index_cache_size_bytes <= l1_capacity); + } + assert!(overall_after.metadata_cache_size_bytes > 0); + assert!(overall_after.metadata_cache_size_bytes <= 16 * 1024 * 1024); + assert!(overall_after.index_cache_hits > overall_before.index_cache_hits); + assert!(!directory.path().join("index").exists()); + + unsafe { lance_session_close(session) }; + assert_eq!(scanned_row_count(dataset), 20_000); + assert!(data_cache_statistics(dataset).bytes_read_from_cache > 0); + unsafe { lance_dataset_close(dataset) }; + + // Closing the final owner drains Foyer asynchronously. Wait for its + // directory lock instead of retrying session construction or sleeping + // for a guessed flush duration. + let lock = std::fs::OpenOptions::new() + .read(true) + .write(true) + .open(directory.path().join("lance-cache-layout")) + .unwrap(); + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + match lock.try_lock() { + Ok(()) => break, + Err(std::fs::TryLockError::WouldBlock) => { + assert!( + std::time::Instant::now() < deadline, + "cache owners did not release directory" + ); + std::thread::sleep(std::time::Duration::from_millis(5)); + } + Err(error) => panic!("failed to check cache directory lock: {error}"), + } + } + lock.unlock().unwrap(); + } + } +} + +#[test] +fn test_session_with_data_cache_serves_repeated_scan() { + let (tmp, uri) = create_large_multi_fragment_dataset(10_000); + let c_uri = file_object_store_uri(&uri); + let (_cache_directory, session) = create_data_cache_session(); + + let dataset = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!(!dataset.is_null(), "dataset open should succeed"); + + assert_eq!(data_cache_statistics(dataset), Default::default()); + assert_eq!(scanned_row_count(dataset), 20_000); + let first_statistics = data_cache_statistics(dataset); + assert!(first_statistics.bytes_read_from_remote > 0); + + let cached_dataset = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!( + !cached_dataset.is_null(), + "second dataset open should succeed" + ); + unsafe { lance_session_close(session) }; + + for entry in std::fs::read_dir(tmp.path().join("large_ds/data")).unwrap() { + std::fs::remove_file(entry.unwrap().path()).unwrap(); + } + assert_eq!(scanned_row_count(cached_dataset), 20_000); + let cached_statistics = data_cache_statistics(cached_dataset); + assert!(cached_statistics.bytes_read_from_cache > 0); + assert_eq!(cached_statistics.bytes_read_from_remote, 0); + + assert_eq!(data_cache_statistics(dataset), first_statistics); + + unsafe { lance_dataset_close(cached_dataset) }; + unsafe { lance_dataset_close(dataset) }; +} + +#[test] +fn test_dataset_data_cache_statistics_validates_inputs_and_defaults_to_zero() { + let (_tmp, uri) = create_test_dataset(); + let c_uri = c_str(&uri); + let dataset = unsafe { lance_dataset_open(c_uri.as_ptr(), ptr::null(), 0) }; + assert!(!dataset.is_null()); + + let mut statistics = LanceDataCacheStatistics::default(); + assert_eq!( + unsafe { lance_dataset_get_data_cache_statistics(dataset, &mut statistics) }, + 0 + ); + assert_eq!(statistics, LanceDataCacheStatistics::default()); + assert_eq!( + unsafe { lance_dataset_get_data_cache_statistics(ptr::null(), &mut statistics) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert_eq!( + unsafe { lance_dataset_get_data_cache_statistics(dataset, ptr::null_mut()) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + + unsafe { lance_dataset_close(dataset) }; +} + #[test] fn test_open_nonexistent() { let c_uri = c_str("memory://nonexistent_dataset_xyz"); @@ -3318,6 +3658,40 @@ fn test_dataset_restore_to_prior_version() { unsafe { lance_dataset_close(ds) }; } +#[test] +fn test_restored_handle_has_independent_data_cache_statistics() { + let (_tmp, uri) = create_large_multi_fragment_dataset(10_000); + let c_uri = file_object_store_uri(&uri); + let (_cache_directory, session) = create_data_cache_session(); + let source = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!(!source.is_null()); + + assert_eq!(scanned_row_count(source), 20_000); + let source_statistics = data_cache_statistics(source); + assert!(source_statistics.bytes_read_from_remote > 0); + + let restored = unsafe { lance_dataset_restore(source, 1) }; + assert!(!restored.is_null()); + assert_eq!(data_cache_statistics(restored), Default::default()); + let source_statistics_after_restore = data_cache_statistics(source); + + assert_eq!(scanned_row_count(restored), 10_000); + let restored_statistics = data_cache_statistics(restored); + assert!(restored_statistics.bytes_read_from_cache > 0); + + assert_eq!( + data_cache_statistics(source), + source_statistics_after_restore + ); + + unsafe { + lance_session_close(session); + lance_dataset_close(restored); + lance_dataset_close(source); + } +} + #[test] fn test_dataset_restore_to_current_latest_writes_new_manifest() { // Restoring to the current latest still writes a new manifest. The diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index 5df9cde..c748015 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -171,6 +171,38 @@ static void test_shared_session(const char *uri) { (unsigned long long)stats.metadata_cache_entries); } +static void test_data_cache_session(const char *uri, const char *write_uri) { + printf(" test_data_cache_session... "); + + char cache_directory[4096]; + int path_len = snprintf(cache_directory, sizeof(cache_directory), + "%s_foyer_cache", write_uri); + ASSERT(path_len > 0 && (size_t)path_len < sizeof(cache_directory), + "cache directory path is too long"); + LanceFoyerCacheOptions options = { + .directory = cache_directory, + .disk_capacity_bytes = 64 * 1024 * 1024, + }; + LanceSession *session = + lance_session_new_with_foyer_cache(0, 16 * 1024 * 1024, &options); + ASSERT(session != NULL, "data-cache session creation failed"); + + LanceDataset *ds = lance_dataset_open_with_session(uri, NULL, 0, session); + ASSERT(ds != NULL, "data-cache session dataset open failed"); + LanceDataCacheStatistics statistics; + memset(&statistics, 0, sizeof(statistics)); + ASSERT(lance_dataset_get_data_cache_statistics(ds, &statistics) == 0, + "data-cache dataset statistics failed"); + LanceIndexDiskCacheStats index_stats = {0}; + ASSERT(lance_session_get_index_disk_cache_stats(session, &index_stats) == 0, + "shared index disk-cache statistics failed"); + lance_session_close(session); + ASSERT(lance_dataset_count_rows(ds) > 0, + "dataset should remain valid after data-cache session close"); + lance_dataset_close(ds); + printf("OK\n"); +} + static void test_scan(const char *uri) { printf(" test_scan... "); @@ -1323,6 +1355,7 @@ int main(int argc, char **argv) { test_open_and_metadata(uri); test_shared_session(uri); + test_data_cache_session(uri, write_uri); test_scan(uri); test_scan_with_limit(uri); test_scanner_blob_handling(blob_uri); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index 0332d1a..82a33ca 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -131,6 +131,28 @@ static void test_shared_session(const std::string& uri) { PASS(); } +static void test_data_cache_session(const std::string& uri, + const std::string& write_uri) { + TEST(test_data_cache_session); + + lance::FoyerCacheOptions options{ + write_uri + "_foyer_cache", + 64 * 1024 * 1024, + }; + auto session = std::make_unique( + 0, 16 * 1024 * 1024, options); + auto ds = lance::Dataset::open_with_session(*session, uri); + auto statistics = ds.data_cache_statistics(); + auto index_statistics = session->index_disk_cache_stats(); + assert(index_statistics.disk_hits == 0); + assert(statistics.bytes_read_from_cache == 0); + assert(statistics.bytes_read_from_remote == 0); + session.reset(); + assert(ds.count_rows() > 0); + + PASS(); +} + static void test_dataset_schema(const std::string& uri) { TEST(test_dataset_schema); @@ -1211,6 +1233,7 @@ int main(int argc, char** argv) { test_dataset_open(uri); test_shared_session(uri); + test_data_cache_session(uri, write_uri); test_dataset_schema(uri); test_scanner_fluent(uri); test_scanner_async_stream_ownership(uri);