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}