[ASan] Move rss_limit_is_exceeded_ flag to ASan.

[alekseyshl]

Summary

Move the OOM decision based on RSS limits out of generic allocator to
ASan allocator, where it makes more sense at the moment.

[kcc]

Why?
It would be lovely to have the same flag work in other sanitizers...

[alekseyshl]

In D34180#779694, @kcc wrote:

Why?
It would be lovely to have the same flag work in other sanitizers...

Yep, other sanitizers need this functionality too, but the current way is pretty ugly (hinted by "bool check_rss_limit = false" parameter). I'd rather have the flag and the callback in sanitizer_allocator.cc, but first I'd like to finish changes addressing allocator_may_return_null issues (in a similar way, trying to move may_return_null_ flag out of the generic allocator) and then revisit RSS limit behavior.

[kcc]

not sure I understand. you want to move this from common to asan and then back?

[alekseyshl]

Nope, not back, but to sanitizer_allocator.cc. Not doing it right now since I want to see how may_return_null_ refactoring will turn out.

[kcc]

LGTM

lib/asan/asan_allocator.cc

static const uptr kMaxAllowedMallocSize =
       FIRST_32_SECOND_64(3UL << 30, 1ULL << 40);
 
   AsanAllocator allocator;
   AsanQuarantine quarantine;
   StaticSpinMutex fallback_mutex;
   AllocatorCache fallback_allocator_cache;
   QuarantineCache fallback_quarantine_cache;
 
+  atomic_uint8_t rss_limit_exceeded;
+
   // ------------------- Options --------------------------
   atomic_uint16_t min_redzone;
   atomic_uint16_t max_redzone;
   atomic_uint8_t alloc_dealloc_mismatch;
 
   // ------------------- Initialization ------------------------
   explicit Allocator(LinkerInitialized)
       : quarantine(LINKER_INITIALIZED),
void SharedInitCode(const AllocatorOptions &options) {
     atomic_store(&max_redzone, options.max_redzone, memory_order_release);
   }
 
   void Initialize(const AllocatorOptions &options) {
     allocator.Init(options.may_return_null, options.release_to_os_interval_ms);
     SharedInitCode(options);
   }
 
