PartitionAlloc: PCScan: Defer PCScan initialization

The CL enables PCScan when
 1) PA-Everywhere,
 2) --enable-features=PartitionAllocPCScan.
and makes sure that it's enabled after FeatureList is initialized.

Bug: 11297512
Change-Id: Ibcf21a503cdf0cb1a085629ab77d4e87061f8fa0
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2473637
Commit-Queue: Anton Bikineev <bikineev@chromium.org>
Reviewed-by: Bartek Nowierski <bartekn@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Cr-Commit-Position: refs/heads/master@{#819429}

base/allocator/allocator_shim.h

@@ -7,6 +7,7 @@
 
 #include <stddef.h>
 
+#include "base/allocator/buildflags.h"
 #include "base/base_export.h"
 #include "build/build_config.h"
 
@@ -150,6 +151,10 @@ BASE_EXPORT void RemoveAllocatorDispatchForTesting(AllocatorDispatch* dispatch);
 BASE_EXPORT void InitializeAllocatorShim();
 #endif  // defined(OS_APPLE)
 
+#if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
+BASE_EXPORT void EnablePCScanIfNeeded();
+#endif
+
 }  // namespace allocator
 }  // namespace base
 

base/allocator/allocator_shim_default_dispatch_to_partition_alloc.cc

@@ -6,6 +6,7 @@
 #include "base/allocator/allocator_shim_internals.h"
 #include "base/allocator/partition_allocator/partition_alloc.h"
 #include "base/allocator/partition_allocator/partition_alloc_constants.h"
+#include "base/allocator/partition_allocator/partition_alloc_features.h"
 #include "base/bits.h"
 #include "base/no_destructor.h"
 #include "build/build_config.h"
@@ -78,7 +79,8 @@ base::ThreadSafePartitionRoot& Allocator() {
 
   auto* new_root = new (g_allocator_buffer) base::ThreadSafePartitionRoot(
       {base::PartitionOptions::Alignment::kRegular,
-       base::PartitionOptions::ThreadCache::kEnabled});
+       base::PartitionOptions::ThreadCache::kEnabled,
+       base::PartitionOptions::PCScan::kDisabledByDefault});
   g_root_.store(new_root, std::memory_order_release);
 
   // Semantically equivalent to base::Lock::Release().
@@ -108,8 +110,10 @@ void* PartitionCalloc(const AllocatorDispatch*,
 base::ThreadSafePartitionRoot* AlignedAllocator() {
   // Since the general-purpose allocator uses the thread cache, this one cannot.
   static base::NoDestructor<base::ThreadSafePartitionRoot> aligned_allocator(
-      base::PartitionOptions{base::PartitionOptions::Alignment::kAlignedAlloc,
-                             base::PartitionOptions::ThreadCache::kDisabled});
+      base::PartitionOptions{
+          base::PartitionOptions::Alignment::kAlignedAlloc,
+          base::PartitionOptions::ThreadCache::kDisabled,
+          base::PartitionOptions::PCScan::kDisabledByDefault});
   return aligned_allocator.get();
 }
 
@@ -186,6 +190,21 @@ size_t PartitionGetSizeEstimate(const AllocatorDispatch*,
 
 }  // namespace
 
