[Sanitizers] 32 bit allocator respects allocator_may_return_null flag

[alekseyshl]

Summary

Make SizeClassAllocator32 return nullptr when it encounters OOM, which
allows the entire sanitizer's allocator to follow allocator_may_return_null=1
policy, even for small allocations (LargeMmapAllocator is already fixed
by D34243).

Will add a test for OOM in primary allocator later, when
SizeClassAllocator64 can gracefully handle OOM too.

[eugenis]

lib/sanitizer_common/sanitizer_allocator_primary32.h

309–311

is it possible for CreateBatch to return nullptr?
It may Allocate().

[alekseyshl]

• Accomodate for CreateBatch returning nullptr.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_primary32.h

309–311

Indeed!

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

147–150

Doing it this way, we are losing the 'unlikeliness' of c->count == 0.
I am not sure this will matter much though.

[alekseyshl]

• Restore UNLIKELY(c->count == 0) condition semantics

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

147–150

Let's bring it back then.

[eugenis]

LGTM

lib/sanitizer_common/sanitizer_allocator_local_cache.h

147–150

unnecessary {}

lib/sanitizer_common/sanitizer_allocator_primary32.h

291

early exit
do we need UNLIKELY here? probably not.

[alekseyshl]

• Address comments

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

147–150

I'd rather keep those, I prefer to add {} when if body is more than one line. This one is.

lib/sanitizer_common/sanitizer_allocator_local_cache.h

