[sanitizer] MmapAlignedOrDie changes to reduce fragmentation

[cryptoad]

Summary

The reasoning behind this change is explained in D33454, which unfortunately
broke the Windows version (due to the platform not supporting partial unmapping
of a memory region).

This new approach changes MmapAlignedOrDie to allow for the specification of
a padding_chunk. If non-null, and the initial allocation is aligned, this
padding chunk will hold the address of the extra memory (of alignment bytes).
This allows AllocateRegion to get 2 regions if the memory is aligned
properly, and thus help reduce fragmentation (and saves on unmapping
operations). As with the initial D33454, we use a stash in the 32-bit Primary
to hold those extra regions and return them on the fast-path.

The Windows version of MmapAlignedOrDie will always return a 0
padding_chunk if one was requested.

[cryptoad]

Correct the top level comment that wasn't accurate anymore.

[cryptoad]

Explicitely initialize num_stashed_regions to 0.

[alekseyshl]

lib/sanitizer_common/sanitizer_common.h

98

How about this:

Since the predominant use case of this function is "size == alignment" and the nature of the
way the alignment requirement is satisfied (by allocating size+alignment bytes of memory), there's a potential
of address space fragmentation. The padding_chunk parameter provides the opportunity
to return the contiguous padding of "size" bytes of the allocated chunk if the initial allocation happened to be
// perfectly aligned and the platform supports partial unmapping of the mapped region.

lib/sanitizer_common/sanitizer_posix.cc

159

Just realized that it should be

