394 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8//9// Concurrent uptr->T hashmap.10//11//===----------------------------------------------------------------------===//12 13#ifndef SANITIZER_ADDRHASHMAP_H14#define SANITIZER_ADDRHASHMAP_H15 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 data32// if h.created() then the element was just created, and the current thread33// has exclusive access to it34// otherwise the current thread has only read access to the data35// }36// {37// Map::Handle h(&m, addr, true);38// this will remove the data from the map in Handle dtor39// the current thread has exclusive access to the data40// if !h.exists() then the element never existed41// }42// {43// Map::Handle h(&m, addr, false, true);44// this will create a new element or return a handle to an existing element45// if !h.created() this thread does *not* have exclusive access to the data46// }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 len59 };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 over99 // elements. Note that this only ensures that the structure of the hashmap is100 // 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(kSize * sizeof(table_[0]), "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(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 available345 // 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 __sanitizer392 393#endif // SANITIZER_ADDRHASHMAP_H394