1 | /* |
2 | * Copyright 2016-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 | #include <folly/io/async/Request.h> |
18 | #include <folly/tracing/StaticTracepoint.h> |
19 | |
20 | #include <glog/logging.h> |
21 | |
22 | #include <folly/MapUtil.h> |
23 | #include <folly/SingletonThreadLocal.h> |
24 | |
25 | namespace folly { |
26 | |
27 | RequestToken::RequestToken(const std::string& str) { |
28 | auto& cache = getCache(); |
29 | { |
30 | auto c = cache.rlock(); |
31 | auto res = c->find(str); |
32 | if (res != c->end()) { |
33 | token_ = res->second; |
34 | return; |
35 | } |
36 | } |
37 | auto c = cache.wlock(); |
38 | auto res = c->find(str); |
39 | if (res != c->end()) { |
40 | token_ = res->second; |
41 | return; |
42 | } |
43 | static uint32_t nextToken{1}; |
44 | |
45 | token_ = nextToken++; |
46 | (*c)[str] = token_; |
47 | } |
48 | |
49 | std::string RequestToken::getDebugString() const { |
50 | auto& cache = getCache(); |
51 | auto c = cache.rlock(); |
52 | for (auto& v : *c) { |
53 | if (v.second == token_) { |
54 | return v.first; |
55 | } |
56 | } |
57 | throw std::logic_error("Could not find debug string in RequestToken" ); |
58 | } |
59 | |
60 | Synchronized<F14FastMap<std::string, uint32_t>>& RequestToken::getCache() { |
61 | static Indestructible<Synchronized<F14FastMap<std::string, uint32_t>>> cache; |
62 | return *cache; |
63 | } |
64 | |
65 | void RequestData::DestructPtr::operator()(RequestData* ptr) { |
66 | if (ptr) { |
67 | auto keepAliveCounter = |
68 | ptr->keepAliveCounter_.fetch_sub(1, std::memory_order_acq_rel); |
69 | // Note: this is the value before decrement, hence == 1 check |
70 | DCHECK(keepAliveCounter > 0); |
71 | if (keepAliveCounter == 1) { |
72 | delete ptr; |
73 | } |
74 | } |
75 | } |
76 | |
77 | /* static */ RequestData::SharedPtr RequestData::constructPtr( |
78 | RequestData* ptr) { |
79 | if (ptr) { |
80 | auto keepAliveCounter = |
81 | ptr->keepAliveCounter_.fetch_add(1, std::memory_order_relaxed); |
82 | DCHECK(keepAliveCounter >= 0); |
83 | } |
84 | return SharedPtr(ptr); |
85 | } |
86 | |
87 | bool RequestContext::doSetContextData( |
88 | const RequestToken& val, |
89 | std::unique_ptr<RequestData>& data, |
90 | DoSetBehaviour behaviour) { |
91 | auto ulock = state_.ulock(); |
92 | |
93 | bool conflict = false; |
94 | auto it = ulock->requestData_.find(val); |
95 | if (it != ulock->requestData_.end()) { |
96 | if (behaviour == DoSetBehaviour::SET_IF_ABSENT) { |
97 | return false; |
98 | } else if (behaviour == DoSetBehaviour::SET) { |
99 | LOG_FIRST_N(WARNING, 1) |
100 | << "Calling RequestContext::setContextData for " |
101 | << val.getDebugString() << " but it is already set" ; |
102 | } |
103 | conflict = true; |
104 | } |
105 | |
106 | auto wlock = ulock.moveFromUpgradeToWrite(); |
107 | if (conflict) { |
108 | if (it->second) { |
109 | if (it->second->hasCallback()) { |
110 | it->second->onUnset(); |
111 | wlock->callbackData_.erase(it->second.get()); |
112 | } |
113 | it->second.reset(nullptr); |
114 | } |
115 | if (behaviour == DoSetBehaviour::SET) { |
116 | return true; |
117 | } |
118 | } |
119 | |
120 | if (data && data->hasCallback()) { |
121 | wlock->callbackData_.insert(data.get()); |
122 | data->onSet(); |
123 | } |
124 | wlock->requestData_[val] = RequestData::constructPtr(data.release()); |
125 | |
126 | return true; |
127 | } |
128 | |
129 | void RequestContext::setContextData( |
130 | const RequestToken& val, |
131 | std::unique_ptr<RequestData> data) { |
132 | doSetContextData(val, data, DoSetBehaviour::SET); |
133 | } |
134 | |
135 | bool RequestContext::setContextDataIfAbsent( |
136 | const RequestToken& val, |
137 | std::unique_ptr<RequestData> data) { |
138 | return doSetContextData(val, data, DoSetBehaviour::SET_IF_ABSENT); |
139 | } |
140 | |
141 | void RequestContext::overwriteContextData( |
142 | const RequestToken& val, |
143 | std::unique_ptr<RequestData> data) { |
144 | doSetContextData(val, data, DoSetBehaviour::OVERWRITE); |
145 | } |
146 | |
147 | bool RequestContext::hasContextData(const RequestToken& val) const { |
148 | return state_.rlock()->requestData_.count(val); |
149 | } |
150 | |
151 | RequestData* RequestContext::getContextData(const RequestToken& val) { |
152 | const RequestData::SharedPtr dflt{nullptr}; |
153 | return get_ref_default(state_.rlock()->requestData_, val, dflt).get(); |
154 | } |
155 | |
156 | const RequestData* RequestContext::getContextData( |
157 | const RequestToken& val) const { |
158 | const RequestData::SharedPtr dflt{nullptr}; |
159 | return get_ref_default(state_.rlock()->requestData_, val, dflt).get(); |
160 | } |
161 | |
162 | void RequestContext::onSet() { |
163 | auto rlock = state_.rlock(); |
164 | for (const auto& data : rlock->callbackData_) { |
165 | data->onSet(); |
166 | } |
167 | } |
168 | |
169 | void RequestContext::onUnset() { |
170 | auto rlock = state_.rlock(); |
171 | for (const auto& data : rlock->callbackData_) { |
172 | data->onUnset(); |
173 | } |
174 | } |
175 | |
176 | void RequestContext::clearContextData(const RequestToken& val) { |
177 | RequestData::SharedPtr requestData; |
178 | // Delete the RequestData after giving up the wlock just in case one of the |
179 | // RequestData destructors will try to grab the lock again. |
180 | { |
181 | auto ulock = state_.ulock(); |
182 | auto it = ulock->requestData_.find(val); |
183 | if (it == ulock->requestData_.end()) { |
184 | return; |
185 | } |
186 | |
187 | auto wlock = ulock.moveFromUpgradeToWrite(); |
188 | if (it->second && it->second->hasCallback()) { |
189 | it->second->onUnset(); |
190 | wlock->callbackData_.erase(it->second.get()); |
191 | } |
192 | |
193 | requestData = std::move(it->second); |
194 | wlock->requestData_.erase(it); |
195 | } |
196 | } |
197 | |
198 | namespace { |
199 | // Execute functor exec for all RequestData in data, which are not in other |
200 | // Similar to std::set_difference but avoid intermediate data structure |
201 | template <typename TData, typename TExec> |
202 | void exec_set_difference(const TData& data, const TData& other, TExec&& exec) { |
203 | auto diter = data.begin(); |
204 | auto dend = data.end(); |
205 | auto oiter = other.begin(); |
206 | auto oend = other.end(); |
207 | while (diter != dend) { |
208 | // Order of "if" optimizes for the 2 common cases: |
209 | // 1) empty other, switching to default context |
210 | // 2) identical other, switching to similar context with same callbacks |
211 | if (oiter == oend) { |
212 | exec(*diter); |
213 | ++diter; |
214 | } else if (*diter == *oiter) { |
215 | ++diter; |
216 | ++oiter; |
217 | } else if (*diter < *oiter) { |
218 | exec(*diter); |
219 | ++diter; |
220 | } else { |
221 | ++oiter; |
222 | } |
223 | } |
224 | } |
225 | } // namespace |
226 | |
227 | std::shared_ptr<RequestContext> RequestContext::setContext( |
228 | std::shared_ptr<RequestContext> newCtx) { |
229 | auto& staticCtx = getStaticContext(); |
230 | if (newCtx == staticCtx) { |
231 | return newCtx; |
232 | } |
233 | |
234 | FOLLY_SDT( |
235 | folly, request_context_switch_before, staticCtx.get(), newCtx.get()); |
236 | |
237 | auto curCtx = staticCtx; |
238 | if (newCtx && curCtx) { |
239 | // Only call set/unset for all request data that differs |
240 | auto ret = folly::acquireLocked( |
241 | as_const(newCtx->state_), as_const(curCtx->state_)); |
242 | auto& newLock = std::get<0>(ret); |
243 | auto& curLock = std::get<1>(ret); |
244 | auto& newData = newLock->callbackData_; |
245 | auto& curData = curLock->callbackData_; |
246 | exec_set_difference( |
247 | curData, newData, [](RequestData* data) { data->onUnset(); }); |
248 | staticCtx = newCtx; |
249 | exec_set_difference( |
250 | newData, curData, [](RequestData* data) { data->onSet(); }); |
251 | } else { |
252 | if (curCtx) { |
253 | curCtx->onUnset(); |
254 | } |
255 | staticCtx = newCtx; |
256 | if (newCtx) { |
257 | newCtx->onSet(); |
258 | } |
259 | } |
260 | return curCtx; |
261 | } |
262 | |
263 | std::shared_ptr<RequestContext>& RequestContext::getStaticContext() { |
264 | using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>; |
265 | return SingletonT::get(); |
266 | } |
267 | |
268 | /* static */ std::shared_ptr<RequestContext> |
269 | RequestContext::setShallowCopyContext() { |
270 | auto& parent = getStaticContext(); |
271 | auto child = std::make_shared<RequestContext>(); |
272 | |
273 | if (parent) { |
274 | auto ret = folly::acquireLocked(as_const(parent->state_), child->state_); |
275 | auto& parentLock = std::get<0>(ret); |
276 | auto& childLock = std::get<1>(ret); |
277 | childLock->callbackData_ = parentLock->callbackData_; |
278 | childLock->requestData_.reserve(parentLock->requestData_.size()); |
279 | for (const auto& entry : parentLock->requestData_) { |
280 | childLock->requestData_.insert(std::make_pair( |
281 | entry.first, RequestData::constructPtr(entry.second.get()))); |
282 | } |
283 | } |
284 | |
285 | // Do not use setContext to avoid global set/unset |
286 | std::swap(child, parent); |
287 | return child; |
288 | } |
289 | |
290 | RequestContext* RequestContext::get() { |
291 | auto& context = getStaticContext(); |
292 | if (!context) { |
293 | static RequestContext defaultContext; |
294 | return std::addressof(defaultContext); |
295 | } |
296 | return context.get(); |
297 | } |
298 | } // namespace folly |
299 | |