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> |
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> |
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> |
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> |
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> |
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> |
189 | bool AddrHashMap<T, kSize>::Handle::created() const { |
190 | return created_; |
191 | } |
192 | |
193 | template<typename T, uptr kSize> |
194 | bool AddrHashMap<T, kSize>::Handle::exists() const { |
195 | return cell_ != nullptr; |
196 | } |
197 | |
198 | template<typename T, uptr kSize> |
199 | AddrHashMap<T, kSize>::AddrHashMap() { |
200 | table_ = (Bucket*)MmapOrDie(size: kSize * sizeof(table_[0]), mem_type: "AddrHashMap" ); |
201 | } |
202 | |
203 | template <typename T, uptr kSize> |
204 | void 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 | |
384 | template<typename T, uptr kSize> |
385 | uptr 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 | |