Die on an OOM situation in many more cases.

Continuation of http://codereview.chromium.org/875004 after a revert.

BUG=http://crbug.com/12673
TEST=run out of memory and die

Review URL: http://codereview.chromium.org/1039007

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@42127 0039d316-1c4b-4281-b951-d872f2087c98

5 files changed, 249 insertions, 19 deletions

diff --git a/base/base.gyp b/base/base.gyp
index aae38b7..802147e 100644
--- a/base/base.gyp
+++ b/base/base.gyp
@@ -94,6 +94,8 @@
         'pickle_unittest.cc',
         'pr_time_unittest.cc',
         'process_util_unittest.cc',
+        'process_util_unittest_mac.h',
+        'process_util_unittest_mac.mm',
         'rand_util_unittest.cc',
         'ref_counted_unittest.cc',
         'scoped_bstr_win_unittest.cc',
diff --git a/base/process_util_mac.mm b/base/process_util_mac.mm
index faae5bf..4cf597c 100644
--- a/base/process_util_mac.mm
+++ b/base/process_util_mac.mm
@@ -11,12 +11,14 @@
 #include <mach/mach_init.h>
 #include <mach/task.h>
 #include <malloc/malloc.h>
+#import <objc/runtime.h>
 #include <spawn.h>
 #include <sys/mman.h>
 #include <sys/sysctl.h>
 #include <sys/types.h>
 #include <sys/wait.h>
 
+#include <new>
 #include <string>
 
 #include "base/debug_util.h"
