lib/sanitizer_common/sanitizer_addrhashmap.h

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