[msan] Refactor allocator instantiation

[stephan.yichao.zhao]

Summary

This is to support D101204.

We introduced a new class AllocatorAPI that can be shared with MSan and DFSan.

[vitalybuka]

compiler-rt/lib/sanitizer_common/sanitizer_allocator_api.h

72

How we going to be sure that it's still linker initialized...

[stephan.yichao.zhao]

update

[stephan.yichao.zhao]

compiler-rt/lib/sanitizer_common/sanitizer_allocator_api.h

72

This template class with virtual method will be extended by clients.
For example, MSan defines a child class MSanAllocatorAPI, then in msan_allocator.cpp, we declare a static variable as before:

static MSanAllocatorAPI allocator_api("MemorySanitizer", kMaxAllowedMallocSize);

[vitalybuka]

Regarding the main point of the patch, I understand that's you are doing this to avoid duplicating code.
As negatives of this approach:

• another abstraction layer
• fixing one sanitizer we will have to think about affect on another

Probably check with @eugenis before fixing my comments if he is ok with this approach.

compiler-rt/lib/sanitizer_common/sanitizer_allocator_api.h

52

please use "using"

55

I guess you can just make it a stadalone function and drop two template arguments from class

template <class Thread, class ThreadLocalMallocStorage>
ThreadLocalMallocStorage *GetAllocatorCache(Thread *t);

62–69

Could you avoid victuals with CRTP, it's template anyway?
Or even better simply assume that such standalone function exist, because they probably do not need to access internals of AllocatorAPI

72

Initialized concern is not about virtuals.
Pre C++20 it's hard to be sure that global of complex type is linker initialized, we just investigated the scudo crash when two globals were linker initialized, but when we moved one inside of another compiler switched to dynamic initialization.

That's should not block the patch, but it's the reason to carefully watch msan bots or reports of unexplained crashes.

152

The point of Metadata to keep sanitizer_common unaware about content of sanitizer specific info there.
This patch breakes this.

In case of DFSan and msan, they have very similar info there, but it's still not nice.

What I see is that you just need something like this on sanitizer derived side:

uptr GetRequestedSize(uptr ptr) {
  return allocator.GetMetaData(ptr)->requested_size;
}

void SetRequestedSize(uptr ptr, uptr size) {
  allocator.GetMetaData(ptr)->requested_size = size;
}

Then you can throw away another template arg Metadata.

I would probably go even to standalone functions:

template<class Allocator>
uptr GetRequestedSize(Allocator& allocator, uptr ptr) {
  return allocator.GetMetaData(ptr)->requested_size;
}

template<class Allocator>
void SetRequestedSize(Allocator& allocator, uptr ptr, uptr size) {
  allocator.GetMetaData(ptr)->requested_size = size;
}

329–332

to my taste its more readable to keep methods definitions inside of the class

[eugenis]

I'm not sure if this is worth the code savings.

One way of thinking about it - can it be used in ASan? TSan? Usually, you do not start factoring out common parts until you are about to add a 3rd copy of a thing.

Perhaps a more incremental approach would be better. For example, most of the Allocate method - handling of allocator cache, fallback allocator, out of memory check - is clearly common code, but Reallocate logic seems a bit too complicated and sanitizer-specific. I especially dislike the hooks, the make the code hard to follow.

requested_size works great as a common allocator concept, but I think you'll have issues if you try applying this refactoring to ASan, which does not use metadata.

compiler-rt/lib/sanitizer_common/sanitizer_allocator_api.h

143

Are these hooks really necessary? If it is called at the end of the function (or at the beginning, like DeallocateHook) MSan/DFSan could wrap the common method instead.

[stephan.yichao.zhao]

Thank you for the feedback. I feel the following takeaways make sense

• another abstraction layer
• fixing one sanitizer we will have to think about affect on another
• can it be used in ASan? TSan? ...: I think it is not straightforward after looking into their use cases

Although Allocate and Deallocate have some similar code, but because of Thread and ThreadLocalMallocStorage, those 10-20 lines of code need a template too...

To support D101204, adding a separate allocator seems better because of the above trade-offs. I would deprecate this one.

