YT: Fix enricher, add from backtrace enricher

From backtrace enricher можно использовать для захвата бектрейсов при создании ошибок из исключений, а так же для перекладывания атрибутов из сложных исключений других библиотек.
commit_hash:76711bd3bb7dbc1e41e43f80d43340d2ce8e4df7

3 files changed, 101 insertions, 11 deletions

diff --git a/library/cpp/yt/error/error.cpp b/library/cpp/yt/error/error.cpp
index ea71a07feed..c6e806fc58d 100644
--- a/library/cpp/yt/error/error.cpp
+++ b/library/cpp/yt/error/error.cpp
@@ -30,6 +30,7 @@ void FormatValue(TStringBuilderBase* builder, TErrorCode code, TStringBuf spec)
 constexpr TStringBuf ErrorMessageTruncatedSuffix = "...<message truncated>";
 
 TError::TEnricher TError::Enricher_;
+TError::TFromExceptionEnricher TError::FromExceptionEnricher_;
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -248,6 +249,7 @@ TError::TErrorOr(const TErrorException& errorEx) noexcept
 {
     *this = errorEx.Error();
     // NB: TErrorException verifies that error not IsOK at throwing end.
+    EnrichFromException(errorEx);
 }
 
 TError::TErrorOr(const std::exception& ex)
@@ -277,8 +279,8 @@ TError::TErrorOr(const std::exception& ex)
         *this = TError(NYT::EErrorCode::Generic, TRuntimeFormat{ex.what()});
         *this <<= TErrorAttribute("exception_type", TypeName(ex));
     }
+    EnrichFromException(ex);
     YT_VERIFY(!IsOK());
-    Enrich();
 }
 
 TError::TErrorOr(std::string message, TDisableFormat)
@@ -644,14 +646,31 @@ void TError::RegisterEnricher(TEnricher enricher)
 {
     // NB: This daisy-chaining strategy is optimal when there's O(1) callbacks. Convert to a vector
     // if the number grows.
-    if (Enricher_) {
-        Enricher_ = [first = std::move(Enricher_), second = std::move(enricher)] (TError& error) {
-            first(error);
-            second(error);
-        };
-    } else {
+    if (!Enricher_) {
         Enricher_ = std::move(enricher);
+        return;
+    }
+    Enricher_ = [first = std::move(Enricher_), second = std::move(enricher)] (TError* error) {
+        first(error);
+        second(error);
+    };
+}
+
+void TError::RegisterFromExceptionEnricher(TFromExceptionEnricher enricher)
+{
+    // NB: This daisy-chaining strategy is optimal when there's O(1) callbacks. Convert to a vector
+    // if the number grows.
+    if (!FromExceptionEnricher_) {
+        FromExceptionEnricher_ = std::move(enricher);
+        return;
     }
+    FromExceptionEnricher_ = [
+        first = std::move(FromExceptionEnricher_),
+        second = std::move(enricher)
+    ] (TError* error, const std::exception& exception) {
+        first(error, exception);
+        second(error, exception);
+    };
 }
 
 TError::TErrorOr(std::unique_ptr<TImpl> impl)
@@ -668,7 +687,14 @@ void TError::MakeMutable()
 void TError::Enrich()
 {
     if (Enricher_) {
-        Enricher_(*this);
+        Enricher_(this);
+    }
+}
+
+void TError::EnrichFromException(const std::exception& exception)
+{
+    if (FromExceptionEnricher_) {
+        FromExceptionEnricher_(this, exception);
     }
 }
 
diff --git a/library/cpp/yt/error/error.h b/library/cpp/yt/error/error.h
index 92f19bc398a..4b329105bcb 100644
--- a/library/cpp/yt/error/error.h
+++ b/library/cpp/yt/error/error.h
@@ -219,13 +219,20 @@ public:
     template <CErrorNestable TValue>
     TError operator << (const std::optional<TValue>& rhs) const &;
 
-    // The |enricher| is called during TError construction and before TErrorOr<> construction. Meant
-    // to enrich the error, e.g. by setting generic attributes. The |RegisterEnricher| method is not
+    // The |enricher| is called during TError initial construction and before TErrorOr<> construction. Meant
+    // to enrich the error, e.g. by setting generic attributes. Copying TError from another TError or TErrorException
+    // doesn't call enrichers. The |RegisterEnricher| method is not
     // threadsafe and is meant to be called from single-threaded bootstrapping code. Multiple
     // enrichers are supported and will be called in order of registration.
-    using TEnricher = std::function<void(TError&)>;
+    using TEnricher = std::function<void(TError*)>;
     static void RegisterEnricher(TEnricher enricher);
 
+    // The |enricher| is called during TError every construction from std::exception (including TErrorException).
+    // The |RegisterFromExceptionEnricher| method is not threadsafe and is meant to be called from single-threaded
+    // bootstrapping code. Multiple enrichers are supported and will be called in order of registration.
+    using TFromExceptionEnricher = std::function<void(TError*, const std::exception&)>;
+    static void RegisterFromExceptionEnricher(TFromExceptionEnricher enricher);
+
 private:
     class TImpl;
     std::unique_ptr<TImpl> Impl_;
@@ -234,10 +241,12 @@ private:
 
     void MakeMutable();
     void Enrich();
+    void EnrichFromException(const std::exception& exception);
 
     friend class TErrorAttributes;
 
     static TEnricher Enricher_;
+    static TFromExceptionEnricher FromExceptionEnricher_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/error/unittests/error_ut.cpp b/library/cpp/yt/error/unittests/error_ut.cpp
index a5576fad58a..33a8e8fa965 100644
--- a/library/cpp/yt/error/unittests/error_ut.cpp
+++ b/library/cpp/yt/error/unittests/error_ut.cpp
@@ -789,6 +789,61 @@ TEST(TErrorTest, MacroStaticAnalysisBrokenFormat)
     // });
 }
 
