[sanitizer] Add a ForEach callback interface for AddrHashMap.

[snehasish]

Summary

This change adds a ForEach method to the AddrHashMap class which can
then be used to iterate over all the key value pairs in the hash map.
I intend to use this in an upcoming change to the memprof runtime.

[snehasish]

I couldn't find any unit tests for this hashmap implementation. @vitalybuka if you can point me to an appropriate place, I would be happy to add a test to cover this new method.

[vitalybuka]

some basic unittest could be nice

compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h

117

ReadLock l(bucket->mtx);

[vitalybuka]

In D111368#3049699, @snehasish wrote:

I couldn't find any unit tests for this hashmap implementation. @vitalybuka if you can point me to an appropriate place, I would be happy to add a test to cover this new method.

I see, just create new sanitizer_addrhashmap_test.cpp

[snehasish]

Added a unit test.

[snehasish]

Added a unit test and addressed the comment. PTAL, thanks!

[vitalybuka]

compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h

130

I believe all 3 atomic_loads needs memory_order_acquire to make sure that val is loaded after we loaded addr1

133

cc @dvyukov
I guess nothing prevents other threads from modifying content of *(c->val) when callback is running ReadLock only guaranties that storage for val will not be released. So other syncronization will be needed on callback side.

It would be nice to see patches which use the new function.

compiler-rt/lib/sanitizer_common/tests/sanitizer_addrhashmap_test.cpp

2

first two lines need formating

10

lines 9-12 can be removed

35

maybe delete them from ref as well

and then at the end

m.ForEach(ExistsInReferenceMap, &reference_map);
EXPECT_TRUE(m.empty());

36

ASSERT_NE()

[snehasish]

Address comments.

[snehasish]

Thanks for the review. PTAL!

compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h

133

I changed the lock type to just a read write lock. The remainder of the code is borrowed from the locked version of acquire (L243-L276) in this revision. If we hold an exclusive lock for the bucket then I don't think we need to ensure that the operations are ordered.

Memprof usage is shown in child revision D111676 - see memprof_allocator.cpp L258. It is called when the profiler destructor is called and the contents of the map, i.e. the profile is written out.

[vitalybuka]

compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h

133

ReadLock is fine.
Issue is that bucket->mtx only locks structure of AddrHashMap only, neither Lock or ReadLock help us here regarding content of the val.

Thread1:
Map::Handle h(&m, pre_existing_address); // returns from acquire on L223 doing essentially nothing
h->update_data();

Thread2:
m.ForEach(Print, ..)

So there is data race between Print and Merge

[snehasish]

Address comments.

[snehasish]

Thanks for taking the time to clarify with an example, Vitaly. I completely missed the point when you raised it previously. I've reverted the change here and added some documentation to clarify expectations for users. I'll update the child revision with additional locking to prevent the data race on accessing the element.

[vitalybuka]

LGTM, but I slightly prefer you avoid this patch and go with simpler solution on another patch.

compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h

 //   otherwise the current thread has only read access to the data
 // }
 // {
 //   Map::Handle h(&m, addr, true);
 //   this will remove the data from the map in Handle dtor
 //   the current thread has exclusive access to the data
 //   if !h.exists() then the element never existed
 // }
+// {
+//   Map::Handle h(&m, addr, false, true);
+//   this will create a new element or return a handle to an existing element
+//   if !h.created() this thread does *not* have exclusive access to the data
+// }
 template<typename T, uptr kSize>
 class AddrHashMap {
  private:
   struct Cell {
     atomic_uintptr_t addr;
     T                val;
   };
 
private:
     Cell                  *cell_;
     uptr                   addr_;
     uptr                   addidx_;
     bool                   created_;
     bool                   remove_;
     bool                   create_;
   };
 
+  typedef void (*ForEachCallback)(const uptr key, const T &val, void *arg);
+  // ForEach acquires a lock on each bucket while iterating over
+  // elements. Note that this only ensures that the structure of the hashmap is
+  // unchanged, there may be a data race to the element itself.
+  void ForEach(ForEachCallback cb, void *arg);
+
  private:
   friend class Handle;
   Bucket *table_;
 
   void acquire(Handle *h);
   void release(Handle *h);
   uptr calcHash(uptr addr);
 };
 
