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
21 namespace __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 // }
47 template<typename T, uptr kSize>
48 class 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
112 template <typename T, uptr kSize>
ForEach(ForEachCallback cb,void * arg)113 void 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
139 template<typename T, uptr kSize>
Handle(AddrHashMap<T,kSize> * map,uptr addr)140 AddrHashMap<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
148 template<typename T, uptr kSize>
Handle(AddrHashMap<T,kSize> * map,uptr addr,bool remove)149 AddrHashMap<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
158 template<typename T, uptr kSize>
Handle(AddrHashMap<T,kSize> * map,uptr addr,bool remove,bool create)159 AddrHashMap<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
168 template<typename T, uptr kSize>
~Handle()169 AddrHashMap<T, kSize>::Handle::~Handle() {
170 map_->release(this);
171 }
172
173 template <typename T, uptr kSize>
174 T *AddrHashMap<T, kSize>::Handle::operator->() {
175 return &cell_->val;
176 }
177
178 template <typename T, uptr kSize>
179 const T &AddrHashMap<T, kSize>::Handle::operator*() const {
180 return cell_->val;
181 }
182
183 template <typename T, uptr kSize>
184 T &AddrHashMap<T, kSize>::Handle::operator*() {
185 return cell_->val;
186 }
187
188 template<typename T, uptr kSize>
created()189 bool AddrHashMap<T, kSize>::Handle::created() const {
190 return created_;
191 }
192
193 template<typename T, uptr kSize>
exists()194 bool AddrHashMap<T, kSize>::Handle::exists() const {
195 return cell_ != nullptr;
196 }
197
198 template<typename T, uptr kSize>
AddrHashMap()199 AddrHashMap<T, kSize>::AddrHashMap() {
200 table_ = (Bucket*)MmapOrDie(kSize * sizeof(table_[0]), "AddrHashMap");
201 }
202
203 template <typename T, uptr kSize>
acquire(Handle * h)204 void AddrHashMap<T, kSize>::acquire(Handle *h) NO_THREAD_SAFETY_ANALYSIS {
205 uptr addr = h->addr_;
206 uptr hash = calcHash(addr);
207 Bucket *b = &table_[hash];
208
209 h->created_ = false;
210 h->addidx_ = -1U;
211 h->bucket_ = b;
212 h->cell_ = nullptr;
213
214 // If we want to remove the element, we need exclusive access to the bucket,
215 // so skip the lock-free phase.
216 if (h->remove_)
217 goto locked;
218
219 retry:
220 // First try to find an existing element w/o read mutex.
221 CHECK(!h->remove_);
222 // Check the embed cells.
223 for (uptr i = 0; i < kBucketSize; i++) {
224 Cell *c = &b->cells[i];
225 uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
226 if (addr1 == addr) {
227 h->cell_ = c;
228 return;
229 }
230 }
231
232 // Check the add cells with read lock.
233 if (atomic_load(&b->add, memory_order_relaxed)) {
234 b->mtx.ReadLock();
235 AddBucket *add = (AddBucket*)atomic_load(&b->add, memory_order_relaxed);
236 for (uptr i = 0; i < add->size; i++) {
237 Cell *c = &add->cells[i];
238 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
239 if (addr1 == addr) {
240 h->addidx_ = i;
241 h->cell_ = c;
242 return;
243 }
244 }
245 b->mtx.ReadUnlock();
246 }
247
248 locked:
249 // Re-check existence under write lock.
250 // Embed cells.
251 b->mtx.Lock();
252 for (uptr i = 0; i < kBucketSize; i++) {
253 Cell *c = &b->cells[i];
254 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
255 if (addr1 == addr) {
256 if (h->remove_) {
257 h->cell_ = c;
258 return;
259 }
260 b->mtx.Unlock();
261 goto retry;
262 }
263 }
264
265 // Add cells.
266 AddBucket *add = (AddBucket*)atomic_load(&b->add, memory_order_relaxed);
267 if (add) {
268 for (uptr i = 0; i < add->size; i++) {
269 Cell *c = &add->cells[i];
270 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
271 if (addr1 == addr) {
272 if (h->remove_) {
273 h->addidx_ = i;
274 h->cell_ = c;
275 return;
276 }
277 b->mtx.Unlock();
278 goto retry;
279 }
280 }
281 }
282
283 // The element does not exist, no need to create it if we want to remove.
284 if (h->remove_ || !h->create_) {
285 b->mtx.Unlock();
286 return;
287 }
288
289 // Now try to create it under the mutex.
290 h->created_ = true;
291 // See if we have a free embed cell.
292 for (uptr i = 0; i < kBucketSize; i++) {
293 Cell *c = &b->cells[i];
294 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
295 if (addr1 == 0) {
296 h->cell_ = c;
297 return;
298 }
299 }
300
301 // Store in the add cells.
302 if (!add) {
303 // Allocate a new add array.
304 const uptr kInitSize = 64;
305 add = (AddBucket*)InternalAlloc(kInitSize);
306 internal_memset(add, 0, kInitSize);
307 add->cap = (kInitSize - sizeof(*add)) / sizeof(add->cells[0]) + 1;
308 add->size = 0;
309 atomic_store(&b->add, (uptr)add, memory_order_relaxed);
310 }
311 if (add->size == add->cap) {
312 // Grow existing add array.
313 uptr oldsize = sizeof(*add) + (add->cap - 1) * sizeof(add->cells[0]);
314 uptr newsize = oldsize * 2;
315 AddBucket *add1 = (AddBucket*)InternalAlloc(newsize);
316 internal_memset(add1, 0, newsize);
317 add1->cap = (newsize - sizeof(*add)) / sizeof(add->cells[0]) + 1;
318 add1->size = add->size;
319 internal_memcpy(add1->cells, add->cells, add->size * sizeof(add->cells[0]));
320 InternalFree(add);
321 atomic_store(&b->add, (uptr)add1, memory_order_relaxed);
322 add = add1;
323 }
324 // Store.
325 uptr i = add->size++;
326 Cell *c = &add->cells[i];
327 CHECK_EQ(atomic_load(&c->addr, memory_order_relaxed), 0);
328 h->addidx_ = i;
329 h->cell_ = c;
330 }
331
332 template <typename T, uptr kSize>
release(Handle * h)333 void AddrHashMap<T, kSize>::release(Handle *h) NO_THREAD_SAFETY_ANALYSIS {
334 if (!h->cell_)
335 return;
336 Bucket *b = h->bucket_;
337 Cell *c = h->cell_;
338 uptr addr1 = atomic_load(&c->addr, memory_order_relaxed);
339 if (h->created_) {
340 // Denote completion of insertion.
341 CHECK_EQ(addr1, 0);
342 // After the following store, the element becomes available
343 // for lock-free reads.
344 atomic_store(&c->addr, h->addr_, memory_order_release);
345 b->mtx.Unlock();
346 } else if (h->remove_) {
347 // Denote that the cell is empty now.
348 CHECK_EQ(addr1, h->addr_);
349 atomic_store(&c->addr, 0, memory_order_release);
350 // See if we need to compact the bucket.
351 AddBucket *add = (AddBucket *)atomic_load(&b->add, memory_order_relaxed);
352 if (h->addidx_ == -1U) {
353 // Removed from embed array, move an add element into the freed cell.
354 if (add && add->size != 0) {
355 uptr last = --add->size;
356 Cell *c1 = &add->cells[last];
357 c->val = c1->val;
358 uptr addr1 = atomic_load(&c1->addr, memory_order_relaxed);
359 atomic_store(&c->addr, addr1, memory_order_release);
360 atomic_store(&c1->addr, 0, memory_order_release);
361 }
362 } else {
363 // Removed from add array, compact it.
364 uptr last = --add->size;
365 Cell *c1 = &add->cells[last];
366 if (c != c1) {
367 *c = *c1;
368 atomic_store(&c1->addr, 0, memory_order_relaxed);
369 }
370 }
371 if (add && add->size == 0) {
372 // FIXME(dvyukov): free add?
373 }
374 b->mtx.Unlock();
375 } else {
376 CHECK_EQ(addr1, h->addr_);
377 if (h->addidx_ != -1U)
378 b->mtx.ReadUnlock();
379 }
380 }
381
382 template<typename T, uptr kSize>
calcHash(uptr addr)383 uptr AddrHashMap<T, kSize>::calcHash(uptr addr) {
384 addr += addr << 10;
385 addr ^= addr >> 6;
386 return addr % kSize;
387 }
388
389 } // namespace __sanitizer
390
391 #endif // SANITIZER_ADDRHASHMAP_H
392