YT-21566: Access thread local variables via noinline functions

970c33b44a7bd166b2716d86d3d2053dcaf05d7d

12 files changed, 83 insertions, 111 deletions

diff --git a/library/cpp/yt/cpu_clock/clock.cpp b/library/cpp/yt/cpu_clock/clock.cpp
index 0ed65e3b92..22f665e1d3 100644
--- a/library/cpp/yt/cpu_clock/clock.cpp
+++ b/library/cpp/yt/cpu_clock/clock.cpp
@@ -31,18 +31,16 @@ double GetTicksToMicroseconds()
     return TicksToMicroseconds;
 }
 
-TCalibrationState GetCalibrationState(TCpuInstant cpuInstant)
+YT_PREVENT_TLS_CACHING TCalibrationState GetCalibrationState(TCpuInstant cpuInstant)
 {
-    YT_THREAD_LOCAL(TCalibrationState) State;
+    thread_local TCalibrationState State;
 
-    auto& state = GetTlsRef(State);
-
-    if (state.CpuInstant + CalibrationCpuPeriod < cpuInstant) {
-        state.CpuInstant = cpuInstant;
-        state.Instant = TInstant::Now();
+    if (State.CpuInstant + CalibrationCpuPeriod < cpuInstant) {
+        State.CpuInstant = cpuInstant;
+        State.Instant = TInstant::Now();
     }
 
-    return state;
+    return State;
 }
 
 TCalibrationState GetCalibrationState()
diff --git a/library/cpp/yt/logging/logger-inl.h b/library/cpp/yt/logging/logger-inl.h
index 5e2b6795dc..d076b3690e 100644
--- a/library/cpp/yt/logging/logger-inl.h
+++ b/library/cpp/yt/logging/logger-inl.h
@@ -90,20 +90,8 @@ protected:
     void DoReserve(size_t newLength) override;
 
 private:
-    struct TPerThreadCache
-    {
-        ~TPerThreadCache();
-
-        TSharedMutableRef Chunk;
-        size_t ChunkOffset = 0;
-    };
-
     TSharedMutableRef Buffer_;
 
-    static YT_THREAD_LOCAL(TPerThreadCache*) Cache_;
-    static YT_THREAD_LOCAL(bool) CacheDestroyed_;
-    static TPerThreadCache* GetCache();
-
     static constexpr size_t ChunkSize = 128_KB - 64;
 };
 
diff --git a/library/cpp/yt/logging/logger.cpp b/library/cpp/yt/logging/logger.cpp
index c11457f836..58add38429 100644
--- a/library/cpp/yt/logging/logger.cpp
+++ b/library/cpp/yt/logging/logger.cpp
@@ -34,15 +34,45 @@ TSharedRef TMessageStringBuilder::Flush()
     return Buffer_.Slice(0, GetLength());
 }
 
-void TMessageStringBuilder::DisablePerThreadCache()
+void TMessageStringBuilder::DoReset()
 {
-    Cache_ = nullptr;
-    CacheDestroyed_ = true;
+    Buffer_.Reset();
 }
 
-void TMessageStringBuilder::DoReset()
+struct TPerThreadCache;
+
+YT_DEFINE_THREAD_LOCAL(TPerThreadCache*, Cache);
+YT_DEFINE_THREAD_LOCAL(bool, CacheDestroyed);
+
+struct TPerThreadCache
 {
-    Buffer_.Reset();
+    TSharedMutableRef Chunk;
+    size_t ChunkOffset = 0;
+
+    ~TPerThreadCache()
+    {
+        TMessageStringBuilder::DisablePerThreadCache();
+    }
+
+    static YT_PREVENT_TLS_CACHING TPerThreadCache* GetCache()
+    {
+        auto& cache = Cache();
+        if (Y_LIKELY(cache)) {
+            return cache;
+        }
+        if (CacheDestroyed()) {
+            return nullptr;
+        }
+        static thread_local TPerThreadCache CacheData;
+        cache = &CacheData;
+        return cache;
+    }
+};
+
+void TMessageStringBuilder::DisablePerThreadCache()
+{
+    Cache() = nullptr;
+    CacheDestroyed() = true;
 }
 
 void TMessageStringBuilder::DoReserve(size_t newCapacity)
