Improve TGenericLocalValue, allow mixed thread-fiber-etc values

Make generic local values more safe and usable:

- instead of factory, register fiber/coroutine-aware "GLS (general local storage) contexts";
- put up to 4 (normally no more than 2 -- threads and one coroutine implementation) different "tls" variables into one generic, choose correct one from current execution context;
- suitable version of variable will be constructed on demand (no dependency from first usage);
- improve unit test too.
commit_hash:8586846a6a775bd66dffcdf58263f78042be2480

2 files changed, 86 insertions, 25 deletions

diff --git a/library/cpp/threading/thread_local/generic.cpp b/library/cpp/threading/thread_local/generic.cpp
index 79161a846da..eca63d76bc6 100644
--- a/library/cpp/threading/thread_local/generic.cpp
+++ b/library/cpp/threading/thread_local/generic.cpp
@@ -14,29 +14,64 @@ namespace {
         NThreading::TThreadLocalValue<TData, NThreading::EThreadLocalImpl::StdThreadLocal> Data_;
     };
 
-    std::atomic<size_t>& DefaultFactoryUsageCounter() {
-        static std::atomic<size_t> v;
-        return v;
-    }
+    class TThreadLocalContext
+        : public NThreading::IGLSContext
+    {
+        bool IsCurrent() const override {
+            return true;
+        }
 
-    auto& genericLocalStorageFactory() {
-        static NThreading::TGenericLocalStorageFactory factory = [] {
-            DefaultFactoryUsageCounter() += 1;
+        THolder<NThreading::IGenericLocalStorage> MakeStorage() const override {
             return MakeHolder<TThreadLocalStorage>();
-        };
+        }
+    };
+
+    class TContextRegistry {
+    public:
+        TContextRegistry() {
+            Register(MakeHolder<TThreadLocalContext>());
+        }
+
+        size_t Count() const {
+            return Count_.load();
+        }
 
-        return factory;
+        const NThreading::IGLSContext& Get(size_t index) const {
+            return *Contexts_[index];
+        }
+
+        void Register(THolder<NThreading::IGLSContext> context) {
+            with_lock (Lock_) {
+                const size_t index = Count_.load();
+                Y_ENSURE(index < NThreading::NDetail::MaxGLSContexts, "Too many generic local contexts registered");
+                Contexts_[index] = std::move(context);
+                Count_.store(index + 1);
+            }
+        }
+    private:
+        TAdaptiveLock Lock_;
+        std::atomic<size_t> Count_ = 0;
+        std::array<THolder<NThreading::IGLSContext>, NThreading::NDetail::MaxGLSContexts> Contexts_ = {};
+    };
+
+    TContextRegistry& Registry() {
+        static TContextRegistry registry;
+        return registry;
     }
 }
 
 namespace NThreading {
-    void SetGenericLocalStorageFactory(TGenericLocalStorageFactory factory) {
-        Y_ENSURE(DefaultFactoryUsageCounter() == 0, "There are some thread local values allocated with default factory");
-
-        genericLocalStorageFactory() = factory;
+    void RegisterGLSContext(THolder<IGLSContext> context) {
+        Registry().Register(std::move(context));
     }
 
-    THolder<IGenericLocalStorage> MakeGenericLocalStorage() {
-        return genericLocalStorageFactory()();
+    namespace NDetail {
+        size_t GLSContextCount() {
+            return Registry().Count();
+        }
+
+        const IGLSContext& GetGLSContext(size_t index) {
+            return Registry().Get(index);
+        }
     }
 }
diff --git a/library/cpp/threading/thread_local/generic.h b/library/cpp/threading/thread_local/generic.h
index f60d5b786dd..b0c341f867e 100644
--- a/library/cpp/threading/thread_local/generic.h
+++ b/library/cpp/threading/thread_local/generic.h
@@ -3,6 +3,7 @@
 #include <util/generic/ptr.h>
 #include <util/generic/vector.h>
 
+#include <array>
 #include <functional>
 #include <mutex>
 
@@ -44,14 +45,32 @@ namespace NThreading {
         virtual TData* GetData() const = 0;
     };
 
-    using TGenericLocalStorageFactory = std::function<THolder<IGenericLocalStorage>()>;
+    class IGLSContext {
+    public:
+        virtual ~IGLSContext() = default;
+
+        virtual bool IsCurrent() const = 0;
+        virtual THolder<IGenericLocalStorage> MakeStorage() const = 0;
+    };
 
-    void SetGenericLocalStorageFactory(TGenericLocalStorageFactory factory);
-    THolder<IGenericLocalStorage> MakeGenericLocalStorage();
+    // Later registrations take priority over earlier ones.
+    void RegisterGLSContext(THolder<IGLSContext> context);
+
+    namespace NDetail {
+        inline constexpr size_t MaxGLSContexts = 4;
+
+        size_t GLSContextCount();
+        const IGLSContext& GetGLSContext(size_t index);
+    }
 
     template <typename T>
     class TGenericLocalValue {
     private:
+        struct TSlot {
+            std::once_flag InitOnce;
+            THolder<IGenericLocalStorage> Storage;
+        };
+
         static const auto& Traits() {
             const static IGenericLocalStorage::TTraits traits = {
                 .Size = sizeof(T),
@@ -63,19 +82,26 @@ namespace NThreading {
         };
     public:
         T* Get() const {
-            std::call_once(InitOnce_, [this]() {
-                Storage_ = MakeGenericLocalStorage();
-            });
+            const size_t count = NDetail::GLSContextCount();
+            for (size_t index = count; index-- > 0;) {
+                const IGLSContext& context = NDetail::GetGLSContext(index);
+                if (context.IsCurrent()) {
+                    return GetMemory(ContextSlots_[index], context);
+                }
+            }
 
-            return static_cast<T*>(Storage_->GetMemory(Traits()));
+            Y_ABORT("unreachable, ContextSlots_[0] always IsCurrent");
         }
 
         T& GetRef() const {
             return *Get();
         }
     private:
-        mutable std::once_flag InitOnce_;
-        mutable THolder<IGenericLocalStorage> Storage_;
+        T* GetMemory(TSlot& slot, const IGLSContext& context) const {
+            std::call_once(slot.InitOnce, [&] { slot.Storage = context.MakeStorage(); });
+            return static_cast<T*>(slot.Storage->GetMemory(Traits()));
+        }
+    private:
+        mutable std::array<TSlot, NDetail::MaxGLSContexts> ContextSlots_;
     };
 }
-