[PartitionAlloc] Optimize the `PartitionAllocZeroFill` option.

When we have satisfied the allocation with a new OS page, we can skip the
`memset`.

Bug: 864462
Change-Id: Ib13539a83a931dbad3e588001aa248ff994d3900
Reviewed-on: https://chromium-review.googlesource.com/1213957
Commit-Queue: Chris Palmer <palmer@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Cr-Commit-Position: refs/heads/master@{#595218}

base/allocator/partition_allocator/partition_alloc.h

@@ -348,9 +348,9 @@ ALWAYS_INLINE void* PartitionAllocGenericFlags(PartitionRootGeneric* root,
                                                size_t size,
                                                const char* type_name) {
   DCHECK_LT(flags, PartitionAllocLastFlag << 1);
-  const bool zero_fill = flags & PartitionAllocZeroFill;
 
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
+  const bool zero_fill = flags & PartitionAllocZeroFill;
   void* result = zero_fill ? calloc(1, size) : malloc(size);
   CHECK(result || flags & PartitionAllocReturnNull);
   return result;
@@ -366,13 +366,6 @@ ALWAYS_INLINE void* PartitionAllocGenericFlags(PartitionRootGeneric* root,
   }
   PartitionAllocHooks::AllocationHookIfEnabled(ret, requested_size, type_name);
 
-  // TODO(crbug.com/864462): This is suboptimal. Change `AllocFromBucket` such
-  // that it tells callers if the allocation was satisfied with a fresh mapping
-  // from the OS, so that we can skip this step and save some time.
-  if (ret && zero_fill) {
-    memset(ret, 0, requested_size);
-  }
-
   return ret;
 #endif
 }

base/allocator/partition_allocator/partition_alloc_unittest.cc

@@ -11,8 +11,11 @@
 #include <vector>
 
 #include "base/allocator/partition_allocator/address_space_randomization.h"
+#include "base/allocator/partition_allocator/partition_alloc.h"
 #include "base/bit_cast.h"
 #include "base/bits.h"
+#include "base/rand_util.h"
+#include "base/stl_util.h"
 #include "base/sys_info.h"
 #include "build/build_config.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -86,6 +89,35 @@ bool ClearAddressSpaceLimit() {
 #endif
 }
 