+  bool RssLimitExceeded() {
+    return atomic_load(&rss_limit_exceeded, memory_order_relaxed);
+  }
+
+  void SetRssLimitExceeded(bool limit_exceeded) {
+    atomic_store(&rss_limit_exceeded, limit_exceeded, memory_order_relaxed);
+  }
+
   void RePoisonChunk(uptr chunk) {
     // This could be a user-facing chunk (with redzones), or some internal
     // housekeeping chunk, like TransferBatch. Start by assuming the former.
     AsanChunk *ac = GetAsanChunk((void *)chunk);
     uptr allocated_size = allocator.GetActuallyAllocatedSize((void *)ac);
     uptr beg = ac->Beg();
     uptr end = ac->Beg() + ac->UsedSize(true);
     uptr chunk_end = chunk + allocated_size;
AsanChunk *ChooseChunk(uptr addr, AsanChunk *left_chunk,
     return right_chunk;
   }
 
   // -------------------- Allocation/Deallocation routines ---------------
   void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
                  AllocType alloc_type, bool can_fill) {
     if (UNLIKELY(!asan_inited))
       AsanInitFromRtl();
+    if (RssLimitExceeded())
+      return allocator.ReturnNullOrDieOnOOM();
     Flags &fl = *flags();
     CHECK(stack);
     const uptr min_alignment = SHADOW_GRANULARITY;
     if (alignment < min_alignment)
       alignment = min_alignment;
     if (size == 0) {
       // We'd be happy to avoid allocating memory for zero-size requests, but
       // some programs/tests depend on this behavior and assume that malloc
void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
     if (size > kMaxAllowedMallocSize || needed_size > kMaxAllowedMallocSize) {
       Report("WARNING: AddressSanitizer failed to allocate 0x%zx bytes\n",
              (void*)size);
       return allocator.ReturnNullOrDieOnBadRequest();
     }
 
     AsanThread *t = GetCurrentThread();
     void *allocated;
-    bool check_rss_limit = true;
     if (t) {
       AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
       allocated =
-          allocator.Allocate(cache, needed_size, 8, false, check_rss_limit);
+          allocator.Allocate(cache, needed_size, 8, false);
     } else {
       SpinMutexLock l(&fallback_mutex);
       AllocatorCache *cache = &fallback_allocator_cache;
       allocated =
-          allocator.Allocate(cache, needed_size, 8, false, check_rss_limit);
+          allocator.Allocate(cache, needed_size, 8, false);
     }
 
     if (!allocated) return allocator.ReturnNullOrDieOnOOM();
 
     if (*(u8 *)MEM_TO_SHADOW((uptr)allocated) == 0 && CanPoisonMemory()) {
       // Heap poisoning is enabled, but the allocator provides an unpoisoned
       // chunk. This is possible if CanPoisonMemory() was false for some
       // time, for example, due to flags()->start_disabled.

 void asan_mz_force_lock() {
   instance.ForceLock();
 }
 
 void asan_mz_force_unlock() {
   instance.ForceUnlock();
 }
 
-void AsanSoftRssLimitExceededCallback(bool exceeded) {
-  instance.allocator.SetRssLimitIsExceeded(exceeded);
+void AsanSoftRssLimitExceededCallback(bool limit_exceeded) {
+  instance.SetRssLimitExceeded(limit_exceeded);
 }
 
 } // namespace __asan
 
 // --- Implementation of LSan-specific functions --- {{{1
 namespace __lsan {
 void LockAllocator() {
   __asan::get_allocator().ForceLock();

lib/sanitizer_common/sanitizer_allocator_combined.h

public:
 
   void Init(bool may_return_null, s32 release_to_os_interval_ms) {
     secondary_.Init(may_return_null);
     stats_.Init();
     InitCommon(may_return_null, release_to_os_interval_ms);
   }
 
   void *Allocate(AllocatorCache *cache, uptr size, uptr alignment,
-                 bool cleared = false, bool check_rss_limit = false) {
+                 bool cleared = false) {
     // Returning 0 on malloc(0) may break a lot of code.
     if (size == 0)
       size = 1;
-    if (size + alignment < size) return ReturnNullOrDieOnBadRequest();
-    if (check_rss_limit && RssLimitIsExceeded()) return ReturnNullOrDieOnOOM();
+    if (size + alignment < size)
+      return ReturnNullOrDieOnBadRequest();
     uptr original_size = size;
     // If alignment requirements are to be fulfilled by the frontend allocator
     // rather than by the primary or secondary, passing an alignment lower than
     // or equal to 8 will prevent any further rounding up, as well as the later
     // alignment check.
     if (alignment > 8)
       size = RoundUpTo(size, alignment);
     void *res;
public:
 
   void *ReturnNullOrDieOnBadRequest() {
     if (MayReturnNull())
       return nullptr;
     ReportAllocatorCannotReturnNull(false);
   }
 
   void *ReturnNullOrDieOnOOM() {
-    if (MayReturnNull()) return nullptr;
+    if (MayReturnNull())
+      return nullptr;
     ReportAllocatorCannotReturnNull(true);
   }
 
   void SetMayReturnNull(bool may_return_null) {
     secondary_.SetMayReturnNull(may_return_null);
     atomic_store(&may_return_null_, may_return_null, memory_order_release);
   }
 
   s32 ReleaseToOSIntervalMs() const {
     return primary_.ReleaseToOSIntervalMs();
   }
 
   void SetReleaseToOSIntervalMs(s32 release_to_os_interval_ms) {
     primary_.SetReleaseToOSIntervalMs(release_to_os_interval_ms);
   }
 
-  bool RssLimitIsExceeded() {
-    return atomic_load(&rss_limit_is_exceeded_, memory_order_acquire);
-  }
-
-  void SetRssLimitIsExceeded(bool rss_limit_is_exceeded) {
-    atomic_store(&rss_limit_is_exceeded_, rss_limit_is_exceeded,
-                 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) {
     secondary_.ForEachChunk(callback, arg);
   }
 
  private:
   PrimaryAllocator primary_;
   SecondaryAllocator secondary_;
   AllocatorGlobalStats stats_;
   atomic_uint8_t may_return_null_;
-  atomic_uint8_t rss_limit_is_exceeded_;
 };