[sanitizer] Random shuffling of chunks for the 32-bit Primary Allocator

Summary:
The 64-bit primary has had random shuffling of chunks for a while, this
implements it for the 32-bit primary. Scudo is currently the only user of
`kRandomShuffleChunks`.

This change consists of a few modifications:
- move the random shuffling functions out of the 64-bit primary to
  `sanitizer_common.h`. Alternatively I could move them to
  `sanitizer_allocator.h` as they are only used in the allocator, I don't feel
  strongly either way;
- small change in the 64-bit primary to make the `rand_state` initialization
  `UNLIKELY`;
- addition of a `rand_state` in the 32-bit primary's `SizeClassInfo` and
  shuffling of chunks when populating the free list.
- enabling the `random_shuffle.cpp` test on platforms using the 32-bit primary
  for Scudo.

Some comments on why the shuffling is done that way. Initially I just
implemented a `Shuffle` function in the `TransferBatch` which was simpler but I
came to realize this wasn't good enough: for chunks of 10000 bytes for example,
with a `CompactSizeClassMap`, a batch holds only 1 chunk, meaning shuffling the
batch has no effect, while a region is usually 1MB, eg: 104 chunks of that size.
So I decided to "stage" the newly gathered chunks in a temporary array that
would be shuffled prior to placing the chunks in batches.
The result is looping twice through n_chunks even if shuffling is not enabled,
but I didn't notice any significant significant performance impact.

Reviewers: alekseyshl

Reviewed By: alekseyshl

Subscribers: srhines, llvm-commits, kubamracek

Differential Revision: https://reviews.llvm.org/D39244

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@316596 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 61 insertions, 24 deletions

