YTORM-1292 Error enrichment via dependency injection

Идея такая: хочу подкладывать атрибуты в ошибки, не протаскивая их через стек (как в orm/server/objects) и не прогоняя все ошибки через специальные методы (как в orm/library/attributes). Для этого завожу fiber-local структурку с ленивым выведением строчек. Поскольку TError и TFls живут в разных мирах, пришлось сделать отдельный трамплин для совсем генеричной доработки ошибок в момент создания.

Игнат посоветовал глянуть на Codicil. Там очень похоже, но нет key/value, поэтому похитил только название. Вообще, можно унифицировать, если есть запрос.
commit_hash:203ec7abe5e8c8484e66d55f16192485db776806

Author: deep

library/cpp/yt/error/error.cpp

@@ -29,6 +29,8 @@ void FormatValue(TStringBuilderBase* builder, TErrorCode code, TStringBuf spec)
 
 constexpr TStringBuf ErrorMessageTruncatedSuffix = "...<message truncated>";
 
+TError::TEnricher TError::Enricher_;
+
 ////////////////////////////////////////////////////////////////////////////////
 
 class TError::TImpl
@@ -275,15 +277,20 @@ TError::TErrorOr(const std::exception& ex)
         *this = TError(NYT::EErrorCode::Generic, TRuntimeFormat{ex.what()});
     }
     YT_VERIFY(!IsOK());
+    Enrich();
 }
 
 TError::TErrorOr(std::string message, TDisableFormat)
     : Impl_(std::make_unique<TImpl>(std::move(message)))
-{ }
+{
+    Enrich();
+}
 
 TError::TErrorOr(TErrorCode code, std::string message, TDisableFormat)
     : Impl_(std::make_unique<TImpl>(code, std::move(message)))
-{ }
+{
+    Enrich();
+}
 
 TError& TError::operator = (const TError& other)
 {
@@ -632,6 +639,20 @@ std::optional<TError> TError::FindMatching(const THashSet<TErrorCode>& codes) co
     });
 }
 
+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 {
+        Enricher_ = std::move(enricher);
+    }
+}
+
 TError::TErrorOr(std::unique_ptr<TImpl> impl)
     : Impl_(std::move(impl))
 { }
@@ -643,6 +664,13 @@ void TError::MakeMutable()
     }
 }
 