void Destroy(SizeClassAllocator *allocator, AllocatorGlobalStats *s) {
     if (s)
       s->Unregister(&stats_);
   }
 
   void *Allocate(SizeClassAllocator *allocator, uptr class_id) {
     CHECK_NE(class_id, 0UL);
     CHECK_LT(class_id, kNumClasses);
     PerClass *c = &per_class_[class_id];
-    if (UNLIKELY(c->count == 0))
-      Refill(allocator, class_id);
+    if (UNLIKELY(c->count == 0)) {
+      if (UNLIKELY(!Refill(allocator, class_id)))
+        return nullptr;
+    }
     stats_.Add(AllocatorStatAllocated, c->class_size);
     void *res = c->batch[--c->count];
     PREFETCH(c->batch[c->count - 1]);
     return res;
   }
 
   void Deallocate(SizeClassAllocator *allocator, uptr class_id, void *p) {
     CHECK_NE(class_id, 0UL);
struct SizeClassAllocator32LocalCache {
   // For small size classes deallocates b to the allocator.
   // Does notthing for large size classes.
   void DestroyBatch(uptr class_id, SizeClassAllocator *allocator,
                     TransferBatch *b) {
     if (uptr batch_class_id = per_class_[class_id].class_id_for_transfer_batch)
       Deallocate(allocator, batch_class_id, b);
   }
 
-  NOINLINE void Refill(SizeClassAllocator *allocator, uptr class_id) {
+  NOINLINE bool Refill(SizeClassAllocator *allocator, uptr class_id) {
     InitCache();
     PerClass *c = &per_class_[class_id];
     TransferBatch *b = allocator->AllocateBatch(&stats_, this, class_id);
+    if (UNLIKELY(!b))
+      return false;
     CHECK_GT(b->Count(), 0);
     b->CopyToArray(c->batch);
     c->count = b->Count();
     DestroyBatch(class_id, allocator, b);
+    return true;
   }
 
   NOINLINE void Drain(SizeClassAllocator *allocator, uptr class_id) {
     InitCache();
     PerClass *c = &per_class_[class_id];
     uptr cnt = Min(c->max_count / 2, c->count);
     uptr first_idx_to_drain = c->count - cnt;
     TransferBatch *b = CreateBatch(
         class_id, allocator, (TransferBatch *)c->batch[first_idx_to_drain]);
+    // Failure to allocate a batch while releasing memory is non recoverable.
+    // TODO(alekseys): Figure out how to do it without allocating a new batch.
+    if (UNLIKELY(!b))
+      DieOnFailure::OnOOM();
     b->SetFromArray(allocator->GetRegionBeginBySizeClass(class_id),
                     &c->batch[first_idx_to_drain], cnt);
     c->count -= cnt;
     allocator->DeallocateBatch(&stats_, class_id, b);
   }
 };
 

lib/sanitizer_common/sanitizer_allocator_primary32.h

 // SizeClassAllocator32 -- allocator for 32-bit address space.
 // This allocator can theoretically be used on 64-bit arch, but there it is less
 // efficient than SizeClassAllocator64.
 //
 // [kSpaceBeg, kSpaceBeg + kSpaceSize) is the range of addresses which can
 // be returned by MmapOrDie().
 //
 // Region:
-//   a result of a single call to MmapAlignedOrDie(kRegionSize, kRegionSize).
+//   a result of a single call to MmapAlignedOrDieOnFatalError(kRegionSize,
+//                                                             kRegionSize).
 // Since the regions are aligned by kRegionSize, there are exactly
 // kNumPossibleRegions possible regions in the address space and so we keep
 // a ByteMap possible_regions to store the size classes of each Region.
 // 0 size class means the region is not used by the allocator.
 //
 // One Region is used to allocate chunks of a single size class.
 // A Region looks like this:
 // UserChunk1 .. UserChunkN <gap> MetaChunkN .. MetaChunk1
void *GetMetaData(const void *p) {
     return reinterpret_cast<void*>(meta);
   }
 
   NOINLINE TransferBatch *AllocateBatch(AllocatorStats *stat, AllocatorCache *c,
                                         uptr class_id) {
     CHECK_LT(class_id, kNumClasses);
     SizeClassInfo *sci = GetSizeClassInfo(class_id);
     SpinMutexLock l(&sci->mutex);
-    if (sci->free_list.empty())
-      PopulateFreeList(stat, c, sci, class_id);
+    if (sci->free_list.empty() &&
+        UNLIKELY(!PopulateFreeList(stat, c, sci, class_id)))
+      return nullptr;
     CHECK(!sci->free_list.empty());
     TransferBatch *b = sci->free_list.front();
     sci->free_list.pop_front();
     return b;
   }
 
   NOINLINE void DeallocateBatch(AllocatorStats *stat, uptr class_id,
                                 TransferBatch *b) {
private:
   }
 
   uptr ComputeRegionBeg(uptr mem) {
     return mem & ~(kRegionSize - 1);
   }
 
   uptr AllocateRegion(AllocatorStats *stat, uptr class_id) {
     CHECK_LT(class_id, kNumClasses);
-    uptr res = reinterpret_cast<uptr>(MmapAlignedOrDie(kRegionSize, kRegionSize,
-                                      "SizeClassAllocator32"));
+    uptr res = reinterpret_cast<uptr>(MmapAlignedOrDieOnFatalError(
+        kRegionSize, kRegionSize, "SizeClassAllocator32"));
+    if (UNLIKELY(!res))
+      return 0;
     MapUnmapCallback().OnMap(res, kRegionSize);
     stat->Add(AllocatorStatMapped, kRegionSize);
     CHECK_EQ(0U, (res & (kRegionSize - 1)));
     possible_regions.set(ComputeRegionId(res), static_cast<u8>(class_id));
     return res;
   }
 
   SizeClassInfo *GetSizeClassInfo(uptr class_id) {
     CHECK_LT(class_id, kNumClasses);
     return &size_class_info_array[class_id];
   }
 
-  void PopulateFreeList(AllocatorStats *stat, AllocatorCache *c,
+  bool PopulateFreeList(AllocatorStats *stat, AllocatorCache *c,
                         SizeClassInfo *sci, uptr class_id) {
     uptr size = ClassIdToSize(class_id);
     uptr reg = AllocateRegion(stat, class_id);
+    if (UNLIKELY(!reg))
+      return false;
     uptr n_chunks = kRegionSize / (size + kMetadataSize);
     uptr max_count = TransferBatch::MaxCached(class_id);
     TransferBatch *b = nullptr;
     for (uptr i = reg; i < reg + n_chunks * size; i += size) {
       if (!b) {
         b = c->CreateBatch(class_id, this, (TransferBatch*)i);
+        if (!b)
+          return false;
         b->Clear();
       }
       b->Add((void*)i);
       if (b->Count() == max_count) {
         CHECK_GT(b->Count(), 0);
         sci->free_list.push_back(b);
         b = nullptr;
       }
     }
     if (b) {
       CHECK_GT(b->Count(), 0);
       sci->free_list.push_back(b);
     }
+    return true;
   }
 
   ByteMap possible_regions;
   SizeClassInfo size_class_info_array[kNumClasses];
 };

lib/sanitizer_common/sanitizer_common.h

 void *MmapOrDieOnFatalError(uptr size, const char *mem_type);
 void *MmapFixedNoReserve(uptr fixed_addr, uptr size,
                          const char *name = nullptr);
 void *MmapNoReserveOrDie(uptr size, const char *mem_type);
 void *MmapFixedOrDie(uptr fixed_addr, uptr size);
 void *MmapFixedNoAccess(uptr fixed_addr, uptr size, const char *name = nullptr);
 void *MmapNoAccess(uptr size);
 // Map aligned chunk of address space; size and alignment are powers of two.
-void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type);
+// Dies on all but out of memory errors, in the latter case returns nullptr.
+void *MmapAlignedOrDieOnFatalError(uptr size, uptr alignment,
+                                   const char *mem_type);
 // Disallow access to a memory range.  Use MmapFixedNoAccess to allocate an
 // unaccessible memory.
 bool MprotectNoAccess(uptr addr, uptr size);
 bool MprotectReadOnly(uptr addr, uptr size);
 
 // Find an available address space.
 uptr FindAvailableMemoryRange(uptr size, uptr alignment, uptr left_padding);
 

lib/sanitizer_common/sanitizer_posix.cc

void *MmapOrDieOnFatalError(uptr size, const char *mem_type) {
   }
   IncreaseTotalMmap(size);
   return (void *)res;
 }
 
 // We want to map a chunk of address space aligned to 'alignment'.
 // We do it by maping a bit more and then unmaping redundant pieces.
 // We probably can do it with fewer syscalls in some OS-dependent way.