+const size_t kTestSizes[] = {
+    1,
+    17,
+    100,
+    base::kSystemPageSize,
+    base::kSystemPageSize + 1,
+    base::internal::PartitionBucket::get_direct_map_size(100),
+    1 << 20,
+    1 << 21,
+};
+constexpr size_t kTestSizesCount = base::size(kTestSizes);
+
+void AllocateRandomly(base::PartitionRootGeneric* root,
+                      size_t count,
+                      int flags) {
+  std::vector<void*> allocations(count, nullptr);
+  for (size_t i = 0; i < count; ++i) {
+    const size_t size = kTestSizes[base::RandGenerator(kTestSizesCount)];
+    allocations[i] = PartitionAllocGenericFlags(root, flags, size, nullptr);
+    EXPECT_NE(nullptr, allocations[i]) << " size: " << size << " i: " << i;
+  }
+
+  for (size_t i = 0; i < count; ++i) {
+    if (allocations[i]) {
+      base::PartitionFree(allocations[i]);
+    }
+  }
+}
+
 }  // namespace
 
 namespace base {
@@ -2156,20 +2188,9 @@ TEST_F(PartitionAllocTest, SmallReallocDoesNotMoveTrailingCookie) {
 }
 
 TEST_F(PartitionAllocTest, ZeroFill) {
-  const size_t test_sizes[] = {
-      1,
-      17,
-      100,
-      kSystemPageSize,
-      kSystemPageSize + 1,
-      internal::PartitionBucket::get_direct_map_size(100),
-      1 << 20,
-      1 << 21,
-  };
-
   constexpr static size_t kAllZerosSentinel =
       std::numeric_limits<size_t>::max();
-  for (size_t size : test_sizes) {
+  for (size_t size : kTestSizes) {
     char* p = static_cast<char*>(PartitionAllocGenericFlags(
         generic_allocator.root(), PartitionAllocZeroFill, size, nullptr));
     size_t non_zero_position = kAllZerosSentinel;
@@ -2183,6 +2204,11 @@ TEST_F(PartitionAllocTest, ZeroFill) {
         << "test allocation size: " << size;
     PartitionFree(p);
   }
+
+  for (int i = 0; i < 10; ++i) {
+    SCOPED_TRACE(i);
+    AllocateRandomly(generic_allocator.root(), 1000, PartitionAllocZeroFill);
+  }
 }
 
 }  // namespace internal

base/allocator/partition_allocator/partition_bucket.cc

@@ -448,11 +448,13 @@ bool PartitionBucket::SetNewActivePage() {
 
 void* PartitionBucket::SlowPathAlloc(PartitionRootBase* root,
                                      int flags,
-                                     size_t size) {
+                                     size_t size,
+                                     bool* is_already_zeroed) {
   // The slow path is called when the freelist is empty.
   DCHECK(!this->active_pages_head->freelist_head);
 
   PartitionPage* new_page = nullptr;
+  *is_already_zeroed = false;
 
   // For the PartitionRootGeneric::Alloc() API, we have a bunch of buckets
   // marked as special cases. We bounce them through to the slow path so that
@@ -474,6 +476,7 @@ void* PartitionBucket::SlowPathAlloc(PartitionRootBase* root,
       PartitionExcessiveAllocationSize();
     }
     new_page = PartitionDirectMap(root, flags, size);
+    *is_already_zeroed = true;
   } else if (LIKELY(this->SetNewActivePage())) {
     // First, did we find an active page in the active pages list?
     new_page = this->active_pages_head;
@@ -505,6 +508,7 @@ void* PartitionBucket::SlowPathAlloc(PartitionRootBase* root,
       void* addr = PartitionPage::ToPointer(new_page);
       root->RecommitSystemPages(addr, new_page->bucket->get_bytes_per_span());
       new_page->Reset();
+      *is_already_zeroed = true;
     }
     DCHECK(new_page);
   } else {
@@ -514,6 +518,7 @@ void* PartitionBucket::SlowPathAlloc(PartitionRootBase* root,
     if (LIKELY(raw_pages != nullptr)) {
       new_page = PartitionPage::FromPointerNoAlignmentCheck(raw_pages);
       InitializeSlotSpan(new_page);
+      *is_already_zeroed = true;
     }
   }
 

base/allocator/partition_allocator/partition_bucket.h

@@ -31,10 +31,17 @@ struct PartitionBucket {
   // Public API.
   void Init(uint32_t new_slot_size);
 
+  // Sets |is_already_zeroed| to true if the allocation was satisfied by
+  // requesting (a) new page(s) from the operating system, or false otherwise.
+  // This enables an optimization for when callers use |PartitionAllocZeroFill|:
+  // there is no need to call memset on fresh pages; the OS has already zeroed
+  // them. (See |PartitionRootBase::AllocFromBucket|.)
+  //
   // Note the matching Free() functions are in PartitionPage.
   BASE_EXPORT NOINLINE void* SlowPathAlloc(PartitionRootBase* root,
                                            int flags,
-                                           size_t size);
+                                           size_t size,
+                                           bool* is_already_zeroed);
 
   ALWAYS_INLINE bool is_direct_mapped() const {
     return !num_system_pages_per_slot_span;

base/allocator/partition_allocator/partition_root_base.h

@@ -86,16 +86,20 @@ struct BASE_EXPORT PartitionRootBase {
 ALWAYS_INLINE void* PartitionRootBase::AllocFromBucket(PartitionBucket* bucket,
                                                        int flags,
                                                        size_t size) {
+  bool zero_fill = flags & PartitionAllocZeroFill;
+  bool is_already_zeroed = false;
+
   PartitionPage* page = bucket->active_pages_head;
   // Check that this page is neither full nor freed.
   DCHECK(page->num_allocated_slots >= 0);
   void* ret = page->freelist_head;
   if (LIKELY(ret != 0)) {
-    // If these DCHECKs fire, you probably corrupted memory.
-    // TODO(palmer): See if we can afford to make this a CHECK.
+    // If these DCHECKs fire, you probably corrupted memory. TODO(palmer): See
+    // if we can afford to make these CHECKs.
     DCHECK(PartitionRootBase::IsValidPage(page));
-    // All large allocations must go through the slow path to correctly
-    // update the size metadata.
+
+    // All large allocations must go through the slow path to correctly update
+    // the size metadata.
     DCHECK(page->get_raw_size() == 0);
     internal::PartitionFreelistEntry* new_head =
         internal::PartitionFreelistEntry::Transform(
@@ -103,15 +107,17 @@ ALWAYS_INLINE void* PartitionRootBase::AllocFromBucket(PartitionBucket* bucket,
     page->freelist_head = new_head;
     page->num_allocated_slots++;
   } else {
-    ret = bucket->SlowPathAlloc(this, flags, size);
+    ret = bucket->SlowPathAlloc(this, flags, size, &is_already_zeroed);
     // TODO(palmer): See if we can afford to make this a CHECK.
     DCHECK(!ret ||
            PartitionRootBase::IsValidPage(PartitionPage::FromPointer(ret)));
   }
+
 #if DCHECK_IS_ON()
-  if (!ret)
-    return 0;
-  // Fill the uninitialized pattern, and write the cookies.
+  if (!ret) {
+    return nullptr;
+  }
+
   page = PartitionPage::FromPointer(ret);
   // TODO(ajwong): Can |page->bucket| ever not be |this|? If not, can this just
   // be bucket->slot_size?
@@ -126,11 +132,20 @@ ALWAYS_INLINE void* PartitionRootBase::AllocFromBucket(PartitionBucket* bucket,
   // The value given to the application is actually just after the cookie.
   ret = char_ret + kCookieSize;
 
-  // Debug fill region kUninitializedByte and surround it with 2 cookies.
+  // Fill the region kUninitializedByte or 0, and surround it with 2 cookies.
   PartitionCookieWriteValue(char_ret);
-  memset(ret, kUninitializedByte, no_cookie_size);
+  if (!zero_fill) {
+    memset(ret, kUninitializedByte, no_cookie_size);
+  } else if (!is_already_zeroed) {
+    memset(ret, 0, no_cookie_size);
+  }
   PartitionCookieWriteValue(char_ret + kCookieSize + no_cookie_size);
+#else
+  if (ret && zero_fill && !is_already_zeroed) {
+    memset(ret, 0, size);
+  }
 #endif
+
   return ret;
 }