[sanitizer] Add a fast version of StackDepotGet() for use in LSan.

Add a class that holds a snapshot of the StackDepot optimized for querying by
ID. This allows us to speed up LSan dramatically.

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@189217 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 129 insertions, 5 deletions

diff --git a/lib/lsan/lsan_common_linux.cc b/lib/lsan/lsan_common_linux.cc
index 08a05958a..38cf88a27 100644
--- a/lib/lsan/lsan_common_linux.cc
+++ b/lib/lsan/lsan_common_linux.cc
@@ -90,26 +90,34 @@ void ProcessGlobalRegions(Frontier *frontier) {
   dl_iterate_phdr(ProcessGlobalRegionsCallback, frontier);
 }
 
-static uptr GetCallerPC(u32 stack_id) {
+static uptr GetCallerPC(u32 stack_id, StackDepotReverseMap *map) {
   CHECK(stack_id);
   uptr size = 0;
-  const uptr *trace = StackDepotGet(stack_id, &size);
+  const uptr *trace = map->Get(stack_id, &size);
   // The top frame is our malloc/calloc/etc. The next frame is the caller.
   if (size >= 2)
     return trace[1];
   return 0;
 }
 
+struct ProcessPlatformAllocParam {
+  Frontier *frontier;
+  StackDepotReverseMap *stack_depot_reverse_map;
+};
+
 // ForEachChunk callback. Identifies unreachable chunks which must be treated as
 // reachable. Marks them as reachable and adds them to the frontier.
 static void ProcessPlatformSpecificAllocationsCb(uptr chunk, void *arg) {
   CHECK(arg);
+  ProcessPlatformAllocParam *param =
+      reinterpret_cast<ProcessPlatformAllocParam *>(arg);
   chunk = GetUserBegin(chunk);
   LsanMetadata m(chunk);
   if (m.allocated() && m.tag() != kReachable) {
-    if (linker->containsAddress(GetCallerPC(m.stack_trace_id()))) {
+    if (linker->containsAddress(
+            GetCallerPC(m.stack_trace_id(), param->stack_depot_reverse_map))) {
       m.set_tag(kReachable);
-      reinterpret_cast<Frontier *>(arg)->push_back(chunk);
+      param->frontier->push_back(chunk);
     }
   }
 }
@@ -119,7 +127,9 @@ static void ProcessPlatformSpecificAllocationsCb(uptr chunk, void *arg) {
 void ProcessPlatformSpecificAllocations(Frontier *frontier) {
   if (!flags()->use_tls) return;
   if (!linker) return;
-  ForEachChunk(ProcessPlatformSpecificAllocationsCb, frontier);
+  StackDepotReverseMap stack_depot_reverse_map;
+  ProcessPlatformAllocParam arg = {frontier, &stack_depot_reverse_map};
+  ForEachChunk(ProcessPlatformSpecificAllocationsCb, &arg);
 }
 
 }  // namespace __lsan
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index 908d11d0e..512e8f6b7 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -380,6 +380,22 @@ void InternalSort(Container *v, uptr size, Compare comp) {
   }
 }
 
+template<class Container, class Value, class Compare>
+uptr InternalBinarySearch(const Container &v, uptr first, uptr last,
+                          const Value &val, Compare comp) {
+  uptr not_found = last + 1;
+  while (last >= first) {
+    uptr mid = (first + last) / 2;
+    if (comp(v[mid], val))
+      first = mid + 1;
+    else if (comp(val, v[mid]))
+      last = mid - 1;
+    else
+      return mid;
+  }
+  return not_found;
+}
+
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_COMMON_H
diff --git a/lib/sanitizer_common/sanitizer_stackdepot.cc b/lib/sanitizer_common/sanitizer_stackdepot.cc
index 08e523832..2793bd071 100644
--- a/lib/sanitizer_common/sanitizer_stackdepot.cc
+++ b/lib/sanitizer_common/sanitizer_stackdepot.cc
@@ -201,4 +201,38 @@ const uptr *StackDepotGet(u32 id, uptr *size) {
   return 0;
 }
 
+bool StackDepotReverseMap::IdDescPair::IdComparator(
+    const StackDepotReverseMap::IdDescPair &a,
+    const StackDepotReverseMap::IdDescPair &b) {
+  return a.id < b.id;
+}
+
+StackDepotReverseMap::StackDepotReverseMap()
+    : map_(StackDepotGetStats()->n_uniq_ids + 100) {
+  for (int idx = 0; idx < kTabSize; idx++) {
+    atomic_uintptr_t *p = &depot.tab[idx];
+    uptr v = atomic_load(p, memory_order_consume);
+    StackDesc *s = (StackDesc*)(v & ~1);
+    for (; s; s = s->link) {
+      IdDescPair pair = {s->id, s};
+      map_.push_back(pair);
+    }
+  }
+  InternalSort(&map_, map_.size(), IdDescPair::IdComparator);
+}
+
+const uptr *StackDepotReverseMap::Get(u32 id, uptr *size) {
+  if (!map_.size()) return 0;
+  IdDescPair pair = {id, 0};
+  uptr idx = InternalBinarySearch(map_, 0, map_.size(), pair,
+                                  IdDescPair::IdComparator);
+  if (idx > map_.size()) {
+    *size = 0;
+    return 0;
+  }
+  StackDesc *desc = map_[idx].desc;
+  *size = desc->size;
+  return desc->stack;
+}
+
 }  // namespace __sanitizer
