pingora_cache/
hashtable.rs1use lru::LruCache;
18use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
19use std::collections::HashMap;
20
21pub struct ConcurrentHashTable<V, const N: usize> {
24 tables: [RwLock<HashMap<u128, V>>; N],
25}
26
27#[inline]
28fn get_shard(key: u128, n_shards: usize) -> usize {
29 (key % n_shards as u128) as usize
30}
31
32impl<V, const N: usize> ConcurrentHashTable<V, N>
33where
34 [RwLock<HashMap<u128, V>>; N]: Default,
35{
36 pub fn new() -> Self {
37 ConcurrentHashTable {
38 tables: Default::default(),
39 }
40 }
41 pub fn get(&self, key: u128) -> &RwLock<HashMap<u128, V>> {
42 &self.tables[get_shard(key, N)]
43 }
44
45 #[allow(dead_code)]
46 pub fn read(&self, key: u128) -> RwLockReadGuard<HashMap<u128, V>> {
47 self.get(key).read()
48 }
49
50 pub fn write(&self, key: u128) -> RwLockWriteGuard<HashMap<u128, V>> {
51 self.get(key).write()
52 }
53
54 }
56
57impl<V, const N: usize> Default for ConcurrentHashTable<V, N>
58where
59 [RwLock<HashMap<u128, V>>; N]: Default,
60{
61 fn default() -> Self {
62 Self::new()
63 }
64}
65
66#[doc(hidden)] pub struct LruShard<V>(RwLock<LruCache<u128, V>>);
68impl<V> Default for LruShard<V> {
69 fn default() -> Self {
70 LruShard(RwLock::new(LruCache::unbounded()))
72 }
73}
74
75pub struct ConcurrentLruCache<V, const N: usize> {
77 lrus: [LruShard<V>; N],
78}
79
80impl<V, const N: usize> ConcurrentLruCache<V, N>
81where
82 [LruShard<V>; N]: Default,
83{
84 pub fn new(shard_capacity: usize) -> Self {
85 use std::num::NonZeroUsize;
86 const ONE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(1) };
88 let mut cache = ConcurrentLruCache {
89 lrus: Default::default(),
90 };
91 for lru in &mut cache.lrus {
92 lru.0
93 .write()
94 .resize(shard_capacity.try_into().unwrap_or(ONE));
95 }
96 cache
97 }
98 pub fn get(&self, key: u128) -> &RwLock<LruCache<u128, V>> {
99 &self.lrus[get_shard(key, N)].0
100 }
101
102 #[allow(dead_code)]
103 pub fn read(&self, key: u128) -> RwLockReadGuard<LruCache<u128, V>> {
104 self.get(key).read()
105 }
106
107 pub fn write(&self, key: u128) -> RwLockWriteGuard<LruCache<u128, V>> {
108 self.get(key).write()
109 }
110
111 }