1 | /* |
2 | * Copyright 2014-present Facebook, Inc. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | */ |
16 | |
17 | #pragma once |
18 | |
19 | #include <memory> |
20 | #include <string> |
21 | |
22 | #include <folly/Synchronized.h> |
23 | #include <folly/container/F14Map.h> |
24 | #include <folly/sorted_vector_types.h> |
25 | |
26 | namespace folly { |
27 | |
28 | /* |
29 | * A token to be used to fetch data from RequestContext. |
30 | * Generally you will want this to be a static, created only once using a |
31 | * string, and then only copied. The string constructor is expensive. |
32 | */ |
33 | class RequestToken { |
34 | public: |
35 | explicit RequestToken(const std::string& str); |
36 | |
37 | bool operator==(const RequestToken& other) const { |
38 | return token_ == other.token_; |
39 | } |
40 | |
41 | // Slow, use only for debug log messages. |
42 | std::string getDebugString() const; |
43 | |
44 | friend struct std::hash<folly::RequestToken>; |
45 | |
46 | private: |
47 | static Synchronized<F14FastMap<std::string, uint32_t>>& getCache(); |
48 | |
49 | uint32_t token_; |
50 | }; |
51 | |
52 | } // namespace folly |
53 | |
54 | namespace std { |
55 | template <> |
56 | struct hash<folly::RequestToken> { |
57 | size_t operator()(const folly::RequestToken& token) const { |
58 | return hash<uint32_t>()(token.token_); |
59 | } |
60 | }; |
61 | } // namespace std |
62 | |
63 | namespace folly { |
64 | |
65 | // Some request context that follows an async request through a process |
66 | // Everything in the context must be thread safe |
67 | |
68 | class RequestData { |
69 | public: |
70 | virtual ~RequestData() = default; |
71 | |
72 | // Avoid calling RequestContext::setContextData, setContextDataIfAbsent, or |
73 | // clearContextData from these callbacks. Doing so will cause deadlock. We |
74 | // could fix these deadlocks, but only at significant performance penalty, so |
75 | // just don't do it! |
76 | |
77 | virtual bool hasCallback() = 0; |
78 | // Callback executed when setting RequestContext. Make sure your RequestData |
79 | // instance overrides the hasCallback method to return true otherwise |
80 | // the callback will not be executed |
81 | virtual void onSet() {} |
82 | // Callback executed when unsetting RequestContext. Make sure your RequestData |
83 | // instance overrides the hasCallback method to return true otherwise |
84 | // the callback will not be executed |
85 | virtual void onUnset() {} |
86 | |
87 | private: |
88 | // Start shallow copy implementation details: |
89 | // For efficiency, RequestContext provides a raw ptr interface. |
90 | // To support shallow copy, we need a shared ptr. |
91 | // To keep it as safe as possible (even if a raw ptr is passed back), |
92 | // the counter lives directly in RequestData. |
93 | |
94 | friend class RequestContext; |
95 | |
96 | // Unique ptr with custom destructor, decrement the counter |
97 | // and only free if 0 |
98 | struct DestructPtr { |
99 | void operator()(RequestData* ptr); |
100 | }; |
101 | using SharedPtr = std::unique_ptr<RequestData, DestructPtr>; |
102 | |
103 | // Initialize the pseudo-shared ptr, increment the counter |
104 | static SharedPtr constructPtr(RequestData* ptr); |
105 | |
106 | std::atomic<int> keepAliveCounter_{0}; |
107 | // End shallow copy |
108 | }; |
109 | |
110 | // If you do not call create() to create a unique request context, |
111 | // this default request context will always be returned, and is never |
112 | // copied between threads. |
113 | class RequestContext { |
114 | public: |
115 | // Create a unique request context for this request. |
116 | // It will be passed between queues / threads (where implemented), |
117 | // so it should be valid for the lifetime of the request. |
118 | static void create() { |
119 | setContext(std::make_shared<RequestContext>()); |
120 | } |
121 | |
122 | // Get the current context. |
123 | static RequestContext* get(); |
124 | |
125 | // The following APIs are used to add, remove and access RequestData instance |
126 | // in the RequestContext instance, normally used for per-RequestContext |
127 | // tracking or callback on set and unset. These APIs are Thread-safe. |
128 | // These APIs are performance sensitive, so please ask if you need help |
129 | // profiling any use of these APIs. |
130 | |
131 | // Add RequestData instance "data" to this RequestContext instance, with |
132 | // string identifier "val". If the same string identifier has already been |
133 | // used, will print a warning message for the first time, clear the existing |
134 | // RequestData instance for "val", and **not** add "data". |
135 | void setContextData( |
136 | const RequestToken& val, |
137 | std::unique_ptr<RequestData> data); |
138 | void setContextData( |
139 | const std::string& val, |
140 | std::unique_ptr<RequestData> data) { |
141 | setContextData(RequestToken(val), std::move(data)); |
142 | } |
143 | |
144 | // Add RequestData instance "data" to this RequestContext instance, with |
145 | // string identifier "val". If the same string identifier has already been |
146 | // used, return false and do nothing. Otherwise add "data" and return true. |
147 | bool setContextDataIfAbsent( |
148 | const RequestToken& val, |
149 | std::unique_ptr<RequestData> data); |
150 | bool setContextDataIfAbsent( |
151 | const std::string& val, |
152 | std::unique_ptr<RequestData> data) { |
153 | return setContextDataIfAbsent(RequestToken(val), std::move(data)); |
154 | } |
155 | |
156 | // Remove the RequestData instance with string identifier "val", if it exists. |
157 | void clearContextData(const RequestToken& val); |
158 | void clearContextData(const std::string& val) { |
159 | clearContextData(RequestToken(val)); |
160 | } |
161 | |
162 | // Returns true if and only if the RequestData instance with string identifier |
163 | // "val" exists in this RequestContext instnace. |
164 | bool hasContextData(const RequestToken& val) const; |
165 | bool hasContextData(const std::string& val) const { |
166 | return hasContextData(RequestToken(val)); |
167 | } |
168 | |
169 | // Get (constant) raw pointer of the RequestData instance with string |
170 | // identifier "val" if it exists, otherwise returns null pointer. |
171 | RequestData* getContextData(const RequestToken& val); |
172 | const RequestData* getContextData(const RequestToken& val) const; |
173 | RequestData* getContextData(const std::string& val) { |
174 | return getContextData(RequestToken(val)); |
175 | } |
176 | const RequestData* getContextData(const std::string& val) const { |
177 | return getContextData(RequestToken(val)); |
178 | } |
179 | |
180 | void onSet(); |
181 | void onUnset(); |
182 | |
183 | // The following API is used to pass the context through queues / threads. |
184 | // saveContext is called to get a shared_ptr to the context, and |
185 | // setContext is used to reset it on the other side of the queue. |
186 | // |
187 | // Whenever possible, use RequestContextScopeGuard instead of setContext |
188 | // to make sure that RequestContext is reset to the original value when |
189 | // we exit the scope. |
190 | // |
191 | // A shared_ptr is used, because many request may fan out across |
192 | // multiple threads, or do post-send processing, etc. |
193 | static std::shared_ptr<RequestContext> setContext( |
194 | std::shared_ptr<RequestContext> ctx); |
195 | |
196 | static std::shared_ptr<RequestContext> saveContext() { |
197 | return getStaticContext(); |
198 | } |
199 | |
200 | private: |
201 | static std::shared_ptr<RequestContext>& getStaticContext(); |
202 | |
203 | // Start shallow copy guard implementation details: |
204 | // All methods are private to encourage proper use |
205 | friend struct ShallowCopyRequestContextScopeGuard; |
206 | |
207 | // This sets a shallow copy of the current context as current, |
208 | // then return the previous context (so it can be reset later). |
209 | static std::shared_ptr<RequestContext> setShallowCopyContext(); |
210 | |
211 | // Similar to setContextData, except it overwrites the data |
212 | // if already set (instead of warn + reset ptr). |
213 | void overwriteContextData( |
214 | const RequestToken& val, |
215 | std::unique_ptr<RequestData> data); |
216 | void overwriteContextData( |
217 | const std::string& val, |
218 | std::unique_ptr<RequestData> data) { |
219 | overwriteContextData(RequestToken(val), std::move(data)); |
220 | } |
221 | // End shallow copy guard |
222 | |
223 | enum class DoSetBehaviour { |
224 | SET, |
225 | SET_IF_ABSENT, |
226 | OVERWRITE, |
227 | }; |
228 | |
229 | bool doSetContextData( |
230 | const RequestToken& val, |
231 | std::unique_ptr<RequestData>& data, |
232 | DoSetBehaviour behaviour); |
233 | bool doSetContextData( |
234 | const std::string& val, |
235 | std::unique_ptr<RequestData>& data, |
236 | DoSetBehaviour behaviour) { |
237 | return doSetContextData(RequestToken(val), data, behaviour); |
238 | } |
239 | |
240 | struct State { |
241 | // This must be optimized for lookup, its hot path is getContextData |
242 | // Efficiency of copying the container also matters in setShallowCopyContext |
243 | F14FastMap<RequestToken, RequestData::SharedPtr> requestData_; |
244 | // This must be optimized for iteration, its hot path is setContext |
245 | // We also use the fact that it's ordered to efficiently compute |
246 | // the difference with previous context |
247 | sorted_vector_set<RequestData*> callbackData_; |
248 | }; |
249 | folly::Synchronized<State> state_; |
250 | }; |
251 | |
252 | /** |
253 | * Note: you probably want to use ShallowCopyRequestContextScopeGuard |
254 | * This resets all other RequestData for the duration of the scope! |
255 | */ |
256 | class RequestContextScopeGuard { |
257 | private: |
258 | std::shared_ptr<RequestContext> prev_; |
259 | |
260 | public: |
261 | RequestContextScopeGuard(const RequestContextScopeGuard&) = delete; |
262 | RequestContextScopeGuard& operator=(const RequestContextScopeGuard&) = delete; |
263 | RequestContextScopeGuard(RequestContextScopeGuard&&) = delete; |
264 | RequestContextScopeGuard& operator=(RequestContextScopeGuard&&) = delete; |
265 | |
266 | // Create a new RequestContext and reset to the original value when |
267 | // this goes out of scope. |
268 | RequestContextScopeGuard() : prev_(RequestContext::saveContext()) { |
269 | RequestContext::create(); |
270 | } |
271 | |
272 | // Set a RequestContext that was previously captured by saveContext(). It will |
273 | // be automatically reset to the original value when this goes out of scope. |
274 | explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx) |
275 | : prev_(RequestContext::setContext(std::move(ctx))) {} |
276 | |
277 | ~RequestContextScopeGuard() { |
278 | RequestContext::setContext(std::move(prev_)); |
279 | } |
280 | }; |
281 | |
282 | /** |
283 | * This guard maintains all the RequestData pointers of the parent. |
284 | * This allows to overwrite a specific RequestData pointer for the |
285 | * scope's duration, without breaking others. |
286 | * |
287 | * Only modified pointers will have their set/onset methods called |
288 | */ |
289 | struct ShallowCopyRequestContextScopeGuard { |
290 | ShallowCopyRequestContextScopeGuard() |
291 | : prev_(RequestContext::setShallowCopyContext()) {} |
292 | |
293 | /** |
294 | * Shallow copy then overwrite one specific RequestData |
295 | * |
296 | * Helper constructor which is a more efficient equivalent to |
297 | * "clearRequestData" then "setRequestData" after the guard. |
298 | */ |
299 | ShallowCopyRequestContextScopeGuard( |
300 | const RequestToken& val, |
301 | std::unique_ptr<RequestData> data) |
302 | : ShallowCopyRequestContextScopeGuard() { |
303 | RequestContext::get()->overwriteContextData(val, std::move(data)); |
304 | } |
305 | ShallowCopyRequestContextScopeGuard( |
306 | const std::string& val, |
307 | std::unique_ptr<RequestData> data) |
308 | : ShallowCopyRequestContextScopeGuard() { |
309 | RequestContext::get()->overwriteContextData(val, std::move(data)); |
310 | } |
311 | |
312 | ~ShallowCopyRequestContextScopeGuard() { |
313 | RequestContext::setContext(std::move(prev_)); |
314 | } |
315 | |
316 | ShallowCopyRequestContextScopeGuard( |
317 | const ShallowCopyRequestContextScopeGuard&) = delete; |
318 | ShallowCopyRequestContextScopeGuard& operator=( |
319 | const ShallowCopyRequestContextScopeGuard&) = delete; |
320 | ShallowCopyRequestContextScopeGuard(ShallowCopyRequestContextScopeGuard&&) = |
321 | delete; |
322 | ShallowCopyRequestContextScopeGuard& operator=( |
323 | ShallowCopyRequestContextScopeGuard&&) = delete; |
324 | |
325 | private: |
326 | std::shared_ptr<RequestContext> prev_; |
327 | }; |
328 | |
329 | } // namespace folly |
330 | |