-void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type) {
+void *MmapAlignedOrDieOnFatalError(uptr size, uptr alignment,
+                                   const char *mem_type) {
   CHECK(IsPowerOfTwo(size));
   CHECK(IsPowerOfTwo(alignment));
   uptr map_size = size + alignment;
-  uptr map_res = (uptr)MmapOrDie(map_size, mem_type);
+  uptr map_res = (uptr)MmapOrDieOnFatalError(map_size, mem_type);
+  if (!map_res)
+    return nullptr;
   uptr map_end = map_res + map_size;
   uptr res = map_res;
   if (res & (alignment - 1))  // Not aligned.
     res = (map_res + alignment) & ~(alignment - 1);
   uptr end = res + size;
   if (res != map_res)
     UnmapOrDie((void*)map_res, res - map_res);
   if (end != map_end)

lib/sanitizer_common/sanitizer_win.cc

if (VirtualFree(addr, size, MEM_DECOMMIT) == 0) {
       Report("ERROR: %s failed to "
              "deallocate 0x%zx (%zd) bytes at address %p (error code: %d)\n",
              SanitizerToolName, size, size, addr, GetLastError());
       CHECK("unable to unmap" && 0);
     }
   }
 }
 
-void *MmapOrDieOnFatalError(uptr size, const char *mem_type) {
-  void *rv = VirtualAlloc(0, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
-  if (rv == 0) {
-    error_t last_error = GetLastError();
-    if (last_error != ERROR_NOT_ENOUGH_MEMORY)
-      ReportMmapFailureAndDie(size, mem_type, "allocate", last_error);
-  }
+static void *ReturnNullptrOnOOMOrDie(uptr size, const char *mem_type,
+                                     const char *mmap_type) {
+  if (last_error == ERROR_NOT_ENOUGH_MEMORY)
+    return nullptr;
+  ReportMmapFailureAndDie(size, mem_type, mmap_type, last_error);
+
+void *MmapOrDieOnFatalError(uptr size, const char *mem_type) {
+  void *rv = VirtualAlloc(0, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+  if (rv == 0)
+    return ReturnNullptrOnOOMOrDie(size, mem_type, "allocate");
   return rv;
 }
 
 // We want to map a chunk of address space aligned to 'alignment'.
-void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type) {
+void *MmapAlignedOrDieOnFatalError(uptr size, uptr alignment,
+                                   const char *mem_type) {
   CHECK(IsPowerOfTwo(size));
   CHECK(IsPowerOfTwo(alignment));
 
   // Windows will align our allocations to at least 64K.
   alignment = Max(alignment, GetMmapGranularity());
 
   uptr mapped_addr =
       (uptr)VirtualAlloc(0, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
   if (!mapped_addr)
-    ReportMmapFailureAndDie(size, mem_type, "allocate aligned", GetLastError());
+    return ReturnNullptrOnOOMOrDie(size, mem_type, "allocate aligned");
 
   // If we got it right on the first try, return. Otherwise, unmap it and go to
   // the slow path.
   if (IsAligned(mapped_addr, alignment))
     return (void*)mapped_addr;
   if (VirtualFree((void *)mapped_addr, 0, MEM_RELEASE) == 0)
     ReportMmapFailureAndDie(size, mem_type, "deallocate", GetLastError());
 
   // If we didn't get an aligned address, overallocate, find an aligned address,
   // unmap, and try to allocate at that aligned address.
   int retries = 0;
   const int kMaxRetries = 10;
   for (; retries < kMaxRetries &&
          (mapped_addr == 0 || !IsAligned(mapped_addr, alignment));
        retries++) {
     // Overallocate size + alignment bytes.
     mapped_addr =
         (uptr)VirtualAlloc(0, size + alignment, MEM_RESERVE, PAGE_NOACCESS);
     if (!mapped_addr)
-      ReportMmapFailureAndDie(size, mem_type, "allocate aligned",
-                              GetLastError());
+      return ReturnNullptrOnOOMOrDie(size, mem_type, "allocate aligned");
 
     // Find the aligned address.
     uptr aligned_addr = RoundUpTo(mapped_addr, alignment);
 
     // Free the overallocation.
     if (VirtualFree((void *)mapped_addr, 0, MEM_RELEASE) == 0)
       ReportMmapFailureAndDie(size, mem_type, "deallocate", GetLastError());
 
     // Attempt to allocate exactly the number of bytes we need at the aligned
     // address. This may fail for a number of reasons, in which case we continue
     // the loop.
     mapped_addr = (uptr)VirtualAlloc((void *)aligned_addr, size,
                                      MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
   }
 
   // Fail if we can't make this work quickly.
   if (retries == kMaxRetries && mapped_addr == 0)
-    ReportMmapFailureAndDie(size, mem_type, "allocate aligned", GetLastError());
+    return ReturnNullptrOnOOMOrDie(size, mem_type, "allocate aligned");
 
   return (void *)mapped_addr;
 }
 
 void *MmapFixedNoReserve(uptr fixed_addr, uptr size, const char *name) {
   // FIXME: is this really "NoReserve"? On Win32 this does not matter much,
   // but on Win64 it does.
   (void)name;  // unsupported

lib/sanitizer_common/tests/sanitizer_common_test.cc

TEST(SanitizerCommon, SortTest) {
   EXPECT_TRUE(IsSorted(array, n));
   // Minimal case - sort three elements.
   array[0] = 1;
   array[1] = 0;
   SortArray(array, 2);
   EXPECT_TRUE(IsSorted(array, 2));
 }
 
-TEST(SanitizerCommon, MmapAlignedOrDie) {
+TEST(SanitizerCommon, MmapAlignedOrDieOnFatalError) {
   uptr PageSize = GetPageSizeCached();
   for (uptr size = 1; size <= 32; size *= 2) {
     for (uptr alignment = 1; alignment <= 32; alignment *= 2) {
       for (int iter = 0; iter < 100; iter++) {
-        uptr res = (uptr)MmapAlignedOrDie(
+        uptr res = (uptr)MmapAlignedOrDieOnFatalError(
             size * PageSize, alignment * PageSize, "MmapAlignedOrDieTest");
         EXPECT_EQ(0U, res % (alignment * PageSize));
         internal_memset((void*)res, 1, size * PageSize);
         UnmapOrDie((void*)res, size * PageSize);
       }
     }
   }
 }