9 files changed, 133 insertions, 7 deletions
diff --git a/include/sanitizer/hwasan_interface.h b/include/sanitizer/hwasan_interface.h
index 79bf23e26..1affd875a 100644
--- a/include/sanitizer/hwasan_interface.h
+++ b/include/sanitizer/hwasan_interface.h
@@ -59,6 +59,9 @@ extern "C" {
   // format.
   void __hwasan_print_shadow(const volatile void *x, size_t size);
 
+  // Print one-line report about the memory usage of the current process.
+  void __hwasan_print_memory_usage();
+
   int __sanitizer_posix_memalign(void **memptr, size_t alignment, size_t size);
   void * __sanitizer_memalign(size_t alignment, size_t size);
   void * __sanitizer_aligned_alloc(size_t alignment, size_t size);
diff --git a/lib/hwasan/hwasan.cc b/lib/hwasan/hwasan.cc
index 37be710b2..02ba6b7b1 100644
--- a/lib/hwasan/hwasan.cc
+++ b/lib/hwasan/hwasan.cc
@@ -171,6 +171,22 @@ static void HWAsanCheckFailed(const char *file, int line, const char *cond,
   Die();
 }
 
+static void HwasanPrintMemoryUsage() {
+  auto thread_stats = Thread::GetThreadStats();
+  auto *sds = StackDepotGetStats();
+  AllocatorStatCounters asc;
+  GetAllocatorStats(asc);
+  Printf(
+      "HWASAN pid: %d rss: %zd threads: %zd stacks: %zd"
+      " thr_aux: %zd stack_depot: %zd uniq_stacks: %zd"
+      " heap: %zd\n",
+      internal_getpid(), GetRSS(), thread_stats.n_live_threads,
+      thread_stats.total_stack_size,
+      thread_stats.n_live_threads * Thread::MemoryUsedPerThread(),
+      sds->allocated, sds->n_uniq_ids,
+      asc[AllocatorStatMapped]);
+}
+
 } // namespace __hwasan
 
 // Interface.
@@ -439,6 +455,8 @@ void __hwasan_handle_longjmp(const void *sp_dst) {
   TagMemory(sp, dst - sp, 0);
 }
 
