[sanitizer] Switch StackStore to 8 MiB blocks

[vitalybuka]

Summary

Larger blocks are more convenient for compressions.
Blocks are allocated with MmapNoReserveOrDie to save some memory.

Also it's 15% faster on StackDepotBenchmarkSuite

Depends on D114464.

[vitalybuka]

fix comment

[dvyukov]

I think this could be a thin wrapper around TwoLevelMap.

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

57

Does this = {} do anything? It confuses me. It produces impression that it initializes something, but does it?

63

Why do we want it to be 0x100000000ull in total? If this assert fails, do I need to fix the code back, or update the 0x100000000ull number?

[dvyukov]

In D114488#3151405, @dvyukov wrote:

I think this could be a thin wrapper around TwoLevelMap.

Looking at the subsequent commits, it probably can't be TwoLevelMap.

[vitalybuka]

In D114488#3151405, @dvyukov wrote:

I think this could be a thin wrapper around TwoLevelMap.

yep. I need to pack and unmap blocks, which is not supported by TwoLevelMap. If we add that into TwoLevelMap it will be more branching or syncronization of fast path of TwoLevelMap which is used is allocator.

In D114488#3151457, @dvyukov wrote:

In D114488#3151405, @dvyukov wrote:

I think this could be a thin wrapper around TwoLevelMap.

Looking at the subsequent commits, it probably can't be TwoLevelMap.

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

57

It does zero init and let you have constexpr constructor
It would not be needed if StaticSpinMutex had one

63

Right, this assert belongs to later patch which makes using Id = u32

[vitalybuka]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

57

Or you saing that it useless because of BlockInfo blocks_[kBlockCount] = {}; below?
Than it's true

[vitalybuka]

remove assert and {}

[vitalybuka]

rebase

[morehouse]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

91

I think relaxed ordering also works here.

[vitalybuka]

order relaxed

[dvyukov]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

91

Relaxed does not work here. Here is proper implementation of double-checked locking in C++:
https://en.wikipedia.org/wiki/Double-checked_locking#Usage_in_C++11

[vitalybuka]

restore release

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

 //===-- sanitizer_stack_store.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 SANITIZER_STACK_STORE_H
 #define SANITIZER_STACK_STORE_H
 
 #include "sanitizer_atomic.h"
