[Sanitizers] 64 bit allocator respects allocator_may_return_null flag

[alekseyshl]

Summary

Make SizeClassAllocator64 return nullptr when it encounters OOM, which
allows the entire sanitizer's allocator to follow
allocator_may_return_null=1 policy
(LargeMmapAllocator: D34243, SizeClassAllocator64: D34433).

[eugenis]

lib/sanitizer_common/sanitizer_allocator_primary64.h

454

Do you want to keep this error message?
I think it was helpful on a couple of occasions.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_primary64.h

454

Well, ok, I'll bring it back. It'll require a flag to show it only once though, otherwise the client which can survive null returned by malloc might be spammed by those messages.

[alekseyshl]

• Bring the oom message back.

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(c, allocator, class_id);
+    if (UNLIKELY(c->count == 0)) {
+      if (UNLIKELY(!Refill(c, allocator, class_id)))
+        return nullptr;
+    }
     stats_.Add(AllocatorStatAllocated, c->class_size);
     CHECK_GT(c->count, 0);
     CompactPtrT chunk = c->chunks[--c->count];
     void *res = reinterpret_cast<void *>(allocator->CompactPtrToPointer(
         allocator->GetRegionBeginBySizeClass(class_id), chunk));
     return res;
   }
 
if (per_class_[1].max_count)
       return;
     for (uptr i = 0; i < kNumClasses; i++) {
       PerClass *c = &per_class_[i];
       c->max_count = 2 * SizeClassMap::MaxCachedHint(i);
       c->class_size = Allocator::ClassIdToSize(i);
     }
   }
 
