[asan] add a flag soft_rss_limit_mb

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

9 files changed, 132 insertions, 6 deletions

diff --git a/lib/asan/asan_allocator.cc b/lib/asan/asan_allocator.cc
index e2e0fb55b..9525cf536 100644
--- a/lib/asan/asan_allocator.cc
+++ b/lib/asan/asan_allocator.cc
@@ -361,15 +361,21 @@ 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);
+      allocated =
+          allocator.Allocate(cache, needed_size, 8, false, check_rss_limit);
     } else {
       SpinMutexLock l(&fallback_mutex);
       AllocatorCache *cache = &fallback_allocator_cache;
-      allocated = allocator.Allocate(cache, needed_size, 8, false);
+      allocated =
+          allocator.Allocate(cache, needed_size, 8, false, check_rss_limit);
     }
 
+    if (!allocated)
+      return allocator.ReturnNullOrDie();
+
     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
@@ -771,6 +777,10 @@ void asan_mz_force_unlock() {
   instance.ForceUnlock();
 }
 
+void AsanSoftRssLimitExceededCallback(bool exceeded) {
+  instance.allocator.SetRssLimitIsExceeded(exceeded);
+}
+
 }  // namespace __asan
 
 // --- Implementation of LSan-specific functions --- {{{1
diff --git a/lib/asan/asan_allocator.h b/lib/asan/asan_allocator.h
index 521d47b0b..3208d1f95 100644
--- a/lib/asan/asan_allocator.h
+++ b/lib/asan/asan_allocator.h
@@ -173,6 +173,7 @@ void asan_mz_force_lock();
 void asan_mz_force_unlock();
 
 void PrintInternalAllocatorStats();
+void AsanSoftRssLimitExceededCallback(bool exceeded);
 
 }  // namespace __asan
 #endif  // ASAN_ALLOCATOR_H
diff --git a/lib/asan/asan_rtl.cc b/lib/asan/asan_rtl.cc
index 32ad0f208..9d994885e 100644
--- a/lib/asan/asan_rtl.cc
+++ b/lib/asan/asan_rtl.cc
@@ -396,6 +396,7 @@ static void AsanInitInternal() {
   InitializeAllocator(allocator_options);
 
   MaybeStartBackgroudThread();
+  SetSoftRssLimitExceededCallback(AsanSoftRssLimitExceededCallback);
 
   // On Linux AsanThread::ThreadStart() calls malloc() that's why asan_inited
   // should be set to 1 prior to initializing the threads.
diff --git a/lib/sanitizer_common/sanitizer_allocator.h b/lib/sanitizer_common/sanitizer_allocator.h
index d749acb90..8ff3f20f1 100644
--- a/lib/sanitizer_common/sanitizer_allocator.h
+++ b/lib/sanitizer_common/sanitizer_allocator.h
@@ -1279,12 +1279,14 @@ class CombinedAllocator {
   }
 
   void *Allocate(AllocatorCache *cache, uptr size, uptr alignment,
-                 bool cleared = false) {
+                 bool cleared = false, bool check_rss_limit = false) {
     // Returning 0 on malloc(0) may break a lot of code.
     if (size == 0)
       size = 1;
     if (size + alignment < size)
       return ReturnNullOrDie();
+    if (check_rss_limit && RssLimitIsExceeded())
+      return ReturnNullOrDie();
     if (alignment > 8)
       size = RoundUpTo(size, alignment);
     void *res;
@@ -1315,6 +1317,15 @@ class CombinedAllocator {
     atomic_store(&may_return_null_, may_return_null, memory_order_release);
   }
 
+  bool RssLimitIsExceeded() {
+    return atomic_load(&rss_limit_is_exceeded_, memory_order_release);
+  }
+
+  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))
@@ -1428,6 +1439,7 @@ class CombinedAllocator {
   SecondaryAllocator secondary_;
   AllocatorGlobalStats stats_;
   atomic_uint8_t may_return_null_;
+  atomic_uint8_t rss_limit_is_exceeded_;
 };
 
 // Returns true if calloc(size, n) should return 0 due to overflow in size*n.
diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index a25bacb3f..71bf49ef5 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -254,6 +254,12 @@ typedef void (*CheckFailedCallbackType)(const char *, int, const char *,
                                        u64, u64);
 void SetCheckFailedCallback(CheckFailedCallbackType callback);
 