@@ -53,7 +83,7 @@ void TMessageStringBuilder::DoReserve(size_t newCapacity)
     auto newChunkSize = std::max(ChunkSize, newCapacity);
     // Hold the old buffer until the data is copied.
     auto oldBuffer = std::move(Buffer_);
-    auto* cache = GetCache();
+    auto* cache = TPerThreadCache::GetCache();
     if (Y_LIKELY(cache)) {
         auto oldCapacity = End_ - Begin_;
         auto deltaCapacity = newCapacity - oldCapacity;
@@ -85,27 +115,6 @@ void TMessageStringBuilder::DoReserve(size_t newCapacity)
     End_ = Begin_ + newCapacity;
 }
 
-TMessageStringBuilder::TPerThreadCache* TMessageStringBuilder::GetCache()
-{
-    if (Y_LIKELY(Cache_)) {
-        return Cache_;
-    }
-    if (CacheDestroyed_) {
-        return nullptr;
-    }
-    static YT_THREAD_LOCAL(TPerThreadCache) Cache;
-    Cache_ = &GetTlsRef(Cache);
-    return Cache_;
-}
-
-TMessageStringBuilder::TPerThreadCache::~TPerThreadCache()
-{
-    TMessageStringBuilder::DisablePerThreadCache();
-}
-
-YT_THREAD_LOCAL(TMessageStringBuilder::TPerThreadCache*) TMessageStringBuilder::Cache_;
-YT_THREAD_LOCAL(bool) TMessageStringBuilder::CacheDestroyed_;
-
 } // namespace NDetail
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -126,16 +135,16 @@ Y_WEAK ILogManager* GetDefaultLogManager()
 
 ////////////////////////////////////////////////////////////////////////////////
 
-YT_THREAD_LOCAL(ELogLevel) ThreadMinLogLevel = ELogLevel::Minimum;
+YT_DEFINE_THREAD_LOCAL(ELogLevel, ThreadMinLogLevel, ELogLevel::Minimum);
 
 void SetThreadMinLogLevel(ELogLevel minLogLevel)
 {
-    ThreadMinLogLevel = minLogLevel;
+    ThreadMinLogLevel() = minLogLevel;
 }
 
 ELogLevel GetThreadMinLogLevel()
 {
-    return ThreadMinLogLevel;
+    return ThreadMinLogLevel();
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -173,7 +182,7 @@ bool TLogger::IsLevelEnabledHeavy(ELogLevel level) const
 
     return
         level >= Category_->MinPlainTextLevel &&
-        level >= ThreadMinLogLevel;
+        level >= ThreadMinLogLevel();
 }
 
 bool TLogger::GetAbortOnAlert() const
diff --git a/library/cpp/yt/logging/logger.h b/library/cpp/yt/logging/logger.h
index 082afcddd8..0913b9ea08 100644
--- a/library/cpp/yt/logging/logger.h
+++ b/library/cpp/yt/logging/logger.h
@@ -324,8 +324,8 @@ void LogStructuredEvent(
             break; \
         } \
         \
