Add Soft/Hard RSS Limits to Scudo Standalone

Reviewed By: vitalybuka

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

compiler-rt/lib/scudo/standalone/CMakeLists.txt

@@ -76,6 +76,7 @@ set(SCUDO_HEADERS
   quarantine.h
   release.h
   report.h
+  rss_limit_checker.h
   secondary.h
   size_class_map.h
   stack_depot.h
@@ -101,6 +102,7 @@ set(SCUDO_SOURCES
   linux.cpp
   release.cpp
   report.cpp
+  rss_limit_checker.cpp
   string_utils.cpp
   )
 

compiler-rt/lib/scudo/standalone/combined.h

@@ -18,6 +18,7 @@
 #include "options.h"
 #include "quarantine.h"
 #include "report.h"
+#include "rss_limit_checker.h"
 #include "secondary.h"
 #include "stack_depot.h"
 #include "string_utils.h"
@@ -147,6 +148,9 @@ public:
     initFlags();
     reportUnrecognizedFlags();
 
+    RssChecker.init(scudo::getFlags()->soft_rss_limit_mb,
+                    scudo::getFlags()->hard_rss_limit_mb);
+
     // Store some flags locally.
     if (getFlags()->may_return_null)
       Primary.Options.set(OptionBit::MayReturnNull);
@@ -346,6 +350,19 @@ public:
     }
     DCHECK_LE(Size, NeededSize);
 
+    switch (RssChecker.getRssLimitExceeded()) {
+    case RssLimitChecker::Neither:
+      break;
+    case RssLimitChecker::Soft:
+      if (Options.get(OptionBit::MayReturnNull))
+        return nullptr;
+      reportSoftRSSLimit(RssChecker.getSoftRssLimit());
+      break;
+    case RssLimitChecker::Hard:
+      reportHardRSSLimit(RssChecker.getHardRssLimit());
+      break;
+    }
+
     void *Block = nullptr;
     uptr ClassId = 0;
     uptr SecondaryBlockEnd = 0;
@@ -856,6 +873,13 @@ public:
            Header.State == Chunk::State::Allocated;
   }
 
+  void setRssLimitsTestOnly(int SoftRssLimitMb, int HardRssLimitMb,
+                            bool MayReturnNull) {
+    RssChecker.init(SoftRssLimitMb, HardRssLimitMb);
+    if (MayReturnNull)
+      Primary.Options.set(OptionBit::MayReturnNull);
+  }
+
   bool useMemoryTaggingTestOnly() const {
     return useMemoryTagging<Params>(Primary.Options.load());
   }
@@ -994,6 +1018,7 @@ private:
   QuarantineT Quarantine;
   TSDRegistryT TSDRegistry;
   pthread_once_t PostInitNonce = PTHREAD_ONCE_INIT;
+  RssLimitChecker RssChecker;
 
 #ifdef GWP_ASAN_HOOKS
   gwp_asan::GuardedPoolAllocator GuardedAlloc;

compiler-rt/lib/scudo/standalone/flags.inc