+// Callback will be called if soft_rss_limit_mb is given and the limit is
+// exceeded (exceeded==true) or if rss went down below the limit
+// (exceeded==false).
+// The callback should be registered once at the tool init time.
+void SetSoftRssLimitExceededCallback(void (*Callback)(bool exceeded));
+
 // Functions related to signal handling.
 typedef void (*SignalHandlerType)(int, void *, void *);
 bool IsDeadlySignal(int signum);
diff --git a/lib/sanitizer_common/sanitizer_common_libcdep.cc b/lib/sanitizer_common/sanitizer_common_libcdep.cc
index ebd693eaa..8c2d3f088 100644
--- a/lib/sanitizer_common/sanitizer_common_libcdep.cc
+++ b/lib/sanitizer_common/sanitizer_common_libcdep.cc
@@ -60,10 +60,18 @@ void ReportErrorSummary(const char *error_type, StackTrace *stack) {
 #endif
 }
 
+static void (*SoftRssLimitExceededCallback)(bool exceeded);
+void SetSoftRssLimitExceededCallback(void (*Callback)(bool exceeded)) {
+  CHECK_EQ(SoftRssLimitExceededCallback, nullptr);
+  SoftRssLimitExceededCallback = Callback;
+}
+
 void BackgroundThread(void *arg) {
   uptr hard_rss_limit_mb = common_flags()->hard_rss_limit_mb;
+  uptr soft_rss_limit_mb = common_flags()->soft_rss_limit_mb;
   uptr prev_reported_rss = 0;
   uptr prev_reported_stack_depot_size = 0;
+  bool reached_soft_rss_limit = false;
   while (true) {
     SleepForMillis(100);
     uptr current_rss_mb = GetRSS() >> 20;
@@ -91,13 +99,28 @@ void BackgroundThread(void *arg) {
       DumpProcessMap();
       Die();
     }
+    if (soft_rss_limit_mb) {
+      if (soft_rss_limit_mb < current_rss_mb && !reached_soft_rss_limit) {
+        reached_soft_rss_limit = true;
+        Report("%s: soft rss limit exhausted (%zdMb vs %zdMb)\n",
+               SanitizerToolName, soft_rss_limit_mb, current_rss_mb);
+        if (SoftRssLimitExceededCallback)
+          SoftRssLimitExceededCallback(true);
+      } else if (soft_rss_limit_mb >= current_rss_mb &&
+                 reached_soft_rss_limit) {
+        reached_soft_rss_limit = false;
+        if (SoftRssLimitExceededCallback)
+          SoftRssLimitExceededCallback(false);
+      }
+    }
   }
 }
 
 void MaybeStartBackgroudThread() {
   if (!SANITIZER_LINUX) return;  // Need to implement/test on other platforms.
-  // Currently, only start the background thread if hard_rss_limit_mb is given.
-  if (!common_flags()->hard_rss_limit_mb) return;
+  // Start the background thread if one of the rss limits is given.
+  if (!common_flags()->hard_rss_limit_mb &&
+      !common_flags()->soft_rss_limit_mb) return;
   if (!real_pthread_create) return;  // Can't spawn the thread anyway.
   internal_start_thread(BackgroundThread, nullptr);
 }
diff --git a/lib/sanitizer_common/sanitizer_flags.inc b/lib/sanitizer_common/sanitizer_flags.inc
index bf7941218..c39e7a698 100644
--- a/lib/sanitizer_common/sanitizer_flags.inc
+++ b/lib/sanitizer_common/sanitizer_flags.inc
@@ -91,10 +91,19 @@ COMMON_FLAG(uptr, mmap_limit_mb, 0,
             "Limit the amount of mmap-ed memory (excluding shadow) in Mb; "
             "not a user-facing flag, used mosly for testing the tools")
 COMMON_FLAG(uptr, hard_rss_limit_mb, 0,
-            "RSS limit in Mb."
+            "Hard RSS limit in Mb."
             " If non-zero, a background thread is spawned at startup"
             " which periodically reads RSS and aborts the process if the"
             " limit is reached")
