diff --git a/Cargo.lock b/Cargo.lock index 9f078c3b1..a8e8270ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3790,6 +3790,15 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "lru" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d2f2f9b4ba7e6b24d95e7e899329d35be83bcded72c8540cdd5368932d1d90a" +dependencies = [ + "hashbrown 0.17.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -4576,6 +4585,7 @@ dependencies = [ "indexmap 2.14.0", "libloading 0.9.0", "log", + "lru 0.18.2", "lz4_flex 0.13.1", "md-5 0.10.6", "opendal-core", @@ -6700,7 +6710,7 @@ dependencies = [ "itertools 0.14.0", "levenshtein_automata", "log", - "lru", + "lru 0.16.4", "lz4_flex 0.13.1", "measure_time", "memmap2", diff --git a/DEPENDENCIES.rust.tsv b/DEPENDENCIES.rust.tsv index f2985fdb5..912dea886 100644 --- a/DEPENDENCIES.rust.tsv +++ b/DEPENDENCIES.rust.tsv @@ -324,6 +324,7 @@ litemap@0.8.2 X lock_api@0.4.14 X X log@0.4.33 X X lru@0.16.4 X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/benchmarks/tpcds/DEPENDENCIES.rust.tsv b/benchmarks/tpcds/DEPENDENCIES.rust.tsv index e31ea76a2..a871ab1e7 100644 --- a/benchmarks/tpcds/DEPENDENCIES.rust.tsv +++ b/benchmarks/tpcds/DEPENDENCIES.rust.tsv @@ -229,6 +229,7 @@ linux-raw-sys@0.12.1 X X X litemap@0.8.2 X lock_api@0.4.14 X X log@0.4.33 X X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/bindings/c/DEPENDENCIES.rust.tsv b/bindings/c/DEPENDENCIES.rust.tsv index a01464e35..9f5c62a93 100644 --- a/bindings/c/DEPENDENCIES.rust.tsv +++ b/bindings/c/DEPENDENCIES.rust.tsv @@ -4,6 +4,7 @@ ahash@0.8.12 X X aho-corasick@1.1.4 X X alloc-no-stdlib@2.0.4 X alloc-stdlib@0.2.4 X +allocator-api2@0.2.21 X X android_system_properties@0.1.5 X X anyhow@1.0.104 X X apache-avro@0.21.0 X @@ -90,6 +91,7 @@ find-msvc-tools@0.1.9 X X flatbuffers@25.12.19 X flate2@1.1.9 X X fnv@1.0.7 X X +foldhash@0.2.0 X foreign-types@0.3.2 X X foreign-types-shared@0.1.1 X X form_urlencoded@1.2.2 X X @@ -165,6 +167,7 @@ libm@0.2.16 X linux-raw-sys@0.12.1 X X X litemap@0.8.2 X log@0.4.33 X X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/bindings/go/DEPENDENCIES.rust.tsv b/bindings/go/DEPENDENCIES.rust.tsv index a01464e35..9f5c62a93 100644 --- a/bindings/go/DEPENDENCIES.rust.tsv +++ b/bindings/go/DEPENDENCIES.rust.tsv @@ -4,6 +4,7 @@ ahash@0.8.12 X X aho-corasick@1.1.4 X X alloc-no-stdlib@2.0.4 X alloc-stdlib@0.2.4 X +allocator-api2@0.2.21 X X android_system_properties@0.1.5 X X anyhow@1.0.104 X X apache-avro@0.21.0 X @@ -90,6 +91,7 @@ find-msvc-tools@0.1.9 X X flatbuffers@25.12.19 X flate2@1.1.9 X X fnv@1.0.7 X X +foldhash@0.2.0 X foreign-types@0.3.2 X X foreign-types-shared@0.1.1 X X form_urlencoded@1.2.2 X X @@ -165,6 +167,7 @@ libm@0.2.16 X linux-raw-sys@0.12.1 X X X litemap@0.8.2 X log@0.4.33 X X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/bindings/python/DEPENDENCIES.rust.tsv b/bindings/python/DEPENDENCIES.rust.tsv index 785a3f063..e3ec6090e 100644 --- a/bindings/python/DEPENDENCIES.rust.tsv +++ b/bindings/python/DEPENDENCIES.rust.tsv @@ -268,6 +268,7 @@ litemap@0.8.2 X lock_api@0.4.14 X X log@0.4.33 X X lru@0.16.4 X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/crates/integration_tests/DEPENDENCIES.rust.tsv b/crates/integration_tests/DEPENDENCIES.rust.tsv index aa9bfc158..73588751e 100644 --- a/crates/integration_tests/DEPENDENCIES.rust.tsv +++ b/crates/integration_tests/DEPENDENCIES.rust.tsv @@ -4,6 +4,7 @@ ahash@0.8.12 X X aho-corasick@1.1.4 X X alloc-no-stdlib@2.0.4 X alloc-stdlib@0.2.4 X +allocator-api2@0.2.21 X X android_system_properties@0.1.5 X X anyhow@1.0.104 X X apache-avro@0.21.0 X @@ -90,6 +91,7 @@ find-msvc-tools@0.1.9 X X flatbuffers@25.12.19 X flate2@1.1.9 X X fnv@1.0.7 X X +foldhash@0.2.0 X foreign-types@0.3.2 X X foreign-types-shared@0.1.1 X X form_urlencoded@1.2.2 X X @@ -165,6 +167,7 @@ libm@0.2.16 X linux-raw-sys@0.12.1 X X X litemap@0.8.2 X log@0.4.33 X X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/crates/integrations/datafusion/DEPENDENCIES.rust.tsv b/crates/integrations/datafusion/DEPENDENCIES.rust.tsv index 5e4ac46d4..3089ed334 100644 --- a/crates/integrations/datafusion/DEPENDENCIES.rust.tsv +++ b/crates/integrations/datafusion/DEPENDENCIES.rust.tsv @@ -286,6 +286,7 @@ litemap@0.8.2 X lock_api@0.4.14 X X log@0.4.33 X X lru@0.16.4 X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv index 39dd29f64..80ae88019 100644 --- a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv +++ b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv @@ -4,6 +4,7 @@ ahash@0.8.12 X X aho-corasick@1.1.4 X X alloc-no-stdlib@2.0.4 X alloc-stdlib@0.2.4 X +allocator-api2@0.2.21 X X android_system_properties@0.1.5 X X anyhow@1.0.104 X X apache-avro@0.21.0 X @@ -93,6 +94,7 @@ find-msvc-tools@0.1.9 X X flatbuffers@25.12.19 X flate2@1.1.9 X X fnv@1.0.7 X X +foldhash@0.2.0 X foreign-types@0.3.2 X X foreign-types-shared@0.1.1 X X form_urlencoded@1.2.2 X X @@ -168,6 +170,7 @@ libm@0.2.16 X linux-raw-sys@0.12.1 X X X litemap@0.8.2 X log@0.4.33 X X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index ee5beacfd..9ce74d201 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -87,6 +87,7 @@ opendal-service-s3 = { version = "0.58.0", optional = true } pretty_assertions = "1" apache-avro = { version = "0.21", features = ["snappy", "zstandard"] } indexmap = "2.5.0" +lru = "0.18.2" roaring = "0.11" crc32fast = "1" zstd = "0.13" diff --git a/crates/paimon/DEPENDENCIES.rust.tsv b/crates/paimon/DEPENDENCIES.rust.tsv index 24855b1c1..89b76f4fc 100644 --- a/crates/paimon/DEPENDENCIES.rust.tsv +++ b/crates/paimon/DEPENDENCIES.rust.tsv @@ -257,6 +257,7 @@ litemap@0.8.2 X lock_api@0.4.14 X X log@0.4.33 X X lru@0.16.4 X +lru@0.18.2 X lz4_flex@0.11.6 X lz4_flex@0.13.1 X lzokay-native@0.1.0 X diff --git a/crates/paimon/src/catalog/filesystem.rs b/crates/paimon/src/catalog/filesystem.rs index cf2f745a5..ae015e254 100644 --- a/crates/paimon/src/catalog/filesystem.rs +++ b/crates/paimon/src/catalog/filesystem.rs @@ -510,17 +510,12 @@ mod tests { #[test] fn test_filesystem_catalog_builds_local_cache_from_catalog_options() { let warehouse = TempDir::new().unwrap(); - let cache = TempDir::new().unwrap(); let mut options = Options::new(); options.set( CatalogOptions::WAREHOUSE, warehouse.path().to_string_lossy(), ); options.set(CatalogOptions::LOCAL_CACHE_ENABLED, "true"); - options.set( - CatalogOptions::LOCAL_CACHE_DIR, - cache.path().to_string_lossy(), - ); let catalog = FileSystemCatalog::new(options).unwrap(); diff --git a/crates/paimon/src/catalog/rest/rest_catalog.rs b/crates/paimon/src/catalog/rest/rest_catalog.rs index 7a19ba12d..401545cb8 100644 --- a/crates/paimon/src/catalog/rest/rest_catalog.rs +++ b/crates/paimon/src/catalog/rest/rest_catalog.rs @@ -548,17 +548,12 @@ mod tests { #[tokio::test] async fn test_rest_catalog_builds_local_cache_from_client_options() { - let cache = tempfile::tempdir().unwrap(); let mut options = Options::new(); options.set(CatalogOptions::URI, "http://localhost:1"); options.set(CatalogOptions::WAREHOUSE, "test-warehouse"); options.set("token.provider", "bear"); options.set("token", "test-token"); options.set(CatalogOptions::LOCAL_CACHE_ENABLED, "true"); - options.set( - CatalogOptions::LOCAL_CACHE_DIR, - cache.path().to_string_lossy(), - ); let catalog = RESTCatalog::new(options, false).await.unwrap(); diff --git a/crates/paimon/src/catalog/rest/rest_token_file_io.rs b/crates/paimon/src/catalog/rest/rest_token_file_io.rs index 8e2accfc7..9d70e1d15 100644 --- a/crates/paimon/src/catalog/rest/rest_token_file_io.rs +++ b/crates/paimon/src/catalog/rest/rest_token_file_io.rs @@ -217,14 +217,9 @@ mod tests { #[tokio::test] async fn test_token_file_io_keeps_catalog_local_cache() { - let cache_directory = tempfile::tempdir().unwrap(); let table_directory = tempfile::tempdir().unwrap(); let mut options = Options::new(); options.set(CatalogOptions::LOCAL_CACHE_ENABLED, "true"); - options.set( - CatalogOptions::LOCAL_CACHE_DIR, - cache_directory.path().to_string_lossy(), - ); let local_cache = create_local_cache(&options).unwrap(); let token_file_io = RESTTokenFileIO::new( Identifier::new("database", "table"), diff --git a/crates/paimon/src/common/options.rs b/crates/paimon/src/common/options.rs index 5ad2147d7..c23cd9318 100644 --- a/crates/paimon/src/common/options.rs +++ b/crates/paimon/src/common/options.rs @@ -76,10 +76,10 @@ impl CatalogOptions { /// Whether to enable local block caching for file reads. pub const LOCAL_CACHE_ENABLED: &'static str = "local-cache.enabled"; - /// Directory for the local disk block cache. + /// Directory for the local disk block cache. If unset, an in-memory cache is used. pub const LOCAL_CACHE_DIR: &'static str = "local-cache.dir"; - /// Maximum total encoded size of the local block cache. + /// Maximum total size of the local block cache. pub const LOCAL_CACHE_MAX_SIZE: &'static str = "local-cache.max-size"; /// Block size used by the local cache. @@ -175,6 +175,40 @@ impl From> for Options { } } +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum ParseMemorySizeError { + Invalid, + Overflow, +} + +/// Parses Java Paimon memory sizes, plus the binary unit aliases previously +/// accepted by the Rust local cache. +pub(crate) fn parse_memory_size(value: &str) -> Result { + let value = value.trim(); + if value.is_empty() { + return Err(ParseMemorySizeError::Invalid); + } + + let pos = value + .find(|character: char| !character.is_ascii_digit()) + .unwrap_or(value.len()); + let (number, unit) = value.split_at(pos); + let number = number + .parse::() + .map_err(|_| ParseMemorySizeError::Invalid)?; + let multiplier = match unit.trim().to_ascii_lowercase().as_str() { + "" | "b" | "bytes" => 1, + "k" | "kb" | "kib" | "kibibytes" => 1024, + "m" | "mb" | "mib" | "mebibytes" => 1024 * 1024, + "g" | "gb" | "gib" | "gibibytes" => 1024 * 1024 * 1024, + "t" | "tb" | "tib" | "tebibytes" => 1024_i64.pow(4), + _ => return Err(ParseMemorySizeError::Invalid), + }; + number + .checked_mul(multiplier) + .ok_or(ParseMemorySizeError::Overflow) +} + #[cfg(test)] mod tests { use super::*; @@ -244,4 +278,23 @@ mod tests { "local-cache.whitelist" ); } + + #[test] + fn test_parse_memory_size_supports_java_and_rust_units() { + for unit in ["k", "kb", "kib", "kibibytes"] { + assert_eq!(parse_memory_size(&format!("2 {unit}")), Ok(2 * 1024)); + } + assert_eq!(parse_memory_size("2 bytes"), Ok(2)); + assert_eq!(parse_memory_size("2 MiB"), Ok(2 * 1024 * 1024)); + assert_eq!(parse_memory_size("2 gib"), Ok(2 * 1024 * 1024 * 1024)); + assert_eq!(parse_memory_size("2 TiB"), Ok(2 * 1024_i64.pow(4))); + assert_eq!( + parse_memory_size("2 unknown"), + Err(ParseMemorySizeError::Invalid) + ); + assert_eq!( + parse_memory_size("9223372036854775807 tb"), + Err(ParseMemorySizeError::Overflow) + ); + } } diff --git a/crates/paimon/src/io/cache/disk.rs b/crates/paimon/src/io/cache/disk.rs index e63a2b5df..bd19e93a2 100644 --- a/crates/paimon/src/io/cache/disk.rs +++ b/crates/paimon/src/io/cache/disk.rs @@ -24,40 +24,15 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex, OnceLock, Weak}; use tokio::io::AsyncWriteExt; +use super::state::{BlockKey, CacheCoordinator, LogicalPath}; + const CACHE_MAGIC: &[u8; 8] = b"PAIMONLC"; const CACHE_FORMAT_VERSION: u8 = 2; const FIXED_HEADER_LEN: usize = CACHE_MAGIC.len() + 1 + 4 + 4 + 8 + 8 + 8; const CHECKSUM_LEN: usize = 4; const MAX_CACHE_KEY_HEADER_LEN: usize = 1024 * 1024; -#[derive(Clone, Debug, Eq, Hash, PartialEq)] -pub(super) struct BlockKey { - namespace: String, - path: String, - block_size: u64, - block_index: u64, -} - impl BlockKey { - #[cfg(test)] - pub(super) fn new(path: impl Into, block_size: u64, block_index: u64) -> Self { - Self::with_namespace("", path, block_size, block_index) - } - - pub(super) fn with_namespace( - namespace: impl Into, - path: impl Into, - block_size: u64, - block_index: u64, - ) -> Self { - Self { - namespace: namespace.into(), - path: path.into(), - block_size, - block_index, - } - } - pub(super) fn cache_relative_path(&self) -> PathBuf { let mut digest = Sha256::new(); digest.update([CACHE_FORMAT_VERSION]); @@ -86,51 +61,17 @@ pub(super) struct DiskCache { root: PathBuf, state: Mutex, recovered: tokio::sync::OnceCell<()>, - in_flight: tokio::sync::Mutex>>>, - path_states: Mutex>>, - prefix_barrier: tokio::sync::RwLock<()>, + coordinator: Arc, } #[derive(Debug, Default)] struct CacheState { entries: IndexMap, paths: HashMap>, - file_sizes: IndexMap, current_size: u64, max_size: Option, } -#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] -struct LogicalPath { - namespace: String, - path: String, -} - -#[derive(Debug)] -struct PathCacheState { - generation: std::sync::atomic::AtomicU64, - publish_gate: tokio::sync::RwLock<()>, -} - -#[derive(Clone)] -pub(in crate::io) struct CacheReadToken { - generation: u64, - state: Arc, -} - -impl CacheReadToken { - pub(super) fn is_current(&self) -> bool { - self.state - .generation - .load(std::sync::atomic::Ordering::SeqCst) - == self.generation - } - - pub(super) async fn publish_guard(&self) -> tokio::sync::RwLockReadGuard<'_, ()> { - self.state.publish_gate.read().await - } -} - impl DiskCache { #[cfg(test)] pub(super) fn new(root: impl AsRef, max_size: Option) -> crate::Result { @@ -167,9 +108,7 @@ impl DiskCache { ..CacheState::default() }), recovered: tokio::sync::OnceCell::new(), - in_flight: tokio::sync::Mutex::new(HashMap::new()), - path_states: Mutex::new(HashMap::new()), - prefix_barrier: tokio::sync::RwLock::new(()), + coordinator: Arc::new(CacheCoordinator::default()), } } @@ -182,7 +121,7 @@ impl DiskCache { }; } - async fn ensure_recovered(&self) { + pub(super) async fn ensure_recovered(&self) { self.recovered .get_or_init(|| async { let root = self.root.clone(); @@ -218,92 +157,8 @@ impl DiskCache { self.evict_over_limit().await; } - pub(super) fn read_token(&self, namespace: &str, path: &str) -> CacheReadToken { - let state = self.path_state(namespace, path); - CacheReadToken { - generation: state.generation.load(std::sync::atomic::Ordering::SeqCst), - state, - } - } - - pub(super) async fn prefix_read_guard(&self) -> tokio::sync::RwLockReadGuard<'_, ()> { - self.prefix_barrier.read().await - } - - fn path_state(&self, namespace: &str, path: &str) -> Arc { - let logical_path = LogicalPath { - namespace: namespace.to_string(), - path: path.to_string(), - }; - let mut states = self - .path_states - .lock() - .unwrap_or_else(|error| error.into_inner()); - if let Some(state) = states.get(&logical_path).and_then(Weak::upgrade) { - return state; - } - if states.len() >= 1024 { - states.retain(|_, state| state.strong_count() > 0); - } - let state = Arc::new(PathCacheState { - generation: std::sync::atomic::AtomicU64::new(0), - publish_gate: tokio::sync::RwLock::new(()), - }); - states.insert(logical_path, Arc::downgrade(&state)); - state - } - - pub(super) async fn block_load_lock(&self, key: &BlockKey) -> Arc> { - let mut in_flight = self.in_flight.lock().await; - if let Some(lock) = in_flight.get(key).and_then(Weak::upgrade) { - return lock; - } - if in_flight.len() >= 1024 { - in_flight.retain(|_, lock| lock.strong_count() > 0); - } - let lock = Arc::new(tokio::sync::Mutex::new(())); - in_flight.insert(key.clone(), Arc::downgrade(&lock)); - lock - } - - pub(super) async fn release_block_load_lock( - &self, - key: &BlockKey, - lock: &Arc>, - ) { - let mut in_flight = self.in_flight.lock().await; - if Arc::strong_count(lock) == 1 - && in_flight - .get(key) - .and_then(Weak::upgrade) - .is_some_and(|current| Arc::ptr_eq(¤t, lock)) - { - in_flight.remove(key); - } - } - - pub(super) fn file_size(&self, namespace: &str, path: &str) -> Option { - let logical_path = LogicalPath { - namespace: namespace.to_string(), - path: path.to_string(), - }; - let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); - let size = state.file_sizes.shift_remove(&logical_path)?; - state.file_sizes.insert(logical_path, size); - Some(size) - } - - pub(super) fn put_file_size(&self, namespace: &str, path: &str, size: u64, capacity: usize) { - let logical_path = LogicalPath { - namespace: namespace.to_string(), - path: path.to_string(), - }; - let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); - state.file_sizes.shift_remove(&logical_path); - state.file_sizes.insert(logical_path, size); - while state.file_sizes.len() > capacity { - state.file_sizes.shift_remove_index(0); - } + pub(super) fn coordinator(&self) -> Arc { + self.coordinator.clone() } pub(super) async fn get_block(&self, key: &BlockKey) -> Option { @@ -411,12 +266,6 @@ impl DiskCache { pub(super) async fn invalidate_path(&self, namespace: &str, path: &str) { self.ensure_recovered().await; - let _prefix_guard = self.prefix_barrier.read().await; - let path_state = self.path_state(namespace, path); - let _publish_guard = path_state.publish_gate.write().await; - path_state - .generation - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); self.invalidate_matching(|logical_path| { logical_path.namespace == namespace && logical_path.path == path }) @@ -439,30 +288,9 @@ impl DiskCache { pub(super) async fn invalidate_prefix(&self, namespace: &str, prefix: &str) { self.ensure_recovered().await; - let _prefix_guard = self.prefix_barrier.write().await; let prefix = prefix.trim_end_matches('/'); - let mut states = { - let states = self - .path_states - .lock() - .unwrap_or_else(|error| error.into_inner()); - states - .iter() - .filter(|(path, _)| logical_path_matches_prefix(path, namespace, prefix)) - .filter_map(|(path, state)| Weak::upgrade(state).map(|state| (path.clone(), state))) - .collect::>() - }; - states.sort_unstable_by(|left, right| left.0.cmp(&right.0)); - let mut publish_guards = Vec::with_capacity(states.len()); - for (_, state) in &states { - publish_guards.push(state.publish_gate.write().await); - state - .generation - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - } self.invalidate_matching(|path| logical_path_matches_prefix(path, namespace, prefix)) .await; - drop(publish_guards); } async fn invalidate_matching(&self, matches: impl Fn(&LogicalPath) -> bool) { @@ -478,7 +306,6 @@ impl DiskCache { for key in &keys { remove_state_entry(&mut state, key); } - state.file_sizes.retain(|path, _| !matches(path)); keys }; for key in keys { @@ -579,10 +406,7 @@ fn collect_evictions(state: &mut CacheState) -> Vec { } fn logical_path(key: &BlockKey) -> LogicalPath { - LogicalPath { - namespace: key.namespace.clone(), - path: key.path.clone(), - } + LogicalPath::from_key(key) } fn logical_path_matches_prefix(path: &LogicalPath, namespace: &str, prefix: &str) -> bool { diff --git a/crates/paimon/src/io/cache/memory.rs b/crates/paimon/src/io/cache/memory.rs new file mode 100644 index 000000000..c50cd91af --- /dev/null +++ b/crates/paimon/src/io/cache/memory.rs @@ -0,0 +1,139 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use bytes::Bytes; +use lru::LruCache; +use std::sync::Mutex; + +use super::state::BlockKey; + +#[derive(Debug)] +pub(super) struct MemoryCache { + state: Mutex, +} + +#[derive(Debug)] +struct MemoryState { + entries: LruCache, + current_size: u64, + max_size: Option, +} + +impl MemoryCache { + pub(super) fn new(max_size: Option) -> Self { + Self { + state: Mutex::new(MemoryState { + entries: LruCache::unbounded(), + current_size: 0, + max_size, + }), + } + } + + pub(super) fn get_block(&self, key: &BlockKey) -> Option { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.entries.get(key).cloned() + } + + pub(super) fn put_block(&self, key: &BlockKey, payload: Bytes) { + let payload_size = payload.len() as u64; + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state + .max_size + .is_some_and(|max_size| payload_size > max_size) + { + return; + } + if let Some(previous) = state.entries.put(key.clone(), payload) { + state.current_size = state.current_size.saturating_sub(previous.len() as u64); + } + state.current_size = state.current_size.saturating_add(payload_size); + while state + .max_size + .is_some_and(|max_size| state.current_size > max_size) + { + let Some((_, payload)) = state.entries.pop_lru() else { + break; + }; + state.current_size = state.current_size.saturating_sub(payload.len() as u64); + } + } + + pub(super) fn remove_block(&self, key: &BlockKey) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(payload) = state.entries.pop(key) { + state.current_size = state.current_size.saturating_sub(payload.len() as u64); + } + } + + pub(super) fn invalidate_path(&self, namespace: &str, path: &str) { + self.invalidate_matching(|key| key.matches_path(namespace, path)); + } + + pub(super) fn invalidate_prefix(&self, namespace: &str, prefix: &str) { + let prefix = prefix.trim_end_matches('/'); + self.invalidate_matching(|key| key.matches_prefix(namespace, prefix)); + } + + fn invalidate_matching(&self, matches: impl Fn(&BlockKey) -> bool) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + let keys = state + .entries + .iter() + .map(|(key, _)| key) + .filter(|key| matches(key)) + .cloned() + .collect::>(); + for key in keys { + if let Some(payload) = state.entries.pop(&key) { + state.current_size = state.current_size.saturating_sub(payload.len() as u64); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_memory_cache_refreshes_lru_and_evicts_by_payload_bytes() { + let cache = MemoryCache::new(Some(8)); + let first = BlockKey::new("snapshot-1", 4, 0); + let second = BlockKey::new("snapshot-1", 4, 1); + let third = BlockKey::new("snapshot-1", 4, 2); + + cache.put_block(&first, Bytes::from_static(b"aaaa")); + cache.put_block(&second, Bytes::from_static(b"bbbb")); + assert_eq!(cache.get_block(&first), Some(Bytes::from_static(b"aaaa"))); + cache.put_block(&third, Bytes::from_static(b"cccc")); + + assert_eq!(cache.get_block(&second), None); + assert_eq!(cache.get_block(&first), Some(Bytes::from_static(b"aaaa"))); + assert_eq!(cache.get_block(&third), Some(Bytes::from_static(b"cccc"))); + } + + #[test] + fn test_memory_cache_skips_block_larger_than_capacity() { + let cache = MemoryCache::new(Some(3)); + let key = BlockKey::new("snapshot-1", 4, 0); + + cache.put_block(&key, Bytes::from_static(b"data")); + + assert_eq!(cache.get_block(&key), None); + } +} diff --git a/crates/paimon/src/io/cache/mod.rs b/crates/paimon/src/io/cache/mod.rs index 66792c534..9560394b8 100644 --- a/crates/paimon/src/io/cache/mod.rs +++ b/crates/paimon/src/io/cache/mod.rs @@ -17,16 +17,20 @@ mod disk; mod file_type; +mod memory; mod reader; +mod state; use self::file_type::FileType; +use self::memory::MemoryCache; +use self::state::{BlockKey, CacheCoordinator, CacheReadToken}; use crate::common::{CatalogOptions, Options}; use sha2::{Digest, Sha256}; use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; -use disk::{BlockKey, CacheReadToken, DiskCache}; +use disk::DiskCache; pub(super) use reader::CachedFileReader; const CACHE_DIRECTORY_NAME: &str = "paimon-local-cache-v2"; @@ -34,13 +38,20 @@ const DEFAULT_FILE_SIZE_CAPACITY: usize = 65_536; #[derive(Debug)] pub(crate) struct LocalCache { - disk: Arc, + backend: CacheBackend, + coordinator: Arc, namespace: String, block_size: u64, whitelist: HashSet, file_size_capacity: usize, } +#[derive(Debug)] +enum CacheBackend { + Memory(MemoryCache), + Disk(Arc), +} + impl LocalCache { pub(super) fn new(config: LocalCacheConfig) -> crate::Result { let file_size_capacity = config @@ -49,8 +60,19 @@ impl LocalCache { .and_then(|capacity| usize::try_from(capacity).ok()) .unwrap_or(DEFAULT_FILE_SIZE_CAPACITY) .clamp(1, DEFAULT_FILE_SIZE_CAPACITY); + let (backend, coordinator) = if let Some(dir) = config.dir { + let disk = DiskCache::shared(dir.join(CACHE_DIRECTORY_NAME), config.max_size)?; + let coordinator = disk.coordinator(); + (CacheBackend::Disk(disk), coordinator) + } else { + ( + CacheBackend::Memory(MemoryCache::new(config.max_size)), + Arc::new(CacheCoordinator::default()), + ) + }; Ok(Self { - disk: DiskCache::shared(config.dir.join(CACHE_DIRECTORY_NAME), config.max_size)?, + backend, + coordinator, namespace: config.namespace, block_size: config.block_size, whitelist: config.whitelist, @@ -71,12 +93,15 @@ impl LocalCache { } async fn get_block(&self, key: &BlockKey, token: &CacheReadToken) -> Option { - let _prefix_guard = self.disk.prefix_read_guard().await; + let _prefix_guard = self.coordinator.prefix_read_guard().await; let _publish_guard = token.publish_guard().await; if !token.is_current() { return None; } - let payload = self.disk.get_block(key).await; + let payload = match &self.backend { + CacheBackend::Memory(memory) => memory.get_block(key), + CacheBackend::Disk(disk) => disk.get_block(key).await, + }; if token.is_current() { payload } else { @@ -85,55 +110,78 @@ impl LocalCache { } async fn put_block(&self, key: &BlockKey, payload: bytes::Bytes, token: &CacheReadToken) { - let _prefix_guard = self.disk.prefix_read_guard().await; + let _prefix_guard = self.coordinator.prefix_read_guard().await; let _publish_guard = token.publish_guard().await; if !token.is_current() { return; } - self.disk.put_block(key, payload).await; + match &self.backend { + CacheBackend::Memory(memory) => memory.put_block(key, payload), + CacheBackend::Disk(disk) => disk.put_block(key, payload).await, + } } async fn remove_block(&self, key: &BlockKey) { - self.disk.remove_block(key).await; + match &self.backend { + CacheBackend::Memory(memory) => memory.remove_block(key), + CacheBackend::Disk(disk) => disk.remove_block(key).await, + } } pub(super) fn read_token(&self, path: &str) -> CacheReadToken { - self.disk.read_token(&self.namespace, path) + self.coordinator.read_token(&self.namespace, path) } async fn block_load_lock(&self, key: &BlockKey) -> Arc> { - self.disk.block_load_lock(key).await + self.coordinator.block_load_lock(key).await } async fn release_block_load_lock(&self, key: &BlockKey, lock: &Arc>) { - self.disk.release_block_load_lock(key, lock).await; + self.coordinator.release_block_load_lock(key, lock).await; } pub(super) async fn file_size(&self, path: &str, token: &CacheReadToken) -> Option { - let _prefix_guard = self.disk.prefix_read_guard().await; + let _prefix_guard = self.coordinator.prefix_read_guard().await; let _publish_guard = token.publish_guard().await; token .is_current() - .then(|| self.disk.file_size(&self.namespace, path)) + .then(|| self.coordinator.file_size(&self.namespace, path)) .flatten() } pub(super) async fn put_file_size(&self, path: &str, size: u64, token: &CacheReadToken) { - let _prefix_guard = self.disk.prefix_read_guard().await; + let _prefix_guard = self.coordinator.prefix_read_guard().await; let _publish_guard = token.publish_guard().await; if !token.is_current() { return; } - self.disk + self.coordinator .put_file_size(&self.namespace, path, size, self.file_size_capacity); } pub(super) async fn invalidate_path(&self, path: &str) { - self.disk.invalidate_path(&self.namespace, path).await; + let _guard = self + .coordinator + .begin_path_invalidation(&self.namespace, path) + .await; + match &self.backend { + CacheBackend::Memory(memory) => memory.invalidate_path(&self.namespace, path), + CacheBackend::Disk(disk) => disk.invalidate_path(&self.namespace, path).await, + } } pub(super) async fn invalidate_prefix(&self, prefix: &str) { - self.disk.invalidate_prefix(&self.namespace, prefix).await; + if let CacheBackend::Disk(disk) = &self.backend { + disk.ensure_recovered().await; + } + let _guard = self + .coordinator + .begin_prefix_invalidation(&self.namespace, prefix) + .await; + match &self.backend { + CacheBackend::Memory(memory) => memory.invalidate_prefix(&self.namespace, prefix), + CacheBackend::Disk(disk) => disk.invalidate_prefix(&self.namespace, prefix).await, + } } } @@ -153,7 +201,7 @@ pub(crate) fn create_local_cache_with_namespace( #[derive(Debug)] pub(crate) struct LocalCacheConfig { - dir: PathBuf, + dir: Option, namespace: String, max_size: Option, block_size: u64, @@ -194,13 +242,7 @@ impl LocalCacheConfig { let dir = options .get(CatalogOptions::LOCAL_CACHE_DIR) .filter(|value| !value.trim().is_empty()) - .ok_or_else(|| crate::Error::ConfigInvalid { - message: format!( - "Missing required option: {}", - CatalogOptions::LOCAL_CACHE_DIR - ), - })? - .into(); + .map(PathBuf::from); let max_size = options .get(CatalogOptions::LOCAL_CACHE_MAX_SIZE) @@ -252,36 +294,17 @@ fn catalog_namespace(options: &Options) -> String { } fn parse_memory_size(key: &str, value: &str) -> crate::Result { - let compact = value - .chars() - .filter(|character| !character.is_ascii_whitespace()) - .collect::() - .to_ascii_lowercase(); - let unit_start = compact - .find(|character: char| !character.is_ascii_digit()) - .unwrap_or(compact.len()); - let (number, unit) = compact.split_at(unit_start); - let number = number - .parse::() - .map_err(|_| crate::Error::ConfigInvalid { - message: format!("Invalid memory size for {key}: '{value}'"), - })?; - let multiplier = match unit { - "" | "b" => 1, - "k" | "kb" | "kib" => 1024, - "m" | "mb" | "mib" => 1024 * 1024, - "g" | "gb" | "gib" => 1024 * 1024 * 1024, - "t" | "tb" | "tib" => 1024_u64.pow(4), - _ => { - return Err(crate::Error::ConfigInvalid { - message: format!("Invalid memory size for {key}: '{value}'"), - }); - } - }; - number - .checked_mul(multiplier) - .ok_or_else(|| crate::Error::ConfigInvalid { - message: format!("Memory size for {key} is too large: '{value}'"), + crate::common::options::parse_memory_size(value) + .map(|size| size as u64) + .map_err(|error| crate::Error::ConfigInvalid { + message: match error { + crate::common::options::ParseMemorySizeError::Invalid => { + format!("Invalid memory size for {key}: '{value}'") + } + crate::common::options::ParseMemorySizeError::Overflow => { + format!("Memory size for {key} is too large: '{value}'") + } + }, }) } @@ -297,13 +320,19 @@ mod tests { } #[test] - fn test_local_cache_config_requires_directory_when_enabled() { + fn test_local_cache_config_uses_memory_when_enabled_without_directory() { let mut options = Options::new(); options.set(crate::common::CatalogOptions::LOCAL_CACHE_ENABLED, "true"); - let error = LocalCacheConfig::from_options(&options).unwrap_err(); - assert!(matches!(error, crate::Error::ConfigInvalid { .. })); - assert!(error.to_string().contains("local-cache.dir")); + let config = LocalCacheConfig::from_options(&options).unwrap().unwrap(); + + assert_eq!(config.dir, None); + assert_eq!(config.max_size, None); + assert_eq!(config.block_size, 1024 * 1024); + assert_eq!( + config.whitelist, + HashSet::from([FileType::Meta, FileType::GlobalIndex]) + ); } #[test] @@ -313,7 +342,10 @@ mod tests { options.set(CatalogOptions::LOCAL_CACHE_DIR, "/tmp/paimon-cache"); let config = LocalCacheConfig::from_options(&options).unwrap().unwrap(); - assert_eq!(config.dir, std::path::Path::new("/tmp/paimon-cache")); + assert_eq!( + config.dir.as_deref(), + Some(std::path::Path::new("/tmp/paimon-cache")) + ); assert_eq!(config.max_size, None); assert_eq!(config.block_size, 1024 * 1024); assert_eq!( @@ -359,8 +391,8 @@ mod tests { .unwrap(); assert_eq!( - effective_config.dir, - std::path::Path::new("/tmp/paimon-cache") + effective_config.dir.as_deref(), + Some(std::path::Path::new("/tmp/paimon-cache")) ); assert_ne!(effective_config.namespace, local_namespace); assert_eq!( @@ -395,7 +427,7 @@ mod tests { async fn test_local_cache_file_size_is_removed_with_path_invalidation() { let directory = tempfile::tempdir().unwrap(); let cache = LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -416,7 +448,7 @@ mod tests { async fn test_local_cache_file_size_is_invalidated_across_shared_instances() { let directory = tempfile::tempdir().unwrap(); let config = || LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -437,30 +469,31 @@ mod tests { #[tokio::test] async fn test_stale_file_size_cannot_republish_after_invalidation() { let directory = tempfile::tempdir().unwrap(); - let cache = LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + let config = || LocalCacheConfig { + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, whitelist: HashSet::from([FileType::Meta]), - }) - .unwrap(); + }; + let first = LocalCache::new(config()).unwrap(); + let second = LocalCache::new(config()).unwrap(); let path = "s3://bucket/table/snapshot/snapshot-1"; - let stale_token = cache.read_token(path); + let stale_token = first.read_token(path); - assert_eq!(cache.file_size(path, &stale_token).await, None); - cache.invalidate_path(path).await; - cache.put_file_size(path, 42, &stale_token).await; + assert_eq!(first.file_size(path, &stale_token).await, None); + second.invalidate_path(path).await; + first.put_file_size(path, 42, &stale_token).await; - let current_token = cache.read_token(path); - assert_eq!(cache.file_size(path, ¤t_token).await, None); + let current_token = first.read_token(path); + assert_eq!(first.file_size(path, ¤t_token).await, None); } #[test] fn test_local_cache_uses_whitelist_and_bypasses_mutable_files() { let directory = tempfile::tempdir().unwrap(); let cache = LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -487,7 +520,7 @@ mod tests { std::fs::write(&nested, b"nested foreign").unwrap(); LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -507,7 +540,7 @@ mod tests { async fn test_local_cache_bounds_file_size_entries() { let directory = tempfile::tempdir().unwrap(); let cache = LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: Some(8), block_size: 4, diff --git a/crates/paimon/src/io/cache/reader.rs b/crates/paimon/src/io/cache/reader.rs index 40aeb7288..72a3b7008 100644 --- a/crates/paimon/src/io/cache/reader.rs +++ b/crates/paimon/src/io/cache/reader.rs @@ -253,14 +253,13 @@ mod tests { #[tokio::test] async fn test_cached_range_reader_reads_unaligned_blocks_and_reuses_them() { - let directory = tempfile::tempdir().unwrap(); let delegate = Arc::new(CountingReader { data: Bytes::from_static(b"abcdefghijkl"), reads: AtomicUsize::new(0), }); let cache = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: None, namespace: "test".to_string(), max_size: None, block_size: 4, @@ -304,14 +303,13 @@ mod tests { #[tokio::test] async fn test_cached_range_single_flight_reads_cold_block_once() { - let directory = tempfile::tempdir().unwrap(); let delegate = Arc::new(SlowCountingReader { data: Bytes::from_static(b"abcdefgh"), reads: AtomicUsize::new(0), }); let cache = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: None, namespace: "test".to_string(), max_size: None, block_size: 4, @@ -341,7 +339,7 @@ mod tests { reads: AtomicUsize::new(0), }); let config = || LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -385,11 +383,10 @@ mod tests { #[tokio::test] async fn test_in_flight_miss_does_not_republish_after_invalidation() { - let directory = tempfile::tempdir().unwrap(); let path = "s3://bucket/table/snapshot/snapshot-1"; let cache = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: None, namespace: "test".to_string(), max_size: None, block_size: 4, @@ -438,7 +435,7 @@ mod tests { let path = "s3://bucket/table/snapshot/snapshot-1"; let cache_a = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -448,7 +445,7 @@ mod tests { ); let cache_b = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -498,16 +495,15 @@ mod tests { let first_path = "s3://bucket/table/snapshot/snapshot-1"; let observed_path = "s3://bucket/table/snapshot/snapshot-2"; let blocked_path = "s3://bucket/table/snapshot/snapshot-3"; - let cache = Arc::new( - LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), - namespace: "test".to_string(), - max_size: None, - block_size: 4, - whitelist: std::collections::HashSet::from([FileType::Meta]), - }) - .unwrap(), - ); + let config = || LocalCacheConfig { + dir: Some(directory.path().to_path_buf()), + namespace: "test".to_string(), + max_size: None, + block_size: 4, + whitelist: std::collections::HashSet::from([FileType::Meta]), + }; + let cache_a = Arc::new(LocalCache::new(config()).unwrap()); + let cache_b = Arc::new(LocalCache::new(config()).unwrap()); let old_reader = CachedFileReader::new( Arc::new(CountingReader { data: Bytes::from_static(b"old!"), @@ -515,21 +511,21 @@ mod tests { }), first_path, 4, - cache.clone(), + cache_a.clone(), ); assert_eq!( old_reader.read(0..4).await.unwrap(), Bytes::from_static(b"old!") ); - let warm_size_token = cache.read_token(first_path); - cache.put_file_size(first_path, 4, &warm_size_token).await; + let warm_size_token = cache_a.read_token(first_path); + cache_a.put_file_size(first_path, 4, &warm_size_token).await; drop(warm_size_token); drop(old_reader); - let observed_token = cache.read_token(observed_path); - let blocked_token = cache.read_token(blocked_path); + let observed_token = cache_a.read_token(observed_path); + let blocked_token = cache_a.read_token(blocked_path); let blocked_guard = blocked_token.publish_guard().await; - let invalidating_cache = cache.clone(); + let invalidating_cache = cache_b; let invalidation = tokio::spawn(async move { invalidating_cache.invalidate_prefix(prefix).await }); tokio::time::timeout(std::time::Duration::from_secs(1), async { @@ -545,9 +541,9 @@ mod tests { reads: AtomicUsize::new(0), }); let current_reader = - CachedFileReader::new(current_delegate.clone(), first_path, 4, cache.clone()); + CachedFileReader::new(current_delegate.clone(), first_path, 4, cache_a.clone()); let mut current_load = tokio::spawn(async move { current_reader.read(0..4).await }); - let size_cache = cache.clone(); + let size_cache = cache_a; let mut current_size = tokio::spawn(async move { let token = size_cache.read_token(first_path); size_cache.file_size(first_path, &token).await @@ -646,7 +642,7 @@ mod tests { }); let cache = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -680,7 +676,7 @@ mod tests { let directory = tempfile::tempdir().unwrap(); let cache = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: Some(directory.path().to_path_buf()), namespace: "test".to_string(), max_size: None, block_size: 4, @@ -695,14 +691,13 @@ mod tests { #[tokio::test] async fn test_cached_full_reader_loads_source_once_then_hits_blocks() { - let directory = tempfile::tempdir().unwrap(); let delegate = Arc::new(CountingReader { data: Bytes::from_static(b"abcdefghijkl"), reads: AtomicUsize::new(0), }); let cache = Arc::new( LocalCache::new(LocalCacheConfig { - dir: directory.path().to_path_buf(), + dir: None, namespace: "test".to_string(), max_size: None, block_size: 4, diff --git a/crates/paimon/src/io/cache/state.rs b/crates/paimon/src/io/cache/state.rs new file mode 100644 index 000000000..7625186b2 --- /dev/null +++ b/crates/paimon/src/io/cache/state.rs @@ -0,0 +1,295 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use indexmap::IndexMap; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, Weak}; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(super) struct BlockKey { + pub(super) namespace: String, + pub(super) path: String, + pub(super) block_size: u64, + pub(super) block_index: u64, +} + +impl BlockKey { + #[cfg(test)] + pub(super) fn new(path: impl Into, block_size: u64, block_index: u64) -> Self { + Self::with_namespace("", path, block_size, block_index) + } + + pub(super) fn with_namespace( + namespace: impl Into, + path: impl Into, + block_size: u64, + block_index: u64, + ) -> Self { + Self { + namespace: namespace.into(), + path: path.into(), + block_size, + block_index, + } + } + + pub(super) fn matches_path(&self, namespace: &str, path: &str) -> bool { + self.namespace == namespace && self.path == path + } + + pub(super) fn matches_prefix(&self, namespace: &str, prefix: &str) -> bool { + self.namespace == namespace && path_matches_prefix(&self.path, prefix) + } +} + +#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub(super) struct LogicalPath { + pub(super) namespace: String, + pub(super) path: String, +} + +impl LogicalPath { + pub(super) fn new(namespace: &str, path: &str) -> Self { + Self { + namespace: namespace.to_string(), + path: path.to_string(), + } + } + + pub(super) fn from_key(key: &BlockKey) -> Self { + Self::new(&key.namespace, &key.path) + } + + fn matches_prefix(&self, namespace: &str, prefix: &str) -> bool { + self.namespace == namespace && path_matches_prefix(&self.path, prefix) + } +} + +fn path_matches_prefix(path: &str, prefix: &str) -> bool { + path == prefix + || path + .strip_prefix(prefix) + .is_some_and(|suffix| suffix.starts_with('/')) +} + +#[derive(Debug)] +struct PathCacheState { + generation: std::sync::atomic::AtomicU64, + publish_gate: Arc>, +} + +#[derive(Clone)] +pub(in crate::io) struct CacheReadToken { + generation: u64, + state: Arc, +} + +impl CacheReadToken { + pub(super) fn is_current(&self) -> bool { + self.state + .generation + .load(std::sync::atomic::Ordering::SeqCst) + == self.generation + } + + pub(super) async fn publish_guard(&self) -> tokio::sync::RwLockReadGuard<'_, ()> { + self.state.publish_gate.read().await + } +} + +#[derive(Debug)] +pub(super) struct CacheCoordinator { + in_flight: tokio::sync::Mutex>>>, + path_states: Mutex>>, + prefix_barrier: Arc>, + file_sizes: Mutex>, +} + +impl Default for CacheCoordinator { + fn default() -> Self { + Self { + in_flight: tokio::sync::Mutex::new(HashMap::new()), + path_states: Mutex::new(HashMap::new()), + prefix_barrier: Arc::new(tokio::sync::RwLock::new(())), + file_sizes: Mutex::new(IndexMap::new()), + } + } +} + +impl CacheCoordinator { + pub(super) fn read_token(&self, namespace: &str, path: &str) -> CacheReadToken { + let state = self.path_state(namespace, path); + CacheReadToken { + generation: state.generation.load(std::sync::atomic::Ordering::SeqCst), + state, + } + } + + pub(super) async fn prefix_read_guard(&self) -> tokio::sync::RwLockReadGuard<'_, ()> { + self.prefix_barrier.read().await + } + + fn path_state(&self, namespace: &str, path: &str) -> Arc { + let logical_path = LogicalPath::new(namespace, path); + let mut states = self + .path_states + .lock() + .unwrap_or_else(|error| error.into_inner()); + if let Some(state) = states.get(&logical_path).and_then(Weak::upgrade) { + return state; + } + if states.len() >= 1024 { + states.retain(|_, state| state.strong_count() > 0); + } + let state = Arc::new(PathCacheState { + generation: std::sync::atomic::AtomicU64::new(0), + publish_gate: Arc::new(tokio::sync::RwLock::new(())), + }); + states.insert(logical_path, Arc::downgrade(&state)); + state + } + + pub(super) async fn block_load_lock(&self, key: &BlockKey) -> Arc> { + let mut in_flight = self.in_flight.lock().await; + if let Some(lock) = in_flight.get(key).and_then(Weak::upgrade) { + return lock; + } + if in_flight.len() >= 1024 { + in_flight.retain(|_, lock| lock.strong_count() > 0); + } + let lock = Arc::new(tokio::sync::Mutex::new(())); + in_flight.insert(key.clone(), Arc::downgrade(&lock)); + lock + } + + pub(super) async fn release_block_load_lock( + &self, + key: &BlockKey, + lock: &Arc>, + ) { + let mut in_flight = self.in_flight.lock().await; + if Arc::strong_count(lock) == 1 + && in_flight + .get(key) + .and_then(Weak::upgrade) + .is_some_and(|current| Arc::ptr_eq(¤t, lock)) + { + in_flight.remove(key); + } + } + + pub(super) fn file_size(&self, namespace: &str, path: &str) -> Option { + let logical_path = LogicalPath::new(namespace, path); + let mut file_sizes = self + .file_sizes + .lock() + .unwrap_or_else(|error| error.into_inner()); + let size = file_sizes.shift_remove(&logical_path)?; + file_sizes.insert(logical_path, size); + Some(size) + } + + pub(super) fn put_file_size(&self, namespace: &str, path: &str, size: u64, capacity: usize) { + let logical_path = LogicalPath::new(namespace, path); + let mut file_sizes = self + .file_sizes + .lock() + .unwrap_or_else(|error| error.into_inner()); + file_sizes.shift_remove(&logical_path); + file_sizes.insert(logical_path, size); + while file_sizes.len() > capacity { + file_sizes.shift_remove_index(0); + } + } + + pub(super) async fn begin_path_invalidation( + &self, + namespace: &str, + path: &str, + ) -> CacheInvalidationGuard { + let prefix_guard = PrefixGuard::Read { + _guard: self.prefix_barrier.clone().read_owned().await, + }; + let state = self.path_state(namespace, path); + let publish_guard = state.publish_gate.clone().write_owned().await; + state + .generation + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.file_sizes + .lock() + .unwrap_or_else(|error| error.into_inner()) + .shift_remove(&LogicalPath::new(namespace, path)); + CacheInvalidationGuard { + _prefix_guard: prefix_guard, + _publish_guards: vec![publish_guard], + } + } + + pub(super) async fn begin_prefix_invalidation( + &self, + namespace: &str, + prefix: &str, + ) -> CacheInvalidationGuard { + let prefix_guard = PrefixGuard::Write { + _guard: self.prefix_barrier.clone().write_owned().await, + }; + let prefix = prefix.trim_end_matches('/'); + let mut states = { + let states = self + .path_states + .lock() + .unwrap_or_else(|error| error.into_inner()); + states + .iter() + .filter(|(path, _)| path.matches_prefix(namespace, prefix)) + .filter_map(|(path, state)| Weak::upgrade(state).map(|state| (path.clone(), state))) + .collect::>() + }; + states.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + let mut publish_guards = Vec::with_capacity(states.len()); + for (_, state) in states { + publish_guards.push(state.publish_gate.clone().write_owned().await); + state + .generation + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + self.file_sizes + .lock() + .unwrap_or_else(|error| error.into_inner()) + .retain(|path, _| !path.matches_prefix(namespace, prefix)); + CacheInvalidationGuard { + _prefix_guard: prefix_guard, + _publish_guards: publish_guards, + } + } +} + +#[derive(Debug)] +enum PrefixGuard { + Read { + _guard: tokio::sync::OwnedRwLockReadGuard<()>, + }, + Write { + _guard: tokio::sync::OwnedRwLockWriteGuard<()>, + }, +} + +#[derive(Debug)] +pub(super) struct CacheInvalidationGuard { + _prefix_guard: PrefixGuard, + _publish_guards: Vec>, +} diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index 0e447a157..3f806fd95 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -1902,6 +1902,10 @@ mod tests { "unit '{unit}' should parse case-insensitively with a space" ); } + + for unit in ["kib", "mib", "gib", "tib"] { + assert_eq!(parse_memory_size(&format!("3{unit}")), None); + } } #[test] diff --git a/docs/src/getting-started.md b/docs/src/getting-started.md index 8d90895e1..c23a9af4c 100644 --- a/docs/src/getting-started.md +++ b/docs/src/getting-started.md @@ -158,24 +158,24 @@ let catalog = CatalogFactory::create(options).await?; | Option | Default | Description | |--------|---------|-------------| | `local-cache.enabled` | `false` | Enable catalog-scoped local block caching. | -| `local-cache.dir` | none | Base cache directory; required when caching is enabled. Paimon stores entries in a private versioned child directory. | -| `local-cache.max-size` | unlimited | Maximum encoded disk usage. Values accept byte units such as `512 MiB` or `20 GiB`. | +| `local-cache.dir` | none | Optional base directory. When set, Paimon uses a persistent disk cache in a private versioned child directory; otherwise it uses memory. | +| `local-cache.max-size` | unlimited | Maximum cache size. Memory caches count payload bytes; disk caches count encoded bytes. Values accept byte units such as `512 MiB` or `20 GiB`. | | `local-cache.block-size` | `1 MiB` | Block size used for cached range reads. | | `local-cache.whitelist` | `meta,global-index` | Comma-separated eligible types: `meta`, `global-index`, `bucket-index`, `data`, and `file-index`. | -The cache is disk-only and is reused after process restarts. Cache keys include -a catalog-configuration fingerprint and the canonical storage object path, so -catalogs can safely use the same base directory without reading one another's -entries. Runtime cache read, write, validation, and eviction failures are -fail-open: the original storage remains the source of truth. Paimon mutable -markers and temporary files always bypass the cache; other eligible files rely -on Paimon's immutable-file convention. Cache managers using the same canonical -directory in one process share LRU and size accounting; if their configured -limits differ, the smallest `local-cache.max-size` is used. Restart recovery -runs on a blocking worker, reads only block headers and file metadata, and -validates payload CRC lazily on the first hit. Use a separate `local-cache.dir` -for each worker or process because processes do not share exact LRU or size -accounting. +Each catalog owns its in-memory cache for the catalog's lifetime. Disk caches +are reused after process restarts. Cache keys include a catalog-configuration +fingerprint and the canonical storage object path, so catalogs can safely use +the same base directory without reading one another's entries. Runtime cache +read, write, validation, and eviction failures are fail-open: the original +storage remains the source of truth. Paimon mutable markers and temporary files +always bypass the cache; other eligible files rely on Paimon's immutable-file +convention. Cache managers using the same canonical directory in one process +share LRU, size accounting, and invalidation state; if their configured limits +differ, the smallest `local-cache.max-size` is used. Restart recovery runs on a +blocking worker, reads only block headers and file metadata, and validates +payload CRC lazily on the first hit. Use a separate `local-cache.dir` for each +worker or process because processes do not share exact LRU or size accounting. ### Manage Databases