handle large alloc reporting locklessly

Which simplifies codes a bit.

Update issue #1159

src/tcmalloc.cc

@@ -180,32 +180,6 @@ DECLARE_int64(tcmalloc_heap_limit_mb);
 #define TC_ALIAS(name) __attribute__((alias(#name)))
 #endif
 
-// For windows, the printf we use to report large allocs is
-// potentially dangerous: it could cause a malloc that would cause an
-// infinite loop.  So by default we set the threshold to a huge number
-// on windows, so this bad situation will never trigger.  You can
-// always set TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD manually if you
-// want this functionality.
-#ifdef _WIN32
-const int64 kDefaultLargeAllocReportThreshold = static_cast<int64>(1) << 62;
-#else
-const int64 kDefaultLargeAllocReportThreshold = static_cast<int64>(1) << 30;
-#endif
-DEFINE_int64(tcmalloc_large_alloc_report_threshold,
-             EnvToInt64("TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD",
-                        kDefaultLargeAllocReportThreshold),
-             "Allocations larger than this value cause a stack "
-             "trace to be dumped to stderr.  The threshold for "
-             "dumping stack traces is increased by a factor of 1.125 "
-             "every time we print a message so that the threshold "
-             "automatically goes up by a factor of ~1000 every 60 "
-             "messages.  This bounds the amount of extra logging "
-             "generated by this flag.  Default value of this flag "
-             "is very large and therefore you should see no extra "
-             "logging unless the flag is overridden.  Set to 0 to "
-             "disable reporting entirely.");
-
-
 // We already declared these functions in tcmalloc.h, but we have to
 // declare them again to give them an ATTRIBUTE_SECTION: we want to
 // put all callers of MallocHook::Invoke* in this module into
@@ -1293,14 +1267,6 @@ void* handle_oom(malloc_fn retry_fn,
   }
 }
 
-// Copy of FLAGS_tcmalloc_large_alloc_report_threshold with
-// automatic increases factored in.
-#ifdef ENABLE_LARGE_ALLOC_REPORT
-static int64_t large_alloc_threshold =
-  (kPageSize > FLAGS_tcmalloc_large_alloc_report_threshold
-   ? kPageSize : FLAGS_tcmalloc_large_alloc_report_threshold);
-#endif
-
 static void ReportLargeAlloc(Length num_pages, void* result) {
   StackTrace stack;
   stack.depth = GetStackTrace(stack.stack, tcmalloc::kMaxStackDepth, 1);
@@ -1318,17 +1284,70 @@ static void ReportLargeAlloc(Length num_pages, void* result) {
   write(STDERR_FILENO, buffer, strlen(buffer));
 }
 
-// Must be called with the page lock held.
-inline bool should_report_large(Length num_pages) {
+static bool should_report_large(Length num_pages) {
 #ifdef ENABLE_LARGE_ALLOC_REPORT
-  const int64 threshold = large_alloc_threshold;
-  if (threshold > 0 && num_pages >= (threshold >> kPageShift)) {
-    // Increase the threshold by 1/8 every time we generate a report.
-    // We cap the threshold at 8GiB to avoid overflow problems.
-    large_alloc_threshold = (threshold + threshold/8 < 8ll<<30
-                             ? threshold + threshold/8 : 8ll<<30);
-    return true;
+// For windows, the printf we use to report large allocs is
+// potentially dangerous: it could cause a malloc that would cause an
+// infinite loop.  So by default we set the threshold to a huge number
+// on windows, so this bad situation will never trigger.  You can
+// always set TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD manually if you
+// want this functionality.
+#ifdef _WIN32
+  constexpr auto kDefaultLargeAllocReportThreshold = int64_t{1} << 62;
+#else
+  constexpr auto kDefaultLargeAllocReportThreshold = int64_t{1} << 30;
+#endif
+
+  // Note, our 'reporting threshold setting' is 64-bit, but we can
+  // only afford size_t threshold variable. I.e. some 32-bit machines
+  // don't support 64-bit atomics. So some care is taken to cast etc.
+  static std::atomic<size_t> large_alloc_threshold;
+  size_t threshold = large_alloc_threshold.load(std::memory_order_relaxed);
+
+  if (threshold == 0) {
+    int64_t value = tcmalloc::commandlineflags::StringToLongLong(
+      TCMallocGetenvSafe("TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD"),
+      kDefaultLargeAllocReportThreshold);
+    if (value < 0) {
+      // Negative limit means disable reporting
+      value = std::numeric_limits<size_t>::max();
+    }
+    value = std::max<int64_t>(kPageSize, value);
+
+    if (sizeof(size_t) < sizeof(int64_t)) {
+      // On 32-bit machines size_t is narrower than int64_t. So lets
+      // make limits larger than size_t's max (i.e. overflowing 32-bit
+      // unsigned int) to be infinity.
+      value = std::min<int64_t>(value, std::numeric_limits<size_t>::max());
+    }
+
+    threshold = static_cast<ssize_t>(value);
+    large_alloc_threshold.store(threshold); // harmless to race
   }
+
+  do {
+    if (PREDICT_TRUE(num_pages < (threshold >> kPageShift))) {
+      return false;
+    }
+
+    // Increase the threshold by 1/8 every time we generate a report.
+    size_t new_threshold = threshold + threshold / 8;
+    if (new_threshold < threshold) {
+      new_threshold = std::numeric_limits<size_t>::max();
+    }
+
+    // Also make new threshold at least as big as the allocation that
+    // triggered the reporting.
+    new_threshold = std::max<size_t>(new_threshold,
+                                     num_pages << kPageShift);
+
+    if (large_alloc_threshold.compare_exchange_strong(
+          threshold, new_threshold,
+          std::memory_order_relaxed, std::memory_order_relaxed)) {
+      return true;
+    }
+  } while (true);
+
 #endif
   return false;
 }
@@ -1336,7 +1355,6 @@ inline bool should_report_large(Length num_pages) {
 // Helper for do_malloc().
 static void* do_malloc_pages(ThreadCache* heap, size_t size) {
   void* result;
-  bool report_large;
 
   Length num_pages = tcmalloc::pages(size);
 
@@ -1348,17 +1366,13 @@ static void* do_malloc_pages(ThreadCache* heap, size_t size) {
   // See https://github.com/gperftools/gperftools/issues/723
   if (heap->SampleAllocation(size)) {
     result = DoSampledAllocation(size);
-
-    SpinLockHolder h(Static::pageheap_lock());
-    report_large = should_report_large(num_pages);
   } else {
     SpinLockHolder h(Static::pageheap_lock());
     Span* span = Static::pageheap()->New(num_pages);
     result = (PREDICT_FALSE(span == NULL) ? NULL : SpanToMallocResult(span));
-    report_large = should_report_large(num_pages);
   }
 
-  if (report_large) {
+  if (should_report_large(num_pages)) {
     ReportLargeAlloc(num_pages, result);
   }
   return result;