[ASan] Move rss_limit_is_exceeded_ flag to ASan.

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

Reviewers: eugenis

Subscribers: kubamracek, llvm-commits

Differential Revision: https://reviews.llvm.org/D34180

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@305342 91177308-0d34-0410-b5e6-96231b3b80d8

Change-Id: I2c7ccdfbfd420b74cbc5c96a50933bace9b88ccf

2 files changed, 21 insertions, 19 deletions

diff --git a/libsanitizer/asan/asan_allocator.cc b/libsanitizer/asan/asan_allocator.cc
index 916d2848b23..b6e2cfa544a 100644
--- a/libsanitizer/asan/asan_allocator.cc
+++ b/libsanitizer/asan/asan_allocator.cc
@@ -233,6 +233,8 @@ struct Allocator {
   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;
@@ -266,6 +268,14 @@ struct Allocator {
     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 a user-facing chunk (with redzones), or some internal
     // housekeeping chunk, like TransferBatch. Start by assuming the former.
@@ -360,6 +370,8 @@ struct Allocator {
                  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;
@@ -397,16 +409,15 @@ struct Allocator {
 
     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();
@@ -833,8 +844,8 @@ 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
diff --git a/libsanitizer/sanitizer_common/sanitizer_allocator_combined.h b/libsanitizer/sanitizer_common/sanitizer_allocator_combined.h
index 159e75525dd..3f643240b9d 100644
--- a/libsanitizer/sanitizer_common/sanitizer_allocator_combined.h
+++ b/libsanitizer/sanitizer_common/sanitizer_allocator_combined.h
@@ -41,12 +41,12 @@ class CombinedAllocator {
   }
 
   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
@@ -87,7 +87,8 @@ class CombinedAllocator {
   }
 
   void *ReturnNullOrDieOnOOM() {
-    if (MayReturnNull()) return nullptr;
+    if (MayReturnNull())
+      return nullptr;
     ReportAllocatorCannotReturnNull(true);
   }
 
@@ -104,15 +105,6 @@ class CombinedAllocator {
     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))
@@ -232,5 +224,4 @@ class CombinedAllocator {
   SecondaryAllocator secondary_;
   AllocatorGlobalStats stats_;
   atomic_uint8_t may_return_null_;
-  atomic_uint8_t rss_limit_is_exceeded_;
 };