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
26namespace 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 */
33class 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
54namespace std {
55template <>
56struct 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
63namespace folly {
64
65// Some request context that follows an async request through a process
66// Everything in the context must be thread safe
67
68class 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.
113class 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 */
256class 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 */
289struct 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