Add a guard for per-fiber extra logging tags

commit_hash:3dfc664a8e26d9bfec6e997c848d5f66ede6e925

7 files changed, 109 insertions, 10 deletions

diff --git a/library/cpp/yt/logging/logger-inl.h b/library/cpp/yt/logging/logger-inl.h
index 4a5df49f30a..c05fc0e3825 100644
--- a/library/cpp/yt/logging/logger-inl.h
+++ b/library/cpp/yt/logging/logger-inl.h
@@ -145,6 +145,9 @@ inline bool HasMessageTags(
     if (!loggingContext.TraceLoggingTag.empty()) {
         return true;
     }
+    if (!GetThreadMessageTag().empty()) {
+        return true;
+    }
     return false;
 }
 
@@ -163,6 +166,14 @@ inline void AppendMessageTags(
             builder->AppendString(TStringBuf(", "));
         }
         builder->AppendString(traceLoggingTag);
+        printComma = true;
+    }
+    if (const auto& threadMessageTag = GetThreadMessageTag(); !threadMessageTag.empty()) {
+        if (printComma) {
+            builder->AppendString(TStringBuf(", "));
+        }
+        builder->AppendString(threadMessageTag);
+        printComma = true;
     }
 }
 
diff --git a/library/cpp/yt/logging/logger.cpp b/library/cpp/yt/logging/logger.cpp
index 7e59a3eb2b2..e0af30382d6 100644
--- a/library/cpp/yt/logging/logger.cpp
+++ b/library/cpp/yt/logging/logger.cpp
@@ -149,6 +149,20 @@ ELogLevel GetThreadMinLogLevel()
 
 ////////////////////////////////////////////////////////////////////////////////
 
+YT_DEFINE_THREAD_LOCAL(std::string, ThreadMessageTag);
+
+void SetThreadMessageTag(std::string messageTag)
+{
+    ThreadMessageTag() = std::move(messageTag);
+}
+
+std::string& GetThreadMessageTag()
+{
+    return ThreadMessageTag();
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 TLogger::TLogger(ILogManager* logManager, TStringBuf categoryName)
     : LogManager_(logManager)
     , Category_(LogManager_ ? LogManager_->GetCategory(categoryName) : nullptr)
diff --git a/library/cpp/yt/logging/logger.h b/library/cpp/yt/logging/logger.h
index a0f2591c030..d46cf10d841 100644
--- a/library/cpp/yt/logging/logger.h
+++ b/library/cpp/yt/logging/logger.h
@@ -156,6 +156,14 @@ ELogLevel GetThreadMinLogLevel();
 
 ////////////////////////////////////////////////////////////////////////////////
 
+//! Sets an extra tag for messages in current thread.
+// NB: Same as above, in fiber environment messages tags
+// are attached to a fiber.
+void SetThreadMessageTag(std::string messageTag);
+std::string& GetThreadMessageTag();
+
+////////////////////////////////////////////////////////////////////////////////
+
 static constexpr auto NullLoggerMinLevel = ELogLevel::Maximum;
 
 // Min level for non-null logger depends on whether we are in debug or release build.
diff --git a/yt/yt/core/concurrency/fiber_scheduler_thread.cpp b/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
index a266c88b33f..0f04d8c6cbd 100644
--- a/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
+++ b/yt/yt/core/concurrency/fiber_scheduler_thread.cpp
@@ -193,6 +193,13 @@ Y_FORCE_INLINE ELogLevel SwapMinLogLevel(ELogLevel minLogLevel)
     return result;
 }
 
+Y_FORCE_INLINE std::string SwapMessageTag(std::string messageTag)
+{
+    auto result = std::move(GetThreadMessageTag());
+    SetThreadMessageTag(std::move(messageTag));
+    return result;
+}
+
 Y_FORCE_INLINE TExceptionSafeContext* GetMachineContext()
 {
     return &FiberContext()->MachineContext;
@@ -806,6 +813,7 @@ protected:
         Fls_ = SwapCurrentFls(Fls_);
         TContextSwitchManager::Get()->OnIn();
         MinLogLevel_ = SwapMinLogLevel(MinLogLevel_);
+        MessageTag_ = SwapMessageTag(MessageTag_);
     }
 
     ~TBaseSwitchHandler()
