[NFC][sanitizer] Add entry point for compression

[vitalybuka]

Summary

Add Compression::Test type which just pretends packing,
but does nothing useful. It's only called from test for now.

Depends on D114493.

[vitalybuka]

comment

[vitalybuka]

rebase

[morehouse]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

146

169

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

149

Don't overflow.

152

122 seems arbitrary. Maybe just check that CountPackedBlocks() decrements after loading a block.

[vitalybuka]

update

[vitalybuka]

update

[kstoimenov]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

97–99

Maybe add thread annotations to those functions also?

[vitalybuka]

rebases

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

97–99

now sure which apply here

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

 namespace __sanitizer {
 
 class StackStore {
   static constexpr uptr kBlockSizeFrames = 0x100000;
   static constexpr uptr kBlockCount = 0x1000;
   static constexpr uptr kBlockSizeBytes = kBlockSizeFrames * sizeof(uptr);
 
  public:
+  enum class Compression : u8 {
+    None = 0,
+    Test,
+  };
+
   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,
            uptr *pack /* number of blocks completed by this call */);
-  StackTrace Load(Id id) const;
+  StackTrace Load(Id id);
   uptr Allocated() const;
 
-  void Pack();
+  // Packs all blocks which don't expect any more writes. A block is going to be
+  // packed once. As soon trace from that block was requested, it will unpack
+  // and stay unpacked after that.
+  // Returns the number of released bytes.
+  uptr Pack(Compression type);
 
   void TestOnlyUnmap();
 
  private:
   friend class StackStoreTest;
   static constexpr uptr GetBlockIdx(uptr frame_idx) {
     return frame_idx / kBlockSizeFrames;
   }
private:
   atomic_uintptr_t total_frames_ = {};
 
   // Each block will hold pointer to exactly kBlockSizeFrames.
   class BlockInfo {
     atomic_uintptr_t data_;
     // Counter to track store progress to know when we can Pack() the block.
     atomic_uint32_t stored_;
     // Protects alloc of new blocks.
-    StaticSpinMutex mtx_;
+    mutable StaticSpinMutex mtx_;
+
+    enum class State : u8 {
+      Storing = 0,
+      Packed,
+      Unpacked,
+    };
+    State state GUARDED_BY(mtx_);
 
     uptr *Create();
 
    public:
     uptr *Get() const;
     uptr *GetOrCreate();
+    uptr *GetOrUnpack();
+    uptr Pack(Compression type);
+    uptr Allocated() const;
     void TestOnlyUnmap();
     bool Stored(uptr n);
+    bool IsPacked() const;
   };
+
   BlockInfo blocks_[kBlockCount] = {};
 };
 
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_STACK_STORE_H

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

StackStore::Id StackStore::Store(const StackTrace &trace, uptr *pack) {
   *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 {
+StackTrace StackStore::Load(Id id) {
   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();
+  const uptr *stack_trace = blocks_[block_idx].GetOrUnpack();
   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 next_block = GetBlockIdx(
+      RoundUpTo(atomic_load_relaxed(&total_frames_), kBlockSizeFrames));
+  uptr res = 0;
+  for (uptr i = 0; i < next_block; ++i) res += blocks_[i].Allocated();
+  return res + sizeof(*this);
 }
 
 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);
for (;;) {
     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
+uptr StackStore::Pack(Compression type) {
+  uptr res = 0;
+  for (BlockInfo &b : blocks_) res += b.Pack(type);
+  return res;
 }
 
 void StackStore::TestOnlyUnmap() {
   for (BlockInfo &b : blocks_) b.TestOnlyUnmap();
   internal_memset(this, 0, sizeof(*this));
 }
 
 uptr *StackStore::BlockInfo::Get() const {

 
 uptr *StackStore::BlockInfo::GetOrCreate() {
   uptr *ptr = Get();
   if (LIKELY(ptr))
     return ptr;
   return Create();
 }
 
+uptr *StackStore::BlockInfo::GetOrUnpack() {
+  SpinMutexLock l(&mtx_);
+  switch (state) {
+    case State::Storing:
+      state = State::Unpacked;
+      FALLTHROUGH;
+    case State::Unpacked:
+      return Get();
+    case State::Packed:
+      break;
+  }
+
+  uptr *ptr = Get();
+  CHECK_NE(nullptr, ptr);
+  // Fake unpacking.
+  for (uptr i = 0; i < kBlockSizeFrames; ++i) ptr[i] = ~ptr[i];
+  state = State::Unpacked;
+  return Get();
+}
+
+uptr StackStore::BlockInfo::Pack(Compression type) {
+  if (type == Compression::None)
+    return 0;
+
+  SpinMutexLock l(&mtx_);
+  switch (state) {
+    case State::Unpacked:
+    case State::Packed:
+      return 0;
+    case State::Storing:
+      break;
+  }
+
+  uptr *ptr = Get();
+  if (!ptr || !Stored(0))
+    return 0;
+
+  // Fake packing.
+  for (uptr i = 0; i < kBlockSizeFrames; ++i) ptr[i] = ~ptr[i];
+  state = State::Packed;
+  return kBlockSizeBytes - kBlockSizeBytes / 10;
+}
+
+uptr StackStore::BlockInfo::Allocated() const {
+  SpinMutexLock l(&mtx_);
+  switch (state) {
+    case State::Packed:
+      return kBlockSizeBytes / 10;
+    case State::Unpacked:
+    case State::Storing:
+      return kBlockSizeBytes;
+  }
+}
+
 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;
 }
 
+bool StackStore::BlockInfo::IsPacked() const {
+  SpinMutexLock l(&mtx_);
+  return state == State::Packed;
+}
+
 }  // 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;
   uptr pack = 0;
   store_id = stackStore.Store(args, &pack);
+  if (pack)
+    stackStore.Pack(StackStore::Compression::None);
 }
 
 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

