pingora_cache/
hashtable.rs

1// Copyright 2025 Cloudflare, Inc.
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//! Concurrent hash tables and LRUs
16
17use lru::LruCache;
18use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
19use std::collections::HashMap;
20
21// There are probably off-the-shelf crates of this, DashMap?
22/// A hash table that shards to a constant number of tables to reduce lock contention
23pub 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    // TODO: work out the lifetimes to provide get/set directly
55}
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)] // not need in public API
67pub struct LruShard<V>(RwLock<LruCache<u128, V>>);
68impl<V> Default for LruShard<V> {
69    fn default() -> Self {
70        // help satisfy default construction of arrays
71        LruShard(RwLock::new(LruCache::unbounded()))
72    }
73}
74
75/// Sharded concurrent data structure for LruCache
76pub 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        // safe, 1 != 0
87        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    // TODO: work out the lifetimes to provide get/set directly
112}