+void __hwasan_print_memory_usage() { HwasanPrintMemoryUsage(); }
+
 static const u8 kFallbackTag = 0xBB;
 
 u8 __hwasan_generate_tag() {
diff --git a/lib/hwasan/hwasan_allocator.cc b/lib/hwasan/hwasan_allocator.cc
index c54117b3e..c9783b7e0 100644
--- a/lib/hwasan/hwasan_allocator.cc
+++ b/lib/hwasan/hwasan_allocator.cc
@@ -48,6 +48,10 @@ static atomic_uint8_t hwasan_allocator_tagging_enabled;
 static const tag_t kFallbackAllocTag = 0xBB;
 static const tag_t kFallbackFreeTag = 0xBC;
 
+void GetAllocatorStats(AllocatorStatCounters s) {
+  allocator.GetStats(s);
+}
+
 void HwasanAllocatorInit() {
   atomic_store_relaxed(&hwasan_allocator_tagging_enabled,
                        !flags()->disable_allocator_tagging);
diff --git a/lib/hwasan/hwasan_allocator.h b/lib/hwasan/hwasan_allocator.h
index 4f8d8a795..6f50e58fc 100644
--- a/lib/hwasan/hwasan_allocator.h
+++ b/lib/hwasan/hwasan_allocator.h
@@ -97,6 +97,7 @@ struct HeapAllocationRecord {
 
 typedef RingBuffer<HeapAllocationRecord> HeapAllocationsRingBuffer;
 
+void GetAllocatorStats(AllocatorStatCounters s);
 
 } // namespace __hwasan
 
diff --git a/lib/hwasan/hwasan_interface_internal.h b/lib/hwasan/hwasan_interface_internal.h
index ddb0d2ae0..50164136a 100644
--- a/lib/hwasan/hwasan_interface_internal.h
+++ b/lib/hwasan/hwasan_interface_internal.h
@@ -143,6 +143,9 @@ void __hwasan_thread_enter();
 SANITIZER_INTERFACE_ATTRIBUTE
 void __hwasan_thread_exit();
 
+SANITIZER_INTERFACE_ATTRIBUTE
+void __hwasan_print_memory_usage();
+
 }  // extern "C"
 
 #endif  // HWASAN_INTERFACE_INTERNAL_H
diff --git a/lib/hwasan/hwasan_thread.cc b/lib/hwasan/hwasan_thread.cc
index f77c37550..f3471e03b 100644
--- a/lib/hwasan/hwasan_thread.cc
+++ b/lib/hwasan/hwasan_thread.cc
@@ -26,14 +26,17 @@ static u32 RandomSeed() {
 
 Thread *Thread::main_thread;
 SpinMutex Thread::thread_list_mutex;
+Thread::ThreadStats Thread::thread_stats;
 
 void Thread::InsertIntoThreadList(Thread *t) {
   CHECK(!t->next_);
+  SpinMutexLock l(&thread_list_mutex);
+  thread_stats.n_live_threads++;
+  thread_stats.total_stack_size += t->stack_size();
   if (!main_thread) {
     main_thread = t;
     return;
   }
-  SpinMutexLock l(&thread_list_mutex);
   Thread *last = main_thread;
   while (last->next_)
     last = last->next_;
@@ -43,6 +46,8 @@ void Thread::InsertIntoThreadList(Thread *t) {
 void Thread::RemoveFromThreadList(Thread *t) {
   CHECK_NE(t, main_thread);
   SpinMutexLock l(&thread_list_mutex);
+  thread_stats.n_live_threads--;
+  thread_stats.total_stack_size -= t->stack_size();
   Thread *prev = main_thread;
   Thread *cur = prev->next_;
   CHECK(cur);
@@ -67,10 +72,17 @@ void Thread::Create() {
   thread->random_state_ =
       flags()->random_tags ? RandomSeed() : thread->unique_id_;
   if (auto sz = flags()->heap_history_size)
-    thread->heap_allocations_ = RingBuffer<HeapAllocationRecord>::New(sz);
-  InsertIntoThreadList(thread);
+    thread->heap_allocations_ = HeapAllocationsRingBuffer::New(sz);
   SetCurrentThread(thread);
   thread->Init();
+  InsertIntoThreadList(thread);
+}
+
+uptr Thread::MemoryUsedPerThread() {
+  uptr res = sizeof(Thread);
+  if (auto sz = flags()->heap_history_size)
+    res += HeapAllocationsRingBuffer::SizeInBytes(sz);
+  return res;
 }
 
 void Thread::Init() {
diff --git a/lib/hwasan/hwasan_thread.h b/lib/hwasan/hwasan_thread.h
index 451baf3c0..6b41f51d8 100644
--- a/lib/hwasan/hwasan_thread.h
+++ b/lib/hwasan/hwasan_thread.h
@@ -26,6 +26,7 @@ class Thread {
 
   uptr stack_top() { return stack_top_; }
   uptr stack_bottom() { return stack_bottom_; }
+  uptr stack_size() { return stack_top() - stack_bottom(); }
   uptr tls_begin() { return tls_begin_; }
   uptr tls_end() { return tls_end_; }
   bool IsMainThread() { return unique_id_ == 0; }
@@ -75,6 +76,18 @@ class Thread {
     Print("Thread: ");
   }
 
+  struct ThreadStats {
+    uptr n_live_threads;
+    uptr total_stack_size;
+  };
+
+  static ThreadStats GetThreadStats() {
+    SpinMutexLock l(&thread_list_mutex);
+    return thread_stats;
+  }
+
+  static uptr MemoryUsedPerThread();
+
  private:
   // NOTE: There is no Thread constructor. It is allocated
   // via mmap() and *must* be valid in zero-initialized state.
@@ -101,6 +114,7 @@ class Thread {
   Thread *next_;  // All live threads form a linked list.
   static SpinMutex thread_list_mutex;
   static Thread *main_thread;
+  static ThreadStats thread_stats;
 
   u64 unique_id_;  // counting from zero.
 
diff --git a/lib/sanitizer_common/sanitizer_ring_buffer.h b/lib/sanitizer_common/sanitizer_ring_buffer.h
index 32f0d3fff..c4649c2be 100644
--- a/lib/sanitizer_common/sanitizer_ring_buffer.h
+++ b/lib/sanitizer_common/sanitizer_ring_buffer.h
@@ -39,6 +39,10 @@ class RingBuffer {
                                  2 * sizeof(T *));
   }
 
+  static uptr SizeInBytes(uptr Size) {
+    return Size * sizeof(T) + 2 * sizeof(T*);
+  }
+
   uptr SizeInBytes() { return SizeInBytes(size()); }
 
   void push(T t) {
@@ -62,10 +66,6 @@ class RingBuffer {
   ~RingBuffer() {}
   RingBuffer(const RingBuffer&) = delete;
 
-  static uptr SizeInBytes(uptr Size) {
-    return Size * sizeof(T) + 2 * sizeof(T*);
-  }
-
   // Data layout:
   // LNDDDDDDDD
   // D: data elements.
diff --git a/test/hwasan/TestCases/print-memory-usage.c b/test/hwasan/TestCases/print-memory-usage.c
new file mode 100644
index 000000000..c38457f65
--- /dev/null
+++ b/test/hwasan/TestCases/print-memory-usage.c
@@ -0,0 +1,71 @@
+// Tests __hwasan_print_memory_usage.
+// RUN: %clang_hwasan %s -o %t
+// RUN: ulimit -s 1000
+// RUN: %run %t 2>&1 | FileCheck %s
+// REQUIRES: stable-runtime
+
+#include <pthread.h>
+#include <stdlib.h>
+#include <stdio.h>
+#include <string.h>
+
+#include <sanitizer/hwasan_interface.h>
+
+int state;
+__thread volatile char *sink;
+
+__attribute__((noinline))
+void *malloc_and_use(int size) {
+  char *x = (char*)malloc(size);
+  for (int i = 0; i < size; i++)
+    x[i] = 42;  // make this memory used.
+  return x;
+}
+
+void *T1(void *arg) {
+
+  for (int i = 1; i <= (1 << 20); i *= 2)
+    sink = malloc_and_use(i);
+
+  __sync_fetch_and_add(&state, 1);
+  while (__sync_fetch_and_add(&state, 0) != 4) {}
+  return NULL;
+}
+
+void *T4(void *arg) { return NULL; }
+
+int main() {
+  __hwasan_enable_allocator_tagging();
+
+  __hwasan_print_memory_usage();
+  // CHECK: HWASAN pid: [[PID:[0-9]*]] rss: {{.*}} threads: 1 stacks: 1024000 thr_aux: {{.*}} stack_depot: 0 uniq_stacks: 0 heap: 0
+
+  void *one_meg = malloc_and_use(1 << 20);
+
+  __hwasan_print_memory_usage();
+  // CHECK: HWASAN pid: [[PID]] rss: {{.*}} threads: 1 stacks: 1024000 thr_aux: {{.*}} stack_depot: {{[1-9][0-9]*}} uniq_stacks: 1 heap: 1052672
+
+  free(one_meg);
+
+  __hwasan_print_memory_usage();
+  // CHECK: HWASAN pid: [[PID]] rss: {{.*}} threads: 1 stacks: 1024000 thr_aux: {{.*}} stack_depot: {{[1-9][0-9]*}} uniq_stacks: 2 heap: 0
+
+  pthread_t t1, t2, t3, t4;
+
+  pthread_create(&t1, NULL, T1, NULL);
+  pthread_create(&t2, NULL, T1, NULL);
+  pthread_create(&t3, NULL, T1, NULL);
+  pthread_create(&t4, NULL, T4, NULL);
+  while (__sync_fetch_and_add(&state, 0) != 3) {}
+  pthread_join(t4, NULL);
+
+  __hwasan_print_memory_usage();
+  // CHECK: HWASAN pid: [[PID]] rss: {{.*}} threads: 4 stacks: 4083520
+
+  __sync_fetch_and_add(&state, 1);
+  pthread_join(t1, NULL);
+  pthread_join(t2, NULL);
+  pthread_join(t3, NULL);
+  __hwasan_print_memory_usage();
+  // CHECK: HWASAN pid: [[PID]] rss: {{.*}} threads: 1 stacks: 1024000
+}