@@ -813,12 +821,14 @@ protected:
         YT_VERIFY(FiberId_ == InvalidFiberId);
         YT_VERIFY(!Fls_);
         YT_VERIFY(MinLogLevel_ == ELogLevel::Minimum);
+        YT_VERIFY(MessageTag_.empty());
     }
 
 private:
     TFls* Fls_ = nullptr;
     TFiberId FiberId_ = InvalidFiberId;
     ELogLevel MinLogLevel_ = ELogLevel::Minimum;
+    std::string MessageTag_;
 };
 
 class TFiberSwitchHandler;
diff --git a/yt/yt/core/logging/log_manager.cpp b/yt/yt/core/logging/log_manager.cpp
index 2e3af02148c..d56b78333a7 100644
--- a/yt/yt/core/logging/log_manager.cpp
+++ b/yt/yt/core/logging/log_manager.cpp
@@ -1579,6 +1579,28 @@ TFiberMinLogLevelGuard::~TFiberMinLogLevelGuard()
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TFiberMessageTagGuard::TFiberMessageTagGuard(std::string messageTag)
+    : OldMessageTag_(std::move(GetThreadMessageTag()))
+{
+    SetThreadMessageTag(std::move(messageTag));
+}
+
+TFiberMessageTagGuard::TFiberMessageTagGuard(TFiberMessageTagGuard&& other)
+    : OldMessageTag_(std::move(other.OldMessageTag_))
+    , Active_(other.Active_)
+{
+    other.Active_ = false;
+}
+
+TFiberMessageTagGuard::~TFiberMessageTagGuard()
+{
+    if (Active_) {
+        SetThreadMessageTag(std::move(OldMessageTag_));
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 #ifndef _win_
 
 ILogManager* GetDefaultLogManager()
diff --git a/yt/yt/core/logging/log_manager.h b/yt/yt/core/logging/log_manager.h
index c9c1a197625..358906e3f8c 100644
--- a/yt/yt/core/logging/log_manager.h
+++ b/yt/yt/core/logging/log_manager.h
@@ -79,6 +79,7 @@ private:
 
 //! Sets the minimum logging level for all messages in current fiber.
 class TFiberMinLogLevelGuard
+    : private TMoveOnly
 {
 public:
     explicit TFiberMinLogLevelGuard(ELogLevel minLogLevel);
@@ -90,6 +91,26 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+//! Sets the minimum logging level for all messages in current fiber.
+class TFiberMessageTagGuard
+    : private TMoveOnly
+{
+public:
+    explicit TFiberMessageTagGuard(std::string messageTag);
+
+    // For use with std::optional in tests.
+    TFiberMessageTagGuard(TFiberMessageTagGuard&& other);
+    TFiberMessageTagGuard& operator=(TFiberMessageTagGuard&& other) = delete;
+
+    ~TFiberMessageTagGuard();
+
+private:
+    // NB: Keeping it non-const to allow moving from in the dtor.
+    std::string OldMessageTag_;
+    bool Active_ = true;
+};
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT::NLogging
 
 template <>
diff --git a/yt/yt/core/logging/unittests/logging_ut.cpp b/yt/yt/core/logging/unittests/logging_ut.cpp
index 4ef91f3d945..ab62740daaa 100644
--- a/yt/yt/core/logging/unittests/logging_ut.cpp
+++ b/yt/yt/core/logging/unittests/logging_ut.cpp
@@ -1323,7 +1323,7 @@ TEST_F(TLoggingTest, MessageLevelOverride)
 ////////////////////////////////////////////////////////////////////////////////
 
 class TLoggingTagsTest
-    : public ::testing::TestWithParam<std::tuple<bool, bool, bool, TString>>
+    : public ::testing::TestWithParam<std::tuple<bool, bool, bool, bool, TString>>
 { };
 
 TEST_P(TLoggingTagsTest, All)
@@ -1331,7 +1331,8 @@ TEST_P(TLoggingTagsTest, All)
     auto hasMessageTag = std::get<0>(GetParam());
     auto hasLoggerTag = std::get<1>(GetParam());
     auto hasTraceContext = std::get<2>(GetParam());
-    auto expected = std::get<3>(GetParam());
+    auto hasThreadMessageTag = std::get<3>(GetParam());
+    auto expected = std::get<4>(GetParam());
 
     auto loggingContext = NLogging::GetLoggingContext();
     if (hasTraceContext) {
@@ -1343,6 +1344,10 @@ TEST_P(TLoggingTagsTest, All)
         logger = logger.WithTag("LoggerTag");
     }
 
+    auto threadLocalTagGuard = hasThreadMessageTag
+        ? std::optional(TFiberMessageTagGuard("ThreadLocalTag"))
+        : std::nullopt;
+
     if (hasMessageTag) {
         EXPECT_EQ(
             expected,
@@ -1363,14 +1368,22 @@ TEST_P(TLoggingTagsTest, All)
 
 INSTANTIATE_TEST_SUITE_P(ValueParametrized, TLoggingTagsTest,
     ::testing::Values(
-        std::tuple(false, false, false, "Log message"),
-        std::tuple(false, false,  true, "Log message (TraceContextTag)"),
-        std::tuple(false,  true, false, "Log message (LoggerTag)"),
-        std::tuple(false,  true,  true, "Log message (LoggerTag, TraceContextTag)"),
-        std::tuple( true, false, false, "Log message (Value: 123)"),
-        std::tuple( true, false,  true, "Log message (Value: 123, TraceContextTag)"),
-        std::tuple( true,  true, false, "Log message (Value: 123, LoggerTag)"),
-        std::tuple( true,  true,  true, "Log message (Value: 123, LoggerTag, TraceContextTag)")));
+        std::tuple(false, false, false, false, "Log message"),
+        std::tuple(false, false, false,  true, "Log message (ThreadLocalTag)"),
+        std::tuple(false, false,  true, false, "Log message (TraceContextTag)"),
+        std::tuple(false, false,  true,  true, "Log message (TraceContextTag, ThreadLocalTag)"),
+        std::tuple(false,  true, false, false, "Log message (LoggerTag)"),
+        std::tuple(false,  true, false,  true, "Log message (LoggerTag, ThreadLocalTag)"),
+        std::tuple(false,  true,  true, false, "Log message (LoggerTag, TraceContextTag)"),
+        std::tuple(false,  true,  true,  true, "Log message (LoggerTag, TraceContextTag, ThreadLocalTag)"),
+        std::tuple( true, false, false, false, "Log message (Value: 123)"),
+        std::tuple( true, false, false,  true, "Log message (Value: 123, ThreadLocalTag)"),
+        std::tuple( true, false,  true, false, "Log message (Value: 123, TraceContextTag)"),
+        std::tuple( true, false,  true,  true, "Log message (Value: 123, TraceContextTag, ThreadLocalTag)"),
+        std::tuple( true,  true, false, false, "Log message (Value: 123, LoggerTag)"),
+        std::tuple( true,  true, false,  true, "Log message (Value: 123, LoggerTag, ThreadLocalTag)"),
+        std::tuple( true,  true,  true, false, "Log message (Value: 123, LoggerTag, TraceContextTag)"),
+        std::tuple( true,  true,  true,  true, "Log message (Value: 123, LoggerTag, TraceContextTag, ThreadLocalTag)")));
 
 ////////////////////////////////////////////////////////////////////////////////