+COMMON_FLAG(uptr, soft_rss_limit_mb, 0,
+            "Soft RSS limit in Mb."
+            " If non-zero, a background thread is spawned at startup"
+            " which periodically reads RSS. If the limit is reached"
+            " all subsequent malloc/new calls will fail or return NULL"
+            " (depending on the value of allocator_may_return_null)"
+            " until the RSS goes below the soft limit."
+            " This limit does not affect memory allocations other than"
+            " malloc/new.")
 COMMON_FLAG(
     bool, coverage, false,
     "If set, coverage information will be dumped at program shutdown (if the "
diff --git a/lib/sanitizer_common/sanitizer_quarantine.h b/lib/sanitizer_common/sanitizer_quarantine.h
index c04dc6fac..404d3753f 100644
--- a/lib/sanitizer_common/sanitizer_quarantine.h
+++ b/lib/sanitizer_common/sanitizer_quarantine.h
@@ -134,6 +134,7 @@ class QuarantineCache {
       size += sizeof(QuarantineBatch);  // Count the batch in Quarantine size.
     }
     QuarantineBatch *b = list_.back();
+    CHECK(b);
     b->batch[b->count++] = ptr;
     b->size += size;
     SizeAdd(size);
@@ -172,6 +173,7 @@ class QuarantineCache {
 
   NOINLINE QuarantineBatch* AllocBatch(Callback cb) {
     QuarantineBatch *b = (QuarantineBatch *)cb.Allocate(sizeof(*b));
+    CHECK(b);
     b->count = 0;
     b->size = 0;
     list_.push_back(b);
diff --git a/test/sanitizer_common/TestCases/Linux/soft_rss_limit_mb_test.cc b/test/sanitizer_common/TestCases/Linux/soft_rss_limit_mb_test.cc
new file mode 100644
index 000000000..187139e1a
--- /dev/null
+++ b/test/sanitizer_common/TestCases/Linux/soft_rss_limit_mb_test.cc
@@ -0,0 +1,62 @@
+// Check soft_rss_limit_mb. Not all sanitizers implement it yet.
+// RUN: %clangxx -O2 %s -o %t
+//
+// Run with limit should fail:
+// RUN: %tool_options=soft_rss_limit_mb=400:quarantine_size=1:allocator_may_return_null=1     %run %t 2>&1 | FileCheck %s -check-prefix=CHECK_MAY_RETURN_1
+// RUN: %tool_options=soft_rss_limit_mb=400:quarantine_size=1:allocator_may_return_null=0 not %run %t 2>&1 | FileCheck %s -check-prefix=CHECK_MAY_RETURN_0
+
+// FIXME: make it work for other sanitizers.
+// XFAIL: lsan
+// XFAIL: tsan
+// XFAIL: msan
+#include <stdlib.h>
+#include <stdio.h>
+#include <string.h>
+#include <unistd.h>
+
+static const int kMaxNumAllocs = 1 << 10;
+static const int kAllocSize = 1 << 20;  // Large enough to go vi mmap.
+
+static char *allocs[kMaxNumAllocs];
+
+int main() {
+  int num_allocs = kMaxNumAllocs / 4;
+  for (int i = 0; i < 3; i++, num_allocs *= 2) {
+    fprintf(stderr, "[%d] allocating %d times\n", i, num_allocs);
+    int zero_results = 0;
+    for (int j = 0; j < num_allocs; j++) {
+      if ((j % (num_allocs / 4)) == 0) {
+        usleep(100000);
+        fprintf(stderr, "  [%d]\n", j);
+      }
+      allocs[j] = (char*)malloc(kAllocSize);
+      if (allocs[j])
+        memset(allocs[j], -1, kAllocSize);
+      else
+        zero_results++;
+    }
+    if (zero_results)
+      fprintf(stderr, "Some of the malloc calls returned null: %d\n",
+              zero_results);
+    if (zero_results != num_allocs)
+      fprintf(stderr, "Some of the malloc calls returned non-null: %d\n",
+              num_allocs - zero_results);
+    for (int j = 0; j < num_allocs; j++) {
+      free(allocs[j]);
+    }
+  }
+}
+
+// CHECK_MAY_RETURN_1: allocating 256 times
+// CHECK_MAY_RETURN_1: Some of the malloc calls returned non-null: 256
+// CHECK_MAY_RETURN_1: allocating 512 times
+// CHECK_MAY_RETURN_1: Some of the malloc calls returned null:
+// CHECK_MAY_RETURN_1: Some of the malloc calls returned non-null:
+// CHECK_MAY_RETURN_1: allocating 1024 times
+// CHECK_MAY_RETURN_1: Some of the malloc calls returned null:
+// CHECK_MAY_RETURN_1: Some of the malloc calls returned non-null:
+
+// CHECK_MAY_RETURN_0: allocating 256 times
+// CHECK_MAY_RETURN_0: Some of the malloc calls returned non-null: 256
+// CHECK_MAY_RETURN_0: allocating 512 times
+// CHECK_MAY_RETURN_0: allocator is terminating the process instead of returning