+TEST(TErrorTest, Enrichers)
+{
+    static auto getAttribute = [] (const TError& error) {
+        return error.Attributes().Get<TString>("test_attribute", "");
+    };
+
+    {
+        static thread_local bool testEnricherEnabled = false;
+        testEnricherEnabled = true;
+
+        TError::RegisterEnricher([](TError* error) {
+            if (testEnricherEnabled) {
+                *error <<= TErrorAttribute("test_attribute", getAttribute(*error) + "X");
+            }
+        });
+
+        // Not from exception.
+        EXPECT_EQ(getAttribute(TError("E")), "X");
+        EXPECT_EQ(getAttribute(TError(NYT::EErrorCode::Generic, "E")), "X");
+
+        // std::exception.
+        EXPECT_EQ(getAttribute(TError(std::runtime_error("E"))), "X");
+
+        // Copying.
+        EXPECT_EQ(getAttribute(TError(TError(std::runtime_error("E")))), "X");
+        EXPECT_EQ(getAttribute(TError(TErrorException() <<= TError(std::runtime_error("E")))), "X");
+
+        testEnricherEnabled = false;
+    }
+
+    {
+        static thread_local bool testFromExceptionEnricherEnabled = false;
+        testFromExceptionEnricherEnabled = true;
+
+        TError::RegisterFromExceptionEnricher([](TError* error, const std::exception&) {
+            if (testFromExceptionEnricherEnabled) {
+                *error <<= TErrorAttribute("test_attribute", getAttribute(*error) + "X");
+            }
+        });
+
+        // Not from exception.
+        EXPECT_EQ(getAttribute(TError("E")), "");
+        EXPECT_EQ(getAttribute(TError(NYT::EErrorCode::Generic, "E")), "");
+
+        // From exception.
+        EXPECT_EQ(getAttribute(TError(std::runtime_error("E"))), "X");
+        EXPECT_EQ(getAttribute(TError(TError(std::runtime_error("E")))), "X");
+
+        // From exception twice.
+        EXPECT_EQ(getAttribute(TError(TErrorException() <<= TError(std::runtime_error("E")))), "XX");
+
+        testFromExceptionEnricherEnabled = false;
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace