1 files changed, 23 insertions, 5 deletions
diff --git a/src/tests/system-alloc_unittest.cc b/src/tests/system-alloc_unittest.cc
index da76285..c006425 100644
--- a/src/tests/system-alloc_unittest.cc
+++ b/src/tests/system-alloc_unittest.cc
@@ -31,6 +31,7 @@
 // Author: Arun Sharma
 
 #include "config_for_unittests.h"
+#include "system-alloc.h"
 #include <stdio.h>
 #if defined HAVE_STDINT_H
 #include <stdint.h>             // to get uintptr_t
@@ -39,9 +40,9 @@
 #endif
 #include <sys/types.h>
 #include <algorithm>
-#include "base/logging.h"
-#include "common.h"
-#include "system-alloc.h"
+#include "base/logging.h"               // for Check_GEImpl, Check_LTImpl, etc
+#include <google/malloc_extension.h>    // for MallocExtension::instance
+#include "common.h"                     // for kAddressBits
 
 class ArraySysAllocator : public SysAllocator {
 public:
@@ -55,6 +56,11 @@ public:
 
   void* Alloc(size_t size, size_t *actual_size, size_t alignment) {
     invoked_ = true;
+
+    if (size > kArraySize) {
+      return NULL;
+    }
+
     void *result = &array_[ptr_];
     uintptr_t ptr = reinterpret_cast<uintptr_t>(result);
 
@@ -75,8 +81,9 @@ public:
     return reinterpret_cast<void *>(ptr);
   }
 
-  void DumpStats(TCMalloc_Printer* printer) {
+  void DumpStats() {
   }
+  void FlagsInitialized() {}
 
 private:
   static const int kArraySize = 8 * 1024 * 1024;
@@ -89,7 +96,7 @@ const int ArraySysAllocator::kArraySize;
 ArraySysAllocator a;
 
 static void TestBasicInvoked() {
-  RegisterSystemAllocator(&a, 0);
+  MallocExtension::instance()->SetSystemAllocator(&a);
 
   // An allocation size that is likely to trigger the system allocator.
   // XXX: this is implementation specific.
@@ -112,8 +119,19 @@ TEST(AddressBits, CpuVirtualBits) {
 }
 #endif
 
+static void TestBasicRetryFailTest() {
+  // Check with the allocator still works after a failed allocation.
+  void* p = malloc(1ULL << 50);  // Asking for 1P ram
+  CHECK(p == NULL);
+
+  char* q = new char[1024];
+  CHECK(q != NULL);
+  delete [] q;
+}
+
 int main(int argc, char** argv) {
   TestBasicInvoked();
+  TestBasicRetryFailTest();
 
   printf("PASS\n");
   return 0;