protected:
   }
 
   uptr CountReadyToPackBlocks() {
     uptr res = 0;
     for (BlockInfo& b : store_.blocks_) res += b.Stored(0);
     return res;
   }
 
+  uptr CountPackedBlocks() const {
+    uptr res = 0;
+    for (const BlockInfo& b : store_.blocks_) res += b.IsPacked();
+    return res;
+  }
+
   uptr IdToOffset(StackStore::Id id) const { return store_.IdToOffset(id); }
 
   static constexpr uptr kBlockSizeFrames = StackStore::kBlockSizeFrames;
+  static constexpr uptr kBlockSizeBytes = StackStore::kBlockSizeBytes;
 
   StackStore store_ = {};
 };
 
 TEST_F(StackStoreTest, Empty) {
   uptr before = store_.Allocated();
   uptr pack = 0;
   EXPECT_EQ(0u, store_.Store({}, &pack));
ForEachTrace(
         }
         total_ready += pack;
         EXPECT_EQ(CountReadyToPackBlocks(), total_ready);
       },
       100000);
   EXPECT_EQ(GetTotalFramesCount() / kBlockSizeFrames, total_ready);
 }
 
+struct StackStorePackTest : public StackStoreTest,
+                            public ::testing::WithParamInterface<
+                                std::pair<StackStore::Compression, uptr>> {};
+
+INSTANTIATE_TEST_SUITE_P(
+    PackUnpacks, StackStorePackTest,
+    ::testing::ValuesIn({
+        StackStorePackTest::ParamType(StackStore::Compression::Test, 4),
+    }));
+
+TEST_P(StackStorePackTest, PackUnpack) {
+  std::vector<StackStore::Id> ids;
+  StackStore::Compression type = GetParam().first;
+  uptr expected_ratio = GetParam().second;
+  ForEachTrace([&](const StackTrace& s) {
+    uptr pack = 0;
+    ids.push_back(store_.Store(s, &pack));
+    if (pack) {
+      uptr before = store_.Allocated();
+      uptr diff = store_.Pack(type);
+      uptr after = store_.Allocated();
+      EXPECT_EQ(before - after, diff);
+      EXPECT_LT(after, before);
+      EXPECT_GE(kBlockSizeBytes / (kBlockSizeBytes - (before - after)),
+                expected_ratio);
+    }
+  });
+  uptr packed_blocks = CountPackedBlocks();
+  // Unpack random block.
+  store_.Load(kBlockSizeFrames * 7 + 123);
+  EXPECT_EQ(packed_blocks - 1, CountPackedBlocks());
+
+  // Unpack all blocks.
+  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));
+  });
+  EXPECT_EQ(0u, CountPackedBlocks());
+
+  EXPECT_EQ(0u, store_.Pack(type));
+  EXPECT_EQ(0u, CountPackedBlocks());
+}
+
 }  // namespace __sanitizer