[NFC][sanitizer] Track progress of populating the block

[vitalybuka]

Summary

In multi-threaded application concurrent StackStore::Store may
finish in order different from assigned Id. So we can't assume
that after we switch writing the next block the previous is done.

The workaround is to count exact number of uptr stored into the block,
including skipped tail/head which were not able to fit entire trace.

Depends on D114490.

[dvyukov]

compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp

77–78

= 0

[vitalybuka]

0

[vitalybuka]

rebase

[morehouse]

compiler-rt/lib/sanitizer_common/tests/sanitizer_stack_store_test.cpp

119

12 seems a bit arbitrary. Maybe we can compute the combined sizes of all stack traces stored and divide by kBlockSizeFrames.

[vitalybuka]

calculated total_ready

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

 
  public:
   constexpr StackStore() = default;
 
   using Id = u32;  // Enough for 2^32 * sizeof(uptr) bytes of traces.
   static_assert(u64(kBlockCount) * kBlockSizeFrames == 1ull << (sizeof(Id) * 8),
                 "");
 
-  Id Store(const StackTrace &trace);
+  Id Store(const StackTrace &trace,
+           uptr *pack /* number of blocks completed by this call */);
   StackTrace Load(Id id) const;
   uptr Allocated() const;
 
+  void Pack();
+
   void TestOnlyUnmap();
 
  private:
+  friend class StackStoreTest;
   static constexpr uptr GetBlockIdx(uptr frame_idx) {
     return frame_idx / kBlockSizeFrames;
   }
 
   static constexpr uptr GetInBlockIdx(uptr frame_idx) {
     return frame_idx % kBlockSizeFrames;
   }
 
   static constexpr uptr IdToOffset(Id id) {
     CHECK_NE(id, 0);
     return id - 1;  // Avoid zero as id.
   }
 
   static constexpr uptr OffsetToId(Id id) {
     // This makes UINT32_MAX to 0 and it will be retrived as and empty stack.
     // But this is not a problem as we will not be able to store anything after
     // that anyway.
     return id + 1;  // Avoid zero as id.
   }
 
-  uptr *Alloc(uptr count, uptr *idx);
+  uptr *Alloc(uptr count, uptr *idx, uptr *pack);
 
   // 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.
+    // Counter to track store progress to know when we can Pack() the block.
+    atomic_uint32_t stored_;
+    // Protects alloc of new blocks.
+    StaticSpinMutex mtx_;
 
     uptr *Create();
 
    public:
     uptr *Get() const;
     uptr *GetOrCreate();
     void TestOnlyUnmap();
+    bool Stored(uptr n);
   };
   BlockInfo blocks_[kBlockCount] = {};
 };
 
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_STACK_STORE_H

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

explicit StackTraceHeader(uptr h)
       : size(h & ((1 << kStackSizeBits) - 1)), tag(h >> kStackSizeBits) {}
 
   uptr ToUptr() const {
     return static_cast<uptr>(size) | (static_cast<uptr>(tag) << kStackSizeBits);
   }
 };
 }  // namespace
 
-StackStore::Id StackStore::Store(const StackTrace &trace) {
+StackStore::Id StackStore::Store(const StackTrace &trace, uptr *pack) {
   if (!trace.size && !trace.tag)
     return 0;
   StackTraceHeader h(trace);
-  uptr idx;
-  uptr *stack_trace = Alloc(h.size + 1, &idx);
+  uptr idx = 0;
+  *pack = 0;
+  uptr *stack_trace = Alloc(h.size + 1, &idx, pack);
   *stack_trace = h.ToUptr();
   internal_memcpy(stack_trace + 1, trace.trace, h.size * sizeof(uptr));
+  *pack += blocks_[GetBlockIdx(idx)].Stored(h.size + 1);
   return OffsetToId(idx);
 }
 
 StackTrace StackStore::Load(Id id) const {
   if (!id)
     return {};
   uptr idx = IdToOffset(id);
   uptr block_idx = GetBlockIdx(idx);
   CHECK_LT(block_idx, ARRAY_SIZE(blocks_));
   const uptr *stack_trace = blocks_[block_idx].Get();
   if (!stack_trace)
     return {};
   stack_trace += GetInBlockIdx(idx);
   StackTraceHeader h(*stack_trace);
   return StackTrace(stack_trace + 1, h.size, h.tag);
 }
 
 uptr StackStore::Allocated() const {
   return RoundUpTo(atomic_load_relaxed(&total_frames_) * sizeof(uptr),
                    GetPageSizeCached()) +
          sizeof(*this);
 }
 
-uptr *StackStore::Alloc(uptr count, uptr *idx) {
+uptr *StackStore::Alloc(uptr count, uptr *idx, uptr *pack) {
   for (;;) {
     // 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))) {
+    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_));
       *idx = start;
       return blocks_[block_idx].GetOrCreate() + GetInBlockIdx(start);
     }
 
     // Retry. We can't use range allocated in two different blocks.
+    CHECK_LE(count, kBlockSizeFrames);
+    uptr in_first = kBlockSizeFrames - GetInBlockIdx(start);
+    // Mark tail/head of these blocks as "stored".to avoid waiting before we can
+    // Pack().
+    *pack += blocks_[block_idx].Stored(in_first);
+    *pack += blocks_[last_idx].Stored(count - in_first);
   }
 }
 
