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
25namespace folly {
26
27RequestToken::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
49std::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
60Synchronized<F14FastMap<std::string, uint32_t>>& RequestToken::getCache() {
61 static Indestructible<Synchronized<F14FastMap<std::string, uint32_t>>> cache;
62 return *cache;
63}
64
65void 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
87bool 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
129void RequestContext::setContextData(
130 const RequestToken& val,
131 std::unique_ptr<RequestData> data) {
132 doSetContextData(val, data, DoSetBehaviour::SET);
133}
134
135bool RequestContext::setContextDataIfAbsent(
136 const RequestToken& val,
137 std::unique_ptr<RequestData> data) {
138 return doSetContextData(val, data, DoSetBehaviour::SET_IF_ABSENT);
139}
140
141void RequestContext::overwriteContextData(
142 const RequestToken& val,
143 std::unique_ptr<RequestData> data) {
144 doSetContextData(val, data, DoSetBehaviour::OVERWRITE);
145}
146
147bool RequestContext::hasContextData(const RequestToken& val) const {
148 return state_.rlock()->requestData_.count(val);
149}
150
151RequestData* RequestContext::getContextData(const RequestToken& val) {
152 const RequestData::SharedPtr dflt{nullptr};
153 return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
154}
155
156const 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
162void RequestContext::onSet() {
163 auto rlock = state_.rlock();
164 for (const auto& data : rlock->callbackData_) {
165 data->onSet();
166 }
167}
168
169void RequestContext::onUnset() {
170 auto rlock = state_.rlock();
171 for (const auto& data : rlock->callbackData_) {
172 data->onUnset();
173 }
174}
175
176void 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
198namespace {
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
201template <typename TData, typename TExec>
202void 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
227std::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
263std::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>
269RequestContext::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
290RequestContext* 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