Add thread local context information in exception (#119)

Summary:

When running with Velox, the thread local context is not captured in
exception.  For example we are missing file path in the Prestissimo query error
20241216_142035_60655_zc54y.

Reviewed By: xiaoxmeng

Differential Revision: D67569514

dwio/nimble/common/CMakeLists.txt

@@ -16,4 +16,4 @@ add_subdirectory(tests)
 add_library(nimble_common Bits.cpp Checksum.cpp FixedBitArray.cpp
                           MetricsLogger.cpp Types.cpp Varint.cpp)
 
-target_link_libraries(nimble_common velox_memory Folly::folly)
+target_link_libraries(nimble_common velox_memory velox_exception Folly::folly)

dwio/nimble/common/Exceptions.h

@@ -15,6 +15,8 @@
  */
 #pragma once
 
+#include "velox/common/base/VeloxException.h"
+
 #include <glog/logging.h>
 #include "folly/FixedString.h"
 #include "folly/experimental/symbolizer/Symbolizer.h"
@@ -92,15 +94,17 @@ class NimbleException : public std::exception {
       std::string_view failingExpression,
       std::string_view errorMessage,
       std::string_view errorCode,
-      bool retryable)
+      bool retryable,
+      velox::VeloxException::Type veloxExceptionType)
       : exceptionName_{std::move(exceptionName)},
         fileName_{fileName},
         fileLine_{fileLine},
         functionName_{functionName},
         failingExpression_{std::move(failingExpression)},
         errorMessage_{std::move(errorMessage)},
         errorCode_{std::move(errorCode)},
-        retryable_{retryable} {
+        retryable_{retryable},
+        context_{captureContextMessage(veloxExceptionType)} {
     captureStackTraceFrames();
   }
 
@@ -147,10 +151,35 @@ class NimbleException : public std::exception {
     return retryable_;
   }
 
+  const std::string& context() const {
+    return context_;
+  }
+
  protected:
   virtual void appendMessage(std::string& /* message */) const {}
 
  private:
+  static std::string captureContextMessage(
+      velox::VeloxException::Type veloxExceptionType) {
+    auto* context = &velox::getExceptionContext();
+    std::string contextMessage = context->message(veloxExceptionType);
+    while (context->parent) {
+      context = context->parent;
+      if (!context->isEssential) {
+        continue;
+      }
+      const auto message = context->message(veloxExceptionType);
+      if (message.empty()) {
+        continue;
+      }
+      if (!contextMessage.empty()) {
+        contextMessage += ' ';
+      }
+      contextMessage += message;
+    }
+    return contextMessage;
+  }
+
   void captureStackTraceFrames() {
     try {
       constexpr size_t skipFrames = 2;
@@ -196,6 +225,11 @@ class NimbleException : public std::exception {
       finalizedMessage_ += failingExpression_;
     }
 
+    if (!context_.empty()) {
+      finalizedMessage_ += "\nContext: ";
+      finalizedMessage_ += context_;
+    }
+
     if (LIKELY(!exceptionFrames_.empty())) {
       std::vector<folly::symbolizer::SymbolizedFrame> symbolizedFrames;
       symbolizedFrames.resize(exceptionFrames_.size());
@@ -226,6 +260,7 @@ class NimbleException : public std::exception {
   const std::string errorMessage_;
   const std::string errorCode_;
   const bool retryable_;
+  const std::string context_;
 
   mutable folly::once_flag once_;
   mutable std::string finalizedMessage_;
@@ -250,7 +285,8 @@ class NimbleUserError : public NimbleException {
             failingExpression,
             errorMessage,
             errorCode,
-            retryable) {}
+            retryable,
+            velox::VeloxException::Type::kUser) {}
 
   const std::string_view errorSource() const override {
     return error_source::User;
@@ -277,7 +313,8 @@ class NimbleInternalError : public NimbleException {
             failingExpression,
             errorMessage,
             errorCode,
-            retryable) {}
+            retryable,
+            velox::VeloxException::Type::kSystem) {}
 
   const std::string_view errorSource() const override {
     return error_source::Internal;
@@ -306,7 +343,8 @@ class NimbleExternalError : public NimbleException {
             failingExpression,
             errorMessage,
             errorCode,
-            retryable),
+            retryable,
+            velox::VeloxException::Type::kSystem),
         externalSource_{externalSource} {}
 
   const std::string_view errorSource() const override {

dwio/nimble/common/tests/ExceptionTests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,8 @@
 #include <gtest/gtest.h>
 #include "dwio/nimble/common/Exceptions.h"
 
-using namespace ::facebook;
+namespace facebook {
+namespace {
 
 template <typename T>
 void verifyException(
@@ -253,3 +254,46 @@ TEST(ExceptionTests, StackTraceThreads) {
       folly::exceptionStr(e).find(
           "facebook::nimble::NimbleException::NimbleException"));
 }
+
+TEST(ExceptionTests, Context) {
+  auto messageFunc = [](velox::VeloxException::Type exceptionType, void* arg) {
+    auto msg = *static_cast<const std::string*>(arg);
+    switch (exceptionType) {
+      case velox::VeloxException::Type::kUser:
+        return fmt::format("USER {}", msg);
+      case velox::VeloxException::Type::kSystem:
+        return fmt::format("SYSTEM {}", msg);
+    }
+  };
+  std::string context1Message = "1";
+  velox::ExceptionContextSetter context1({messageFunc, &context1Message, true});
+  std::string context2Message = "2";
+  velox::ExceptionContextSetter context2(
+      {messageFunc, &context2Message, false});
+  std::string context3Message = "3";
+  velox::ExceptionContextSetter context3(
+      {messageFunc, &context3Message, false});
+
+  try {
+    NIMBLE_NOT_SUPPORTED("");
+    FAIL();
+  } catch (const nimble::NimbleException& e) {
+    ASSERT_EQ(e.context(), "USER 3 USER 1");
+    ASSERT_NE(
+        std::string(e.what()).find("Context: " + e.context()),
+        std::string::npos);
+  }
+
+  try {
+    NIMBLE_UNKNOWN("");
+    FAIL();
+  } catch (const nimble::NimbleException& e) {
+    ASSERT_EQ(e.context(), "SYSTEM 3 SYSTEM 1");
+    ASSERT_NE(
+        std::string(e.what()).find("Context: " + e.context()),
+        std::string::npos);
+  }
+}
+
+} // namespace
+} // namespace facebook