Set tcmalloc heap limit prior to testing oom

Otherwise it can take long time to OOM on osex.

2 files changed, 18 insertions, 0 deletions

diff --git a/src/tcmalloc.cc b/src/tcmalloc.cc
index 0fa880c..9ec663e 100644
--- a/src/tcmalloc.cc
+++ b/src/tcmalloc.cc
@@ -163,6 +163,7 @@ using tcmalloc::Static;
 using tcmalloc::ThreadCache;
 
 DECLARE_double(tcmalloc_release_rate);
+DECLARE_int64(tcmalloc_heap_limit_mb);
 
 // Those common architectures are known to be safe w.r.t. aliasing function
 // with "extra" unused args to function with fewer arguments (e.g.
@@ -835,6 +836,12 @@ class TCMallocImplementation : public MallocExtension {
       return true;
     }
 
+    if (strcmp(name, "tcmalloc.heap_limit_mb") == 0) {
+      SpinLockHolder l(Static::pageheap_lock());
+      *value = FLAGS_tcmalloc_heap_limit_mb;
+      return true;
+    }
+
     return false;
   }
 
@@ -853,6 +860,12 @@ class TCMallocImplementation : public MallocExtension {
       return true;
     }
 
+    if (strcmp(name, "tcmalloc.heap_limit_mb") == 0) {
+      SpinLockHolder l(Static::pageheap_lock());
+      FLAGS_tcmalloc_heap_limit_mb = value;
+      return true;
+    }
+
     return false;
   }
 
diff --git a/src/tests/tcmalloc_unittest.cc b/src/tests/tcmalloc_unittest.cc
index f1bc078..658772f 100644
--- a/src/tests/tcmalloc_unittest.cc
+++ b/src/tests/tcmalloc_unittest.cc
@@ -1615,11 +1615,16 @@ static int RunAllTests(int argc, char** argv) {
   // Check that large allocations fail with NULL instead of crashing
 #ifndef DEBUGALLOCATION    // debug allocation takes forever for huge allocs
   fprintf(LOGSTREAM, "Testing out of memory\n");
+  size_t old_limit;
+  CHECK(MallocExtension::instance()->GetNumericProperty("tcmalloc.heap_limit_mb", &old_limit));
+  // Don't exercise more than 1 gig, no need to.
+  CHECK(MallocExtension::instance()->SetNumericProperty("tcmalloc.heap_limit_mb", 1 << 10));
   for (int s = 0; ; s += (10<<20)) {
     void* large_object = rnd.alloc(s);
     if (large_object == NULL) break;
     free(large_object);
   }
+  CHECK(MallocExtension::instance()->SetNumericProperty("tcmalloc.heap_limit_mb", old_limit));
 #endif
 
   TestHugeThreadCache();