+#include "sanitizer_common.h"
 #include "sanitizer_internal_defs.h"
 #include "sanitizer_mutex.h"
 #include "sanitizer_stacktrace.h"
 
 namespace __sanitizer {
 
 class StackStore {
+  static constexpr uptr kBlockSizeFrames = 0x100000;
+  static constexpr uptr kBlockCount = 0x1000;
+  static constexpr uptr kBlockSizeBytes = kBlockSizeFrames * sizeof(uptr);
+
  public:
   constexpr StackStore() = default;
 
   using Id = uptr;
 
   Id Store(const StackTrace &trace);
   StackTrace Load(Id id) const;
   uptr Allocated() const;
 
   void TestOnlyUnmap();
 
  private:
-  uptr *Alloc(uptr count = 1);
-  uptr *TryAlloc(uptr count);
-  uptr *RefillAndAlloc(uptr count);
-  mutable StaticSpinMutex mtx_ = {};  // Protects alloc of new blocks.
-  atomic_uintptr_t region_pos_ = {};  // Region allocator for Node's.
-  atomic_uintptr_t region_end_ = {};
-  atomic_uintptr_t mapped_size_ = {};
-
-  struct BlockInfo {
-    const BlockInfo *next;
-    uptr ptr;
-    uptr size;
+  static constexpr uptr GetBlockIdx(uptr frame_idx) {
+    return frame_idx / kBlockSizeFrames;
+  }
+
+  static constexpr uptr GetInBlockIdx(uptr frame_idx) {
+    return frame_idx % kBlockSizeFrames;
+  }
+
+  uptr *Alloc(uptr count);
+
+  // Total number of allocated frames.
+  atomic_uintptr_t total_frames_ = {};
+
+  // Each block will hold pointer to exactly kBlockSizeFrames.
+  class BlockInfo {
+    atomic_uintptr_t data_;
+    StaticSpinMutex mtx_;  // Protects alloc of new blocks.
+
+    uptr *Create();
+    uptr *Get() const;
+
+   public:
+    uptr *GetOrCreate();
+    void TestOnlyUnmap();
   };
-  const BlockInfo *curr_ = nullptr;
+  BlockInfo blocks_[kBlockCount] = {};
 };
 
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_STACK_STORE_H

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

StackTrace StackStore::Load(Id id) const {
   if (!id)
     return {};
   const uptr *stack_trace = reinterpret_cast<const uptr *>(id);
   StackTraceHeader h(*stack_trace);
   return StackTrace(stack_trace + 1, h.size, h.tag);
 }
 
 uptr StackStore::Allocated() const {
-  return atomic_load_relaxed(&mapped_size_);
+  return RoundUpTo(atomic_load_relaxed(&total_frames_) * sizeof(uptr),
+                   GetPageSizeCached()) +
+         sizeof(*this);
 }
 
-uptr *StackStore::TryAlloc(uptr count) {
-  // Optimisic lock-free allocation, essentially try to bump the region ptr.
+uptr *StackStore::Alloc(uptr count) {
   for (;;) {
-    uptr cmp = atomic_load(&region_pos_, memory_order_acquire);
-    uptr end = atomic_load(&region_end_, memory_order_acquire);
-    uptr size = count * sizeof(uptr);
-    if (cmp == 0 || cmp + size > end)
-      return nullptr;
-    if (atomic_compare_exchange_weak(&region_pos_, &cmp, cmp + size,
-                                     memory_order_acquire))
-      return reinterpret_cast<uptr *>(cmp);
+    // Optimisic lock-free allocation, essentially try to bump the
+    // total_frames_.
+    uptr start = atomic_fetch_add(&total_frames_, count, memory_order_relaxed);
+    uptr block_idx = GetBlockIdx(start);
+    if (LIKELY(block_idx == GetBlockIdx(start + count - 1))) {
+      // Fits into the a single block.
+      CHECK_LT(block_idx, ARRAY_SIZE(blocks_));
+      return blocks_[block_idx].GetOrCreate() + GetInBlockIdx(start);
+    }
+
+    // Retry. We can't use range allocated in two different blocks.
   }
 }
 
-uptr *StackStore::Alloc(uptr count) {
-  // First, try to allocate optimisitically.
-  uptr *s = TryAlloc(count);
-  if (LIKELY(s))
-    return s;
-  return RefillAndAlloc(count);
+void StackStore::TestOnlyUnmap() {
+  for (BlockInfo &b : blocks_) b.TestOnlyUnmap();
+  internal_memset(this, 0, sizeof(*this));
 }
 
-uptr *StackStore::RefillAndAlloc(uptr count) {
-  // If failed, lock, retry and alloc new superblock.
+uptr *StackStore::BlockInfo::Get() const {
+  return reinterpret_cast<uptr *>(atomic_load(&data_, memory_order_relaxed));
+}
+
+uptr *StackStore::BlockInfo::Create() {
   SpinMutexLock l(&mtx_);
-  for (;;) {
-    uptr *s = TryAlloc(count);
-    if (s)
-      return s;
-    atomic_store(&region_pos_, 0, memory_order_relaxed);
-    uptr size = count * sizeof(uptr) + sizeof(BlockInfo);
-    uptr allocsz = RoundUpTo(Max<uptr>(size, 64u * 1024u), GetPageSizeCached());
-    uptr mem = (uptr)MmapOrDie(allocsz, "stack depot");
-    BlockInfo *new_block = (BlockInfo *)(mem + allocsz) - 1;
-    new_block->next = curr_;
-    new_block->ptr = mem;
-    new_block->size = allocsz;
-    curr_ = new_block;
-
-    atomic_fetch_add(&mapped_size_, allocsz, memory_order_relaxed);
-
-    allocsz -= sizeof(BlockInfo);
-    atomic_store(&region_end_, mem + allocsz, memory_order_release);
-    atomic_store(&region_pos_, mem, memory_order_release);
+  uptr *ptr = Get();
+  if (!ptr) {
+    ptr = reinterpret_cast<uptr *>(
+        MmapNoReserveOrDie(kBlockSizeBytes, "StackStore"));
+    atomic_store(&data_, reinterpret_cast<uptr>(ptr), memory_order_release);
   }
+  return ptr;
 }
 
-void StackStore::TestOnlyUnmap() {
-  while (curr_) {
-    uptr mem = curr_->ptr;
-    uptr allocsz = curr_->size;
-    curr_ = curr_->next;
-    UnmapOrDie((void *)mem, allocsz);
-  }
-  internal_memset(this, 0, sizeof(*this));
+uptr *StackStore::BlockInfo::GetOrCreate() {
+  uptr *ptr = Get();
+  if (LIKELY(ptr))
+    return ptr;
+  return Create();
+
+void StackStore::BlockInfo::TestOnlyUnmap() {
+  if (uptr *ptr = Get())
+    UnmapOrDie(ptr, StackStore::kBlockSizeBytes);
 }
 
 }  // namespace __sanitizer