diff --git a/quickwit/quickwit-indexing/src/actors/uploader.rs b/quickwit/quickwit-indexing/src/actors/uploader.rs index 94db9379ab2..d641ac6c271 100644 --- a/quickwit/quickwit-indexing/src/actors/uploader.rs +++ b/quickwit/quickwit-indexing/src/actors/uploader.rs @@ -338,6 +338,7 @@ impl Handler for Uploader { report_splits.push(ReportSplit { storage_uri: split_store.remote_uri().to_string(), split_id: packaged_split.split_id().to_string(), + num_bytes: split_metadata.footer_offsets.end, }); split_metadata_list.push(split_metadata); diff --git a/quickwit/quickwit-proto/protos/quickwit/search.proto b/quickwit/quickwit-proto/protos/quickwit/search.proto index 7cc168a858b..3244f227c6f 100644 --- a/quickwit/quickwit-proto/protos/quickwit/search.proto +++ b/quickwit/quickwit-proto/protos/quickwit/search.proto @@ -107,6 +107,8 @@ message ReportSplit { string split_id = 2; // The storage uri. This URI does NOT include the split id. string storage_uri = 1; + // The size of the split file, in bytes. + uint64 num_bytes = 3; } message ReportSplitsRequest { diff --git a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs index 5fef3bfdc29..c69a5b597a5 100644 --- a/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs +++ b/quickwit/quickwit-proto/src/codegen/quickwit/quickwit.search.rs @@ -43,6 +43,9 @@ pub struct ReportSplit { /// The storage uri. This URI does NOT include the split id. #[prost(string, tag = "1")] pub storage_uri: ::prost::alloc::string::String, + /// The size of the split file, in bytes. + #[prost(uint64, tag = "3")] + pub num_bytes: u64, } #[derive(serde::Serialize, serde::Deserialize, utoipa::ToSchema)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/quickwit/quickwit-search/src/leaf.rs b/quickwit/quickwit-search/src/leaf.rs index f184375438e..e6fe80ebf22 100644 --- a/quickwit/quickwit-search/src/leaf.rs +++ b/quickwit/quickwit-search/src/leaf.rs @@ -48,6 +48,7 @@ use tantivy::{DateTime, Index, ReloadPolicy, Searcher, TantivyError, Term}; use tantivy_fst::DisjunctionRegex; use tokio::task::{JoinError, JoinSet}; use tracing::*; +use ulid::Ulid; use crate::collector::{IncrementalCollector, make_collector_for_split, make_merge_collector}; use crate::metrics::SplitSearchOutcomeCounters; @@ -109,12 +110,21 @@ pub(crate) async fn open_split_bundle( // We wrap the top-level storage with the split cache. // 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() { - SplitCache::wrap_storage(split_cache.clone(), index_storage.clone()) - } else { - index_storage.clone() - }; + let index_storage_with_split_cache = if let Some(split_cache) = + searcher_context.split_cache_opt.as_ref() + { + let split_ulid = Ulid::from_str(&split_and_footer_offsets.split_id).with_context(|| { + format!("invalid split ulid `{}`", split_and_footer_offsets.split_id) + })?; + SplitCache::wrap_storage( + split_cache.clone(), + index_storage.clone(), + split_ulid, + split_and_footer_offsets.split_footer_end, + ) + } else { + index_storage.clone() + }; let (hotcache_bytes, bundle_storage) = BundleStorage::open_from_split_data( index_storage_with_split_cache, diff --git a/quickwit/quickwit-storage/src/split_cache/download_task.rs b/quickwit/quickwit-storage/src/split_cache/download_task.rs index 9af0390bfb5..cab490f7f36 100644 --- a/quickwit/quickwit-storage/src/split_cache/download_task.rs +++ b/quickwit/quickwit-storage/src/split_cache/download_task.rs @@ -32,6 +32,7 @@ async fn download_split( split_ulid, storage_uri, living_token: _, + num_bytes: _, } = candidate_split; let split_filename = split_file(*split_ulid); let target_filepath = root_path.join(&split_filename); diff --git a/quickwit/quickwit-storage/src/split_cache/mod.rs b/quickwit/quickwit-storage/src/split_cache/mod.rs index 5ee6ea2f1a5..0f84d5b6b85 100644 --- a/quickwit/quickwit-storage/src/split_cache/mod.rs +++ b/quickwit/quickwit-storage/src/split_cache/mod.rs @@ -126,10 +126,22 @@ impl SplitCache { } /// Wraps a storage with our split cache. - pub fn wrap_storage(self_arc: Arc, storage: Arc) -> Arc { - let cache = Arc::new(SplitCacheBackingStorage { + /// + /// The returned storage is specific to the given `split_id`: it is only + /// ever meant to be used to access that single split's file. `num_bytes` is + /// the split's total file size, it is only used to register that split as a + /// download candidate with a known size. + pub fn wrap_storage( + self_arc: Arc, + storage: Arc, + split_id: Ulid, + num_bytes: u64, + ) -> Arc { + let cache = Arc::new(SplitScopedStorageCache { split_cache: self_arc, storage_root_uri: storage.uri().clone(), + split_id, + num_bytes, }); wrap_storage_with_cache(cache, storage) } @@ -146,22 +158,28 @@ impl SplitCache { error!(storage_uri=%report_split.storage_uri, "received invalid storage uri: ignoring"); continue; }; - split_table.report(split_ulid, storage_uri); + split_table.report(split_ulid, storage_uri, report_split.num_bytes); } } // Returns a split guard object. As long as it is not dropped, the // split won't be evicted from the cache. - async fn get_split_file(&self, split_id: Ulid, storage_uri: &Uri) -> Option { - // We touch before even checking the fd cache in order to update the file's last access time - // for the file cache. - let num_bytes_opt: Option = self + async fn get_split_file( + &self, + split_id: Ulid, + storage_uri: &Uri, + num_bytes: u64, + ) -> Option { + // We touch before even checking the fd cache in order to update the file's last access + // time for the file cache. + let is_on_disk = self .split_table .lock() .unwrap() - .touch(split_id, storage_uri); - - let num_bytes = num_bytes_opt?; + .touch(split_id, storage_uri, num_bytes); + if !is_on_disk { + return None; + } self.fd_cache .get_or_open_split_file(&self.root_path, split_id, num_bytes) .await @@ -192,26 +210,32 @@ fn split_id_from_path(split_path: &Path) -> Option { Ulid::from_str(split_id_str).ok() } -struct SplitCacheBackingStorage { +/// A `StorageCache` implementation that is specific to a single split: it is only ever meant to +/// be used to access `split_id`'s own file. It is not a general-purpose cache for arbitrary splits +/// living under `storage_root_uri`, since `num_bytes` would otherwise no longer make sense for the +/// paths it doesn't apply to. +struct SplitScopedStorageCache { split_cache: Arc, storage_root_uri: Uri, + split_id: Ulid, + num_bytes: u64, } -impl SplitCacheBackingStorage { +impl SplitScopedStorageCache { async fn get_impl(&self, path: &Path, byte_range: Range) -> Option { - let split_id = split_id_from_path(path)?; + self.verify_cache_isnt_used_for_other_splits(path)?; let split_file: SplitFile = self .split_cache - .get_split_file(split_id, &self.storage_root_uri) + .get_split_file(self.split_id, &self.storage_root_uri, self.num_bytes) .await?; split_file.get_range(byte_range).await.ok() } async fn get_all_impl(&self, path: &Path) -> Option { - let split_id = split_id_from_path(path)?; + self.verify_cache_isnt_used_for_other_splits(path)?; let split_file = self .split_cache - .get_split_file(split_id, &self.storage_root_uri) + .get_split_file(self.split_id, &self.storage_root_uri, self.num_bytes) .await?; split_file.get_all().await.ok() } @@ -225,10 +249,19 @@ impl SplitCacheBackingStorage { split_metrics.misses_num_items.inc(); } } + + fn verify_cache_isnt_used_for_other_splits(&self, path: &Path) -> Option<()> { + debug_assert_eq!(split_id_from_path(path), Some(self.split_id)); + if split_id_from_path(path) != Some(self.split_id) { + error!("split scoped cache used on an unexpected split, please report"); + return None; + } + Some(()) + } } #[async_trait] -impl StorageCache for SplitCacheBackingStorage { +impl StorageCache for SplitScopedStorageCache { async fn get(&self, path: &Path, byte_range: Range) -> Option { let result = self.get_impl(path, byte_range).await; self.record_hit_metrics(result.as_ref()); diff --git a/quickwit/quickwit-storage/src/split_cache/split_table.rs b/quickwit/quickwit-storage/src/split_cache/split_table.rs index e97557f03ed..5cb3bc9978d 100644 --- a/quickwit/quickwit-storage/src/split_cache/split_table.rs +++ b/quickwit/quickwit-storage/src/split_cache/split_table.rs @@ -17,6 +17,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::{Arc, Weak}; use std::time::{Duration, Instant}; +use quickwit_common::rate_limited_warn; use quickwit_common::uri::Uri; use quickwit_config::SplitCacheLimits; use ulid::Ulid; @@ -58,8 +59,13 @@ impl Eq for SplitKey {} #[derive(Clone, Debug)] enum Status { Candidate(CandidateSplit), - Downloading { alive_token: Weak<()> }, - OnDisk { num_bytes: u64 }, + Downloading { + alive_token: Weak<()>, + num_bytes: u64, + }, + OnDisk { + num_bytes: u64, + }, } impl PartialEq for Status { @@ -97,6 +103,11 @@ pub struct SplitInfo { /// /// It is possible for the split table size in bytes to exceed its limits, by at /// most one split. +/// +/// This accounts for the size of splits that are currently downloading as well as splits +/// already on disk: we know a split's size as soon as it becomes a download candidate, so we +/// reserve that space upfront. This way, several concurrent downloads can never complete and +/// push the total disk usage over the configured limits. pub struct SplitTable { on_disk_splits: BTreeSet, downloading_splits: BTreeSet, @@ -105,6 +116,7 @@ pub struct SplitTable { origin_time: Instant, limits: SplitCacheLimits, on_disk_bytes: u64, + downloading_bytes: u64, } impl SplitTable { @@ -121,6 +133,7 @@ impl SplitTable { origin_time, limits, on_disk_bytes: 0u64, + downloading_bytes: 0u64, }; split_table.acknowledge_on_disk_splits(existing_filepaths); split_table @@ -180,7 +193,10 @@ impl SplitTable { let split_info = self.split_to_status.remove(&split_ulid)?; let split_queue: &mut BTreeSet = match split_info.status { Status::Candidate { .. } => &mut self.candidate_splits, - Status::Downloading { .. } => &mut self.downloading_splits, + Status::Downloading { num_bytes, .. } => { + self.downloading_bytes -= num_bytes; + &mut self.downloading_splits + } Status::OnDisk { num_bytes } => { self.on_disk_bytes -= num_bytes; &mut self.on_disk_splits @@ -188,7 +204,7 @@ impl SplitTable { }; let is_in_queue = split_queue.remove(&split_info.split_key); assert!(is_in_queue); - if let Status::Downloading { alive_token } = &split_info.status + if let Status::Downloading { alive_token, .. } = &split_info.status && alive_token.strong_count() == 0 { return None; @@ -205,7 +221,7 @@ impl SplitTable { let mut splits_to_remove = Vec::new(); for split in &self.downloading_splits { if let Some(split_info) = self.split_to_status.get(&split.split_ulid) - && let Status::Downloading { alive_token } = &split_info.status + && let Status::Downloading { alive_token, .. } = &split_info.status && alive_token.strong_count() == 0 { splits_to_remove.push(split.split_ulid); @@ -228,7 +244,10 @@ impl SplitTable { self.truncate_candidate_list(); self.candidate_splits.insert(split_info.split_key) } - Status::Downloading { .. } => self.downloading_splits.insert(split_info.split_key), + Status::Downloading { num_bytes, .. } => { + self.downloading_bytes += num_bytes; + self.downloading_splits.insert(split_info.split_key) + } Status::OnDisk { num_bytes } => { self.on_disk_bytes += num_bytes; self.on_disk_splits.insert(split_info.split_key) @@ -245,13 +264,14 @@ impl SplitTable { assert!(split_ulid_was_absent); } - /// Touch the file, updating its last access time, possibly extending its life in the - /// cache (if in cache). + /// Touches the file, updating its last access time, possibly extending its + /// life in the cache (if in cache). The split's file size is used when + /// registering a brand new download candidate. /// - /// If the file is already on the disk cache, return `Some(num_bytes)`. - /// If the file is not in cache, register the file in the candidate for download list, and - /// return `None`. - pub fn touch(&mut self, split_ulid: Ulid, storage_uri: &Uri) -> Option { + /// Returns `true` if the file is already on the disk cache. + /// Returns `false` if the file is not in cache (the file is registered as + /// candidate in the download list). + pub fn touch(&mut self, split_ulid: Ulid, storage_uri: &Uri, num_bytes: u64) -> bool { let timestamp = compute_timestamp(self.origin_time); let status = self.mutate_split(split_ulid, |old_split_info| { if let Some(mut split_info) = old_split_info { @@ -266,15 +286,12 @@ impl SplitTable { storage_uri: storage_uri.clone(), split_ulid, living_token: Arc::new(()), + num_bytes, }); Some(SplitInfo { split_key, status }) } }); - if let Some(Status::OnDisk { num_bytes }) = status { - Some(num_bytes) - } else { - None - } + matches!(status, Some(Status::OnDisk { .. })) } /// Mutates a split ulid. @@ -293,25 +310,7 @@ impl SplitTable { Some(new_status) } - fn change_split_status(&mut self, split_ulid: Ulid, status: Status) { - let start_time = self.origin_time; - self.mutate_split(split_ulid, move |split_info_opt| { - if let Some(mut split_info) = split_info_opt { - split_info.status = status; - Some(split_info) - } else { - Some(SplitInfo { - split_key: SplitKey { - last_accessed: compute_timestamp(start_time), - split_ulid, - }, - status, - }) - } - }); - } - - pub(crate) fn report(&mut self, split_ulid: Ulid, storage_uri: Uri) { + pub(crate) fn report(&mut self, split_ulid: Ulid, storage_uri: Uri, num_bytes: u64) { let origin_time = self.origin_time; self.mutate_split(split_ulid, move |split_info_opt| { if let Some(split_info) = split_info_opt { @@ -326,6 +325,7 @@ impl SplitTable { storage_uri, split_ulid, living_token: Arc::new(()), + num_bytes, }); Some(SplitInfo { split_key, status }) }); @@ -341,7 +341,44 @@ impl SplitTable { } pub(crate) fn register_as_downloaded(&mut self, split_ulid: Ulid, num_bytes: u64) { - self.change_split_status(split_ulid, Status::OnDisk { num_bytes }); + let start_time = self.origin_time; + self.mutate_split(split_ulid, move |split_info_opt| { + let new_status = Status::OnDisk { num_bytes }; + if let Some(mut split_info) = split_info_opt { + let was_expected_download = matches!( + split_info.status, + Status::Downloading { num_bytes: expected_dl_bytes, .. } + if expected_dl_bytes == num_bytes + ); + if !was_expected_download { + // This can happen during initial upgrade (indexers sending + // 0 as size) but shouldn't happen afterwards. + rate_limited_warn!( + limit_per_min = 1, + downloaded_bytes = num_bytes, + previous_status = ?split_info.status, + %split_ulid, + "unexpected split status when reported in cache" + ); + } + split_info.status = new_status; + Some(split_info) + } else { + // This is not expected to happen. + rate_limited_warn!( + limit_per_min = 1, + %split_ulid, + "split registered in cache without being reported first" + ); + Some(SplitInfo { + split_key: SplitKey { + last_accessed: compute_timestamp(start_time), + split_ulid, + }, + status: new_status, + }) + } + }); record_split_added_to_disk(num_bytes); } @@ -359,7 +396,10 @@ impl SplitTable { let alive_token = Arc::downgrade(&candidate_split.living_token); self.insert(SplitInfo { split_key: split_info.split_key, - status: Status::Downloading { alive_token }, + status: Status::Downloading { + alive_token, + num_bytes: candidate_split.num_bytes, + }, }); Some(candidate_split) } @@ -369,7 +409,9 @@ impl SplitTable { } fn is_out_of_limits(&self) -> bool { - if self.on_disk_splits.is_empty() { + if self.on_disk_splits.is_empty() && self.downloading_splits.is_empty() { + // Always allow at least one split to be downloaded/kept, even if it alone exceeds + // the configured limits. return false; } if self.on_disk_splits.len() + self.downloading_splits.len() @@ -377,7 +419,7 @@ impl SplitTable { { return true; } - if self.on_disk_bytes > self.limits.max_num_bytes.as_u64() { + if self.on_disk_bytes + self.downloading_bytes > self.limits.max_num_bytes.as_u64() { return true; } false @@ -453,6 +495,7 @@ pub(crate) struct CandidateSplit { pub storage_uri: Uri, pub split_ulid: Ulid, pub living_token: Arc<()>, + pub num_bytes: u64, } pub(crate) struct DownloadOpportunity { @@ -499,8 +542,8 @@ mod tests { let ulids = sorted_split_ulids(2); let ulid1 = ulids[0]; let ulid2 = ulids[1]; - split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI)); - split_table.report(ulid2, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI), 10_000); + split_table.report(ulid2, Uri::for_test(TEST_STORAGE_URI), 20_000); let candidate = split_table.best_candidate().unwrap(); assert_eq!(candidate.split_ulid, ulid2); } @@ -519,10 +562,10 @@ mod tests { let ulids = sorted_split_ulids(2); let ulid1 = ulids[0]; let ulid2 = ulids[1]; - split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI)); - split_table.report(ulid2, Uri::for_test(TEST_STORAGE_URI)); - let num_bytes_opt = split_table.touch(ulid1, &Uri::for_test("s3://test1/")); - assert!(num_bytes_opt.is_none()); + split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI), 10_000); + split_table.report(ulid2, Uri::for_test(TEST_STORAGE_URI), 20_000); + let is_on_disk = split_table.touch(ulid1, &Uri::for_test("s3://test1/"), 15_000); + assert!(!is_on_disk); let candidate = split_table.best_candidate().unwrap(); assert_eq!(candidate.split_ulid, ulid1); } @@ -539,19 +582,16 @@ mod tests { Default::default(), ); let ulid1 = Ulid::new(); - split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI), 10_000_000); assert_eq!(split_table.num_bytes(), 0); let download = split_table.start_download(ulid1); assert!(download.is_some()); assert!(split_table.start_download(ulid1).is_none()); split_table.register_as_downloaded(ulid1, 10_000_000); assert_eq!(split_table.num_bytes(), 10_000_000); - assert_eq!( - split_table.touch(ulid1, &Uri::for_test(TEST_STORAGE_URI)), - Some(10_000_000) - ); + assert!(split_table.touch(ulid1, &Uri::for_test(TEST_STORAGE_URI), 10_000_000)); let ulid2 = Ulid::new(); - split_table.report(ulid2, Uri::for_test("s3://test`/")); + split_table.report(ulid2, Uri::for_test("s3://test`/"), 3_000_000); let download = split_table.start_download(ulid2); assert!(download.is_some()); assert!(split_table.start_download(ulid2).is_none()); @@ -560,6 +600,37 @@ mod tests { assert_eq!(split_table.num_bytes(), 13_000_000); } + #[test] + fn test_downloading_splits_count_towards_byte_limit() { + let mut split_table = SplitTable::with_limits_and_existing_splits( + SplitCacheLimits { + max_num_bytes: ByteSize::kb(100), + max_num_splits: NonZeroU32::new(10).unwrap(), + num_concurrent_downloads: NonZeroU32::new(2).unwrap(), + max_file_descriptors: NonZeroU32::new(100).unwrap(), + }, + Default::default(), + ); + let ulids = sorted_split_ulids(2); + let ulid1 = ulids[0]; + let ulid2 = ulids[1]; + split_table.report(ulid2, Uri::for_test(TEST_STORAGE_URI), 10_000); + // Reported last, so it is the best (most recently touched) candidate. Its size + // alone exceeds the configured byte limit. + split_table.report(ulid1, Uri::for_test(TEST_STORAGE_URI), 150_000); + + // The cache always makes progress: the best candidate is allowed to start + // downloading even though its size alone exceeds the limit. + let opportunity = split_table.find_download_opportunity().unwrap(); + assert_eq!(opportunity.split_to_download.split_ulid, ulid1); + assert_eq!(split_table.downloading_bytes, 150_000); + + // A second, concurrent download is blocked: the first split's size (still only + // "downloading", not yet on disk) already exceeds the byte limit on its own, and + // there is nothing on disk yet to evict to make room for the new candidate. + assert!(split_table.find_download_opportunity().is_none()); + } + #[test] fn test_eviction_due_to_size() { let mut split_table = SplitTable::with_limits_and_existing_splits( @@ -582,11 +653,11 @@ mod tests { (split_ulids[5], 300_000), ]; for (split_ulid, num_bytes) in splits { - split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI), num_bytes); split_table.register_as_downloaded(split_ulid, num_bytes); } let new_ulid = Ulid::new(); - split_table.report(new_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(new_ulid, Uri::for_test(TEST_STORAGE_URI), 50_000); let DownloadOpportunity { splits_to_delete, split_to_download, @@ -620,11 +691,11 @@ mod tests { (split_ulids[5], 300_000), ]; for (split_ulid, num_bytes) in splits { - split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI), num_bytes); split_table.register_as_downloaded(split_ulid, num_bytes); } let new_ulid = Ulid::new(); - split_table.report(new_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(new_ulid, Uri::for_test(TEST_STORAGE_URI), 50_000); let DownloadOpportunity { splits_to_delete, split_to_download, @@ -645,10 +716,10 @@ mod tests { Default::default(), ); let split_ulid = Ulid::new(); - split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI), 42_000); let candidate = split_table.start_download(split_ulid).unwrap(); // This report should be cancelled as we have a download currently running. - split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI), 42_000); assert!(split_table.start_download(split_ulid).is_none()); std::mem::drop(candidate); @@ -657,7 +728,7 @@ mod tests { assert!(split_table.start_download(split_ulid).is_none()); // This report should be considered as our candidate (and its alive token has been dropped) - split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI), 42_000); let candidate2 = split_table.start_download(split_ulid).unwrap(); assert_eq!(candidate2.split_ulid, split_ulid); @@ -676,7 +747,7 @@ mod tests { ); for i in 1..2_000 { let split_ulid = Ulid::new(); - split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI)); + split_table.report(split_ulid, Uri::for_test(TEST_STORAGE_URI), 1_024); assert_eq!( split_table.candidate_splits.len(), i.min(super::MAX_NUM_CANDIDATES) @@ -702,6 +773,7 @@ mod tests { storage_uri: Uri::for_test(TEST_STORAGE_URI), split_ulid, living_token: Arc::new(()), + num_bytes: 4_096, }; let split_info = SplitInfo { split_key: SplitKey {