[sanitizer] Minor 32-bit primary improvements

[cryptoad]

Summary

For the 32-bit TransferBatch:

• SetFromArray callers have bounds count, so relax the CHECK to DCHECK;
• same for Add;
• mark CopyToArray as const;

For the 32-bit Primary:

• {Dea,A}llocateBatch are only called from places that check class_id, relax the CHECK to DCHECK;
• same for AllocateRegion;
• remove GetRegionBeginBySizeClass that is not used;
• use a local variable for the random shuffle state, so that the compiler can use a register instead of reading and writing to the SizeClassInfo at every iteration;

For the 32-bit local cache:

• pass the count to drain instead of doing a Min everytime which is at times superfluous.

[cryptoad]

And I forgot:

• the SizeClassInfo mutex should be a StaticSpinMutex since it's zero initialized and doesn't need to be "constructed";

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

183

Drain(c, allocator, class_id, c->count / 2); is the same, but reads a bit easier, more logical ("drain half of the elements").

192

Why the loop, why don't we drain the class in one call?

Drain(c, allocator, i, c->count);

We used to have to have a loop because of the Drain API, but now we can specify the number of elements to drain.

lib/sanitizer_common/sanitizer_allocator_primary32.h

305

Add a comment why you need a local variable here, otherwise someone will "optimize" it back later or, better yet, add a local variable in RandomShuffle and let the compiler do its job generating the best code possible.

[cryptoad]

Addressing a couple of comments from Aleksey.

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

192

Since the TransferBatch that we are using can only hold max_count / 2, we can't directly drain count.
I could change it to something like if count > max_count / 2 then we can drain max_count / 2, then drain count elements.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

192

Ah, I see.

This restriction was not obvious before and it is even more obfuscated now. Was Min() call in Drain so taxing that it is worth complicating the already not so easy to follow logic of the code? If it is not, I'd rather not change the code. If it is though, well, do it, but add a DCHECK to Drain to document the restriction of this API: DCHECK_LE(count, c->max_count / 2); and ignore my previous comment, Drain(c, allocator, class_id, c->max_count / 2);

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_local_cache.h

192

Yeah you are right it's not worth it. I'll revert this part to its original state.

[cryptoad]

The changes to the local cache were not worth doing. Reverting those.

lib/sanitizer_common/sanitizer_allocator_local_cache.h

struct SizeClassAllocator32LocalCache {
   void Deallocate(SizeClassAllocator *allocator, uptr class_id, void *p) {
     CHECK_NE(class_id, 0UL);
     CHECK_LT(class_id, kNumClasses);
     // If the first allocator call on a new thread is a deallocation, then
     // max_count will be zero, leading to check failure.
     PerClass *c = &per_class_[class_id];
     InitCache(c);
     if (UNLIKELY(c->count == c->max_count))
-      Drain(c, allocator, class_id);
+      Drain(c, allocator, class_id, c->max_count / 2);
     c->batch[c->count++] = p;
     stats_.Sub(AllocatorStatAllocated, c->class_size);
   }
 
   void Drain(SizeClassAllocator *allocator) {
     for (uptr i = 1; i < kNumClasses; i++) {
       PerClass *c = &per_class_[i];
       while (c->count > 0)
-        Drain(c, allocator, i);
+        Drain(c, allocator, i, Min(c->max_count / 2, c->count));
     }
   }
 
  private:
   typedef typename Allocator::SizeClassMapT SizeClassMap;
   static const uptr kBatchClassID = SizeClassMap::kBatchClassID;
   static const uptr kNumClasses = SizeClassMap::kNumClasses;
   // If kUseSeparateSizeClassForBatch is true, all TransferBatch objects are
NOINLINE bool Refill(PerClass *c, SizeClassAllocator *allocator,
     CHECK_GT(b->Count(), 0);
     b->CopyToArray(c->batch);
     c->count = b->Count();
     DestroyBatch(class_id, allocator, b);
     return true;
   }
 
   NOINLINE void Drain(PerClass *c, SizeClassAllocator *allocator,
-                      uptr class_id) {
-    const uptr count = Min(c->max_count / 2, c->count);
+                      uptr class_id, uptr count) {
     const uptr first_idx_to_drain = c->count - count;
     TransferBatch *b = CreateBatch(
         class_id, allocator, (TransferBatch *)c->batch[first_idx_to_drain]);
     // Failure to allocate a batch while releasing memory is non recoverable.
     // TODO(alekseys): Figure out how to do it without allocating a new batch.
     if (UNLIKELY(!b))
       DieOnFailure::OnOOM();
     b->SetFromArray(&c->batch[first_idx_to_drain], count);
     c->count -= count;
     allocator->DeallocateBatch(&stats_, class_id, b);
   }
 };

lib/sanitizer_common/sanitizer_allocator_primary32.h

