[sanitizer] 64-bit allocator's PopulateFreeArray partial refactor

[cryptoad]

Summary

This is an attempt at making PopulateFreeArray less obscure, more consistent,
and a tiny bit faster in some circumstances:

• use more consistent variable names, that work both for the user & the metadata portions of the code; the purpose of the code is mostly the same for both regions, so it makes sense that the code should be mostly similar as well;
• replace the while sum loops with a single RoundUpTo;
• mask most of the metadata computations behind kMetadataSize, allowing some blocks to be completely optimized out if not use metadata;
• const the constant variables;
• add a LIKELY as the branch it applies to will almost always be taken.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_primary64.h

732

I'd suggest to wrap the entire metadata allocating block (from // Calculate the required space for metadata) with:

if (kMetadataSize) {
  ...
}

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_primary64.h

732

So I thought about it, but I was wondering about something: currently the check for exhaustion will be done whether or not we use metadata.
I wanted to leave it like that at first, because it sort of made sense.
But then I figure that if we are not using metadata, and are exhausted, then the initial mmap would fail anyway since we'd hit the free array.
It seems to be safe to place all of this within if (kMetadataSize), but I wanted to raise this to your attention and get your feedback.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_primary64.h

732

Yep, now we will return false due to MapWithCallback failure and the OOM error handling will kick in. I think, it's just what we want.

[cryptoad]

lib/sanitizer_common/sanitizer_allocator_primary64.h

732

Tried out wrapping everything within a if (kMetadataSize).
The SizeClassAllocator64PopulateFreeListOOM test fails.
If we don't keep the exhaustion check out of the if, the map is successful, and we "overflow" into the free array space (it's empty and mappable).
So I think the current way is the way to go.

[alekseyshl]

lib/sanitizer_common/sanitizer_allocator_primary64.h

732

That means, when kMetadataSize == 0 and exhaustion check is true, we already stepped into the free array zone, right? We have not used any of it yet, but still, sloppy. Ok, let's take care of it in a separate patch.

[cryptoad]

Introduce IsRegionExhausted that checks if the region has enough room prior
to calling Mmap for the user & metadata portions of memory.
Fold all the metadata work in a if (kMetadataSize).

[alekseyshl]

Great, thanks!

lib/sanitizer_common/sanitizer_allocator_primary64.h

if (region->mapped_free_array < needed_space) {
       uptr new_map_size = new_mapped_free_array - region->mapped_free_array;
       if (UNLIKELY(!MapWithCallback(current_map_end, new_map_size)))
         return false;
       region->mapped_free_array = new_mapped_free_array;
     }
     return true;
   }
 
+  // Check whether this size class is exhausted.
+  bool IsRegionExhausted(RegionInfo *region, uptr class_id,
+                         uptr additional_map_size) {
+    if (LIKELY(region->mapped_user + region->mapped_meta +
+               additional_map_size <= kRegionSize - kFreeArraySize))
+      return false;
+    if (!region->exhausted) {
+      region->exhausted = true;
+      Printf("%s: Out of memory. ", SanitizerToolName);
+      Printf("The process has exhausted %zuMB for size class %zu.\n",
+             kRegionSize >> 20, ClassIdToSize(class_id));
+    }
+    return true;
+  }
+
   NOINLINE bool PopulateFreeArray(AllocatorStats *stat, uptr class_id,
                                   RegionInfo *region, uptr requested_count) {
     // region->mutex is held.
-    const uptr size = ClassIdToSize(class_id);
-    const uptr new_space_beg = region->allocated_user;
-    const uptr new_space_end = new_space_beg + requested_count * size;
     const uptr region_beg = GetRegionBeginBySizeClass(class_id);
+    const uptr size = ClassIdToSize(class_id);
 
+    const uptr total_user_bytes =
+        region->allocated_user + requested_count * size;
     // Map more space for chunks, if necessary.
-    if (new_space_end > region->mapped_user) {
+    if (LIKELY(total_user_bytes > region->mapped_user)) {
       if (UNLIKELY(region->mapped_user == 0)) {
         if (!kUsingConstantSpaceBeg && kRandomShuffleChunks)
           // The random state is initialized from ASLR.
           region->rand_state = static_cast<u32>(region_beg >> 12);
         // Postpone the first release to OS attempt for ReleaseToOSIntervalMs,
         // preventing just allocated memory from being released sooner than
         // necessary and also preventing extraneous ReleaseMemoryPagesToOS calls
         // for short lived processes.
         // Do it only when the feature is turned on, to avoid a potentially
         // extraneous syscall.
         if (ReleaseToOSIntervalMs() >= 0)
           region->rtoi.last_release_at_ns = NanoTime();
       }
       // Do the mmap for the user memory.
-      uptr map_size = kUserMapSize;
-      while (new_space_end > region->mapped_user + map_size)
-        map_size += kUserMapSize;
-      CHECK_GE(region->mapped_user + map_size, new_space_end);
+      const uptr user_map_size =
+          RoundUpTo(total_user_bytes - region->mapped_user, kUserMapSize);
+      if (UNLIKELY(IsRegionExhausted(region, class_id, user_map_size)))
+        return false;
       if (UNLIKELY(!MapWithCallback(region_beg + region->mapped_user,
-                                    map_size)))
+                                    user_map_size)))
         return false;