if (is_aligned && padding_chunk && size == alignment) {

...

The 'second chunk' logic does not work when size != alignment. With some improvements, it might help with size < alignment case, but, fortunately, our only use case for now is size == alignment, let's not complicate it more than necessary.

lib/sanitizer_common/tests/sanitizer_common_test.cc

130

Since the structure of the test is the same, maybe we should move #if#else#endinf into the loop?I mean, ifdef'ing only lines 100-103 and 118-129?

      EXPECT_EQ(0U, res % (alignment * PageSize));
      internal_memset((void*)res, 1, size * PageSize);
#if SANITIZER_WINDOWS
      // padding_chunk is not supported on Windows.
      EXPECT_EQ(0U, padding_chunk);
      UnmapOrDie((void*)res, size * PageSize);
#else
      if (size == 1 && alignment == 1) {
        // mmap returns PageSize aligned chunks, so this is a specific case
        // where we can check that padding_chunk will never be 0.
        EXPECT_NE(0U, padding_chunk);
      }
      if (padding_chunk) {
        EXPECT_EQ(res + size * PageSize, padding_chunk);
        internal_memset((void*)padding_chunk, 1, alignment * PageSize);
        UnmapOrDie((void*)padding_chunk, alignment * PageSize);
      }
#endif

[kcc]

lib/sanitizer_common/tests/sanitizer_common_test.cc

130

Try to avoid #if at all costs, even in tests.
In this case, I am pretty sure you can do

if(SANITIZER_WINDOWS)

instead.

[cryptoad]

Addressing review comments:

• modify MmapAlignedOrDie's comment in sanitizer_common.h;
• require size == alignment for a padding_chunk to be returned, adjust the test accordingly;
• change the structure of the new test with an if (SANITIZER_WINDOWS).

lib/sanitizer_common/sanitizer_allocator_primary32.h

 // SizeClassAllocator32 -- allocator for 32-bit address space.
 // This allocator can theoretically be used on 64-bit arch, but there it is less
 // efficient than SizeClassAllocator64.
 //
 // [kSpaceBeg, kSpaceBeg + kSpaceSize) is the range of addresses which can
 // be returned by MmapOrDie().
 //
 // Region:
-//   a result of a single call to MmapAlignedOrDie(kRegionSize, kRegionSize).
+//   a result of an allocation of kRegionSize bytes aligned on kRegionSize.
 // Since the regions are aligned by kRegionSize, there are exactly
 // kNumPossibleRegions possible regions in the address space and so we keep
 // a ByteMap possible_regions to store the size classes of each Region.
 // 0 size class means the region is not used by the allocator.
 //
 // One Region is used to allocate chunks of a single size class.
 // A Region looks like this:
 // UserChunk1 .. UserChunkN <gap> MetaChunkN .. MetaChunk1
uptr ComputeRegionId(uptr mem) {
     CHECK_LT(res, kNumPossibleRegions);
     return res;
   }
 
   uptr ComputeRegionBeg(uptr mem) {
     return mem & ~(kRegionSize - 1);
   }
 
+  // Allocates a region of kRegionSize bytes, aligned on kRegionSize. If we get
+  // more than one region back (in the event the allocation is aligned on the
+  // first try), attempt to store the second region into a stash. If the stash
+  // is full, just unmap the superfluous memory.
+  uptr AllocateRegionSlow(AllocatorStats *stat) {
+    uptr map_size = kRegionSize;
+    uptr padding_chunk;
+    uptr region = reinterpret_cast<uptr>(
+        MmapAlignedOrDie(kRegionSize, kRegionSize, "SizeClassAllocator32",
+                         &padding_chunk));
+    if (padding_chunk) {
+      // We have an extra region, attempt to stash it.
+      CHECK_EQ(padding_chunk, region + kRegionSize);
+      bool trim_extra = true;
+      {
+        SpinMutexLock l(&regions_stash_mutex);
+        if (num_stashed_regions < kMaxStashedRegions) {
+          regions_stash[num_stashed_regions++] = padding_chunk;
+          map_size = 2 * kRegionSize;
+          trim_extra = false;
+        }
+      }
+      if (trim_extra)
+        UnmapOrDie((void*)padding_chunk, kRegionSize);
+    }
+    MapUnmapCallback().OnMap(region, map_size);
+    stat->Add(AllocatorStatMapped, map_size);
+    return region;
+  }
+
   uptr AllocateRegion(AllocatorStats *stat, uptr class_id) {
     CHECK_LT(class_id, kNumClasses);
-    uptr res = reinterpret_cast<uptr>(MmapAlignedOrDie(kRegionSize, kRegionSize,
-                                      "SizeClassAllocator32"));
-    MapUnmapCallback().OnMap(res, kRegionSize);
-    stat->Add(AllocatorStatMapped, kRegionSize);
-    CHECK_EQ(0U, (res & (kRegionSize - 1)));
-    possible_regions.set(ComputeRegionId(res), static_cast<u8>(class_id));
-    return res;
+    uptr region = 0;
+    {
+      SpinMutexLock l(&regions_stash_mutex);
+      if (num_stashed_regions > 0)
+        region = regions_stash[--num_stashed_regions];
+    }
+    if (!region)
+      region = AllocateRegionSlow(stat);
+    CHECK(IsAligned(region, kRegionSize));
+    possible_regions.set(ComputeRegionId(region), static_cast<u8>(class_id));
+    return region;
   }
 
   SizeClassInfo *GetSizeClassInfo(uptr class_id) {
     CHECK_LT(class_id, kNumClasses);
     return &size_class_info_array[class_id];
   }
 
   void PopulateFreeList(AllocatorStats *stat, AllocatorCache *c,
for (uptr i = reg; i < reg + n_chunks * size; i += size) {
       }
     }
     if (b) {
       CHECK_GT(b->Count(), 0);
       sci->free_list.push_back(b);
     }
   }
 
+  // Unless several threads request regions simultaneously from different size
+  // classes, the stash rarely contains more than 1 entry.
+  static const uptr kMaxStashedRegions = 8;
+  SpinMutex regions_stash_mutex;
+  uptr num_stashed_regions;
+  uptr regions_stash[kMaxStashedRegions];
+
   ByteMap possible_regions;
   SizeClassInfo size_class_info_array[kNumClasses];
 };

