YT-19731: Whitelist now prevents dropping inner errors with whitelisted attributes

No tests for now
9e6aa6815b8d892d1e76281e95f5d044196801e1

5 files changed, 222 insertions, 14 deletions

diff --git a/library/cpp/yt/misc/optional.h b/library/cpp/yt/misc/optional.h
index cfae5c36d5..d5e3c07fbe 100644
--- a/library/cpp/yt/misc/optional.h
+++ b/library/cpp/yt/misc/optional.h
@@ -13,6 +13,16 @@ struct TOptionalTraits
 {
     using TOptional = std::optional<T>;
     using TValue = T;
+
+    static constexpr bool HasValue(const TOptional& opt)
+    {
+        return opt.has_value();
+    }
+
+    static constexpr TOptional Empty()
+    {
+        return std::nullopt;
+    }
 };
 
 template <class T>
@@ -20,6 +30,16 @@ struct TOptionalTraits<std::optional<T>>
 {
     using TOptional = std::optional<T>;
     using TValue = T;
+
+    static constexpr bool HasValue(const TOptional& opt)
+    {
+        return opt.has_value();
+    }
+
+    static constexpr TOptional Empty()
+    {
+        return std::nullopt;
+    }
 };
 
 template <class T>
@@ -27,6 +47,16 @@ struct TOptionalTraits<T*>
 {
     using TOptional = T*;
     using TValue = T*;
+
+    static constexpr bool HasValue(const TOptional& opt)
+    {
+        return opt != nullptr;
+    }
+
+    static constexpr TOptional Empty()
+    {
+        return nullptr;
+    }
 };
 
 template <class T>