compiler-rt/lib/msan/msan.h

 void *msan_realloc(void *ptr, uptr size, StackTrace *stack);
 void *msan_reallocarray(void *ptr, uptr nmemb, uptr size, StackTrace *stack);
 void *msan_valloc(uptr size, StackTrace *stack);
 void *msan_pvalloc(uptr size, StackTrace *stack);
 void *msan_aligned_alloc(uptr alignment, uptr size, StackTrace *stack);
 void *msan_memalign(uptr alignment, uptr size, StackTrace *stack);
 int msan_posix_memalign(void **memptr, uptr alignment, uptr size,
                         StackTrace *stack);
+void msan_dealloc(void *ptr, StackTrace *stack);
 
 void InstallTrapHandler();
 void InstallAtExitHandler();
 
 const char *GetStackOriginDescr(u32 id, uptr *pc);
 
 void EnterSymbolizer();
 void ExitSymbolizer();

compiler-rt/lib/msan/msan_allocator.cpp

 //===-- msan_allocator.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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of MemorySanitizer.
 //
 // MemorySanitizer allocator.
 //===----------------------------------------------------------------------===//
 
-#include "sanitizer_common/sanitizer_allocator.h"
-#include "sanitizer_common/sanitizer_allocator_checks.h"
-#include "sanitizer_common/sanitizer_allocator_interface.h"
-#include "sanitizer_common/sanitizer_allocator_report.h"
-#include "sanitizer_common/sanitizer_errno.h"
-#include "msan.h"
 #include "msan_allocator.h"
+
+#include "msan.h"
 #include "msan_origin.h"
-#include "msan_thread.h"
 #include "msan_poisoning.h"
