xref: /freebsd/contrib/llvm-project/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h (revision 7ef62cebc2f965b0f640263e179276928885e33d)
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(kSize * sizeof(table_[0]), "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(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(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