[asan] new flag: hard_rss_limit_mb

llvm-svn: 224353

compiler-rt/lib/asan/asan_interceptors.cc

@@ -236,6 +236,11 @@ INTERCEPTOR(int, pthread_create, void *thread,
   }
   return result;
 }
+
+INTERCEPTOR(int, pthread_join, void *t, void **arg) {
+  return real_pthread_join(t, arg);
+}
+DEFINE_REAL_PTHREAD_FUNCTIONS;
 #endif  // ASAN_INTERCEPT_PTHREAD_CREATE
 
 #if ASAN_INTERCEPT_SIGNAL_AND_SIGACTION
@@ -902,6 +907,7 @@ void InitializeAsanInterceptors() {
   // Intercept threading-related functions
 #if ASAN_INTERCEPT_PTHREAD_CREATE
   ASAN_INTERCEPT_FUNC(pthread_create);
+  ASAN_INTERCEPT_FUNC(pthread_join);
 #endif
 
   // Intercept atexit function.

compiler-rt/lib/asan/asan_rtl.cc

@@ -658,6 +658,8 @@ static void AsanInitInternal() {
   InitializeAllocator(common_flags()->allocator_may_return_null,
                       flags()->quarantine_size);
 
+  MaybeStartBackgroudThread();
+
   // On Linux AsanThread::ThreadStart() calls malloc() that's why asan_inited
   // should be set to 1 prior to initializing the threads.
   asan_inited = 1;

compiler-rt/lib/sanitizer_common/sanitizer_common.h

@@ -557,6 +557,10 @@ INLINE void GetExtraActivationFlags(char *buf, uptr size) { *buf = '\0'; }
 INLINE void SanitizerInitializeUnwinder() {}
 #endif
 
+void *internal_start_thread(void(*func)(void*), void *arg);
+void internal_join_thread(void *th);
+void MaybeStartBackgroudThread();
+
 // Make the compiler think that something is going on there.
 // Use this inside a loop that looks like memset/memcpy/etc to prevent the
 // compiler from recognising it and turning it into an actual call to

compiler-rt/lib/sanitizer_common/sanitizer_common_libcdep.cc

@@ -13,6 +13,7 @@
 
 #include "sanitizer_common.h"
 #include "sanitizer_flags.h"
+#include "sanitizer_stackdepot.h"
 #include "sanitizer_stacktrace.h"
 #include "sanitizer_symbolizer.h"
 
@@ -59,6 +60,48 @@ void ReportErrorSummary(const char *error_type, StackTrace *stack) {
 #endif
 }
 
+void BackgroundThread(void *arg) {
+  uptr hard_rss_limit_mb = common_flags()->hard_rss_limit_mb;
+  uptr prev_reported_rss = 0;
+  uptr prev_reported_stack_depot_size = 0;
+  while (true) {
+    SleepForMillis(100);
+    uptr current_rss_mb = GetRSS() >> 20;
+    if (common_flags()->verbosity) {
+      // If RSS has grown 10% since last time, print some information.
+      if (prev_reported_rss * 11 / 10 < current_rss_mb) {
+        Printf("%s: RSS: %zdMb\n", SanitizerToolName, current_rss_mb);
+        prev_reported_rss = current_rss_mb;
+      }
+      // If stack depot has grown 10% since last time, print it too.
+      StackDepotStats *stack_depot_stats = StackDepotGetStats();
+      if (prev_reported_stack_depot_size * 11 / 10 <
+          stack_depot_stats->allocated) {
+        Printf("%s: StackDepot: %zd ids; %zdM allocated\n",
+               SanitizerToolName,
+               stack_depot_stats->n_uniq_ids,
+               stack_depot_stats->allocated >> 20);
+        prev_reported_stack_depot_size = stack_depot_stats->allocated;
+      }
+    }
+    // Check RSS against the limit.
+    if (hard_rss_limit_mb && hard_rss_limit_mb < current_rss_mb) {
+      Report("%s: hard rss limit exhausted (%zdMb vs %zdMb)\n",
+             SanitizerToolName, hard_rss_limit_mb, current_rss_mb);
+      DumpProcessMap();
+      Die();
+    }
+  }
+}
+
+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;
+  if (!real_pthread_create) return;  // Can't spawn the thread anyway.
+  internal_start_thread(BackgroundThread, nullptr);
+}
+
 }  // namespace __sanitizer
 
 void NOINLINE

