mito2/cache/
range_cache.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Coordinates range-result cache loads and SST retirement. No I/O or decoding
16//! runs under the coordinator lock; loaders use the shared query runtime.
17
18use std::collections::HashMap;
19use std::future::Future;
20use std::sync::{Arc, Mutex, Weak, mpsc};
21
22use moka::sync::Cache;
23use store_api::storage::FileId;
24use tokio::sync::{OwnedSemaphorePermit, watch};
25
26use crate::cache::{
27    RANGE_RESULT_TYPE, RangeResultMemoryLimiter, range_result_cache_weight, removal_cause_str,
28};
29use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, RANGE_CACHE_WORKER_EVENTS};
30use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue};
31use crate::sst::file::FileHandle;
32
33/// Admission, not a queue: excess distinct misses keep using ordinary scans.
34const MAX_IN_FLIGHT_RANGES: usize = 32;
35
36type SharedResult = Option<Arc<RangeCacheDelivery>>;
37
38/// Keeps completed results charged while subscribers or replay streams own them.
39/// Cache entries retain only a weak reference: after the last reader leaves,
40/// the resident value is covered solely by the cache's capacity accounting.
41pub(crate) struct RangeCacheDelivery {
42    value: Arc<RangeScanCacheValue>,
43    _permit: OwnedSemaphorePermit,
44}
45
46impl RangeCacheDelivery {
47    /// Shares a decoded result together with its memory reservation.
48    pub(crate) fn new(value: Arc<RangeScanCacheValue>, permit: OwnedSemaphorePermit) -> Arc<Self> {
49        Arc::new(Self {
50            value,
51            _permit: permit,
52        })
53    }
54
55    /// Borrows the result without detaching it from the delivery reservation.
56    pub(crate) fn value(&self) -> &RangeScanCacheValue {
57        &self.value
58    }
59}
60
61pub(crate) enum RangeCacheLookup {
62    Hit(Arc<RangeCacheDelivery>),
63    Loading(watch::Receiver<SharedResult>),
64    Bypass,
65}
66
67struct Entry {
68    generation: u64,
69    value: Arc<RangeScanCacheValue>,
70    delivery: Mutex<Weak<RangeCacheDelivery>>,
71}
72
73impl Entry {
74    fn deliver(&self, memory: &RangeResultMemoryLimiter) -> SharedResult {
75        // Serialize cache-hit reservations so readers of the same result share
76        // one charge, including when the last previous reader just left.
77        let mut shared_delivery = self
78            .delivery
79            .lock()
80            .unwrap_or_else(std::sync::PoisonError::into_inner);
81        if let Some(delivery) = shared_delivery.upgrade() {
82            return Some(delivery);
83        }
84        let permit = memory.try_acquire(self.value.estimated_size())?;
85        let delivery = RangeCacheDelivery::new(self.value.clone(), permit);
86        *shared_delivery = Arc::downgrade(&delivery);
87        Some(delivery)
88    }
89}
90
91struct Flight {
92    generation: u64,
93    sender: watch::Sender<SharedResult>,
94    invalidated: bool,
95}
96
97#[derive(Default)]
98struct State {
99    generation: u64,
100    flights: HashMap<RangeScanCacheKey, Flight>,
101    /// File IDs are globally unique. Index by physical file, not query region:
102    /// a copied region can read an SST whose original region differs from its own.
103    files: HashMap<FileId, HashMap<Arc<RangeScanCacheKey>, u64>>,
104}
105
106impl State {
107    fn next_generation(&mut self) -> u64 {
108        self.generation += 1;
109        self.generation
110    }
111
112    fn forget(&mut self, key: &RangeScanCacheKey, generation: u64) {
113        for (file, _) in &key.row_groups {
114            if let Some(keys) = self.files.get_mut(file) {
115                // Eviction callbacks may arrive after the same key is refilled.
116                if keys.get(key) == Some(&generation) {
117                    keys.remove(key);
118                }
119                if keys.is_empty() {
120                    self.files.remove(file);
121                }
122            }
123        }
124    }
125}
126
127pub(crate) struct RangeCacheWorker {
128    cache: Cache<RangeScanCacheKey, Arc<Entry>>,
129    state: Mutex<State>,
130    /// Moka can call listeners during a cache operation. Queue only cleanup
131    /// metadata instead of reentering `state` from its eviction listener.
132    evictions: Mutex<mpsc::Receiver<(Arc<RangeScanCacheKey>, u64)>>,
133    /// Covers loads and completed deliveries, separately from cache residency
134    /// and the temporary concat-memory limiter. Never wait for this budget.
135    pub(crate) load_memory: RangeResultMemoryLimiter,
136}
137
138impl RangeCacheWorker {
139    pub(crate) fn new(capacity: u64) -> Arc<Self> {
140        let (tx, rx) = mpsc::channel();
141        let cache = Cache::builder()
142            .max_capacity(capacity)
143            .weigher(|k: &RangeScanCacheKey, v: &Arc<Entry>| range_result_cache_weight(k, &v.value))
144            .eviction_listener(move |key, entry: Arc<Entry>, cause| {
145                CACHE_BYTES
146                    .with_label_values(&[RANGE_RESULT_TYPE])
147                    .sub(range_result_cache_weight(&key, &entry.value).into());
148                CACHE_EVICTION
149                    .with_label_values(&[RANGE_RESULT_TYPE, removal_cause_str(cause)])
150                    .inc();
151                let _ = tx.send((key, entry.generation));
152            })
153            .build();
154        Arc::new(Self {
155            cache,
156            state: Mutex::new(State::default()),
157            evictions: Mutex::new(rx),
158            load_memory: RangeResultMemoryLimiter::new(capacity as usize, 1024),
159        })
160    }
161
162    fn drain_evictions(&self, state: &mut State) {
163        for (key, generation) in self.evictions.lock().unwrap().try_iter() {
164            state.forget(&key, generation);
165        }
166    }
167
168    #[cfg(test)]
169    pub(crate) fn get(&self, key: &RangeScanCacheKey) -> Option<Arc<RangeScanCacheValue>> {
170        self.cache.get(key).map(|entry| entry.value.clone())
171    }
172
173    fn get_for_delivery(&self, key: &RangeScanCacheKey) -> Option<RangeCacheLookup> {
174        let entry = self.cache.get(key)?;
175        // A hit without delivery budget must bypass, not start another load.
176        Some(match entry.deliver(&self.load_memory) {
177            Some(delivery) => RangeCacheLookup::Hit(delivery),
178            None => {
179                RANGE_CACHE_WORKER_EVENTS
180                    .with_label_values(&["bypass_memory"])
181                    .inc();
182                RangeCacheLookup::Bypass
183            }
184        })
185    }
186
187    /// For test/benchmark producers that do not hold live SST handles.
188    #[cfg(any(test, feature = "test"))]
189    pub(crate) fn insert(&self, key: RangeScanCacheKey, value: Arc<RangeScanCacheValue>) {
190        let Some(permit) = self.load_memory.try_acquire(value.estimated_size()) else {
191            return;
192        };
193        let delivery = RangeCacheDelivery::new(value, permit);
194        let mut state = self.state.lock().unwrap();
195        self.drain_evictions(&mut state);
196        self.publish(&mut state, key, &delivery);
197    }
198
199    fn publish(
200        &self,
201        state: &mut State,
202        key: RangeScanCacheKey,
203        delivery: &Arc<RangeCacheDelivery>,
204    ) {
205        let generation = state.next_generation();
206        let indexed_key = Arc::new(key.clone());
207        for (file, _) in &key.row_groups {
208            state
209                .files
210                .entry(*file)
211                .or_default()
212                .insert(indexed_key.clone(), generation);
213        }
214        CACHE_BYTES
215            .with_label_values(&[RANGE_RESULT_TYPE])
216            .add(range_result_cache_weight(&key, &delivery.value).into());
217        self.cache.insert(
218            key,
219            Arc::new(Entry {
220                generation,
221                value: delivery.value.clone(),
222                delivery: Mutex::new(Arc::downgrade(delivery)),
223            }),
224        );
225        self.drain_evictions(state);
226    }
227
228    /// Called at logical retirement, before the last snapshot releases the SST.
229    /// Existing readers may finish, but their results must not refill this cache.
230    pub(crate) fn invalidate_file(&self, file: FileId) {
231        RANGE_CACHE_WORKER_EVENTS
232            .with_label_values(&["invalidate"])
233            .inc();
234        let mut state = self.state.lock().unwrap();
235        self.drain_evictions(&mut state);
236        if let Some(keys) = state.files.remove(&file) {
237            for (key, generation) in keys {
238                self.cache.invalidate(key.as_ref());
239                state.forget(&key, generation);
240            }
241        }
242        for (key, flight) in &mut state.flights {
243            if key.row_groups.iter().any(|(id, _)| *id == file) {
244                flight.invalidated = true;
245            }
246        }
247        self.drain_evictions(&mut state);
248    }
249
250    pub(crate) fn get_or_load<F>(
251        self: &Arc<Self>,
252        key: RangeScanCacheKey,
253        files: Vec<FileHandle>,
254        load: F,
255    ) -> RangeCacheLookup
256    where
257        F: Future<Output = SharedResult> + Send + 'static,
258    {
259        if files.iter().any(FileHandle::is_deleted) {
260            RANGE_CACHE_WORKER_EVENTS
261                .with_label_values(&["bypass_retired"])
262                .inc();
263            return RangeCacheLookup::Bypass;
264        }
265        if let Some(lookup) =
266            crate::cache::update_hit_miss(self.get_for_delivery(&key), RANGE_RESULT_TYPE)
267        {
268            return lookup;
269        }
270        let mut state = self.state.lock().unwrap();
271        self.drain_evictions(&mut state);
272        // Close the miss/register race with completion and invalidation.
273        if files.iter().any(FileHandle::is_deleted) {
274            RANGE_CACHE_WORKER_EVENTS
275                .with_label_values(&["bypass_retired"])
276                .inc();
277            return RangeCacheLookup::Bypass;
278        }
279        if let Some(lookup) = self.get_for_delivery(&key) {
280            return lookup;
281        }
282        if let Some(flight) = state.flights.get(&key) {
283            RANGE_CACHE_WORKER_EVENTS
284                .with_label_values(&["coalesced"])
285                .inc();
286            return RangeCacheLookup::Loading(flight.sender.subscribe());
287        }
288        if state.flights.len() >= MAX_IN_FLIGHT_RANGES {
289            RANGE_CACHE_WORKER_EVENTS
290                .with_label_values(&["bypass_capacity"])
291                .inc();
292            return RangeCacheLookup::Bypass;
293        }
294        let generation = state.next_generation();
295        let (sender, receiver) = watch::channel(None);
296        state.flights.insert(
297            key.clone(),
298            Flight {
299                generation,
300                sender: sender.clone(),
301                invalidated: false,
302            },
303        );
304        drop(state);
305
306        RANGE_CACHE_WORKER_EVENTS
307            .with_label_values(&["started"])
308            .inc();
309        let guard = LoadGuard {
310            worker: self.clone(),
311            key,
312            generation,
313        };
314        common_runtime::spawn_query(async move {
315            // Losing the first requester does not cancel other subscribers.
316            // Losing every requester drops the loader and unregisters the flight.
317            let value = tokio::select! {
318                biased;
319                _ = sender.closed() => return,
320                value = load => value,
321            };
322            guard.complete(value, &files);
323        });
324        RangeCacheLookup::Loading(receiver)
325    }
326}
327
328/// Also unregisters a cancelled or panicking load. Generation matching keeps a
329/// stale task from unregistering a later retry of the same key.
330struct LoadGuard {
331    worker: Arc<RangeCacheWorker>,
332    key: RangeScanCacheKey,
333    generation: u64,
334}
335
336impl LoadGuard {
337    fn complete(&self, value: SharedResult, files: &[FileHandle]) {
338        let mut state = self.worker.state.lock().unwrap();
339        self.worker.drain_evictions(&mut state);
340        let Some(flight) = state.flights.get(&self.key) else {
341            return;
342        };
343        if flight.generation != self.generation {
344            return;
345        }
346        let flight = state.flights.remove(&self.key).unwrap();
347        let event = if let Some(value) = &value {
348            if !flight.invalidated && !files.iter().any(FileHandle::is_deleted) {
349                self.worker.publish(&mut state, self.key.clone(), value);
350                "published"
351            } else {
352                "late_fill_skipped"
353            }
354        } else {
355            "load_skipped"
356        };
357        RANGE_CACHE_WORKER_EVENTS.with_label_values(&[event]).inc();
358        // Fan out independently of Moka admission, including to old snapshots.
359        // Do not hold the coordinator lock while notifying readers.
360        drop(state);
361        let _ = flight.sender.send(value);
362    }
363}
364
365impl Drop for LoadGuard {
366    fn drop(&mut self) {
367        let mut state = self.worker.state.lock().unwrap();
368        if state
369            .flights
370            .get(&self.key)
371            .is_some_and(|flight| flight.generation == self.generation)
372        {
373            state.flights.remove(&self.key);
374            RANGE_CACHE_WORKER_EVENTS
375                .with_label_values(&["abandoned"])
376                .inc();
377        }
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use std::sync::atomic::{AtomicUsize, Ordering};
384    use std::time::Duration;
385
386    use tokio::sync::oneshot;
387
388    use super::*;
389    use crate::cache::test_util::range_cache_key;
390    use crate::sst::file::FileMeta;
391    use crate::sst::file_purger::NoopFilePurger;
392
393    fn empty_result() -> Arc<RangeScanCacheValue> {
394        Arc::new(RangeScanCacheValue::new(vec![], 0))
395    }
396
397    fn empty_delivery() -> Arc<RangeCacheDelivery> {
398        let value = empty_result();
399        let memory = RangeResultMemoryLimiter::new(1024, 1024);
400        let permit = memory.try_acquire(value.estimated_size()).unwrap();
401        RangeCacheDelivery::new(value, permit)
402    }
403
404    fn receiver(lookup: RangeCacheLookup) -> watch::Receiver<SharedResult> {
405        match lookup {
406            RangeCacheLookup::Loading(receiver) => receiver,
407            _ => panic!("expected a shared load"),
408        }
409    }
410
411    async fn loaded(mut receiver: watch::Receiver<SharedResult>) -> SharedResult {
412        tokio::time::timeout(Duration::from_secs(5), async {
413            receiver
414                .wait_for(|value| value.is_some())
415                .await
416                .ok()
417                .and_then(|value| value.clone())
418        })
419        .await
420        .expect("load must finish or cancel")
421    }
422
423    #[tokio::test]
424    async fn same_key_misses_share_one_load_even_after_first_reader_leaves() {
425        let worker = RangeCacheWorker::new(1024 * 1024);
426        let key = range_cache_key(&[FileId::random()]);
427        let count = Arc::new(AtomicUsize::new(0));
428        let (release, gate) = oneshot::channel();
429        let first_count = count.clone();
430        let first = receiver(worker.get_or_load(key.clone(), vec![], async move {
431            first_count.fetch_add(1, Ordering::Relaxed);
432            gate.await.unwrap();
433            Some(empty_delivery())
434        }));
435        let mut followers = Vec::new();
436        for _ in 0..16 {
437            let count = count.clone();
438            followers.push(receiver(worker.get_or_load(
439                key.clone(),
440                vec![],
441                async move {
442                    count.fetch_add(1, Ordering::Relaxed);
443                    Some(empty_delivery())
444                },
445            )));
446        }
447        drop(first);
448        release.send(()).unwrap();
449        let mut results = Vec::new();
450        for follower in followers {
451            results.push(loaded(follower).await.unwrap());
452        }
453        assert_eq!(count.load(Ordering::Relaxed), 1);
454        assert!(results.iter().all(|value| Arc::ptr_eq(value, &results[0])));
455        assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &results[0].value));
456    }
457
458    #[test]
459    fn retirement_removes_all_dependent_keys_but_preserves_other_files() {
460        let worker = RangeCacheWorker::new(1024 * 1024);
461        let a = FileId::random();
462        let b = FileId::random();
463        let combined = range_cache_key(&[a, b, a]);
464        let mut copied = range_cache_key(&[a]);
465        copied.region_id = store_api::storage::RegionId::new(2, 1);
466        let unaffected = range_cache_key(&[b]);
467        for key in [&combined, &copied, &unaffected] {
468            worker.insert(key.clone(), empty_result());
469        }
470        worker.invalidate_file(a);
471        worker.invalidate_file(a);
472        assert!(worker.get(&combined).is_none());
473        assert!(worker.get(&copied).is_none());
474        assert!(worker.get(&unaffected).is_some());
475        worker.invalidate_file(b);
476        assert!(worker.get(&unaffected).is_none());
477        assert!(worker.state.lock().unwrap().files.is_empty());
478    }
479
480    #[test]
481    fn delayed_replacement_eviction_cannot_hide_a_new_entry_from_invalidation() {
482        let worker = RangeCacheWorker::new(1024 * 1024);
483        let file = FileId::random();
484        let key = range_cache_key(&[file]);
485        worker.insert(key.clone(), empty_result());
486        worker.cache.run_pending_tasks();
487        let replacement = empty_result();
488        worker.insert(key.clone(), replacement.clone());
489        worker.cache.run_pending_tasks();
490        assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &replacement));
491        worker.invalidate_file(file);
492        assert!(worker.get(&key).is_none());
493    }
494
495    #[tokio::test]
496    async fn retired_inflight_result_serves_old_readers_without_refilling() {
497        let worker = RangeCacheWorker::new(1024 * 1024);
498        let file = FileId::random();
499        let key = range_cache_key(&[file]);
500        let (release, gate) = oneshot::channel();
501        let pending = receiver(worker.get_or_load(key.clone(), vec![], async move {
502            gate.await.unwrap();
503            Some(empty_delivery())
504        }));
505        worker.invalidate_file(file);
506        release.send(()).unwrap();
507        assert!(loaded(pending).await.is_some());
508        assert!(worker.get(&key).is_none());
509    }
510
511    #[test]
512    fn old_snapshot_cannot_start_a_fill_after_retirement() {
513        let worker = RangeCacheWorker::new(1024 * 1024);
514        let file = FileHandle::new(FileMeta::default(), Arc::new(NoopFilePurger));
515        file.mark_deleted();
516        let key = range_cache_key(&[file.file_id().file_id()]);
517        assert!(matches!(
518            worker.get_or_load(key, vec![file], async { Some(empty_delivery()) }),
519            RangeCacheLookup::Bypass
520        ));
521    }
522
523    #[rstest::rstest]
524    #[tokio::test]
525    async fn index_replacement_retirement_invalidates_and_blocks_old_reader_fills(
526        #[values(false, true)] gc_enabled: bool,
527    ) {
528        use common_base::readable_size::ReadableSize;
529        use store_api::region_engine::RegionEngine;
530        use store_api::region_request::RegionRequest;
531        use store_api::storage::RegionId;
532
533        use crate::config::MitoConfig;
534        use crate::gc::GcConfig;
535        use crate::sst::file::IndexType;
536        use crate::sst::version::SstVersion;
537        use crate::test_util::{CreateRequestBuilder, TestEnv};
538
539        let mut env = TestEnv::new().await;
540        let engine = env
541            .create_engine(MitoConfig {
542                range_result_cache_size: ReadableSize::mb(1),
543                gc: GcConfig {
544                    enable: gc_enabled,
545                    ..Default::default()
546                },
547                ..Default::default()
548            })
549            .await;
550        let region_id = RegionId::new(1, 1);
551        engine
552            .handle_request(
553                region_id,
554                RegionRequest::Create(CreateRequestBuilder::new().build()),
555            )
556            .await
557            .unwrap();
558        let region = engine.get_region(region_id).unwrap();
559        let cache = engine.cache_manager();
560        let worker = cache.range_result_cache.as_ref().unwrap();
561        let meta = FileMeta {
562            region_id,
563            file_id: FileId::random(),
564            ..Default::default()
565        };
566        let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test());
567        version.add_files(region.file_purger.clone(), [meta.clone()].into_iter());
568        let old_reader = version.levels()[0].files[&meta.file_id].clone();
569        let key = range_cache_key(&[meta.file_id]);
570        worker.insert(key.clone(), empty_result());
571        assert!(worker.get(&key).is_some());
572        let loading_key = key.clone().for_cross_range_merge();
573        let (release, gate) = oneshot::channel();
574        let pending = receiver(worker.get_or_load(
575            loading_key.clone(),
576            vec![old_reader.clone()],
577            async move {
578                gate.await.unwrap();
579                Some(empty_delivery())
580            },
581        ));
582        let updated = FileMeta {
583            index_version: 1,
584            available_indexes: smallvec::smallvec![IndexType::InvertedIndex],
585            ..meta
586        };
587        version.add_files(region.file_purger.clone(), [updated.clone()].into_iter());
588        version.remove_files([updated].into_iter());
589
590        // R1-C1: neither a pre-existing flight nor a new request through the old
591        // index handle may refill a range after the current handle is retired.
592        assert!(old_reader.is_deleted());
593        assert!(worker.get(&key).is_none());
594        assert!(matches!(
595            worker.get_or_load(key, vec![old_reader], async { Some(empty_delivery()) }),
596            RangeCacheLookup::Bypass
597        ));
598        release.send(()).unwrap();
599        assert!(loaded(pending).await.is_some());
600        assert!(worker.get(&loading_key).is_none());
601    }
602
603    #[tokio::test]
604    async fn admission_rejection_does_not_discard_the_shared_result() {
605        let worker = RangeCacheWorker::new(1);
606        let key = range_cache_key(&[FileId::random()]);
607        let value = empty_result();
608        let permit = worker
609            .load_memory
610            .try_acquire(value.estimated_size())
611            .unwrap();
612        let delivery = RangeCacheDelivery::new(value, permit);
613        let pending =
614            receiver(worker.get_or_load(key.clone(), vec![], async move { Some(delivery) }));
615        let delivery = loaded(pending).await.unwrap();
616        worker.cache.run_pending_tasks();
617        assert!(worker.get(&key).is_none());
618        {
619            let mut state = worker.state.lock().unwrap();
620            worker.drain_evictions(&mut state);
621            assert!(state.files.is_empty());
622        }
623        assert_eq!(worker.load_memory.available_permits(), 0);
624        drop(delivery);
625        // Publishing wakes readers before the producer drops its watch sender.
626        // Its LoadGuard also needs the coordinator lock to finish cleaning up.
627        tokio::time::timeout(Duration::from_secs(5), async {
628            while worker.load_memory.available_permits() != 1 {
629                tokio::task::yield_now().await;
630            }
631        })
632        .await
633        .expect("a rejected value must release its reservation after all owners leave");
634    }
635
636    #[test]
637    fn cache_hits_share_delivery_budget_until_the_last_reader_leaves() {
638        let worker = RangeCacheWorker::new(1024);
639        let file = FileId::random();
640        let key = range_cache_key(&[file]);
641        worker.insert(key.clone(), empty_result());
642        assert_eq!(worker.load_memory.available_permits(), 1);
643
644        let blocker = worker.load_memory.try_acquire(1024).unwrap();
645        assert!(matches!(
646            worker.get_or_load(key.clone(), vec![], async {
647                panic!("must not load on a hit")
648            }),
649            RangeCacheLookup::Bypass
650        ));
651        assert!(worker.get(&key).is_some());
652        drop(blocker);
653
654        let RangeCacheLookup::Hit(first) =
655            worker.get_or_load(key.clone(), vec![], async { panic!("must hit the cache") })
656        else {
657            panic!("expected an accounted cache hit");
658        };
659        let RangeCacheLookup::Hit(second) =
660            worker.get_or_load(key.clone(), vec![], async { panic!("must reuse delivery") })
661        else {
662            panic!("readers of the same value must share its reservation");
663        };
664        assert!(Arc::ptr_eq(&first, &second));
665        assert_eq!(worker.load_memory.available_permits(), 0);
666
667        worker.invalidate_file(file);
668        worker.cache.run_pending_tasks();
669        assert!(worker.get(&key).is_none());
670        drop(first);
671        assert_eq!(worker.load_memory.available_permits(), 0);
672        drop(second);
673        assert_eq!(worker.load_memory.available_permits(), 1);
674    }
675
676    #[tokio::test]
677    async fn failed_load_releases_followers_and_allows_a_retry() {
678        let worker = RangeCacheWorker::new(1024 * 1024);
679        let key = range_cache_key(&[FileId::random()]);
680        let failed = receiver(worker.get_or_load(key.clone(), vec![], async { None }));
681        assert!(loaded(failed).await.is_none());
682        let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) }));
683        assert!(loaded(retry).await.is_some());
684    }
685
686    #[tokio::test]
687    async fn abandoning_all_subscribers_cancels_and_unregisters_the_load() {
688        let worker = RangeCacheWorker::new(1024 * 1024);
689        let key = range_cache_key(&[FileId::random()]);
690        let (started, start_rx) = oneshot::channel();
691        let marker = Arc::new(());
692        let weak_marker = Arc::downgrade(&marker);
693        let pending = receiver(worker.get_or_load(key.clone(), vec![], async move {
694            let _marker = marker;
695            started.send(()).unwrap();
696            std::future::pending().await
697        }));
698        start_rx.await.unwrap();
699        drop(pending);
700        tokio::time::timeout(Duration::from_secs(5), async {
701            while weak_marker.upgrade().is_some()
702                || !worker.state.lock().unwrap().flights.is_empty()
703            {
704                tokio::task::yield_now().await;
705            }
706        })
707        .await
708        .expect("cancelled loaders must release resources and the flight key");
709        assert!(worker.get(&key).is_none());
710        let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) }));
711        assert!(loaded(retry).await.is_some());
712    }
713
714    #[tokio::test]
715    async fn distinct_loads_are_bounded_without_queueing_misses() {
716        let worker = RangeCacheWorker::new(1024 * 1024);
717        let mut subscribers = Vec::new();
718        for _ in 0..MAX_IN_FLIGHT_RANGES {
719            subscribers.push(receiver(worker.get_or_load(
720                range_cache_key(&[FileId::random()]),
721                vec![],
722                std::future::pending(),
723            )));
724        }
725        assert!(matches!(
726            worker.get_or_load(range_cache_key(&[FileId::random()]), vec![], async {
727                Some(empty_delivery())
728            }),
729            RangeCacheLookup::Bypass
730        ));
731        drop(subscribers);
732    }
733}
mito2/cache/range_cache.rs

mito2/cache/
range_cache.rs