-  NOINLINE void Refill(PerClass *c, SizeClassAllocator *allocator,
+  NOINLINE bool Refill(PerClass *c, SizeClassAllocator *allocator,
                        uptr class_id) {
     InitCache();
     uptr num_requested_chunks = c->max_count / 2;
-    allocator->GetFromAllocator(&stats_, class_id, c->chunks,
-                                num_requested_chunks);
+    if (UNLIKELY(!allocator->GetFromAllocator(&stats_, class_id, c->chunks,
+                                              num_requested_chunks)))
+      return false;
     c->count = num_requested_chunks;
+    return true;
   }
 
   NOINLINE void Drain(PerClass *c, SizeClassAllocator *allocator, uptr class_id,
                       uptr count) {
     InitCache();
     CHECK_GE(c->count, count);
     uptr first_idx_to_drain = c->count - count;
     c->count -= count;

lib/sanitizer_common/sanitizer_allocator_primary64.h

if (kUsingConstantSpaceBeg) {
       CHECK_EQ(kSpaceBeg, reinterpret_cast<uptr>(
                               MmapFixedNoAccess(kSpaceBeg, TotalSpaceSize)));
     } else {
       NonConstSpaceBeg =
           reinterpret_cast<uptr>(MmapNoAccess(TotalSpaceSize));
       CHECK_NE(NonConstSpaceBeg, ~(uptr)0);
     }
     SetReleaseToOSIntervalMs(release_to_os_interval_ms);
-    MapWithCallback(SpaceEnd(), AdditionalSize());
+    MapWithCallbackOrDie(SpaceEnd(), AdditionalSize());
   }
 
   s32 ReleaseToOSIntervalMs() const {
     return atomic_load(&release_to_os_interval_ms_, memory_order_relaxed);
   }
 
   void SetReleaseToOSIntervalMs(s32 release_to_os_interval_ms) {
     atomic_store(&release_to_os_interval_ms_, release_to_os_interval_ms,
                  memory_order_relaxed);
   }
 
-  void MapWithCallback(uptr beg, uptr size) {
-    CHECK_EQ(beg, reinterpret_cast<uptr>(MmapFixedOrDie(beg, size)));
-    MapUnmapCallback().OnMap(beg, size);
-  }
-
-  void UnmapWithCallback(uptr beg, uptr size) {
-    MapUnmapCallback().OnUnmap(beg, size);
-    UnmapOrDie(reinterpret_cast<void *>(beg), size);
-  }
-
   static bool CanAllocate(uptr size, uptr alignment) {
     return size <= SizeClassMap::kMaxSize &&
       alignment <= SizeClassMap::kMaxSize;
   }
 
   NOINLINE void ReturnToAllocator(AllocatorStats *stat, uptr class_id,
                                   const CompactPtrT *chunks, uptr n_chunks) {
     RegionInfo *region = GetRegionInfo(class_id);
     uptr region_beg = GetRegionBeginBySizeClass(class_id);
     CompactPtrT *free_array = GetFreeArray(region_beg);
 
     BlockingMutexLock l(&region->mutex);
     uptr old_num_chunks = region->num_freed_chunks;
     uptr new_num_freed_chunks = old_num_chunks + n_chunks;
-    EnsureFreeArraySpace(region, region_beg, new_num_freed_chunks);
+    // Failure to allocate free array space while releasing memory is non
+    // recoverable.
+    if (UNLIKELY(!EnsureFreeArraySpace(region, region_beg,
+                                       new_num_freed_chunks)))
+      DieOnFailure::OnOOM();
     for (uptr i = 0; i < n_chunks; i++)
       free_array[old_num_chunks + i] = chunks[i];
     region->num_freed_chunks = new_num_freed_chunks;
     region->n_freed += n_chunks;
 
     MaybeReleaseToOS(class_id);
   }
 
-  NOINLINE void GetFromAllocator(AllocatorStats *stat, uptr class_id,
+  NOINLINE bool GetFromAllocator(AllocatorStats *stat, uptr class_id,
                                  CompactPtrT *chunks, uptr n_chunks) {
     RegionInfo *region = GetRegionInfo(class_id);
     uptr region_beg = GetRegionBeginBySizeClass(class_id);
     CompactPtrT *free_array = GetFreeArray(region_beg);
 
     BlockingMutexLock l(&region->mutex);
     if (UNLIKELY(region->num_freed_chunks < n_chunks)) {
-      PopulateFreeArray(stat, class_id, region,
-                        n_chunks - region->num_freed_chunks);
+      if (UNLIKELY(!PopulateFreeArray(stat, class_id, region,
+                                      n_chunks - region->num_freed_chunks)))
+        return false;
       CHECK_GE(region->num_freed_chunks, n_chunks);
     }
     region->num_freed_chunks -= n_chunks;
     uptr base_idx = region->num_freed_chunks;
     for (uptr i = 0; i < n_chunks; i++)
       chunks[i] = free_array[base_idx + i];
     region->n_allocated += n_chunks;
+    return true;
   }
 
-
   bool PointerIsMine(const void *p) {
     uptr P = reinterpret_cast<uptr>(p);
     if (kUsingConstantSpaceBeg && (kSpaceBeg % kSpaceSize) == 0)
       return P / kSpaceSize == kSpaceBeg / kSpaceSize;
     return P >= SpaceBeg() && P < SpaceEnd();
   }
 
   uptr GetRegionBegin(const void *p) {
uptr TotalMemoryUsed() {
     uptr res = 0;
     for (uptr i = 0; i < kNumClasses; i++)
       res += GetRegionInfo(i)->allocated_user;
     return res;
   }
 
   // Test-only.
   void TestOnlyUnmap() {
-    UnmapWithCallback(SpaceBeg(), kSpaceSize + AdditionalSize());
+    UnmapWithCallbackOrDie(SpaceBeg(), kSpaceSize + AdditionalSize());
   }
 
   static void FillMemoryProfile(uptr start, uptr rss, bool file, uptr *stats,
                            uptr stats_size) {
     for (uptr class_id = 0; class_id < stats_size; class_id++)
       if (stats[class_id] == start)
         stats[class_id] = rss;
   }
uptr GetChunkIdx(uptr chunk, uptr size) {
     return (u32)offset / (u32)size;
   }
 
   CompactPtrT *GetFreeArray(uptr region_beg) {
     return reinterpret_cast<CompactPtrT *>(region_beg + kRegionSize -
                                            kFreeArraySize);
   }
 
-  void EnsureFreeArraySpace(RegionInfo *region, uptr region_beg,
+  bool MapWithCallback(uptr beg, uptr size) {
+    uptr mapped = reinterpret_cast<uptr>(MmapFixedOrDieOnFatalError(beg, size));
+    if (!mapped)
+      return false;
+    CHECK_EQ(beg, mapped);
+    MapUnmapCallback().OnMap(beg, size);
+    return true;
+  }
+
+  void MapWithCallbackOrDie(uptr beg, uptr size) {
+    CHECK_EQ(beg, reinterpret_cast<uptr>(MmapFixedOrDie(beg, size)));
+    MapUnmapCallback().OnMap(beg, size);
+  }
+
+  void UnmapWithCallbackOrDie(uptr beg, uptr size) {
+    MapUnmapCallback().OnUnmap(beg, size);
+    UnmapOrDie(reinterpret_cast<void *>(beg), size);
+  }
+
+  bool EnsureFreeArraySpace(RegionInfo *region, uptr region_beg,
                             uptr num_freed_chunks) {
     uptr needed_space = num_freed_chunks * sizeof(CompactPtrT);
     if (region->mapped_free_array < needed_space) {
       CHECK_LE(needed_space, kFreeArraySize);
       uptr new_mapped_free_array = RoundUpTo(needed_space, kFreeArrayMapSize);
       uptr current_map_end = reinterpret_cast<uptr>(GetFreeArray(region_beg)) +
                              region->mapped_free_array;
       uptr new_map_size = new_mapped_free_array - region->mapped_free_array;
-      MapWithCallback(current_map_end, new_map_size);
+      if (UNLIKELY(!MapWithCallback(current_map_end, new_map_size)))
+        return false;
       region->mapped_free_array = new_mapped_free_array;
     }
+    return true;
   }
 
-
-  NOINLINE void PopulateFreeArray(AllocatorStats *stat, uptr class_id,
+  NOINLINE bool PopulateFreeArray(AllocatorStats *stat, uptr class_id,
                                   RegionInfo *region, uptr requested_count) {
     // region->mutex is held.
-    uptr size = ClassIdToSize(class_id);
-    uptr beg_idx = region->allocated_user;
-    uptr end_idx = beg_idx + requested_count * size;
-    uptr region_beg = GetRegionBeginBySizeClass(class_id);
-    if (end_idx > region->mapped_user) {
+    const uptr size = ClassIdToSize(class_id);
+    const uptr new_space_beg = region->allocated_user;
+    const uptr new_space_end = new_space_beg + requested_count * size;
+    const uptr region_beg = GetRegionBeginBySizeClass(class_id);
+
+    // Map more space for chunks, if necessary.
+    if (new_space_end > region->mapped_user) {
       if (!kUsingConstantSpaceBeg && region->mapped_user == 0)
         region->rand_state = static_cast<u32>(region_beg >> 12);  // From ASLR.
       // Do the mmap for the user memory.
       uptr map_size = kUserMapSize;
-      while (end_idx > region->mapped_user + map_size)
+      while (new_space_end > region->mapped_user + map_size)
         map_size += kUserMapSize;
-      CHECK_GE(region->mapped_user + map_size, end_idx);
-      MapWithCallback(region_beg + region->mapped_user, map_size);
+      CHECK_GE(region->mapped_user + map_size, new_space_end);
+      if (UNLIKELY(!MapWithCallback(region_beg + region->mapped_user,
+                                    map_size)))
+        return false;
       stat->Add(AllocatorStatMapped, map_size);
       region->mapped_user += map_size;
     }
+    const uptr new_chunks_count = (region->mapped_user - new_space_beg) / size;
+
+    // Calculate the required space for metadata.
+    const uptr requested_allocated_meta =
+        region->allocated_meta + new_chunks_count * kMetadataSize;
+    uptr requested_mapped_meta = region->mapped_meta;
+    while (requested_allocated_meta > requested_mapped_meta)
+      requested_mapped_meta += kMetaMapSize;
+    // Check whether this size class is exhausted.
+    if (region->mapped_user + requested_mapped_meta >
+        kRegionSize - kFreeArraySize)
+      return false;
+    // Map more space for metadata, if necessary.
+    if (requested_mapped_meta > region->mapped_meta) {
+      if (UNLIKELY(!MapWithCallback(
+              GetMetadataEnd(region_beg) - requested_mapped_meta,
+              requested_mapped_meta - region->mapped_meta)))
+        return false;
+      region->mapped_meta = requested_mapped_meta;
+    }
+
+    // If necessary, allocate more space for the free array and populate it with
+    // newly allocated chunks.
+    const uptr total_freed_chunks = region->num_freed_chunks + new_chunks_count;
+    if (UNLIKELY(!EnsureFreeArraySpace(region, region_beg, total_freed_chunks)))
+      return false;
     CompactPtrT *free_array = GetFreeArray(region_beg);
-    uptr total_count = (region->mapped_user - beg_idx) / size;
-    uptr num_freed_chunks = region->num_freed_chunks;
-    EnsureFreeArraySpace(region, region_beg, num_freed_chunks + total_count);
-    for (uptr i = 0; i < total_count; i++) {
-      uptr chunk = beg_idx + i * size;
-      free_array[num_freed_chunks + total_count - 1 - i] =
-          PointerToCompactPtr(0, chunk);
-    }
+    for (uptr i = 0, chunk = new_space_beg; i < new_chunks_count;
+         i++, chunk += size)
+      free_array[total_freed_chunks - 1 - i] = PointerToCompactPtr(0, chunk);
     if (kRandomShuffleChunks)
-      RandomShuffle(&free_array[num_freed_chunks], total_count,
+      RandomShuffle(&free_array[region->num_freed_chunks], new_chunks_count,
                     &region->rand_state);
-    region->num_freed_chunks += total_count;
-    region->allocated_user += total_count * size;
-    CHECK_LE(region->allocated_user, region->mapped_user);
 
-    region->allocated_meta += total_count * kMetadataSize;
-    if (region->allocated_meta > region->mapped_meta) {
-      uptr map_size = kMetaMapSize;
-      while (region->allocated_meta > region->mapped_meta + map_size)
-        map_size += kMetaMapSize;
-      // Do the mmap for the metadata.
-      CHECK_GE(region->mapped_meta + map_size, region->allocated_meta);
-      MapWithCallback(GetMetadataEnd(region_beg) -
-                      region->mapped_meta - map_size, map_size);
-      region->mapped_meta += map_size;
-    }
+    // All necessary memory is mapped and now it is safe to advance all
+    // 'allocated_*' counters.
+    region->num_freed_chunks += new_chunks_count;
+    region->allocated_user += new_chunks_count * size;
+    CHECK_LE(region->allocated_user, region->mapped_user);
+    region->allocated_meta = requested_allocated_meta;
     CHECK_LE(region->allocated_meta, region->mapped_meta);
-    if (region->mapped_user + region->mapped_meta >
-        kRegionSize - kFreeArraySize) {
-      Printf("%s: Out of memory. Dying. ", SanitizerToolName);
-      Printf("The process has exhausted %zuMB for size class %zu.\n",
-          kRegionSize / 1024 / 1024, size);
-      Die();
-    }
+
+    return true;
   }
 
   void MaybeReleaseChunkRange(uptr region_beg, uptr chunk_size,
                               CompactPtrT first, CompactPtrT last) {
     uptr beg_ptr = CompactPtrToPointer(region_beg, first);
     uptr end_ptr = CompactPtrToPointer(region_beg, last) + chunk_size;
     ReleaseMemoryPagesToOS(beg_ptr, end_ptr);
   }

lib/sanitizer_common/sanitizer_common.h

 void UnmapOrDie(void *addr, uptr size);
 // Behaves just like MmapOrDie, but tolerates out of memory condition, in that
 // case returns nullptr.
 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);
+// Behaves just like MmapFixedOrDie, but tolerates out of memory condition, in
+// that case returns nullptr.
+void *MmapFixedOrDieOnFatalError(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.
 // 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.

lib/sanitizer_common/sanitizer_posix.cc

uptr p = internal_mmap(nullptr,
                          -1, 0);
   int reserrno;
   if (internal_iserror(p, &reserrno))
     ReportMmapFailureAndDie(size, mem_type, "allocate noreserve", reserrno);
   IncreaseTotalMmap(size);
   return (void *)p;
 }
 
-void *MmapFixedOrDie(uptr fixed_addr, uptr size) {
+void *MmapFixedImpl(uptr fixed_addr, uptr size, bool tolerate_enomem) {
   uptr PageSize = GetPageSizeCached();
   uptr p = internal_mmap((void*)(fixed_addr & ~(PageSize - 1)),
       RoundUpTo(size, PageSize),
       PROT_READ | PROT_WRITE,
       MAP_PRIVATE | MAP_ANON | MAP_FIXED,
       -1, 0);
   int reserrno;
   if (internal_iserror(p, &reserrno)) {
+    if (tolerate_enomem && reserrno == ENOMEM)
+      return nullptr;
     char mem_type[30];
     internal_snprintf(mem_type, sizeof(mem_type), "memory at address 0x%zx",
                       fixed_addr);
     ReportMmapFailureAndDie(size, mem_type, "allocate", reserrno);
   }
   IncreaseTotalMmap(size);
   return (void *)p;
 }
 
+void *MmapFixedOrDie(uptr fixed_addr, uptr size) {
+  return MmapFixedImpl(fixed_addr, size, false /*tolerate_enomem*/);
+}
+
+void *MmapFixedOrDieOnFatalError(uptr fixed_addr, uptr size) {
+  return MmapFixedImpl(fixed_addr, size, true /*tolerate_enomem*/);
+}
+
 bool MprotectNoAccess(uptr addr, uptr size) {
   return 0 == internal_mprotect((void*)addr, size, PROT_NONE);
 }
 
 bool MprotectReadOnly(uptr addr, uptr size) {
   return 0 == internal_mprotect((void *)addr, size, PROT_READ);
 }
 

lib/sanitizer_common/sanitizer_win.cc

if (p == 0) {
     char mem_type[30];
     internal_snprintf(mem_type, sizeof(mem_type), "memory at address 0x%zx",
                       fixed_addr);
     ReportMmapFailureAndDie(size, mem_type, "allocate", GetLastError());
   }
   return p;
 }
 
+void *MmapFixedOrDieOnFatalError(uptr fixed_addr, uptr size) {
+  void *p = VirtualAlloc((LPVOID)fixed_addr, size,
+      MEM_COMMIT, PAGE_READWRITE);
+  if (p == 0) {
+    char mem_type[30];
+    internal_snprintf(mem_type, sizeof(mem_type), "memory at address 0x%zx",
+                      fixed_addr);
+    return ReturnNullptrOnOOMOrDie(size, mem_type, "allocate");
+  }
+  return p;
+}
+
 void *MmapNoReserveOrDie(uptr size, const char *mem_type) {
   // FIXME: make this really NoReserve?
   return MmapOrDie(size, mem_type);
 }
 
 void *MmapFixedNoAccess(uptr fixed_addr, uptr size, const char *name) {
   (void)name; // unsupported
   void *res = VirtualAlloc((LPVOID)fixed_addr, size,

lib/sanitizer_common/tests/sanitizer_allocator_test.cc

TEST(SanitizerCommon, LargeMmapAllocatorMapUnmapCallback) {
   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);
 }
 
-template<class Allocator>
-void FailInAssertionOnOOM() {
-  Allocator a;
+// Don't test OOM conditions on Win64 because it causes other tests on the same
+// machine to OOM.
+#if SANITIZER_CAN_USE_ALLOCATOR64 && !SANITIZER_WINDOWS64 && !SANITIZER_ANDROID
+TEST(SanitizerCommon, SizeClassAllocator64Overflow) {
+  Allocator64 a;
   a.Init(kReleaseToOSIntervalNever);
-  SizeClassAllocatorLocalCache<Allocator> cache;
+  SizeClassAllocatorLocalCache<Allocator64> cache;
   memset(&cache, 0, sizeof(cache));
   cache.Init(0);
   AllocatorStats stats;
   stats.Init();
+
   const size_t kNumChunks = 128;
   uint32_t chunks[kNumChunks];
+  bool allocation_failed = false;
   for (int i = 0; i < 1000000; i++) {
-    a.GetFromAllocator(&stats, 52, chunks, kNumChunks);
+    if (!a.GetFromAllocator(&stats, 52, chunks, kNumChunks)) {
+      allocation_failed = true;
+      break;
     }
-
-  a.TestOnlyUnmap();
   }
+  EXPECT_EQ(allocation_failed, true);
 
-// Don't test OOM conditions on Win64 because it causes other tests on the same
-// machine to OOM.
-#if SANITIZER_CAN_USE_ALLOCATOR64 && !SANITIZER_WINDOWS64 && !SANITIZER_ANDROID
-TEST(SanitizerCommon, SizeClassAllocator64Overflow) {
-  EXPECT_DEATH(FailInAssertionOnOOM<Allocator64>(), "Out of memory");
+  a.TestOnlyUnmap();
 }
 #endif
 
 TEST(SanitizerCommon, LargeMmapAllocator) {
   LargeMmapAllocator<NoOpMapUnmapCallback, DieOnFailure> a;
   a.Init();
   AllocatorStats stats;
   stats.Init();
TEST(SanitizerCommon, SizeClassAllocator64PopulateFreeListOOM) {
 
   // ...one man is on a mission to overflow a region with a series of
   // successive allocations.
 
   const uptr kClassID = 107;
   const uptr kAllocationSize = SpecialSizeClassMap::Size(kClassID);
   ASSERT_LT(2 * kAllocationSize, kRegionSize);
   ASSERT_GT(3 * kAllocationSize, kRegionSize);
-  cache.Allocate(a, kClassID);
-  EXPECT_DEATH(cache.Allocate(a, kClassID) && cache.Allocate(a, kClassID),
-               "The process has exhausted");
+  EXPECT_NE(cache.Allocate(a, kClassID), nullptr);
+  EXPECT_NE(cache.Allocate(a, kClassID), nullptr);
+  EXPECT_EQ(cache.Allocate(a, kClassID), nullptr);
 
   const uptr Class2 = 100;
   const uptr Size2 = SpecialSizeClassMap::Size(Class2);
   ASSERT_EQ(Size2 * 8, kRegionSize);
   char *p[7];
   for (int i = 0; i < 7; i++) {
     p[i] = (char*)cache.Allocate(a, Class2);
+    EXPECT_NE(p[i], nullptr);
     fprintf(stderr, "p[%d] %p s = %lx\n", i, (void*)p[i], Size2);
     p[i][Size2 - 1] = 42;
     if (i) ASSERT_LT(p[i - 1], p[i]);
   }
-  EXPECT_DEATH(cache.Allocate(a, Class2), "The process has exhausted");
+  EXPECT_EQ(cache.Allocate(a, Class2), nullptr);
   cache.Deallocate(a, Class2, p[0]);
   cache.Drain(a);
   ASSERT_EQ(p[6][Size2 - 1], 42);
   a->TestOnlyUnmap();
   delete a;
 }
 
 #endif