lib/sanitizer_common/sanitizer_common.h

 void UnmapOrDie(void *addr, uptr size);
 void *MmapFixedNoReserve(uptr fixed_addr, uptr size,
                          const char *name = nullptr);
 void *MmapNoReserveOrDie(uptr size, const char *mem_type);
 void *MmapFixedOrDie(uptr fixed_addr, uptr size);
 void *MmapFixedNoAccess(uptr fixed_addr, uptr size, const char *name = nullptr);
 void *MmapNoAccess(uptr size);
 // Map aligned chunk of address space; size and alignment are powers of two.
-void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type);
+// The padding_chunk parameter leaves us the opportunity to return a second
+// contiguous chunk if the initial allocation is aligned, in order to save on
+// some unmapping operations and reduce memory fragmentation. That chunk will
+// hold 'alignment' bytes.
+void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type,
+                       uptr *padding_chunk);
 // Disallow access to a memory range.  Use MmapFixedNoAccess to allocate an
 // unaccessible memory.
 bool MprotectNoAccess(uptr addr, uptr size);
 bool MprotectReadOnly(uptr addr, uptr size);
 
 // Find an available address space.
 uptr FindAvailableMemoryRange(uptr size, uptr alignment, uptr left_padding);
 

lib/sanitizer_common/sanitizer_posix.cc

if (internal_iserror(res)) {
     Report("ERROR: %s failed to deallocate 0x%zx (%zd) bytes at address %p\n",
            SanitizerToolName, size, size, addr);
     CHECK("unable to unmap" && 0);
   }
   DecreaseTotalMmap(size);
 }
 
 // We want to map a chunk of address space aligned to 'alignment'.
-// We do it by maping a bit more and then unmaping redundant pieces.
+// We do it by mapping a bit more and then unmapping redundant pieces.
 // We probably can do it with fewer syscalls in some OS-dependent way.
-void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type) {
+void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type,
+                       uptr* padding_chunk) {
   CHECK(IsPowerOfTwo(size));
   CHECK(IsPowerOfTwo(alignment));
   uptr map_size = size + alignment;
   uptr map_res = (uptr)MmapOrDie(map_size, mem_type);
   uptr map_end = map_res + map_size;
+  bool is_aligned = IsAligned(map_res, alignment);
+  if (is_aligned && padding_chunk) {
+    *padding_chunk = map_res + size;
+    return (void *)map_res;
+  }
+  if (padding_chunk)
+    *padding_chunk = 0;
   uptr res = map_res;
-  if (res & (alignment - 1))  // Not aligned.
-    res = (map_res + alignment) & ~(alignment - 1);
-  uptr end = res + size;
-  if (res != map_res)
+  if (!is_aligned) {
+    res = (map_res + alignment - 1) & ~(alignment - 1);
     UnmapOrDie((void*)map_res, res - map_res);
-  if (end != map_end)
-    UnmapOrDie((void*)end, map_end - end);
+  }
+  uptr end = res + size;
   return (void*)res;
 }
 
 void *MmapNoReserveOrDie(uptr size, const char *mem_type) {
   uptr PageSize = GetPageSizeCached();
   uptr p = internal_mmap(nullptr,
                          RoundUpTo(size, PageSize),
                          PROT_READ | PROT_WRITE,

lib/sanitizer_common/sanitizer_win.cc

if (VirtualFree(addr, size, MEM_DECOMMIT) == 0) {
              "deallocate 0x%zx (%zd) bytes at address %p (error code: %d)\n",
              SanitizerToolName, size, size, addr, GetLastError());
       CHECK("unable to unmap" && 0);
     }
   }
 }
 
 // We want to map a chunk of address space aligned to 'alignment'.
