mac: Periodically shim new malloc zones immediately after startup.

base/allocator/allocator_interception_mac.mm

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 // This file contains all the logic necessary to intercept allocations on
 // macOS. "malloc zones" are an abstraction that allows the process to intercept
 // all malloc-related functions.  There is no good mechanism [short of
 // interposition] to determine new malloc zones are added, so there's no clean
 // mechanism to intercept all malloc zones. This file contains logic to
 // intercept the default and purgeable zones, which always exist. A cursory
@@ skipping 12 matching lines @@
 #include <mach/mach.h>
 #include <mach/mach_vm.h>
 #import <objc/runtime.h>
 #include <stddef.h>
 
 #include <new>
 
 #include "base/allocator/allocator_shim.h"
 #include "base/allocator/features.h"
 #include "base/allocator/malloc_zone_functions_mac.h"
+#include "base/base_switches.h"
+#include "base/bind.h"
+#include "base/command_line.h"
 #include "base/logging.h"
 #include "base/mac/mac_util.h"
 #include "base/mac/mach_logging.h"
 #include "base/process/memory.h"
 #include "base/scoped_clear_errno.h"
+#include "base/threading/sequenced_task_runner_handle.h"
 #include "build/build_config.h"
 #include "third_party/apple_apsl/CFBase.h"
 
 namespace base {
 namespace allocator {
 
 bool g_replaced_default_zone = false;
 
 namespace {
 
@@ skipping 443 matching lines @@
   Method orig_method =
       class_getClassMethod(nsobject_class, @selector(allocWithZone:));
   g_old_allocWithZone =
       reinterpret_cast<allocWithZone_t>(method_getImplementation(orig_method));
   CHECK(g_old_allocWithZone)
       << "Failed to get allocWithZone allocation function.";
   method_setImplementation(orig_method,
                            reinterpret_cast<IMP>(oom_killer_allocWithZone));
 }
 
+namespace {
+
+void ShimNewMallocZones(base::Time end_time, base::TimeDelta delay) {
+  StoreFunctionsForAllZones();
+
+  // Use the functions for the default zone as a template to replace those
+  // new zones.
+  ChromeMallocZone* default_zone =
+      reinterpret_cast<ChromeMallocZone*>(malloc_default_zone());
+  DCHECK(IsMallocZoneAlreadyStored(default_zone));
+
+  MallocZoneFunctions new_functions;
+  StoreZoneFunctions(default_zone, &new_functions);
+  ReplaceFunctionsForStoredZones(&new_functions);
+
+  if (base::Time::Now() > end_time)
+    return;
+
+  base::TimeDelta next_delay = delay * 2;
+  SequencedTaskRunnerHandle::Get()->PostDelayedTask(
+      FROM_HERE, base::Bind(&ShimNewMallocZones, end_time, next_delay), delay);
+}
+
+}  // namespace
+
+void PeriodicallyShimNewMallocZones() {
+  if (!CommandLine::InitializedForCurrentProcess() ||
+      !CommandLine::ForCurrentProcess()->HasSwitch(
+          switches::kEnableHeapProfiling)) {
+    return;
+  }
+  base::Time end_time = base::Time::Now() + base::TimeDelta::FromMinutes(1);
+  base::TimeDelta initial_delay = base::TimeDelta::FromSeconds(1);
+  ShimNewMallocZones(end_time, initial_delay);
+}
+
 }  // namespace allocator
 }  // namespace base