@@ -34,6 +64,16 @@ struct TOptionalTraits<TIntrusivePtr<T>>
 {
     using TOptional = TIntrusivePtr<T>;
     using TValue = TIntrusivePtr<T>;
+
+    static bool HasValue(const TOptional& opt)
+    {
+        return opt.Get() != nullptr;
+    }
+
+    static constexpr TOptional Empty()
+    {
+        return TIntrusivePtr<T>{};
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/core/misc/error.cpp b/yt/yt/core/misc/error.cpp
index 7784292c70..e84b9da94b 100644
--- a/yt/yt/core/misc/error.cpp
+++ b/yt/yt/core/misc/error.cpp
@@ -355,6 +355,51 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+bool IsWhitelisted(const TError& error, const THashSet<TStringBuf>& attributeWhitelist)
+{
+    for (const auto& key : error.Attributes().ListKeys()) {
+        if (attributeWhitelist.contains(key)) {
+            return true;
+        }
+    }
+
+    for (const auto& innerError : error.InnerErrors()) {
+        if (IsWhitelisted(innerError, attributeWhitelist)) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
+//! Returns vector which consists of objects from errors such that:
+//! if N is the number of objects in errors s.t. IsWhitelisted is true
+//! then first N objects of returned vector are the ones for which IsWhitelisted is true
+//! followed by std::max(0, maxInnerErrorCount - N - 1) remaining objects
+//! finally followed by errors.back().
+std::vector<TError>& ApplyWhitelist(std::vector<TError>& errors, const THashSet<TStringBuf>& attributeWhitelist, int maxInnerErrorCount)
+{
+    if (std::ssize(errors) < std::max(2, maxInnerErrorCount)) {
+        return errors;
+    }
+
+    auto firstNotWhitelisted = std::partition(
+        errors.begin(),
+        std::prev(errors.end()),
+        [&attributeWhitelist] (const TError& error) {
+            return IsWhitelisted(error, attributeWhitelist);
+        });
+
+    int lastErrorOffset = std::max<int>(firstNotWhitelisted - errors.begin(), maxInnerErrorCount - 1);
+
+    *(errors.begin() + lastErrorOffset) = std::move(errors.back());
+    errors.resize(lastErrorOffset + 1);
+
+    return errors;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 TError::TErrorOr() = default;
 
 TError::~TErrorOr() = default;
@@ -629,7 +674,10 @@ std::vector<TError>* TError::MutableInnerErrors()
 
 const TString InnerErrorsTruncatedKey("inner_errors_truncated");
 
-TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit, const THashSet<TStringBuf>& attributeWhitelist) const &
+TError TError::Truncate(
+    int maxInnerErrorCount,
+    i64 stringLimit,
+    const THashSet<TStringBuf>& attributeWhitelist) const &
 {
     if (!Impl_) {
         return TError();
@@ -666,22 +714,37 @@ TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit, const THashSet<
     }
     result->CopyBuiltinAttributesFrom(*Impl_);
 
-    if (std::ssize(InnerErrors()) <= maxInnerErrorCount) {
-        for (const auto& innerError : InnerErrors()) {
-            result->MutableInnerErrors()->push_back(truncateInnerError(innerError));
+    const auto& innerErrors = InnerErrors();
+    auto& copiedInnerErrors = *result->MutableInnerErrors();
+
+    if (std::ssize(innerErrors) <= maxInnerErrorCount) {
+        for (const auto& innerError : innerErrors) {
+            copiedInnerErrors.push_back(truncateInnerError(innerError));
         }
     } else {
         result->MutableAttributes()->Set(InnerErrorsTruncatedKey, true);
-        for (int i = 0; i + 1 < maxInnerErrorCount; ++i) {
-            result->MutableInnerErrors()->push_back(truncateInnerError(InnerErrors()[i]));
+
+        // NB(arkady-e1ppa): We want to always keep the last inner error,
+        // so we make room for it and do not check if it is whitelisted.
+        for (int idx = 0; idx < std::ssize(innerErrors) - 1; ++idx) {
+            const auto& innerError = innerErrors[idx];
+            if (
+                IsWhitelisted(innerError, attributeWhitelist) ||
+                std::ssize(copiedInnerErrors) < maxInnerErrorCount - 1)
+            {
+                copiedInnerErrors.push_back(truncateInnerError(innerError));
+            }
         }
-        result->MutableInnerErrors()->push_back(truncateInnerError(InnerErrors().back()));
+        copiedInnerErrors.push_back(truncateInnerError(innerErrors.back()));
     }
 
     return TError(std::move(result));
 }
 
-TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit, const THashSet<TStringBuf>& attributeWhitelist) &&
+TError TError::Truncate(
+    int maxInnerErrorCount,
+    i64 stringLimit,
+    const THashSet<TStringBuf>& attributeWhitelist) &&
 {
     if (!Impl_) {
         return TError();
@@ -711,14 +774,12 @@ TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit, const THashSet<
             truncateInnerError(innerError);
         }
     } else {
-        auto& innerErrors = *MutableInnerErrors();
+        auto& innerErrors = ApplyWhitelist(*MutableInnerErrors(), attributeWhitelist, maxInnerErrorCount);
         MutableAttributes()->Set(InnerErrorsTruncatedKey, true);
-        for (int i = 0; i + 1 < maxInnerErrorCount; ++i) {
-            truncateInnerError(innerErrors[i]);
+
+        for (auto& innerError : innerErrors) {
+            truncateInnerError(innerError);
         }
-        truncateInnerError(innerErrors.back());
-        innerErrors[maxInnerErrorCount - 1] = std::move(innerErrors.back());
-        innerErrors.resize(maxInnerErrorCount);
     }
 
     return std::move(*this);
diff --git a/yt/yt/core/misc/error_helpers-inl.h b/yt/yt/core/misc/error_helpers-inl.h
new file mode 100644
index 0000000000..d15c5f6f93
--- /dev/null
+++ b/yt/yt/core/misc/error_helpers-inl.h
@@ -0,0 +1,37 @@
+#ifndef ERROR_HELPERS_INL_H_
+#error "Direct inclusion of this file is not allowed, include error_helpers.h"
+// For the sake of sane code completion.
+#include "error_helpers.h"
+#endif
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+typename TOptionalTraits<T>::TOptional FindAttribute(const TError& error, TStringBuf key)
+{
+    return error.Attributes().Find<T>(key);
+}
+
+template <class T>
+typename TOptionalTraits<T>::TOptional FindAttributeRecursive(const TError& error, TStringBuf key)
+{
+    auto attr = FindAttribute<T>(error, key);
+    if (TOptionalTraits<T>::HasValue(attr)) {
+        return attr;
+    }
+
+    for (const auto& inner : error.InnerErrors()) {
+        attr = FindAttribute<T>(inner, key);
+        if (TOptionalTraits<T>::HasValue(attr)) {
+            return attr;
+        }
+    }
+
+    return TOptionalTraits<T>::Empty();
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/yt/yt/core/misc/error_helpers.h b/yt/yt/core/misc/error_helpers.h
new file mode 100644
index 0000000000..5eda19e4bf
--- /dev/null
+++ b/yt/yt/core/misc/error_helpers.h
@@ -0,0 +1,27 @@
+#pragma once
+
+#include "error.h"
+
+#include <yt/yt/core/ytree/attributes.h>
+
+#include <library/cpp/yt/misc/optional.h>
+
+namespace NYT {
+
+// NB: Methods below are listed in a separate file and not in error.h to prevent
+// circular includes cause by the fact that attributes include error.
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+typename TOptionalTraits<T>::TOptional FindAttribute(const TError& error, TStringBuf key);
+
+template <class T>
+typename TOptionalTraits<T>::TOptional FindAttributeRecursive(const TError& error, TStringBuf key);
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define ERROR_HELPERS_INL_H_
+#include "error_helpers-inl.h"
+#undef ERROR_HELPERS_INL_H_
diff --git a/yt/yt/core/misc/unittests/error_ut.cpp b/yt/yt/core/misc/unittests/error_ut.cpp
index 6fc88ef8ab..9b33cb935c 100644
--- a/yt/yt/core/misc/unittests/error_ut.cpp
+++ b/yt/yt/core/misc/unittests/error_ut.cpp
@@ -2,6 +2,7 @@
 #include <yt/yt/core/test_framework/framework.h>
 
 #include <yt/yt/core/misc/error.h>
+#include <yt/yt/core/misc/error_helpers.h>
 
 #include <yt/yt/core/yson/string.h>
 
@@ -478,6 +479,25 @@ TEST(TErrorTest, FormatCtor)
     EXPECT_EQ("Some error hello", TError("Some error %v", "hello").GetMessage());
 }
 
+TEST(TErrorTest, FindRecursive)
+{
+    auto inner = TError("Inner")
+        << TErrorAttribute("inner_attr", 42);
+    auto error = TError("Error")
+        << inner
+        << TErrorAttribute("attr", 8);
+
+    auto attr = FindAttribute<int>(error, "attr");
+    EXPECT_TRUE(attr);
+    EXPECT_EQ(*attr, 8);
+
+    EXPECT_FALSE(FindAttribute<int>(error, "inner_attr"));
+
+    auto innerAttr = FindAttributeRecursive<int>(error, "inner_attr");
+    EXPECT_TRUE(innerAttr);
+    EXPECT_EQ(*innerAttr, 42);
+}
+
 TEST(TErrorTest, TruncateSimple)
 {
     auto error = TError("Some error")
@@ -653,6 +673,29 @@ TEST(TErrorTest, TruncateWhitelistInnerErrorsRValue)
     EXPECT_EQ("Some long long attr", truncatedInnerError.Attributes().Get<TString>("attr2"));
 }
 
+TEST(TErrorTest, TruncateWhitelistSaveInnerError)
+{
+    auto genericInner = TError("GenericInner");
+    auto whitelistedInner = TError("Inner")
+        << TErrorAttribute("whitelisted_key", 42);
+
+    auto error = TError("Error")
+        << (genericInner << TErrorAttribute("foo", "bar"))
+        << whitelistedInner
+        << genericInner;
+
+    error = std::move(error).Truncate(1, 20, {
+        "whitelisted_key"
+    });
+    EXPECT_TRUE(!error.IsOK());
+    EXPECT_EQ(error.InnerErrors().size(), 2u);
+    EXPECT_EQ(error.InnerErrors()[0], whitelistedInner);
+    EXPECT_EQ(error.InnerErrors()[1], genericInner);
+
+    EXPECT_TRUE(FindAttributeRecursive<int>(error, "whitelisted_key"));
+    EXPECT_FALSE(FindAttributeRecursive<int>(error, "foo"));
+}
+
 TEST(TErrorTest, YTExceptionToError)
 {
     try {