diff --git a/lib/sanitizer_common/sanitizer_stackdepot.h b/lib/sanitizer_common/sanitizer_stackdepot.h
index 5915fdbb4..8fd0cae72 100644
--- a/lib/sanitizer_common/sanitizer_stackdepot.h
+++ b/lib/sanitizer_common/sanitizer_stackdepot.h
@@ -13,6 +13,7 @@
 #ifndef SANITIZER_STACKDEPOT_H
 #define SANITIZER_STACKDEPOT_H
 
+#include "sanitizer_common.h"
 #include "sanitizer_internal_defs.h"
 
 namespace __sanitizer {
@@ -31,6 +32,31 @@ struct StackDepotStats {
 
 StackDepotStats *StackDepotGetStats();
 
+struct StackDesc;
+
+// Instantiating this class creates a snapshot of StackDepot which can be
+// efficiently queried with StackDepotGet(). You can use it concurrently with
+// StackDepot, but the snapshot is only guaranteed to contain those stack traces
+// which were stored before it was instantiated.
+class StackDepotReverseMap {
+ public:
+  StackDepotReverseMap();
+  const uptr *Get(u32 id, uptr *size);
+
+ private:
+  struct IdDescPair {
+   u32 id;
+   StackDesc *desc;
+
+   static bool IdComparator(const IdDescPair &a, const IdDescPair &b);
+  };
+
+  InternalMmapVector<IdDescPair> map_;
+
+  // Disallow evil constructors.
+  StackDepotReverseMap(const StackDepotReverseMap&);
+  void operator=(const StackDepotReverseMap&);
+};
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_STACKDEPOT_H
diff --git a/lib/sanitizer_common/tests/sanitizer_common_test.cc b/lib/sanitizer_common/tests/sanitizer_common_test.cc
index ee913dcbb..cefaf6cf4 100644
--- a/lib/sanitizer_common/tests/sanitizer_common_test.cc
+++ b/lib/sanitizer_common/tests/sanitizer_common_test.cc
@@ -158,4 +158,19 @@ TEST(SanitizerCommon, ThreadStackTlsWorker) {
   pthread_join(t, 0);
 }
 
+bool UptrLess(uptr a, uptr b) {
+  return a < b;
+}
+
+TEST(SanitizerCommon, InternalBinarySearch) {
+  const uptr sz = 5;
+  uptr arr[sz];
+  for (uptr i = 0; i < sz; i++) arr[i] = i * i;
+
+  for (uptr i = 0; i < sz; i++)
+    ASSERT_EQ(InternalBinarySearch(arr, 0, sz, i * i, UptrLess), i);
+
+  ASSERT_EQ(InternalBinarySearch(arr, 0, sz, 7, UptrLess), sz + 1);
+}
+
 }  // namespace __sanitizer
diff --git a/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cc b/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cc
index 5350c2ab8..5c075d53e 100644
--- a/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cc
+++ b/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cc
@@ -66,4 +66,27 @@ TEST(SanitizerCommon, StackDepotSeveral) {
   EXPECT_NE(i1, i2);
 }
 
+TEST(SanitizerCommon, StackDepotReverseMap) {
+  uptr s1[] = {1, 2, 3, 4, 5};
+  uptr s2[] = {7, 1, 3, 0};
+  uptr s3[] = {10, 2, 5, 3};
+  uptr s4[] = {1, 3, 2, 5};
+  u32 ids[4] = {0};
+  ids[0] = StackDepotPut(s1, ARRAY_SIZE(s1));
+  ids[1] = StackDepotPut(s2, ARRAY_SIZE(s2));
+  ids[2] = StackDepotPut(s3, ARRAY_SIZE(s3));
+  ids[3] = StackDepotPut(s4, ARRAY_SIZE(s4));
+
+  StackDepotReverseMap map;
+
+  for (uptr i = 0; i < 4; i++) {
+    uptr sz_depot, sz_map;
+    const uptr *sp_depot, *sp_map;
+    sp_depot = StackDepotGet(ids[i], &sz_depot);
+    sp_map = map.Get(ids[i], &sz_map);
+    EXPECT_EQ(sz_depot, sz_map);
+    EXPECT_EQ(sp_depot, sp_map);
+  }
+}
+
 }  // namespace __sanitizer