From 4305f640f0e5d4c27b7787407bc84c80552fed11 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Wed, 24 May 2023 10:35:34 -0700
Subject: [PATCH] [msan] Implement __sanitizer_get_current_allocated_bytes

__sanitizer_get_current_allocated_bytes had as body, but allocator
caches were not registered to collect stats. It's done by
SizeClassAllocator64LocalCache::Init().

Reviewed By: kstoimenov

Differential Revision: https://reviews.llvm.org/D151352
---
 compiler-rt/lib/msan/msan_allocator.cpp                            | 7 +++++++
 compiler-rt/lib/msan/msan_allocator.h                              | 1 +
 compiler-rt/lib/msan/msan_thread.cpp                               | 1 +
 .../test/sanitizer_common/TestCases/allocator_interface.cpp        | 2 +-
 4 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/compiler-rt/lib/msan/msan_allocator.cpp b/compiler-rt/lib/msan/msan_allocator.cpp
index c32a029..96fdf7b 100644
--- a/compiler-rt/lib/msan/msan_allocator.cpp
+++ b/compiler-rt/lib/msan/msan_allocator.cpp
@@ -145,8 +145,13 @@ AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
   return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
 }
 
+void MsanThreadLocalMallocStorage::Init() {
+  allocator.InitCache(GetAllocatorCache(this));
+}
+
 void MsanThreadLocalMallocStorage::CommitBack() {
   allocator.SwallowCache(GetAllocatorCache(this));
+  allocator.DestroyCache(GetAllocatorCache(this));
 }
 
 static void *MsanAllocate(StackTrace *stack, uptr size, uptr alignment,
@@ -393,3 +398,5 @@ const void *__sanitizer_get_allocated_begin(const void *p) {
 }
 
 uptr __sanitizer_get_allocated_size(const void *p) { return AllocationSize(p); }
+
+void __sanitizer_purge_allocator() { allocator.ForceReleaseToOS(); }
diff --git a/compiler-rt/lib/msan/msan_allocator.h b/compiler-rt/lib/msan/msan_allocator.h
index 365af4d..364331d 100644
--- a/compiler-rt/lib/msan/msan_allocator.h
+++ b/compiler-rt/lib/msan/msan_allocator.h
@@ -20,6 +20,7 @@ namespace __msan {
 struct MsanThreadLocalMallocStorage {
   // Allocator cache contains atomic_uint64_t which must be 8-byte aligned.
   ALIGNED(8) uptr allocator_cache[96 * (512 * 8 + 16)];  // Opaque.
+  void Init();
   void CommitBack();
 
  private:
diff --git a/compiler-rt/lib/msan/msan_thread.cpp b/compiler-rt/lib/msan/msan_thread.cpp
index 40ad6a5..ff9b90b 100644
--- a/compiler-rt/lib/msan/msan_thread.cpp
+++ b/compiler-rt/lib/msan/msan_thread.cpp
@@ -47,6 +47,7 @@ void MsanThread::Init() {
   CHECK(MEM_IS_APP(stack_.bottom));
   CHECK(MEM_IS_APP(stack_.top - 1));
   ClearShadowForThreadStackAndTLS();
+  malloc_storage().Init();
 }
 
 void MsanThread::TSDDtor(void *tsd) {
diff --git a/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp b/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp
index d43c47f..e086a3c 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp
+++ b/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp
@@ -5,7 +5,7 @@
 // UNSUPPORTED: ubsan
 
 // FIXME: implementation is incomplete.
-// XFAIL: msan, lsan, hwasan
+// XFAIL: lsan, hwasan
 
 #include <assert.h>
 #include <sanitizer/allocator_interface.h>
-- 
2.7.4