Add data-structures for representing async stack-traces

Summary:
Defines new data-structures folly::AsyncStackFrame and folly::AsyncStackRoot
that can be used to build intrusive data-structures that represent
asynchronous call-chains.

These data-structures can be walked at either runtime from within the process
or by external tooling, such as debuggers or profilers, allowing them to capture
stack-traces that correctly attribute the calls to asynchonous callbacks to the
asynchronous caller that initiaded the operation rather than to the immediate
caller of the callback (typically an executor's event-loop).

These data-structures are initially intended to be used by folly::coro coroutines
to allow chains of these coroutines to be walked, but should also general enough
to be applied to other types of asynchronous operations, such as futures or even
asynchronous operations from other languages (eg. Rust).

Reviewed By: andriigrynenko

Differential Revision: D21118885

fbshipit-source-id: 0d6130b40fd04ef330800eae4feb93bb79f48105

folly/tracing/AsyncStack-inl.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
+
+namespace folly {
+
+inline void checkAsyncStackFrameIsActive(
+    FOLLY_MAYBE_UNUSED const folly::AsyncStackFrame& frame) noexcept {
+  assert(frame.stackRoot != nullptr);
+  assert(tryGetCurrentAsyncStackRoot() == frame.stackRoot);
+  assert(frame.stackRoot->topFrame.load(std::memory_order_relaxed) == &frame);
+}
+
+inline void activateAsyncStackFrame(
+    folly::AsyncStackRoot& root,
+    folly::AsyncStackFrame& frame) noexcept {
+  assert(tryGetCurrentAsyncStackRoot() == &root);
+  root.setTopFrame(frame);
+}
+
+inline void deactivateAsyncStackFrame(folly::AsyncStackFrame& frame) noexcept {
+  checkAsyncStackFrameIsActive(frame);
+  frame.stackRoot->topFrame.store(nullptr, std::memory_order_relaxed);
+  frame.stackRoot = nullptr;
+}
+
+inline void pushAsyncStackFrameCallerCallee(
+    folly::AsyncStackFrame& callerFrame,
+    folly::AsyncStackFrame& calleeFrame) noexcept {
+  checkAsyncStackFrameIsActive(callerFrame);
+  calleeFrame.stackRoot = callerFrame.stackRoot;
+  calleeFrame.parentFrame = &callerFrame;
+  calleeFrame.stackRoot->topFrame.store(
+      &calleeFrame, std::memory_order_release);
+
+  // Clearing out non-top-frame's stackRoot is not strictly necessary
+  // but it may help with debugging.
+  callerFrame.stackRoot = nullptr;
+}
+
+inline void popAsyncStackFrameCallee(
+    folly::AsyncStackFrame& calleeFrame) noexcept {
+  checkAsyncStackFrameIsActive(calleeFrame);
+  assert(calleeFrame.parentFrame != nullptr);
+  auto& callerFrame = *calleeFrame.parentFrame;
+  auto& stackRoot = *calleeFrame.stackRoot;
+  callerFrame.stackRoot = &stackRoot;
+  stackRoot.topFrame.store(&callerFrame, std::memory_order_release);
+
+  // Clearing out non-top-frame's stackRoot is not strictly necessary
+  // but it may help with debugging.
+  calleeFrame.stackRoot = nullptr;
+}
+
+#if FOLLY_HAS_COROUTINES
+
+template <typename Promise>
+void resumeCoroutineWithNewAsyncStackRoot(
+    std::experimental::coroutine_handle<Promise> h) noexcept {
+  resumeCoroutineWithNewAsyncStackRoot(h, h.promise().getAsyncFrame());
+}
+
+#endif
+
+inline AsyncStackFrame* AsyncStackFrame::getParentFrame() noexcept {
+  return parentFrame;
+}
+
+inline const AsyncStackFrame* AsyncStackFrame::getParentFrame() const noexcept {
+  return parentFrame;
+}
+
+inline void AsyncStackFrame::setParentFrame(AsyncStackFrame& frame) noexcept {
+  parentFrame = &frame;
+}
+
+inline AsyncStackRoot* AsyncStackFrame::getStackRoot() noexcept {
+  return stackRoot;
+}
+
+inline void AsyncStackFrame::setReturnAddress(void* p) noexcept {
+  instructionPointer = p;
+}
+
+inline void* AsyncStackFrame::getReturnAddress() const noexcept {
+  return instructionPointer;
+}
+
+inline void AsyncStackRoot::setTopFrame(AsyncStackFrame& frame) noexcept {
+  assert(this->topFrame.load(std::memory_order_relaxed) == nullptr);
+  assert(frame.stackRoot == nullptr);
+  frame.stackRoot = this;
+  this->topFrame.store(&frame, std::memory_order_release);
+}
+
+inline AsyncStackFrame* AsyncStackRoot::getTopFrame() noexcept {
+  return topFrame.load(std::memory_order_relaxed);
+}
+
+inline void AsyncStackRoot::setStackFrameContext(
+    void* framePtr,
+    void* ip) noexcept {
+  stackFramePtr = framePtr;
+  returnAddress = ip;
+}
+
+inline void* AsyncStackRoot::getStackFramePointer() const noexcept {
+  return stackFramePtr;
+}
+
+inline void* AsyncStackRoot::getReturnAddress() const noexcept {
+  return returnAddress;
+}
+
+inline const AsyncStackRoot* AsyncStackRoot::getNextRoot() const noexcept {
+  return nextRoot;
+}
+
+inline void AsyncStackRoot::setNextRoot(AsyncStackRoot* next) noexcept {
+  nextRoot = next;
+}
+
+} // namespace folly

folly/tracing/AsyncStack.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/tracing/AsyncStack.h>
+#include <folly/Likely.h>
+
+#include <glog/logging.h>
+
+#include <atomic>
+#include <cassert>
+#include <mutex>
+
+#if defined(__linux__)
+#define FOLLY_ASYNC_STACK_ROOT_USE_PTHREAD 1
+#else
+#define FOLLY_ASYNC_STACK_ROOT_USE_PTHREAD 0
+#endif
+
+#if FOLLY_ASYNC_STACK_ROOT_USE_PTHREAD
+
+#include <folly/portability/PThread.h>
+
+// Use a global TLS key variable to make it easier for profilers/debuggers
+// to lookup the current thread's AsyncStackRoot by walking the pthread
+// TLS structures.
+extern "C" {
+// Current pthread implementation has valid keys in range 0 .. 1023.
+// Initialise to some value that will be interpreted as an invalid key.
+pthread_key_t folly_async_stack_root_tls_key = 0xFFFF'FFFFu;
+}
+
+#endif // FOLLY_ASYNC_STACK_ROOT_USE_PTHREAD
+
+namespace folly {
+
+namespace {
+
+#if FOLLY_ASYNC_STACK_ROOT_USE_PTHREAD
+static pthread_once_t initialiseTlsKeyFlag = PTHREAD_ONCE_INIT;
+
+static void ensureAsyncRootTlsKeyIsInitialised() noexcept {
+  (void)pthread_once(&initialiseTlsKeyFlag, []() noexcept {
+    int result = pthread_key_create(&folly_async_stack_root_tls_key, nullptr);
+    if (UNLIKELY(result != 0)) {
+      LOG(FATAL)
+          << "Failed to initialise folly_async_stack_root_tls_key: (error:"
+          << result << ")";
+      std::terminate();
+    }
+
+    VLOG(2) << "Initialising folly_async_stack_root_tls_key at address "
+            << (void*)(&folly_async_stack_root_tls_key)
+            << " with pthread_key_t " << folly_async_stack_root_tls_key;
+  });
+}
+
+#endif
+
+struct AsyncStackRootHolder {
+#if FOLLY_ASYNC_STACK_ROOT_USE_PTHREAD
+  AsyncStackRootHolder() noexcept {
+    ensureAsyncRootTlsKeyIsInitialised();
+    const int result =
+        pthread_setspecific(folly_async_stack_root_tls_key, this);
+    if (FOLLY_UNLIKELY(result != 0)) {
+      LOG(FATAL) << "Failed to set current thread's AsyncStackRoot";
+      std::terminate();
+    }
+  }
+#endif
+
+  AsyncStackRoot* get() const noexcept {
+    return value.load(std::memory_order_relaxed);
+  }
+
+  void set(AsyncStackRoot* root) noexcept {
+    value.store(root, std::memory_order_release);
+  }
+
+  void set_relaxed(AsyncStackRoot* root) noexcept {
+    value.store(root, std::memory_order_relaxed);
+  }
+
+  std::atomic<AsyncStackRoot*> value{nullptr};
+};
+
+static thread_local AsyncStackRootHolder currentThreadAsyncStackRoot;
+
+} // namespace
+
+AsyncStackRoot* tryGetCurrentAsyncStackRoot() noexcept {
+  return currentThreadAsyncStackRoot.get();
+}
+
+AsyncStackRoot* exchangeCurrentAsyncStackRoot(
+    AsyncStackRoot* newRoot) noexcept {
+  auto* oldStackRoot = currentThreadAsyncStackRoot.get();
+  currentThreadAsyncStackRoot.set(newRoot);
+  return oldStackRoot;
+}
+
+namespace detail {
+
+ScopedAsyncStackRoot::ScopedAsyncStackRoot(
+    void* framePointer,
+    void* returnAddress) noexcept {
+  root_.setStackFrameContext(framePointer, returnAddress);
+  root_.nextRoot = currentThreadAsyncStackRoot.get();
+  currentThreadAsyncStackRoot.set(&root_);
+}
+
+ScopedAsyncStackRoot::~ScopedAsyncStackRoot() {
+  assert(currentThreadAsyncStackRoot.get() == &root_);
+  assert(root_.topFrame.load(std::memory_order_relaxed) == nullptr);
+  currentThreadAsyncStackRoot.set_relaxed(root_.nextRoot);
+}
+
+} // namespace detail
+} // namespace folly
+
+namespace folly {
+
+[[noreturn]] static void detached_task() {
+  LOG(FATAL) << "The detached_task() dummy function should never be called";
+}
+
+AsyncStackRoot& getCurrentAsyncStackRoot() noexcept {
+  auto* root = tryGetCurrentAsyncStackRoot();
+  assert(root != nullptr);
+  return *root;
+}
+
+static AsyncStackFrame makeDetachedRootFrame() noexcept {
+  AsyncStackFrame frame;
+  frame.setReturnAddress(reinterpret_cast<char*>(&detached_task) + 2);
+  return frame;
+}
+
+static AsyncStackFrame detachedRootFrame = makeDetachedRootFrame();
+
+AsyncStackFrame& getDetachedRootAsyncStackFrame() noexcept {
+  return detachedRootFrame;
+}
+
+#if FOLLY_HAS_COROUTINES
+
+FOLLY_NOINLINE void resumeCoroutineWithNewAsyncStackRoot(
+    std::experimental::coroutine_handle<> h,
+    folly::AsyncStackFrame& frame) noexcept {
+  detail::ScopedAsyncStackRoot root;
+  root.activateFrame(frame);
+  h.resume();
+}
+
+#endif // FOLLY_HAS_COROUTINES
+
+} // namespace folly

folly/tracing/test/AsyncStackTest.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/tracing/AsyncStack.h>
+
+#include <folly/portability/GTest.h>
+#include <glog/logging.h>
+
+TEST(AsyncStack, ScopedAsyncStackRoot) {
+  void* const stackFramePtr = FOLLY_ASYNC_STACK_FRAME_POINTER();
+  void* const returnAddress = FOLLY_ASYNC_STACK_RETURN_ADDRESS();
+
+  CHECK(folly::tryGetCurrentAsyncStackRoot() == nullptr);
+
+  {
+    folly::detail::ScopedAsyncStackRoot scopedRoot{stackFramePtr,
+                                                   returnAddress};
+    auto* root = folly::tryGetCurrentAsyncStackRoot();
+    CHECK_NOTNULL(root);
+
+    folly::AsyncStackFrame frame;
+    scopedRoot.activateFrame(frame);
+
+    CHECK_EQ(root, frame.getStackRoot());
+    CHECK_EQ(stackFramePtr, root->getStackFramePointer());
+    CHECK_EQ(returnAddress, root->getReturnAddress());
+    CHECK_EQ(&frame, root->getTopFrame());
+
+    folly::deactivateAsyncStackFrame(frame);
+
+    CHECK(frame.getStackRoot() == nullptr);
+    CHECK(root->getTopFrame() == nullptr);
+  }
+
+  CHECK(folly::tryGetCurrentAsyncStackRoot() == nullptr);
+}
+
+TEST(AsyncStack, PushPop) {
+  folly::detail::ScopedAsyncStackRoot scopedRoot{nullptr};
+
+  auto& root = folly::getCurrentAsyncStackRoot();
+
+  folly::AsyncStackFrame frame1;
+  folly::AsyncStackFrame frame2;
+  folly::AsyncStackFrame frame3;
+
+  scopedRoot.activateFrame(frame1);
+
+  CHECK_EQ(&frame1, root.getTopFrame());
+  CHECK_EQ(&root, frame1.getStackRoot());
+
+  folly::pushAsyncStackFrameCallerCallee(frame1, frame2);
+
+  CHECK_EQ(&frame2, root.getTopFrame());
+  CHECK_EQ(&frame1, frame2.getParentFrame());
+  CHECK_EQ(&root, frame2.getStackRoot());
+  CHECK(frame1.getStackRoot() == nullptr);
+
+  folly::pushAsyncStackFrameCallerCallee(frame2, frame3);
+
+  CHECK_EQ(&frame3, root.getTopFrame());
+  CHECK_EQ(&frame2, frame3.getParentFrame());
+  CHECK_EQ(&frame1, frame2.getParentFrame());
+  CHECK(frame1.getParentFrame() == nullptr);
+  CHECK(frame2.getStackRoot() == nullptr);
+
+  folly::deactivateAsyncStackFrame(frame3);
+
+  CHECK(root.getTopFrame() == nullptr);
+  CHECK(frame3.getStackRoot() == nullptr);
+
+  folly::activateAsyncStackFrame(root, frame3);
+
+  CHECK_EQ(&frame3, root.getTopFrame());
+  CHECK_EQ(&root, frame3.getStackRoot());
+
+  folly::popAsyncStackFrameCallee(frame3);
+
+  CHECK_EQ(&frame2, root.getTopFrame());
+  CHECK_EQ(&root, frame2.getStackRoot());
+  CHECK(frame3.getStackRoot() == nullptr);
+
+  folly::popAsyncStackFrameCallee(frame2);
+
+  CHECK_EQ(&frame1, root.getTopFrame());
+  CHECK_EQ(&root, frame1.getStackRoot());
+  CHECK(frame2.getStackRoot() == nullptr);
+
+  folly::deactivateAsyncStackFrame(frame1);
+
+  CHECK(root.getTopFrame() == nullptr);
+  CHECK(frame1.getStackRoot() == nullptr);
+}