Skip to main content

pingora_cache/
hashtable.rs

1// Copyright 2026 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
23#[derive(Debug)]
24pub struct ConcurrentHashTable<V, const N: usize> {
25    tables: [RwLock<HashMap<u128, V>>; N],
26}
27
28#[inline]
29fn get_shard(key: u128, n_shards: usize) -> usize {
30    (key % n_shards as u128) as usize
31}
32
33impl<V, const N: usize> ConcurrentHashTable<V, N> {
34    pub fn new() -> Self {
35        // Build the per-shard array element-by-element via `arrayvec`. The
36        // stdlib only auto-derives `Default` for `[T; N]` up to N=32, so the
37        // previous `Default::default()`-based init silently capped this type
38        // at 32 shards. Mirrors the same `arrayvec::ArrayVec` pattern that
39        // `pingora_lru::Lru<T, N>::with_capacity_and_watermark` uses to lift
40        // the same constraint, so the two sharded structures now support
41        // identical shard counts (callers can pick any `N`).
42        let mut tables = arrayvec::ArrayVec::<_, N>::new();
43        for _ in 0..N {
44            tables.push(RwLock::new(HashMap::new()));
45        }
46        ConcurrentHashTable {
47            // `into_inner` is infallible here because the loop above pushed
48            // exactly N elements. `.ok().expect(...)` avoids requiring the
49            // element type to be `Debug` (which `into_inner`'s `Err` payload
50            // would otherwise demand for `.expect`).
51            tables: tables
52                .into_inner()
53                .ok()
54                .expect("ArrayVec pushed N times, into_inner is infallible"),
55        }
56    }
57    pub fn get(&self, key: u128) -> &RwLock<HashMap<u128, V>> {
58        &self.tables[get_shard(key, N)]
59    }
60
61    #[allow(dead_code)]
62    pub fn get_shard_at_idx(&self, idx: usize) -> Option<&RwLock<HashMap<u128, V>>> {
63        self.tables.get(idx)
64    }
65
66    #[allow(dead_code)]
67    pub fn read(&self, key: u128) -> RwLockReadGuard<'_, HashMap<u128, V>> {
68        self.get(key).read()
69    }
70
71    pub fn write(&self, key: u128) -> RwLockWriteGuard<'_, HashMap<u128, V>> {
72        self.get(key).write()
73    }
74
75    #[allow(dead_code)]
76    pub fn for_each<F>(&self, mut f: F)
77    where
78        F: FnMut(&u128, &V),
79    {
80        for shard in &self.tables {
81            let guard = shard.read();
82            for (key, value) in guard.iter() {
83                f(key, value);
84            }
85        }
86    }
87
88    // TODO: work out the lifetimes to provide get/set directly
89}
90
91impl<V, const N: usize> Default for ConcurrentHashTable<V, N> {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97#[doc(hidden)] // not need in public API
98pub struct LruShard<V>(RwLock<LruCache<u128, V>>);
99
100/// Sharded concurrent data structure for LruCache
101pub struct ConcurrentLruCache<V, const N: usize> {
102    lrus: [LruShard<V>; N],
103}
104
105impl<V, const N: usize> ConcurrentLruCache<V, N> {
106    pub fn new(shard_capacity: usize) -> Self {
107        use std::num::NonZeroUsize;
108        // safe, 1 != 0
109        const ONE: NonZeroUsize = NonZeroUsize::new(1).unwrap();
110        // Same `arrayvec` element-by-element init as `ConcurrentHashTable::new`
111        // and `pingora_lru::Lru` — works for any `N`, not just `N <= 32`.
112        let cap = shard_capacity.try_into().unwrap_or(ONE);
113        let mut lrus = arrayvec::ArrayVec::<_, N>::new();
114        for _ in 0..N {
115            lrus.push(LruShard(RwLock::new(LruCache::new(cap))));
116        }
117        ConcurrentLruCache {
118            // Same as `ConcurrentHashTable::new` — `into_inner` cannot fail
119            // because we pushed exactly N elements; `.ok().expect(...)` keeps
120            // the message readable without forcing the element type to be
121            // `Debug`.
122            lrus: lrus
123                .into_inner()
124                .ok()
125                .expect("ArrayVec pushed N times, into_inner is infallible"),
126        }
127    }
128    pub fn get(&self, key: u128) -> &RwLock<LruCache<u128, V>> {
129        &self.lrus[get_shard(key, N)].0
130    }
131
132    #[allow(dead_code)]
133    pub fn read(&self, key: u128) -> RwLockReadGuard<'_, LruCache<u128, V>> {
134        self.get(key).read()
135    }
136
137    pub fn write(&self, key: u128) -> RwLockWriteGuard<'_, LruCache<u128, V>> {
138        self.get(key).write()
139    }
140
141    // TODO: work out the lifetimes to provide get/set directly
142}