SmartExceptionTracer: add ability to print async stack trace if exception is thrown

Summary: This adds the ability to print the async stack trace from where the exception is thrown.

Reviewed By: yfeldblum

Differential Revision: D26989150

fbshipit-source-id: 2edfadfa617d407df947987d962617129d626040

folly/experimental/exception_tracer/SmartExceptionTracer.cpp

@@ -29,7 +29,10 @@ namespace {
 
 struct ExceptionMeta {
   void (*deleter)(void*);
+  // normal stack trace
   StackTrace trace;
+  // async stack trace
+  StackTrace traceAsync;
 };
 
 Synchronized<std::unordered_map<void*, ExceptionMeta>>& getMeta() {
@@ -80,6 +83,11 @@ void throwCallback(
       if (n != -1) {
         meta.trace.frameCount = n;
       }
+      ssize_t nAsync = symbolizer::getAsyncStackTraceSafe(
+          meta.traceAsync.addresses, kMaxFrames);
+      if (nAsync != -1) {
+        meta.traceAsync.frameCount = nAsync;
+      }
     } catch (const std::bad_alloc&) {
     }
   });
@@ -90,40 +98,86 @@ struct Initialize {
 };
 
 Initialize initialize;
-} // namespace
 
-ExceptionInfo getTrace(const std::exception_ptr& ptr) {
+// ExceptionMetaFunc takes a `const ExceptionMeta&` and
+// returns a pair of iterators to represent a range of addresses for
+// the stack frames to use.
+template <typename ExceptionMetaFunc>
+ExceptionInfo getTraceWithFunc(
+    const std::exception& ex, ExceptionMetaFunc func) {
+  ExceptionInfo info;
+  info.type = &typeid(ex);
+  getMeta().withRLock([&](auto& locked) noexcept {
+    auto* meta = get_ptr(locked, (void*)&ex);
+    // If we can't find the exception, return an empty stack trace.
+    if (!meta) {
+      return;
+    }
+
+    auto [traceBeginIt, traceEndIt] = func(*meta);
+    info.frames.assign(traceBeginIt, traceEndIt);
+  });
+  return info;
+}
+
+template <typename ExceptionMetaFunc>
+ExceptionInfo getTraceWithFunc(
+    const std::exception_ptr& ptr, ExceptionMetaFunc func) {
   try {
     // To get a pointer to the actual exception we need to rethrow the
     // exception_ptr and catch it.
     std::rethrow_exception(ptr);
   } catch (std::exception& ex) {
-    return getTrace(ex);
+    return getTraceWithFunc(ex, std::move(func));
   } catch (...) {
     return ExceptionInfo();
   }
 }
 
-ExceptionInfo getTrace(const exception_wrapper& ew) {
+template <typename ExceptionMetaFunc>
+ExceptionInfo getTraceWithFunc(
+    const exception_wrapper& ew, ExceptionMetaFunc func) {
   if (auto* ex = ew.get_exception()) {
-    return getTrace(*ex);
+    return getTraceWithFunc(*ex, std::move(func));
   }
   return ExceptionInfo();
 }
 
+auto getAsyncStackTraceItPair(const ExceptionMeta& meta) {
+  return std::make_pair(
+      meta.traceAsync.addresses,
+      meta.traceAsync.addresses + meta.traceAsync.frameCount);
+}
+
+auto getNormalStackTraceItPair(const ExceptionMeta& meta) {
+  return std::make_pair(
+      meta.trace.addresses, meta.trace.addresses + meta.trace.frameCount);
+}
+
+} // namespace
+
+ExceptionInfo getTrace(const std::exception_ptr& ptr) {
+  return getTraceWithFunc(ptr, getNormalStackTraceItPair);
+}
+
+ExceptionInfo getTrace(const exception_wrapper& ew) {
+  return getTraceWithFunc(ew, getNormalStackTraceItPair);
+}
+
 ExceptionInfo getTrace(const std::exception& ex) {
-  ExceptionInfo info;
-  info.type = &typeid(ex);
-  getMeta().withRLock([&](auto& locked) noexcept {
-    auto* meta = get_ptr(locked, (void*)&ex);
-    // If we can't find the exception, return an empty stack trace.
-    if (!meta) {
-      return;
-    }
-    info.frames.assign(
-        meta->trace.addresses, meta->trace.addresses + meta->trace.frameCount);
-  });
-  return info;
+  return getTraceWithFunc(ex, getNormalStackTraceItPair);
+}
+
+ExceptionInfo getAsyncTrace(const std::exception_ptr& ptr) {
+  return getTraceWithFunc(ptr, getAsyncStackTraceItPair);
+}
+
+ExceptionInfo getAsyncTrace(const exception_wrapper& ew) {
+  return getTraceWithFunc(ew, getAsyncStackTraceItPair);
+}
+
+ExceptionInfo getAsyncTrace(const std::exception& ex) {
+  return getTraceWithFunc(ex, getAsyncStackTraceItPair);
 }
 
 } // namespace exception_tracer

folly/experimental/exception_tracer/SmartExceptionTracer.h

@@ -28,5 +28,11 @@ ExceptionInfo getTrace(const std::exception& ex);
 
 ExceptionInfo getTrace(const exception_wrapper& ew);
 
+ExceptionInfo getAsyncTrace(const std::exception_ptr& ex);
+
+ExceptionInfo getAsyncTrace(const std::exception& ex);
+
+ExceptionInfo getAsyncTrace(const exception_wrapper& ew);
+
 } // namespace exception_tracer
 } // namespace folly

folly/experimental/exception_tracer/test/SmartExceptionTracerTest.cpp

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include <folly/experimental/coro/BlockingWait.h>
+#include <folly/experimental/coro/Task.h>
 #include <folly/experimental/exception_tracer/SmartExceptionTracer.h>
 #include <folly/portability/GTest.h>
 
@@ -88,3 +90,56 @@ TEST(SmartExceptionTracer, UnthrownException) {
       ss.str().find("Exception type: std::runtime_error (0 frames)") !=
       std::string::npos);
 }
+
+namespace {
+
+[[noreturn]] void funcD() {
+  throw std::runtime_error("test ex");
+}
+
+FOLLY_NOINLINE folly::coro::Task<void> co_funcC() {
+  funcD();
+  co_return;
+}
+
+FOLLY_NOINLINE folly::coro::Task<void> co_funcB() {
+  co_await co_funcC();
+}
+
+FOLLY_NOINLINE folly::coro::Task<void> co_funcA() {
+  co_await co_funcB();
+}
+
+} // namespace
+
+TEST(SmartExceptionTracer, AsyncStackTrace) {
+  try {
+    folly::coro::blockingWait(co_funcA());
+    FAIL() << "exception should be thrown";
+  } catch (const std::exception& ex) {
+    ASSERT_EQ(std::string("test ex"), ex.what());
+    auto info = folly::exception_tracer::getAsyncTrace(ex);
+    LOG(INFO) << "async stack trace: " << info;
+
+    std::ostringstream ss;
+    ss << info;
+
+    // Symbols should appear in this relative order
+    std::vector<std::string> symbols{
+        "funcD",
+        "co_funcC",
+        "co_funcB",
+        "co_funcA",
+    };
+    std::vector<size_t> positions;
+    for (const auto& symbol : symbols) {
+      SCOPED_TRACE(symbol);
+      auto pos = ss.str().find(symbol);
+      ASSERT_NE(std::string::npos, pos);
+      positions.push_back(pos);
+    }
+    for (size_t i = 0; i < positions.size() - 1; ++i) {
+      ASSERT_LT(positions[i], positions[i + 1]);
+    }
+  }
+}