compiler-rt/lib/sanitizer_common/sanitizer_flags.cc

@@ -140,6 +140,11 @@ void ParseCommonFlagsFromString(CommonFlags *f, const char *str) {
   ParseFlag(str, &f->mmap_limit_mb, "mmap_limit_mb",
             "Limit the amount of mmap-ed memory (excluding shadow) in Mb; "
             "not a user-facing flag, used mosly for testing the tools");
+  ParseFlag(str, &f->hard_rss_limit_mb, "hard_rss_limit_mb",
+            "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");
   ParseFlag(str, &f->coverage, "coverage",
       "If set, coverage information will be dumped at program shutdown (if the "
       "coverage instrumentation was enabled at compile time).");

compiler-rt/lib/sanitizer_common/sanitizer_flags.h

@@ -54,6 +54,7 @@ struct CommonFlags {
   bool intercept_tls_get_addr;
   bool help;
   uptr mmap_limit_mb;
+  uptr hard_rss_limit_mb;
   bool coverage;
   bool coverage_direct;
   const char *coverage_dir;

compiler-rt/lib/sanitizer_common/sanitizer_libc.h

@@ -98,6 +98,25 @@ int internal_fork();
 // Threading
 uptr internal_sched_yield();
 
+// These functions call appropriate pthread_ functions directly, bypassing
+// the interceptor. They are weak and may not be present in some tools.
+SANITIZER_WEAK_ATTRIBUTE
+int real_pthread_create(void *th, void *attr, void *(*callback)(void *),
+                        void *param);
+SANITIZER_WEAK_ATTRIBUTE
+int real_pthread_join(void *th, void **ret);
+
+#define DEFINE_REAL_PTHREAD_FUNCTIONS                                          \
+  namespace __sanitizer {                                                      \
+  int real_pthread_create(void *th, void *attr, void *(*callback)(void *),     \
+                          void *param) {                                       \
+    return REAL(pthread_create)(th, attr, callback, param);                    \
+  }                                                                            \
+  int real_pthread_join(void *th, void **ret) {                                \
+    return REAL(pthread_join(th, ret));                                        \
+  }                                                                            \
+  }  // namespace __sanitizer
+
 // Error handling
 bool internal_iserror(uptr retval, int *rverrno = 0);
 

compiler-rt/lib/sanitizer_common/sanitizer_linux.cc

@@ -938,6 +938,21 @@ bool IsDeadlySignal(int signum) {
   return (signum == SIGSEGV) && common_flags()->handle_segv;
 }
 
+void *internal_start_thread(void(*func)(void *arg), void *arg) {
+  // Start the thread with signals blocked, otherwise it can steal user signals.
+  __sanitizer_sigset_t set, old;
+  internal_sigfillset(&set);
+  internal_sigprocmask(SIG_SETMASK, &set, &old);
+  void *th;
+  real_pthread_create(&th, 0, (void*(*)(void *arg))func, arg);
+  internal_sigprocmask(SIG_SETMASK, &old, 0);
+  return th;
+}
+
+void internal_join_thread(void *th) {
+  real_pthread_join(th, 0);
+}
+
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_FREEBSD || SANITIZER_LINUX

compiler-rt/lib/sanitizer_common/sanitizer_mac.cc

@@ -325,6 +325,9 @@ uptr GetRSS() {
   return 0;
 }
 
+void *internal_start_thread(void (*func)(void *arg), void *arg) { return 0; }
+void internal_join_thread(void *th) { return 0; }
+
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_MAC

compiler-rt/lib/sanitizer_common/sanitizer_win.cc

@@ -379,6 +379,9 @@ uptr GetRSS() {
   return 0;
 }
 
+void *internal_start_thread(void (*func)(void *arg), void *arg) { return 0; }
+void internal_join_thread(void *th) { return 0; }
+
 // ---------------------- BlockingMutex ---------------- {{{1
 const uptr LOCK_UNINITIALIZED = 0;
 const uptr LOCK_READY = (uptr)-1;

compiler-rt/lib/tsan/rtl/tsan_interceptors.cc

@@ -948,6 +948,8 @@ TSAN_INTERCEPTOR(int, pthread_join, void *th, void **ret) {
   return res;
 }
 
+DEFINE_REAL_PTHREAD_FUNCTIONS;
+
 TSAN_INTERCEPTOR(int, pthread_detach, void *th) {
   SCOPED_TSAN_INTERCEPTOR(pthread_detach, th);
   int tid = ThreadTid(thr, pc, (uptr)th);
@@ -2565,19 +2567,4 @@ void InitializeInterceptors() {
   FdInit();
 }
 
-void *internal_start_thread(void(*func)(void *arg), void *arg) {
-  // Start the thread with signals blocked, otherwise it can steal user signals.
-  __sanitizer_sigset_t set, old;
-  internal_sigfillset(&set);
-  internal_sigprocmask(SIG_SETMASK, &set, &old);
-  void *th;
-  REAL(pthread_create)(&th, 0, (void*(*)(void *arg))func, arg);
-  internal_sigprocmask(SIG_SETMASK, &old, 0);
-  return th;
-}
-
-void internal_join_thread(void *th) {
-  REAL(pthread_join)(th, 0);
-}
-
 }  // namespace __tsan

compiler-rt/lib/tsan/rtl/tsan_platform.h

@@ -252,9 +252,6 @@ void InitializePlatform();
 void FlushShadowMemory();
 void WriteMemoryProfile(char *buf, uptr buf_size, uptr nthread, uptr nlive);
 
-void *internal_start_thread(void(*func)(void*), void *arg);
-void internal_join_thread(void *th);
-
 // Says whether the addr relates to a global var.
 // Guesses with high probability, may yield both false positives and negatives.
 bool IsGlobalVar(uptr addr);

compiler-rt/test/sanitizer_common/TestCases/Linux/hard_rss_limit_mb_test.cc

@@ -0,0 +1,35 @@
+// Check hard_rss_limit_mb. Not all sanitizers implement it yet.
+// RUN: %clangxx -O2 %s -o %t
+//
+// Run with limit should fail:
+// RUN: %tool_options=hard_rss_limit_mb=100 not %run %t 2>&1 | FileCheck %s
+//
+// Run w/o limit or with a large enough limit should pass:
+// RUN: %tool_options=hard_rss_limit_mb=1000 %run %t
+// RUN: %run %t
+//
+// FIXME: make it work for other sanitizers.
+// XFAIL: lsan
+// XFAIL: tsan
+// XFAIL: msan
+
+#include <string.h>
+#include <stdio.h>
+#include <unistd.h>
+
+const int kNumAllocs = 200 * 1000;
+const int kAllocSize = 1000;
+volatile char *sink[kNumAllocs];
+
+int main(int argc, char **argv) {
+  for (int i = 0; i < kNumAllocs; i++) {
+    if ((i % 1000) == 0) {
+      fprintf(stderr, "[%d]\n", i);
+    }
+    char *x = new char[kAllocSize];
+    memset(x, 0, kAllocSize);
+    sink[i] = x;
+  }
+  sleep(1);  // Make sure the background thread has time to kill the process.
+// CHECK: hard rss limit exhausted
+}