[Sanitizer] Introduce Allocator::may_return_null bool flag.

[samsonov]

Summary

Turn "allocator_may_return_null" common flag into an
Allocator::may_return_null bool flag. We want to make sure
that common flags are immutable after initialization. There
are cases when we want to change this flag in the allocator
at runtime: e.g. in unit tests and during ASan activation
on Android.

[kcc]

LGTM
Please try to test this on something other than just the reg tests

[eugenis]

LGTM

[samsonov]

Ran some more tests, didn't observe any problems. Submitting this now, thanks for the review!

lib/asan/asan_activation.cc

 }
 
 void AsanActivate() {
   if (!asan_is_deactivated) return;
   VReport(1, "Activating ASan\n");
 
   // Restore flag values.
   // FIXME: this is not atomic, and there may be other threads alive.
-  flags()->quarantine_size = asan_deactivated_flags.quarantine_size;
   flags()->max_redzone = asan_deactivated_flags.max_redzone;
   flags()->poison_heap = asan_deactivated_flags.poison_heap;
   common_flags()->malloc_context_size =
       asan_deactivated_flags.malloc_context_size;
   flags()->alloc_dealloc_mismatch =
       asan_deactivated_flags.alloc_dealloc_mismatch;
-  common_flags()->allocator_may_return_null =
-      asan_deactivated_flags.allocator_may_return_null;
 
   ParseExtraActivationFlags();
 
-  ReInitializeAllocator();
+  ReInitializeAllocator(asan_deactivated_flags.allocator_may_return_null,
+                        asan_deactivated_flags.quarantine_size);
 
   asan_is_deactivated = false;
-  VReport(
-      1,
-      "quarantine_size %d, max_redzone %d, poison_heap %d, "
+  VReport(1, "quarantine_size %d, max_redzone %d, poison_heap %d, "
              "malloc_context_size %d, alloc_dealloc_mismatch %d, "
              "allocator_may_return_null %d\n",
-      flags()->quarantine_size, flags()->max_redzone, flags()->poison_heap,
-      common_flags()->malloc_context_size, flags()->alloc_dealloc_mismatch,
-      common_flags()->allocator_may_return_null);
+          asan_deactivated_flags.quarantine_size, flags()->max_redzone,
+          flags()->poison_heap, common_flags()->malloc_context_size,
+          flags()->alloc_dealloc_mismatch,
+          asan_deactivated_flags.allocator_may_return_null);
 }
 
 }  // namespace __asan

lib/asan/asan_allocator.h

enum AllocType {
   FROM_MALLOC = 1,  // Memory block came from malloc, calloc, realloc, etc.
   FROM_NEW = 2,     // Memory block came from operator new.
   FROM_NEW_BR = 3   // Memory block came from operator new [ ]
 };
 
 static const uptr kNumberOfSizeClasses = 255;
 struct AsanChunk;
 