-template<typename T, uptr kSize>
+template <typename T, uptr kSize>
+void AddrHashMap<T, kSize>::ForEach(ForEachCallback cb, void *arg) {
+  for (uptr n = 0; n < kSize; n++) {
+    Bucket *bucket = &table_[n];
+
+    ReadLock lock(&bucket->mtx);
+
+    for (uptr i = 0; i < kBucketSize; i++) {
+      Cell *c = &bucket->cells[i];
+      uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
+      if (addr1 != 0)
+        cb(addr1, c->val, arg);
+    }
+
+    // Iterate over any additional cells.
+    if (AddBucket *add =
+            (AddBucket *)atomic_load(&bucket->add, memory_order_acquire)) {
+      for (uptr i = 0; i < add->size; i++) {
+        Cell *c = &add->cells[i];
+        uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
+        if (addr1 != 0)
+          cb(addr1, c->val, arg);
+      }
+    }
+  }
+}
+
+template<typename T, uptr kSize>
 AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr) {
   map_ = map;
   addr_ = addr;
   remove_ = false;
   create_ = true;
   map_->acquire(this);
 }
 

compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt

 include(CompilerRTCompile)
 
 clang_compiler_add_cxx_check()
 
 # FIXME: use SANITIZER_COMMON_SUPPORTED_ARCH here
 filter_available_targets(SANITIZER_UNITTEST_SUPPORTED_ARCH x86_64 i386 mips64 mips64el riscv64)
 if(APPLE)
   darwin_filter_host_archs(SANITIZER_UNITTEST_SUPPORTED_ARCH SANITIZER_UNITTEST_SUPPORTED_ARCH)
 endif()
 
 set(SANITIZER_UNITTESTS
+  sanitizer_addrhashmap_test.cpp
   sanitizer_allocator_test.cpp
   sanitizer_atomic_test.cpp
   sanitizer_bitvector_test.cpp
   sanitizer_bvgraph_test.cpp
   sanitizer_chained_origin_depot_test.cpp
   sanitizer_common_test.cpp
   sanitizer_deadlock_detector_test.cpp
   sanitizer_flags_test.cpp

compiler-rt/lib/sanitizer_common/tests/sanitizer_addrhashmap_test.cpp

+//===-- sanitizer_addrhashmap_test.cpp ------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#include "sanitizer_common/sanitizer_addrhashmap.h"
+
+#include <unordered_map>
+
+#include "gtest/gtest.h"
+
+namespace __sanitizer {
+
+struct Value {
+  int payload;
+  inline bool operator==(const Value& rhs) const {
+    return payload == rhs.payload;
+  }
+};
+
+using MapTy = AddrHashMap<Value, 11>;
+using HandleTy = MapTy::Handle;
+using RefMapTy = std::unordered_map<uptr, Value>;
+
+static void ExistsInReferenceMap(const uptr key, const Value& val, void* arg) {
+  RefMapTy* ref = reinterpret_cast<RefMapTy*>(arg);
+  const RefMapTy::iterator iter = ref->find(key);
+  ASSERT_NE(iter, ref->end());
+  EXPECT_EQ(iter->second, val);
+  ref->erase(iter);
+}
+
+TEST(AddrHashMap, Basic) {
+  // Use a reference implementation to compare with.
+  RefMapTy reference_map{
+      {0x1000, {1}},
+      {0x2000, {2}},
+      {0x3000, {3}},
+  };
+
+  MapTy m;
+
+  for (const auto& key_val : reference_map) {
+    const uptr key = key_val.first;
+    const Value val = key_val.second;
+
+    // Insert all the elements.
+    {
+      HandleTy h(&m, key);
+      ASSERT_TRUE(h.created());
+      h->payload = val.payload;
+    }
+  }
+
+  // Now check that all the elements are present.
+  m.ForEach(ExistsInReferenceMap, &reference_map);
+  EXPECT_TRUE(reference_map.empty());
+}
+
+}  // namespace __sanitizer