+void StackStore::Pack() {
+  // TODO
+}
+
 void StackStore::TestOnlyUnmap() {
   for (BlockInfo &b : blocks_) b.TestOnlyUnmap();
   internal_memset(this, 0, sizeof(*this));
 }
 
 uptr *StackStore::BlockInfo::Get() const {
   // Idiomatic double-checked locking uses memory_order_acquire here. But
   // relaxed is find for us, justification is similar to
uptr *StackStore::BlockInfo::GetOrCreate() {
   return Create();
 }
 
 void StackStore::BlockInfo::TestOnlyUnmap() {
   if (uptr *ptr = Get())
     UnmapOrDie(ptr, StackStore::kBlockSizeBytes);
 }
 
+bool StackStore::BlockInfo::Stored(uptr n) {
+  return n + atomic_fetch_add(&stored_, n, memory_order_release) ==
+         kBlockSizeFrames;
+}
+
 }  // namespace __sanitizer

compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp

 }
 
 uptr StackDepotNode::allocated() {
   return stackStore.Allocated() + useCounts.MemoryUsage();
 }
 
 void StackDepotNode::store(u32 id, const args_type &args, hash_type hash) {
   stack_hash = hash;
-  store_id = stackStore.Store(args);
+  uptr pack = 0;
+  store_id = stackStore.Store(args, &pack);
 }
 
 StackDepotNode::args_type StackDepotNode::load(u32 id) const {
   if (!store_id)
     return {};
   return stackStore.Load(store_id);
 }
 

compiler-rt/lib/sanitizer_common/tests/sanitizer_stack_store_test.cpp

 //===-- sanitizer_stack_store_test.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 "sanitizer_common/sanitizer_stack_store.h"
 
 #include <algorithm>
 #include <numeric>
 #include <vector>
 
 #include "gtest/gtest.h"
+#include "sanitizer_atomic.h"
 #include "sanitizer_hash.h"
 #include "sanitizer_stacktrace.h"
 
 namespace __sanitizer {
 
 class StackStoreTest : public testing::Test {
  protected:
   void SetUp() override {}
for (uptr i = 0; i < n; ++i) {
       StackTrace s(frames.data(), size, tag);
       if (!s.size && !s.tag)
         continue;
       fn(s);
       std::next_permutation(frames.begin(), frames.end());
     };
   }
 
+  using BlockInfo = StackStore::BlockInfo;
+
+  uptr GetTotalFramesCount() const {
+    return atomic_load_relaxed(&store_.total_frames_);
+  }
+
+  uptr CountReadyToPackBlocks() {
+    uptr res = 0;
+    for (BlockInfo& b : store_.blocks_) res += b.Stored(0);
+    return res;
+  }
+
+  uptr IdToOffset(StackStore::Id id) const { return store_.IdToOffset(id); }
+
+  static constexpr uptr kBlockSizeFrames = StackStore::kBlockSizeFrames;
+
   StackStore store_ = {};
 };
 
 TEST_F(StackStoreTest, Empty) {
   uptr before = store_.Allocated();
-  EXPECT_EQ(0u, store_.Store({}));
+  uptr pack = 0;
+  EXPECT_EQ(0u, store_.Store({}, &pack));
   uptr after = store_.Allocated();
   EXPECT_EQ(before, after);
 }
 
 TEST_F(StackStoreTest, Basic) {
   std::vector<StackStore::Id> ids;
-  ForEachTrace([&](const StackTrace& s) { ids.push_back(store_.Store(s)); });
+  ForEachTrace([&](const StackTrace& s) {
+    uptr pack = 0;
+    ids.push_back(store_.Store(s, &pack));
+  });
 
   auto id = ids.begin();
   ForEachTrace([&](const StackTrace& s) {
     StackTrace trace = store_.Load(*(id++));
     EXPECT_EQ(s.size, trace.size);
     EXPECT_EQ(s.tag, trace.tag);
     EXPECT_EQ(std::vector<uptr>(s.trace, s.trace + s.size),
               std::vector<uptr>(trace.trace, trace.trace + trace.size));
   });
 }
 
 TEST_F(StackStoreTest, Allocated) {
   EXPECT_LE(store_.Allocated(), 0x100000u);
   std::vector<StackStore::Id> ids;
-  ForEachTrace([&](const StackTrace& s) { ids.push_back(store_.Store(s)); });
+  ForEachTrace([&](const StackTrace& s) {
+    uptr pack = 0;
+    ids.push_back(store_.Store(s, &pack));
+  });
   EXPECT_NEAR(store_.Allocated(), FIRST_32_SECOND_64(500000000u, 1000000000u),
               FIRST_32_SECOND_64(50000000u, 100000000u));
   store_.TestOnlyUnmap();
   EXPECT_LE(store_.Allocated(), 0x100000u);
 }
 
+TEST_F(StackStoreTest, ReadyToPack) {
+  uptr next_pack = kBlockSizeFrames;
+  uptr total_ready = 0;
+  ForEachTrace(
+      [&](const StackTrace& s) {
+        uptr pack = 0;
+        StackStore::Id id = store_.Store(s, &pack);
+        uptr end_idx = IdToOffset(id) + 1 + s.size;
+        if (end_idx >= next_pack) {
+          EXPECT_EQ(1u, pack);
+          next_pack += kBlockSizeFrames;
+        } else {
+          EXPECT_EQ(0u, pack);
+        }
+        total_ready += pack;
+        EXPECT_EQ(CountReadyToPackBlocks(), total_ready);
+      },
+      100000);
+  EXPECT_EQ(GetTotalFramesCount() / kBlockSizeFrames, total_ready);
+}
+
 }  // namespace __sanitizer