[compiler-rt] [sanitizer] Add graceful handling of exceeding StackStore limit. (PR #76115)

[via llvm-commits]

https://github.com/kda created https://github.com/llvm/llvm-project/pull/76115

None

>From 32c0cf9fe908343ccdaa0d0814eb62cca326ad99 Mon Sep 17 00:00:00 2001
From: Kevin Athey <kda at google.com>
Date: Wed, 20 Dec 2023 17:35:10 -0800
Subject: [PATCH] [sanitizer] Add graceful handling of exceeding StackStore
 limit.

---
 .../lib/sanitizer_common/sanitizer_stack_store.cpp       | 9 +++++++--
 .../sanitizer_common/tests/sanitizer_stackdepot_test.cpp | 4 ++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp
index 148470943b47b3..c11df0ddfde438 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp
@@ -44,6 +44,9 @@ StackStore::Id StackStore::Store(const StackTrace &trace, uptr *pack) {
   uptr idx = 0;
   *pack = 0;
   uptr *stack_trace = Alloc(h.size + 1, &idx, pack);
+  // No more space.
+  if (stack_trace == nullptr)
+    return 0;
   *stack_trace = h.ToUptr();
   internal_memcpy(stack_trace + 1, trace.trace, h.size * sizeof(uptr));
   *pack += blocks_[GetBlockIdx(idx)].Stored(h.size + 1);
@@ -76,8 +79,10 @@ uptr *StackStore::Alloc(uptr count, uptr *idx, uptr *pack) {
     uptr block_idx = GetBlockIdx(start);
     uptr last_idx = GetBlockIdx(start + count - 1);
     if (LIKELY(block_idx == last_idx)) {
-      // Fits into the a single block.
-      CHECK_LT(block_idx, ARRAY_SIZE(blocks_));
+      // Fits into a single block.
+      // No more available blocks.  Indicate inability to allocate more memory.
+      if (block_idx >= ARRAY_SIZE(blocks_))
+        return nullptr;
       *idx = start;
       return blocks_[block_idx].GetOrCreate(this) + GetInBlockIdx(start);
     }
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cpp
index 3835ce26c4d54b..479e4a0c184f74 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cpp
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_stackdepot_test.cpp
@@ -148,6 +148,10 @@ static struct StackDepotBenchmarkParams {
     {500000, 10, 16, true, false},
     {1500000, 10, 4, true, true},
     {800000, 10, 16, true, true},
+    // Go crazy, and create too many unique stacks, such that StackStore runs
+    // out of space.
+    {1000000, 1, 128, true, true},
+    {100000000, 1, 1, true, true},
 };
 
 static std::string PrintStackDepotBenchmarkParams(