-        static YT_THREAD_LOCAL(i64) localByteCounter__; \
-        static YT_THREAD_LOCAL(ui8) localMessageCounter__; \
+        static thread_local i64 localByteCounter__; \
+        static thread_local ui8 localMessageCounter__; \
         \
         localByteCounter__ += message__.MessageRef.Size(); \
         if (Y_UNLIKELY(++localMessageCounter__ == 0)) { \
diff --git a/library/cpp/yt/memory/memory_tag.cpp b/library/cpp/yt/memory/memory_tag.cpp
index ebb0b43370..ee6a1a4ad0 100644
--- a/library/cpp/yt/memory/memory_tag.cpp
+++ b/library/cpp/yt/memory/memory_tag.cpp
@@ -8,16 +8,16 @@ namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-YT_THREAD_LOCAL(TMemoryTag) CurrentMemoryTag;
+YT_DEFINE_THREAD_LOCAL(TMemoryTag, CurrentMemoryTag);
 
 Y_WEAK TMemoryTag GetCurrentMemoryTag()
 {
-    return CurrentMemoryTag;
+    return CurrentMemoryTag();
 }
 
 Y_WEAK void SetCurrentMemoryTag(TMemoryTag tag)
 {
-    CurrentMemoryTag = tag;
+    CurrentMemoryTag() = tag;
 }
 
 Y_WEAK size_t GetMemoryUsageForTag(TMemoryTag /*tag*/)
diff --git a/library/cpp/yt/memory/tls_scratch-inl.h b/library/cpp/yt/memory/tls_scratch-inl.h
index 32c26d67c0..6984fe755b 100644
--- a/library/cpp/yt/memory/tls_scratch-inl.h
+++ b/library/cpp/yt/memory/tls_scratch-inl.h
@@ -11,14 +11,13 @@ namespace NYT {
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class T>
-TMutableRange<T> GetTlsScratchBuffer(size_t size)
+YT_PREVENT_TLS_CACHING TMutableRange<T> GetTlsScratchBuffer(size_t size)
 {
     // This is a workround for std::vector<bool>.
     using TBoxed = std::array<T, 1>;
-    YT_THREAD_LOCAL(std::vector<TBoxed>) tlsVector;
-    auto& vector = GetTlsRef(tlsVector);
-    vector.reserve(size);
-    auto range = TMutableRange(reinterpret_cast<T*>(vector.data()), size);
+    thread_local std::vector<TBoxed> tlsVector;
+    tlsVector.reserve(size);
+    auto range = TMutableRange(reinterpret_cast<T*>(tlsVector.data()), size);
     std::fill(range.begin(), range.end(), T());
     return range;
 }
diff --git a/library/cpp/yt/misc/thread_name.cpp b/library/cpp/yt/misc/thread_name.cpp
index d59043db51..bb1737b08a 100644
--- a/library/cpp/yt/misc/thread_name.cpp
+++ b/library/cpp/yt/misc/thread_name.cpp
@@ -29,19 +29,18 @@ TThreadName::TThreadName(const TString& name)
 ////////////////////////////////////////////////////////////////////////////////
 
 // This function uses cached TThread::CurrentThreadName() result
-TThreadName GetCurrentThreadName()
+YT_PREVENT_TLS_CACHING TThreadName GetCurrentThreadName()
 {
-    static YT_THREAD_LOCAL(TThreadName) ThreadName;
-    auto& threadName = GetTlsRef(ThreadName);
+    static thread_local TThreadName ThreadName;
 
-    if (threadName.Length == 0) {
+    if (ThreadName.Length == 0) {
         if (auto name = TThread::CurrentThreadName()) {
             auto length = std::min<int>(TThreadName::BufferCapacity - 1, name.length());
-            threadName.Length = length;
-            ::memcpy(threadName.Buffer.data(), name.data(), length);
+            ThreadName.Length = length;
+            ::memcpy(ThreadName.Buffer.data(), name.data(), length);
         }
     }
-    return threadName;
+    return ThreadName;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/misc/tls-inl.h b/library/cpp/yt/misc/tls-inl.h
deleted file mode 100644
index 51ef91517f..0000000000
--- a/library/cpp/yt/misc/tls-inl.h
+++ /dev/null
@@ -1,18 +0,0 @@
-#ifndef TLS_INL_H_
-#error "Direct inclusion of this file is not allowed, include tls.h"
-// For the sake of sane code completion.
-#include "tls.h"
-#endif
-
-namespace NYT {
-
-////////////////////////////////////////////////////////////////////////////////
-
-template <class T>
-T& GetTlsRef(volatile T& arg)
-{
-    return const_cast<T&>(arg);
-}
-////////////////////////////////////////////////////////////////////////////////
-
-} // namespace NYT
diff --git a/library/cpp/yt/misc/tls.h b/library/cpp/yt/misc/tls.h
index d1bda8ac68..a0656dedc1 100644
--- a/library/cpp/yt/misc/tls.h
+++ b/library/cpp/yt/misc/tls.h
@@ -2,21 +2,17 @@
 
 #include <util/system/compiler.h>
 
-// Workaround for fiber (un)friendly TLS.
-// Volatile qualifier prevents caching access to thread local variables.
-#define YT_THREAD_LOCAL(...) thread_local __VA_ARGS__ volatile
-
-namespace NYT {
-
-////////////////////////////////////////////////////////////////////////////////
+#define YT_PREVENT_TLS_CACHING Y_NO_INLINE
 
-template <class T>
-Y_FORCE_INLINE T& GetTlsRef(volatile T& arg);
-
-////////////////////////////////////////////////////////////////////////////////
+// Workaround for fiber (un)friendly TLS.
+#define YT_DECLARE_THREAD_LOCAL(type, name) \
+    type& name();
 
-} // namespace NYT
+#define YT_DEFINE_THREAD_LOCAL(type, name, ...) \
+    thread_local type name##Data { __VA_ARGS__ }; \
+    Y_NO_INLINE type& name() \
+    { \
+        asm volatile(""); \
+        return name##Data; \
+    }
 
-#define TLS_INL_H_
-#include "tls-inl.h"
-#undef TLS_INL_H_
diff --git a/library/cpp/yt/system/thread_id-inl.h b/library/cpp/yt/system/thread_id-inl.h
index 73434145e3..916d70436e 100644
--- a/library/cpp/yt/system/thread_id-inl.h
+++ b/library/cpp/yt/system/thread_id-inl.h
@@ -14,15 +14,16 @@ namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-extern YT_THREAD_LOCAL(TSequentialThreadId) CachedSequentialThreadId;
+YT_DECLARE_THREAD_LOCAL(TSequentialThreadId, CachedSequentialThreadId);
 extern std::atomic<TSequentialThreadId> SequentialThreadIdGenerator;
 
 inline TSequentialThreadId GetSequentialThreadId()
 {
-    if (Y_UNLIKELY(CachedSequentialThreadId == InvalidSequentialThreadId)) {
-        CachedSequentialThreadId = ++SequentialThreadIdGenerator;
+    auto& cachedSequentialThreadId = CachedSequentialThreadId();
+    if (Y_UNLIKELY(cachedSequentialThreadId == InvalidSequentialThreadId)) {
+        cachedSequentialThreadId = ++SequentialThreadIdGenerator;
     }
-    return CachedSequentialThreadId;
+    return cachedSequentialThreadId;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/system/thread_id.cpp b/library/cpp/yt/system/thread_id.cpp
index 1ecf1fe9bd..a1971dd0dd 100644
--- a/library/cpp/yt/system/thread_id.cpp
+++ b/library/cpp/yt/system/thread_id.cpp
@@ -6,7 +6,7 @@ namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-YT_THREAD_LOCAL(TSequentialThreadId) CachedSequentialThreadId = InvalidSequentialThreadId;
+YT_DEFINE_THREAD_LOCAL(TSequentialThreadId, CachedSequentialThreadId, InvalidSequentialThreadId);
 std::atomic<TSequentialThreadId> SequentialThreadIdGenerator = InvalidSequentialThreadId;
 
 TSystemThreadId GetSystemThreadId()
diff --git a/library/cpp/yt/threading/spin_lock_count.cpp b/library/cpp/yt/threading/spin_lock_count.cpp
index 501a8bd73e..b1670e92d4 100644
--- a/library/cpp/yt/threading/spin_lock_count.cpp
+++ b/library/cpp/yt/threading/spin_lock_count.cpp
@@ -13,26 +13,26 @@ namespace NYT::NThreading::NPrivate {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-YT_THREAD_LOCAL(i64) ActiveSpinLockCount = 0;
+YT_DEFINE_THREAD_LOCAL(i64, ActiveSpinLockCount, 0);
 
 ////////////////////////////////////////////////////////////////////////////////
 
 void RecordSpinLockAcquired(bool isAcquired)
 {
     if (isAcquired) {
-        ActiveSpinLockCount++;
+        ActiveSpinLockCount()++;
     }
 }
 
 void RecordSpinLockReleased()
 {
-    YT_VERIFY(ActiveSpinLockCount > 0);
-    ActiveSpinLockCount--;
+    YT_VERIFY(ActiveSpinLockCount() > 0);
+    ActiveSpinLockCount()--;
 }
 
 void VerifyNoSpinLockAffinity()
 {
-    YT_VERIFY(ActiveSpinLockCount == 0);
+    YT_VERIFY(ActiveSpinLockCount() == 0);
 }
 
 #endif