userver: userver/engine/impl/task_local_storage.hpp Source File
Loading...
Searching...
No Matches
task_local_storage.hpp
1#pragma once
2
3#include <atomic>
4#include <cstddef>
5#include <memory>
6#include <type_traits>
7#include <typeinfo>
8
9#include <userver/utils/fast_pimpl.hpp>
10
11USERVER_NAMESPACE_BEGIN
12
13namespace engine::impl::task_local {
14
15using Key = std::size_t;
16
17enum class VariableKind { kNormal, kInherited };
18
19class DataBase {
20 public:
21 using Deleter = void (*)(DataBase&) noexcept;
22
23 void DeleteSelf() noexcept;
24
25 protected:
26 explicit DataBase(Deleter deleter);
27
28 private:
29 const Deleter deleter_;
30};
31
32class NormalDataBase : public DataBase {
33 protected:
34 template <typename Derived>
35 NormalDataBase(std::in_place_type_t<Derived>, Key /*key*/)
36 : DataBase([](DataBase& base) noexcept {
37 delete &static_cast<Derived&>(base);
38 }) {}
39};
40
41class InheritedDataBase : public DataBase {
42 public:
43 void AddRef() noexcept;
44
45 Key GetKey() const noexcept;
46
47 protected:
48 template <typename Derived>
49 InheritedDataBase(std::in_place_type_t<Derived>, Key key)
50 : DataBase([](DataBase& base) noexcept {
51 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-static-cast-downcast)
52 if (--static_cast<InheritedDataBase&>(base).ref_counter_ == 0) {
53 delete &static_cast<Derived&>(base);
54 }
55 }),
56 key_(key) {}
57
58 private:
59 std::atomic<std::size_t> ref_counter_{1};
60 const Key key_{};
61};
62
63template <VariableKind Kind>
64using ConditionalDataBase =
65 std::conditional_t<Kind == VariableKind::kInherited, InheritedDataBase,
66 NormalDataBase>;
67
68template <typename T, VariableKind Kind>
69class DataImpl final : public ConditionalDataBase<Kind> {
70 using Base = ConditionalDataBase<Kind>;
71
72 public:
73 template <typename... Args>
74 explicit DataImpl(Key key, Args&&... args)
75 : Base(std::in_place_type<DataImpl>, key),
76 variable_(std::forward<Args>(args)...) {}
77
78 T& Get() noexcept { return variable_; }
79
80 private:
81 T variable_;
82};
83
84[[noreturn]] void ReportVariableNotSet(const std::type_info& type);
85
86class Storage final {
87 public:
88 Storage();
89
90 Storage(Storage&&) = delete;
91 Storage& operator=(Storage&&) = delete;
92 ~Storage();
93
94 // Copies pointers to inherited variables from 'other'
95 // 'this' must not contain any variables
96 void InheritFrom(Storage& other);
97
98 // Copies pointers to specific inherited variables from 'other'
99 // 'this' must not contain the variable being copied
100 // does nothing if there is nothing to copy
101 void InheritNodeIfExists(Storage& other, Key key);
102
103 // Moves other's variables into 'this', leaving 'other' in an empty state
104 // 'this' must not contain any variables
105 void InitializeFrom(Storage&& other) noexcept;
106
107 // Variable accessors must be called with the same T, Kind, key.
108 // Otherwise it is UB.
109 template <typename T, VariableKind Kind>
110 T& GetOrEmplace(Key key) {
111 DataBase* const old_data = GetGeneric(key);
112 if (!old_data) {
113 const bool has_existing_variable = false;
114 return DoEmplace<T, Kind>(key, has_existing_variable);
115 }
116 return static_cast<DataImpl<T, Kind>&>(*old_data).Get();
117 }
118
119 template <typename T, VariableKind Kind>
120 T* GetOptional(Key key) noexcept {
121 DataBase* const data = GetGeneric(key);
122 if (!data) return nullptr;
123 return &static_cast<DataImpl<T, Kind>&>(*data).Get();
124 }
125
126 template <typename T, VariableKind Kind>
127 T& Get(Key key) {
128 T* const result = GetOptional<T, Kind>(key);
129 if (!result) ReportVariableNotSet(typeid(T));
130 return *result;
131 }
132
133 template <typename T, VariableKind Kind, typename... Args>
134 T& Emplace(Key key, Args&&... args) {
135 DataBase* const old_data = GetGeneric(key);
136 const bool has_existing_variable = old_data != nullptr;
137 auto& result = DoEmplace<T, Kind>(key, has_existing_variable,
138 std::forward<Args>(args)...);
139 if (old_data) old_data->DeleteSelf();
140 return result;
141 }
142
143 template <typename T, VariableKind Kind>
144 void Erase(Key key) noexcept {
145 static_assert(Kind == VariableKind::kInherited);
146 EraseInherited(key);
147 }
148
149 private:
150 DataBase* GetGeneric(Key key) noexcept;
151
152 void SetGeneric(Key key, NormalDataBase& node, bool has_existing_variable);
153
154 void SetGeneric(Key key, InheritedDataBase& node, bool has_existing_variable);
155
156 void EraseInherited(Key key) noexcept;
157
158 void InheritNode(InheritedDataBase&);
159
160 // Provides strong exception guarantee. Does not delete the old data, if any.
161 template <typename T, VariableKind Kind, typename... Args>
162 T& DoEmplace(Key key, bool has_existing_variable, Args&&... args) {
163 auto new_data =
164 std::make_unique<DataImpl<T, Kind>>(key, std::forward<Args>(args)...);
165 SetGeneric(key, *new_data, has_existing_variable);
166 return new_data.release()->Get();
167 }
168
169 struct Impl;
170 utils::FastPimpl<Impl, 40, 8> impl_;
171};
172
173class Variable final {
174 public:
175 Variable();
176
177 Variable(const Variable&) = delete;
178 Variable(Variable&&) = delete;
179 Variable& operator=(const Variable&) = delete;
180 Variable& operator=(Variable&&) = delete;
181
182 Key GetKey() const noexcept;
183
184 private:
185 const Key key_;
186};
187
188Storage& GetCurrentStorage() noexcept;
189
190struct InternalTag final {
191 explicit InternalTag() = default;
192};
193
194} // namespace engine::impl::task_local
195
196USERVER_NAMESPACE_END