diff --git a/src/config.rs b/src/config.rs index e91ed764..59025a1e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -4,6 +4,11 @@ use serde::Deserialize; static CONFIG: OnceLock = OnceLock::new(); +/// Bytes per MiB / GiB — used by the `*_bytes()` size accessors below so the +/// `* 1024 * 1024` chains don't repeat (and read as the unit they mean). +const MIB: usize = 1024 * 1024; +const GIB: usize = 1024 * 1024 * 1024; + /// Load config from environment variables. pub fn load_config_from_env() -> Result { // Load each sub-config separately to avoid #[serde(flatten)] issues with envy @@ -145,8 +150,13 @@ const_default!(d_pressure_flush_pct: u32 = 75); // "none" — never fsync (test/throwaway data only) const_default!(d_wal_fsync_mode: String = "ms"); const_default!(d_wal_max_files: usize = 200); -const_default!(d_foyer_memory_mb: usize = 512); -const_default!(d_foyer_disk_gb: usize = 100); +const_default!(d_foyer_memory_mb: usize = 1024); +// Local disk is cheap and fast relative to S3 GETs, so default the cache large +// — servers run 500GB–1TB cache volumes. foyer creates the backing file sparse +// (no upfront allocation), but this is the logical ceiling at which it starts +// evicting, so it MUST stay <= the cache volume's free space or writes hit +// ENOSPC before eviction kicks in. Lower it on smaller disks. +const_default!(d_foyer_disk_gb: usize = 500); const_default!(d_foyer_ttl: u64 = 604_800); // 7 days const_default!(d_foyer_shards: usize = 8); const_default!(d_foyer_file_size_mb: usize = 32); @@ -155,6 +165,10 @@ const_default!(d_metadata_size_hint: usize = 1_048_576); const_default!(d_metadata_memory_mb: usize = 512); const_default!(d_metadata_disk_gb: usize = 5); const_default!(d_metadata_shards: usize = 4); +const_default!(d_warm_inline_max_mb: usize = 0); +const_default!(d_foyer_block_size_mb: usize = 256); +const_default!(d_l1_max_entry_mb: usize = 16); +const_default!(d_cache_recent_days: usize = 8); const_default!(d_page_rows: usize = 20_000); const_default!(d_zstd_level: i32 = 3); // Tiered compression by partition age. Hot writes prioritize ingest latency; @@ -167,13 +181,20 @@ const_default!(d_row_group_size: usize = 134_217_728); // 128MB const_default!(d_checkpoint_interval: u64 = 10); const_default!(d_optimize_target: i64 = 128 * 1024 * 1024); const_default!(d_stats_cache_size: usize = 50); -const_default!(d_vacuum_retention: u64 = 72); +// Observability data is high-churn and rarely time-traveled; the only hard +// floor is that retention must outlive any in-flight query (which holds a Delta +// snapshot referencing files vacuum would delete). With no query running beyond +// ~1h, 48h is a 48x safety margin while reclaiming tombstoned bytes far sooner +// than the old 72h default. +const_default!(d_vacuum_retention: u64 = 48); const_default!(d_optimize_window_hours: u64 = 48); const_default!(d_compact_min_files: usize = 5); const_default!(d_light_optimize_target: i64 = 16 * 1024 * 1024); const_default!(d_light_schedule: String = "0 */5 * * * *"); const_default!(d_optimize_schedule: String = "0 */30 * * * *"); const_default!(d_vacuum_schedule: String = "0 0 2 * * *"); +const_default!(d_warm_recency_days: u64 = 2); +const_default!(d_warm_concurrency: usize = 4); const_default!(d_mem_gb: usize = 8); const_default!(d_mem_fraction: f64 = 0.9); const_default!(d_otlp_endpoint: String = "http://localhost:4317"); @@ -477,6 +498,40 @@ pub struct CacheConfig { pub timefusion_foyer_metadata_disk_gb: usize, #[serde(default = "d_metadata_shards")] pub timefusion_foyer_metadata_shards: usize, + /// Disk block size (MB) for the main data cache. The block is foyer's + /// minimal eviction unit AND its size caps the largest entry that can land + /// on disk — so it must be >= the largest file we want cached locally. This + /// acts as a *floor*: `from_app_config` automatically raises the effective + /// block size to 2x the compaction target size, so the two can't drift out + /// of sync if an operator bumps the target. Default 256MB. + /// + /// Memory note: this also bounds the transient buffer each multipart-write + /// warm holds in heap (see `timefusion_warm_inline_max_mb`). Up to + /// `timefusion_warm_concurrency` compactions can run at once, so the worst + /// case is `block_size_mb * warm_concurrency` of transient heap during a + /// busy maintenance window (e.g. 256MB x 4 = 1GB). On smaller-memory + /// instances set `timefusion_warm_inline_max_mb` to cap this independently + /// of the on-disk block size. + #[serde(default = "d_foyer_block_size_mb")] + pub timefusion_foyer_block_size_mb: usize, + /// Entries larger than this (MB) are inserted disk-only (foyer + /// `Location::OnDisk`) so warming a big compaction output doesn't evict the + /// hot small-entry working set from L1 memory. Small entries keep the + /// default L1+disk placement for fastest repeat reads. 0 = always use L1. + #[serde(default = "d_l1_max_entry_mb")] + pub timefusion_foyer_l1_max_entry_mb: usize, + /// Don't admit writes whose `date=` partition is older than this many days + /// (e.g. cold-tier recompress rewrites) — recent data stays local, old data + /// is served from S3. 0 = no age limit. Pairs with the cache TTL, which + /// governs how long an admitted entry survives before falling back to S3. + #[serde(default = "d_cache_recent_days")] + pub timefusion_cache_recent_days: usize, + /// Optional extra cap (MB) on the in-flight buffer used to warm the cache + /// directly from a multipart write (so we don't re-download a file we just + /// streamed to S3). Always bounded by the disk block size; 0 = bound only + /// by the block size. + #[serde(default = "d_warm_inline_max_mb")] + pub timefusion_warm_inline_max_mb: usize, #[serde(default)] pub timefusion_foyer_disabled: bool, } @@ -492,20 +547,28 @@ impl CacheConfig { self.timefusion_foyer_stats.eq_ignore_ascii_case("true") } pub fn memory_size_bytes(&self) -> usize { - self.timefusion_foyer_memory_mb * 1024 * 1024 + self.timefusion_foyer_memory_mb * MIB } pub fn disk_size_bytes(&self) -> usize { - self.timefusion_foyer_disk_mb.map_or(self.timefusion_foyer_disk_gb * 1024 * 1024 * 1024, |mb| mb * 1024 * 1024) + self.timefusion_foyer_disk_mb.map_or(self.timefusion_foyer_disk_gb * GIB, |mb| mb * MIB) } pub fn file_size_bytes(&self) -> usize { - self.timefusion_foyer_file_size_mb * 1024 * 1024 + self.timefusion_foyer_file_size_mb * MIB } pub fn metadata_memory_size_bytes(&self) -> usize { - self.timefusion_foyer_metadata_memory_mb * 1024 * 1024 + self.timefusion_foyer_metadata_memory_mb * MIB + } + pub fn warm_inline_max_bytes(&self) -> usize { + self.timefusion_warm_inline_max_mb * MIB + } + pub fn block_size_bytes(&self) -> usize { + self.timefusion_foyer_block_size_mb * MIB + } + pub fn l1_max_entry_bytes(&self) -> usize { + self.timefusion_foyer_l1_max_entry_mb * MIB } pub fn metadata_disk_size_bytes(&self) -> usize { - self.timefusion_foyer_metadata_disk_mb - .map_or(self.timefusion_foyer_metadata_disk_gb * 1024 * 1024 * 1024, |mb| mb * 1024 * 1024) + self.timefusion_foyer_metadata_disk_mb.map_or(self.timefusion_foyer_metadata_disk_gb * GIB, |mb| mb * MIB) } } @@ -553,6 +616,32 @@ pub struct MaintenanceConfig { pub timefusion_vacuum_schedule: String, #[serde(default = "d_recompress_schedule")] pub timefusion_recompress_schedule: String, + /// Proactively warm the Foyer cache for files written by a flush/optimize + /// commit, so recent partitions dashboards read don't cold-start after + /// every compaction. Footers are always warmed when enabled. + #[serde(default = "d_true")] + pub timefusion_warm_after_compaction: bool, + /// In addition to footers, warm the full file contents into the main + /// (full-file) cache. Off by default — footers carry most of the + /// planning-latency win at a fraction of the bytes; enable for data-read + /// warmth on the hottest partitions. + #[serde(default)] + pub timefusion_warm_full_files: bool, + /// Only warm files whose `date=` partition is within this many days of + /// today. Bounds warming to the partitions dashboards actually query. + /// 0 = no recency limit. + #[serde(default = "d_warm_recency_days")] + pub timefusion_warm_recency_days: u64, + /// Max concurrent warm fetches per commit. Bounds the S3 GET burst a + /// warm job adds right after a compaction. + #[serde(default = "d_warm_concurrency")] + pub timefusion_warm_concurrency: usize, + /// After a compaction commit, proactively evict the cached full-file bytes + /// of the files it tombstoned (no longer in the live set), instead of + /// waiting for VACUUM / TTL / LRU to reclaim them. Cheap (in-cache only, no + /// S3) and keeps the cache from filling with dead compaction outputs. + #[serde(default = "d_true")] + pub timefusion_evict_after_compaction: bool, } /// Which DataFusion `MemoryPool` to back the runtime with. @@ -632,7 +721,19 @@ mod tests { let config = AppConfig::default(); assert_eq!(config.core.pgwire_port, 5432); assert_eq!(config.buffer.timefusion_flush_interval_secs, 600); - assert_eq!(config.cache.timefusion_foyer_memory_mb, 512); + assert_eq!(config.cache.timefusion_foyer_memory_mb, 1024); + assert_eq!(config.cache.timefusion_foyer_disk_gb, 500); + assert_eq!(config.cache.disk_size_bytes(), 500 * 1024 * 1024 * 1024); + assert_eq!(config.cache.timefusion_warm_inline_max_mb, 0); + assert_eq!(config.cache.timefusion_foyer_block_size_mb, 256); + assert_eq!(config.cache.block_size_bytes(), 256 * 1024 * 1024); + assert_eq!(config.cache.timefusion_foyer_l1_max_entry_mb, 16); + assert_eq!(config.cache.timefusion_cache_recent_days, 8); + assert!(config.maintenance.timefusion_warm_after_compaction); + assert!(config.maintenance.timefusion_evict_after_compaction); + assert!(!config.maintenance.timefusion_warm_full_files); + assert_eq!(config.maintenance.timefusion_warm_recency_days, 2); + assert_eq!(config.maintenance.timefusion_warm_concurrency, 4); } #[test] diff --git a/src/database.rs b/src/database.rs index ff9f8e8e..5f0e4181 100644 --- a/src/database.rs +++ b/src/database.rs @@ -76,6 +76,32 @@ fn should_refresh_table(current_version: Option, last_written_version: Opti } } +/// Whether `uri` belongs to a partition no older than `cutoff` (inclusive). +/// Parses the `date=YYYY-MM-DD` Hive partition segment; if absent or +/// unparseable, returns `true` (warm rather than silently skip a file we can't +/// classify). A `None` cutoff means "no recency limit". +fn within_recency(uri: &str, cutoff: Option) -> bool { + // Single source of truth for `date=` partition recency parsing, shared with + // the object-store cache admission window. + crate::object_store_cache::date_partition_within(uri, cutoff) +} + +/// The cache-key prefix for a table: its URI minus any `?endpoint=...` query +/// string (`table_url()` may carry one; `get_file_uris()` omits it) and trailing +/// slash. File URIs are relativized against this to form cache keys. +fn table_cache_prefix(table_uri: &str) -> &str { + table_uri.split('?').next().unwrap_or(table_uri).trim_end_matches('/') +} + +/// Relativize an absolute file URI against a `table_cache_prefix`, yielding the +/// bucket-relative path the cached object store keys full files by. `None` on +/// prefix mismatch (trailing-slash or query-string drift between `table_url()` +/// and `get_file_uris()`). Shared by the warm and evict paths so a single-char +/// difference can't desync which key was warmed vs. evicted. +fn relativize_to_prefix(prefix: &str, uri: &str) -> Option { + uri.strip_prefix(prefix).map(|rel| object_store::path::Path::from(rel.trim_start_matches('/'))) +} + // Helper function to extract project_id from a batch pub fn extract_project_id(batch: &RecordBatch) -> Option { use datafusion::arrow::array::{StringArray, StringViewArray}; @@ -1604,6 +1630,208 @@ impl Database { Ok(uris) } + /// Best-effort warm of the Foyer cache for parquet files just written by a + /// flush or optimize commit. Reuses the read path so the recent partitions + /// dashboards query don't cold-start after every compaction: a ranged GET + /// of each new footer primes the metadata cache (query planning pays zero + /// S3 round-trips), and — when `timefusion_warm_full_files` is set — a full + /// GET primes the main cache for data reads. + /// + /// Non-blocking and strictly best-effort: the whole job runs in a detached, + /// concurrency-bounded task and never affects the commit. Files are filtered + /// to partitions within `timefusion_warm_recency_days` so we don't spend S3 + /// GETs (and evict useful entries) warming cold partitions nobody reads. + async fn warm_cache_for_uris(&self, object_store: Arc, table_uri: String, uris: Vec) { + let maint = &self.config.maintenance; + if !maint.timefusion_warm_after_compaction || uris.is_empty() { + return; + } + let warm_full_files = maint.timefusion_warm_full_files; + let recency_days = maint.timefusion_warm_recency_days; + let concurrency = maint.timefusion_warm_concurrency.max(1); + let metadata_size_hint = self.config.cache.timefusion_parquet_metadata_size_hint as u64; + let stats_cache = self.object_store_cache.clone(); + + // Relativize absolute s3:// URIs against the table root: the cached + // object store consumes bucket-relative paths. + let prefix = table_cache_prefix(&table_uri); + // Cap the day count before the i64 cast — recency_days is a config + // value so overflow can't happen in practice, but a silent wrap would + // turn a misconfiguration into "warm nothing". 3650d (~10y) is well + // past any partition we'd query. + let cutoff = (recency_days > 0).then(|| Utc::now().date_naive() - chrono::Duration::days(recency_days.min(3650) as i64)); + + let mut dropped = 0usize; + let paths: Vec = uris + .into_iter() + .filter(|u| u.ends_with(".parquet")) + .filter(|u| within_recency(u, cutoff)) + .filter_map(|u| match relativize_to_prefix(prefix, &u) { + Some(path) => Some(path), + None => { + // Prefix mismatch (e.g. trailing-slash or query-string + // drift between table_url() and get_file_uris()). Warming + // this file would address the wrong key, so skip it — but + // log so a systematic mismatch is diagnosable instead of a + // silent no-op. + if dropped == 0 { + debug!("warm: URI {} does not start with table prefix {}; skipping (warm only)", u, prefix); + } + dropped += 1; + None + } + }) + .collect(); + + if dropped > 0 { + debug!("warm: skipped {} file(s) that did not relativize against prefix {}", dropped, prefix); + } + if paths.is_empty() { + return; + } + + let count = paths.len(); + // Baseline the cache stats *before* warming: the warm GETs are all + // misses (they fetch from the inner store to populate Foyer), so a + // post-warm hit rate would read artificially low. The real + // beneficiary is the next dashboard query — log the pre-warm + // steady-state rate as the relevant baseline. + let baseline = match &stats_cache { + Some(cache) => { + let s = cache.get_stats().await.main; + let rate = if s.hits + s.misses > 0 { (s.hits as f64 / (s.hits + s.misses) as f64) * 100.0 } else { 0.0 }; + Some(rate) + } + None => None, + }; + + // Labelled scope rather than `full=true/false` so warm logs are easy to + // filter (e.g. in Loki) by what was actually primed. + let scope = if warm_full_files { "full" } else { "footer-only" }; + futures::stream::iter(paths) + .for_each_concurrent(concurrency, |path| { + let store = object_store.clone(); + async move { + let _ = crate::object_store_cache::warm_footer(store.as_ref(), &path, metadata_size_hint).await; + if warm_full_files { + let _ = crate::object_store_cache::warm_full(store.as_ref(), &path).await; + } + } + }) + .await; + + match baseline { + Some(rate) => info!( + "Cache warm complete: {} files warmed (scope={}); foyer main hit rate before warm was {:.2}% (next query benefits)", + count, scope, rate + ), + None => info!("Cache warm complete: {} files warmed (scope={})", count, scope), + } + } + + /// Proactively evict the cached full-file bytes of files a compaction + /// tombstoned (present pre-commit, gone post-commit), so dead compaction + /// outputs don't linger in the cache until VACUUM / TTL / LRU reclaims them. + /// + /// Correctness is unaffected: the files still exist in S3 until VACUUM, so a + /// straggler query holding the old Delta snapshot just re-reads them from S3 + /// (a cache miss), never a wrong result. Cheap and in-cache only (no S3), + /// so it runs inline. + fn evict_cache_for_uris(&self, table_uri: &str, removed: &[String]) { + if !self.config.maintenance.timefusion_evict_after_compaction || removed.is_empty() { + return; + } + let Some(cache) = self.object_store_cache.as_ref() else { + return; + }; + // Same relativization as warm_cache_for_uris: the cache keys full files + // by their object-store-relative path. + let prefix = table_cache_prefix(table_uri); + let mut evicted = 0usize; + let mut dropped = 0usize; + for u in removed { + if let Some(path) = relativize_to_prefix(prefix, u) { + cache.evict_data_entry(path.as_ref()); + evicted += 1; + } else { + // Prefix mismatch (trailing-slash or query-string drift between + // table_url() and get_file_uris()) — we'd evict the wrong key, so + // skip. Log like the warm path: a systematic mismatch here means + // tombstoned files linger in cache until TTL/LRU, which is worth + // diagnosing rather than silently swallowing. + if dropped == 0 { + debug!("evict: URI {} does not start with table prefix {}; skipping (evict only)", u, prefix); + } + dropped += 1; + } + } + if evicted > 0 { + debug!("Evicted {} tombstoned file(s) from cache after compaction", evicted); + } + if dropped > 0 { + debug!("evict: skipped {} file(s) that did not relativize against prefix {}", dropped, prefix); + } + } + + /// Warm the cache for files added by a just-committed flush/optimize on the + /// given logical table. Fire-and-forget: resolving the table (which may + /// issue a rate-limited PG roundtrip) and taking the read lock both happen + /// inside a spawned task, so the caller — notably the flush callback — is + /// never blocked. No-op when warming is disabled or the list is empty. + pub fn warm_cache_for_table(&self, project_id: &str, table_name: &str, uris: Vec) { + if uris.is_empty() || !self.config.maintenance.timefusion_warm_after_compaction { + return; + } + let db = self.clone(); + let project_id = project_id.to_string(); + let table_name = table_name.to_string(); + tokio::spawn(async move { + if let Ok(table_ref) = db.resolve_table(&project_id, &table_name).await { + let (store, table_uri) = { + let t = table_ref.read().await; + (t.log_store().object_store(None), t.table_url().to_string()) + }; + // Already inside a detached task — await the warm directly + // instead of spawning a second nested task. + db.warm_cache_for_uris(store, table_uri, uris).await; + } + }); + } + + /// Atomically swap a freshly-optimized `new_table` in under the write lock, + /// then refresh the cache for the file-set delta vs `pre_uris`: warm the + /// files this optimize added and evict the ones it tombstoned. Returns the + /// new table's live file URIs (captured before the swap) for callers that + /// need them (e.g. the tantivy GC hook). + /// + /// Both optimize paths — full Z-order and light — funnel through here so the + /// warm/evict pair can't drift; the evict call was once missing from the + /// light path, and a single helper keeps them in lockstep. + async fn swap_and_refresh_cache( + &self, table_ref: &Arc>, new_table: DeltaTable, pre_uris: &std::collections::HashSet, + ) -> Vec { + // Capture live URIs off `new_table` *before* the swap moves it in. + let live_uris: Vec = new_table.get_file_uris().map(|it| it.collect()).unwrap_or_default(); + let live_set: std::collections::HashSet<&String> = live_uris.iter().collect(); + let added: Vec = live_uris.iter().filter(|u| !pre_uris.contains(*u)).cloned().collect(); + let removed: Vec = pre_uris.iter().filter(|u| !live_set.contains(u)).cloned().collect(); + let warm_store = new_table.log_store().object_store(None); + let warm_table_uri = new_table.table_url().to_string(); + { + let mut table = table_ref.write().await; + *table = new_table; + } + // Eviction is in-cache only (cheap), so run it inline. Warming issues S3 + // GETs, so detach it — the maintenance loop shouldn't block on priming + // the cache (preserves the previous in-`warm_cache_for_uris` spawn). + self.evict_cache_for_uris(&warm_table_uri, &removed); + let db = self.clone(); + tokio::spawn(async move { + db.warm_cache_for_uris(warm_store, warm_table_uri, added).await; + }); + live_uris + } + pub async fn get_or_create_table(&self, project_id: &str, table_name: &str) -> Result>> { // Route to appropriate table based on whether project has custom storage if self.has_custom_storage(project_id, table_name).await { @@ -1933,6 +2161,15 @@ impl Database { table.clone() }; + // Pre-state file set, used to derive the files this optimize *adds* + // (to warm) and *removes* (to evict) — see warm/evict_cache_for_uris. + let track_files = self.config.maintenance.timefusion_warm_after_compaction || self.config.maintenance.timefusion_evict_after_compaction; + let pre_uris: std::collections::HashSet = if track_files { + table_clone.get_file_uris().map(|it| it.collect()).unwrap_or_default() + } else { + Default::default() + }; + let target_size = self.config.parquet.timefusion_optimize_target_size; // Generate partition filters for each date in the configurable window @@ -1996,13 +2233,10 @@ impl Database { let compression_ratio = metrics.num_files_removed as f64 / metrics.num_files_added as f64; info!("Optimization compression ratio: {:.2}x", compression_ratio); } - // Capture live file URIs from the new table *before* taking - // the write lock to swap it in — used by the tantivy GC hook - // below to drop indexes whose covered files no longer exist. - let live_uris: Vec = new_table.get_file_uris().map(|it| it.collect()).unwrap_or_default(); - let mut table = table_ref.write().await; - *table = new_table; - drop(table); + // Swap the optimized table in and refresh the cache (warm + // newly-added files, evict tombstoned ones). Returns the new + // live file URIs for the tantivy GC hook below. + let live_uris = self.swap_and_refresh_cache(table_ref, new_table, &pre_uris).await; // Tantivy compaction GC — drop sidecar indexes for files that // were rewritten away. Best-effort: errors are logged. if let Some(svc) = self.tantivy_indexer().cloned() { @@ -2212,6 +2446,14 @@ impl Database { let table = table_ref.read().await; table.clone() }; + // Pre-state file set for deriving the files this optimize adds (to + // warm) and removes (to evict). + let track_files = self.config.maintenance.timefusion_warm_after_compaction || self.config.maintenance.timefusion_evict_after_compaction; + let pre_uris: std::collections::HashSet = if track_files { + table_clone.get_file_uris().map(|it| it.collect()).unwrap_or_default() + } else { + Default::default() + }; if attempt == 0 { info!("Light optimizing files from date: {}", today); } else { @@ -2248,8 +2490,10 @@ impl Database { metrics.num_files_removed, metrics.num_files_added ); - let mut table = table_ref.write().await; - *table = new_table; + // Swap the optimized table in and refresh the cache (warm + // freshly-compacted files, evict the small files just + // tombstoned) via the shared helper. + let _ = self.swap_and_refresh_cache(table_ref, new_table, &pre_uris).await; return Ok(()); } Err(e) => { @@ -3284,6 +3528,30 @@ mod tests { use super::*; use crate::{config::AppConfig, test_utils::test_helpers::*}; + #[test] + fn test_within_recency() { + let cutoff = chrono::NaiveDate::from_ymd_opt(2026, 6, 4); + + // Files on/after the cutoff date are warmed. + assert!(within_recency("s3://b/t/date=2026-06-06/part-0.parquet", cutoff)); + assert!(within_recency("s3://b/t/date=2026-06-04/part-0.parquet", cutoff), "cutoff is inclusive"); + // Older partitions are skipped. + assert!(!within_recency("s3://b/t/date=2026-06-01/part-0.parquet", cutoff)); + + // No `date=` segment → warm (don't silently skip an unclassifiable file). + assert!(within_recency("s3://b/t/part-0.parquet", cutoff)); + // Unparseable date → warm. + assert!(within_recency("s3://b/t/date=not-a-date/part-0.parquet", cutoff)); + // Truncated date (segment shorter than YYYY-MM-DD) → warm. + assert!(within_recency("s3://b/t/date=2026-06", cutoff)); + + // None cutoff → no recency limit, always warm even very old partitions. + assert!(within_recency("s3://b/t/date=2000-01-01/part-0.parquet", None)); + + // Nested partitioning (project_id then date) still locates `date=`. + assert!(!within_recency("s3://b/t/project_id=default/date=2026-05-01/part.parquet", cutoff)); + } + /// Roundtrip the watermark through serialize → JSON → parse. Pins the /// on-disk format so a future change to `serialize_watermark_to_json` /// can't silently break `derive_wal_cursors_from_delta`. Absent shards diff --git a/src/main.rs b/src/main.rs index a740220f..8cf97227 100644 --- a/src/main.rs +++ b/src/main.rs @@ -92,6 +92,11 @@ async fn async_main(cfg: &'static AppConfig) -> anyhow::Result<()> { let post = db.list_file_uris(&project_id, &table_name).await.unwrap_or_default(); let pre_set: std::collections::HashSet = pre.into_iter().collect(); let added: Vec = post.into_iter().filter(|u| !pre_set.contains(u)).collect(); + // Warm the just-flushed files into the cache so the first + // dashboard read of this partition hits Foyer, not S3. + // Fire-and-forget (spawns internally); clone since `added` is + // also returned for the sidecar tantivy indexer. + db.warm_cache_for_table(&project_id, &table_name, added.clone()); Ok(added) }) }, diff --git a/src/object_store_cache.rs b/src/object_store_cache.rs index 04e9e18f..cf6dbb33 100644 --- a/src/object_store_cache.rs +++ b/src/object_store_cache.rs @@ -9,7 +9,9 @@ use async_trait::async_trait; use bytes::Bytes; use chrono::{DateTime, Utc}; use dashmap::DashSet; -use foyer::{BlockEngineConfig, DeviceBuilder, FsDeviceBuilder, HybridCache, HybridCacheBuilder, HybridCachePolicy, PsyncIoEngineConfig}; +use foyer::{ + BlockEngineConfig, DeviceBuilder, FsDeviceBuilder, HybridCache, HybridCacheBuilder, HybridCachePolicy, HybridCacheProperties, Location, PsyncIoEngineConfig, +}; use futures::stream::BoxStream; use object_store::{ Attributes, CopyOptions, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore, ObjectStoreExt, @@ -112,6 +114,20 @@ pub struct FoyerCacheConfig { pub metadata_disk_size_bytes: usize, /// Number of shards for metadata cache pub metadata_shards: usize, + /// Optional extra cap on bytes buffered to warm the cache inline from a + /// multipart write (see `CachingMultipartUpload`). Always bounded by + /// `block_size_bytes`; 0 = bound only by the block size. + pub warm_inline_max_bytes: usize, + /// Disk block size for the main data cache — foyer's eviction unit and the + /// hard cap on the largest entry that can persist to disk. Must be >= the + /// largest file we want cached (compaction target size). + pub block_size_bytes: usize, + /// Entries larger than this are inserted disk-only (`Location::OnDisk`) so + /// they don't evict the hot L1 working set. 0 = always use L1. + pub l1_max_entry_bytes: usize, + /// Don't admit writes whose `date=` partition is older than this many days. + /// 0 = no age limit. + pub cache_recent_days: usize, } impl Default for FoyerCacheConfig { @@ -128,24 +144,48 @@ impl Default for FoyerCacheConfig { metadata_memory_size_bytes: 67_108_864, // 64MB metadata_disk_size_bytes: 536_870_912, // 512MB metadata_shards: 4, // Fewer shards for metadata cache + warm_inline_max_bytes: 0, // bound by block size + block_size_bytes: 268_435_456, // 256MB — fits 128MB compaction outputs + l1_max_entry_bytes: 16_777_216, // 16MB + cache_recent_days: 8, } } } impl FoyerCacheConfig { pub fn from_app_config(cfg: &crate::config::AppConfig) -> Self { + // The disk block size caps the largest file that can be cached locally, + // and compaction writes files at ~the optimize target size. Floor the + // block at 2x that target so the two stay in lockstep automatically — + // an operator can raise timefusion_optimize_target_size without + // silently losing the ability to cache the bigger outputs. The + // configured block size acts as a lower bound / explicit override. + let optimize_target = cfg.parquet.timefusion_optimize_target_size.max(0) as usize; + let block_size_bytes = cfg.cache.block_size_bytes().max(optimize_target.saturating_mul(2)); + let disk_size_bytes = cfg.cache.disk_size_bytes(); + if block_size_bytes > disk_size_bytes { + tracing::warn!( + "Foyer disk block size ({}MB) exceeds disk capacity ({}MB) — large files won't persist to disk. Raise timefusion_foyer_disk_gb or lower the optimize target.", + block_size_bytes / 1024 / 1024, + disk_size_bytes / 1024 / 1024 + ); + } Self { - memory_size_bytes: cfg.cache.memory_size_bytes(), - disk_size_bytes: cfg.cache.disk_size_bytes(), - ttl: cfg.cache.ttl(), - cache_dir: cfg.core.cache_dir(), - shards: cfg.cache.timefusion_foyer_shards, - file_size_bytes: cfg.cache.file_size_bytes(), - enable_stats: cfg.cache.stats_enabled(), + memory_size_bytes: cfg.cache.memory_size_bytes(), + disk_size_bytes, + ttl: cfg.cache.ttl(), + cache_dir: cfg.core.cache_dir(), + shards: cfg.cache.timefusion_foyer_shards, + file_size_bytes: cfg.cache.file_size_bytes(), + enable_stats: cfg.cache.stats_enabled(), parquet_metadata_size_hint: cfg.cache.timefusion_parquet_metadata_size_hint, metadata_memory_size_bytes: cfg.cache.metadata_memory_size_bytes(), - metadata_disk_size_bytes: cfg.cache.metadata_disk_size_bytes(), - metadata_shards: cfg.cache.timefusion_foyer_metadata_shards, + metadata_disk_size_bytes: cfg.cache.metadata_disk_size_bytes(), + metadata_shards: cfg.cache.timefusion_foyer_metadata_shards, + warm_inline_max_bytes: cfg.cache.warm_inline_max_bytes(), + block_size_bytes, + l1_max_entry_bytes: cfg.cache.l1_max_entry_bytes(), + cache_recent_days: cfg.cache.timefusion_cache_recent_days, } } @@ -164,6 +204,10 @@ impl FoyerCacheConfig { metadata_memory_size_bytes: 10 * 1024 * 1024, // 10MB for tests metadata_disk_size_bytes: 50 * 1024 * 1024, // 50MB for tests metadata_shards: 2, + warm_inline_max_bytes: 0, // bound by block size + block_size_bytes: 4 * 1024 * 1024, // 4MB — must be <= test disk size + l1_max_entry_bytes: 1024 * 1024, // 1MB + cache_recent_days: 0, // no age limit in tests (avoid date flakiness) } } @@ -209,6 +253,21 @@ impl CacheStats { type FoyerCache = Arc>; type StatsRef = Arc>; +/// Floor for the foyer disk block (region) size. Matches the legacy default +/// (`timefusion_foyer_file_size_mb`), small enough that even a modest disk +/// budget yields several regions. +const MIN_DISK_BLOCK_BYTES: usize = 4 * 1024 * 1024; + +/// Cap a desired foyer disk block (region) size to the device. Foyer carves the +/// device into block-sized regions, so a block >= the device leaves zero usable +/// regions and every disk insert stalls (a 256MB block on a 50MB device wedged +/// CI). Keep several regions by capping at a quarter of the device, floored at +/// the legacy 4MB granularity and never above the device itself. Shared by both +/// cache builders so neither can silently wedge on a small disk. +fn capped_block_size(desired: usize, disk_size: usize) -> usize { + desired.min(disk_size / 4).max(MIN_DISK_BLOCK_BYTES).min(disk_size) +} + /// Shared Foyer cache that can be used across multiple object stores #[derive(Debug)] pub struct SharedFoyerCache { @@ -223,9 +282,10 @@ impl SharedFoyerCache { /// Create a new shared Foyer cache pub async fn new(config: FoyerCacheConfig) -> anyhow::Result { info!( - "Initializing shared Foyer hybrid cache (memory: {}MB, disk: {}GB, ttl: {}s, parquet_metadata_hint: {}KB)", + "Initializing shared Foyer hybrid cache (memory: {}MB, disk: {}GB, block: {}MB, ttl: {}s, parquet_metadata_hint: {}KB)", config.memory_size_bytes / 1024 / 1024, config.disk_size_bytes / 1024 / 1024 / 1024, + config.block_size_bytes / 1024 / 1024, config.ttl.as_secs(), config.parquet_metadata_size_hint / 1024 ); @@ -241,6 +301,10 @@ impl SharedFoyerCache { let metadata_cache_dir = config.cache_dir.join("metadata"); std::fs::create_dir_all(&metadata_cache_dir)?; + // The main data cache wants a block big enough to hold full compaction + // outputs (128MB) so they persist, capped to the device so it can't wedge. + let data_block_size = capped_block_size(config.block_size_bytes, config.disk_size_bytes); + let cache = HybridCacheBuilder::new() .with_policy(HybridCachePolicy::WriteOnInsertion) .memory(config.memory_size_bytes) @@ -249,8 +313,7 @@ impl SharedFoyerCache { .storage() .with_io_engine_config(PsyncIoEngineConfig::new()) .with_engine_config( - BlockEngineConfig::new(FsDeviceBuilder::new(&config.cache_dir).with_capacity(config.disk_size_bytes).build()?) - .with_block_size(config.file_size_bytes), + BlockEngineConfig::new(FsDeviceBuilder::new(&config.cache_dir).with_capacity(config.disk_size_bytes).build()?).with_block_size(data_block_size), ) .build() .await?; @@ -264,7 +327,7 @@ impl SharedFoyerCache { .with_io_engine_config(PsyncIoEngineConfig::new()) .with_engine_config( BlockEngineConfig::new(FsDeviceBuilder::new(&metadata_cache_dir).with_capacity(config.metadata_disk_size_bytes).build()?) - .with_block_size(config.file_size_bytes), + .with_block_size(capped_block_size(config.file_size_bytes, config.metadata_disk_size_bytes)), ) .build() .await?; @@ -311,6 +374,14 @@ impl SharedFoyerCache { info!("Invalidating _last_checkpoint cache for table: {}", table_path); self.cache.remove(&last_checkpoint_key); } + + /// Best-effort eviction of a main (full-file) cache entry by its key — the + /// relativized object path, matching `make_cache_key`. Used to proactively + /// drop the (now dead) full-file bytes of a file a compaction tombstoned, + /// instead of waiting for VACUUM / TTL / LRU to reclaim them. + pub fn evict_data_entry(&self, key: &str) { + self.cache.remove(key); + } } /// Strip the `scheme://` prefix and trailing slashes from a table URI, yielding @@ -326,6 +397,119 @@ fn is_parquet_file(location: &Path) -> bool { location.as_ref().ends_with(".parquet") } +/// Best-effort: warm the Parquet footer of `location` into the cache by issuing +/// a ranged GET of the last `metadata_size_hint` bytes through `store`. When +/// `store` is a [`FoyerObjectStoreCache`], that ranged GET lands in the +/// metadata cache, so subsequent query planning (footer parse, row-group +/// stats, schema, pruning) pays zero S3 round-trips. The single `head` resolves +/// the file size needed to address the tail. +/// +/// Strictly best-effort: every error is swallowed and reported via the return +/// value. Warming must never affect correctness or a caller's commit. Returns +/// `true` if the footer range was fetched. +pub async fn warm_footer(store: &dyn ObjectStore, location: &Path, metadata_size_hint: u64) -> bool { + // Single suffix GET: the response carries the resolved absolute range + total + // size, so a `FoyerObjectStoreCache` caches the footer under the same + // absolute key a later bounded footer read requests — one round-trip, no + // separate HEAD. Falls back to HEAD + bounded GET for stores that don't + // support suffix ranges. + let opts = GetOptions { + range: Some(GetRange::Suffix(metadata_size_hint.max(1))), + ..Default::default() + }; + match store.get_opts(location, opts).await { + Ok(result) => result.bytes().await.is_ok(), + Err(_) => warm_footer_via_head(store, location, metadata_size_hint).await, + } +} + +/// HEAD + bounded-GET fallback for [`warm_footer`] when the store doesn't +/// support suffix ranges. Two round-trips, but always correct. +async fn warm_footer_via_head(store: &dyn ObjectStore, location: &Path, metadata_size_hint: u64) -> bool { + let size = match store.head(location).await { + Ok(meta) => meta.size, + Err(_) => return false, + }; + if size == 0 { + return false; + } + let start = size.saturating_sub(metadata_size_hint.max(1)); + let opts = GetOptions { + range: Some(GetRange::Bounded(start..size)), + ..Default::default() + }; + store.get_opts(location, opts).await.is_ok() +} + +/// Best-effort: warm the full contents of `location` into the cache via a plain +/// GET through `store`. For a [`FoyerObjectStoreCache`] this populates the main +/// (full-file) cache so ranged data reads — DataFusion row-group scans — hit +/// Foyer instead of S3. Errors are swallowed; see [`warm_footer`]. +pub async fn warm_full(store: &dyn ObjectStore, location: &Path) -> bool { + // Explicitly drain the body: `GetResult`'s payload is a stream, and a + // generic store may not have read it yet by the time `get_opts` returns. + // (`FoyerObjectStoreCache` populates the cache eagerly inside `get_opts`, + // but consuming the bytes keeps this correct for any inner store and is + // a no-op cost there.) + match store.get_opts(location, GetOptions::default()).await { + Ok(result) => result.bytes().await.is_ok(), + Err(_) => false, + } +} + +/// Parse the `date=YYYY-MM-DD` partition segment from `s` and return whether it +/// is on or after `cutoff`. Strings without a parseable date segment (Delta log, +/// checkpoints) are always within the window; a `None` cutoff means no age limit. +/// +/// Shared by the cache-admission window here and the compaction-warm recency +/// filter in `database.rs` so the two parsers can't drift. +pub fn date_partition_within(s: &str, cutoff: Option) -> bool { + let Some(cutoff) = cutoff else { return true }; + match s.find("date=").and_then(|i| s.get(i + 5..i + 15)).and_then(|d| chrono::NaiveDate::parse_from_str(d, "%Y-%m-%d").ok()) { + Some(date) => date >= cutoff, + None => true, + } +} + +/// Whether `location` should be admitted to the cache given the recent-days +/// window. Paths without a `date=YYYY-MM-DD` segment (Delta log, checkpoints) +/// are always admitted. 0 days = no age limit. +/// +/// Keeps cold-tier rewrites (recompress of week+-old partitions) out of the +/// cache so recent data stays local and old data is served from S3. +fn is_within_recent_window(location: &Path, recent_days: usize) -> bool { + if recent_days == 0 { + return true; + } + let cutoff = Utc::now().date_naive() - chrono::Duration::days(recent_days as i64); + date_partition_within(location.as_ref(), Some(cutoff)) +} + +/// Insert into the main full-file cache, steering large entries to disk-only +/// (`Location::OnDisk`, which makes them phantom in L1) so warming a 128MB +/// compaction output doesn't evict the hot small-entry working set from memory. +/// Small entries keep the default L1+disk placement for fastest repeat reads. +fn insert_main(cache: &FoyerCache, key: String, value: CacheValue, l1_max_entry_bytes: usize) { + if l1_max_entry_bytes > 0 && value.data.len() > l1_max_entry_bytes { + cache.insert_with_properties(key, value, HybridCacheProperties::default().with_location(Location::OnDisk)); + } else { + cache.insert(key, value); + } +} + +/// Synthesize an `ObjectMeta` for a just-written object from its `PutResult` +/// (e_tag/version) and known size — lets the write path warm the cache without +/// a post-write GET just to learn the metadata. +fn put_result_meta(location: Path, size: u64, result: &PutResult) -> ObjectMeta { + ObjectMeta { + location, + last_modified: Utc::now(), + size, + e_tag: result.e_tag.clone(), + version: result.version.clone(), + } +} + /// Foyer-based hybrid cache implementation for object store #[derive(derive_more::Display, derive_more::Debug)] #[display("FoyerHybridCachedObjectStore({})", inner)] @@ -413,6 +597,14 @@ impl FoyerObjectStoreCache { format!("{}#range:{}-{}", location, range.start, range.end) } + /// Key for a path's `ObjectMeta` in the metadata cache, kept distinct from + /// range keys (`#range:`). Delta data files are immutable, so a path's + /// size/etag is stable for the cache TTL — caching it lets footer reads skip + /// the per-read HEAD. + fn make_meta_cache_key(location: &Path) -> String { + format!("{}#meta", location) + } + /// Invalidate all metadata cache entries for a given file async fn invalidate_metadata_cache(&self, location: &Path) { // We can't enumerate all possible range keys, but we can at least @@ -582,7 +774,9 @@ impl FoyerObjectStoreCache { current_millis().saturating_sub(value.timestamp_millis), value.data.len() ); - return Ok(Self::make_get_result(Bytes::from(value.data.clone()), value.meta.clone())); + let result = Self::make_get_result(Bytes::from(value.data.clone()), value.meta.clone()); + self.maybe_touch(&self.cache, &cache_key, entry.clone(), self.config.l1_max_entry_bytes); + return Ok(result); } } @@ -633,7 +827,7 @@ impl FoyerObjectStoreCache { } }; - self.cache.insert(cache_key, CacheValue::new(data.clone(), result.meta.clone())); + self.insert_main_value(location, CacheValue::new(data.clone(), result.meta.clone())); Ok(Self::make_get_result(Bytes::from(data), result.meta)) } @@ -671,14 +865,44 @@ impl FoyerObjectStoreCache { is_parquet, current_millis().saturating_sub(value.timestamp_millis) ); - return Ok(Bytes::from(value.data[range.start as usize..range.end as usize].to_vec())); + let sliced = Bytes::from(value.data[range.start as usize..range.end as usize].to_vec()); + self.maybe_touch(&self.cache, &full_cache_key, entry.clone(), self.config.l1_max_entry_bytes); + return Ok(sliced); } } // For Parquet files, implement smart caching based on the range if is_parquet { - // First get the file size to determine if this is a metadata request - let file_meta = match self.inner.head(location).await { + // Probe the metadata range cache *before* any HEAD: its key is just + // (location, range), so a steady-state footer read served from cache + // pays zero S3 round-trips. Data ranges aren't stored here and fall + // through to the size-based classification below. + let range_cache_key = Self::make_range_cache_key(location, &range); + if let Ok(Some(entry)) = self.metadata_cache.get(&range_cache_key).await { + let value = entry.value(); + if !value.is_expired(self.config.ttl) { + self.update_metadata_stats(|s| s.hits += 1).await; + span.record("cache_hit", true); + span.record("is_metadata", true); + debug!( + "Metadata cache HIT for: {} (range: {}..{}, size: {} bytes, age={}ms)", + location, + range.start, + range.end, + value.data.len(), + current_millis().saturating_sub(value.timestamp_millis) + ); + let sliced = Bytes::from(value.data.clone()); + // l1_max=0: metadata entries are tiny, always keep in L1. + self.maybe_touch(&self.metadata_cache, &range_cache_key, entry.clone(), 0); + return Ok(sliced); + } + } + + // Range-cache miss: we need the file size to classify the request and + // to stamp the cached range's meta. Use the cached ObjectMeta + // (immutable Delta files) so this HEAD is paid at most once per file. + let file_meta = match self.head_cached(location).await { Ok(meta) => meta, Err(e) => { debug!("Failed to get metadata for {}: {}", location, e); @@ -694,28 +918,6 @@ impl FoyerObjectStoreCache { span.record("is_metadata", is_metadata_request); if is_metadata_request { - // For metadata requests, use the metadata cache - let range_cache_key = Self::make_range_cache_key(location, &range); - - // Check if we have this specific range cached in the metadata cache - if let Ok(Some(entry)) = self.metadata_cache.get(&range_cache_key).await { - let value = entry.value(); - let ttl = self.config.ttl; // Use unified TTL - if !value.is_expired(ttl) { - self.update_metadata_stats(|s| s.hits += 1).await; - span.record("cache_hit", true); - debug!( - "Metadata cache HIT for: {} (range: {}..{}, size: {} bytes, age={}ms)", - location, - range.start, - range.end, - value.data.len(), - current_millis().saturating_sub(value.timestamp_millis) - ); - return Ok(Bytes::from(value.data.clone())); - } - } - // Cache miss for metadata range - fetch just the range span.record("cache_hit", false); self.update_metadata_stats(|s| { @@ -830,6 +1032,27 @@ impl FoyerObjectStoreCache { Ok(result) } + /// Resolve a path's `ObjectMeta` from cache only (no S3). Checks the + /// full-file cache, then — for immutable parquet data files — the dedicated + /// meta cache. Returns `None` if neither has a live entry. + async fn cached_meta(&self, location: &Path) -> Option { + if let Ok(Some(entry)) = self.cache.get(&Self::make_cache_key(location)).await { + let value = entry.value(); + if !value.is_expired(self.get_ttl_for_path(location)) { + return Some(value.meta.clone()); + } + } + if is_parquet_file(location) + && let Ok(Some(entry)) = self.metadata_cache.get(&Self::make_meta_cache_key(location)).await + { + let value = entry.value(); + if !value.is_expired(self.config.ttl) { + return Some(value.meta.clone()); + } + } + None + } + #[instrument( name = "foyer_cache.head", skip_all, @@ -840,20 +1063,20 @@ impl FoyerObjectStoreCache { )] async fn head_cached(&self, location: &Path) -> ObjectStoreResult { let span = tracing::Span::current(); - let cache_key = Self::make_cache_key(location); - - if let Ok(Some(entry)) = self.cache.get(&cache_key).await { - let value = entry.value(); - let ttl = self.get_ttl_for_path(location); - if !value.is_expired(ttl) { - span.record("cache_hit", true); - return Ok(value.meta.clone()); - } + if let Some(meta) = self.cached_meta(location).await { + span.record("cache_hit", true); + return Ok(meta); } span.record("cache_hit", false); let inner_span = tracing::trace_span!(parent: &span, "s3.head", location = %location); - self.inner.head(location).instrument(inner_span).await + let meta = self.inner.head(location).instrument(inner_span).await?; + // Cache immutable parquet meta so later footer reads skip the HEAD. Skip + // mutable paths (Delta log / _last_checkpoint can be rewritten in place). + if is_parquet_file(location) { + self.metadata_cache.insert(Self::make_meta_cache_key(location), CacheValue::new(Vec::new(), meta.clone())); + } + Ok(meta) } /// Core put logic: writes to inner store, then caches the new data @@ -862,6 +1085,11 @@ impl FoyerObjectStoreCache { let payload_size = payload.content_length(); let is_parquet = is_parquet_file(location); + // Keep a cheap (Arc-backed) handle to the payload so we can warm the + // cache from the bytes we already hold — no need to re-download what we + // just wrote to S3. + let payload_for_cache = payload.clone(); + debug!("S3 PUT request starting: {} (size: {} bytes, parquet: {})", location, payload_size, is_parquet); let start_time = std::time::Instant::now(); let result = self.inner.put_opts(location, payload, opts).await?; @@ -873,17 +1101,24 @@ impl FoyerObjectStoreCache { is_parquet ); - // After successful write, update the cache with the new data - self.update_stats(|s| s.inner_gets += 1).await; - if let Ok(get_result) = self.inner.get(location).await { - let (data, meta) = Self::collect_payload(get_result).await; - if !data.is_empty() { - let size = meta.size; - self.cache.insert(Self::make_cache_key(location), CacheValue::new(data, meta)); - debug!("Updated cache after write: {} (size: {} bytes)", location, size); + // Warm the cache directly from the just-written bytes — a range-agnostic + // full-file entry, so any subsequent ranged read is served by a slice. + // ObjectMeta is reconstructed from the PutResult (e_tag/version) and the + // known payload size; no post-write GET. insert_main_value applies the + // recent-days window and large-entry disk steering. + if payload_size > 0 { + let mut data = Vec::with_capacity(payload_size); + for chunk in payload_for_cache.iter() { + data.extend_from_slice(chunk); } + let meta = put_result_meta(location.clone(), payload_size as u64, &result); + self.insert_main_value(location, CacheValue::new(data, meta)); + debug!("Warmed cache from write payload: {} (size: {} bytes)", location, payload_size); } + // Overwrites land a fresh full-file entry (checked first on read), but + // stale per-range metadata entries from a previous version of this key + // must still be dropped. if is_parquet { self.invalidate_metadata_cache(location).await; } @@ -897,6 +1132,123 @@ impl FoyerObjectStoreCache { self.invalidate_metadata_cache(location).await; } } + + /// Admit a full-file entry to the main cache, honoring the recent-days + /// window (cold/old partitions are skipped → served from S3) and steering + /// large entries to disk-only so they don't evict the L1 hot set. + fn insert_main_value(&self, location: &Path, value: CacheValue) { + if !is_within_recent_window(location, self.config.cache_recent_days) { + return; + } + insert_main(&self.cache, Self::make_cache_key(location), value, self.config.l1_max_entry_bytes); + } + + /// Sliding-TTL refresh: keep an entry at most `ttl` past its *last query* + /// rather than its insertion. On a hit, once an entry is more than halfway + /// to expiry, re-insert it with a fresh timestamp so frequently-queried + /// data survives indefinitely while cold data still ages out after `ttl`. + /// + /// Throttled (the halfway gate + one in-flight refresh per key) so a hot + /// entry is rewritten at most once per `ttl/2`, and run in the background + /// off a cheap `entry` clone so the read never blocks on the re-insert (the + /// data clone happens in the spawned task, not on the query path). + fn maybe_touch(&self, cache: &FoyerCache, key: &str, entry: foyer::HybridCacheEntry, l1_max_entry_bytes: usize) { + let age = current_millis().saturating_sub(entry.value().timestamp_millis); + // `as_millis()` is u128; clamp before the u64 cast so an absurdly large + // configured TTL can't silently truncate into a tiny value. + let ttl_millis = self.config.ttl.as_millis().min(u64::MAX as u128) as u64; + if age.saturating_mul(2) <= ttl_millis { + return; // still fresh enough — don't churn the cache + } + if !self.refreshing.insert(key.to_string()) { + return; // a refresh is already in flight for this key + } + let cache = cache.clone(); + let refreshing = self.refreshing.clone(); + let key = key.to_string(); + let handle = tokio::spawn(async move { + let v = entry.value(); + insert_main(&cache, key.clone(), CacheValue::new(v.data.clone(), v.meta.clone()), l1_max_entry_bytes); + refreshing.remove(&key); + }); + // Best-effort join registration: if the lock is contended we drop the + // handle and the refresh task simply detaches — it still runs to + // completion, it just won't be awaited by `background_tasks` on + // shutdown. So `background_tasks` is not an exhaustive registry of + // in-flight refreshes; don't assume it joins every one. + if let Ok(mut tasks) = self.background_tasks.try_lock() { + tasks.spawn(async move { + let _ = handle.await; + }); + } + } +} + +/// Wraps an inner [`MultipartUpload`] to tee written bytes into a bounded +/// buffer, so the completed file can be inserted into the cache directly — we +/// never re-download a file we just streamed to S3. If the upload grows past +/// `max_warm_bytes` the buffer is dropped and the rest streams through +/// un-captured (large compaction outputs fall back to the selective +/// post-commit warm path); this bounds both transient memory and L1 cache +/// pressure. Strictly best-effort: failure to capture never affects the write. +struct CachingMultipartUpload { + inner: Box, + location: Path, + cache: FoyerCache, + /// `None` once the cap was exceeded (capture abandoned for this upload). + buffer: Option>, + max_warm_bytes: usize, + l1_max_entry_bytes: usize, +} + +impl std::fmt::Debug for CachingMultipartUpload { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachingMultipartUpload").field("location", &self.location).finish() + } +} + +#[async_trait] +impl MultipartUpload for CachingMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> object_store::UploadPart { + if let Some(buf) = self.buffer.as_mut() { + if buf.len().saturating_add(data.content_length()) > self.max_warm_bytes { + // Too big to warm without risking memory / L1 eviction — give + // up capturing for this upload. + self.buffer = None; + } else { + for chunk in data.iter() { + buf.extend_from_slice(chunk); + } + } + } + self.inner.put_part(data) + } + + async fn complete(&mut self) -> ObjectStoreResult { + let result = self.inner.complete().await?; + if let Some(buf) = self.buffer.take() + && !buf.is_empty() + { + let size = buf.len() as u64; + let meta = put_result_meta(self.location.clone(), size, &result); + // Use the same key derivation as the read path so a multipart-warmed + // entry is found by a later GET even if `make_cache_key` ever does + // more than `location.to_string()`. + insert_main( + &self.cache, + FoyerObjectStoreCache::make_cache_key(&self.location), + CacheValue::new(buf, meta), + self.l1_max_entry_bytes, + ); + debug!("Warmed cache from multipart write: {} (size: {} bytes)", self.location, size); + } + Ok(result) + } + + async fn abort(&mut self) -> ObjectStoreResult<()> { + self.buffer = None; + self.inner.abort().await + } } #[async_trait] @@ -906,7 +1258,29 @@ impl ObjectStore for FoyerObjectStoreCache { } async fn put_multipart_opts(&self, location: &Path, opts: PutMultipartOptions) -> ObjectStoreResult> { - self.inner.put_multipart_opts(location, opts).await + let inner = self.inner.put_multipart_opts(location, opts).await?; + // Skip capture for cold-partition rewrites (e.g. tier recompress of + // week+-old data) — recent data stays local, old data is served from S3. + if !is_within_recent_window(location, self.config.cache_recent_days) { + return Ok(inner); + } + // Parquet writers (flush + compaction outputs) stream large files via + // multipart. Tee the written bytes into a bounded buffer so the + // completed file warms the cache directly — no re-download of what we + // just uploaded. Cap the buffer at the disk block size (the largest + // entry foyer can persist), optionally tightened by warm_inline_max_bytes. + let mut cap = self.config.block_size_bytes; + if self.config.warm_inline_max_bytes > 0 { + cap = cap.min(self.config.warm_inline_max_bytes); + } + Ok(Box::new(CachingMultipartUpload { + inner, + location: location.clone(), + cache: self.cache.clone(), + buffer: Some(Vec::new()), + max_warm_bytes: cap, + l1_max_entry_bytes: self.config.l1_max_entry_bytes, + })) } async fn get_opts(&self, location: &Path, options: GetOptions) -> ObjectStoreResult { @@ -934,6 +1308,71 @@ impl ObjectStore for FoyerObjectStoreCache { range: range.start..range.start + data_len, }); } + // Suffix range (footer warm + any suffix reader): resolve to an absolute + // range so it shares cache keys with bounded footer reads. If we already + // know the size (cached meta), reuse the bounded path for free; otherwise + // a single suffix GET to the inner store learns the absolute range + size + // from the response — one round-trip, no separate HEAD. + if let Some(GetRange::Suffix(n)) = options.range + && options.if_match.is_none() + && options.if_none_match.is_none() + && options.if_modified_since.is_none() + && options.if_unmodified_since.is_none() + { + let n = n.max(1); + if let Some(meta) = self.cached_meta(location).await { + let range = meta.size.saturating_sub(n)..meta.size; + let bytes = self.get_range_cached(location, range.clone()).await?; + let data_len = bytes.len() as u64; + return Ok(GetResult { + payload: GetResultPayload::Stream(Box::pin(futures::stream::once(async move { Ok(bytes) }))), + meta, + attributes: Attributes::new(), + range: range.start..range.start + data_len, + }); + } + let result = self + .inner + .get_opts( + location, + GetOptions { + range: Some(GetRange::Suffix(n)), + ..Default::default() + }, + ) + .await?; + let meta = result.meta.clone(); + let abs_range = result.range.clone(); + let attributes = result.attributes.clone(); + let bytes = result.bytes().await?; + self.update_metadata_stats(|s| { + s.misses += 1; + s.inner_gets += 1; + }) + .await; + // Populate both the footer-range cache (under the absolute key bounded + // reads use) and the immutable-meta cache, so the next footer read is + // a pure cache hit. + if is_parquet_file(location) { + let range_meta = ObjectMeta { + location: location.clone(), + last_modified: meta.last_modified, + size: bytes.len() as u64, + e_tag: meta.e_tag.clone(), + version: meta.version.clone(), + }; + self.metadata_cache + .insert(Self::make_range_cache_key(location, &abs_range), CacheValue::new(bytes.to_vec(), range_meta)); + self.metadata_cache.insert(Self::make_meta_cache_key(location), CacheValue::new(Vec::new(), meta.clone())); + } + let data_len = bytes.len() as u64; + return Ok(GetResult { + payload: GetResultPayload::Stream(Box::pin(futures::stream::once(async move { Ok(bytes) }))), + meta, + attributes, + range: abs_range.start..abs_range.start + data_len, + }); + } // Bypass cache for complex (conditional / non-bounded) requests if options.range.is_some() || options.if_match.is_some() @@ -1004,7 +1443,7 @@ mod tests { let stats = cache.get_stats().await; assert_eq!(stats.main.inner_puts, 1); - assert_eq!(stats.main.inner_gets, 1); // We fetch after write to cache it + assert_eq!(stats.main.inner_gets, 0); // Cached directly from the write payload — no re-fetch // First get - cache hit (since we cache on write) let result = cache.get(&path).await?; @@ -1016,7 +1455,7 @@ mod tests { assert_eq!(bytes[0], data); let stats = cache.get_stats().await; - assert_eq!(stats.main.inner_gets, 1); // No additional fetch needed + assert_eq!(stats.main.inner_gets, 0); // No fetch needed - cached from write payload assert_eq!(stats.main.misses, 0); assert_eq!(stats.main.hits, 1); @@ -1029,7 +1468,7 @@ mod tests { assert_eq!(bytes2[0], data); let stats = cache.get_stats().await; - assert_eq!(stats.main.inner_gets, 1); // Still just the one from write + assert_eq!(stats.main.inner_gets, 0); // Still no fetch - served from cache assert_eq!(stats.main.hits, 2); // Two cache hits total assert_eq!(stats.main.misses, 0); @@ -1083,7 +1522,7 @@ mod tests { } let stats = cache.get_stats().await; - assert_eq!(stats.main.inner_gets, 3); // From the writes + assert_eq!(stats.main.inner_gets, 0); // Cached from write payloads — no re-fetch assert_eq!(stats.main.misses, 0); assert_eq!(stats.main.hits, 3); @@ -1100,7 +1539,7 @@ mod tests { } let stats = cache.get_stats().await; - assert_eq!(stats.main.inner_gets, 3); // No new inner gets + assert_eq!(stats.main.inner_gets, 0); // No inner gets at all assert_eq!(stats.main.hits, 6); // Total 6 hits (3 per read) info!("Cache successfully prevented {} S3 accesses", stats.main.hits); @@ -1169,7 +1608,7 @@ mod tests { assert_eq!(bytes[0].len(), large_data.len()); let stats = cache.get_stats().await; - assert_eq!(stats.main.inner_gets, 1); // From the write + assert_eq!(stats.main.inner_gets, 0); // Cached from write payload — no re-fetch assert_eq!(stats.main.hits, 1); // Second get - cache hit @@ -1181,7 +1620,7 @@ mod tests { assert_eq!(bytes2[0].len(), large_data.len()); let stats = cache.get_stats().await; - assert_eq!(stats.main.inner_gets, 1); // Still just from the write + assert_eq!(stats.main.inner_gets, 0); // Still no fetch - served from cache assert_eq!(stats.main.hits, 2); // Two cache hits total info!("Large file test - main cache hits: {}, misses: {}", stats.main.hits, stats.main.misses); @@ -1406,4 +1845,425 @@ mod tests { let _ = std::fs::remove_dir_all(&cache_dir); Ok(()) } + + #[tokio::test] + async fn test_multipart_capture_warms_cache() -> anyhow::Result<()> { + use object_store::MultipartUpload; + + let test_id = format!("mpu_capture_{}", std::process::id()); + let inner = Arc::new(InMemory::new()); + // Tighten the inline cap to 1MB (below the 4MB block size) so we can + // exercise both the captured and skipped paths. + let config = FoyerCacheConfig::test_config_with(&test_id, |c| { + c.warm_inline_max_bytes = 1024 * 1024; + }); + let cache_dir = config.cache_dir.clone(); + let _ = std::fs::remove_dir_all(&cache_dir); + + let cache = FoyerObjectStoreCache::new(inner.clone(), config).await?; + cache.reset_stats().await; + + // Small multipart write (under the cap) → captured into the cache on + // complete, with no re-download. + let small_path = Path::from("table/date=2026-06-05/small.parquet"); + let small_data = Bytes::from(vec![b'a'; 256 * 1024]); + let mut upload = cache.put_multipart(&small_path).await?; + upload.put_part(small_data.clone().into()).await?; + upload.complete().await?; + + // A read is served entirely from cache — the multipart write warmed it. + let result = cache.get(&small_path).await?; + use futures::TryStreamExt; + let bytes: Vec = match result.payload { + GetResultPayload::Stream(s) => s.try_collect().await?, + _ => panic!("Expected stream"), + }; + assert_eq!(bytes.concat().len(), small_data.len()); + let stats = cache.get_stats().await; + assert_eq!(stats.main.hits, 1, "small multipart write should warm the cache"); + assert_eq!(stats.main.misses, 0, "no S3 read needed after multipart capture"); + + // Large multipart write (over the cap) → capture abandoned, streams + // through, so the first read is a genuine miss. + cache.reset_stats().await; + let big_path = Path::from("table/date=2026-06-05/big.parquet"); + let big_chunk = Bytes::from(vec![b'b'; 768 * 1024]); + let mut upload = cache.put_multipart(&big_path).await?; + upload.put_part(big_chunk.clone().into()).await?; // 768KB + upload.put_part(big_chunk.clone().into()).await?; // 1.5MB total > 1MB cap + upload.complete().await?; + + let _ = cache.get(&big_path).await?; + let stats = cache.get_stats().await; + assert_eq!(stats.main.misses, 1, "over-cap multipart write should not be cached inline"); + + cache.shutdown().await?; + let _ = std::fs::remove_dir_all(&cache_dir); + Ok(()) + } + + #[test] + fn test_block_size_tracks_optimize_target() { + use crate::config::AppConfig; + let mut cfg = AppConfig::default(); + + // Defaults: 2x the 128MB target == the 256MB configured floor. + assert_eq!(FoyerCacheConfig::from_app_config(&cfg).block_size_bytes, 256 * 1024 * 1024); + + // Raise the optimize target past the floor → block auto-tracks to 2x, + // so big outputs stay cacheable without touching the cache config. + cfg.parquet.timefusion_optimize_target_size = 512 * 1024 * 1024; + assert_eq!( + FoyerCacheConfig::from_app_config(&cfg).block_size_bytes, + 1024 * 1024 * 1024, + "block size should track 2x the optimize target" + ); + + // With a small target, the configured block size is the floor. + cfg.parquet.timefusion_optimize_target_size = 16 * 1024 * 1024; + cfg.cache.timefusion_foyer_block_size_mb = 256; + assert_eq!( + FoyerCacheConfig::from_app_config(&cfg).block_size_bytes, + 256 * 1024 * 1024, + "configured block size acts as a floor" + ); + } + + #[tokio::test] + async fn test_sliding_ttl_refresh_on_query() -> anyhow::Result<()> { + let test_id = format!("sliding_ttl_{}", std::process::id()); + let config = FoyerCacheConfig::test_config_with(&test_id, |c| { + c.ttl = Duration::from_millis(1000); + }); + let cache_dir = config.cache_dir.clone(); + let _ = std::fs::remove_dir_all(&cache_dir); + let inner = Arc::new(InMemory::new()); + let cache = FoyerObjectStoreCache::new(inner, config).await?; + + let path = Path::from("table/part-hot.parquet"); + cache.put(&path, PutPayload::from(Bytes::from(vec![b'h'; 4096]))).await?; + + // Query past the halfway point (ttl/2 = 500ms) → triggers a sliding-TTL + // refresh that re-stamps the entry to "now". + tokio::time::sleep(Duration::from_millis(600)).await; + let _ = cache.get(&path).await?; // hit + background touch + tokio::time::sleep(Duration::from_millis(200)).await; // let the re-insert land + + // Now ~1200ms since the original insert (> base TTL) but well within the + // refreshed window — a non-sliding TTL would have expired this entry. + tokio::time::sleep(Duration::from_millis(400)).await; + cache.reset_stats().await; + let _ = cache.get(&path).await?; + let stats = cache.get_stats().await; + assert_eq!(stats.main.hits, 1, "queried entry should survive past base TTL via sliding refresh"); + assert_eq!(stats.main.misses, 0); + + cache.shutdown().await?; + let _ = std::fs::remove_dir_all(&cache_dir); + Ok(()) + } + + #[tokio::test] + async fn test_evict_data_entry_removes_cached_file() -> anyhow::Result<()> { + let inner = Arc::new(InMemory::new()); + let shared = SharedFoyerCache::new(FoyerCacheConfig::test_config("evict_entry")).await?; + let cache = FoyerObjectStoreCache::new_with_shared_cache(inner, &shared); + cache.reset_stats().await; + + let path = Path::from("table/date=2026-06-05/part.parquet"); + cache.put(&path, PutPayload::from(Bytes::from(vec![b'a'; 4096]))).await?; + let _ = cache.get(&path).await?; + assert_eq!(cache.get_stats().await.main.hits, 1, "freshly written file should be cached"); + assert!( + shared.cache.memory().contains(&path.to_string()), + "freshly read file should be in the in-memory cache" + ); + + // Proactive eviction (what the compaction path does for tombstoned + // files) drops the entry from the in-memory cache immediately. foyer's + // HybridCache::remove deletes the on-disk copy asynchronously, so we + // assert on the memory layer for a deterministic result; the dead bytes + // are reclaimed from disk shortly after rather than waiting for VACUUM. + shared.evict_data_entry(path.as_ref()); + assert!( + !shared.cache.memory().contains(&path.to_string()), + "evicted entry should be dropped from the in-memory cache" + ); + + cache.shutdown().await?; + Ok(()) + } + + #[test] + fn test_is_within_recent_window() { + let today = Utc::now().date_naive(); + let recent = Path::from(format!("t/date={}/part.parquet", today)); + let old = Path::from(format!("t/date={}/part.parquet", today - chrono::Duration::days(30))); + + // Recent partitions are admitted; week+-old ones are skipped. + assert!(is_within_recent_window(&recent, 8)); + assert!(!is_within_recent_window(&old, 8)); + // 0 = no age limit → everything admitted. + assert!(is_within_recent_window(&old, 0)); + // No date= segment (Delta log, checkpoints) → always admitted. + assert!(is_within_recent_window(&Path::from("t/_delta_log/00001.json"), 8)); + } + + #[tokio::test] + async fn test_recent_window_skips_old_partition_writes() -> anyhow::Result<()> { + let test_id = format!("recent_window_{}", std::process::id()); + let inner = Arc::new(InMemory::new()); + let config = FoyerCacheConfig::test_config_with(&test_id, |c| { + c.cache_recent_days = 8; // enforce the window in this test + }); + let cache_dir = config.cache_dir.clone(); + let _ = std::fs::remove_dir_all(&cache_dir); + + let cache = FoyerObjectStoreCache::new(inner.clone(), config).await?; + cache.reset_stats().await; + + let today = Utc::now().date_naive(); + let recent = Path::from(format!("t/date={}/part.parquet", today)); + let old = Path::from(format!("t/date={}/part.parquet", today - chrono::Duration::days(30))); + let data = Bytes::from(vec![b'a'; 4096]); + + // Recent write is admitted → served from cache (no S3 read). + cache.put(&recent, PutPayload::from(data.clone())).await?; + let _ = cache.get(&recent).await?; + assert_eq!(cache.get_stats().await.main.hits, 1, "recent write should be cached"); + + // Old-partition write is NOT admitted → read falls through to S3 (miss). + cache.reset_stats().await; + cache.put(&old, PutPayload::from(data.clone())).await?; + let _ = cache.get(&old).await?; + let stats = cache.get_stats().await; + assert_eq!(stats.main.hits, 0, "old-partition write should not be cached"); + assert_eq!(stats.main.misses, 1, "old partition served from S3"); + + cache.shutdown().await?; + let _ = std::fs::remove_dir_all(&cache_dir); + Ok(()) + } + + #[tokio::test] + async fn test_warm_footer_primes_metadata_cache() -> anyhow::Result<()> { + let test_id = format!("warm_footer_{}", std::process::id()); + let inner = Arc::new(InMemory::new()); + let config = FoyerCacheConfig::test_config_with(&test_id, |c| { + c.parquet_metadata_size_hint = 1024; // 1KB footer + }); + let cache_dir = config.cache_dir.clone(); + let _ = std::fs::remove_dir_all(&cache_dir); + + let cache = FoyerObjectStoreCache::new(inner.clone(), config).await?; + + // Write the file straight to the inner store so nothing is cached yet — + // this simulates a multipart compaction output that bypassed put_cached. + let file_size = 10 * 1024; + let path = Path::from("table/date=2026-06-05/part-0.parquet"); + inner.put(&path, PutPayload::from(Bytes::from(vec![b'x'; file_size]))).await?; + cache.reset_stats().await; + + // Warm the footer. The ranged GET should populate the metadata cache. + assert!(warm_footer(&cache, &path, 1024).await); + + let stats = cache.get_stats().await; + assert_eq!(stats.metadata.misses, 1, "warm should fetch the footer once"); + assert_eq!(stats.metadata.hits, 0); + + // A subsequent read of the same footer range is now a metadata HIT — + // i.e. query planning pays zero S3 round-trips post-warm. + let metadata_range = (file_size - 1024) as u64..file_size as u64; + let _ = cache.get_range(&path, metadata_range).await?; + let stats = cache.get_stats().await; + assert_eq!(stats.metadata.hits, 1, "footer read should hit after warm"); + + cache.shutdown().await?; + let _ = std::fs::remove_dir_all(&cache_dir); + Ok(()) + } + + #[tokio::test] + async fn test_warm_full_primes_main_cache() -> anyhow::Result<()> { + let test_id = format!("warm_full_{}", std::process::id()); + let inner = Arc::new(InMemory::new()); + let config = FoyerCacheConfig::test_config(&test_id); + let cache_dir = config.cache_dir.clone(); + let _ = std::fs::remove_dir_all(&cache_dir); + + let cache = FoyerObjectStoreCache::new(inner.clone(), config).await?; + + let file_size = 8 * 1024; + let path = Path::from("table/date=2026-06-05/part-1.parquet"); + inner.put(&path, PutPayload::from(Bytes::from(vec![b'y'; file_size]))).await?; + cache.reset_stats().await; + + // Warm the full file into the main cache. The warm itself incurs one + // miss (the fetch from the inner store); reset so we isolate the + // post-warm read behavior. + assert!(warm_full(&cache, &path).await); + let stats = cache.get_stats().await; + assert_eq!(stats.main.misses, 1, "warm should fetch the full file once"); + cache.reset_stats().await; + + // Any subsequent data read is served from the full-file cache (main HIT), + // never falling back to S3. + let _ = cache.get_range(&path, 0..1024).await?; + let stats = cache.get_stats().await; + assert_eq!(stats.main.hits, 1, "data read should hit main cache after full warm"); + assert_eq!(stats.main.misses, 0); + + cache.shutdown().await?; + let _ = std::fs::remove_dir_all(&cache_dir); + Ok(()) + } + + /// Guards key consistency across the three cache paths that derive a key + /// independently: the multipart-write warm (`complete()`), the read path + /// (`make_cache_key`), and the compaction eviction path (`evict_data_entry` + /// on a relativized object path). If any of them diverged, a multipart-warmed + /// entry would either never be read back or never be evicted. + #[tokio::test] + async fn test_multipart_warm_read_and_evict_key_consistency() -> anyhow::Result<()> { + use object_store::MultipartUpload; + + let inner = Arc::new(InMemory::new()); + let shared = SharedFoyerCache::new(FoyerCacheConfig::test_config("mpu_key_consistency")).await?; + let cache = FoyerObjectStoreCache::new_with_shared_cache(inner, &shared); + cache.reset_stats().await; + + // Warm via the multipart-write path. + let path = Path::from("table/date=2026-06-05/part.parquet"); + let data = Bytes::from(vec![b'z'; 64 * 1024]); + let mut upload = cache.put_multipart(&path).await?; + upload.put_part(data.clone().into()).await?; + upload.complete().await?; + + // Read path: a plain GET must find the entry the multipart write warmed — + // i.e. `complete()` inserted under the same key the read derives. A key + // mismatch would surface here as a miss + an S3 fetch. + let _ = cache.get(&path).await?; + let stats = cache.get_stats().await; + assert_eq!(stats.main.hits, 1, "multipart-warmed entry must be found by a plain GET (warm/read key match)"); + assert_eq!(stats.main.misses, 0, "no S3 read needed after multipart capture"); + + // Eviction path: the compaction hook evicts by the relativized object + // path. It must target the same key warming/reads use, or tombstoned + // files would linger. Assert at the in-memory layer — foyer removes the + // on-disk copy asynchronously, so the memory layer is the deterministic + // signal (mirrors test_evict_data_entry_removes_cached_file). + assert!( + shared.cache.memory().contains(&path.to_string()), + "warmed entry should be in the in-memory cache" + ); + shared.evict_data_entry(path.as_ref()); + assert!( + !shared.cache.memory().contains(&path.to_string()), + "evict must drop the same key warming/reads use" + ); + + cache.shutdown().await?; + Ok(()) + } + + /// Wraps an `InMemory` store and counts S3-equivalent round-trips, so tests + /// can assert that warming + reads issue the expected number of HEADs/GETs. + /// `head()` is an extension method that routes through `get_opts(head:true)`, + /// so we count it there. + #[derive(Debug)] + struct CountingStore { + inner: Arc, + heads: Arc, + gets: Arc, + } + + impl std::fmt::Display for CountingStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "CountingStore") + } + } + + #[async_trait] + impl ObjectStore for CountingStore { + async fn put_opts(&self, location: &Path, payload: PutPayload, opts: PutOptions) -> ObjectStoreResult { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts(&self, location: &Path, opts: PutMultipartOptions) -> ObjectStoreResult> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> ObjectStoreResult { + use std::sync::atomic::Ordering; + if options.head { + self.heads.fetch_add(1, Ordering::Relaxed); + } else { + self.gets.fetch_add(1, Ordering::Relaxed); + } + self.inner.get_opts(location, options).await + } + + fn delete_stream(&self, locations: BoxStream<'static, ObjectStoreResult>) -> BoxStream<'static, ObjectStoreResult> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, ObjectStoreResult> { + self.inner.list(prefix) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> ObjectStoreResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> ObjectStoreResult<()> { + self.inner.copy_opts(from, to, options).await + } + } + + /// Locks in both performance wins: a suffix-based footer warm is a single GET + /// (no HEAD), and a later footer read of a warmed file is a pure cache hit — + /// zero S3 round-trips (no HEAD to classify, no GET). + #[tokio::test] + async fn test_warm_footer_eliminates_read_path_heads() -> anyhow::Result<()> { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let mem = Arc::new(InMemory::new()); + let file_size = 10 * 1024usize; + let path = Path::from("table/date=2026-06-05/part-heads.parquet"); + mem.put(&path, PutPayload::from(Bytes::from(vec![b'x'; file_size]))).await?; + + let heads = Arc::new(AtomicUsize::new(0)); + let gets = Arc::new(AtomicUsize::new(0)); + let counting = Arc::new(CountingStore { + inner: mem.clone(), + heads: heads.clone(), + gets: gets.clone(), + }); + + let config = FoyerCacheConfig::test_config_with("warm_footer_heads", |c| { + c.parquet_metadata_size_hint = 1024; + }); + let cache_dir = config.cache_dir.clone(); + let _ = std::fs::remove_dir_all(&cache_dir); + let cache = FoyerObjectStoreCache::new(counting, config).await?; + cache.reset_stats().await; + + // Footer warm: a single suffix GET, no HEAD. + assert!(warm_footer(&cache, &path, 1024).await); + assert_eq!(heads.load(Ordering::Relaxed), 0, "suffix warm must not issue a HEAD"); + assert_eq!(gets.load(Ordering::Relaxed), 1, "suffix warm is a single GET"); + + // A later footer read of the warmed file is served entirely from cache: + // no HEAD to classify the range, no GET for the bytes. + let footer = (file_size - 1024) as u64..file_size as u64; + let bytes = cache.get_range(&path, footer).await?; + assert_eq!(bytes.len(), 1024); + assert_eq!(heads.load(Ordering::Relaxed), 0, "warmed footer read must not HEAD"); + assert_eq!(gets.load(Ordering::Relaxed), 1, "warmed footer read must not GET (still just the warm)"); + assert_eq!(cache.get_stats().await.metadata.hits, 1, "footer served from the metadata cache"); + + cache.shutdown().await?; + let _ = std::fs::remove_dir_all(&cache_dir); + Ok(()) + } } diff --git a/tests/cache_performance_test.rs b/tests/cache_performance_test.rs index 32f57721..a31bc426 100644 --- a/tests/cache_performance_test.rs +++ b/tests/cache_performance_test.rs @@ -45,10 +45,9 @@ async fn test_cache_performance_and_s3_bypass() -> Result<()> { // Get baseline stats after writes let stats_after_write = shared_cache.get_stats().await; assert_eq!(stats_after_write.main.inner_puts, 3, "Should have written to inner store 3 times"); - assert_eq!( - stats_after_write.main.inner_gets, 3, - "Should have fetched from inner store 3 times during write" - ); + // Writes warm the cache directly from the put payload (no post-write + // re-fetch), so no inner GETs are issued during writes. + assert_eq!(stats_after_write.main.inner_gets, 0, "Writes warm from payload — no inner GET during write"); // First read - should hit cache since we cache on write let start = Instant::now(); @@ -81,7 +80,7 @@ async fn test_cache_performance_and_s3_bypass() -> Result<()> { let stats = shared_cache.get_stats().await; assert_eq!(stats.main.hits, 6, "Should have 6 cache hits total (3 per read iteration)"); assert_eq!(stats.main.misses, 0, "Should have no cache misses since files were cached on write"); - assert_eq!(stats.main.inner_gets, 3, "Should have fetched from inner store 3 times during write"); + assert_eq!(stats.main.inner_gets, 0, "Writes warm from payload and reads hit cache — no inner GETs at all"); assert_eq!(stats.main.inner_puts, 3, "Should have written to inner store 3 times"); // Test cache invalidation on write @@ -186,17 +185,18 @@ async fn test_parquet_metadata_cache_performance() -> Result<()> { // Configure cache with metadata optimization let config = FoyerCacheConfig { - memory_size_bytes: 50 * 1024 * 1024, // 50MB - disk_size_bytes: 100 * 1024 * 1024, // 100MB - ttl: std::time::Duration::from_secs(300), - cache_dir: std::path::PathBuf::from("/tmp/test_parquet_metadata_perf"), - shards: 4, - file_size_bytes: 4 * 1024 * 1024, // 4MB - enable_stats: true, + memory_size_bytes: 50 * 1024 * 1024, // 50MB + disk_size_bytes: 100 * 1024 * 1024, // 100MB + ttl: std::time::Duration::from_secs(300), + cache_dir: std::path::PathBuf::from("/tmp/test_parquet_metadata_perf"), + shards: 4, + file_size_bytes: 4 * 1024 * 1024, // 4MB + enable_stats: true, parquet_metadata_size_hint: 1_048_576, // 1MB metadata_memory_size_bytes: 20 * 1024 * 1024, // 20MB - metadata_disk_size_bytes: 50 * 1024 * 1024, // 50MB - metadata_shards: 2, + metadata_disk_size_bytes: 50 * 1024 * 1024, // 50MB + metadata_shards: 2, + ..Default::default() }; // Clean up cache directory