diff --git a/lib/sanitizer_common/sanitizer_allocator.h b/lib/sanitizer_common/sanitizer_allocator.h
index 8c5696ea7..38368361d 100644
--- a/lib/sanitizer_common/sanitizer_allocator.h
+++ b/lib/sanitizer_common/sanitizer_allocator.h
@@ -56,6 +56,19 @@ struct NoOpMapUnmapCallback {
 // Callback type for iterating over chunks.
 typedef void (*ForEachChunkCallback)(uptr chunk, void *arg);
 
+INLINE u32 Rand(u32 *state) {  // ANSI C linear congruential PRNG.
+  return (*state = *state * 1103515245 + 12345) >> 16;
+}
+
+INLINE u32 RandN(u32 *state, u32 n) { return Rand(state) % n; }  // [0, n)
+
+template<typename T>
+INLINE void RandomShuffle(T *a, u32 n, u32 *rand_state) {
+  if (n <= 1) return;
+  for (u32 i = n - 1; i > 0; i--)
+    Swap(a[i], a[RandN(rand_state, i + 1)]);
+}
+
 #include "sanitizer_allocator_size_class_map.h"
 #include "sanitizer_allocator_stats.h"
 #include "sanitizer_allocator_primary64.h"
diff --git a/lib/sanitizer_common/sanitizer_allocator_primary32.h b/lib/sanitizer_common/sanitizer_allocator_primary32.h
index 8abf2f9f5..6c682aff6 100644
--- a/lib/sanitizer_common/sanitizer_allocator_primary32.h
+++ b/lib/sanitizer_common/sanitizer_allocator_primary32.h
@@ -268,7 +268,8 @@ class SizeClassAllocator32 {
   struct SizeClassInfo {
     SpinMutex mutex;
     IntrusiveList<TransferBatch> free_list;
-    char padding[kCacheLineSize - sizeof(uptr) -
+    u32 rand_state;
+    char padding[kCacheLineSize - 2 * sizeof(uptr) -
                  sizeof(IntrusiveList<TransferBatch>)];
   };
   COMPILER_CHECK(sizeof(SizeClassInfo) == kCacheLineSize);
@@ -301,29 +302,62 @@ class SizeClassAllocator32 {
     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);
+    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();
+      }
+      b->Add((void*)pointers_array[i]);
+      if (b->Count() == max_count) {
+        sci->free_list.push_back(b);
+        b = nullptr;
+      }
+    }
+    *current_batch = b;
+    return true;
+  }
+
   bool PopulateFreeList(AllocatorStats *stat, AllocatorCache *c,
                         SizeClassInfo *sci, uptr class_id) {
     uptr size = ClassIdToSize(class_id);
     uptr reg = AllocateRegion(stat, class_id);
     if (UNLIKELY(!reg))
       return false;
+    if (kRandomShuffleChunks)
+      if (UNLIKELY(sci->rand_state == 0))
+        // The random state is initialized from ASLR (PIE) and time.
+        sci->rand_state = reinterpret_cast<uptr>(sci) ^ NanoTime();
     uptr n_chunks = kRegionSize / (size + kMetadataSize);
     uptr max_count = TransferBatch::MaxCached(class_id);
     CHECK_GT(max_count, 0);
     TransferBatch *b = nullptr;
+    const uptr kShuffleArraySize = 48;
+    uptr shuffle_array[kShuffleArraySize];
+    uptr count = 0;
     for (uptr i = reg; i < reg + n_chunks * size; i += size) {
-      if (!b) {
-        b = c->CreateBatch(class_id, this, (TransferBatch*)i);
-        if (UNLIKELY(!b))
+      shuffle_array[count++] = i;
+      if (count == kShuffleArraySize) {
+        if (UNLIKELY(!PopulateBatches(c, sci, class_id, &b, max_count,
+                                      shuffle_array, count)))
           return false;
-        b->Clear();
-      }
-      b->Add((void*)i);
-      if (b->Count() == max_count) {
-        sci->free_list.push_back(b);
-        b = nullptr;
+        count = 0;
       }
     }
+    if (count) {
+      if (UNLIKELY(!PopulateBatches(c, sci, class_id, &b, max_count,
+                                    shuffle_array, count)))
+        return false;
+    }
     if (b) {
       CHECK_GT(b->Count(), 0);
       sci->free_list.push_back(b);
diff --git a/lib/sanitizer_common/sanitizer_allocator_primary64.h b/lib/sanitizer_common/sanitizer_allocator_primary64.h
index bb2b917f6..8b7610b2c 100644
--- a/lib/sanitizer_common/sanitizer_allocator_primary64.h
+++ b/lib/sanitizer_common/sanitizer_allocator_primary64.h
@@ -597,18 +597,6 @@ class SizeClassAllocator64 {
   };
   COMPILER_CHECK(sizeof(RegionInfo) >= kCacheLineSize);
 
-  u32 Rand(u32 *state) {  // ANSI C linear congruential PRNG.
-    return (*state = *state * 1103515245 + 12345) >> 16;
-  }
-
-  u32 RandN(u32 *state, u32 n) { return Rand(state) % n; }  // [0, n)
-
-  void RandomShuffle(u32 *a, u32 n, u32 *rand_state) {
-    if (n <= 1) return;
-    for (u32 i = n - 1; i > 0; i--)
-      Swap(a[i], a[RandN(rand_state, i + 1)]);
-  }
-
   RegionInfo *GetRegionInfo(uptr class_id) const {
     CHECK_LT(class_id, kNumClasses);
     RegionInfo *regions =
@@ -681,8 +669,10 @@ class SizeClassAllocator64 {
 
     // Map more space for chunks, if necessary.
     if (new_space_end > region->mapped_user) {
-      if (!kUsingConstantSpaceBeg && region->mapped_user == 0)
-        region->rand_state = static_cast<u32>(region_beg >> 12);  // From ASLR.
+      if (!kUsingConstantSpaceBeg && kRandomShuffleChunks)
+        if (UNLIKELY(region->mapped_user == 0))
+          // The random state is initialized from ASLR.
+          region->rand_state = static_cast<u32>(region_beg >> 12);
       // Do the mmap for the user memory.
       uptr map_size = kUserMapSize;
       while (new_space_end > region->mapped_user + map_size)