Skip to main content

foyer_memory/
raw.rs

1// Copyright 2026 foyer Project Authors
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
15use std::{
16    any::Any,
17    fmt::Debug,
18    future::Future,
19    hash::Hash,
20    ops::Deref,
21    pin::Pin,
22    sync::{
23        Arc,
24        atomic::{AtomicBool, Ordering},
25    },
26    task::{Context, Poll},
27};
28
29use equivalent::Equivalent;
30use foyer_common::{
31    code::HashBuilder,
32    error::{Error, ErrorKind, Result},
33    event::{Event, EventListener},
34    metrics::Metrics,
35    properties::{Location, Properties, Source},
36    spawn::Spawner,
37    strict_assert,
38    utils::scope::Scope,
39};
40use futures_util::FutureExt as _;
41use itertools::Itertools;
42use parking_lot::{Mutex, RwLock};
43use pin_project::{pin_project, pinned_drop};
44
45use crate::{
46    Piece,
47    eviction::{Eviction, Op},
48    indexer::{Indexer, hash_table::HashTableIndexer, sentry::Sentry},
49    inflight::{
50        Enqueue, FetchOrTake, FetchTarget, InflightManager, Notifier, OptionalFetch, OptionalFetchBuilder,
51        RequiredFetch, RequiredFetchBuilder, Waiter,
52    },
53    pipe::{ArcPipe, NoopPipe},
54    record::{Data, Record},
55};
56
57/// The weighter for the in-memory cache.
58///
59/// The weighter is used to calculate the weight of the cache entry.
60pub trait Weighter<K, V>: Fn(&K, &V) -> usize + Send + Sync + 'static {}
61impl<K, V, T> Weighter<K, V> for T where T: Fn(&K, &V) -> usize + Send + Sync + 'static {}
62
63/// The filter for the in-memory cache.
64///
65/// The filter is used to decide whether to admit or reject an entry based on its key and value.
66///
67/// If the filter returns true, the key value can be inserted into the in-memory cache;
68/// otherwise, the key value cannot be inserted.
69///
70/// To ensure API consistency, the in-memory cache will still return a cache entry,
71/// but it will not count towards the in-memory cache usage,
72/// and it will be immediately reclaimed when the cache entry is dropped.
73pub trait Filter<K, V>: Fn(&K, &V) -> bool + Send + Sync + 'static {}
74impl<K, V, T> Filter<K, V> for T where T: Fn(&K, &V) -> bool + Send + Sync + 'static {}
75
76pub struct RawCacheConfig<E, S>
77where
78    E: Eviction,
79    S: HashBuilder,
80{
81    pub capacity: usize,
82    pub shards: usize,
83    pub eviction_config: E::Config,
84    pub hash_builder: S,
85    pub weighter: Arc<dyn Weighter<E::Key, E::Value>>,
86    pub filter: Arc<dyn Filter<E::Key, E::Value>>,
87    pub event_listener: Option<Arc<dyn EventListener<Key = E::Key, Value = E::Value>>>,
88    pub metrics: Arc<Metrics>,
89}
90
91struct RawCacheShard<E, S, I>
92where
93    E: Eviction,
94    S: HashBuilder,
95    I: Indexer<Eviction = E>,
96{
97    eviction: E,
98    indexer: Sentry<I>,
99
100    usage: usize,
101    entries: usize,
102    capacity: usize,
103
104    inflights: Arc<Mutex<InflightManager<E, S, I>>>,
105
106    metrics: Arc<Metrics>,
107    _event_listener: Option<Arc<dyn EventListener<Key = E::Key, Value = E::Value>>>,
108}
109
110impl<E, S, I> RawCacheShard<E, S, I>
111where
112    E: Eviction,
113    S: HashBuilder,
114    I: Indexer<Eviction = E>,
115{
116    /// Evict entries to fit the target usage.
117    fn evict(&mut self, target: usize, garbages: &mut Vec<(Event, Arc<Record<E>>)>) {
118        // Evict overflow records.
119        while self.usage > target {
120            let evicted = match self.eviction.pop() {
121                Some(evicted) => evicted,
122                None => break,
123            };
124            self.metrics.memory_evict.increase(1);
125
126            let e = self.indexer.remove(evicted.hash(), evicted.key()).unwrap();
127            assert_eq!(Arc::as_ptr(&evicted), Arc::as_ptr(&e));
128
129            strict_assert!(!evicted.as_ref().is_in_indexer());
130            strict_assert!(!evicted.as_ref().is_in_eviction());
131
132            self.usage -= evicted.weight();
133            self.entries -= 1;
134            self.metrics.memory_entries.decrease(1);
135
136            garbages.push((Event::Evict, evicted));
137        }
138    }
139
140    #[expect(clippy::type_complexity)]
141    fn emplace(
142        &mut self,
143        record: Arc<Record<E>>,
144        garbages: &mut Vec<(Event, Arc<Record<E>>)>,
145        notifiers: &mut Vec<Notifier<Option<RawCacheEntry<E, S, I>>>>,
146    ) {
147        *notifiers = self
148            .inflights
149            .lock()
150            .take(record.hash(), record.key(), None)
151            .unwrap_or_default();
152
153        if record.properties().phantom().unwrap_or_default() {
154            if let Some(old) = self.indexer.remove(record.hash(), record.key()) {
155                strict_assert!(!old.is_in_indexer());
156
157                if old.is_in_eviction() {
158                    self.eviction.remove(&old);
159                }
160                strict_assert!(!old.is_in_eviction());
161
162                self.usage -= old.weight();
163                self.entries -= 1;
164                self.metrics.memory_entries.decrease(1);
165
166                garbages.push((Event::Replace, old));
167            }
168            record.inc_refs(notifiers.len() + 1);
169            garbages.push((Event::Remove, record));
170            self.metrics.memory_insert.increase(1);
171            return;
172        }
173
174        let weight = record.weight();
175        let old_usage = self.usage;
176
177        // Evict overflow records.
178        self.evict(self.capacity.saturating_sub(weight), garbages);
179
180        // Insert new record
181        if let Some(old) = self.indexer.insert(record.clone()) {
182            self.metrics.memory_replace.increase(1);
183
184            strict_assert!(!old.is_in_indexer());
185
186            if old.is_in_eviction() {
187                self.eviction.remove(&old);
188            }
189            strict_assert!(!old.is_in_eviction());
190
191            self.usage -= old.weight();
192
193            garbages.push((Event::Replace, old));
194        } else {
195            self.metrics.memory_insert.increase(1);
196            self.entries += 1;
197            self.metrics.memory_entries.increase(1);
198        }
199        strict_assert!(record.is_in_indexer());
200
201        strict_assert!(!record.is_in_eviction());
202        self.eviction.push(record.clone());
203        strict_assert!(record.is_in_eviction());
204
205        self.usage += weight;
206        // Increase the reference count within the lock section.
207        // The reference count of the new record must be at the moment.
208        record.inc_refs(notifiers.len() + 1);
209
210        match self.usage.cmp(&old_usage) {
211            std::cmp::Ordering::Greater => self.metrics.memory_usage.increase((self.usage - old_usage) as _),
212            std::cmp::Ordering::Less => self.metrics.memory_usage.decrease((old_usage - self.usage) as _),
213            std::cmp::Ordering::Equal => {}
214        }
215    }
216
217    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::shard::remove"))]
218    fn remove<Q>(&mut self, hash: u64, key: &Q) -> Option<Arc<Record<E>>>
219    where
220        Q: Hash + Equivalent<E::Key> + ?Sized,
221    {
222        let record = self.indexer.remove(hash, key)?;
223
224        if record.is_in_eviction() {
225            self.eviction.remove(&record);
226        }
227        strict_assert!(!record.is_in_indexer());
228        strict_assert!(!record.is_in_eviction());
229
230        self.usage -= record.weight();
231        self.entries -= 1;
232
233        self.metrics.memory_remove.increase(1);
234        self.metrics.memory_usage.decrease(record.weight() as _);
235        self.metrics.memory_entries.decrease(1);
236
237        record.inc_refs(1);
238
239        Some(record)
240    }
241
242    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::shard::get_noop"))]
243    fn get_noop<Q>(&self, hash: u64, key: &Q) -> Option<Arc<Record<E>>>
244    where
245        Q: Hash + Equivalent<E::Key> + ?Sized,
246    {
247        self.get_inner(hash, key)
248    }
249
250    #[cfg_attr(
251        feature = "tracing",
252        fastrace::trace(name = "foyer::memory::raw::shard::get_immutable")
253    )]
254    fn get_immutable<Q>(&self, hash: u64, key: &Q) -> Option<Arc<Record<E>>>
255    where
256        Q: Hash + Equivalent<E::Key> + ?Sized,
257    {
258        self.get_inner(hash, key)
259            .inspect(|record| self.acquire_immutable(record))
260    }
261
262    #[cfg_attr(
263        feature = "tracing",
264        fastrace::trace(name = "foyer::memory::raw::shard::get_mutable")
265    )]
266    fn get_mutable<Q>(&mut self, hash: u64, key: &Q) -> Option<Arc<Record<E>>>
267    where
268        Q: Hash + Equivalent<E::Key> + ?Sized,
269    {
270        self.get_inner(hash, key).inspect(|record| self.acquire_mutable(record))
271    }
272
273    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::shard::get_inner"))]
274    fn get_inner<Q>(&self, hash: u64, key: &Q) -> Option<Arc<Record<E>>>
275    where
276        Q: Hash + Equivalent<E::Key> + ?Sized,
277    {
278        let record = match self.indexer.get(hash, key).cloned() {
279            Some(record) => {
280                self.metrics.memory_hit.increase(1);
281                record
282            }
283            None => {
284                self.metrics.memory_miss.increase(1);
285                return None;
286            }
287        };
288
289        strict_assert!(record.is_in_indexer());
290
291        record.inc_refs(1);
292
293        Some(record)
294    }
295
296    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::shard::clear"))]
297    fn clear(&mut self, garbages: &mut Vec<Arc<Record<E>>>) {
298        let records = self.indexer.drain().collect_vec();
299        self.eviction.clear();
300
301        let mut count = 0;
302
303        for record in records {
304            count += 1;
305            strict_assert!(!record.is_in_indexer());
306            strict_assert!(!record.is_in_eviction());
307
308            garbages.push(record);
309        }
310
311        self.entries = 0;
312        if count > 0 {
313            self.metrics.memory_entries.decrease(count);
314            self.metrics.memory_remove.increase(count);
315        }
316    }
317
318    #[cfg_attr(
319        feature = "tracing",
320        fastrace::trace(name = "foyer::memory::raw::shard::acquire_immutable")
321    )]
322    fn acquire_immutable(&self, record: &Arc<Record<E>>) {
323        match E::acquire() {
324            Op::Immutable(f) => f(&self.eviction, record),
325            _ => unreachable!(),
326        }
327    }
328
329    #[cfg_attr(
330        feature = "tracing",
331        fastrace::trace(name = "foyer::memory::raw::shard::acquire_mutable")
332    )]
333    fn acquire_mutable(&mut self, record: &Arc<Record<E>>) {
334        match E::acquire() {
335            Op::Mutable(mut f) => f(&mut self.eviction, record),
336            _ => unreachable!(),
337        }
338    }
339
340    #[cfg_attr(
341        feature = "tracing",
342        fastrace::trace(name = "foyer::memory::raw::shard::release_immutable")
343    )]
344    fn release_immutable(&self, record: &Arc<Record<E>>) {
345        match E::release() {
346            Op::Immutable(f) => f(&self.eviction, record),
347            _ => unreachable!(),
348        }
349    }
350
351    #[cfg_attr(
352        feature = "tracing",
353        fastrace::trace(name = "foyer::memory::raw::shard::release_mutable")
354    )]
355    fn release_mutable(&mut self, record: &Arc<Record<E>>) {
356        match E::release() {
357            Op::Mutable(mut f) => f(&mut self.eviction, record),
358            _ => unreachable!(),
359        }
360    }
361}
362
363impl<E, S> RawCacheShard<E, S, HashTableIndexer<E>>
364where
365    E: Eviction,
366    S: HashBuilder,
367{
368    fn evict_if<F>(&mut self, predicate: &mut F, garbages: &mut Vec<(Event, Arc<Record<E>>)>)
369    where
370        F: FnMut(&E::Key, &E::Value) -> bool + ?Sized,
371    {
372        let Self {
373            eviction,
374            indexer,
375            usage,
376            entries,
377            metrics,
378            ..
379        } = self;
380
381        for record in indexer.extract_if(|record| predicate(record.key(), record.value())) {
382            if record.is_in_eviction() {
383                eviction.remove(&record);
384            }
385            strict_assert!(!record.is_in_indexer());
386            strict_assert!(!record.is_in_eviction());
387
388            *usage -= record.weight();
389            *entries -= 1;
390            metrics.memory_evict.increase(1);
391            metrics.memory_usage.decrease(record.weight() as _);
392            metrics.memory_entries.decrease(1);
393
394            garbages.push((Event::Evict, record));
395        }
396    }
397}
398
399struct RawCacheInner<E, S, I>
400where
401    E: Eviction,
402    S: HashBuilder,
403    I: Indexer<Eviction = E>,
404{
405    shards: Vec<RwLock<RawCacheShard<E, S, I>>>,
406
407    capacity: usize,
408
409    hash_builder: Arc<S>,
410    weighter: Arc<dyn Weighter<E::Key, E::Value>>,
411    filter: Arc<dyn Filter<E::Key, E::Value>>,
412
413    metrics: Arc<Metrics>,
414    event_listener: Option<Arc<dyn EventListener<Key = E::Key, Value = E::Value>>>,
415}
416
417impl<E, S, I> RawCacheInner<E, S, I>
418where
419    E: Eviction,
420    S: HashBuilder,
421    I: Indexer<Eviction = E>,
422{
423    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::inner::clear"))]
424    fn clear(&self) {
425        let mut garbages = vec![];
426
427        self.shards
428            .iter()
429            .map(|shard| shard.write())
430            .for_each(|mut shard| shard.clear(&mut garbages));
431
432        // Do not deallocate data within the lock section.
433        if let Some(listener) = self.event_listener.as_ref() {
434            for record in garbages {
435                listener.on_leave(Event::Clear, record.key(), record.value());
436            }
437        }
438    }
439}
440
441pub struct RawCache<E, S, I>
442where
443    E: Eviction,
444    S: HashBuilder,
445    I: Indexer<Eviction = E>,
446{
447    pipe: ArcPipe<E::Key, E::Value, E::Properties>,
448    inner: Arc<RawCacheInner<E, S, I>>,
449}
450
451impl<E, S, I> Clone for RawCache<E, S, I>
452where
453    E: Eviction,
454    S: HashBuilder,
455    I: Indexer<Eviction = E>,
456{
457    fn clone(&self) -> Self {
458        Self {
459            pipe: self.pipe.clone(),
460            inner: self.inner.clone(),
461        }
462    }
463}
464
465impl<E, S, I> Drop for RawCacheInner<E, S, I>
466where
467    E: Eviction,
468    S: HashBuilder,
469    I: Indexer<Eviction = E>,
470{
471    fn drop(&mut self) {
472        self.clear();
473    }
474}
475
476impl<E, S, I> RawCache<E, S, I>
477where
478    E: Eviction,
479    S: HashBuilder,
480    I: Indexer<Eviction = E>,
481{
482    pub fn new(config: RawCacheConfig<E, S>) -> Self {
483        assert!(config.shards > 0, "shards must be greater than zero.");
484
485        let shard_capacities = (0..config.shards)
486            .map(|index| Self::shard_capacity_for(config.capacity, config.shards, index))
487            .collect_vec();
488
489        let shards = shard_capacities
490            .into_iter()
491            .map(|shard_capacity| RawCacheShard {
492                eviction: E::new(shard_capacity, &config.eviction_config),
493                indexer: Sentry::default(),
494                usage: 0,
495                entries: 0,
496                capacity: shard_capacity,
497                inflights: Arc::new(Mutex::new(InflightManager::new())),
498                metrics: config.metrics.clone(),
499                _event_listener: config.event_listener.clone(),
500            })
501            .map(RwLock::new)
502            .collect_vec();
503
504        Self {
505            pipe: Arc::new(NoopPipe::default()),
506            inner: Arc::new(RawCacheInner {
507                shards,
508                capacity: config.capacity,
509                hash_builder: Arc::new(config.hash_builder),
510                weighter: config.weighter,
511                filter: config.filter,
512                metrics: config.metrics,
513                event_listener: config.event_listener,
514            }),
515        }
516    }
517
518    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::resize"))]
519    pub fn resize(&self, capacity: usize) -> Result<()> {
520        let shards = self.inner.shards.len();
521        assert!(shards > 0, "shards must be greater than zero.");
522
523        let shard_capacities = (0..shards)
524            .map(|index| Self::shard_capacity_for(capacity, shards, index))
525            .collect_vec();
526
527        let handles = shard_capacities
528            .into_iter()
529            .enumerate()
530            .map(|(i, shard_capacity)| {
531                let pipe = self.pipe.clone();
532                let inner = self.inner.clone();
533                std::thread::spawn(move || {
534                    let mut garbages = vec![];
535                    let res = inner.shards[i].write().with(|mut shard| {
536                        shard.eviction.update(shard_capacity, None).inspect(|_| {
537                            shard.capacity = shard_capacity;
538                            shard.evict(shard_capacity, &mut garbages)
539                        })
540                    });
541                    // Deallocate data out of the lock critical section.
542                    let piped = pipe.is_enabled();
543                    if inner.event_listener.is_some() || piped {
544                        for (event, record) in garbages {
545                            if let Some(listener) = inner.event_listener.as_ref() {
546                                listener.on_leave(event, record.key(), record.value())
547                            }
548                            if piped && event == Event::Evict {
549                                pipe.send(Piece::new(record));
550                            }
551                        }
552                    }
553                    res
554                })
555            })
556            .collect_vec();
557
558        let errs = handles
559            .into_iter()
560            .map(|handle| handle.join().unwrap())
561            .filter(|res| res.is_err())
562            .map(|res| res.unwrap_err())
563            .collect_vec();
564        if !errs.is_empty() {
565            let mut e = Error::new(ErrorKind::Config, "resize raw cache failed");
566            for err in errs {
567                e = e.with_context("reason", format!("{err}"));
568            }
569            return Err(e);
570        }
571
572        Ok(())
573    }
574
575    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::insert"))]
576    pub fn insert(&self, key: E::Key, value: E::Value) -> RawCacheEntry<E, S, I> {
577        self.insert_with_properties(key, value, Default::default())
578    }
579
580    #[cfg_attr(
581        feature = "tracing",
582        fastrace::trace(name = "foyer::memory::raw::insert_with_properties")
583    )]
584    pub fn insert_with_properties(
585        &self,
586        key: E::Key,
587        value: E::Value,
588        properties: E::Properties,
589    ) -> RawCacheEntry<E, S, I> {
590        self.insert_with_properties_inner(key, value, properties, Source::Outer)
591    }
592
593    fn insert_with_properties_inner(
594        &self,
595        key: E::Key,
596        value: E::Value,
597        mut properties: E::Properties,
598        source: Source,
599    ) -> RawCacheEntry<E, S, I> {
600        let hash = self.inner.hash_builder.hash_one(&key);
601        let weight = (self.inner.weighter)(&key, &value);
602        if !(self.inner.filter)(&key, &value) {
603            properties = properties.with_phantom(true);
604        }
605        if let Some(location) = properties.location()
606            && location == Location::OnDisk
607        {
608            properties = properties.with_phantom(true);
609        }
610        let record = Arc::new(Record::new(Data {
611            key,
612            value,
613            properties,
614            hash,
615            weight,
616        }));
617        self.insert_inner(record, source)
618    }
619
620    #[doc(hidden)]
621    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::insert_piece"))]
622    pub fn insert_piece(&self, piece: Piece<E::Key, E::Value, E::Properties>) -> RawCacheEntry<E, S, I> {
623        self.insert_inner(piece.into_record(), Source::Memory)
624    }
625
626    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::insert_inner"))]
627    fn insert_inner(&self, record: Arc<Record<E>>, source: Source) -> RawCacheEntry<E, S, I> {
628        let mut garbages = vec![];
629        let mut notifiers = vec![];
630
631        self.inner.shards[self.shard(record.hash())]
632            .write()
633            .with(|mut shard| shard.emplace(record.clone(), &mut garbages, &mut notifiers));
634
635        // Notify waiters out of the lock critical section.
636        for notifier in notifiers {
637            let _ = notifier.send(Ok(Some(RawCacheEntry {
638                pipe: self.pipe.clone(),
639                record: record.clone(),
640                inner: self.inner.clone(),
641                source,
642            })));
643        }
644
645        // Deallocate data out of the lock critical section.
646        let piped = self.pipe.is_enabled();
647        if self.inner.event_listener.is_some() || piped {
648            for (event, record) in garbages {
649                if let Some(listener) = self.inner.event_listener.as_ref() {
650                    listener.on_leave(event, record.key(), record.value())
651                }
652                if piped && event == Event::Evict {
653                    self.pipe.send(Piece::new(record));
654                }
655            }
656        }
657
658        RawCacheEntry {
659            record,
660            pipe: self.pipe.clone(),
661            inner: self.inner.clone(),
662            source,
663        }
664    }
665
666    /// Evict all entries in the cache and offload them into the disk cache via the pipe if needed.
667    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::evict_all"))]
668    pub fn evict_all(&self) {
669        let mut garbages = vec![];
670        for shard in self.inner.shards.iter() {
671            shard.write().evict(0, &mut garbages);
672        }
673
674        // Deallocate data out of the lock critical section.
675        let piped = self.pipe.is_enabled();
676        if self.inner.event_listener.is_some() || piped {
677            for (event, record) in garbages {
678                if let Some(listener) = self.inner.event_listener.as_ref() {
679                    listener.on_leave(event, record.key(), record.value())
680                }
681                if piped && event == Event::Evict {
682                    self.pipe.send(Piece::new(record));
683                }
684            }
685        }
686    }
687
688    /// Evict all entries in the cache and offload them into the disk cache via the pipe if needed.
689    ///
690    /// This function obeys the io throttler of the disk cache and make sure all entries will be offloaded.
691    /// Therefore, this function is asynchronous.
692    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::flush"))]
693    pub async fn flush(&self) {
694        let mut garbages = vec![];
695        for shard in self.inner.shards.iter() {
696            shard.write().evict(0, &mut garbages);
697        }
698
699        self.flush_evicted(garbages).await;
700    }
701
702    async fn flush_evicted(&self, garbages: Vec<(Event, Arc<Record<E>>)>) {
703        // Deallocate data out of the lock critical section.
704        let piped = self.pipe.is_enabled();
705
706        if let Some(listener) = self.inner.event_listener.as_ref() {
707            for (event, record) in garbages.iter() {
708                listener.on_leave(*event, record.key(), record.value());
709            }
710        }
711        if piped {
712            let pieces = garbages.into_iter().map(|(_, record)| Piece::new(record)).collect_vec();
713            self.pipe.flush(pieces).await;
714        }
715    }
716
717    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::remove"))]
718    pub fn remove<Q>(&self, key: &Q) -> Option<RawCacheEntry<E, S, I>>
719    where
720        Q: Hash + Equivalent<E::Key> + ?Sized,
721    {
722        let hash = self.inner.hash_builder.hash_one(key);
723
724        self.inner.shards[self.shard(hash)]
725            .write()
726            .with(|mut shard| {
727                shard.remove(hash, key).map(|record| RawCacheEntry {
728                    pipe: self.pipe.clone(),
729                    inner: self.inner.clone(),
730                    record,
731                    source: Source::Memory,
732                })
733            })
734            .inspect(|record| {
735                // Deallocate data out of the lock critical section.
736                if let Some(listener) = self.inner.event_listener.as_ref() {
737                    listener.on_leave(Event::Remove, record.key(), record.value());
738                }
739            })
740    }
741
742    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::get"))]
743    pub fn get<Q>(&self, key: &Q) -> Option<RawCacheEntry<E, S, I>>
744    where
745        Q: Hash + Equivalent<E::Key> + ?Sized,
746    {
747        let hash = self.inner.hash_builder.hash_one(key);
748
749        let record = match E::acquire() {
750            Op::Noop => self.inner.shards[self.shard(hash)].read().get_noop(hash, key),
751            Op::Immutable(_) => self.inner.shards[self.shard(hash)]
752                .read()
753                .with(|shard| shard.get_immutable(hash, key)),
754            Op::Mutable(_) => self.inner.shards[self.shard(hash)]
755                .write()
756                .with(|mut shard| shard.get_mutable(hash, key)),
757        }?;
758
759        Some(RawCacheEntry {
760            pipe: self.pipe.clone(),
761            inner: self.inner.clone(),
762            record,
763            source: Source::Memory,
764        })
765    }
766
767    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::contains"))]
768    pub fn contains<Q>(&self, key: &Q) -> bool
769    where
770        Q: Hash + Equivalent<E::Key> + ?Sized,
771    {
772        let hash = self.inner.hash_builder.hash_one(key);
773
774        self.inner.shards[self.shard(hash)]
775            .read()
776            .with(|shard| shard.indexer.get(hash, key).is_some())
777    }
778
779    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::touch"))]
780    pub fn touch<Q>(&self, key: &Q) -> bool
781    where
782        Q: Hash + Equivalent<E::Key> + ?Sized,
783    {
784        let hash = self.inner.hash_builder.hash_one(key);
785
786        match E::acquire() {
787            Op::Noop => self.inner.shards[self.shard(hash)].read().get_noop(hash, key),
788            Op::Immutable(_) => self.inner.shards[self.shard(hash)]
789                .read()
790                .with(|shard| shard.get_immutable(hash, key)),
791            Op::Mutable(_) => self.inner.shards[self.shard(hash)]
792                .write()
793                .with(|mut shard| shard.get_mutable(hash, key)),
794        }
795        .is_some()
796    }
797
798    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::clear"))]
799    pub fn clear(&self) {
800        self.inner.clear();
801    }
802
803    pub fn capacity(&self) -> usize {
804        self.inner.capacity
805    }
806
807    pub fn usage(&self) -> usize {
808        self.inner.shards.iter().map(|shard| shard.read().usage).sum()
809    }
810
811    pub fn entries(&self) -> usize {
812        self.inner.shards.iter().map(|shard| shard.read().entries).sum()
813    }
814
815    pub fn metrics(&self) -> &Metrics {
816        &self.inner.metrics
817    }
818
819    pub fn hash_builder(&self) -> &Arc<S> {
820        &self.inner.hash_builder
821    }
822
823    pub fn shards(&self) -> usize {
824        self.inner.shards.len()
825    }
826
827    pub(crate) fn with_pipe(mut self, pipe: ArcPipe<E::Key, E::Value, E::Properties>) -> Self {
828        self.pipe = pipe;
829        self
830    }
831
832    fn shard(&self, hash: u64) -> usize {
833        hash as usize % self.inner.shards.len()
834    }
835
836    fn shard_capacity_for(total: usize, shards: usize, index: usize) -> usize {
837        let base = total / shards;
838        let remainder = total % shards;
839        base + usize::from(index < remainder)
840    }
841}
842
843impl<E, S> RawCache<E, S, HashTableIndexer<E>>
844where
845    E: Eviction,
846    S: HashBuilder,
847{
848    /// Evict entries matching the predicate and offload them into the disk cache via the pipe if needed.
849    ///
850    /// This function obeys the io throttler of the disk cache and makes sure all matching entries are offloaded.
851    ///
852    /// The predicate is called while holding a shard lock and must not access this cache.
853    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::flush_if"))]
854    pub async fn flush_if<F>(&self, mut predicate: F)
855    where
856        F: FnMut(&E::Key, &E::Value) -> bool,
857    {
858        let mut garbages = vec![];
859        for shard in self.inner.shards.iter() {
860            shard.write().evict_if(&mut predicate, &mut garbages);
861        }
862
863        drop(predicate);
864        self.flush_evicted(garbages).await;
865    }
866}
867
868pub struct RawCacheEntry<E, S, I>
869where
870    E: Eviction,
871    S: HashBuilder,
872    I: Indexer<Eviction = E>,
873{
874    pipe: ArcPipe<E::Key, E::Value, E::Properties>,
875    inner: Arc<RawCacheInner<E, S, I>>,
876    record: Arc<Record<E>>,
877    source: Source,
878}
879
880impl<E, S, I> Debug for RawCacheEntry<E, S, I>
881where
882    E: Eviction,
883    S: HashBuilder,
884    I: Indexer<Eviction = E>,
885{
886    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
887        f.debug_struct("RawCacheEntry").field("record", &self.record).finish()
888    }
889}
890
891impl<E, S, I> Drop for RawCacheEntry<E, S, I>
892where
893    E: Eviction,
894    S: HashBuilder,
895    I: Indexer<Eviction = E>,
896{
897    fn drop(&mut self) {
898        let hash = self.record.hash();
899        let shard = &self.inner.shards[hash as usize % self.inner.shards.len()];
900
901        if self.record.dec_refs(1) == 0 {
902            if self.record.properties().phantom().unwrap_or_default() {
903                if let Some(listener) = self.inner.event_listener.as_ref() {
904                    listener.on_leave(Event::Evict, self.record.key(), self.record.value());
905                }
906                if self.pipe.is_enabled() {
907                    self.pipe.send(Piece::new(self.record.clone()));
908                }
909                return;
910            }
911
912            match E::release() {
913                Op::Noop => {}
914                Op::Immutable(_) => shard.read().with(|shard| shard.release_immutable(&self.record)),
915                Op::Mutable(_) => shard.write().with(|mut shard| shard.release_mutable(&self.record)),
916            }
917        }
918    }
919}
920
921impl<E, S, I> Clone for RawCacheEntry<E, S, I>
922where
923    E: Eviction,
924    S: HashBuilder,
925    I: Indexer<Eviction = E>,
926{
927    fn clone(&self) -> Self {
928        self.record.inc_refs(1);
929        Self {
930            pipe: self.pipe.clone(),
931            inner: self.inner.clone(),
932            record: self.record.clone(),
933            source: self.source,
934        }
935    }
936}
937
938impl<E, S, I> Deref for RawCacheEntry<E, S, I>
939where
940    E: Eviction,
941    S: HashBuilder,
942    I: Indexer<Eviction = E>,
943{
944    type Target = E::Value;
945
946    fn deref(&self) -> &Self::Target {
947        self.value()
948    }
949}
950
951unsafe impl<E, S, I> Send for RawCacheEntry<E, S, I>
952where
953    E: Eviction,
954    S: HashBuilder,
955    I: Indexer<Eviction = E>,
956{
957}
958
959unsafe impl<E, S, I> Sync for RawCacheEntry<E, S, I>
960where
961    E: Eviction,
962    S: HashBuilder,
963    I: Indexer<Eviction = E>,
964{
965}
966
967impl<E, S, I> RawCacheEntry<E, S, I>
968where
969    E: Eviction,
970    S: HashBuilder,
971    I: Indexer<Eviction = E>,
972{
973    pub fn hash(&self) -> u64 {
974        self.record.hash()
975    }
976
977    pub fn key(&self) -> &E::Key {
978        self.record.key()
979    }
980
981    pub fn value(&self) -> &E::Value {
982        self.record.value()
983    }
984
985    pub fn properties(&self) -> &E::Properties {
986        self.record.properties()
987    }
988
989    pub fn weight(&self) -> usize {
990        self.record.weight()
991    }
992
993    pub fn refs(&self) -> usize {
994        self.record.refs()
995    }
996
997    pub fn is_outdated(&self) -> bool {
998        !self.record.is_in_indexer()
999    }
1000
1001    pub fn piece(&self) -> Piece<E::Key, E::Value, E::Properties> {
1002        Piece::new(self.record.clone())
1003    }
1004
1005    pub fn source(&self) -> Source {
1006        self.source
1007    }
1008}
1009
1010impl<E, S, I> RawCache<E, S, I>
1011where
1012    E: Eviction,
1013    S: HashBuilder,
1014    I: Indexer<Eviction = E>,
1015{
1016    #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::get_or_fetch"))]
1017    pub fn get_or_fetch<Q, F, FU, IT, ER>(&self, key: &Q, fetch: F) -> RawGetOrFetch<E, S, I>
1018    where
1019        Q: Hash + Equivalent<E::Key> + ?Sized + ToOwned<Owned = E::Key>,
1020        F: FnOnce() -> FU,
1021        FU: Future<Output = std::result::Result<IT, ER>> + Send + 'static,
1022        IT: Into<FetchTarget<E::Key, E::Value, E::Properties>>,
1023        ER: Into<anyhow::Error>,
1024    {
1025        let fut = fetch();
1026        self.get_or_fetch_inner(
1027            key,
1028            || None,
1029            || {
1030                Some(Box::new(|_| {
1031                    async {
1032                        match fut.await {
1033                            Ok(it) => Ok(it.into()),
1034                            Err(e) => Err(Error::new(ErrorKind::External, "fetch failed").with_source(e)),
1035                        }
1036                    }
1037                    .boxed()
1038                }))
1039            },
1040            (),
1041            &Spawner::current(),
1042        )
1043    }
1044
1045    /// Advanced fetch with specified runtime.
1046    ///
1047    /// This function is for internal usage and the doc is hidden.
1048    #[doc(hidden)]
1049    #[cfg_attr(
1050        feature = "tracing",
1051        fastrace::trace(name = "foyer::memory::raw::get_or_fetch_inner")
1052    )]
1053    pub fn get_or_fetch_inner<Q, C, FO, FR>(
1054        &self,
1055        key: &Q,
1056        fo: FO,
1057        fr: FR,
1058        ctx: C,
1059        spawner: &Spawner,
1060    ) -> RawGetOrFetch<E, S, I>
1061    where
1062        Q: Hash + Equivalent<E::Key> + ?Sized + ToOwned<Owned = E::Key>,
1063        C: Any + Send + Sync + 'static,
1064        FO: FnOnce() -> Option<OptionalFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1065        FR: FnOnce() -> Option<RequiredFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1066    {
1067        let hash = self.inner.hash_builder.hash_one(key);
1068
1069        // Make sure cache query and inflight query are in the same lock critical section.
1070        let extract = |key: &Q, opt: Option<Arc<Record<E>>>, inflights: &Arc<Mutex<InflightManager<E, S, I>>>| {
1071            opt.map(|record| {
1072                RawGetOrFetch::Hit(Some(RawCacheEntry {
1073                    pipe: self.pipe.clone(),
1074                    inner: self.inner.clone(),
1075                    record,
1076                    source: Source::Memory,
1077                }))
1078            })
1079            .unwrap_or_else(|| match inflights.lock().enqueue(hash, key, fr()) {
1080                Enqueue::Lead {
1081                    id,
1082                    close,
1083                    waiter,
1084                    required_fetch_builder,
1085                } => {
1086                    let fetch = RawFetch {
1087                        state: RawFetchState::Init {
1088                            optional_fetch_builder: fo(),
1089                            required_fetch_builder,
1090                        },
1091                        id,
1092                        hash,
1093                        key: Some(key.to_owned()),
1094                        ctx,
1095                        cache: self.clone(),
1096                        inflights: inflights.clone(),
1097                        close,
1098                    };
1099                    spawner.spawn(fetch);
1100                    RawGetOrFetch::Miss(RawWait { waiter })
1101                }
1102                Enqueue::Wait(waiter) => RawGetOrFetch::Miss(RawWait { waiter }),
1103            })
1104        };
1105
1106        match E::acquire() {
1107            Op::Noop => self.inner.shards[self.shard(hash)]
1108                .read()
1109                .with(|shard| extract(key, shard.get_noop(hash, key), &shard.inflights)),
1110            Op::Immutable(_) => self.inner.shards[self.shard(hash)]
1111                .read()
1112                .with(|shard| extract(key, shard.get_immutable(hash, key), &shard.inflights)),
1113            Op::Mutable(_) => self.inner.shards[self.shard(hash)]
1114                .write()
1115                .with(|mut shard| extract(key, shard.get_mutable(hash, key), &shard.inflights)),
1116        }
1117    }
1118}
1119
1120#[must_use]
1121#[pin_project(project = RawGetOrFetchProj)]
1122pub enum RawGetOrFetch<E, S, I>
1123where
1124    E: Eviction,
1125    S: HashBuilder,
1126    I: Indexer<Eviction = E>,
1127{
1128    Hit(Option<RawCacheEntry<E, S, I>>),
1129    Miss(#[pin] RawWait<E, S, I>),
1130}
1131
1132impl<E, S, I> Debug for RawGetOrFetch<E, S, I>
1133where
1134    E: Eviction,
1135    S: HashBuilder,
1136    I: Indexer<Eviction = E>,
1137{
1138    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1139        match self {
1140            Self::Hit(e) => f.debug_tuple("Hit").field(e).finish(),
1141            Self::Miss(fut) => f.debug_tuple("Miss").field(fut).finish(),
1142        }
1143    }
1144}
1145
1146impl<E, S, I> Future for RawGetOrFetch<E, S, I>
1147where
1148    E: Eviction,
1149    S: HashBuilder,
1150    I: Indexer<Eviction = E>,
1151{
1152    type Output = Result<Option<RawCacheEntry<E, S, I>>>;
1153
1154    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1155        let this = self.project();
1156        match this {
1157            RawGetOrFetchProj::Hit(opt) => {
1158                assert!(opt.is_some(), "entry is already taken");
1159                Poll::Ready(Ok(opt.take()))
1160            }
1161            RawGetOrFetchProj::Miss(fut) => fut.poll(cx),
1162        }
1163    }
1164}
1165
1166impl<E, S, I> RawGetOrFetch<E, S, I>
1167where
1168    E: Eviction,
1169    S: HashBuilder,
1170    I: Indexer<Eviction = E>,
1171{
1172    pub fn need_await(&self) -> bool {
1173        matches!(self, Self::Miss(_))
1174    }
1175
1176    #[expect(clippy::allow_attributes)]
1177    #[allow(clippy::result_large_err)]
1178    pub fn try_unwrap(self) -> std::result::Result<Option<RawCacheEntry<E, S, I>>, Self> {
1179        match self {
1180            Self::Hit(opt) => {
1181                assert!(opt.is_some(), "entry is already taken");
1182                Ok(opt)
1183            }
1184            Self::Miss(_) => Err(self),
1185        }
1186    }
1187}
1188
1189type Once<T> = Option<T>;
1190
1191#[must_use]
1192enum Try<E, S, I, C>
1193where
1194    E: Eviction,
1195    S: HashBuilder,
1196    I: Indexer<Eviction = E>,
1197    C: Any + Send + 'static,
1198{
1199    Noop,
1200    SetStateAndContinue(RawFetchState<E, S, I, C>),
1201    Ready,
1202}
1203
1204macro_rules! handle_try {
1205    ($state:expr, $method:ident($($args:expr),* $(,)?)) => {
1206        handle_try! { $state, Self::$method($($args),*) }
1207    };
1208
1209    ($state:expr, $try:expr) => {
1210        match $try {
1211            Try::Noop => {}
1212            Try::SetStateAndContinue(state) => {
1213                $state = state;
1214                continue;
1215            },
1216            Try::Ready => {
1217                $state = RawFetchState::Ready;
1218                return Poll::Ready(())
1219            },
1220        }
1221    };
1222}
1223
1224#[expect(clippy::type_complexity)]
1225pub enum RawFetchState<E, S, I, C>
1226where
1227    E: Eviction,
1228    S: HashBuilder,
1229    I: Indexer<Eviction = E>,
1230{
1231    Init {
1232        optional_fetch_builder: Option<OptionalFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1233        required_fetch_builder: Option<RequiredFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1234    },
1235    FetchOptional {
1236        optional_fetch: OptionalFetch<FetchTarget<E::Key, E::Value, E::Properties>>,
1237        required_fetch_builder: Option<RequiredFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1238    },
1239    FetchRequired {
1240        required_fetch: RequiredFetch<FetchTarget<E::Key, E::Value, E::Properties>>,
1241    },
1242    Notify {
1243        res: Option<Result<Option<RawCacheEntry<E, S, I>>>>,
1244        notifiers: Vec<Notifier<Option<RawCacheEntry<E, S, I>>>>,
1245    },
1246    Ready,
1247}
1248
1249impl<E, S, I, C> Debug for RawFetchState<E, S, I, C>
1250where
1251    E: Eviction,
1252    S: HashBuilder,
1253    I: Indexer<Eviction = E>,
1254{
1255    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1256        match self {
1257            Self::Init { .. } => f.debug_struct("Init").finish(),
1258            Self::FetchOptional { .. } => f.debug_struct("Optional").finish(),
1259            Self::FetchRequired { .. } => f.debug_struct("Required").finish(),
1260            Self::Notify { res, .. } => f.debug_struct("Notify").field("res", res).finish(),
1261            Self::Ready => f.debug_struct("Ready").finish(),
1262        }
1263    }
1264}
1265
1266#[pin_project]
1267pub struct RawWait<E, S, I>
1268where
1269    E: Eviction,
1270    S: HashBuilder,
1271    I: Indexer<Eviction = E>,
1272{
1273    waiter: Waiter<Option<RawCacheEntry<E, S, I>>>,
1274}
1275
1276impl<E, S, I> Debug for RawWait<E, S, I>
1277where
1278    E: Eviction,
1279    S: HashBuilder,
1280    I: Indexer<Eviction = E>,
1281{
1282    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1283        f.debug_struct("RawWait").field("waiter", &self.waiter).finish()
1284    }
1285}
1286
1287impl<E, S, I> Future for RawWait<E, S, I>
1288where
1289    E: Eviction,
1290    S: HashBuilder,
1291    I: Indexer<Eviction = E>,
1292{
1293    type Output = Result<Option<RawCacheEntry<E, S, I>>>;
1294
1295    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1296        let this = self.project();
1297        // TODO(MrCroxx): Switch to `Result::flatten` after MSRV is 1.89+
1298        // return waiter.poll_unpin(cx).map(|r| r.map_err(|e| e.into()).flatten());
1299        this.waiter.poll_unpin(cx).map(|r| match r {
1300            Ok(r) => r,
1301            Err(e) => Err(Error::new(ErrorKind::ChannelClosed, "waiter channel closed").with_source(e)),
1302        })
1303    }
1304}
1305
1306#[pin_project(PinnedDrop)]
1307pub struct RawFetch<E, S, I, C>
1308where
1309    E: Eviction,
1310    S: HashBuilder,
1311    I: Indexer<Eviction = E>,
1312{
1313    state: RawFetchState<E, S, I, C>,
1314    id: usize,
1315    hash: u64,
1316    key: Once<E::Key>,
1317    ctx: C,
1318    cache: RawCache<E, S, I>,
1319    inflights: Arc<Mutex<InflightManager<E, S, I>>>,
1320    close: Arc<AtomicBool>,
1321}
1322
1323impl<E, S, I, C> Debug for RawFetch<E, S, I, C>
1324where
1325    E: Eviction,
1326    S: HashBuilder,
1327    I: Indexer<Eviction = E>,
1328{
1329    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1330        f.debug_struct("RawFetch")
1331            .field("state", &self.state)
1332            .field("id", &self.id)
1333            .field("hash", &self.hash)
1334            .finish()
1335    }
1336}
1337
1338impl<E, S, I, C> Future for RawFetch<E, S, I, C>
1339where
1340    E: Eviction,
1341    S: HashBuilder,
1342    I: Indexer<Eviction = E>,
1343    C: Any + Send + 'static,
1344{
1345    type Output = ();
1346
1347    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1348        let this = self.as_mut().project();
1349        loop {
1350            match this.state {
1351                RawFetchState::Init {
1352                    optional_fetch_builder,
1353                    required_fetch_builder,
1354                } => {
1355                    handle_try! { *this.state, try_set_optional(optional_fetch_builder, required_fetch_builder, this.ctx) }
1356                    handle_try! { *this.state, try_set_required(required_fetch_builder, this.ctx, *this.id, *this.hash, this.key.as_ref().unwrap(), this.inflights, Ok(None)) }
1357                }
1358                RawFetchState::FetchOptional {
1359                    optional_fetch,
1360                    required_fetch_builder,
1361                } => {
1362                    if this.close.load(Ordering::Relaxed) {
1363                        return Poll::Ready(());
1364                    }
1365                    match optional_fetch.poll_unpin(cx) {
1366                        Poll::Pending => return Poll::Pending,
1367                        Poll::Ready(Ok(Some(target))) => {
1368                            handle_try! {*this.state, handle_target(target, this.key, this.cache, Source::Disk) }
1369                        }
1370                        Poll::Ready(Ok(None)) => {
1371                            handle_try! { *this.state, try_set_required(required_fetch_builder, this.ctx, *this.id, *this.hash, this.key.as_ref().unwrap(), &this.inflights, Ok(None)) }
1372                        }
1373                        Poll::Ready(Err(e)) => {
1374                            handle_try! { *this.state, try_set_required(required_fetch_builder, this.ctx, *this.id, *this.hash, this.key.as_ref().unwrap(), &this.inflights, Err(e)) }
1375                        }
1376                    }
1377                }
1378                RawFetchState::FetchRequired { required_fetch } => {
1379                    if this.close.load(Ordering::Relaxed) {
1380                        return Poll::Ready(());
1381                    }
1382                    match required_fetch.poll_unpin(cx) {
1383                        Poll::Pending => return Poll::Pending,
1384                        Poll::Ready(Ok(target)) => {
1385                            handle_try! { *this.state, handle_target(target, this.key, this.cache, Source::Outer) }
1386                        }
1387                        Poll::Ready(Err(e)) => {
1388                            handle_try! { *this.state, handle_error(e, *this.id, *this.hash, this.key.as_ref().unwrap(), this.inflights) }
1389                        }
1390                    }
1391                }
1392                RawFetchState::Notify { res, notifiers } => {
1393                    handle_try! { *this.state, handle_notify(res.take().unwrap(), notifiers) }
1394                }
1395                RawFetchState::Ready => return Poll::Ready(()),
1396            }
1397        }
1398    }
1399}
1400
1401impl<E, S, I, C> RawFetch<E, S, I, C>
1402where
1403    E: Eviction,
1404    S: HashBuilder,
1405    I: Indexer<Eviction = E>,
1406    C: Any + Send + 'static,
1407{
1408    #[expect(clippy::type_complexity)]
1409    fn try_set_optional(
1410        optional_fetch_builder: &mut Option<OptionalFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1411        required_fetch_builder: &mut Option<RequiredFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1412        ctx: &mut C,
1413    ) -> Try<E, S, I, C> {
1414        match optional_fetch_builder.take() {
1415            None => Try::Noop,
1416            Some(optional_fetch_builder) => {
1417                let optional_fetch = optional_fetch_builder(ctx);
1418                Try::SetStateAndContinue(RawFetchState::FetchOptional {
1419                    optional_fetch,
1420                    required_fetch_builder: required_fetch_builder.take(),
1421                })
1422            }
1423        }
1424    }
1425
1426    #[expect(clippy::type_complexity)]
1427    fn try_set_required(
1428        required_fetch_builder: &mut Option<RequiredFetchBuilder<E::Key, E::Value, E::Properties, C>>,
1429        ctx: &mut C,
1430        id: usize,
1431        hash: u64,
1432        key: &E::Key,
1433        inflights: &Arc<Mutex<InflightManager<E, S, I>>>,
1434        res_no_fetch: Result<Option<RawCacheEntry<E, S, I>>>,
1435    ) -> Try<E, S, I, C> {
1436        // Fast path if the required fetch builder is provided.
1437        match required_fetch_builder.take() {
1438            None => {}
1439            Some(required_fetch_builder) => {
1440                let required_fetch = required_fetch_builder(ctx);
1441                return Try::SetStateAndContinue(RawFetchState::FetchRequired { required_fetch });
1442            }
1443        }
1444        // Slow path if the leader has no optional fetch.
1445        let fetch_or_take = match inflights.lock().fetch_or_take(hash, key, id) {
1446            Some(fetch_or_take) => fetch_or_take,
1447            None => return Try::Ready,
1448        };
1449        match fetch_or_take {
1450            FetchOrTake::Fetch(required_fetch_builder) => {
1451                let required_fetch = required_fetch_builder(ctx);
1452                Try::SetStateAndContinue(RawFetchState::FetchRequired { required_fetch })
1453            }
1454            FetchOrTake::Notifiers(notifiers) => Try::SetStateAndContinue(RawFetchState::Notify {
1455                res: Some(res_no_fetch),
1456                notifiers,
1457            }),
1458        }
1459    }
1460
1461    fn handle_target(
1462        target: FetchTarget<E::Key, E::Value, E::Properties>,
1463        key: &mut Once<E::Key>,
1464        cache: &RawCache<E, S, I>,
1465        source: Source,
1466    ) -> Try<E, S, I, C> {
1467        match target {
1468            FetchTarget::Entry { value, properties } => {
1469                let key = key.take().unwrap();
1470                cache.insert_with_properties_inner(key, value, properties, source);
1471            }
1472            FetchTarget::Piece(piece) => {
1473                cache.insert_piece(piece);
1474            }
1475        }
1476        Try::Ready
1477    }
1478
1479    fn handle_error(
1480        e: Error,
1481        id: usize,
1482        hash: u64,
1483        key: &E::Key,
1484        inflights: &Arc<Mutex<InflightManager<E, S, I>>>,
1485    ) -> Try<E, S, I, C> {
1486        let notifiers = match inflights.lock().take(hash, key, Some(id)) {
1487            Some(notifiers) => notifiers,
1488            None => {
1489                return Try::Ready;
1490            }
1491        };
1492        Try::SetStateAndContinue(RawFetchState::Notify {
1493            res: Some(Err(e)),
1494            notifiers,
1495        })
1496    }
1497
1498    #[expect(clippy::type_complexity)]
1499    fn handle_notify(
1500        res: Result<Option<RawCacheEntry<E, S, I>>>,
1501        notifiers: &mut Vec<Notifier<Option<RawCacheEntry<E, S, I>>>>,
1502    ) -> Try<E, S, I, C> {
1503        match res {
1504            Ok(e) => {
1505                for notifier in notifiers.drain(..) {
1506                    let _ = notifier.send(Ok(e.clone()));
1507                }
1508            }
1509            Err(e) => {
1510                for notifier in notifiers.drain(..) {
1511                    let _ = notifier.send(Err(e.clone()));
1512                }
1513            }
1514        }
1515        Try::Ready
1516    }
1517}
1518
1519#[pinned_drop]
1520impl<E, S, I, C> PinnedDrop for RawFetch<E, S, I, C>
1521where
1522    E: Eviction,
1523    S: HashBuilder,
1524    I: Indexer<Eviction = E>,
1525{
1526    fn drop(self: Pin<&mut Self>) {
1527        let this = self.project();
1528        match this.state {
1529            RawFetchState::Notify { .. } | RawFetchState::Ready => return,
1530            RawFetchState::Init { .. } | RawFetchState::FetchOptional { .. } | RawFetchState::FetchRequired { .. } => {}
1531        }
1532        if let Some(notifiers) = this
1533            .inflights
1534            .lock()
1535            .take(*this.hash, this.key.as_ref().unwrap(), Some(*this.id))
1536        {
1537            for notifier in notifiers {
1538                let _ =
1539                    notifier
1540                        .send(Err(Error::new(ErrorKind::TaskCancelled, "fetch task cancelled")
1541                            .with_context("hash", *this.hash)));
1542            }
1543        }
1544    }
1545}
1546
1547#[cfg(test)]
1548mod tests {
1549    use foyer_common::hasher::ModHasher;
1550    use rand::{Rng, SeedableRng, rngs::SmallRng, seq::IndexedRandom};
1551
1552    use super::*;
1553    use crate::{
1554        eviction::{
1555            fifo::{Fifo, FifoConfig},
1556            lfu::{Lfu, LfuConfig},
1557            lru::{Lru, LruConfig},
1558            s3fifo::{S3Fifo, S3FifoConfig},
1559            sieve::{Sieve, SieveConfig},
1560            test_utils::TestProperties,
1561        },
1562        indexer::hash_table::HashTableIndexer,
1563        test_utils::PiecePipe,
1564    };
1565
1566    fn is_send_sync_static<T: Send + Sync + 'static>() {}
1567
1568    #[test]
1569    fn test_send_sync_static() {
1570        is_send_sync_static::<RawCache<Fifo<(), (), TestProperties>, ModHasher, HashTableIndexer<_>>>();
1571        is_send_sync_static::<RawCache<S3Fifo<(), (), TestProperties>, ModHasher, HashTableIndexer<_>>>();
1572        is_send_sync_static::<RawCache<Lfu<(), (), TestProperties>, ModHasher, HashTableIndexer<_>>>();
1573        is_send_sync_static::<RawCache<Lru<(), (), TestProperties>, ModHasher, HashTableIndexer<_>>>();
1574        is_send_sync_static::<RawCache<Sieve<(), (), TestProperties>, ModHasher, HashTableIndexer<_>>>();
1575    }
1576
1577    #[expect(clippy::type_complexity)]
1578    fn fifo_cache_for_test()
1579    -> RawCache<Fifo<u64, u64, TestProperties>, ModHasher, HashTableIndexer<Fifo<u64, u64, TestProperties>>> {
1580        RawCache::new(RawCacheConfig {
1581            capacity: 256,
1582            shards: 4,
1583            eviction_config: FifoConfig::default(),
1584            hash_builder: Default::default(),
1585            weighter: Arc::new(|_, _| 1),
1586            filter: Arc::new(|_, _| true),
1587            event_listener: None,
1588            metrics: Arc::new(Metrics::noop()),
1589        })
1590    }
1591
1592    #[expect(clippy::type_complexity)]
1593    fn s3fifo_cache_for_test()
1594    -> RawCache<S3Fifo<u64, u64, TestProperties>, ModHasher, HashTableIndexer<S3Fifo<u64, u64, TestProperties>>> {
1595        RawCache::new(RawCacheConfig {
1596            capacity: 256,
1597            shards: 4,
1598            eviction_config: S3FifoConfig::default(),
1599            hash_builder: Default::default(),
1600            weighter: Arc::new(|_, _| 1),
1601            filter: Arc::new(|_, _| true),
1602            event_listener: None,
1603            metrics: Arc::new(Metrics::noop()),
1604        })
1605    }
1606
1607    #[expect(clippy::type_complexity)]
1608    fn lru_cache_for_test()
1609    -> RawCache<Lru<u64, u64, TestProperties>, ModHasher, HashTableIndexer<Lru<u64, u64, TestProperties>>> {
1610        RawCache::new(RawCacheConfig {
1611            capacity: 256,
1612            shards: 4,
1613            eviction_config: LruConfig::default(),
1614            hash_builder: Default::default(),
1615            weighter: Arc::new(|_, _| 1),
1616            filter: Arc::new(|_, _| true),
1617            event_listener: None,
1618            metrics: Arc::new(Metrics::noop()),
1619        })
1620    }
1621
1622    #[expect(clippy::type_complexity)]
1623    fn lfu_cache_for_test()
1624    -> RawCache<Lfu<u64, u64, TestProperties>, ModHasher, HashTableIndexer<Lfu<u64, u64, TestProperties>>> {
1625        RawCache::new(RawCacheConfig {
1626            capacity: 256,
1627            shards: 4,
1628            eviction_config: LfuConfig::default(),
1629            hash_builder: Default::default(),
1630            weighter: Arc::new(|_, _| 1),
1631            filter: Arc::new(|_, _| true),
1632            event_listener: None,
1633            metrics: Arc::new(Metrics::noop()),
1634        })
1635    }
1636
1637    #[expect(clippy::type_complexity)]
1638    fn sieve_cache_for_test()
1639    -> RawCache<Sieve<u64, u64, TestProperties>, ModHasher, HashTableIndexer<Sieve<u64, u64, TestProperties>>> {
1640        RawCache::new(RawCacheConfig {
1641            capacity: 256,
1642            shards: 4,
1643            eviction_config: SieveConfig {},
1644            hash_builder: Default::default(),
1645            weighter: Arc::new(|_, _| 1),
1646            filter: Arc::new(|_, _| true),
1647            event_listener: None,
1648            metrics: Arc::new(Metrics::noop()),
1649        })
1650    }
1651
1652    #[test_log::test]
1653    fn test_insert_phantom() {
1654        let fifo = fifo_cache_for_test();
1655
1656        let e1 = fifo.insert_with_properties(1, 1, TestProperties::default().with_phantom(true));
1657        assert_eq!(fifo.usage(), 0);
1658        drop(e1);
1659        assert_eq!(fifo.usage(), 0);
1660
1661        let e2a = fifo.insert_with_properties(2, 2, TestProperties::default().with_phantom(true));
1662        assert_eq!(fifo.usage(), 0);
1663        assert!(fifo.get(&2).is_none());
1664        assert_eq!(fifo.usage(), 0);
1665        drop(e2a);
1666        assert_eq!(fifo.usage(), 0);
1667
1668        let fifo = fifo_cache_for_test();
1669        fifo.insert(1, 1);
1670        assert_eq!(fifo.usage(), 1);
1671        assert_eq!(fifo.get(&1).unwrap().value(), &1);
1672        let e2 = fifo.insert_with_properties(1, 100, TestProperties::default().with_phantom(true));
1673        assert_eq!(fifo.usage(), 0);
1674        drop(e2);
1675        assert_eq!(fifo.usage(), 0);
1676        assert!(fifo.get(&1).is_none());
1677    }
1678
1679    #[expect(clippy::type_complexity)]
1680    #[test_log::test]
1681    fn test_insert_filter() {
1682        let fifo: RawCache<
1683            Fifo<u64, u64, TestProperties>,
1684            ModHasher,
1685            HashTableIndexer<Fifo<u64, u64, TestProperties>>,
1686        > = RawCache::new(RawCacheConfig {
1687            capacity: 256,
1688            shards: 4,
1689            eviction_config: FifoConfig::default(),
1690            hash_builder: Default::default(),
1691            weighter: Arc::new(|_, _| 1),
1692            filter: Arc::new(|k, _| !matches!(*k, 42)),
1693            event_listener: None,
1694            metrics: Arc::new(Metrics::noop()),
1695        });
1696
1697        fifo.insert(1, 1);
1698        fifo.insert(2, 2);
1699        fifo.insert(42, 42);
1700        assert_eq!(fifo.usage(), 2);
1701        assert_eq!(fifo.get(&1).unwrap().value(), &1);
1702        assert_eq!(fifo.get(&2).unwrap().value(), &2);
1703        assert!(fifo.get(&42).is_none());
1704    }
1705
1706    #[test]
1707    fn test_evict_all() {
1708        let pipe = Arc::new(PiecePipe::default());
1709
1710        let fifo = fifo_cache_for_test().with_pipe(pipe.clone());
1711        for i in 0..fifo.capacity() as _ {
1712            fifo.insert(i, i);
1713        }
1714        assert_eq!(fifo.usage(), fifo.capacity());
1715
1716        fifo.evict_all();
1717        let mut pieces = pipe
1718            .pieces()
1719            .iter()
1720            .map(|p| (p.hash(), *p.key(), *p.value()))
1721            .collect_vec();
1722        pieces.sort_by_key(|t| t.0);
1723        let expected = (0..fifo.capacity() as u64).map(|i| (i, i, i)).collect_vec();
1724        assert_eq!(pieces, expected);
1725    }
1726
1727    async fn assert_flush_if<E>(cache: RawCache<E, ModHasher, HashTableIndexer<E>>)
1728    where
1729        E: Eviction<Key = u64, Value = u64, Properties = TestProperties>,
1730    {
1731        let pipe = Arc::new(PiecePipe::default());
1732        let cache = cache.with_pipe(pipe.clone());
1733
1734        for key in 0..8 {
1735            cache.insert(key, key * 10);
1736        }
1737        let pinned = cache.insert(8, 80);
1738
1739        cache
1740            .flush_if(|key, value| (key % 2 == 0 && *value < 60) || *key == 8)
1741            .await;
1742
1743        let mut pieces = pipe
1744            .pieces()
1745            .iter()
1746            .map(|piece| (*piece.key(), *piece.value()))
1747            .collect_vec();
1748        pieces.sort_unstable();
1749        assert_eq!(pieces, vec![(0, 0), (2, 20), (4, 40), (8, 80)]);
1750        assert_eq!(cache.entries(), 5);
1751        assert_eq!(cache.usage(), 5);
1752        assert!(pinned.is_outdated());
1753
1754        for key in [0, 2, 4, 8] {
1755            assert!(cache.get(&key).is_none());
1756        }
1757        for key in [1, 3, 5, 6, 7] {
1758            assert_eq!(*cache.get(&key).unwrap().value(), key * 10);
1759        }
1760    }
1761
1762    #[tokio::test]
1763    async fn test_flush_if() {
1764        assert_flush_if(fifo_cache_for_test()).await;
1765        assert_flush_if(s3fifo_cache_for_test()).await;
1766        assert_flush_if(lru_cache_for_test()).await;
1767        assert_flush_if(lfu_cache_for_test()).await;
1768        assert_flush_if(sieve_cache_for_test()).await;
1769    }
1770
1771    #[test]
1772    fn test_insert_size_over_capacity() {
1773        let cache: RawCache<Fifo<Vec<u8>, Vec<u8>, TestProperties>, ModHasher, HashTableIndexer<_>> =
1774            RawCache::new(RawCacheConfig {
1775                capacity: 4 * 1024, // 4KB
1776                shards: 1,
1777                eviction_config: FifoConfig::default(),
1778                hash_builder: Default::default(),
1779                weighter: Arc::new(|k, v| k.len() + v.len()),
1780                filter: Arc::new(|_, _| true),
1781                event_listener: None,
1782                metrics: Arc::new(Metrics::noop()),
1783            });
1784
1785        let key = vec![b'k'; 1024]; // 1KB
1786        let value = vec![b'v'; 5 * 1024]; // 5KB
1787
1788        cache.insert(key.clone(), value.clone());
1789        assert_eq!(cache.usage(), 6 * 1024);
1790        assert_eq!(cache.get(&key).unwrap().value(), &value);
1791    }
1792
1793    #[test]
1794    fn test_capacity_distribution_without_loss() {
1795        let cache: RawCache<Fifo<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
1796            RawCache::new(RawCacheConfig {
1797                capacity: 3,
1798                shards: 2,
1799                eviction_config: FifoConfig::default(),
1800                hash_builder: Default::default(),
1801                weighter: Arc::new(|_, _| 1),
1802                filter: Arc::new(|_, _| true),
1803                event_listener: None,
1804                metrics: Arc::new(Metrics::noop()),
1805            });
1806
1807        for key in 0..3 {
1808            let entry = cache.insert(key, key);
1809            drop(entry);
1810        }
1811
1812        assert_eq!(cache.usage(), 3);
1813
1814        for key in 0..3 {
1815            let entry = cache.get(&key).expect("entry should exist");
1816            assert_eq!(*entry, key);
1817            drop(entry);
1818        }
1819    }
1820
1821    #[test]
1822    fn test_capacity_distribution_with_more_shards_than_capacity() {
1823        let cache: RawCache<Fifo<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
1824            RawCache::new(RawCacheConfig {
1825                capacity: 2,
1826                shards: 4,
1827                eviction_config: FifoConfig::default(),
1828                hash_builder: Default::default(),
1829                weighter: Arc::new(|_, _| 1),
1830                filter: Arc::new(|_, _| true),
1831                event_listener: None,
1832                metrics: Arc::new(Metrics::noop()),
1833            });
1834
1835        for key in 0..2 {
1836            let entry = cache.insert(key, key);
1837            drop(entry);
1838        }
1839
1840        assert_eq!(cache.usage(), 2);
1841
1842        for key in 0..2 {
1843            let entry = cache.get(&key).expect("entry should exist");
1844            assert_eq!(*entry, key);
1845            drop(entry);
1846        }
1847
1848        assert!(cache.get(&2).is_none());
1849    }
1850
1851    fn test_resize<E>(cache: &RawCache<E, ModHasher, HashTableIndexer<E>>)
1852    where
1853        E: Eviction<Key = u64, Value = u64>,
1854    {
1855        let capacity = cache.capacity();
1856        for i in 0..capacity as u64 * 2 {
1857            cache.insert(i, i);
1858        }
1859        assert_eq!(cache.usage(), capacity);
1860        cache.resize(capacity / 2).unwrap();
1861        assert_eq!(cache.usage(), capacity / 2);
1862        for i in 0..capacity as u64 * 2 {
1863            cache.insert(i, i);
1864        }
1865        assert_eq!(cache.usage(), capacity / 2);
1866    }
1867
1868    #[test]
1869    fn test_fifo_cache_resize() {
1870        let cache = fifo_cache_for_test();
1871        test_resize(&cache);
1872    }
1873
1874    #[test]
1875    fn test_s3fifo_cache_resize() {
1876        let cache = s3fifo_cache_for_test();
1877        test_resize(&cache);
1878    }
1879
1880    #[test]
1881    fn test_lru_cache_resize() {
1882        let cache = lru_cache_for_test();
1883        test_resize(&cache);
1884    }
1885
1886    #[test]
1887    fn test_lfu_cache_resize() {
1888        let cache = lfu_cache_for_test();
1889        test_resize(&cache);
1890    }
1891
1892    #[test]
1893    fn test_sieve_cache_resize() {
1894        let cache = sieve_cache_for_test();
1895        test_resize(&cache);
1896    }
1897
1898    mod fuzzy {
1899        use foyer_common::properties::Hint;
1900
1901        use super::*;
1902
1903        fn fuzzy<E, S>(cache: RawCache<E, S, HashTableIndexer<E>>, hints: Vec<Hint>)
1904        where
1905            E: Eviction<Key = u64, Value = u64, Properties = TestProperties>,
1906            S: HashBuilder,
1907        {
1908            let handles = (0..8)
1909                .map(|i| {
1910                    let c = cache.clone();
1911                    let hints = hints.clone();
1912                    std::thread::spawn(move || {
1913                        let mut rng = SmallRng::seed_from_u64(i);
1914                        for _ in 0..100000 {
1915                            let key = rng.next_u64();
1916                            if let Some(entry) = c.get(&key) {
1917                                assert_eq!(key, *entry);
1918                                drop(entry);
1919                                continue;
1920                            }
1921                            let hint = hints.choose(&mut rng).cloned().unwrap();
1922                            c.insert_with_properties(key, key, TestProperties::default().with_hint(hint));
1923                        }
1924                    })
1925                })
1926                .collect_vec();
1927
1928            handles.into_iter().for_each(|handle| handle.join().unwrap());
1929
1930            assert_eq!(cache.usage(), cache.capacity());
1931        }
1932
1933        #[test_log::test]
1934        fn test_fifo_cache_fuzzy() {
1935            let cache: RawCache<Fifo<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
1936                RawCache::new(RawCacheConfig {
1937                    capacity: 256,
1938                    shards: 4,
1939                    eviction_config: FifoConfig::default(),
1940                    hash_builder: Default::default(),
1941                    weighter: Arc::new(|_, _| 1),
1942                    filter: Arc::new(|_, _| true),
1943                    event_listener: None,
1944                    metrics: Arc::new(Metrics::noop()),
1945                });
1946            let hints = vec![Hint::Normal];
1947            fuzzy(cache, hints);
1948        }
1949
1950        #[test_log::test]
1951        fn test_s3fifo_cache_fuzzy() {
1952            let cache: RawCache<S3Fifo<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
1953                RawCache::new(RawCacheConfig {
1954                    capacity: 256,
1955                    shards: 4,
1956                    eviction_config: S3FifoConfig::default(),
1957                    hash_builder: Default::default(),
1958                    weighter: Arc::new(|_, _| 1),
1959                    filter: Arc::new(|_, _| true),
1960                    event_listener: None,
1961                    metrics: Arc::new(Metrics::noop()),
1962                });
1963            let hints = vec![Hint::Normal];
1964            fuzzy(cache, hints);
1965        }
1966
1967        #[test_log::test]
1968        fn test_lru_cache_fuzzy() {
1969            let cache: RawCache<Lru<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
1970                RawCache::new(RawCacheConfig {
1971                    capacity: 256,
1972                    shards: 4,
1973                    eviction_config: LruConfig::default(),
1974                    hash_builder: Default::default(),
1975                    weighter: Arc::new(|_, _| 1),
1976                    filter: Arc::new(|_, _| true),
1977                    event_listener: None,
1978                    metrics: Arc::new(Metrics::noop()),
1979                });
1980            let hints = vec![Hint::Normal, Hint::Low];
1981            fuzzy(cache, hints);
1982        }
1983
1984        #[test_log::test]
1985        fn test_lfu_cache_fuzzy() {
1986            let cache: RawCache<Lfu<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
1987                RawCache::new(RawCacheConfig {
1988                    capacity: 256,
1989                    shards: 4,
1990                    eviction_config: LfuConfig::default(),
1991                    hash_builder: Default::default(),
1992                    weighter: Arc::new(|_, _| 1),
1993                    filter: Arc::new(|_, _| true),
1994                    event_listener: None,
1995                    metrics: Arc::new(Metrics::noop()),
1996                });
1997            let hints = vec![Hint::Normal];
1998            fuzzy(cache, hints);
1999        }
2000
2001        #[test_log::test]
2002        fn test_sieve_cache_fuzzy() {
2003            let cache: RawCache<Sieve<u64, u64, TestProperties>, ModHasher, HashTableIndexer<_>> =
2004                RawCache::new(RawCacheConfig {
2005                    capacity: 256,
2006                    shards: 4,
2007                    eviction_config: SieveConfig {},
2008                    hash_builder: Default::default(),
2009                    weighter: Arc::new(|_, _| 1),
2010                    filter: Arc::new(|_, _| true),
2011                    event_listener: None,
2012                    metrics: Arc::new(Metrics::noop()),
2013                });
2014            let hints = vec![Hint::Normal];
2015            fuzzy(cache, hints);
2016        }
2017    }
2018}