@@ -45,3 +45,12 @@ SCUDO_FLAG(bool, may_return_null, true,
 SCUDO_FLAG(int, release_to_os_interval_ms, SCUDO_ANDROID ? INT32_MIN : 5000,
            "Interval (in milliseconds) at which to attempt release of unused "
            "memory to the OS. Negative values disable the feature.")
+
+SCUDO_FLAG(int, hard_rss_limit_mb, 0,
+           "Hard RSS Limit in Mb. If non-zero, once the limit is achieved, "
+           "abort the process")
+
+SCUDO_FLAG(int, soft_rss_limit_mb, 0,
+           "Soft RSS Limit in Mb. If non-zero, once the limit is reached, all "
+           "subsequent calls will fail or return NULL until the RSS goes below "
+           "the soft limit")

compiler-rt/lib/scudo/standalone/report.cpp

@@ -36,6 +36,18 @@ private:
 
 inline void NORETURN trap() { __builtin_trap(); }
 
+void NORETURN reportSoftRSSLimit(uptr RssLimitMb) {
+  ScopedErrorReport Report;
+  Report.append("Soft RSS limit of %zu MB exhausted, current RSS is %zu MB\n",
+                RssLimitMb, GetRSS() >> 20);
+}
+
+void NORETURN reportHardRSSLimit(uptr RssLimitMb) {
+  ScopedErrorReport Report;
+  Report.append("Hard RSS limit of %zu MB exhausted, current RSS is %zu MB\n",
+                RssLimitMb, GetRSS() >> 20);
+}
+
 // This could potentially be called recursively if a CHECK fails in the reports.
 void NORETURN reportCheckFailed(const char *File, int Line,
                                 const char *Condition, u64 Value1, u64 Value2) {

compiler-rt/lib/scudo/standalone/report.h

@@ -33,6 +33,8 @@ void NORETURN reportAlignmentTooBig(uptr Alignment, uptr MaxAlignment);
 void NORETURN reportAllocationSizeTooBig(uptr UserSize, uptr TotalSize,
                                          uptr MaxSize);
 void NORETURN reportOutOfMemory(uptr RequestedSize);
+void NORETURN reportSoftRSSLimit(uptr RssLimitMb);
+void NORETURN reportHardRSSLimit(uptr RssLimitMb);
 enum class AllocatorAction : u8 {
   Recycling,
   Deallocating,

compiler-rt/lib/scudo/standalone/rss_limit_checker.cpp

@@ -0,0 +1,37 @@
+//===-- common.cpp ----------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "rss_limit_checker.h"
+#include "atomic_helpers.h"
+#include "string_utils.h"
+
+namespace scudo {
+
+void RssLimitChecker::check(u64 NextCheck) {
+  // The interval for the checks is 250ms.
+  static constexpr u64 CheckInterval = 250 * 1000000;
+
+  // Early return in case another thread already did the calculation.
+  if (!atomic_compare_exchange_strong(&RssNextCheckAtNS, &NextCheck,
+                                      getMonotonicTime() + CheckInterval,
+                                      memory_order_relaxed)) {
+    return;
+  }
+
+  const uptr CurrentRssMb = GetRSS() >> 20;
+
+  RssLimitExceeded Result = RssLimitExceeded::Neither;
+  if (UNLIKELY(HardRssLimitMb && HardRssLimitMb < CurrentRssMb))
+    Result = RssLimitExceeded::Hard;
+  else if (UNLIKELY(SoftRssLimitMb && SoftRssLimitMb < CurrentRssMb))
+    Result = RssLimitExceeded::Soft;
+
+  atomic_store_relaxed(&RssLimitStatus, static_cast<u8>(Result));
+}
+
+} // namespace scudo

compiler-rt/lib/scudo/standalone/rss_limit_checker.h

@@ -0,0 +1,63 @@
+//===-- common.h ------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCUDO_RSS_LIMIT_CHECKER_H_
+#define SCUDO_RSS_LIMIT_CHECKER_H_
+
+#include "atomic_helpers.h"
+#include "common.h"
+#include "internal_defs.h"
+
+namespace scudo {
+
+class RssLimitChecker {
+public:
+  enum RssLimitExceeded {
+    Neither,
+    Soft,
+    Hard,
+  };
+
+  void init(int SoftRssLimitMb, int HardRssLimitMb) {
+    CHECK_GE(SoftRssLimitMb, 0);
+    CHECK_GE(HardRssLimitMb, 0);
+    this->SoftRssLimitMb = static_cast<uptr>(SoftRssLimitMb);
+    this->HardRssLimitMb = static_cast<uptr>(HardRssLimitMb);
+  }
+
+  // Opportunistic RSS limit check. This will update the RSS limit status, if
+  // it can, every 250ms, otherwise it will just return the current one.
+  RssLimitExceeded getRssLimitExceeded() {
+    if (!HardRssLimitMb && !SoftRssLimitMb)
+      return RssLimitExceeded::Neither;
+
+    u64 NextCheck = atomic_load_relaxed(&RssNextCheckAtNS);
+    u64 Now = getMonotonicTime();
+
+    if (UNLIKELY(Now >= NextCheck))
+      check(NextCheck);
+
+    return static_cast<RssLimitExceeded>(atomic_load_relaxed(&RssLimitStatus));
+  }
+
+  uptr getSoftRssLimit() const { return SoftRssLimitMb; }
+  uptr getHardRssLimit() const { return HardRssLimitMb; }
+
+private:
+  void check(u64 NextCheck);
+
+  uptr SoftRssLimitMb = 0;
+  uptr HardRssLimitMb = 0;
+
+  atomic_u64 RssNextCheckAtNS = {};
+  atomic_u8 RssLimitStatus = {};
+};
+
+} // namespace scudo
+
+#endif // SCUDO_RSS_LIMIT_CHECKER_H_

compiler-rt/lib/scudo/standalone/tests/combined_test.cpp

@@ -731,3 +731,41 @@ TEST(ScudoCombinedTest, BasicTrustyConfig) {
 
 #endif
 #endif
+
+#if SCUDO_LINUX
+
+SCUDO_TYPED_TEST(ScudoCombinedTest, SoftRssLimit) {
+  auto *Allocator = this->Allocator.get();
+  Allocator->setRssLimitsTestOnly(1, 0, true);
+
+  size_t Megabyte = 1024 * 1024;
+  size_t ChunkSize = 16;
+  size_t Error = 256;
+
+  std::vector<void *> Ptrs;
+  for (size_t index = 0; index < Megabyte + Error; index += ChunkSize) {
+    void *Ptr = Allocator->allocate(ChunkSize, Origin);
+    Ptrs.push_back(Ptr);
+  }
+
+  EXPECT_EQ(nullptr, Allocator->allocate(ChunkSize, Origin));
+
+  for (void *Ptr : Ptrs)
+    Allocator->deallocate(Ptr, Origin);
+}
+
+SCUDO_TYPED_TEST(ScudoCombinedTest, HardRssLimit) {
+  auto *Allocator = this->Allocator.get();
+  Allocator->setRssLimitsTestOnly(0, 1, false);
+
+  size_t Megabyte = 1024 * 1024;
+
+  EXPECT_DEATH(
+      {
+        disableDebuggerdMaybe();
+        Allocator->allocate(Megabyte, Origin);
+      },
+      "");
+}
+
+#endif