-      stat->Add(AllocatorStatMapped, map_size);
-      region->mapped_user += map_size;
+      stat->Add(AllocatorStatMapped, user_map_size);
+      region->mapped_user += user_map_size;
     }
-    const uptr new_chunks_count = (region->mapped_user - new_space_beg) / size;
+    const uptr new_chunks_count =
+        (region->mapped_user - region->allocated_user) / size;
 
+    if (kMetadataSize) {
       // Calculate the required space for metadata.
-    const uptr requested_allocated_meta =
+      const uptr total_meta_bytes =
           region->allocated_meta + new_chunks_count * kMetadataSize;
-    uptr requested_mapped_meta = region->mapped_meta;
-    while (requested_allocated_meta > requested_mapped_meta)
-      requested_mapped_meta += kMetaMapSize;
-    // Check whether this size class is exhausted.
-    if (region->mapped_user + requested_mapped_meta >
-        kRegionSize - kFreeArraySize) {
-      if (!region->exhausted) {
-        region->exhausted = true;
-        Printf("%s: Out of memory. ", SanitizerToolName);
-        Printf("The process has exhausted %zuMB for size class %zu.\n",
-               kRegionSize >> 20, size);
-      }
-      return false;
-    }
+      const uptr meta_map_size = (total_meta_bytes > region->mapped_meta) ?
+          RoundUpTo(total_meta_bytes - region->mapped_meta, kMetaMapSize) : 0;
       // Map more space for metadata, if necessary.
-    if (requested_mapped_meta > region->mapped_meta) {
+      if (meta_map_size) {
+        if (UNLIKELY(IsRegionExhausted(region, class_id, meta_map_size)))
+          return false;
         if (UNLIKELY(!MapWithCallback(
-              GetMetadataEnd(region_beg) - requested_mapped_meta,
-              requested_mapped_meta - region->mapped_meta)))
+            GetMetadataEnd(region_beg) - region->mapped_meta - meta_map_size,
+            meta_map_size)))
           return false;
-      region->mapped_meta = requested_mapped_meta;
+        region->mapped_meta += meta_map_size;
+      }
     }
 
     // If necessary, allocate more space for the free array and populate it with
     // newly allocated chunks.
     const uptr total_freed_chunks = region->num_freed_chunks + new_chunks_count;
     if (UNLIKELY(!EnsureFreeArraySpace(region, region_beg, total_freed_chunks)))
       return false;
     CompactPtrT *free_array = GetFreeArray(region_beg);
-    for (uptr i = 0, chunk = new_space_beg; i < new_chunks_count;
+    for (uptr i = 0, chunk = region->allocated_user; i < new_chunks_count;
          i++, chunk += size)
       free_array[total_freed_chunks - 1 - i] = PointerToCompactPtr(0, chunk);
     if (kRandomShuffleChunks)
       RandomShuffle(&free_array[region->num_freed_chunks], new_chunks_count,
                     &region->rand_state);
 
     // All necessary memory is mapped and now it is safe to advance all
     // 'allocated_*' counters.
     region->num_freed_chunks += new_chunks_count;
     region->allocated_user += new_chunks_count * size;
     CHECK_LE(region->allocated_user, region->mapped_user);
-    region->allocated_meta = requested_allocated_meta;
+    region->allocated_meta += new_chunks_count * kMetadataSize;
     CHECK_LE(region->allocated_meta, region->mapped_meta);
     region->exhausted = false;
 
     // TODO(alekseyshl): Consider bumping last_release_at_ns here to prevent
     // MaybeReleaseToOS from releasing just allocated pages or protect these
     // not yet used chunks some other way.
 
     return true;