+void TError::Enrich()
+{
+    if (Enricher_) {
+        Enricher_(*this);
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 TError& TError::operator <<= (const TErrorAttribute& attribute) &

library/cpp/yt/error/error.h

@@ -219,15 +219,25 @@ class [[nodiscard]] TErrorOr<void>
     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
+    // 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&)>;
+    static void RegisterEnricher(TEnricher enricher);
+
 private:
     class TImpl;
     std::unique_ptr<TImpl> Impl_;
 
     explicit TErrorOr(std::unique_ptr<TImpl> impl);
 
     void MakeMutable();
+    void Enrich();
 
     friend class TErrorAttributes;
+
+    static TEnricher Enricher_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////

yt/yt/core/concurrency/fls-inl.h

@@ -92,6 +92,12 @@ Y_FORCE_INLINE T* TFlsSlot<T>::GetOrCreate() const
     return static_cast<T*>(cookie);
 }
 
+template <class T>
+Y_FORCE_INLINE T* TFlsSlot<T>::MaybeGet() const
+{
+    return static_cast<T*>(GetCurrentFls()->Get(Index_));
+}
+
 template <class T>
 T* TFlsSlot<T>::Create() const
 {

yt/yt/core/concurrency/fls.h

@@ -42,6 +42,7 @@ class TFlsSlot
     T* operator->();
 
     T* GetOrCreate() const;
+    T* MaybeGet() const;
     const T* Get(const TFls& fls) const;
 
     bool IsInitialized() const;

yt/yt/core/misc/error.cpp

@@ -1,10 +1,11 @@
 #include "error.h"
 #include "serialize.h"
 
-#include <yt/yt/core/concurrency/public.h>
+#include <yt/yt/core/concurrency/fls.h>
 
 #include <yt/yt/core/net/local_address.h>
 
+#include <yt/yt/core/misc/collection_helpers.h>
 #include <yt/yt/core/misc/protobuf_helpers.h>
 
 #include <yt/yt/core/tracing/trace_context.h>
@@ -18,6 +19,8 @@
 
 #include <library/cpp/yt/global/variable.h>
 
+#include <library/cpp/yt/misc/global.h>
+
 namespace NYT {
 
 using namespace NTracing;
@@ -31,6 +34,8 @@ using NYT::ToProto;
 
 constexpr TStringBuf OriginalErrorDepthAttribute = "original_error_depth";
 
+bool TErrorCodicils::Initialized_ = false;
+
 ////////////////////////////////////////////////////////////////////////////////
 
 namespace NDetail {
@@ -656,4 +661,96 @@ void TErrorSerializer::Load(TStreamLoadContext& context, TError& error)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+static YT_DEFINE_GLOBAL(NConcurrency::TFlsSlot<TErrorCodicils>, ErrorCodicilsSlot);
+
+TErrorCodicils::TGuard::~TGuard()
+{
+    TErrorCodicils::GetOrCreate().Set(std::move(Key_), std::move(OldGetter_));
+}
+
+TErrorCodicils::TGuard::TGuard(
+    std::string key,
+    TGetter oldGetter)
+    : Key_(std::move(key))
+    , OldGetter_(std::move(oldGetter))
+{ }
+
+void TErrorCodicils::Initialize()
+{
+    if (Initialized_) {
+        // Multiple calls are OK.
+        return;
+    }
+    Initialized_ = true;
+
+    ErrorCodicilsSlot(); // Warm up the slot.
+    TError::RegisterEnricher([] (TError& error) {
+        if (auto* codicils = TErrorCodicils::MaybeGet()) {
+            codicils->Apply(error);
+        }
+    });
+}
+
+TErrorCodicils& TErrorCodicils::GetOrCreate()
+{
+    return *ErrorCodicilsSlot().GetOrCreate();
+}
+
+TErrorCodicils* TErrorCodicils::MaybeGet()
+{
+    return ErrorCodicilsSlot().MaybeGet();
+}
+
+std::optional<std::string> TErrorCodicils::MaybeEvaluate(const std::string& key)
+{
+    auto* instance = MaybeGet();
+    if (!instance) {
+        return {};
+    }
+
+    auto getter = instance->Get(key);
+    if (!getter) {
+        return {};
+    }
+
+    return getter();
+}
+
+auto TErrorCodicils::Guard(std::string key, TGetter getter) -> TGuard
+{
+    auto& instance = GetOrCreate();
+    auto [it, added] = instance.Getters_.try_emplace(key, getter);
+    TGetter oldGetter;
+    if (!added) {
+        oldGetter = std::move(it->second);
+        it->second = std::move(getter);
+    }
+    return TGuard(std::move(key), std::move(oldGetter));
+}
+
+void TErrorCodicils::Apply(TError& error) const
+{
+    for (const auto& [key, getter] : Getters_) {
+        error <<= TErrorAttribute(key, getter());
+    }
+}
+
+void TErrorCodicils::Set(std::string key, TGetter getter)
+{
+    // We could enforce Initialized_, but that could make an error condition worse at runtime.
+    // Instead, let's keep enrichment optional.
+    if (getter) {
+        Getters_.insert_or_assign(std::move(key), std::move(getter));
+    } else {
+        Getters_.erase(std::move(key));
+    }
+}
+
+auto TErrorCodicils::Get(const std::string& key) const -> TGetter
+{
+    return GetOrDefault(Getters_, std::move(key));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT

yt/yt/core/misc/error.h

@@ -13,6 +13,8 @@
 #include <library/cpp/yt/error/error.h>
 #include <library/cpp/yt/error/origin_attributes.h>
 
+#include <util/generic/noncopyable.h>
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -90,6 +92,56 @@ struct TSerializerTraits<
 
 ////////////////////////////////////////////////////////////////////////////////
 
+// A fiber-local set of attributes to enrich all errors with.
+class TErrorCodicils
+{
+public:
+    using TGetter = std::function<std::string()>;
+
+    class TGuard
+        : public TNonCopyable
+    {
+    public:
+        ~TGuard();
+
+    private:
+        friend class TErrorCodicils;
+        TGuard(std::string key, TGetter oldGetter);
+        std::string Key_;
+        TGetter OldGetter_;
+    };
+
+    // Call from single-threaded bootstrapping code. Errors will not be enriched otherwise.
+    static void Initialize();
+
+    // Gets or creates an instance for this fiber.
+    static TErrorCodicils& GetOrCreate();
+
+    // Gets the instance for this fiber if one was created previously.
+    static TErrorCodicils* MaybeGet();
+
+    // Evaluates the codicil for the key if one was set.
+    static std::optional<std::string> MaybeEvaluate(const std::string& key);
+
+    // Sets the getter and returns an RAII object to restore the previous value on desctruction.
+    static TGuard Guard(std::string key, TGetter getter);
+
+    // Adds error attributes.
+    void Apply(TError& error) const;
+
+    // Sets the getter (or erases if the getter is empty).
+    void Set(std::string key, TGetter getter);
+
+    // Gets the getter.
+    TGetter Get(const std::string& key) const;
+
+private:
+    THashMap<std::string, TGetter> Getters_;
+    static bool Initialized_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT
 
 #define ERROR_INL_H_

yt/yt/core/misc/unittests/error_ut.cpp

@@ -228,6 +228,20 @@ TEST(TErrorTest, NativeFiberId)
         .ThrowOnError();
 }
 
+TEST(TErrorTest, ErrorCodicils)
+{
+    EXPECT_FALSE(TError("ErrorCodicils").Attributes().Contains("test_attribute"));
+    {
+        auto guard = TErrorCodicils::Guard("test_attribute", [] () -> std::string {
+            return "test_value";
+        });
+        EXPECT_EQ("test_value",
+            TError("ErrorCodicils").Attributes().Get<std::string>("test_attribute"));
+        EXPECT_FALSE(TError().Attributes().Contains("test_attribute"));
+    }
+    EXPECT_FALSE(TError("ErrorCodicils").Attributes().Contains("test_attribute"));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace

yt/yt/core/test_framework/framework.cpp

@@ -156,6 +156,8 @@ Y_TEST_HOOK_BEFORE_RUN(GTEST_YT_SETUP)
     NYT::TSignalRegistry::Get()->PushDefaultSignalHandler(NYT::AllCrashSignals);
 #endif
 
+    NYT::TErrorCodicils::Initialize();
+
     auto config = NYT::NLogging::TLogManagerConfig::CreateYTServer("unittester", GetOutputPath().GetPath());
     NYT::NLogging::TLogManager::Get()->Configure(config);
     NYT::NLogging::TLogManager::Get()->EnableReopenOnSighup();