2 * Copyright 2004-present Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #include <folly/io/async/Request.h>
18 #include <folly/tracing/StaticTracepoint.h>
20 #include <glog/logging.h>
22 #include <folly/MapUtil.h>
23 #include <folly/SingletonThreadLocal.h>
27 bool RequestContext::doSetContextData(
28 const std::string& val,
29 std::unique_ptr<RequestData>& data,
31 auto ulock = state_.ulock();
33 bool conflict = false;
34 auto it = ulock->requestData_.find(val);
35 if (it != ulock->requestData_.end()) {
39 LOG_FIRST_N(WARNING, 1) << "Calling RequestContext::setContextData for "
40 << val << " but it is already set";
45 auto wlock = ulock.moveFromUpgradeToWrite();
48 if (it->second->hasCallback()) {
49 wlock->callbackData_.erase(it->second.get());
51 it->second.reset(nullptr);
56 if (data && data->hasCallback()) {
57 wlock->callbackData_.insert(data.get());
59 wlock->requestData_[val] = std::move(data);
64 void RequestContext::setContextData(
65 const std::string& val,
66 std::unique_ptr<RequestData> data) {
67 doSetContextData(val, data, false /* strict */);
70 bool RequestContext::setContextDataIfAbsent(
71 const std::string& val,
72 std::unique_ptr<RequestData> data) {
73 return doSetContextData(val, data, true /* strict */);
76 bool RequestContext::hasContextData(const std::string& val) const {
77 return state_.rlock()->requestData_.count(val);
80 RequestData* RequestContext::getContextData(const std::string& val) {
81 const std::unique_ptr<RequestData> dflt{nullptr};
82 return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
85 const RequestData* RequestContext::getContextData(
86 const std::string& val) const {
87 const std::unique_ptr<RequestData> dflt{nullptr};
88 return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
91 void RequestContext::onSet() {
92 auto rlock = state_.rlock();
93 for (const auto& data : rlock->callbackData_) {
98 void RequestContext::onUnset() {
99 auto rlock = state_.rlock();
100 for (const auto& data : rlock->callbackData_) {
105 void RequestContext::clearContextData(const std::string& val) {
106 std::unique_ptr<RequestData> requestData;
107 // Delete the RequestData after giving up the wlock just in case one of the
108 // RequestData destructors will try to grab the lock again.
110 auto ulock = state_.ulock();
111 auto it = ulock->requestData_.find(val);
112 if (it == ulock->requestData_.end()) {
116 auto wlock = ulock.moveFromUpgradeToWrite();
117 if (it->second && it->second->hasCallback()) {
118 wlock->callbackData_.erase(it->second.get());
121 requestData = std::move(it->second);
122 wlock->requestData_.erase(it);
126 std::shared_ptr<RequestContext> RequestContext::setContext(
127 std::shared_ptr<RequestContext> ctx) {
128 auto& curCtx = getStaticContext();
130 FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
143 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
144 using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
145 static SingletonT singleton;
147 return singleton.get();
150 RequestContext* RequestContext::get() {
151 auto& context = getStaticContext();
153 static RequestContext defaultContext;
154 return std::addressof(defaultContext);
156 return context.get();