@@ -355,6 +357,10 @@ size_t GetSystemCommitCharge() {
 
 namespace {
 
+bool g_oom_killer_enabled;
+
+// === C malloc/calloc/valloc/realloc ===
+
 typedef void* (*malloc_type)(struct _malloc_zone_t* zone,
                              size_t size);
 typedef void* (*calloc_type)(struct _malloc_zone_t* zone,
@@ -405,25 +411,86 @@ void* oom_killer_realloc(struct _malloc_zone_t* zone,
   return result;
 }
 
+// === C++ operator new ===
+
+void oom_killer_new() {
+  DebugUtil::BreakDebugger();
+}
+
+// === Core Foundation CFAllocators ===
+
+// This is the real structure of a CFAllocatorRef behind the scenes. See
+// http://opensource.apple.com/source/CF/CF-550/CFBase.c for details.
+struct ChromeCFAllocator {
+  _malloc_zone_t fake_malloc_zone;
+  void* allocator;
+  CFAllocatorContext context;
+};
+typedef ChromeCFAllocator* ChromeCFAllocatorRef;
+
+CFAllocatorAllocateCallBack g_old_cfallocator_system_default;
+CFAllocatorAllocateCallBack g_old_cfallocator_malloc;
+CFAllocatorAllocateCallBack g_old_cfallocator_malloc_zone;
+
+void* oom_killer_cfallocator_system_default(CFIndex alloc_size,
+                                            CFOptionFlags hint,
+                                            void* info) {
+  void* result = g_old_cfallocator_system_default(alloc_size, hint, info);
+  if (!result)
+    DebugUtil::BreakDebugger();
+  return result;
+}
+
+void* oom_killer_cfallocator_malloc(CFIndex alloc_size,
+                                    CFOptionFlags hint,
+                                    void* info) {
+  void* result = g_old_cfallocator_malloc(alloc_size, hint, info);
+  if (!result)
+    DebugUtil::BreakDebugger();
+  return result;
+}
+
+void* oom_killer_cfallocator_malloc_zone(CFIndex alloc_size,
+                                         CFOptionFlags hint,
+                                         void* info) {
+  void* result = g_old_cfallocator_malloc_zone(alloc_size, hint, info);
+  if (!result)
+    DebugUtil::BreakDebugger();
+  return result;
+}
+
+// === Cocoa NSObject allocation ===
+
+typedef id (*allocWithZone_t)(id, SEL, NSZone*);
+allocWithZone_t g_old_allocWithZone;
+
+id oom_killer_allocWithZone(id self, SEL _cmd, NSZone* zone)
+{
+  id result = g_old_allocWithZone(self, _cmd, zone);
+  if (!result)
+    DebugUtil::BreakDebugger();
+  return result;
+}
+
 }  // namespace
 
 void EnableTerminationOnOutOfMemory() {
+  if (g_oom_killer_enabled)
+    return;
+
+  g_oom_killer_enabled = true;
+
+  // === C malloc/calloc/valloc/realloc ===
+
+  // This approach is not perfect, as requests for amounts of memory larger than
+  // MALLOC_ABSOLUTE_MAX_SIZE (currently SIZE_T_MAX - (2 * PAGE_SIZE)) will
+  // still fail with a NULL rather than dying (see
+  // http://opensource.apple.com/source/Libc/Libc-583/gen/malloc.c for details).
+  // Unfortunately, it's the best we can do. Also note that this does not affect
+  // allocations from non-default zones.
+
   CHECK(!g_old_malloc && !g_old_calloc && !g_old_valloc && !g_old_realloc)
-      << "EnableTerminationOnOutOfMemory() called twice!";
-
-  // This approach is sub-optimal:
-  // - Requests for amounts of memory larger than MALLOC_ABSOLUTE_MAX_SIZE
-  //   (currently SIZE_T_MAX - (2 * PAGE_SIZE)) will still fail with a NULL
-  //   rather than dying (see
-  //   http://opensource.apple.com/source/Libc/Libc-583/gen/malloc.c for
-  //   details).
-  // - It is unclear whether allocations via the C++ operator new() are affected
-  //   by this (although it is likely).
-  // - This does not affect allocations from non-default zones.
-  // - It is unclear whether allocations from CoreFoundation's
-  //   kCFAllocatorDefault or +[NSObject alloc] are affected by this.
-  // Nevertheless this is better than nothing for now.
-  // TODO(avi):Do better. http://crbug.com/12673
+      << "Old allocators unexpectedly non-null";
 
   int32 major;
   int32 minor;
@@ -459,6 +526,61 @@ void EnableTerminationOnOutOfMemory() {
   if (zone_allocators_protected) {
     mprotect(reinterpret_cast<void*>(page_start), len, PROT_READ);
   }
+
+  // === C++ operator new ===
+
+  // Yes, operator new does call through to malloc, but this will catch failures
+  // that our imperfect handling of malloc cannot.
+
+  std::set_new_handler(oom_killer_new);
+
+  // === Core Foundation CFAllocators ===
+
+  // This will not catch allocation done by custom allocators, but will catch
+  // all allocation done by system-provided ones.
+
+  CHECK(!g_old_cfallocator_system_default && !g_old_cfallocator_malloc &&
+        !g_old_cfallocator_malloc_zone)
+      << "Old allocators unexpectedly non-null";
+
+  ChromeCFAllocatorRef allocator = const_cast<ChromeCFAllocatorRef>(
+      reinterpret_cast<const ChromeCFAllocator*>(kCFAllocatorSystemDefault));
+  g_old_cfallocator_system_default = allocator->context.allocate;
+  CHECK(g_old_cfallocator_system_default)
+      << "Failed to get kCFAllocatorSystemDefault allocation function.";
+  allocator->context.allocate = oom_killer_cfallocator_system_default;
+
+  allocator = const_cast<ChromeCFAllocatorRef>(
+      reinterpret_cast<const ChromeCFAllocator*>(kCFAllocatorMalloc));
+  g_old_cfallocator_malloc = allocator->context.allocate;
+  CHECK(g_old_cfallocator_malloc)
+      << "Failed to get kCFAllocatorMalloc allocation function.";
+  allocator->context.allocate = oom_killer_cfallocator_malloc;
+
+  allocator = const_cast<ChromeCFAllocatorRef>(
+      reinterpret_cast<const ChromeCFAllocator*>(kCFAllocatorMallocZone));
+  g_old_cfallocator_malloc_zone = allocator->context.allocate;
+  CHECK(g_old_cfallocator_malloc_zone)
+      << "Failed to get kCFAllocatorMallocZone allocation function.";
+  allocator->context.allocate = oom_killer_cfallocator_malloc_zone;
+
+  // === Cocoa NSObject allocation ===
+
+  // Note that both +[NSObject new] and +[NSObject alloc] call through to
+  // +[NSObject allocWithZone:].
+
+  CHECK(!g_old_allocWithZone)
+      << "Old allocator unexpectedly non-null";
+
+  Class nsobject_class = [NSObject class];
+  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 base
diff --git a/base/process_util_unittest.cc b/base/process_util_unittest.cc
index 7b9a7af..296aaa9 100644
--- a/base/process_util_unittest.cc
+++ b/base/process_util_unittest.cc
@@ -29,6 +29,9 @@
 #if defined(OS_WIN)
 #include <windows.h>
 #endif
+#if defined(OS_MACOSX)
+#include "base/process_util_unittest_mac.h"
+#endif
 
 namespace base {
 
@@ -506,8 +509,8 @@ TEST_F(ProcessUtilTest, ParseProcStatCPU) {
 
 #endif  // defined(OS_POSIX)
 
-// TODO(vandebo) make this work on Windows and Mac too.
-#if defined(OS_LINUX)
+// TODO(vandebo) make this work on Windows too.
+#if !defined(OS_WIN)
 
 #if defined(USE_TCMALLOC)
 extern "C" {
@@ -521,7 +524,8 @@ class OutOfMemoryTest : public testing::Test {
       : value_(NULL),
         // Make test size as large as possible minus a few pages so
         // that alignment or other rounding doesn't make it wrap.
-        test_size_(std::numeric_limits<std::size_t>::max() - 8192) {
+        test_size_(std::numeric_limits<std::size_t>::max() - 12 * 1024),
+        signed_test_size_(std::numeric_limits<ssize_t>::max()) {
   }
 
   virtual void SetUp() {
@@ -539,9 +543,14 @@ class OutOfMemoryTest : public testing::Test {
 
   void* value_;
   size_t test_size_;
+  ssize_t signed_test_size_;
 };
 
 TEST_F(OutOfMemoryTest, New) {
+  ASSERT_DEATH(value_ = operator new(test_size_), "");
+}
+
+TEST_F(OutOfMemoryTest, NewArray) {
   ASSERT_DEATH(value_ = new char[test_size_], "");
 }
 
@@ -561,6 +570,7 @@ TEST_F(OutOfMemoryTest, Valloc) {
   ASSERT_DEATH(value_ = valloc(test_size_), "");
 }
 
+#if defined(OS_LINUX)
 TEST_F(OutOfMemoryTest, Pvalloc) {
   ASSERT_DEATH(value_ = pvalloc(test_size_), "");
 }
@@ -585,7 +595,35 @@ TEST_F(OutOfMemoryTest, Posix_memalign) {
   // value, since we're asserting death.
   ASSERT_DEATH(EXPECT_EQ(ENOMEM, posix_memalign(&value_, 8, test_size_)), "");
 }
+#endif  // OS_LINUX
+
+#if defined(OS_MACOSX)
+
+// Since these allocation functions take a signed size, it's possible that
+// calling them just once won't be enough to exhaust memory.
+
+TEST_F(OutOfMemoryTest, CFAllocatorSystemDefault) {
+  ASSERT_DEATH(while ((value_ =
+      AllocateViaCFAllocatorSystemDefault(signed_test_size_))) {}, "");
+}
+
+TEST_F(OutOfMemoryTest, CFAllocatorMalloc) {
+  ASSERT_DEATH(while ((value_ =
+      AllocateViaCFAllocatorMalloc(signed_test_size_))) {}, "");
+}
+
+TEST_F(OutOfMemoryTest, CFAllocatorMallocZone) {
+  ASSERT_DEATH(while ((value_ =
+      AllocateViaCFAllocatorMallocZone(signed_test_size_))) {}, "");
+}
+
+TEST_F(OutOfMemoryTest, PsychoticallyBigObjCObject) {
+  ASSERT_DEATH(while ((value_ =
+      AllocatePsychoticallyBigObjCObject())) {}, "");
+}
+
+#endif  // OS_MACOSX
 
-#endif  // defined(OS_LINUX)
+#endif  // !defined(OS_WIN)
 
 }  // namespace base
diff --git a/base/process_util_unittest_mac.h b/base/process_util_unittest_mac.h
new file mode 100644
index 0000000..8402e85
--- /dev/null
+++ b/base/process_util_unittest_mac.h
@@ -0,0 +1,26 @@
+// Copyright (c) 2010 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 helpers for the process_util_unittest to allow it to fully
+// test the Mac code.
+
+#ifndef BASE_PROCESS_UTIL_UNITTEST_MAC_H_
+#define BASE_PROCESS_UTIL_UNITTEST_MAC_H_
+
+#include "base/basictypes.h"
+
+namespace base {
+
+// Allocates memory via system allocators. Alas, they take a _signed_ size for
+// allocation.
+void* AllocateViaCFAllocatorSystemDefault(int32 size);
+void* AllocateViaCFAllocatorMalloc(int32 size);
+void* AllocateViaCFAllocatorMallocZone(int32 size);
+
+// Allocates a huge Objective C object.
+void* AllocatePsychoticallyBigObjCObject();
+
+} // namespace base
+
+#endif  // BASE_PROCESS_UTIL_UNITTEST_MAC_H_
diff --git a/base/process_util_unittest_mac.mm b/base/process_util_unittest_mac.mm
new file mode 100644
index 0000000..3cd7e16
--- /dev/null
+++ b/base/process_util_unittest_mac.mm
@@ -0,0 +1,42 @@
+// Copyright (c) 2010 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.
+
+#include "base/process_util_unittest_mac.h"
+
+#import <Cocoa/Cocoa.h>
+
+@interface PsychoticallyBigObjCObject : NSObject
+{
+  // On 32 bits, the compiler limits Objective C objects to < 2G in size, and on
+  // 64 bits, the ObjC2 Runtime Reference says that sizeof(anInstance) is
+  // constrained to 32 bits. Keep it < 2G for simplicity.
+  int justUnder2Gigs_[(2U * 1024 * 1024 * 1024 - 1) / sizeof(int)];
+}
+
+@end
+
+@implementation PsychoticallyBigObjCObject
+
+@end
+
+
+namespace base {
+
+void* AllocateViaCFAllocatorSystemDefault(int32 size) {
+  return CFAllocatorAllocate(kCFAllocatorSystemDefault, size, 0);
+}
+
+void* AllocateViaCFAllocatorMalloc(int32 size) {
+  return CFAllocatorAllocate(kCFAllocatorMalloc, size, 0);
+}
+
+void* AllocateViaCFAllocatorMallocZone(int32 size) {
+  return CFAllocatorAllocate(kCFAllocatorMallocZone, size, 0);
+}
+
+void* AllocatePsychoticallyBigObjCObject() {
+  return [[PsychoticallyBigObjCObject alloc] init];
+}
+
+}  // namespace base