+#include "msan_thread.h"
+#include "sanitizer_common/sanitizer_allocator.h"
+#include "sanitizer_common/sanitizer_allocator_api.h"
+#include "sanitizer_common/sanitizer_allocator_interface.h"
 
 namespace __msan {
 
 struct Metadata {
   uptr requested_size;
 };
 
 struct MsanMapUnmapCallback {
struct AP32 {
   typedef __sanitizer::CompactSizeClassMap SizeClassMap;
   static const uptr kRegionSizeLog = 20;
   using AddressSpaceView = LocalAddressSpaceView;
   typedef MsanMapUnmapCallback MapUnmapCallback;
   static const uptr kFlags = 0;
 };
 typedef SizeClassAllocator32<AP32> PrimaryAllocator;
 #endif
-typedef CombinedAllocator<PrimaryAllocator> Allocator;
-typedef Allocator::AllocatorCache AllocatorCache;
 
-static Allocator allocator;
-static AllocatorCache fallback_allocator_cache;
-static StaticSpinMutex fallback_mutex;
-
-static uptr max_malloc_size;
-
-void MsanAllocatorInit() {
-  SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
-  allocator.Init(common_flags()->allocator_release_to_os_interval_ms);
-  if (common_flags()->max_allocation_size_mb)
-    max_malloc_size = Min(common_flags()->max_allocation_size_mb << 20,
-                          kMaxAllowedMallocSize);
-  else
-    max_malloc_size = kMaxAllowedMallocSize;
-}
-
-AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
-  CHECK(ms);
-  CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));
-  return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
-}
-
-void MsanThreadLocalMallocStorage::CommitBack() {
-  allocator.SwallowCache(GetAllocatorCache(this));
-}
-
-static void *MsanAllocate(StackTrace *stack, uptr size, uptr alignment,
-                          bool zeroise) {
-  if (size > max_malloc_size) {
-    if (AllocatorMayReturnNull()) {
-      Report("WARNING: MemorySanitizer failed to allocate 0x%zx bytes\n", size);
-      return nullptr;
-    }
-    ReportAllocationSizeTooBig(size, max_malloc_size, stack);
-  }
-  MsanThread *t = GetCurrentThread();
-  void *allocated;
-  if (t) {
-    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
-    allocated = allocator.Allocate(cache, size, alignment);
-  } else {
-    SpinMutexLock l(&fallback_mutex);
-    AllocatorCache *cache = &fallback_allocator_cache;
-    allocated = allocator.Allocate(cache, size, alignment);
-  }
-  if (UNLIKELY(!allocated)) {
-    SetAllocatorOutOfMemory();
-    if (AllocatorMayReturnNull())
-      return nullptr;
-    ReportOutOfMemory(size, stack);
-  }
-  Metadata *meta =
-      reinterpret_cast<Metadata *>(allocator.GetMetaData(allocated));
-  meta->requested_size = size;
+class MSanAllocatorAPI
+    : public AllocatorAPI<PrimaryAllocator, MsanThread,
+                          MsanThreadLocalMallocStorage, Metadata> {
+ public:
+  MSanAllocatorAPI(const char *sanitizer_name_, uptr max_allowed_malloc_size_)
+      : AllocatorAPI<PrimaryAllocator, MsanThread, MsanThreadLocalMallocStorage,
+                     Metadata>(sanitizer_name_, max_allowed_malloc_size_) {}
+
+ private:
+  MsanThread *GetCurrentThread() override { return __msan::GetCurrentThread(); }
+  void AllocateHook(StackTrace *stack, void *allocated, uptr size,
+                    bool zeroise) override {
     if (zeroise) {
       __msan_clear_and_unpoison(allocated, size);
     } else if (flags()->poison_in_malloc) {
       __msan_poison(allocated, size);
       if (__msan_get_track_origins()) {
         stack->tag = StackTrace::TAG_ALLOC;
         Origin o = Origin::CreateHeapOrigin(stack);
         __msan_set_origin(allocated, size, o.raw_id());
       }
     }
     MSAN_MALLOC_HOOK(allocated, size);
-  return allocated;
   }
-
-void MsanDeallocate(StackTrace *stack, void *p) {
-  CHECK(p);
+  void DeallocateHook(StackTrace *stack, void *p, uptr size) override {
     MSAN_FREE_HOOK(p);
-  Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(p));
-  uptr size = meta->requested_size;
-  meta->requested_size = 0;
     // This memory will not be reused by anyone else, so we are free to keep it
     // poisoned.
     if (flags()->poison_in_free) {
       __msan_poison(p, size);
       if (__msan_get_track_origins()) {
         stack->tag = StackTrace::TAG_DEALLOC;
         Origin o = Origin::CreateHeapOrigin(stack);
         __msan_set_origin(p, size, o.raw_id());
       }
     }
-  MsanThread *t = GetCurrentThread();
-  if (t) {
-    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
-    allocator.Deallocate(cache, p);
-  } else {
-    SpinMutexLock l(&fallback_mutex);
-    AllocatorCache *cache = &fallback_allocator_cache;
-    allocator.Deallocate(cache, p);
   }
-}
-
-static void *MsanReallocate(StackTrace *stack, void *old_p, uptr new_size,
-                            uptr alignment) {
-  Metadata *meta = reinterpret_cast<Metadata*>(allocator.GetMetaData(old_p));
-  uptr old_size = meta->requested_size;
-  uptr actually_allocated_size = allocator.GetActuallyAllocatedSize(old_p);
-  if (new_size <= actually_allocated_size) {
-    // We are not reallocating here.
-    meta->requested_size = new_size;
-    if (new_size > old_size) {
-      if (flags()->poison_in_malloc) {
-        stack->tag = StackTrace::TAG_ALLOC;
-        PoisonMemory((char *)old_p + old_size, new_size - old_size, stack);
-      }
-    }
-    return old_p;
-  }
-  uptr memcpy_size = Min(new_size, old_size);
-  void *new_p = MsanAllocate(stack, new_size, alignment, false /*zeroise*/);
-  if (new_p) {
+  void RellocateReuseHook(StackTrace *stack, void *old_p, uptr old_size,
+                          uptr new_size) override {
+  void RellocateNewHook(StackTrace *stack, void *old_p, void *new_p,
+                        uptr memcpy_size) override {
     CopyMemory(new_p, old_p, memcpy_size, stack);
-    MsanDeallocate(stack, old_p);
-  }
-  return new_p;
   }
+};
 
-static void *MsanCalloc(StackTrace *stack, uptr nmemb, uptr size) {
-  if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
-    if (AllocatorMayReturnNull())
-      return nullptr;
-    ReportCallocOverflow(nmemb, size, stack);
-  }
-  return MsanAllocate(stack, nmemb * size, sizeof(u64), true);
-}
+static MSanAllocatorAPI allocator_api("MemorySanitizer", kMaxAllowedMallocSize);
+
+void MsanAllocatorInit() { allocator_api.Init(); }
 
-static uptr AllocationSize(const void *p) {
-  if (!p) return 0;
-  const void *beg = allocator.GetBlockBegin(p);
-  if (beg != p) return 0;
-  Metadata *b = (Metadata *)allocator.GetMetaData(p);
-  return b->requested_size;
+void MsanThreadLocalMallocStorage::CommitBack() {
+  allocator_api.CommitBack(this);
 }
 
 void *msan_malloc(uptr size, StackTrace *stack) {
-  return SetErrnoOnNull(MsanAllocate(stack, size, sizeof(u64), false));
+  return allocator_api.Malloc(size, stack);
 }
 
 void *msan_calloc(uptr nmemb, uptr size, StackTrace *stack) {
-  return SetErrnoOnNull(MsanCalloc(stack, nmemb, size));
+  return allocator_api.Calloc(nmemb, size, stack);
 }
 
 void *msan_realloc(void *ptr, uptr size, StackTrace *stack) {
-  if (!ptr)
-    return SetErrnoOnNull(MsanAllocate(stack, size, sizeof(u64), false));
-  if (size == 0) {
-    MsanDeallocate(stack, ptr);
-    return nullptr;
-  }
-  return SetErrnoOnNull(MsanReallocate(stack, ptr, size, sizeof(u64)));
+  return allocator_api.Realloc(ptr, size, stack);
 }
 
 void *msan_reallocarray(void *ptr, uptr nmemb, uptr size, StackTrace *stack) {
-  if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
-    errno = errno_ENOMEM;
-    if (AllocatorMayReturnNull())
-      return nullptr;
-    ReportReallocArrayOverflow(nmemb, size, stack);
-  }
-  return msan_realloc(ptr, nmemb * size, stack);
+  return allocator_api.Reallocarray(ptr, nmemb, size, stack);
 }
 
 void *msan_valloc(uptr size, StackTrace *stack) {
-  return SetErrnoOnNull(MsanAllocate(stack, size, GetPageSizeCached(), false));
+  return allocator_api.Valloc(size, stack);
 }
 
 void *msan_pvalloc(uptr size, StackTrace *stack) {
-  uptr PageSize = GetPageSizeCached();
-  if (UNLIKELY(CheckForPvallocOverflow(size, PageSize))) {
-    errno = errno_ENOMEM;
-    if (AllocatorMayReturnNull())
-      return nullptr;
-    ReportPvallocOverflow(size, stack);
-  }
-  // pvalloc(0) should allocate one page.
-  size = size ? RoundUpTo(size, PageSize) : PageSize;
-  return SetErrnoOnNull(MsanAllocate(stack, size, PageSize, false));
+  return allocator_api.Pvalloc(size, stack);
 }
 
 void *msan_aligned_alloc(uptr alignment, uptr size, StackTrace *stack) {
-  if (UNLIKELY(!CheckAlignedAllocAlignmentAndSize(alignment, size))) {
-    errno = errno_EINVAL;
-    if (AllocatorMayReturnNull())
-      return nullptr;
-    ReportInvalidAlignedAllocAlignment(size, alignment, stack);
-  }
-  return SetErrnoOnNull(MsanAllocate(stack, size, alignment, false));
+  return allocator_api.AlignedAlloc(alignment, size, stack);
 }
 
 void *msan_memalign(uptr alignment, uptr size, StackTrace *stack) {
-  if (UNLIKELY(!IsPowerOfTwo(alignment))) {
-    errno = errno_EINVAL;
-    if (AllocatorMayReturnNull())
-      return nullptr;
-    ReportInvalidAllocationAlignment(alignment, stack);
-  }
-  return SetErrnoOnNull(MsanAllocate(stack, size, alignment, false));
+  return allocator_api.Memalign(alignment, size, stack);
 }
 
 int msan_posix_memalign(void **memptr, uptr alignment, uptr size,
                         StackTrace *stack) {
-  if (UNLIKELY(!CheckPosixMemalignAlignment(alignment))) {
-    if (AllocatorMayReturnNull())
-      return errno_EINVAL;
-    ReportInvalidPosixMemalignAlignment(alignment, stack);
-  }
-  void *ptr = MsanAllocate(stack, size, alignment, false);
-  if (UNLIKELY(!ptr))
-    // OOM error is already taken care of by MsanAllocate.
-    return errno_ENOMEM;
-  CHECK(IsAligned((uptr)ptr, alignment));
-  *memptr = ptr;
-  return 0;
+  return allocator_api.PosixMemalign(memptr, alignment, size, stack);
+}
+
+void msan_dealloc(void *p, StackTrace *stack) {
+  return allocator_api.Deallocate(p, stack);
 }
 
 } // namespace __msan
 
 using namespace __msan;
 
 uptr __sanitizer_get_current_allocated_bytes() {
-  uptr stats[AllocatorStatCount];
-  allocator.GetStats(stats);
-  return stats[AllocatorStatAllocated];
+  return allocator_api.GetCurrentAllocatedBytes();
 }
 
-uptr __sanitizer_get_heap_size() {
-  uptr stats[AllocatorStatCount];
-  allocator.GetStats(stats);
-  return stats[AllocatorStatMapped];
-}
+uptr __sanitizer_get_heap_size() { return allocator_api.GetHeapSize(); }
 
 uptr __sanitizer_get_free_bytes() { return 1; }
 
 uptr __sanitizer_get_unmapped_bytes() { return 1; }
 
 uptr __sanitizer_get_estimated_allocated_size(uptr size) { return size; }
 
-int __sanitizer_get_ownership(const void *p) { return AllocationSize(p) != 0; }
+int __sanitizer_get_ownership(const void *p) {
+  return allocator_api.AllocationSize(p) != 0;
+}
 
-uptr __sanitizer_get_allocated_size(const void *p) { return AllocationSize(p); }
+uptr __sanitizer_get_allocated_size(const void *p) {
+  return allocator_api.AllocationSize(p);
+}

compiler-rt/lib/msan/msan_interceptors.cpp

 #define MSAN_MAYBE_INTERCEPT_PVALLOC INTERCEPT_FUNCTION(pvalloc)
 #else
 #define MSAN_MAYBE_INTERCEPT_PVALLOC
 #endif
 
 INTERCEPTOR(void, free, void *ptr) {
   GET_MALLOC_STACK_TRACE;
   if (!ptr || UNLIKELY(IsInDlsymAllocPool(ptr))) return;
-  MsanDeallocate(&stack, ptr);
+  msan_dealloc(ptr, &stack);
 }
 
 #if !SANITIZER_FREEBSD && !SANITIZER_NETBSD
 INTERCEPTOR(void, cfree, void *ptr) {
   GET_MALLOC_STACK_TRACE;
   if (!ptr || UNLIKELY(IsInDlsymAllocPool(ptr))) return;
-  MsanDeallocate(&stack, ptr);
+  msan_dealloc(ptr, &stack);
 }
 #define MSAN_MAYBE_INTERCEPT_CFREE INTERCEPT_FUNCTION(cfree)
 #else
 #define MSAN_MAYBE_INTERCEPT_CFREE
 #endif
 
 #if !SANITIZER_NETBSD
 INTERCEPTOR(uptr, malloc_usable_size, void *ptr) {

compiler-rt/lib/msan/msan_new_delete.cpp

 INTERCEPTOR_ATTRIBUTE
 void *operator new(size_t size, std::align_val_t align, std::nothrow_t const&)
 { OPERATOR_NEW_BODY_ALIGN(true /*nothrow*/); }
 INTERCEPTOR_ATTRIBUTE
 void *operator new[](size_t size, std::align_val_t align, std::nothrow_t const&)
 { OPERATOR_NEW_BODY_ALIGN(true /*nothrow*/); }
 
 #define OPERATOR_DELETE_BODY \
-  GET_MALLOC_STACK_TRACE; \
-  if (ptr) MsanDeallocate(&stack, ptr)
+  GET_MALLOC_STACK_TRACE;    \
+  if (ptr)                   \
+  msan_dealloc(ptr, &stack)
 
 INTERCEPTOR_ATTRIBUTE
 void operator delete(void *ptr) NOEXCEPT { OPERATOR_DELETE_BODY; }
 INTERCEPTOR_ATTRIBUTE
 void operator delete[](void *ptr) NOEXCEPT { OPERATOR_DELETE_BODY; }
 INTERCEPTOR_ATTRIBUTE
 void operator delete(void *ptr, std::nothrow_t const&) { OPERATOR_DELETE_BODY; }
 INTERCEPTOR_ATTRIBUTE

compiler-rt/lib/sanitizer_common/CMakeLists.txt

 # Explicitly list all sanitizer_common headers. Not all of these are
 # included in sanitizer_common source files, but we need to depend on
 # headers when building our custom unit tests.
 set(SANITIZER_IMPL_HEADERS
   sancov_flags.h
   sancov_flags.inc
   sanitizer_addrhashmap.h
   sanitizer_allocator.h
+  sanitizer_allocator_api.h
   sanitizer_allocator_bytemap.h
   sanitizer_allocator_checks.h
   sanitizer_allocator_combined.h
   sanitizer_allocator_interface.h
   sanitizer_allocator_internal.h
   sanitizer_allocator_local_cache.h
   sanitizer_allocator_primary32.h
   sanitizer_allocator_primary64.h

compiler-rt/lib/sanitizer_common/sanitizer_allocator_api.h

+//===-- sanitizer_allocator_api.h -------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// Shared allocator API definitions for MemorySanitizer, DataFlowSanitizer.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SANITIZER_ALLOCATOR_API_H
+#define SANITIZER_ALLOCATOR_API_H
+
+#include "sanitizer_common/sanitizer_allocator.h"
+#include "sanitizer_common/sanitizer_allocator_checks.h"
+#include "sanitizer_common/sanitizer_allocator_report.h"
+#include "sanitizer_common/sanitizer_errno.h"
+#include "sanitizer_common/sanitizer_stacktrace.h"
+
+namespace __sanitizer {
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+class AllocatorAPI {
+ public:
+  AllocatorAPI(const char *sanitizer_name_, uptr max_allowed_malloc_size_)
+      : sanitizer_name(sanitizer_name_),
+        kMaxAllowedMallocSize(max_allowed_malloc_size_) {}
+  void Init();
+
+  void *Malloc(uptr size, StackTrace *stack = nullptr);
+  void Deallocate(void *p, StackTrace *stack = nullptr);
+  void *Calloc(uptr nmemb, uptr size, StackTrace *stack = nullptr);
+  void *Realloc(void *ptr, uptr size, StackTrace *stack = nullptr);
+  void *Reallocarray(void *ptr, uptr nmemb, uptr size,
+                     StackTrace *stack = nullptr);
+  void *Valloc(uptr size, StackTrace *stack = nullptr);
+  void *Pvalloc(uptr size, StackTrace *stack = nullptr);
+  void *AlignedAlloc(uptr alignment, uptr size, StackTrace *stack = nullptr);
+  void *Memalign(uptr alignment, uptr size, StackTrace *stack = nullptr);
+  int PosixMemalign(void **memptr, uptr alignment, uptr size,
+                    StackTrace *stack = nullptr);
+
+  void CommitBack(ThreadLocalMallocStorage *ms);
+  uptr AllocationSize(const void *p);
+  uptr GetCurrentAllocatedBytes();
+  uptr GetHeapSize();
+
+ private:
+  typedef CombinedAllocator<PrimaryAllocator> Allocator;
+  typedef typename Allocator::AllocatorCache AllocatorCache;
+
+  AllocatorCache *GetAllocatorCache(ThreadLocalMallocStorage *ms);
+
+  void *Allocate(StackTrace *stack, uptr size, uptr alignment, bool zeroise);
+  void *InternalCalloc(StackTrace *stack, uptr nmemb, uptr size);
+  void *InternalReallocate(StackTrace *stack, void *old_p, uptr new_size,
+                           uptr alignment);
+
+  virtual Thread *GetCurrentThread() = 0;
+  virtual void AllocateHook(StackTrace *stack, void *allocated, uptr size,
+                            bool zeroise) = 0;
+  virtual void DeallocateHook(StackTrace *stack, void *p, uptr size) = 0;
+  virtual void RellocateReuseHook(StackTrace *stack, void *old_p, uptr old_size,
+                                  uptr new_size) = 0;
+  virtual void RellocateNewHook(StackTrace *stack, void *old_p, void *new_p,
+                                uptr memcpy_size) = 0;
+
+  Allocator allocator;
+  AllocatorCache fallback_allocator_cache;
+  StaticSpinMutex fallback_mutex;
+  uptr max_malloc_size;
+
+  const char *sanitizer_name;
+  const uptr kMaxAllowedMallocSize;
+};
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                  Metadata>::Init() {
+  SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
+  allocator.Init(common_flags()->allocator_release_to_os_interval_ms);
+  if (common_flags()->max_allocation_size_mb)
+    max_malloc_size = Min(common_flags()->max_allocation_size_mb << 20,
+                          kMaxAllowedMallocSize);
+  else
+    max_malloc_size = kMaxAllowedMallocSize;
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+typename AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                      Metadata>::AllocatorCache *
+AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+             Metadata>::GetAllocatorCache(ThreadLocalMallocStorage *ms) {
+  CHECK(ms);
+  CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));
+  return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                  Metadata>::CommitBack(ThreadLocalMallocStorage *ms) {
+  allocator.SwallowCache(GetAllocatorCache(ms));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Allocate(StackTrace *stack, uptr size,
+                                       uptr alignment, bool zeroise) {
+  if (size > max_malloc_size) {
+    if (AllocatorMayReturnNull()) {
+      Report("WARNING: %s failed to allocate 0x%zx bytes\n", sanitizer_name,
+             size);
+      return nullptr;
+    }
+    ReportAllocationSizeTooBig(size, max_malloc_size, stack);
+  }
+  Thread *t = GetCurrentThread();
+  void *allocated;
+  if (t) {
+    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
+    allocated = allocator.Allocate(cache, size, alignment);
+  } else {
+    SpinMutexLock l(&fallback_mutex);
+    AllocatorCache *cache = &fallback_allocator_cache;
+    allocated = allocator.Allocate(cache, size, alignment);
+  }
+  if (UNLIKELY(!allocated)) {
+    SetAllocatorOutOfMemory();
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    ReportOutOfMemory(size, stack);
+  }
+  Metadata *meta =
+      reinterpret_cast<Metadata *>(allocator.GetMetaData(allocated));
+  meta->requested_size = size;
+  AllocateHook(stack, allocated, size, zeroise);
+  return allocated;
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                  Metadata>::Deallocate(void *p, StackTrace *stack) {
+  CHECK(p);
+  Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(p));
+  uptr size = meta->requested_size;
+  DeallocateHook(stack, p, size);
+  meta->requested_size = 0;
+  Thread *t = GetCurrentThread();
+  if (t) {
+    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
+    allocator.Deallocate(cache, p);
+  } else {
+    SpinMutexLock l(&fallback_mutex);
+    AllocatorCache *cache = &fallback_allocator_cache;
+    allocator.Deallocate(cache, p);
+  }
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::InternalReallocate(StackTrace *stack, void *old_p,
+                                                 uptr new_size,
+                                                 uptr alignment) {
+  Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(old_p));
+  uptr old_size = meta->requested_size;
+  uptr actually_allocated_size = allocator.GetActuallyAllocatedSize(old_p);
+  if (new_size <= actually_allocated_size) {
+    // We are not reallocating here.
+    meta->requested_size = new_size;
+    if (new_size > old_size)
+      RellocateReuseHook(stack, old_p, old_size, new_size);
+    return old_p;
+  }
+  uptr memcpy_size = Min(new_size, old_size);
+  void *new_p = Allocate(stack, new_size, alignment, false /*zeroise*/);
+  if (new_p) {
+    RellocateNewHook(stack, old_p, new_p, memcpy_size);
+    Deallocate(old_p, stack);
+  }
+  return new_p;
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::InternalCalloc(StackTrace *stack, uptr nmemb,
+                                             uptr size) {
+  if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    ReportCallocOverflow(nmemb, size, stack);
+  }
+  return Allocate(stack, nmemb * size, sizeof(u64), true);
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+uptr AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                  Metadata>::AllocationSize(const void *p) {
+  if (!p)
+    return 0;
+  const void *beg = allocator.GetBlockBegin(p);
+  if (beg != p)
+    return 0;
+  Metadata *b = (Metadata *)allocator.GetMetaData(p);
+  return b->requested_size;
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Malloc(uptr size, StackTrace *stack) {
+  return SetErrnoOnNull(Allocate(stack, size, sizeof(u64), false));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Calloc(uptr nmemb, uptr size, StackTrace *stack) {
+  return SetErrnoOnNull(InternalCalloc(stack, nmemb, size));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Realloc(void *ptr, uptr size, StackTrace *stack) {
+  if (!ptr)
+    return SetErrnoOnNull(Allocate(stack, size, sizeof(u64), false));
+  if (size == 0) {
+    Deallocate(ptr, stack);
+    return nullptr;
+  }
+  return SetErrnoOnNull(InternalReallocate(stack, ptr, size, sizeof(u64)));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Reallocarray(void *ptr, uptr nmemb, uptr size,
+                                           StackTrace *stack) {
+  if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
+    errno = errno_ENOMEM;
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    ReportReallocArrayOverflow(nmemb, size, stack);
+  }
+  return Realloc(ptr, nmemb * size, stack);
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Valloc(uptr size, StackTrace *stack) {
+  return SetErrnoOnNull(Allocate(stack, size, GetPageSizeCached(), false));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Pvalloc(uptr size, StackTrace *stack) {
+  uptr PageSize = GetPageSizeCached();
+  if (UNLIKELY(CheckForPvallocOverflow(size, PageSize))) {
+    errno = errno_ENOMEM;
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    ReportPvallocOverflow(size, stack);
+  }
+  // pvalloc(0) should allocate one page.
+  size = size ? RoundUpTo(size, PageSize) : PageSize;
+  return SetErrnoOnNull(Allocate(stack, size, PageSize, false));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::AlignedAlloc(uptr alignment, uptr size,
+                                           StackTrace *stack) {
+  if (UNLIKELY(!CheckAlignedAllocAlignmentAndSize(alignment, size))) {
+    errno = errno_EINVAL;
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    ReportInvalidAlignedAllocAlignment(size, alignment, stack);
+  }
+  return SetErrnoOnNull(Allocate(stack, size, alignment, false));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+void *AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                   Metadata>::Memalign(uptr alignment, uptr size,
+                                       StackTrace *stack) {
+  if (UNLIKELY(!IsPowerOfTwo(alignment))) {
+    errno = errno_EINVAL;
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    ReportInvalidAllocationAlignment(alignment, stack);
+  }
+  return SetErrnoOnNull(Allocate(stack, size, alignment, false));
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+int AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                 Metadata>::PosixMemalign(void **memptr, uptr alignment,
+                                          uptr size, StackTrace *stack) {
+  if (UNLIKELY(!CheckPosixMemalignAlignment(alignment))) {
+    if (AllocatorMayReturnNull())
+      return errno_EINVAL;
+    ReportInvalidPosixMemalignAlignment(alignment, stack);
+  }
+  void *ptr = Allocate(stack, size, alignment, false);
+  if (UNLIKELY(!ptr))
+    // OOM error is already taken care of by MsanAllocate.
+    return errno_ENOMEM;
+  CHECK(IsAligned((uptr)ptr, alignment));
+  *memptr = ptr;
+  return 0;
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+uptr AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                  Metadata>::GetCurrentAllocatedBytes() {
+  uptr stats[AllocatorStatCount];
+  allocator.GetStats(stats);
+  return stats[AllocatorStatAllocated];
+}
+
+template <class PrimaryAllocator, class Thread, class ThreadLocalMallocStorage,
+          class Metadata>
+uptr AllocatorAPI<PrimaryAllocator, Thread, ThreadLocalMallocStorage,
+                  Metadata>::GetHeapSize() {
+  uptr stats[AllocatorStatCount];
+  allocator.GetStats(stats);
+  return stats[AllocatorStatMapped];
+}
+
+}  // namespace __sanitizer
+
+#endif  // SANITIZER_ALLOCATOR_API_H