diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index 2b324c975a4..79a70cb7db0 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -9469,6 +9469,7 @@ dependencies = [ "serde_json", "tantivy", "tantivy-fst", + "tempfile", "thiserror 2.0.18", "tokio", "tokio-util", @@ -9578,6 +9579,7 @@ dependencies = [ "base64 0.22.1", "bytes", "bytesize", + "fail", "foyer", "futures", "http 1.4.2", diff --git a/quickwit/quickwit-lambda-server/src/context.rs b/quickwit/quickwit-lambda-server/src/context.rs index d3b9167414f..e8faad760a5 100644 --- a/quickwit/quickwit-lambda-server/src/context.rs +++ b/quickwit/quickwit-lambda-server/src/context.rs @@ -33,8 +33,11 @@ impl LambdaSearcherContext { info!("initializing lambda searcher context"); let searcher_config = try_searcher_config_from_env()?; - let searcher_context = - Arc::new(SearcherContext::new_without_invoker(searcher_config, None)); + let searcher_context = Arc::new(SearcherContext::new_without_invoker( + searcher_config, + None, + None, + )); let storage_resolver = StorageResolver::configured(&Default::default()); Ok(Self { diff --git a/quickwit/quickwit-search/Cargo.toml b/quickwit/quickwit-search/Cargo.toml index cbbe2b269d5..4c95fc62214 100644 --- a/quickwit/quickwit-search/Cargo.toml +++ b/quickwit/quickwit-search/Cargo.toml @@ -53,6 +53,7 @@ assert-json-diff = { workspace = true } proptest = { workspace = true } rand = { workspace = true } serde_json = { workspace = true } +tempfile = { workspace = true } quickwit-indexing = { workspace = true, features = ["testsuite"] } quickwit-metastore = { workspace = true, features = ["testsuite"] } diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index 99ff06c312b..36ae3986f9c 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -46,6 +46,7 @@ use quickwit_query::tokenizers::TokenizerManager; use quickwit_storage::{ BundleStorage, ByteRangeCache, CountingStorage, MemorySizedCache, OwnedBytes, SearchSplitCache, Storage, StorageResolver, TimeoutAndRetryStorage, wrap_storage_with_cache, + wrap_storage_with_split_range_cache, }; use tantivy::aggregation::AggContextParams; use tantivy::aggregation::agg_req::{AggregationVariants, Aggregations}; @@ -160,16 +161,19 @@ async fn get_split_footer_from_cache_or_fetch( Ok(footer_data_opt) } -/// Returns hotcache_bytes and the split directory (`BundleStorage`) with cache layer: -/// - A split footer cache given by `SearcherContext.split_footer_cache`. +/// Returns hotcache_bytes and the split directory (`BundleStorage`). pub(crate) async fn open_split_bundle( searcher_context: &SearcherContext, index_storage: Arc, split_and_footer_offsets: &SplitIdAndFooterOffsets, ) -> anyhow::Result<(FileSlice, BundleStorage)> { let split_file = PathBuf::from(format!("{}.split", split_and_footer_offsets.split_id)); + let foyer_storage: Arc = match &searcher_context.split_range_disk_cache_opt { + Some(cache) => wrap_storage_with_split_range_cache(cache.clone(), index_storage.clone()), + None => index_storage.clone(), + }; let footer_data = get_split_footer_from_cache_or_fetch( - index_storage.clone(), + foyer_storage.clone(), split_and_footer_offsets, &searcher_context.split_footer_cache, ) @@ -179,9 +183,9 @@ pub(crate) async fn open_split_bundle( // This is before the bundle storage: at this point, this storage is reading `.split` files. let index_storage_with_split_cache = if let Some(split_cache) = searcher_context.split_cache_opt.as_ref() { - SearchSplitCache::wrap_storage(split_cache.clone(), index_storage.clone()) + SearchSplitCache::wrap_storage(split_cache.clone(), foyer_storage) } else { - index_storage.clone() + foyer_storage }; let (hotcache_bytes, bundle_storage) = BundleStorage::open_from_split_data( @@ -3113,7 +3117,8 @@ mod tests { offload_threshold: 3, ..LambdaConfig::for_test() }); - let searcher_context = SearcherContext::new(config, None, Some(Arc::new(DummyInvoker))); + let searcher_context = + SearcherContext::new(config, None, None, Some(Arc::new(DummyInvoker))); let splits = make_splits_with_requests(7); let result = super::schedule_search_tasks(splits, &searcher_context).await; assert_eq!(result.local_search_tasks.len(), 3); @@ -3133,7 +3138,8 @@ mod tests { offload_threshold: 0, ..LambdaConfig::for_test() }); - let searcher_context = SearcherContext::new(config, None, Some(Arc::new(DummyInvoker))); + let searcher_context = + SearcherContext::new(config, None, None, Some(Arc::new(DummyInvoker))); let splits = make_splits_with_requests(5); let result = super::schedule_search_tasks(splits, &searcher_context).await; assert!(result.local_search_tasks.is_empty()); @@ -3147,7 +3153,8 @@ mod tests { offload_threshold: 100, ..LambdaConfig::for_test() }); - let searcher_context = SearcherContext::new(config, None, Some(Arc::new(DummyInvoker))); + let searcher_context = + SearcherContext::new(config, None, None, Some(Arc::new(DummyInvoker))); let splits = make_splits_with_requests(5); let result = super::schedule_search_tasks(splits, &searcher_context).await; assert_eq!(result.local_search_tasks.len(), 5); diff --git a/quickwit/quickwit-search/src/lib.rs b/quickwit/quickwit-search/src/lib.rs index 2d891dbfa65..9d15c7f6c12 100644 --- a/quickwit/quickwit-search/src/lib.rs +++ b/quickwit/quickwit-search/src/lib.rs @@ -42,6 +42,8 @@ pub(crate) mod top_k_collector; mod metrics; mod search_permit_provider; +#[cfg(test)] +mod split_range_cache_layer_tests; #[cfg(test)] mod tests; @@ -289,7 +291,11 @@ pub async fn single_node_search( let search_job_placer = SearchJobPlacer::new(searcher_pool.clone()); let cluster_client = ClusterClient::new(search_job_placer); let searcher_config = SearcherConfig::default(); - let searcher_context = Arc::new(SearcherContext::new_without_invoker(searcher_config, None)); + let searcher_context = Arc::new(SearcherContext::new_without_invoker( + searcher_config, + None, + None, + )); let search_service = Arc::new(SearchServiceImpl::new( metastore.clone(), storage_resolver, diff --git a/quickwit/quickwit-search/src/service.rs b/quickwit/quickwit-search/src/service.rs index 52bfc846696..518b330583a 100644 --- a/quickwit/quickwit-search/src/service.rs +++ b/quickwit/quickwit-search/src/service.rs @@ -29,7 +29,8 @@ use quickwit_proto::search::{ SearchPlanResponse, SearchRequest, SearchResponse, SnippetRequest, }; use quickwit_storage::{ - MemorySizedCache, QuickwitCache, SearchSplitCache, StorageCache, StorageResolver, + FoyerSplitRangeCache, MemorySizedCache, QuickwitCache, SearchSplitCache, StorageCache, + StorageResolver, }; use tantivy::aggregation::AggregationLimitsGuard; @@ -417,6 +418,8 @@ pub struct SearcherContext { pub predicate_cache: Arc, /// Search split cache. `None` if no split cache is configured. pub split_cache_opt: Option>, + /// Process-wide split range disk cache. `None` if not configured. + pub split_range_disk_cache_opt: Option>, /// List fields cache. Caches the raw fields-metadata blob for a given split. pub list_fields_cache: ListFieldsCache, /// The aggregation limits are passed to limit the memory usage. @@ -439,17 +442,19 @@ impl SearcherContext { #[cfg(test)] pub fn for_test() -> SearcherContext { let searcher_config = SearcherConfig::default(); - SearcherContext::new_without_invoker(searcher_config, None) + SearcherContext::new_without_invoker(searcher_config, None, None) } /// Creates a new searcher context without a lambda invoker. pub fn new_without_invoker( searcher_config: SearcherConfig, split_cache_opt: Option>, + split_range_disk_cache_opt: Option>, ) -> Self { Self::new( searcher_config, split_cache_opt, + split_range_disk_cache_opt, None::>, ) } @@ -458,6 +463,7 @@ impl SearcherContext { pub fn new( searcher_config: SearcherConfig, split_cache_opt: Option>, + split_range_disk_cache_opt: Option>, lambda_invoker: Option, ) -> Self { let global_split_footer_cache = MemorySizedCache::from_config( @@ -490,6 +496,7 @@ impl SearcherContext { leaf_search_cache, list_fields_cache, split_cache_opt, + split_range_disk_cache_opt, aggregation_limit, lambda_invoker, } diff --git a/quickwit/quickwit-search/src/split_range_cache_layer_tests.rs b/quickwit/quickwit-search/src/split_range_cache_layer_tests.rs new file mode 100644 index 00000000000..df123855385 --- /dev/null +++ b/quickwit/quickwit-search/src/split_range_cache_layer_tests.rs @@ -0,0 +1,225 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed 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 std::num::NonZeroU32; +use std::path::Path; +use std::sync::Arc; + +use bytesize::ByteSize; +use quickwit_config::{ + CachePolicy, DiskCompression, RecoverMode, SearcherConfig, SplitCacheLimits, + SplitRangeCacheWritePolicy, SplitRangeDiskCacheConfig, +}; +use quickwit_proto::search::SplitIdAndFooterOffsets; +use quickwit_storage::{ + CountingStorage, DownloadCounters, FoyerSplitRangeCache, OwnedBytes, PutPayload, + RamStorageBuilder, SearchSplitCache, SplitPayloadBuilder, Storage, StorageResolver, +}; + +use crate::SearcherContext; +use crate::leaf::open_split_bundle; + +const SPLIT_ID: &str = "range-cache-split"; +const BODY_FILE: &str = "segment.fast"; +const BODY_BYTES: &[u8] = b"FASTDATA"; +const HOTCACHE_BYTES: &[u8] = b"HOT"; + +fn range_cache_config(path: impl AsRef) -> SplitRangeDiskCacheConfig { + SplitRangeDiskCacheConfig { + path: path.as_ref().to_path_buf(), + disk_capacity: ByteSize::mb(64), + memory_capacity: ByteSize::mb(8), + buffer_pool_size: ByteSize::mb(4), + submit_queue_size_threshold: ByteSize::mb(8), + memory_eviction_policy: CachePolicy::S3Fifo, + write_policy: SplitRangeCacheWritePolicy::WriteOnEviction, + compression: DiskCompression::Lz4, + recover_mode: RecoverMode::Quiet, + block_size: ByteSize::mb(4), + max_entry_size: ByteSize::mb(2), + flushers: 1, + reclaimers: 1, + } +} + +fn lower_reads(counters: &DownloadCounters) -> u64 { + counters.snapshot().1 +} + +struct SplitBundle { + split_bytes: OwnedBytes, + footer_offsets: SplitIdAndFooterOffsets, +} + +async fn build_split() -> SplitBundle { + let temp_dir = tempfile::tempdir().unwrap(); + let body_path = temp_dir.path().join(BODY_FILE); + std::fs::write(&body_path, BODY_BYTES).unwrap(); + let payload = + SplitPayloadBuilder::get_split_payload(&[body_path], &[], HOTCACHE_BYTES).unwrap(); + let footer_range = payload.footer_range.clone(); + let split_bytes = payload.read_all().await.unwrap(); + SplitBundle { + split_bytes, + footer_offsets: SplitIdAndFooterOffsets { + split_id: SPLIT_ID.to_string(), + split_footer_start: footer_range.start, + split_footer_end: footer_range.end, + ..Default::default() + }, + } +} + +fn split_file_name() -> String { + format!("{SPLIT_ID}.split") +} + +async fn open_range_cache(dir: &Path) -> Arc { + Arc::new( + FoyerSplitRangeCache::open(&range_cache_config(dir)) + .await + .unwrap(), + ) +} + +fn wrap_counted_ram(split_bytes: &OwnedBytes) -> (Arc, Arc) { + let ram = RamStorageBuilder::default() + .put(&split_file_name(), split_bytes.as_slice()) + .build(); + CountingStorage::instrument_storage(Arc::new(ram)) +} + +fn context_with_range_cache(cache: Arc) -> SearcherContext { + SearcherContext::new_without_invoker(SearcherConfig::default(), None, Some(cache)) +} + +#[tokio::test] +async fn test_open_split_bundle_footer_ram_hit_bypasses_lower_tiers() { + let split = build_split().await; + let cache_dir = tempfile::tempdir().unwrap(); + let cache = open_range_cache(cache_dir.path()).await; + let context = context_with_range_cache(cache.clone()); + let footer = split.split_bytes.slice( + split.footer_offsets.split_footer_start as usize + ..split.footer_offsets.split_footer_end as usize, + ); + context.split_footer_cache.put(SPLIT_ID.to_string(), footer); + + let (counted, counters) = wrap_counted_ram(&split.split_bytes); + let (_hotcache, _bundle) = open_split_bundle(&context, counted, &split.footer_offsets) + .await + .unwrap(); + assert_eq!(lower_reads(&counters), 0); + cache.close().await.unwrap(); +} + +#[tokio::test] +async fn test_open_split_bundle_footer_miss_uses_foyer_then_reuses_storage_for_body() { + let split = build_split().await; + let cache_dir = tempfile::tempdir().unwrap(); + let cache = open_range_cache(cache_dir.path()).await; + let context = context_with_range_cache(cache.clone()); + let (counted, counters) = wrap_counted_ram(&split.split_bytes); + + let (_hotcache, bundle) = open_split_bundle(&context, counted, &split.footer_offsets) + .await + .unwrap(); + assert_eq!(lower_reads(&counters), 1, "cold footer is one lower read"); + + bundle.get_slice(Path::new(BODY_FILE), 0..4).await.unwrap(); + bundle.get_slice(Path::new(BODY_FILE), 0..4).await.unwrap(); + assert_eq!( + lower_reads(&counters), + 2, + "second exact body range must hit Foyer" + ); + cache.close().await.unwrap(); +} + +#[tokio::test] +async fn test_open_split_bundle_footer_skips_whole_split_cache() { + let split = build_split().await; + let split_cache_dir = tempfile::tempdir().unwrap(); + std::fs::write( + split_cache_dir.path().join(split_file_name()), + split.split_bytes.as_slice(), + ) + .unwrap(); + let split_cache = SearchSplitCache::with_root_path( + split_cache_dir.path().to_path_buf(), + StorageResolver::unconfigured(), + SplitCacheLimits { + max_num_bytes: ByteSize::mb(64), + max_num_splits: NonZeroU32::new(8).unwrap(), + num_concurrent_downloads: NonZeroU32::new(1).unwrap(), + max_file_descriptors: NonZeroU32::new(8).unwrap(), + }, + ) + .unwrap(); + + let range_cache_dir = tempfile::tempdir().unwrap(); + let range_cache = open_range_cache(range_cache_dir.path()).await; + let context = SearcherContext::new_without_invoker( + SearcherConfig::default(), + Some(split_cache), + Some(range_cache.clone()), + ); + let (counted, counters) = wrap_counted_ram(&split.split_bytes); + let (_hotcache, bundle) = open_split_bundle(&context, counted, &split.footer_offsets) + .await + .unwrap(); + assert_eq!( + lower_reads(&counters), + 1, + "footer fetch skips SplitCache and reads through Foyer" + ); + + bundle.get_slice(Path::new(BODY_FILE), 0..4).await.unwrap(); + assert_eq!( + lower_reads(&counters), + 1, + "body read must hit the on-disk whole-split cache" + ); + range_cache.close().await.unwrap(); +} + +#[tokio::test] +async fn test_open_split_bundle_recovers_footer_from_foyer() { + let split = build_split().await; + let cache_dir = tempfile::tempdir().unwrap(); + let config = range_cache_config(cache_dir.path()); + { + let cache = Arc::new(FoyerSplitRangeCache::open(&config).await.unwrap()); + let context = context_with_range_cache(cache.clone()); + let (counted, counters) = wrap_counted_ram(&split.split_bytes); + open_split_bundle(&context, counted, &split.footer_offsets) + .await + .unwrap(); + assert_eq!(lower_reads(&counters), 1); + cache.close().await.unwrap(); + } + + let recovered = Arc::new(FoyerSplitRangeCache::open(&config).await.unwrap()); + let context = context_with_range_cache(recovered.clone()); + let (counted, counters) = wrap_counted_ram(&split.split_bytes); + open_split_bundle(&context, counted, &split.footer_offsets) + .await + .unwrap(); + assert_eq!( + lower_reads(&counters), + 0, + "recovered footer range must not read lower storage" + ); + recovered.close().await.unwrap(); +} diff --git a/quickwit/quickwit-search/src/tests.rs b/quickwit/quickwit-search/src/tests.rs index 94183d16c2a..63c6c3ed08a 100644 --- a/quickwit/quickwit-search/src/tests.rs +++ b/quickwit/quickwit-search/src/tests.rs @@ -1031,6 +1031,7 @@ async fn test_search_util(test_sandbox: &TestSandbox, query: &str) -> Vec { let searcher_context: Arc = Arc::new(SearcherContext::new_without_invoker( SearcherConfig::default(), None, + None, )); let search_response = single_doc_mapping_leaf_search( @@ -1671,6 +1672,7 @@ async fn test_single_node_list_terms() -> anyhow::Result<()> { let searcher_context = Arc::new(SearcherContext::new_without_invoker( SearcherConfig::default(), None, + None, )); { diff --git a/quickwit/quickwit-serve/src/lib.rs b/quickwit/quickwit-serve/src/lib.rs index dd45d678d95..951ea64e111 100644 --- a/quickwit/quickwit-serve/src/lib.rs +++ b/quickwit/quickwit-serve/src/lib.rs @@ -119,7 +119,7 @@ use quickwit_search::{ SearchJobPlacer, SearchService, SearchServiceClient, SearcherContext, SearcherPool, create_search_client_from_channel, start_searcher_service, }; -use quickwit_storage::{SearchSplitCache, StorageResolver}; +use quickwit_storage::{FoyerSplitRangeCache, SearchSplitCache, StorageResolver}; pub use quickwit_telemetry_exporters::{EnvFilterReloadFn, do_nothing_env_filter_reload_fn}; pub use quickwit_transport::reload_tls_cert; use tcp_listener::TcpListenerResolver; @@ -729,6 +729,20 @@ pub async fn serve_quickwit( None }; + let split_range_disk_cache_opt = if node_config.is_service_enabled(QuickwitService::Searcher) { + match &node_config.searcher_config.split_range_disk_cache { + Some(config) => Some(Arc::new( + FoyerSplitRangeCache::open(config) + .await + .context("failed to open searcher split range disk cache")?, + )), + None => None, + } + } else { + None + }; + let split_range_disk_cache_for_shutdown = split_range_disk_cache_opt.clone(); + // Initialize Lambda invoker if enabled and searcher service is running let searcher_context = if node_config.is_service_enabled(QuickwitService::Searcher) { if let Some(lambda_config) = &node_config.searcher_config.lambda { @@ -741,6 +755,7 @@ pub async fn serve_quickwit( Arc::new(SearcherContext::new( node_config.searcher_config.clone(), search_split_cache_opt, + split_range_disk_cache_opt, Some(invoker), )) } @@ -753,12 +768,14 @@ pub async fn serve_quickwit( Arc::new(SearcherContext::new_without_invoker( node_config.searcher_config.clone(), search_split_cache_opt, + split_range_disk_cache_opt, )) } } else { Arc::new(SearcherContext::new_without_invoker( node_config.searcher_config.clone(), search_split_cache_opt, + split_range_disk_cache_opt, )) }; @@ -1077,6 +1094,12 @@ pub async fn serve_quickwit( let actor_exit_statuses = shutdown_handle .await .context("failed to gracefully shutdown services")?; + if let Some(cache) = split_range_disk_cache_for_shutdown { + cache + .close() + .await + .context("failed to close searcher split range disk cache")?; + } Ok(actor_exit_statuses) } @@ -2086,6 +2109,7 @@ mod tests { let searcher_context = Arc::new(SearcherContext::new_without_invoker( SearcherConfig::default(), None, + None, )); let metastore = metastore_for_test(); let (change_stream, change_stream_tx) = ClusterChangeStream::new_unbounded(); diff --git a/quickwit/quickwit-storage/Cargo.toml b/quickwit/quickwit-storage/Cargo.toml index 6866f84be7b..c6b640da215 100644 --- a/quickwit/quickwit-storage/Cargo.toml +++ b/quickwit/quickwit-storage/Cargo.toml @@ -37,6 +37,7 @@ stable_deref_trait = { workspace = true } tantivy = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } +fail = { workspace = true } tokio = { workspace = true, features = ["test-util"] } tokio-stream = { workspace = true } tokio-util = { workspace = true } @@ -97,6 +98,7 @@ azure = [ ] gcs = ["dep:opendal", "opendal/services-gcs"] ci-test = [] +failpoints = ["fail/failpoints"] integration-testsuite = [ "azure", "azure_core/azurite_workaround", diff --git a/quickwit/quickwit-storage/src/lib.rs b/quickwit/quickwit-storage/src/lib.rs index 0aa79c1d21a..894a18711fe 100644 --- a/quickwit/quickwit-storage/src/lib.rs +++ b/quickwit/quickwit-storage/src/lib.rs @@ -58,7 +58,9 @@ mod versioned_component; use quickwit_common::uri::Uri; pub use split_cache::SearchSplitCache; -pub use split_range_cache::FoyerSplitRangeCache; +pub use split_range_cache::{ + FoyerSplitRangeCache, FoyerSplitRangeStorage, wrap_storage_with_split_range_cache, +}; pub use tantivy::directory::OwnedBytes; pub use versioned_component::VersionedComponent; diff --git a/quickwit/quickwit-storage/src/split_range_cache/key.rs b/quickwit/quickwit-storage/src/split_range_cache/key.rs index 0aea374f33b..1bac1c61bfe 100644 --- a/quickwit/quickwit-storage/src/split_range_cache/key.rs +++ b/quickwit/quickwit-storage/src/split_range_cache/key.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#![allow(dead_code)] - use std::io::{Read, Write}; use std::ops::Range; diff --git a/quickwit/quickwit-storage/src/split_range_cache/mod.rs b/quickwit/quickwit-storage/src/split_range_cache/mod.rs index 18265ca9c93..c122f186c5e 100644 --- a/quickwit/quickwit-storage/src/split_range_cache/mod.rs +++ b/quickwit/quickwit-storage/src/split_range_cache/mod.rs @@ -13,6 +13,7 @@ // limitations under the License. mod key; +mod storage; #[cfg(test)] mod tests; @@ -28,13 +29,12 @@ use quickwit_config::{ CachePolicy, DiskCompression, RecoverMode, SplitRangeCacheWritePolicy, SplitRangeDiskCacheConfig, }; +pub use storage::{FoyerSplitRangeStorage, wrap_storage_with_split_range_cache}; /// Process-wide Foyer hybrid cache for exact split byte-range payloads. pub struct FoyerSplitRangeCache { pub(crate) cache: foyer::HybridCache, - #[allow(dead_code)] pub(crate) max_entry_size: usize, - #[allow(dead_code)] pub(crate) block_size: usize, } diff --git a/quickwit/quickwit-storage/src/split_range_cache/storage.rs b/quickwit/quickwit-storage/src/split_range_cache/storage.rs new file mode 100644 index 00000000000..2b2c9d85227 --- /dev/null +++ b/quickwit/quickwit-storage/src/split_range_cache/storage.rs @@ -0,0 +1,289 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed 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 std::future::Future; +use std::ops::Range; +use std::path::Path; +use std::sync::Arc; +use std::{fmt, io}; + +use async_trait::async_trait; +use bytes::Bytes; +use fail::fail_point; +use foyer::Code; +use quickwit_common::uri::Uri; +use tokio::io::AsyncRead; +use tracing::{error, warn}; + +use super::{FoyerSplitRangeCache, SplitRangeCacheKey}; +use crate::stable_deref_bytes::into_owned_bytes; +use crate::storage::SendableAsync; +use crate::{ + BulkDeleteError, OwnedBytes, PutPayload, Storage, StorageError, StorageErrorKind, StorageResult, +}; + +/// Foyer hybrid-cache entry header size in the 0.22.3 block engine. +pub(crate) const FOYER_ENTRY_HEADER_SIZE: usize = 36; +/// Foyer blob index reserved at the end of each block. +pub(crate) const FOYER_BLOB_INDEX_SIZE: usize = 4 * 1024; +/// Foyer disk page size used to align encoded entries. +pub(crate) const FOYER_PAGE_SIZE: usize = 4 * 1024; + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub(crate) enum AdmissionBypass { + MaxEntrySize, + EncodedTooLarge, +} + +pub(crate) fn admission_bypass_reason( + key_size: usize, + value: &Bytes, + max_entry_size: usize, + block_size: usize, +) -> Option { + if value.len() > max_entry_size { + return Some(AdmissionBypass::MaxEntrySize); + } + // `max_entry_size < block_size` is not enough: the disk slot is + // `block_size - blob index` after header, key, and page alignment. + let encoded_len = FOYER_ENTRY_HEADER_SIZE + key_size + Bytes::estimated_size(value); + let aligned_len = encoded_len.div_ceil(FOYER_PAGE_SIZE) * FOYER_PAGE_SIZE; + if aligned_len > block_size - FOYER_BLOB_INDEX_SIZE { + return Some(AdmissionBypass::EncodedTooLarge); + } + None +} + +#[derive(Debug, thiserror::Error)] +#[error(transparent)] +struct LowerStorageError(StorageError); + +pub(crate) enum CacheFetchError { + Lower(StorageError), + Foyer, +} + +impl FoyerSplitRangeCache { + pub(crate) async fn get_or_fetch( + &self, + key: SplitRangeCacheKey, + fetch: F, + ) -> Result + where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + let key_size = key.estimated_size(); + let max_entry_size = self.max_entry_size; + let block_size = self.block_size; + match self + .cache + .get_or_fetch(&key, || async move { + let bytes = fetch().await.map_err(LowerStorageError)?; + if admission_bypass_reason(key_size, &bytes, max_entry_size, block_size).is_some() { + // Foyer keeps this tag on the RAM entry and skips disk enqueue + // on eviction (write-on-eviction). + Ok::<_, LowerStorageError>(( + bytes, + foyer::HybridCacheProperties::default() + .with_location(foyer::Location::InMem), + )) + } else { + Ok((bytes, foyer::HybridCacheProperties::default())) + } + }) + .await + { + Ok(entry) => Ok(entry.value().clone()), + Err(error) => { + if let Some(lower_error) = error.downcast_ref::() { + Err(CacheFetchError::Lower(lower_error.0.clone())) + } else { + warn!( + error = ?error, + "split range cache fetch failed, reading from storage" + ); + Err(CacheFetchError::Foyer) + } + } + } + } +} + +/// Read-only [`Storage`] decorator that caches exact split byte-range payloads. +#[derive(Clone)] +pub struct FoyerSplitRangeStorage { + inner: Arc, + cache: Arc, +} + +/// Wraps `storage` so [`Storage::get_slice`] is served from `cache` on an exact +/// `{object URI, byte range}` key. +pub fn wrap_storage_with_split_range_cache( + cache: Arc, + storage: Arc, +) -> Arc { + Arc::new(FoyerSplitRangeStorage { + inner: storage, + cache, + }) +} + +impl fmt::Debug for FoyerSplitRangeStorage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FoyerSplitRangeStorage") + .field("uri", self.inner.uri()) + .finish() + } +} + +fn unsupported_operation(paths: &[&Path]) -> StorageError { + let msg = "Unsupported operation. FoyerSplitRangeStorage only supports async reads"; + error!(paths=?paths, msg); + io::Error::other(format!("{msg}: {paths:?}")).into() +} + +#[async_trait] +impl Storage for FoyerSplitRangeStorage { + async fn check_connectivity(&self) -> anyhow::Result<()> { + self.inner.check_connectivity().await + } + + async fn put(&self, path: &Path, _payload: Box) -> StorageResult<()> { + Err(unsupported_operation(&[path])) + } + + async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { + self.inner.copy_to(path, output).await + } + + async fn get_slice(&self, path: &Path, byte_range: Range) -> StorageResult { + if byte_range.is_empty() { + return Ok(OwnedBytes::empty()); + } + if should_bypass_cache() { + return self.inner.get_slice(path, byte_range).await; + } + let object_uri = self + .inner + .uri() + .join(path) + .map_err(|error| StorageErrorKind::Internal.with_error(error))? + .into_string(); + let key = SplitRangeCacheKey { + object_uri, + byte_range: byte_range.clone(), + }; + let inner = self.inner.clone(); + let owned_path = path.to_owned(); + let fetch_range = byte_range.clone(); + let fetch_result = self + .cache + .get_or_fetch(key, move || async move { + inner + .get_slice(&owned_path, fetch_range) + .await + .map(Bytes::from_owner) + }) + .await; + match fetch_result { + Ok(bytes) => Ok(into_owned_bytes(bytes)), + Err(CacheFetchError::Lower(storage_error)) => Err(storage_error), + Err(CacheFetchError::Foyer) => self.inner.get_slice(path, byte_range).await, + } + } + + async fn get_slice_stream( + &self, + path: &Path, + range: Range, + ) -> StorageResult> { + self.inner.get_slice_stream(path, range).await + } + + async fn get_all(&self, path: &Path) -> StorageResult { + self.inner.get_all(path).await + } + + async fn delete(&self, path: &Path) -> StorageResult<()> { + Err(unsupported_operation(&[path])) + } + + async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { + Err(BulkDeleteError { + error: Some(unsupported_operation(paths)), + ..Default::default() + }) + } + + async fn file_num_bytes(&self, path: &Path) -> StorageResult { + self.inner.file_num_bytes(path).await + } + + fn uri(&self) -> &Uri { + self.inner.uri() + } +} + +fn should_bypass_cache() -> bool { + fail_point!("split-range-cache-before-get", |_| true); + false +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_admission_bypass_pinned_foyer_block_format() { + let key_size = 40; + let block_size = 2 * FOYER_PAGE_SIZE; + // encoded = 36 + 40 + (usize_len + value_len) = 84 + value_len on 64-bit. + // 4012 => encoded 4096, one page, fits in block_size - blob index. + // 4013 => encoded 4097, two pages, exceeds that slot. + assert_eq!( + admission_bypass_reason( + key_size, + &Bytes::from(vec![0; 4012]), + usize::MAX, + block_size + ), + None + ); + assert_eq!( + admission_bypass_reason( + key_size, + &Bytes::from(vec![0; 4013]), + usize::MAX, + block_size + ), + Some(AdmissionBypass::EncodedTooLarge) + ); + assert_eq!( + admission_bypass_reason(key_size, &Bytes::from(vec![0; 101]), 100, 4 * 1024 * 1024), + Some(AdmissionBypass::MaxEntrySize) + ); + // 5 KiB < max_entry_size 7 KiB < block_size 8 KiB, but the disk slot is + // only 4 KiB after the blob index. + assert_eq!( + admission_bypass_reason( + key_size, + &Bytes::from(vec![0; 5 * 1024]), + 7 * 1024, + 8 * 1024 + ), + Some(AdmissionBypass::EncodedTooLarge) + ); + } +} diff --git a/quickwit/quickwit-storage/src/split_range_cache/tests.rs b/quickwit/quickwit-storage/src/split_range_cache/tests.rs index e39b5750546..d1454e51cf8 100644 --- a/quickwit/quickwit-storage/src/split_range_cache/tests.rs +++ b/quickwit/quickwit-storage/src/split_range_cache/tests.rs @@ -12,9 +12,28 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt; +use std::ops::Range; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use quickwit_common::uri::Uri; use quickwit_config::SplitRangeCacheWritePolicy; +use tokio::io::AsyncRead; +use tokio::sync::watch; use super::*; +use crate::storage::SendableAsync; +use crate::{ + BulkDeleteError, OwnedBytes, PutPayload, RamStorageBuilder, Storage, StorageErrorKind, + StorageResult, wrap_storage_with_split_range_cache, +}; + +const SPLIT_PATH: &str = "a.split"; +const SPLIT_BYTES: &[u8] = b"abcde"; #[test] fn test_flush_on_close_pairs_with_write_policy() { @@ -51,3 +70,359 @@ async fn test_split_range_cache_builder_write_on_insertion() { ); cache.close().await.unwrap(); } + +struct LowerProbe { + inner: Arc, + get_slice_calls: AtomicUsize, + get_slice_completed: AtomicUsize, + gate: watch::Receiver, +} + +impl fmt::Debug for LowerProbe { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("LowerProbe") + .field("uri", self.inner.uri()) + .finish() + } +} + +#[async_trait] +impl Storage for LowerProbe { + async fn check_connectivity(&self) -> anyhow::Result<()> { + self.inner.check_connectivity().await + } + + async fn put(&self, path: &Path, payload: Box) -> StorageResult<()> { + self.inner.put(path, payload).await + } + + async fn copy_to(&self, path: &Path, output: &mut dyn SendableAsync) -> StorageResult<()> { + self.inner.copy_to(path, output).await + } + + async fn get_slice(&self, path: &Path, range: Range) -> StorageResult { + self.get_slice_calls.fetch_add(1, Ordering::Relaxed); + let mut gate = self.gate.clone(); + let _ = gate.wait_for(|open| *open).await; + let result = self.inner.get_slice(path, range).await; + self.get_slice_completed.fetch_add(1, Ordering::Relaxed); + result + } + + async fn get_slice_stream( + &self, + path: &Path, + range: Range, + ) -> StorageResult> { + self.inner.get_slice_stream(path, range).await + } + + async fn get_all(&self, path: &Path) -> StorageResult { + self.inner.get_all(path).await + } + + async fn delete(&self, path: &Path) -> StorageResult<()> { + self.inner.delete(path).await + } + + async fn bulk_delete<'a>(&self, paths: &[&'a Path]) -> Result<(), BulkDeleteError> { + self.inner.bulk_delete(paths).await + } + + async fn file_num_bytes(&self, path: &Path) -> StorageResult { + self.inner.file_num_bytes(path).await + } + + fn uri(&self) -> &Uri { + self.inner.uri() + } +} + +struct Fixture { + storage: Arc, + cache: Arc, + lower: Arc, + gate_tx: watch::Sender, + _temp_dir: tempfile::TempDir, +} + +impl Fixture { + async fn new() -> Self { + Self::with_payload(SPLIT_BYTES, true).await + } + + async fn new_with_blocked_lower_read() -> Self { + Self::with_payload(SPLIT_BYTES, false).await + } + + async fn with_payload(payload: &[u8], gate_open: bool) -> Self { + let temp_dir = tempfile::tempdir().unwrap(); + let cache = Arc::new( + FoyerSplitRangeCache::open(&config_for_test(temp_dir.path())) + .await + .unwrap(), + ); + let ram: Arc = Arc::new( + RamStorageBuilder::default() + .put(SPLIT_PATH, payload) + .build(), + ); + let (gate_tx, gate_rx) = watch::channel(gate_open); + let lower = Arc::new(LowerProbe { + inner: ram, + get_slice_calls: AtomicUsize::new(0), + get_slice_completed: AtomicUsize::new(0), + gate: gate_rx, + }); + let storage = wrap_storage_with_split_range_cache(cache.clone(), lower.clone()); + Self { + storage, + cache, + lower, + gate_tx, + _temp_dir: temp_dir, + } + } + + fn release_lower_read(&self) { + self.gate_tx.send(true).unwrap(); + } + + fn lower_reads(&self) -> usize { + self.lower.get_slice_calls.load(Ordering::Relaxed) + } + + fn lower_completed(&self) -> usize { + self.lower.get_slice_completed.load(Ordering::Relaxed) + } + + async fn wait_until_lower_read_started(&self) { + wait_until(|| self.lower_reads() > 0, "lower read start").await; + } + + async fn wait_until_lower_read_completed(&self) { + wait_until(|| self.lower_completed() > 0, "lower read completion").await; + } + + async fn close(&self) { + self.cache.close().await.unwrap(); + } +} + +async fn wait_until(predicate: impl Fn() -> bool, what: &str) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + while !predicate() { + if tokio::time::Instant::now() >= deadline { + panic!("timed out waiting for {what}"); + } + tokio::time::sleep(Duration::from_millis(1)).await; + } +} + +#[tokio::test] +async fn test_empty_range_and_exact_hit_behavior() { + let fixture = Fixture::new().await; + let path = Path::new(SPLIT_PATH); + assert!( + fixture + .storage + .get_slice(path, 4..4) + .await + .unwrap() + .is_empty() + ); + assert_eq!(fixture.lower_reads(), 0); + assert_eq!( + fixture + .storage + .get_slice(path, 1..4) + .await + .unwrap() + .as_slice(), + b"bcd" + ); + assert_eq!( + fixture + .storage + .get_slice(path, 1..4) + .await + .unwrap() + .as_slice(), + b"bcd" + ); + assert_eq!(fixture.lower_reads(), 1); + fixture.storage.get_slice(path, 0..5).await.unwrap(); + assert_eq!( + fixture.lower_reads(), + 2, + "covering ranges are distinct keys" + ); + fixture.close().await; +} + +#[tokio::test] +async fn test_identical_concurrent_misses_fetch_once() { + let fixture = Fixture::new_with_blocked_lower_read().await; + let path = Path::new(SPLIT_PATH); + let first = fixture.storage.get_slice(path, 0..4); + let second = fixture.storage.get_slice(path, 0..4); + let release = async { + fixture.wait_until_lower_read_started().await; + fixture.release_lower_read(); + }; + let (first_result, second_result, _) = tokio::join!(first, second, release); + assert_eq!(first_result.unwrap(), second_result.unwrap()); + assert_eq!(fixture.lower_reads(), 1); + fixture.close().await; +} + +#[tokio::test] +async fn test_remote_error_is_not_cached_or_rewritten() { + let fixture = Fixture::new().await; + for _ in 0..2 { + let error = fixture + .storage + .get_slice(Path::new("missing.split"), 0..4) + .await + .unwrap_err(); + assert_eq!(error.kind(), StorageErrorKind::NotFound); + } + assert_eq!(fixture.lower_reads(), 2); + fixture.close().await; +} + +#[tokio::test] +async fn test_writes_are_unsupported() { + let fixture = Fixture::new().await; + let path = Path::new(SPLIT_PATH); + let put_error = fixture + .storage + .put(path, Box::new(b"x".to_vec())) + .await + .unwrap_err(); + assert_eq!(put_error.kind(), StorageErrorKind::Io); + assert!( + put_error + .to_string() + .contains("Unsupported operation. FoyerSplitRangeStorage only supports async reads") + ); + let delete_error = fixture.storage.delete(path).await.unwrap_err(); + assert_eq!(delete_error.kind(), StorageErrorKind::Io); + let bulk_error = fixture.storage.bulk_delete(&[path]).await.unwrap_err(); + assert_eq!( + bulk_error.error.as_ref().unwrap().kind(), + StorageErrorKind::Io + ); + fixture.close().await; +} + +#[tokio::test] +async fn test_get_all_is_not_cached() { + let fixture = Fixture::new().await; + let path = Path::new(SPLIT_PATH); + assert_eq!( + fixture.storage.get_all(path).await.unwrap().as_slice(), + SPLIT_BYTES + ); + assert_eq!(fixture.lower_reads(), 0); + fixture.storage.get_slice(path, 0..5).await.unwrap(); + assert_eq!(fixture.lower_reads(), 1); + fixture.close().await; +} + +#[tokio::test] +async fn test_initiating_caller_drop_surviving_waiter_succeeds() { + let fixture = Fixture::new_with_blocked_lower_read().await; + let path = Path::new(SPLIT_PATH); + let mut initiating = Box::pin(fixture.storage.get_slice(path, 0..4)); + tokio::select! { + biased; + result = &mut initiating => panic!("fetch completed before release: {result:?}"), + () = fixture.wait_until_lower_read_started() => {} + } + drop(initiating); + let waiter = fixture.storage.get_slice(path, 0..4); + fixture.release_lower_read(); + assert_eq!(waiter.await.unwrap().as_slice(), b"abcd"); + assert_eq!(fixture.lower_reads(), 1); + fixture.close().await; +} + +#[tokio::test] +async fn test_waiter_drop_does_not_cancel_fetch() { + let fixture = Fixture::new_with_blocked_lower_read().await; + let path = Path::new(SPLIT_PATH); + let mut initiating = Box::pin(fixture.storage.get_slice(path, 0..4)); + tokio::select! { + biased; + result = &mut initiating => panic!("fetch completed before release: {result:?}"), + () = fixture.wait_until_lower_read_started() => {} + } + let mut waiter = Box::pin(fixture.storage.get_slice(path, 0..4)); + for _ in 0..16 { + tokio::select! { + biased; + result = &mut waiter => panic!("waiter completed before release: {result:?}"), + () = tokio::task::yield_now() => {} + } + } + drop(waiter); + fixture.release_lower_read(); + assert_eq!(initiating.await.unwrap().as_slice(), b"abcd"); + assert_eq!(fixture.lower_reads(), 1); + fixture.close().await; +} + +#[tokio::test] +async fn test_all_callers_dropped_detached_completion() { + let fixture = Fixture::new_with_blocked_lower_read().await; + let path = Path::new(SPLIT_PATH); + let mut initiating = Box::pin(fixture.storage.get_slice(path, 0..4)); + tokio::select! { + biased; + result = &mut initiating => panic!("fetch completed before release: {result:?}"), + () = fixture.wait_until_lower_read_started() => {} + } + drop(initiating); + fixture.release_lower_read(); + fixture.wait_until_lower_read_completed().await; + assert_eq!( + fixture + .storage + .get_slice(path, 0..4) + .await + .unwrap() + .as_slice(), + b"abcd" + ); + assert_eq!(fixture.lower_reads(), 1); + fixture.close().await; +} + +#[tokio::test] +async fn test_oversized_value_is_memory_only_and_returned() { + let payload = vec![7u8; 3 * 1024 * 1024]; + let fixture = Fixture::with_payload(&payload, true).await; + let path = Path::new(SPLIT_PATH); + let range = 0..payload.len(); + assert_eq!( + fixture + .storage + .get_slice(path, range.clone()) + .await + .unwrap() + .as_slice(), + payload.as_slice() + ); + assert_eq!( + fixture + .storage + .get_slice(path, range) + .await + .unwrap() + .as_slice(), + payload.as_slice() + ); + assert_eq!(fixture.lower_reads(), 1); + fixture.close().await; +}