-void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type) {
+void *MmapAlignedOrDie(uptr size, uptr alignment, const char *mem_type,
+                       uptr *padding_chunk) {
   CHECK(IsPowerOfTwo(size));
   CHECK(IsPowerOfTwo(alignment));
 
+  if (padding_chunk)
+    *padding_chunk = 0;
+
   // Windows will align our allocations to at least 64K.
   alignment = Max(alignment, GetMmapGranularity());
 
   uptr mapped_addr =
       (uptr)VirtualAlloc(0, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
   if (!mapped_addr)
     ReportMmapFailureAndDie(size, mem_type, "allocate aligned", GetLastError());
 

lib/sanitizer_common/tests/sanitizer_common_test.cc

TEST(SanitizerCommon, SortTest) {
   EXPECT_TRUE(IsSorted(array, 2));
 }
 
 TEST(SanitizerCommon, MmapAlignedOrDie) {
   uptr PageSize = GetPageSizeCached();
   for (uptr size = 1; size <= 32; size *= 2) {
     for (uptr alignment = 1; alignment <= 32; alignment *= 2) {
       for (int iter = 0; iter < 100; iter++) {
-        uptr res = (uptr)MmapAlignedOrDie(
-            size * PageSize, alignment * PageSize, "MmapAlignedOrDieTest");
+        uptr res = (uptr)MmapAlignedOrDie(size * PageSize, alignment * PageSize,
+            "MmapAlignedOrDieTest", nullptr);
         EXPECT_EQ(0U, res % (alignment * PageSize));
         internal_memset((void*)res, 1, size * PageSize);
         UnmapOrDie((void*)res, size * PageSize);
       }
     }
   }
 }
 
+#if SANITIZER_WINDOWS
+TEST(SanitizerCommon, MmapAlignedOrDiePaddingChunk) {
+  uptr PageSize = GetPageSizeCached();
+  for (uptr size = 1; size <= 32; size *= 2) {
+    for (uptr alignment = 1; alignment <= 32; alignment *= 2) {
+      for (int iter = 0; iter < 100; iter++) {
+        uptr padding_chunk;
+        uptr res = (uptr)MmapAlignedOrDie(size * PageSize, alignment * PageSize,
+            "MmapAlignedOrDiePaddingChunkTest", &padding_chunk);
+        EXPECT_EQ(0U, res % (alignment * PageSize));
+        // padding_chunk is not supported on Windows.
+        EXPECT_EQ(0U, padding_chunk);
+        internal_memset((void*)res, 1, size * PageSize);
+        UnmapOrDie((void*)res, size * PageSize);
+      }
+    }
+  }
+}
+#else
+TEST(SanitizerCommon, MmapAlignedOrDiePaddingChunk) {
+  uptr PageSize = GetPageSizeCached();
+  for (uptr size = 1; size <= 32; size *= 2) {
+    for (uptr alignment = 1; alignment <= 32; alignment *= 2) {
+      for (int iter = 0; iter < 100; iter++) {
+        uptr padding_chunk;
+        uptr res = (uptr)MmapAlignedOrDie(size * PageSize, alignment * PageSize,
+            "MmapAlignedOrDiePaddingChunkTest", &padding_chunk);
+        EXPECT_EQ(0U, res % (alignment * PageSize));
+        internal_memset((void*)res, 1, size * PageSize);
+        UnmapOrDie((void*)res, size * PageSize);
+        if (size == 1 && alignment == 1) {
+          // mmap returns PageSize aligned chunks, so this is a specific case
+          // where we can check that padding_chunk will never be 0.
+          EXPECT_NE(0U, padding_chunk);
+        }
+        if (padding_chunk) {
+          EXPECT_EQ(res + size * PageSize, padding_chunk);
+          internal_memset((void*)padding_chunk, 1, alignment * PageSize);
+          UnmapOrDie((void*)padding_chunk, alignment * PageSize);
+        }
+      }
+    }
+  }
+}
+#endif
+
 #if SANITIZER_LINUX
 TEST(SanitizerCommon, SanitizerSetThreadName) {
   const char *names[] = {
     "0123456789012",
     "01234567890123",
     "012345678901234",  // Larger names will be truncated on linux.
   };