Add coroutine-frame memory allocation hooks to enable identifying async...

Add coroutine-frame memory allocation hooks to enable identifying async coroutine allocations in traces

Summary:
Customises the coroutine-frame allocation for all async coroutine types to call through
the new `folly_coro_async_malloc` and `folly_coro_async_free` functions when heap-allocating
coroutine frames.

This should allow the identification of samples/traces that are calling into memory allocation
for allocating coroutine frames and should enable better quantification of the CPU cycles
and memory usage attributable to coroutine-frame allocations by filtering to samples that
contain these functions in their stack-traces.

Reviewed By: davidtgoldblatt

Differential Revision: D20929042

fbshipit-source-id: a94691377b92fab42736942ec8c7316c82a4205d

folly/experimental/coro/AsyncGenerator.h

@@ -23,6 +23,7 @@
 #include <folly/experimental/coro/Utils.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
 #include <folly/experimental/coro/WithCancellation.h>
+#include <folly/experimental/coro/detail/Malloc.h>
 #include <folly/experimental/coro/detail/ManualLifetime.h>
 
 #include <glog/logging.h>
@@ -60,6 +61,14 @@ class AsyncGeneratorPromise {
     }
   }
 
+  static void* operator new(std::size_t size) {
+    return ::folly_coro_async_malloc(size);
+  }
+
+  static void operator delete(void* ptr, std::size_t size) {
+    ::folly_coro_async_free(ptr, size);
+  }
+
   AsyncGenerator<Reference, Value> get_return_object() noexcept;
 
   std::experimental::suspend_always initial_suspend() noexcept {

folly/experimental/coro/BlockingWait.h

@@ -21,6 +21,7 @@
 #include <folly/experimental/coro/Task.h>
 #include <folly/experimental/coro/Traits.h>
 #include <folly/experimental/coro/ViaIfAsync.h>
+#include <folly/experimental/coro/detail/Malloc.h>
 #include <folly/experimental/coro/detail/Traits.h>
 #include <folly/fibers/Baton.h>
 #include <folly/synchronization/Baton.h>
@@ -56,6 +57,14 @@ class BlockingWaitPromiseBase {
  public:
   BlockingWaitPromiseBase() noexcept = default;
 
+  static void* operator new(std::size_t size) {
+    return ::folly_coro_async_malloc(size);
+  }
+
+  static void operator delete(void* ptr, std::size_t size) {
+    ::folly_coro_async_free(ptr, size);
+  }
+
   std::experimental::suspend_always initial_suspend() {
     return {};
   }

folly/experimental/coro/Task.h

@@ -35,6 +35,7 @@
 #include <folly/experimental/coro/ViaIfAsync.h>
 #include <folly/experimental/coro/WithCancellation.h>
 #include <folly/experimental/coro/detail/InlineTask.h>
+#include <folly/experimental/coro/detail/Malloc.h>
 #include <folly/experimental/coro/detail/Traits.h>
 #include <folly/futures/Future.h>
 #include <folly/io/async/Request.h>
@@ -73,6 +74,14 @@ class TaskPromiseBase {
   TaskPromiseBase() noexcept {}
 
  public:
+  static void* operator new(std::size_t size) {
+    return ::folly_coro_async_malloc(size);
+  }
+
+  static void operator delete(void* ptr, std::size_t size) {
+    ::folly_coro_async_free(ptr, size);
+  }
+
   std::experimental::suspend_always initial_suspend() noexcept {
     return {};
   }

folly/experimental/coro/ViaIfAsync.h

@@ -23,6 +23,7 @@
 #include <folly/Traits.h>
 #include <folly/experimental/coro/Traits.h>
 #include <folly/experimental/coro/WithCancellation.h>
+#include <folly/experimental/coro/detail/Malloc.h>
 #include <folly/io/async/Request.h>
 #include <folly/lang/CustomizationPoint.h>
 
@@ -44,6 +45,14 @@ class ViaCoroutine {
     promise_type(folly::Executor::KeepAlive<>& executor) noexcept
         : executor_(std::move(executor)) {}
 
+    static void* operator new(std::size_t size) {
+      return ::folly_coro_async_malloc(size);
+    }
+
+    static void operator delete(void* ptr, std::size_t size) {
+      ::folly_coro_async_free(ptr, size);
+    }
+
     ViaCoroutine get_return_object() noexcept {
       return ViaCoroutine{
           std::experimental::coroutine_handle<promise_type>::from_promise(

folly/experimental/coro/Wait.h

@@ -26,6 +26,14 @@ class Wait {
  public:
   class promise_type {
    public:
+    static void* operator new(std::size_t size) {
+      return ::folly_coro_async_malloc(size);
+    }
+
+    void operator delete(void* ptr, std::size_t size) {
+      ::folly_coro_async_malloc(ptr, size);
+    }
+
     Wait get_return_object() {
       return Wait(promise_.get_future());
     }

folly/experimental/coro/detail/BarrierTask.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <folly/experimental/coro/detail/Barrier.h>
+#include <folly/experimental/coro/detail/Malloc.h>
 
 #include <cassert>
 #include <experimental/coroutine>
@@ -43,6 +44,14 @@ class BarrierTask {
     };
 
    public:
+    static void* operator new(std::size_t size) {
+      return ::folly_coro_async_malloc(size);
+    }
+
+    static void operator delete(void* ptr, std::size_t size) {
+      ::folly_coro_async_free(ptr, size);
+    }
+
     BarrierTask get_return_object() noexcept {
       return BarrierTask{
           std::experimental::coroutine_handle<promise_type>::from_promise(

folly/experimental/coro/detail/InlineTask.h

@@ -18,6 +18,7 @@
 
 #include <folly/ScopeGuard.h>
 #include <folly/Try.h>
+#include <folly/experimental/coro/detail/Malloc.h>
 
 #include <cassert>
 #include <experimental/coroutine>
@@ -67,6 +68,14 @@ class InlineTaskPromiseBase {
   InlineTaskPromiseBase& operator=(InlineTaskPromiseBase&&) = delete;
 
  public:
+  static void* operator new(std::size_t size) {
+    return ::folly_coro_async_malloc(size);
+  }
+
+  static void operator delete(void* ptr, std::size_t size) {
+    ::folly_coro_async_free(ptr, size);
+  }
+
   std::experimental::suspend_always initial_suspend() noexcept {
     return {};
   }

folly/experimental/coro/detail/Malloc.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/experimental/coro/detail/Malloc.h>
+#include <folly/Benchmark.h>
+
+#include <new>
+
+extern "C" {
+
+FOLLY_NOINLINE
+void* folly_coro_async_malloc(std::size_t size) {
+  void* p = ::operator new(size);
+
+  // Add this after the call to prevent the compiler from
+  // turning the call to operator new() into a tailcall.
+  folly::doNotOptimizeAway(p);
+
+  return p;
+}
+
+FOLLY_NOINLINE
+void folly_coro_async_free(void* ptr, std::size_t size) {
+  ::operator delete(ptr, size);
+
+  // Add this after the call to prevent the compiler from
+  // turning the call to operator delete() into a tailcall.
+  folly::doNotOptimizeAway(size);
+}
+} // extern "C"

folly/experimental/coro/detail/Malloc.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/CPortability.h>
+
+#include <cstddef>
+
+extern "C" {
+
+// Heap allocations for coroutine-frames for all async coroutines
+// (Task, AsyncGenerator, etc.) should be funneled through these
+// functions to allow better tracing/profiling of coroutine allocations.
+FOLLY_NOINLINE
+void* folly_coro_async_malloc(std::size_t size);
+
+FOLLY_NOINLINE
+void folly_coro_async_free(void* ptr, std::size_t size);
+} // extern "C"