public:
   static const bool kRandomShuffleChunks = Params::kFlags &
       SizeClassAllocator32FlagMasks::kRandomShuffleChunks;
   static const bool kUseSeparateSizeClassForBatch = Params::kFlags &
       SizeClassAllocator32FlagMasks::kUseSeparateSizeClassForBatch;
 
   struct TransferBatch {
     static const uptr kMaxNumCached = SizeClassMap::kMaxNumCachedHint - 2;
     void SetFromArray(void *batch[], uptr count) {
+      DCHECK_LE(count, kMaxNumCached);
       count_ = count;
-      CHECK_LE(count_, kMaxNumCached);
       for (uptr i = 0; i < count; i++)
         batch_[i] = batch[i];
     }
     uptr Count() const { return count_; }
     void Clear() { count_ = 0; }
     void Add(void *ptr) {
       batch_[count_++] = ptr;
-      CHECK_LE(count_, kMaxNumCached);
+      DCHECK_LE(count_, kMaxNumCached);
     }
-    void CopyToArray(void *to_batch[]) {
+    void CopyToArray(void *to_batch[]) const {
       for (uptr i = 0, n = Count(); i < n; i++)
         to_batch[i] = batch_[i];
     }
 
     // How much memory do we need for a batch containing n elements.
     static uptr AllocationSizeRequiredForNElements(uptr n) {
       return sizeof(uptr) * 2 + sizeof(void *) * n;
     }
void *GetMetaData(const void *p) {
     u32 offset = mem - beg;
     uptr n = offset / (u32)size;  // 32-bit division
     uptr meta = (beg + kRegionSize) - (n + 1) * kMetadataSize;
     return reinterpret_cast<void*>(meta);
   }
 
   NOINLINE TransferBatch *AllocateBatch(AllocatorStats *stat, AllocatorCache *c,
                                         uptr class_id) {
-    CHECK_LT(class_id, kNumClasses);
+    DCHECK_LT(class_id, kNumClasses);
     SizeClassInfo *sci = GetSizeClassInfo(class_id);
     SpinMutexLock l(&sci->mutex);
     if (sci->free_list.empty()) {
       if (UNLIKELY(!PopulateFreeList(stat, c, sci, class_id)))
         return nullptr;
       DCHECK(!sci->free_list.empty());
     }
     TransferBatch *b = sci->free_list.front();
     sci->free_list.pop_front();
     return b;
   }
 
   NOINLINE void DeallocateBatch(AllocatorStats *stat, uptr class_id,
                                 TransferBatch *b) {
-    CHECK_LT(class_id, kNumClasses);
+    DCHECK_LT(class_id, kNumClasses);
     CHECK_GT(b->Count(), 0);
     SizeClassInfo *sci = GetSizeClassInfo(class_id);
     SpinMutexLock l(&sci->mutex);
     sci->free_list.push_front(b);
   }
 
-  uptr GetRegionBeginBySizeClass(uptr class_id) { return 0; }
-
   bool PointerIsMine(const void *p) {
     uptr mem = reinterpret_cast<uptr>(p);
     if (mem < kSpaceBeg || mem >= kSpaceBeg + kSpaceSize)
       return false;
     return GetSizeClass(p) != 0;
   }
 
   uptr GetSizeClass(const void *p) {
for (uptr region = 0; region < kNumPossibleRegions; region++)
              chunk < region_beg + max_chunks_in_region * chunk_size;
              chunk += chunk_size) {
           // Too slow: CHECK_EQ((void *)chunk, GetBlockBegin((void *)chunk));
           callback(chunk, arg);
         }
       }
   }
 
-  void PrintStats() {
-  }
+  void PrintStats() {}
 
-  static uptr AdditionalSize() {
-    return 0;
-  }
+  static uptr AdditionalSize() { return 0; }
 
   typedef SizeClassMap SizeClassMapT;
   static const uptr kNumClasses = SizeClassMap::kNumClasses;
 
  private:
   static const uptr kRegionSize = 1 << kRegionSizeLog;
   static const uptr kNumPossibleRegions = kSpaceSize / kRegionSize;
 
   struct ALIGNED(SANITIZER_CACHE_LINE_SIZE) SizeClassInfo {
-    SpinMutex mutex;
+    StaticSpinMutex mutex;
     IntrusiveList<TransferBatch> free_list;
     u32 rand_state;
   };
   COMPILER_CHECK(sizeof(SizeClassInfo) % kCacheLineSize == 0);
 
   uptr ComputeRegionId(uptr mem) {
     const uptr res = mem >> kRegionSizeLog;
     CHECK_LT(res, kNumPossibleRegions);
     return res;
   }
 
   uptr ComputeRegionBeg(uptr mem) {
     return mem & ~(kRegionSize - 1);
   }
 
   uptr AllocateRegion(AllocatorStats *stat, uptr class_id) {
-    CHECK_LT(class_id, kNumClasses);
-    uptr res = reinterpret_cast<uptr>(MmapAlignedOrDieOnFatalError(
+    DCHECK_LT(class_id, kNumClasses);
+    const uptr res = reinterpret_cast<uptr>(MmapAlignedOrDieOnFatalError(
         kRegionSize, kRegionSize, PrimaryAllocatorName));
     if (UNLIKELY(!res))
       return 0;
     MapUnmapCallback().OnMap(res, kRegionSize);
     stat->Add(AllocatorStatMapped, kRegionSize);
     CHECK(IsAligned(res, kRegionSize));
     possible_regions.set(ComputeRegionId(res), static_cast<u8>(class_id));
     return res;
   }
 
   SizeClassInfo *GetSizeClassInfo(uptr class_id) {
     DCHECK_LT(class_id, kNumClasses);
     return &size_class_info_array[class_id];
   }
 
   bool PopulateBatches(AllocatorCache *c, SizeClassInfo *sci, uptr class_id,
                        TransferBatch **current_batch, uptr max_count,
                        uptr *pointers_array, uptr count) {
     // If using a separate class for batches, we do not need to shuffle it.
     if (kRandomShuffleChunks && (!kUseSeparateSizeClassForBatch ||
-        class_id != SizeClassMap::kBatchClassID))
-      RandomShuffle(pointers_array, count, &sci->rand_state);
+        class_id != SizeClassMap::kBatchClassID)) {
+      u32 rand_state = sci->rand_state;
+      RandomShuffle(pointers_array, count, &rand_state);
+      sci->rand_state = rand_state;
+    }
     TransferBatch *b = *current_batch;
     for (uptr i = 0; i < count; i++) {
       if (!b) {
         b = c->CreateBatch(class_id, this, (TransferBatch*)pointers_array[i]);
         if (UNLIKELY(!b))
           return false;
         b->Clear();
       }