-void InitializeAllocator();
-void ReInitializeAllocator();
+void InitializeAllocator(bool may_return_null, uptr quarantine_size);
+void ReInitializeAllocator(bool may_return_null, uptr quarantine_size);
 
 class AsanChunkView {
  public:
   explicit AsanChunkView(AsanChunk *chunk) : chunk_(chunk) {}
   bool IsValid();   // Checks if AsanChunkView points to a valid allocated
                     // or quarantined chunk.
   uptr Beg();       // First byte of user memory.
   uptr End();       // Last byte of user memory.

lib/asan/asan_allocator2.cc

struct QuarantineCallback {
 
   void Deallocate(void *p) {
     allocator.Deallocate(cache_, p);
   }
 
   AllocatorCache *cache_;
 };
 
-void InitializeAllocator() {
-  allocator.Init();
-  quarantine.Init((uptr)flags()->quarantine_size, kMaxThreadLocalQuarantine);
+void InitializeAllocator(bool may_return_null, uptr quarantine_size) {
+  allocator.Init(may_return_null);
+  quarantine.Init(quarantine_size, kMaxThreadLocalQuarantine);
 }
 
-void ReInitializeAllocator() {
-  quarantine.Init((uptr)flags()->quarantine_size, kMaxThreadLocalQuarantine);
+void ReInitializeAllocator(bool may_return_null, uptr quarantine_size) {
+  allocator.SetMayReturnNull(may_return_null);
+  quarantine.Init(quarantine_size, kMaxThreadLocalQuarantine);
 }
 
 static void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
                       AllocType alloc_type, bool can_fill) {
   if (UNLIKELY(!asan_inited))
     AsanInitFromRtl();
   Flags &fl = *flags();
   CHECK(stack);
static void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
   if (!PrimaryAllocator::CanAllocate(needed_size, alignment)) {
     needed_size += rz_size;
     using_primary_allocator = false;
   }
   CHECK(IsAligned(needed_size, min_alignment));
   if (size > kMaxAllowedMallocSize || needed_size > kMaxAllowedMallocSize) {
     Report("WARNING: AddressSanitizer failed to allocate %p bytes\n",
            (void*)size);
-    return AllocatorReturnNull();
+    return allocator.ReturnNullOrDie();
   }
 
   AsanThread *t = GetCurrentThread();
   void *allocated;
   if (t) {
     AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
     allocated = allocator.Allocate(cache, needed_size, 8, false);
   } else {

 }
 
 void *asan_malloc(uptr size, BufferedStackTrace *stack) {
   return Allocate(size, 8, stack, FROM_MALLOC, true);
 }
 
 void *asan_calloc(uptr nmemb, uptr size, BufferedStackTrace *stack) {
   if (CallocShouldReturnNullDueToOverflow(size, nmemb))
-    return AllocatorReturnNull();
+    return allocator.ReturnNullOrDie();
   void *ptr = Allocate(nmemb * size, 8, stack, FROM_MALLOC, false);
   // If the memory comes from the secondary allocator no need to clear it
   // as it comes directly from mmap.
   if (ptr && allocator.FromPrimary(ptr))
     REAL(memset)(ptr, 0, nmemb * size);
   return ptr;
 }
 

lib/asan/asan_rtl.cc

Report("Shadow memory range interleaves with an existing memory mapping. "
            "ASan cannot proceed correctly. ABORTING.\n");
     DumpProcessMap();
     Die();
   }
 
   AsanTSDInit(PlatformTSDDtor);
   InstallDeadlySignalHandlers(AsanOnSIGSEGV);
 
-  InitializeAllocator();
+  InitializeAllocator(common_flags()->allocator_may_return_null,
+                      flags()->quarantine_size);
 
   // On Linux AsanThread::ThreadStart() calls malloc() that's why asan_inited
   // should be set to 1 prior to initializing the threads.
   asan_inited = 1;
   asan_init_is_running = false;
 
   if (flags()->atexit)
     Atexit(asan_atexit);

lib/lsan/lsan_allocator.cc

 typedef LargeMmapAllocator<> SecondaryAllocator;
 typedef CombinedAllocator<PrimaryAllocator, AllocatorCache,
           SecondaryAllocator> Allocator;
 
 static Allocator allocator;
 static THREADLOCAL AllocatorCache cache;
 
 void InitializeAllocator() {
-  allocator.Init();
+  allocator.Init(common_flags()->allocator_may_return_null);
 }
 
 void AllocatorThreadFinish() {
   allocator.SwallowCache(&cache);
 }
 
 static ChunkMetadata *Metadata(const void *p) {
   return reinterpret_cast<ChunkMetadata *>(allocator.GetMetaData(p));

lib/msan/msan.h

 extern int msan_report_count;
 
 bool ProtectRange(uptr beg, uptr end);
 bool InitShadow(bool map_shadow, bool init_origins);
 char *GetProcSelfMaps();
 void InitializeInterceptors();
 
 void MsanAllocatorThreadFinish();
+void *MsanCalloc(StackTrace *stack, uptr nmemb, uptr size);
 void *MsanReallocate(StackTrace *stack, void *oldp, uptr size,
                      uptr alignment, bool zeroise);
 void MsanDeallocate(StackTrace *stack, void *ptr);
 void InstallTrapHandler();
 void InstallAtExitHandler();
 void ReplaceOperatorsNewAndDelete();
 
 const char *GetStackOriginDescr(u32 id, uptr *pc);

lib/msan/msan_allocator.cc

 static SpinMutex fallback_mutex;
 
 static int inited = 0;
 
 static inline void Init() {
   if (inited) return;
   __msan_init();
   inited = true;  // this must happen before any threads are created.
-  allocator.Init();
+  allocator.Init(common_flags()->allocator_may_return_null);
 }
 
 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) {
   Init();
   if (size > kMaxAllowedMallocSize) {
     Report("WARNING: MemorySanitizer failed to allocate %p bytes\n",
            (void *)size);
-    return AllocatorReturnNull();
+    return allocator.ReturnNullOrDie();
   }
   MsanThread *t = GetCurrentThread();
   void *allocated;
   if (t) {
     AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
     allocated = allocator.Allocate(cache, size, alignment, false);
   } else {
     SpinMutexLock l(&fallback_mutex);
if (t) {
     allocator.Deallocate(cache, p);
   } else {
     SpinMutexLock l(&fallback_mutex);
     AllocatorCache *cache = &fallback_allocator_cache;
     allocator.Deallocate(cache, p);
   }
 }
 
+void *MsanCalloc(StackTrace *stack, uptr nmemb, uptr size) {
+  Init();
+  if (CallocShouldReturnNullDueToOverflow(size, nmemb))
+    return allocator.ReturnNullOrDie();
+  return MsanReallocate(stack, 0, nmemb * size, sizeof(u64), true);
+}
+
 void *MsanReallocate(StackTrace *stack, void *old_p, uptr new_size,
                      uptr alignment, bool zeroise) {
   if (!old_p)
     return MsanAllocate(stack, new_size, alignment, zeroise);
   if (!new_size) {
     MsanDeallocate(stack, old_p);
     return 0;
   }

lib/msan/msan_interceptors.cc

if (srcaddr) {
       SIZE_T sz = *addrlen;
       __msan_unpoison(srcaddr, (sz < srcaddr_sz) ? sz : srcaddr_sz);
     }
   }
   return res;
 }
 
 INTERCEPTOR(void *, calloc, SIZE_T nmemb, SIZE_T size) {
-  if (CallocShouldReturnNullDueToOverflow(size, nmemb))
-    return AllocatorReturnNull();
   GET_MALLOC_STACK_TRACE;
-  if (!msan_inited) {
+  if (UNLIKELY(!msan_inited)) {
     // Hack: dlsym calls calloc before REAL(calloc) is retrieved from dlsym.
     const SIZE_T kCallocPoolSize = 1024;
     static uptr calloc_memory_for_dlsym[kCallocPoolSize];
     static SIZE_T allocated;
     SIZE_T size_in_words = ((nmemb * size) + kWordSize - 1) / kWordSize;
     void *mem = (void*)&calloc_memory_for_dlsym[allocated];
     allocated += size_in_words;
     CHECK(allocated < kCallocPoolSize);
     return mem;
   }
-  return MsanReallocate(&stack, 0, nmemb * size, sizeof(u64), true);
+  return MsanCalloc(&stack, nmemb, size);
 }
 
 INTERCEPTOR(void *, realloc, void *ptr, SIZE_T size) {
   GET_MALLOC_STACK_TRACE;
   return MsanReallocate(&stack, ptr, size, sizeof(u64), false);
 }
 
 INTERCEPTOR(void *, malloc, SIZE_T size) {

lib/sanitizer_common/sanitizer_allocator.h

 #include "sanitizer_common.h"
 #include "sanitizer_libc.h"
 #include "sanitizer_list.h"
 #include "sanitizer_mutex.h"
 #include "sanitizer_lfstack.h"
 
 namespace __sanitizer {
 
-// Depending on allocator_may_return_null either return 0 or crash.
-void *AllocatorReturnNull();
+// Prints error message and kills the program.
+void NORETURN ReportAllocatorCannotReturnNull();
 
 // SizeClassMap maps allocation sizes into size classes and back.
 // Class 0 corresponds to size 0.
 // Classes 1 - 16 correspond to sizes 16 to 256 (size = class_id * 16).
 // Next 4 classes: 256 + i * 64  (i = 1 to 4).
 // Next 4 classes: 512 + i * 128 (i = 1 to 4).
 // ...
 // Next 4 classes: 2^k + i * 2^(k-2) (i = 1 to 4).

 };
 
 // This class can (de)allocate only large chunks of memory using mmap/unmap.
 // The main purpose of this allocator is to cover large and rare allocation
 // sizes not covered by more efficient allocators (e.g. SizeClassAllocator64).
 template <class MapUnmapCallback = NoOpMapUnmapCallback>
 class LargeMmapAllocator {
  public:
-  void Init() {
+  void Init(bool may_return_null) {
     internal_memset(this, 0, sizeof(*this));
     page_size_ = GetPageSizeCached();
+    atomic_store(&may_return_null_, may_return_null, memory_order_relaxed);
   }
 
   void *Allocate(AllocatorStats *stat, uptr size, uptr alignment) {
     CHECK(IsPowerOfTwo(alignment));
     uptr map_size = RoundUpMapSize(size);
     if (alignment > page_size_)
       map_size += alignment;
-    if (map_size < size) return AllocatorReturnNull();  // Overflow.
+    // Overflow.
+    if (map_size < size)
+      return ReturnNullOrDie();
     uptr map_beg = reinterpret_cast<uptr>(
         MmapOrDie(map_size, "LargeMmapAllocator"));
     CHECK(IsAligned(map_beg, page_size_));
     MapUnmapCallback().OnMap(map_beg, map_size);
     uptr map_end = map_beg + map_size;
     uptr res = map_beg + page_size_;
     if (res & (alignment - 1))  // Align.
       res += alignment - (res & (alignment - 1));
CHECK_LT(size_log, ARRAY_SIZE(stats.by_size_log));
       stats.max_allocated = Max(stats.max_allocated, stats.currently_allocated);
       stats.by_size_log[size_log]++;
       stat->Add(AllocatorStatAllocated, map_size);
       stat->Add(AllocatorStatMapped, map_size);
     }
     return reinterpret_cast<void*>(res);
   }
 
+  void *ReturnNullOrDie() {
+    if (atomic_load(&may_return_null_, memory_order_acquire))
+      return 0;
+    ReportAllocatorCannotReturnNull();
+  }
+
+  void SetMayReturnNull(bool may_return_null) {
+    atomic_store(&may_return_null_, may_return_null, memory_order_release);
+  }
+
   void Deallocate(AllocatorStats *stat, void *p) {
     Header *h = GetHeader(p);
     {
       SpinMutexLock l(&mutex_);
       uptr idx = h->chunk_idx;
       CHECK_EQ(chunks_[idx], h);
       CHECK_LT(idx, n_chunks_);
       chunks_[idx] = chunks_[n_chunks_ - 1];
private:
   uptr page_size_;
   Header *chunks_[kMaxNumChunks];
   uptr n_chunks_;
   uptr min_mmap_, max_mmap_;
   bool chunks_sorted_;
   struct Stats {
     uptr n_allocs, n_frees, currently_allocated, max_allocated, by_size_log[64];
   } stats;
+  atomic_uint8_t may_return_null_;
   SpinMutex mutex_;
 };
 
 // This class implements a complete memory allocator by using two
 // internal allocators:
 // PrimaryAllocator is efficient, but may not allocate some sizes (alignments).
 //  When allocating 2^x bytes it should return 2^x aligned chunk.
 // PrimaryAllocator is used via a local AllocatorCache.
 // SecondaryAllocator can allocate anything, but is not efficient.
 template <class PrimaryAllocator, class AllocatorCache,
           class SecondaryAllocator>  // NOLINT
 class CombinedAllocator {
  public:
-  void Init() {
+  void Init(bool may_return_null) {
     primary_.Init();
-    secondary_.Init();
+    secondary_.Init(may_return_null);
     stats_.Init();
+    atomic_store(&may_return_null_, may_return_null, memory_order_relaxed);
   }
 
   void *Allocate(AllocatorCache *cache, uptr size, uptr alignment,
                  bool cleared = false) {
     // Returning 0 on malloc(0) may break a lot of code.
     if (size == 0)
       size = 1;
     if (size + alignment < size)
-      return AllocatorReturnNull();
+      return ReturnNullOrDie();
     if (alignment > 8)
       size = RoundUpTo(size, alignment);
     void *res;
     bool from_primary = primary_.CanAllocate(size, alignment);
     if (from_primary)
       res = cache->Allocate(&primary_, primary_.ClassID(size));
     else
       res = secondary_.Allocate(&stats_, size, alignment);
     if (alignment > 8)
       CHECK_EQ(reinterpret_cast<uptr>(res) & (alignment - 1), 0);
     if (cleared && res && from_primary)
       internal_bzero_aligned16(res, RoundUpTo(size, 16));
     return res;
   }
 
+  void *ReturnNullOrDie() {
+    if (atomic_load(&may_return_null_, memory_order_acquire))
+      return 0;
+    ReportAllocatorCannotReturnNull();
+  }
+
+  void SetMayReturnNull(bool may_return_null) {
+    secondary_.SetMayReturnNull(may_return_null);
+    atomic_store(&may_return_null_, may_return_null, memory_order_release);
+  }
+
   void Deallocate(AllocatorCache *cache, void *p) {
     if (!p) return;
     if (primary_.PointerIsMine(p))
       cache->Deallocate(&primary_, primary_.GetSizeClass(p), p);
     else
       secondary_.Deallocate(&stats_, p);
   }
 
void ForEachChunk(ForEachChunkCallback callback, void *arg) {
     primary_.ForEachChunk(callback, arg);
     secondary_.ForEachChunk(callback, arg);
   }
 
  private:
   PrimaryAllocator primary_;
   SecondaryAllocator secondary_;
   AllocatorGlobalStats stats_;
+  atomic_uint8_t may_return_null_;
 };
 
 // Returns true if calloc(size, n) should return 0 due to overflow in size*n.
 bool CallocShouldReturnNullDueToOverflow(uptr size, uptr n);
 
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_ALLOCATOR_H
 

lib/sanitizer_common/sanitizer_allocator.cc

 //===-- sanitizer_allocator.cc --------------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is shared between AddressSanitizer and ThreadSanitizer
 // run-time libraries.
 // This allocator is used inside run-times.
 //===----------------------------------------------------------------------===//
 #include "sanitizer_allocator.h"
 #include "sanitizer_allocator_internal.h"
 #include "sanitizer_common.h"
-#include "sanitizer_flags.h"
 
 namespace __sanitizer {
 
 // ThreadSanitizer for Go uses libc malloc/free.
 #if defined(SANITIZER_GO) || defined(SANITIZER_USE_MALLOC)
 # if SANITIZER_LINUX && !SANITIZER_ANDROID
 extern "C" void *__libc_malloc(uptr size);
 extern "C" void __libc_free(void *ptr);

 
 InternalAllocator *internal_allocator() {
   InternalAllocator *internal_allocator_instance =
       reinterpret_cast<InternalAllocator *>(&internal_alloc_placeholder);
   if (atomic_load(&internal_allocator_initialized, memory_order_acquire) == 0) {
     SpinMutexLock l(&internal_alloc_init_mu);
     if (atomic_load(&internal_allocator_initialized, memory_order_relaxed) ==
         0) {
-      internal_allocator_instance->Init();
+      internal_allocator_instance->Init(/* may_return_null*/ false);
       atomic_store(&internal_allocator_initialized, 1, memory_order_release);
     }
   }
   return internal_allocator_instance;
 }
 
 static void *RawInternalAlloc(uptr size, InternalAllocatorCache *cache) {
   if (cache == 0) {

 }
 
 bool CallocShouldReturnNullDueToOverflow(uptr size, uptr n) {
   if (!size) return false;
   uptr max = (uptr)-1L;
   return (max / size) < n;
 }
 
-void *AllocatorReturnNull() {
-  if (common_flags()->allocator_may_return_null)
-    return 0;
+void NORETURN ReportAllocatorCannotReturnNull() {
   Report("%s's allocator is terminating the process instead of returning 0\n",
          SanitizerToolName);
   Report("If you don't like this behavior set allocator_may_return_null=1\n");
   CHECK(0);
-  return 0;
+  Die();
 }
 
 }  // namespace __sanitizer

lib/sanitizer_common/tests/sanitizer_allocator_test.cc

 //===-- sanitizer_allocator_test.cc ---------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of ThreadSanitizer/AddressSanitizer runtime.
 // Tests for sanitizer_allocator.h.
 //
 //===----------------------------------------------------------------------===//
 #include "sanitizer_common/sanitizer_allocator.h"
 #include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_common.h"
-#include "sanitizer_common/sanitizer_flags.h"
 
 #include "sanitizer_test_utils.h"
 #include "sanitizer_pthread_wrappers.h"
 
 #include "gtest/gtest.h"
 
 #include <stdlib.h>
 #include <algorithm>
TEST(SanitizerCommon, SizeClassAllocator32MapUnmapCallback) {
   //         TestMapUnmapCallback::map_count,
   //         TestMapUnmapCallback::unmap_count);
 }
 
 TEST(SanitizerCommon, LargeMmapAllocatorMapUnmapCallback) {
   TestMapUnmapCallback::map_count = 0;
   TestMapUnmapCallback::unmap_count = 0;
   LargeMmapAllocator<TestMapUnmapCallback> a;
-  a.Init();
+  a.Init(/* may_return_null */ false);
   AllocatorStats stats;
   stats.Init();
   void *x = a.Allocate(&stats, 1 << 20, 1);
   EXPECT_EQ(TestMapUnmapCallback::map_count, 1);
   a.Deallocate(&stats, x);
   EXPECT_EQ(TestMapUnmapCallback::unmap_count, 1);
 }
 

 TEST(SanitizerCommon, SizeClassAllocator64Overflow) {
   EXPECT_DEATH(FailInAssertionOnOOM<Allocator64>(), "Out of memory");
 }
 #endif
 
 #if !defined(_WIN32)  // FIXME: This currently fails on Windows.
 TEST(SanitizerCommon, LargeMmapAllocator) {
   LargeMmapAllocator<> a;
-  a.Init();
+  a.Init(/* may_return_null */ false);
   AllocatorStats stats;
   stats.Init();
 
   static const int kNumAllocs = 1000;
   char *allocated[kNumAllocs];
   static const uptr size = 4000;
   // Allocate some.
   for (int i = 0; i < kNumAllocs; i++) {

 
 template
 <class PrimaryAllocator, class SecondaryAllocator, class AllocatorCache>
 void TestCombinedAllocator() {
   typedef
       CombinedAllocator<PrimaryAllocator, AllocatorCache, SecondaryAllocator>
       Allocator;
   Allocator *a = new Allocator;
-  a->Init();
+  a->Init(/* may_return_null */ true);
 
   AllocatorCache cache;
   memset(&cache, 0, sizeof(cache));
   a->InitCache(&cache);
 
-  bool allocator_may_return_null = common_flags()->allocator_may_return_null;
-  common_flags()->allocator_may_return_null = true;
   EXPECT_EQ(a->Allocate(&cache, -1, 1), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, -1, 1024), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, (uptr)-1 - 1024, 1), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, (uptr)-1 - 1024, 1024), (void*)0);
   EXPECT_EQ(a->Allocate(&cache, (uptr)-1 - 1023, 1024), (void*)0);
 
-  common_flags()->allocator_may_return_null = false;
+  // Set to false
+  a->SetMayReturnNull(false);
   EXPECT_DEATH(a->Allocate(&cache, -1, 1),
                "allocator is terminating the process");
-  // Restore the original value.
-  common_flags()->allocator_may_return_null = allocator_may_return_null;
 
   const uptr kNumAllocs = 100000;
   const uptr kNumIter = 10;
   for (uptr iter = 0; iter < kNumIter; iter++) {
     std::vector<void*> allocated;
     for (uptr i = 0; i < kNumAllocs; i++) {
       uptr size = (i % (1 << 14)) + 1;
       if ((i % 1024) == 0)

 #endif
 
 TEST(SanitizerCommon, SizeClassAllocator32Iteration) {
   TestSizeClassAllocatorIteration<Allocator32Compact>();
 }
 
 TEST(SanitizerCommon, LargeMmapAllocatorIteration) {
   LargeMmapAllocator<> a;
-  a.Init();
+  a.Init(/* may_return_null */ false);
   AllocatorStats stats;
   stats.Init();
 
   static const uptr kNumAllocs = 1000;
   char *allocated[kNumAllocs];
   static const uptr size = 40;
   // Allocate some.
   for (uptr i = 0; i < kNumAllocs; i++)
ASSERT_NE(reported_chunks.find(reinterpret_cast<uptr>(allocated[i])),
               reported_chunks.end());
   }
   for (uptr i = 0; i < kNumAllocs; i++)
     a.Deallocate(&stats, allocated[i]);
 }
 
 TEST(SanitizerCommon, LargeMmapAllocatorBlockBegin) {
   LargeMmapAllocator<> a;
-  a.Init();
+  a.Init(/* may_return_null */ false);
   AllocatorStats stats;
   stats.Init();
 
   static const uptr kNumAllocs = 1024;
   static const uptr kNumExpectedFalseLookups = 10000000;
   char *allocated[kNumAllocs];
   static const uptr size = 4096;
   // Allocate some.

lib/tsan/rtl/tsan_interceptors.cc

 TSAN_INTERCEPTOR(void*, __libc_memalign, uptr align, uptr sz) {
   SCOPED_TSAN_INTERCEPTOR(__libc_memalign, align, sz);
   return user_alloc(thr, pc, sz, align);
 }
 
 TSAN_INTERCEPTOR(void*, calloc, uptr size, uptr n) {
   if (cur_thread()->in_symbolizer)
     return __libc_calloc(size, n);
-  if (__sanitizer::CallocShouldReturnNullDueToOverflow(size, n))
-    return AllocatorReturnNull();
   void *p = 0;
   {
     SCOPED_INTERCEPTOR_RAW(calloc, size, n);
-    p = user_alloc(thr, pc, n * size);
-    if (p)
-      internal_memset(p, 0, n * size);
+    p = user_calloc(thr, pc, size, n);
   }
   invoke_malloc_hook(p, n * size);
   return p;
 }
 
 TSAN_INTERCEPTOR(void*, realloc, void *p, uptr size) {
   if (cur_thread()->in_symbolizer)
     return __libc_realloc(p, size);

lib/tsan/rtl/tsan_mman.h

 void InitializeAllocator();
 void AllocatorThreadStart(ThreadState *thr);
 void AllocatorThreadFinish(ThreadState *thr);
 void AllocatorPrintStats();
 
 // For user allocations.
 void *user_alloc(ThreadState *thr, uptr pc, uptr sz,
                  uptr align = kDefaultAlignment, bool signal = true);
+void *user_calloc(ThreadState *thr, uptr pc, uptr sz, uptr n);
 // Does not accept NULL.
 void user_free(ThreadState *thr, uptr pc, void *p, bool signal = true);
 void *user_realloc(ThreadState *thr, uptr pc, void *p, uptr sz);
 void *user_alloc_aligned(ThreadState *thr, uptr pc, uptr sz, uptr align);
 uptr user_alloc_usable_size(const void *p);
 
 // Invoking malloc/free hooks that may be installed by the user.
 void invoke_malloc_hook(void *ptr, uptr size);

lib/tsan/rtl/tsan_mman.cc

 };
 
 static char allocator_placeholder[sizeof(Allocator)] ALIGNED(64);
 Allocator *allocator() {
   return reinterpret_cast<Allocator*>(&allocator_placeholder);
 }
 
 void InitializeAllocator() {
-  allocator()->Init();
+  allocator()->Init(common_flags()->allocator_may_return_null);
 }
 
 void AllocatorThreadStart(ThreadState *thr) {
   allocator()->InitCache(&thr->alloc_cache);
   internal_allocator()->InitCache(&thr->internal_alloc_cache);
 }
 
 void AllocatorThreadFinish(ThreadState *thr) {
static void SignalUnsafeCall(ThreadState *thr, uptr pc) {
   if (!IsFiredSuppression(ctx, rep, stack)) {
     rep.AddStack(stack, true);
     OutputReport(thr, rep);
   }
 }
 
 void *user_alloc(ThreadState *thr, uptr pc, uptr sz, uptr align, bool signal) {
   if ((sz >= (1ull << 40)) || (align >= (1ull << 40)))
-    return AllocatorReturnNull();
+    return allocator()->ReturnNullOrDie();
   void *p = allocator()->Allocate(&thr->alloc_cache, sz, align);
   if (p == 0)
     return 0;
   if (ctx && ctx->initialized)
     OnUserAlloc(thr, pc, (uptr)p, sz, true);
   if (signal)
     SignalUnsafeCall(thr, pc);
   return p;
 }
 
+void *user_calloc(ThreadState *thr, uptr pc, uptr size, uptr n) {
+  if (CallocShouldReturnNullDueToOverflow(size, n))
+    return allocator()->ReturnNullOrDie();
+  void *p = user_alloc(thr, pc, n * size);
+  if (p)
+    internal_memset(p, 0, n * size);
+  return p;
+}
+
 void user_free(ThreadState *thr, uptr pc, void *p, bool signal) {
   if (ctx && ctx->initialized)
     OnUserFree(thr, pc, (uptr)p, true);
   allocator()->Deallocate(&thr->alloc_cache, p);
   if (signal)
     SignalUnsafeCall(thr, pc);
 }