1//===-- sanitizer_addrhashmap.h ---------------------------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// Concurrent uptr->T hashmap.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef SANITIZER_ADDRHASHMAP_H
14#define SANITIZER_ADDRHASHMAP_H
15
16#include "sanitizer_common.h"
17#include "sanitizer_mutex.h"
18#include "sanitizer_atomic.h"
19#include "sanitizer_allocator_internal.h"
20
21namespace __sanitizer {
22
23// Concurrent uptr->T hashmap.
24// T must be a POD type, kSize is preferably a prime but can be any number.
25// Usage example:
26//
27// typedef AddrHashMap<uptr, 11> Map;
28// Map m;
29// {
30// Map::Handle h(&m, addr);
31// use h.operator->() to access the data
32// if h.created() then the element was just created, and the current thread
33// has exclusive access to it
34// otherwise the current thread has only read access to the data
35// }
36// {
37// Map::Handle h(&m, addr, true);
38// this will remove the data from the map in Handle dtor
39// the current thread has exclusive access to the data
40// if !h.exists() then the element never existed
41// }
42// {
43// Map::Handle h(&m, addr, false, true);
44// this will create a new element or return a handle to an existing element
45// if !h.created() this thread does *not* have exclusive access to the data
46// }
47template<typename T, uptr kSize>
48class AddrHashMap {
49 private:
50 struct Cell {
51 atomic_uintptr_t addr;
52 T val;
53 };
54
55 struct AddBucket {
56 uptr cap;
57 uptr size;
58 Cell cells[1]; // variable len
59 };
60
61 static const uptr kBucketSize = 3;
62
63 struct Bucket {
64 Mutex mtx;
65 atomic_uintptr_t add;
66 Cell cells[kBucketSize];
67 };
68
69 public:
70 AddrHashMap();
71
72 class Handle {
73 public:
74 Handle(AddrHashMap<T, kSize> *map, uptr addr);
75 Handle(AddrHashMap<T, kSize> *map, uptr addr, bool remove);
76 Handle(AddrHashMap<T, kSize> *map, uptr addr, bool remove, bool create);
77
78 ~Handle();
79 T *operator->();
80 T &operator*();
81 const T &operator*() const;
82 bool created() const;
83 bool exists() const;
84
85 private:
86 friend AddrHashMap<T, kSize>;
87 AddrHashMap<T, kSize> *map_;
88 Bucket *bucket_;
89 Cell *cell_;
90 uptr addr_;
91 uptr addidx_;
92 bool created_;
93 bool remove_;
94 bool create_;
95 };
96
97 typedef void (*ForEachCallback)(const uptr key, const T &val, void *arg);
98 // ForEach acquires a lock on each bucket while iterating over
99 // elements. Note that this only ensures that the structure of the hashmap is
100 // unchanged, there may be a data race to the element itself.
101 void ForEach(ForEachCallback cb, void *arg);
102
103 private:
104 friend class Handle;
105 Bucket *table_;
106
107 void acquire(Handle *h);
108 void release(Handle *h);
109 uptr calcHash(uptr addr);
110};
111
112template <typename T, uptr kSize>
113void AddrHashMap<T, kSize>::ForEach(ForEachCallback cb, void *arg) {
114 for (uptr n = 0; n < kSize; n++) {
115 Bucket *bucket = &table_[n];
116
117 ReadLock lock(&bucket->mtx);
118
119 for (uptr i = 0; i < kBucketSize; i++) {
120 Cell *c = &bucket->cells[i];
121 uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
122 if (addr1 != 0)
123 cb(addr1, c->val, arg);
124 }
125
126 // Iterate over any additional cells.
127 if (AddBucket *add =
128 (AddBucket *)atomic_load(&bucket->add, memory_order_acquire)) {
129 for (uptr i = 0; i < add->size; i++) {
130 Cell *c = &add->cells[i];
131 uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
132 if (addr1 != 0)
133 cb(addr1, c->val, arg);
134 }
135 }
136 }
137}
138
139template<typename T, uptr kSize>
140AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr) {
141 map_ = map;
142 addr_ = addr;
143 remove_ = false;
144 create_ = true;
145 map_->acquire(this);
146}
147
148template<typename T, uptr kSize>
149AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr,
150 bool remove) {
151 map_ = map;
152 addr_ = addr;
153 remove_ = remove;
154 create_ = true;
155 map_->acquire(this);
156}
157
158template<typename T, uptr kSize>
159AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr,
160 bool remove, bool create) {
161 map_ = map;
162 addr_ = addr;
163 remove_ = remove;
164 create_ = create;
165 map_->acquire(this);
166}
167
168template<typename T, uptr kSize>
169AddrHashMap<T, kSize>::Handle::~Handle() {
170 map_->release(this);
171}
172
173template <typename T, uptr kSize>
174T *AddrHashMap<T, kSize>::Handle::operator->() {
175 return &cell_->val;
176}
177
178template <typename T, uptr kSize>
179const T &AddrHashMap<T, kSize>::Handle::operator*() const {
180 return cell_->val;
181}
182
183template <typename T, uptr kSize>
184T &AddrHashMap<T, kSize>::Handle::operator*() {
185 return cell_->val;
186}
187
188template<typename T, uptr kSize>
189bool AddrHashMap<T, kSize>::Handle::created() const {
190 return created_;
191}
192
193template<typename T, uptr kSize>
194bool AddrHashMap<T, kSize>::Handle::exists() const {
195 return cell_ != nullptr;
196}
197
198template<typename T, uptr kSize>
199AddrHashMap<T, kSize>::AddrHashMap() {
200 table_ = (Bucket*)MmapOrDie(size: kSize * sizeof(table_[0]), mem_type: "AddrHashMap");
201}
202
203template <typename T, uptr kSize>
204void AddrHashMap<T, kSize>::acquire(Handle *h)
205 SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
206 uptr addr = h->addr_;
207 uptr hash = calcHash(addr);
208 Bucket *b = &table_[hash];
209
210 h->created_ = false;
211 h->addidx_ = -1U;
212 h->bucket_ = b;
213 h->cell_ = nullptr;
214
215 // If we want to remove the element, we need exclusive access to the bucket,
216 // so skip the lock-free phase.
217 if (h->remove_)
218 goto locked;
219
220 retry:
221 // First try to find an existing element w/o read mutex.
222 CHECK(!h->remove_);
223 // Check the embed cells.
224 for (uptr i = 0; i < kBucketSize; i++) {
225 Cell *c = &b->cells[i];
226 uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
227 if (addr1 == addr) {
228 h->cell_ = c;
229 return;
230 }
231 }
232
233 // Check the add cells with read lock.
234 if (atomic_load(&b->add, memory_order_relaxed)) {
235 b->mtx.ReadLock();
236 AddBucket *add = (AddBucket*)atomic_load(&b->add, memory_order_relaxed);
237 for (uptr i = 0; i < add->size; i++) {
238 Cell *c = &add->cells[i];
239 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
240 if (addr1 == addr) {
241 h->addidx_ = i;
242 h->cell_ = c;
243 return;
244 }
245 }
246 b->mtx.ReadUnlock();
247 }
248
249 locked:
250 // Re-check existence under write lock.
251 // Embed cells.
252 b->mtx.Lock();
253 for (uptr i = 0; i < kBucketSize; i++) {
254 Cell *c = &b->cells[i];
255 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
256 if (addr1 == addr) {
257 if (h->remove_) {
258 h->cell_ = c;
259 return;
260 }
261 b->mtx.Unlock();
262 goto retry;
263 }
264 }
265
266 // Add cells.
267 AddBucket *add = (AddBucket*)atomic_load(&b->add, memory_order_relaxed);
268 if (add) {
269 for (uptr i = 0; i < add->size; i++) {
270 Cell *c = &add->cells[i];
271 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
272 if (addr1 == addr) {
273 if (h->remove_) {
274 h->addidx_ = i;
275 h->cell_ = c;
276 return;
277 }
278 b->mtx.Unlock();
279 goto retry;
280 }
281 }
282 }
283
284 // The element does not exist, no need to create it if we want to remove.
285 if (h->remove_ || !h->create_) {
286 b->mtx.Unlock();
287 return;
288 }
289
290 // Now try to create it under the mutex.
291 h->created_ = true;
292 // See if we have a free embed cell.
293 for (uptr i = 0; i < kBucketSize; i++) {
294 Cell *c = &b->cells[i];
295 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
296 if (addr1 == 0) {
297 h->cell_ = c;
298 return;
299 }
300 }
301
302 // Store in the add cells.
303 if (!add) {
304 // Allocate a new add array.
305 const uptr kInitSize = 64;
306 add = (AddBucket*)InternalAlloc(size: kInitSize);
307 internal_memset(add, 0, kInitSize);
308 add->cap = (kInitSize - sizeof(*add)) / sizeof(add->cells[0]) + 1;
309 add->size = 0;
310 atomic_store(&b->add, (uptr)add, memory_order_relaxed);
311 }
312 if (add->size == add->cap) {
313 // Grow existing add array.
314 uptr oldsize = sizeof(*add) + (add->cap - 1) * sizeof(add->cells[0]);
315 uptr newsize = oldsize * 2;
316 AddBucket *add1 = (AddBucket*)InternalAlloc(size: newsize);
317 internal_memset(add1, 0, newsize);
318 add1->cap = (newsize - sizeof(*add)) / sizeof(add->cells[0]) + 1;
319 add1->size = add->size;
320 internal_memcpy(add1->cells, add->cells, add->size * sizeof(add->cells[0]));
321 InternalFree(add);
322 atomic_store(&b->add, (uptr)add1, memory_order_relaxed);
323 add = add1;
324 }
325 // Store.
326 uptr i = add->size++;
327 Cell *c = &add->cells[i];
328 CHECK_EQ(atomic_load(&c->addr, memory_order_relaxed), 0);
329 h->addidx_ = i;
330 h->cell_ = c;
331 }
332
333 template <typename T, uptr kSize>
334 void AddrHashMap<T, kSize>::release(Handle *h)
335 SANITIZER_NO_THREAD_SAFETY_ANALYSIS {
336 if (!h->cell_)
337 return;
338 Bucket *b = h->bucket_;
339 Cell *c = h->cell_;
340 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
341 if (h->created_) {
342 // Denote completion of insertion.
343 CHECK_EQ(addr1, 0);
344 // After the following store, the element becomes available
345 // for lock-free reads.
346 atomic_store(&c->addr, h->addr_, memory_order_release);
347 b->mtx.Unlock();
348 } else if (h->remove_) {
349 // Denote that the cell is empty now.
350 CHECK_EQ(addr1, h->addr_);
351 atomic_store(&c->addr, 0, memory_order_release);
352 // See if we need to compact the bucket.
353 AddBucket *add = (AddBucket *)atomic_load(&b->add, memory_order_relaxed);
354 if (h->addidx_ == -1U) {
355 // Removed from embed array, move an add element into the freed cell.
356 if (add && add->size != 0) {
357 uptr last = --add->size;
358 Cell *c1 = &add->cells[last];
359 c->val = c1->val;
360 uptr addr1 = atomic_load(&c1->addr, memory_order_relaxed);
361 atomic_store(&c->addr, addr1, memory_order_release);
362 atomic_store(&c1->addr, 0, memory_order_release);
363 }
364 } else {
365 // Removed from add array, compact it.
366 uptr last = --add->size;
367 Cell *c1 = &add->cells[last];
368 if (c != c1) {
369 *c = *c1;
370 atomic_store(&c1->addr, 0, memory_order_relaxed);
371 }
372 }
373 if (add && add->size == 0) {
374 // FIXME(dvyukov): free add?
375 }
376 b->mtx.Unlock();
377 } else {
378 CHECK_EQ(addr1, h->addr_);
379 if (h->addidx_ != -1U)
380 b->mtx.ReadUnlock();
381 }
382 }
383
384template<typename T, uptr kSize>
385uptr AddrHashMap<T, kSize>::calcHash(uptr addr) {
386 addr += addr << 10;
387 addr ^= addr >> 6;
388 return addr % kSize;
389}
390
391} // namespace __sanitizer
392
393#endif // SANITIZER_ADDRHASHMAP_H
394