[sanitizer] Switch StackStore from pointers to 32bit IDs

[vitalybuka]

Summary

Depends on D114488.

[kstoimenov]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

48

Can we add a check that id > 0? Also I am usually nervous about using '-' for unsigned type.

[kstoimenov]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

72

I noticed that you have to cast to u8 * in D114503. Does it have to return uptr *?

[vitalybuka]

CHECK

[vitalybuka]

rebase

[vitalybuka]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

72

Than I need cast to uptr* from u8*, which is worse as alignment will be less clear
GetAs<uptr>() and GetAs<u8>() ?

[morehouse]

LGTM

[kstoimenov]

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

72

Your call. I just notice that you have to cast and wanted to point it out.

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.h

 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;
+  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);
   StackTrace Load(Id id) const;
   uptr Allocated() const;
 
   void TestOnlyUnmap();
 
  private:
   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);
+  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);
 
   // 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 *Get() const;
     uptr *GetOrCreate();
     void TestOnlyUnmap();
   };
   BlockInfo blocks_[kBlockCount] = {};
 };
 
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_STACK_STORE_H

compiler-rt/lib/sanitizer_common/sanitizer_stack_store.cpp

struct StackTraceHeader {
   }
 };
 }  // namespace
 
 StackStore::Id StackStore::Store(const StackTrace &trace) {
   if (!trace.size && !trace.tag)
     return 0;
   StackTraceHeader h(trace);
-  uptr *stack_trace = Alloc(h.size + 1);
+  uptr idx;
+  uptr *stack_trace = Alloc(h.size + 1, &idx);
   *stack_trace = h.ToUptr();
   internal_memcpy(stack_trace + 1, trace.trace, h.size * sizeof(uptr));
-  return reinterpret_cast<StackStore::Id>(stack_trace);
+  return OffsetToId(idx);
 }
 
 StackTrace StackStore::Load(Id id) const {
   if (!id)
     return {};
-  const uptr *stack_trace = reinterpret_cast<const uptr *>(id);
+  uptr idx = IdToOffset(id);
+  uptr block_idx = GetBlockIdx(idx);
+  CHECK_LT(block_idx, ARRAY_SIZE(blocks_));
+  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 *StackStore::Alloc(uptr count, uptr *idx) {
   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))) {
       // 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.
   }
 }
 
 void StackStore::TestOnlyUnmap() {