[Asan] Cleanup kAllocBegMagic setup

[vitalybuka]

Summary

Make it atomic.
Wrap it into class.
Set it late after chunk is initialized.
Reset it soon when the chunk is still valid.

Depends on D87645.

[morehouse]

compiler-rt/lib/asan/asan_allocator.cpp

166

I'm not sure which style guide we follow, but Google style is to have constants before other members.

[vitalybuka]

reorder members

compiler-rt/lib/asan/asan_allocator.cpp

 // If the left redzone is greater than the ChunkHeader size we store a magic
 // value in the first uptr word of the memory block and store the address of
 // ChunkBase in the next uptr.
 // M B L L L L L L L L L  H H U U U U U U
 //   |                    ^
 //   ---------------------|
 //   M -- magic value kAllocBegMagic
 //   B -- address of ChunkHeader pointing to the first 'H'
-static const uptr kAllocBegMagic = 0xCC6E96B9;
 
 class ChunkHeader {
  public:
   atomic_uint8_t chunk_state;
   u8 alloc_type : 2;
   u8 lsan_tag : 2;
 
   // align < 8 -> 0
enum {
   CHUNK_QUARANTINE = 3,
 };
 
 class AsanChunk : public ChunkBase {
  public:
   uptr Beg() { return reinterpret_cast<uptr>(this) + kChunkHeaderSize; }
 };
 
+class LargeChunkHeader {
+  atomic_uint64_t magic;
+  AsanChunk *chunk_header;
+  static constexpr uptr kAllocBegMagic = 0xCC6E96B9;
+
+ public:
+  AsanChunk *Get() {
+    return atomic_load(&magic, memory_order_acquire) == kAllocBegMagic
+               ? chunk_header
+               : reinterpret_cast<AsanChunk *>(this);
+  }
+
+  void Set(AsanChunk *p) {
+    if (p) {
+      chunk_header = p;
+      atomic_store(&magic, kAllocBegMagic, memory_order_release);
+      return;
+    }
+
+    u64 old = kAllocBegMagic;
+    if (!atomic_compare_exchange_strong(&magic, &old, 0,
+                                        memory_order_release)) {
+      CHECK_EQ(old, kAllocBegMagic);
+    }
+  }
+};
+
 struct QuarantineCallback {
   QuarantineCallback(AllocatorCache *cache, BufferedStackTrace *stack)
       : cache_(cache),
         stack_(stack) {
   }
 
   void Recycle(AsanChunk *m) {
+    void *p = get_allocator().GetBlockBegin(m);
+    if (p != m) {
+      // Clear the magic value, as allocator internals may overwrite the
+      // contents of deallocated chunk, confusing GetAsanChunk lookup.
+      reinterpret_cast<LargeChunkHeader *>(p)->Set(nullptr);
+    }
+
     u8 old_chunk_state = CHUNK_QUARANTINE;
     if (!atomic_compare_exchange_strong(&m->chunk_state, &old_chunk_state,
                                         CHUNK_INVALID, memory_order_acquire)) {
       CHECK_EQ(old_chunk_state, CHUNK_QUARANTINE);
     }
 
     PoisonShadow(m->Beg(),
                  RoundUpTo(m->UsedSize(), SHADOW_GRANULARITY),
                  kAsanHeapLeftRedzoneMagic);
-    void *p = get_allocator().GetBlockBegin(m);
-    if (p != m) {
-      uptr *alloc_magic = reinterpret_cast<uptr *>(p);
-      CHECK_EQ(alloc_magic[0], kAllocBegMagic);
-      // Clear the magic value, as allocator internals may overwrite the
-      // contents of deallocated chunk, confusing GetAsanChunk lookup.
-      alloc_magic[0] = 0;
-      CHECK_EQ(alloc_magic[1], reinterpret_cast<uptr>(m));
-    }
 
     // Statistics.
     AsanStats &thread_stats = GetCurrentThreadStats();
     thread_stats.real_frees++;
     thread_stats.really_freed += m->UsedSize();
 
     get_allocator().Deallocate(cache_, p);
   }
void *Allocate(uptr size, uptr alignment, BufferedStackTrace *stack,
     uptr user_beg = alloc_beg + rz_size;
     if (!IsAligned(user_beg, alignment))
       user_beg = RoundUpTo(user_beg, alignment);
     uptr user_end = user_beg + size;
     CHECK_LE(user_end, alloc_end);
     uptr chunk_beg = user_beg - kChunkHeaderSize;
     AsanChunk *m = reinterpret_cast<AsanChunk *>(chunk_beg);
     m->alloc_type = alloc_type;
-    if (alloc_beg != chunk_beg) {
-      CHECK_LE(alloc_beg + 2 * sizeof(uptr), chunk_beg);
-      reinterpret_cast<uptr *>(alloc_beg)[0] = kAllocBegMagic;
-      reinterpret_cast<uptr *>(alloc_beg)[1] = chunk_beg;
-    }
     CHECK(size);
     m->SetUsedSize(size);
     if (using_primary_allocator) {
       CHECK(allocator.FromPrimary(allocated));
     } else {
       CHECK(!allocator.FromPrimary(allocated));
       uptr *meta = reinterpret_cast<uptr *>(allocator.GetMetaData(allocated));
       meta[1] = chunk_beg;
if (can_fill && fl.max_malloc_fill_size) {
       REAL(memset)(res, fl.malloc_fill_byte, fill_size);
     }
 #if CAN_SANITIZE_LEAKS
     m->lsan_tag = __lsan::DisabledInThisThread() ? __lsan::kIgnored
                                                  : __lsan::kDirectlyLeaked;
 #endif
     // Must be the last mutation of metadata in this function.
     atomic_store(&m->chunk_state, CHUNK_ALLOCATED, memory_order_release);
+    if (alloc_beg != chunk_beg) {
+      CHECK_LE(alloc_beg + sizeof(LargeChunkHeader), chunk_beg);
+      reinterpret_cast<LargeChunkHeader *>(alloc_beg)->Set(m);
+    }
     ASAN_MALLOC_HOOK(res, size);
     return res;
   }
 
   // Set quarantine flag if chunk is allocated, issue ASan error report on
   // available and quarantined chunks. Return true on success, false otherwise.
   bool AtomicallySetQuarantineFlagIfAllocated(AsanChunk *m, void *ptr,
                                               BufferedStackTrace *stack) {
#endif
   AsanChunk *GetAsanChunk(void *alloc_beg) {
     if (!alloc_beg)
       return nullptr;
     AsanChunk *p = nullptr;
     if (!allocator.FromPrimary(alloc_beg)) {
       uptr *meta = reinterpret_cast<uptr *>(allocator.GetMetaData(alloc_beg));
       p = reinterpret_cast<AsanChunk *>(meta[1]);
     } else {
-      uptr *alloc_magic = reinterpret_cast<uptr *>(alloc_beg);
-      if (alloc_magic[0] == kAllocBegMagic)
-        p = reinterpret_cast<AsanChunk *>(alloc_magic[1]);
-      else
-        p = reinterpret_cast<AsanChunk *>(alloc_beg);
+      p = reinterpret_cast<LargeChunkHeader *>(alloc_beg)->Get();
     }
     if (!p)
       return nullptr;
     u8 state = atomic_load(&p->chunk_state, memory_order_relaxed);
     // It does not guaranty that Chunk is initialized, but it's
     // definitely not for any other value.
     if (state == CHUNK_ALLOCATED || state == CHUNK_QUARANTINE)
       return p;