1use std::ops::Range;
16use std::sync::Arc;
17use std::time::Instant;
18
19use api::v1::index::{BloomFilterLoc, BloomFilterMeta};
20use async_trait::async_trait;
21use bytes::Bytes;
22use index::bloom_filter::error::{Error, IoSnafu, Result, SharedReadSnafu};
23use index::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader};
24use index::target::IndexTarget;
25use snafu::IntoError;
26use store_api::storage::{FileId, IndexVersion};
27
28use crate::cache::index::{INDEX_METADATA_TYPE, IndexCache};
29use crate::metrics::{CACHE_HIT, CACHE_MISS};
30
31const INDEX_TYPE_BLOOM_FILTER_INDEX: &str = "bloom_filter_index";
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
35pub enum Tag {
36 Skipping,
37 Fulltext,
38}
39
40pub type BloomFilterIndexKey = (FileId, IndexVersion, Arc<IndexTarget>, Tag);
41
42pub type BloomFilterIndexCache = IndexCache<BloomFilterIndexKey, BloomFilterMeta>;
44pub type BloomFilterIndexCacheRef = Arc<BloomFilterIndexCache>;
45
46impl BloomFilterIndexCache {
47 pub fn new(index_metadata_cap: u64, index_content_cap: u64, page_size: u64) -> Self {
49 Self::new_with_weighter(
50 index_metadata_cap,
51 index_content_cap,
52 page_size,
53 INDEX_TYPE_BLOOM_FILTER_INDEX,
54 "bloom_filter_index_metadata",
55 bloom_filter_index_metadata_weight,
56 bloom_filter_index_content_weight,
57 )
58 }
59
60 pub fn invalidate_file(&self, file_id: FileId) {
62 self.invalidate_if(move |key| key.0 == file_id);
63 }
64}
65
66fn bloom_filter_index_metadata_weight(k: &BloomFilterIndexKey, meta: &Arc<BloomFilterMeta>) -> u32 {
68 let base = k.0.as_bytes().len()
69 + std::mem::size_of::<IndexVersion>()
70 + k.2.approx_mem_size()
71 + std::mem::size_of::<Tag>()
72 + std::mem::size_of::<BloomFilterMeta>();
73
74 let vec_estimated = meta.segment_loc_indices.len() * std::mem::size_of::<u64>()
75 + meta.bloom_filter_locs.len() * std::mem::size_of::<BloomFilterLoc>();
76
77 (base + vec_estimated) as u32
78}
79
80fn bloom_filter_index_content_weight((k, _): &(BloomFilterIndexKey, u64), v: &Bytes) -> u32 {
82 (k.0.as_bytes().len() + k.2.approx_mem_size() + v.len()) as u32
83}
84
85pub struct CachedBloomFilterIndexBlobReader<R> {
87 file_id: FileId,
88 index_version: IndexVersion,
89 target: Arc<IndexTarget>,
90 tag: Tag,
91 blob_size: u64,
92 inner: R,
93 cache: BloomFilterIndexCacheRef,
94}
95
96impl<R> CachedBloomFilterIndexBlobReader<R> {
97 pub fn new(
99 file_id: FileId,
100 index_version: IndexVersion,
101 target: IndexTarget,
102 tag: Tag,
103 blob_size: u64,
104 inner: R,
105 cache: BloomFilterIndexCacheRef,
106 ) -> Self {
107 Self {
108 file_id,
109 index_version,
110 target: Arc::new(target),
111 tag,
112 blob_size,
113 inner,
114 cache,
115 }
116 }
117}
118
119#[async_trait]
120impl<R: BloomFilterReader + Send> BloomFilterReader for CachedBloomFilterIndexBlobReader<R> {
121 async fn range_read(
122 &self,
123 offset: u64,
124 size: u32,
125 metrics: Option<&mut BloomFilterReadMetrics>,
126 ) -> Result<Bytes> {
127 let start = metrics.as_ref().map(|_| Instant::now());
128 let inner = &self.inner;
129 let (result, cache_metrics) = self
130 .cache
131 .get_or_load(
132 (
133 self.file_id,
134 self.index_version,
135 self.target.clone(),
136 self.tag,
137 ),
138 self.blob_size,
139 offset,
140 size,
141 move |ranges| async move { inner.read_vec(&ranges, None).await },
142 )
143 .await
144 .map_err(map_page_error)?;
145
146 if let Some(m) = metrics {
147 m.total_ranges += cache_metrics.num_pages;
148 m.total_bytes += cache_metrics.page_bytes;
149 m.cache_hit += cache_metrics.cache_hit;
150 m.cache_miss += cache_metrics.cache_miss;
151 if let Some(start) = start {
152 m.fetch_elapsed += start.elapsed();
153 }
154 }
155
156 Ok(result)
157 }
158
159 async fn read_vec(
160 &self,
161 ranges: &[Range<u64>],
162 metrics: Option<&mut BloomFilterReadMetrics>,
163 ) -> Result<Vec<Bytes>> {
164 let start = metrics.as_ref().map(|_| Instant::now());
165
166 let inner = &self.inner;
167 let (pages, total_cache_metrics) = self
168 .cache
169 .get_or_load_vec(
170 (
171 self.file_id,
172 self.index_version,
173 self.target.clone(),
174 self.tag,
175 ),
176 self.blob_size,
177 ranges,
178 move |ranges| async move { inner.read_vec(&ranges, None).await },
179 )
180 .await
181 .map_err(map_page_error)?;
182 if let Some(m) = metrics {
183 m.total_ranges += total_cache_metrics.num_pages;
184 m.total_bytes += total_cache_metrics.page_bytes;
185 m.cache_hit += total_cache_metrics.cache_hit;
186 m.cache_miss += total_cache_metrics.cache_miss;
187 if let Some(start) = start {
188 m.fetch_elapsed += start.elapsed();
189 }
190 }
191
192 Ok(pages)
193 }
194
195 async fn metadata(
197 &self,
198 metrics: Option<&mut BloomFilterReadMetrics>,
199 ) -> Result<Arc<BloomFilterMeta>> {
200 let key = (
201 self.file_id,
202 self.index_version,
203 self.target.clone(),
204 self.tag,
205 );
206 if let Some(cached) = self.cache.get_metadata(key.clone()).await {
207 CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc();
208 if let Some(m) = metrics {
209 m.cache_hit += 1;
210 }
211 return Ok(cached);
212 }
213
214 let start = metrics.as_ref().map(|_| Instant::now());
215 let collect_metrics = metrics.is_some();
216 let mut loaded = BloomFilterReadMetrics::default();
217 let result = self
218 .cache
219 .get_or_load_metadata(key, async {
220 self.inner
221 .metadata(collect_metrics.then_some(&mut loaded))
222 .await
223 })
224 .await;
225 if let Some(m) = metrics {
226 m.total_bytes += loaded.total_bytes;
230 m.total_ranges += loaded.total_ranges;
231 m.cache_hit += loaded.cache_hit;
232 m.cache_miss += loaded.cache_miss;
233 if let Some(start) = start {
234 m.fetch_elapsed += start.elapsed();
235 }
236 }
237 let meta = result.map_err(|source| SharedReadSnafu.into_error(source))?;
238 CACHE_MISS.with_label_values(&[INDEX_METADATA_TYPE]).inc();
239 Ok(meta)
240 }
241}
242
243fn map_page_error(error: Arc<std::io::Error>) -> Error {
245 if let Some(source) = crate::cache::find_shared_page_source::<Error>(&error) {
246 SharedReadSnafu.into_error(source)
247 } else {
248 IoSnafu.into_error(std::io::Error::other(error))
249 }
250}
251
252#[cfg(test)]
253mod test {
254 use rand::{Rng, RngCore};
255 use store_api::storage::ColumnId;
256
257 use super::*;
258
259 const FUZZ_REPEAT_TIMES: usize = 100;
260
261 struct MetadataReader {
262 release: tokio::sync::Notify,
263 fail: std::sync::atomic::AtomicBool,
264 }
265
266 #[async_trait]
267 impl BloomFilterReader for MetadataReader {
268 async fn range_read(
269 &self,
270 _: u64,
271 _: u32,
272 _: Option<&mut BloomFilterReadMetrics>,
273 ) -> Result<Bytes> {
274 unreachable!("metadata test does not read content pages")
275 }
276
277 async fn read_vec(
278 &self,
279 _: &[Range<u64>],
280 _: Option<&mut BloomFilterReadMetrics>,
281 ) -> Result<Vec<Bytes>> {
282 unreachable!("metadata test does not read content pages")
283 }
284
285 async fn metadata(
286 &self,
287 metrics: Option<&mut BloomFilterReadMetrics>,
288 ) -> Result<Arc<BloomFilterMeta>> {
289 self.release.notified().await;
290 if let Some(metrics) = metrics {
291 metrics.total_bytes += 8192;
292 metrics.total_ranges += 1;
293 metrics.fetch_elapsed += std::time::Duration::from_secs(3600);
295 }
296 if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
297 return Err(IoSnafu.into_error(std::io::Error::other("metadata failed")));
298 }
299 Ok(Arc::new(BloomFilterMeta::default()))
300 }
301 }
302
303 #[tokio::test]
304 async fn metadata_waiters_do_not_duplicate_io_metrics() {
305 for fail in [false, true] {
306 let cache = Arc::new(BloomFilterIndexCache::new(4096, 4096, 1024));
307 let reader = CachedBloomFilterIndexBlobReader::new(
308 FileId::random(),
309 0,
310 IndexTarget::ColumnId(42),
311 Tag::Skipping,
312 8192,
313 MetadataReader {
314 release: tokio::sync::Notify::new(),
315 fail: std::sync::atomic::AtomicBool::new(fail),
316 },
317 cache,
318 );
319 let mut owner_metrics = BloomFilterReadMetrics::default();
320 let mut waiter_metrics = BloomFilterReadMetrics::default();
321 let mut owner = Box::pin(reader.metadata(Some(&mut owner_metrics)));
322 assert!(futures::poll!(&mut owner).is_pending());
323 let mut waiter = Box::pin(reader.metadata(Some(&mut waiter_metrics)));
324 assert!(futures::poll!(&mut waiter).is_pending());
325 reader.inner.release.notify_one();
326 let owner_result = owner.await;
327 let waiter_result = waiter.await;
328 assert_eq!(owner_result.is_err(), fail);
329 assert_eq!(waiter_result.is_err(), fail);
330 assert_eq!(
331 (owner_metrics.total_bytes, owner_metrics.total_ranges),
332 (8192, 1)
333 );
334 assert_eq!(
335 (waiter_metrics.total_bytes, waiter_metrics.total_ranges),
336 (0, 0)
337 );
338 assert_eq!((owner_metrics.cache_hit, waiter_metrics.cache_hit), (0, 0));
339 for metrics in [&owner_metrics, &waiter_metrics] {
340 assert!(metrics.fetch_elapsed > std::time::Duration::ZERO);
343 assert!(metrics.fetch_elapsed < std::time::Duration::from_secs(3600));
344 }
345 if fail {
346 let Error::SharedRead { source: owner, .. } = owner_result.unwrap_err() else {
347 panic!("owner must preserve the typed shared read error");
348 };
349 let Error::SharedRead { source: waiter, .. } = waiter_result.unwrap_err() else {
350 panic!("waiter must preserve the typed shared read error");
351 };
352 assert!(Arc::ptr_eq(&owner, &waiter));
353 } else {
354 let mut hot_metrics = BloomFilterReadMetrics::default();
355 reader.metadata(Some(&mut hot_metrics)).await.unwrap();
356 assert_eq!(hot_metrics.cache_hit, 1);
357 assert_eq!(hot_metrics.total_bytes, 0);
358 }
359 }
360 }
361
362 #[tokio::test]
363 async fn json_paths_have_independent_cache_entries() {
364 let cache = BloomFilterIndexCache::new(1024 * 1024, 1024 * 1024, 1024);
365 let file_id = FileId::random();
366 let first = (
367 file_id,
368 0,
369 Arc::new(IndexTarget::new_json_path(1, vec!["a".into(), "b".into()]).unwrap()),
370 Tag::Skipping,
371 );
372 let second = (
373 file_id,
374 0,
375 Arc::new(IndexTarget::new_json_path(1, vec!["a.b".into()]).unwrap()),
376 Tag::Skipping,
377 );
378 cache
379 .put_metadata(
380 first.clone(),
381 Arc::new(BloomFilterMeta {
382 row_count: 1,
383 ..Default::default()
384 }),
385 )
386 .await;
387 cache
388 .put_metadata(
389 second.clone(),
390 Arc::new(BloomFilterMeta {
391 row_count: 2,
392 ..Default::default()
393 }),
394 )
395 .await;
396 assert_eq!(
397 cache.get_metadata(first.clone()).await.unwrap().row_count,
398 1
399 );
400 assert_eq!(
401 cache.get_metadata(second.clone()).await.unwrap().row_count,
402 2
403 );
404 cache.invalidate_file(file_id);
405 assert!(cache.get_metadata(first).await.is_none());
406 assert!(cache.get_metadata(second).await.is_none());
407 }
408
409 #[test]
410 fn bloom_filter_metadata_weight_counts_vec_contents() {
411 let file_id = FileId::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
412 let version = 0;
413 let column_id: ColumnId = 42;
414 let tag = Tag::Skipping;
415
416 let meta = BloomFilterMeta {
417 rows_per_segment: 128,
418 segment_count: 2,
419 row_count: 256,
420 bloom_filter_size: 1024,
421 segment_loc_indices: vec![0, 64, 128, 192],
422 bloom_filter_locs: vec![
423 BloomFilterLoc {
424 offset: 0,
425 size: 512,
426 element_count: 1000,
427 },
428 BloomFilterLoc {
429 offset: 512,
430 size: 512,
431 element_count: 1000,
432 },
433 ],
434 };
435
436 let weight = bloom_filter_index_metadata_weight(
437 &(
438 file_id,
439 version,
440 Arc::new(IndexTarget::ColumnId(column_id)),
441 tag,
442 ),
443 &Arc::new(meta.clone()),
444 );
445
446 let base = file_id.as_bytes().len()
447 + std::mem::size_of::<IndexVersion>()
448 + std::mem::size_of::<IndexTarget>()
449 + std::mem::size_of::<Tag>()
450 + std::mem::size_of::<BloomFilterMeta>();
451 let expected_dynamic = meta.segment_loc_indices.len() * std::mem::size_of::<u64>()
452 + meta.bloom_filter_locs.len() * std::mem::size_of::<BloomFilterLoc>();
453
454 assert_eq!(weight as usize, base + expected_dynamic);
455 }
456
457 #[test]
458 fn fuzz_index_calculation() {
459 use common_cache::page::PageLayout;
460 use std::num::NonZeroU64;
461
462 let mut rng = rand::rng();
463 let mut data = vec![0u8; 1024 * 1024];
464 rng.fill_bytes(&mut data);
465 for _ in 0..FUZZ_REPEAT_TIMES {
466 let offset = rng.random_range(0..data.len() as u64);
467 let size = rng.random_range(0..data.len() as u64 - offset);
468 let page_size = NonZeroU64::new(rng.random_range(1_u64..1024)).unwrap();
469 let layout = PageLayout::new(page_size, data.len() as u64);
470 let requested = offset..offset + size;
471 let mut read = Vec::with_capacity(size as usize);
472 for id in layout.page_ids(requested.clone()).unwrap() {
473 let range = layout.page_range(id).unwrap();
474 let page = &data[range.start as usize..range.end as usize];
475 let slice = layout
476 .request_slice(id, requested.clone())
477 .unwrap()
478 .unwrap();
479 read.extend_from_slice(&page[slice]);
480 }
481 assert_eq!(read, data[offset as usize..(offset + size) as usize]);
482 }
483 }
484}