+namespace base {
+namespace allocator {
+
+#if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
+void EnablePCScanIfNeeded() {
+  if (!features::IsPartitionAllocPCScanEnabled())
+    return;
+  Allocator().EnablePCScan();
+  AlignedAllocator()->EnablePCScan();
+}
+#endif
+
+}  // namespace allocator
+}  // namespace base
+
 constexpr AllocatorDispatch AllocatorDispatch::default_dispatch = {
     &PartitionMalloc,          /* alloc_function */
     &PartitionMallocUnchecked, /* alloc_unchecked_function */

base/allocator/partition_allocator/partition_alloc.cc

@@ -18,6 +18,7 @@
 #include "base/allocator/partition_allocator/partition_direct_map_extent.h"
 #include "base/allocator/partition_allocator/partition_oom.h"
 #include "base/allocator/partition_allocator/partition_page.h"
+#include "base/allocator/partition_allocator/partition_root.h"
 #include "base/check_op.h"
 #include "base/no_destructor.h"
 #include "base/synchronization/lock.h"
@@ -278,6 +279,7 @@ void PartitionRoot<thread_safe>::Init(PartitionOptions opts) {
   // the beginning of the slot.
   allow_extras = (opts.alignment != PartitionOptions::Alignment::kAlignedAlloc);
 
+  scannable = (opts.pcscan != PartitionOptions::PCScan::kAlwaysDisabled);
   // Concurrent freeing in PCScan can only safely work on thread-safe
   // partitions.
   if (thread_safe && opts.pcscan == PartitionOptions::PCScan::kEnabled)

base/allocator/partition_allocator/partition_alloc.h

@@ -80,7 +80,6 @@
 #include "base/compiler_specific.h"
 #include "base/gtest_prod_util.h"
 #include "base/notreached.h"
-#include "base/optional.h"
 #include "base/partition_alloc_buildflags.h"
 #include "base/stl_util.h"
 #include "base/synchronization/lock.h"

base/allocator/partition_allocator/partition_bucket.cc

@@ -291,12 +291,12 @@ ALWAYS_INLINE void* PartitionBucket<thread_safe>::AllocNewSlotSpan(
   char* tag_bitmap = super_page + PartitionPageSize();
   char* quarantine_bitmaps = tag_bitmap + ReservedTagBitmapSize();
   const size_t quarantine_bitmaps_size =
-      root->pcscan ? 2 * sizeof(QuarantineBitmap) : 0;
+      root->scannable ? 2 * sizeof(QuarantineBitmap) : 0;
   PA_DCHECK(quarantine_bitmaps_size % PartitionPageSize() == 0);
   char* ret = quarantine_bitmaps + quarantine_bitmaps_size;
   root->next_partition_page = ret + total_size;
   root->next_partition_page_end = root->next_super_page - PartitionPageSize();
-  PA_DCHECK(ret == SuperPagePayloadBegin(super_page, root->pcscan.has_value()));
+  PA_DCHECK(ret == SuperPagePayloadBegin(super_page, root->scannable));
   PA_DCHECK(root->next_partition_page_end == SuperPagePayloadEnd(super_page));
 
   // The first slot span is accessible. The given committed_size is equal to

base/allocator/partition_allocator/partition_root.h

@@ -13,6 +13,7 @@
 #include "base/allocator/partition_allocator/partition_tag.h"
 #include "base/allocator/partition_allocator/pcscan.h"
 #include "base/allocator/partition_allocator/thread_cache.h"
+#include "base/optional.h"
 
 namespace base {
 
@@ -40,13 +41,20 @@ struct PartitionOptions {
   };
 
   enum class PCScan {
-    kDisabled,
+    // Should be used for value partitions, i.e. partitions that are known to
+    // not have pointers. No metadata (quarantine bitmaps) is allocated for such
+    // partitions.
+    kAlwaysDisabled,
+    // PCScan is disabled by default, but can be enabled by calling
+    // PartitionRoot::EnablePCScan().
+    kDisabledByDefault,
+    // PCScan is always enabled.
     kEnabled,
   };
 
   Alignment alignment = Alignment::kRegular;
   ThreadCache thread_cache = ThreadCache::kDisabled;
-  PCScan pcscan = PCScan::kDisabled;
+  PCScan pcscan = PCScan::kAlwaysDisabled;
 };
 
 // Never instantiate a PartitionRoot directly, instead use
@@ -70,6 +78,7 @@ struct BASE_EXPORT PartitionRoot {
   // nothing) instead of true|false, so that we can just add or subtract the
   // size instead of having an if branch on the hot paths.
   bool allow_extras;
+  bool scannable = false;
   bool initialized = false;
 
 #if ENABLE_TAG_FOR_CHECKED_PTR2 || ENABLE_TAG_FOR_MTE_CHECKED_PTR
@@ -225,6 +234,13 @@ struct BASE_EXPORT PartitionRoot {
     return features::IsPartitionAllocGigaCageEnabled() && allow_extras;
   }
 
+  void EnablePCScan() {
+    // TODO(bikineev): Make CHECK once PCScan is enabled.
+    if (!scannable || pcscan.has_value())
+      return;
+    pcscan.emplace(this);
+  }
+
  private:
   // Allocates memory, without initializing extras.
   //

base/allocator/partition_allocator/pcscan.cc

@@ -34,7 +34,7 @@ ThreadSafePartitionRoot& PCScanMetadataAllocator() {
   static base::NoDestructor<ThreadSafePartitionRoot> allocator{
       PartitionOptions{PartitionOptions::Alignment::kRegular,
                        PartitionOptions::ThreadCache::kDisabled,
-                       PartitionOptions::PCScan::kDisabled}};
+                       PartitionOptions::PCScan::kAlwaysDisabled}};
   return *allocator;
 }
 

content/app/content_main_runner_impl.cc

@@ -15,6 +15,7 @@
 
 #include "base/allocator/allocator_check.h"
 #include "base/allocator/allocator_extension.h"
+#include "base/allocator/allocator_shim.h"
 #include "base/allocator/buildflags.h"
 #include "base/at_exit.h"
 #include "base/base_switches.h"
@@ -232,6 +233,13 @@ void InitializeV8IfNeeded(const base::CommandLine& command_line,
 #endif  // V8_USE_EXTERNAL_STARTUP_DATA
 }
 
+void EnablePCScanForMallocPartitionsIfNeeded() {
+#if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
+  CHECK(base::FeatureList::GetInstance());
+  base::allocator::EnablePCScanIfNeeded();
+#endif
+}
+
 #if BUILDFLAG(USE_ZYGOTE_HANDLE)
 pid_t LaunchZygoteHelper(base::CommandLine* cmd_line,
                          base::ScopedFD* control_fd) {
@@ -479,6 +487,10 @@ int RunZygote(ContentMainDelegate* delegate) {
   InitializeFieldTrialAndFeatureList();
   delegate->PostFieldTrialInitialization();
 
+  // After feature list has been initialized, enable pcscan on malloc
+  // partitions.
+  EnablePCScanForMallocPartitionsIfNeeded();
+
   for (size_t i = 0; i < base::size(kMainFunctions); ++i) {
     if (process_type == kMainFunctions[i].name)
       return kMainFunctions[i].function(main_params);
@@ -529,6 +541,11 @@ int RunOtherNamedProcessTypeMain(const std::string& process_type,
     }
   }
 
+  if (process_type != switches::kZygoteProcess) {
+    // Zygote processes are handled in RunZygote.
+    EnablePCScanForMallocPartitionsIfNeeded();
+  }
+
 #if BUILDFLAG(USE_ZYGOTE_HANDLE)
   // Zygote startup is special -- see RunZygote comments above
   // for why we don't use ZygoteMain directly.
@@ -826,6 +843,10 @@ int ContentMainRunnerImpl::Run(bool start_service_manager_only) {
       // has been updated.
       InitializeFieldTrialAndFeatureList();
       delegate_->PostFieldTrialInitialization();
+
+      // After feature list has been initialized, enable pcscan on malloc
+      // partitions.
+      EnablePCScanForMallocPartitionsIfNeeded();
     }
 
 #if defined(OS_LINUX) || defined(OS_CHROMEOS)
@@ -965,6 +986,9 @@ int ContentMainRunnerImpl::RunServiceManager(MainFunctionParams& main_params,
 #endif
   }
 
+  // Enable PCScan once we are certain that FeatureList was initialized.
+  EnablePCScanForMallocPartitionsIfNeeded();
+
   if (should_start_service_manager_only) {
     DVLOG(0) << "Chrome is running in ServiceManager only mode.";
     return -1;