async_graphql/dataloader/
cache.rs

1use std::{
2    borrow::Cow,
3    collections::{hash_map::RandomState, HashMap},
4    hash::{BuildHasher, Hash},
5    marker::PhantomData,
6    num::NonZeroUsize,
7};
8
9/// Factory for creating cache storage.
10pub trait CacheFactory: Send + Sync + 'static {
11    /// Create a cache storage.
12    ///
13    /// TODO: When GAT is stable, this memory allocation can be optimized away.
14    fn create<K, V>(&self) -> Box<dyn CacheStorage<Key = K, Value = V>>
15    where
16        K: Send + Sync + Clone + Eq + Hash + 'static,
17        V: Send + Sync + Clone + 'static;
18}
19
20/// Cache storage for [DataLoader](crate::dataloader::DataLoader).
21pub trait CacheStorage: Send + Sync + 'static {
22    /// The key type of the record.
23    type Key: Send + Sync + Clone + Eq + Hash + 'static;
24
25    /// The value type of the record.
26    type Value: Send + Sync + Clone + 'static;
27
28    /// Returns a reference to the value of the key in the cache or None if it
29    /// is not present in the cache.
30    fn get(&mut self, key: &Self::Key) -> Option<&Self::Value>;
31
32    /// Puts a key-value pair into the cache. If the key already exists in the
33    /// cache, then it updates the key's value.
34    fn insert(&mut self, key: Cow<'_, Self::Key>, val: Cow<'_, Self::Value>);
35
36    /// Removes the value corresponding to the key from the cache.
37    fn remove(&mut self, key: &Self::Key);
38
39    /// Clears the cache, removing all key-value pairs.
40    fn clear(&mut self);
41
42    /// Returns an iterator over the key-value pairs in the cache.
43    fn iter(&self) -> Box<dyn Iterator<Item = (&'_ Self::Key, &'_ Self::Value)> + '_>;
44}
45
46/// No cache.
47pub struct NoCache;
48
49impl CacheFactory for NoCache {
50    fn create<K, V>(&self) -> Box<dyn CacheStorage<Key = K, Value = V>>
51    where
52        K: Send + Sync + Clone + Eq + Hash + 'static,
53        V: Send + Sync + Clone + 'static,
54    {
55        Box::new(NoCacheImpl {
56            _mark1: PhantomData,
57            _mark2: PhantomData,
58        })
59    }
60}
61
62struct NoCacheImpl<K, V> {
63    _mark1: PhantomData<K>,
64    _mark2: PhantomData<V>,
65}
66
67impl<K, V> CacheStorage for NoCacheImpl<K, V>
68where
69    K: Send + Sync + Clone + Eq + Hash + 'static,
70    V: Send + Sync + Clone + 'static,
71{
72    type Key = K;
73    type Value = V;
74
75    #[inline]
76    fn get(&mut self, _key: &K) -> Option<&V> {
77        None
78    }
79
80    #[inline]
81    fn insert(&mut self, _key: Cow<'_, Self::Key>, _val: Cow<'_, Self::Value>) {}
82
83    #[inline]
84    fn remove(&mut self, _key: &K) {}
85
86    #[inline]
87    fn clear(&mut self) {}
88
89    fn iter(&self) -> Box<dyn Iterator<Item = (&'_ Self::Key, &'_ Self::Value)> + '_> {
90        Box::new(std::iter::empty())
91    }
92}
93
94/// [std::collections::HashMap] cache.
95pub struct HashMapCache<S = RandomState> {
96    _mark: PhantomData<S>,
97}
98
99impl<S: Send + Sync + BuildHasher + Default + 'static> HashMapCache<S> {
100    /// Use specified `S: BuildHasher` to create a `HashMap` cache.
101    pub fn new() -> Self {
102        Self { _mark: PhantomData }
103    }
104}
105
106impl Default for HashMapCache<RandomState> {
107    fn default() -> Self {
108        Self { _mark: PhantomData }
109    }
110}
111
112impl<S: Send + Sync + BuildHasher + Default + 'static> CacheFactory for HashMapCache<S> {
113    fn create<K, V>(&self) -> Box<dyn CacheStorage<Key = K, Value = V>>
114    where
115        K: Send + Sync + Clone + Eq + Hash + 'static,
116        V: Send + Sync + Clone + 'static,
117    {
118        Box::new(HashMapCacheImpl::<K, V, S>(HashMap::<K, V, S>::default()))
119    }
120}
121
122struct HashMapCacheImpl<K, V, S>(HashMap<K, V, S>);
123
124impl<K, V, S> CacheStorage for HashMapCacheImpl<K, V, S>
125where
126    K: Send + Sync + Clone + Eq + Hash + 'static,
127    V: Send + Sync + Clone + 'static,
128    S: Send + Sync + BuildHasher + 'static,
129{
130    type Key = K;
131    type Value = V;
132
133    #[inline]
134    fn get(&mut self, key: &Self::Key) -> Option<&Self::Value> {
135        self.0.get(key)
136    }
137
138    #[inline]
139    fn insert(&mut self, key: Cow<'_, Self::Key>, val: Cow<'_, Self::Value>) {
140        self.0.insert(key.into_owned(), val.into_owned());
141    }
142
143    #[inline]
144    fn remove(&mut self, key: &Self::Key) {
145        self.0.remove(key);
146    }
147
148    #[inline]
149    fn clear(&mut self) {
150        self.0.clear();
151    }
152
153    fn iter(&self) -> Box<dyn Iterator<Item = (&'_ Self::Key, &'_ Self::Value)> + '_> {
154        Box::new(self.0.iter())
155    }
156}
157
158/// LRU cache.
159pub struct LruCache {
160    cap: usize,
161}
162
163impl LruCache {
164    /// Creates a new LRU Cache that holds at most `cap` items.
165    pub fn new(cap: usize) -> Self {
166        Self { cap }
167    }
168}
169
170impl CacheFactory for LruCache {
171    fn create<K, V>(&self) -> Box<dyn CacheStorage<Key = K, Value = V>>
172    where
173        K: Send + Sync + Clone + Eq + Hash + 'static,
174        V: Send + Sync + Clone + 'static,
175    {
176        Box::new(LruCacheImpl(lru::LruCache::new(
177            NonZeroUsize::new(self.cap).unwrap(),
178        )))
179    }
180}
181
182struct LruCacheImpl<K, V>(lru::LruCache<K, V>);
183
184impl<K, V> CacheStorage for LruCacheImpl<K, V>
185where
186    K: Send + Sync + Clone + Eq + Hash + 'static,
187    V: Send + Sync + Clone + 'static,
188{
189    type Key = K;
190    type Value = V;
191
192    #[inline]
193    fn get(&mut self, key: &Self::Key) -> Option<&Self::Value> {
194        self.0.get(key)
195    }
196
197    #[inline]
198    fn insert(&mut self, key: Cow<'_, Self::Key>, val: Cow<'_, Self::Value>) {
199        self.0.put(key.into_owned(), val.into_owned());
200    }
201
202    #[inline]
203    fn remove(&mut self, key: &Self::Key) {
204        self.0.pop(key);
205    }
206
207    #[inline]
208    fn clear(&mut self) {
209        self.0.clear();
210    }
211
212    fn iter(&self) -> Box<dyn Iterator<Item = (&'_ Self::Key, &'_ Self::Value)> + '_> {
213        Box::new(self.0.iter())
214    }
215}
async_graphql/dataloader/cache.rs

async_graphql/dataloader/
cache.rs
