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.
16 #include <folly/io/async/Request.h>
22 #include <glog/logging.h>
24 #include <folly/MapUtil.h>
25 #include <folly/SingletonThreadLocal.h>
26 #include <folly/tracing/StaticTracepoint.h>
30 void RequestContext::setContextData(
31 const std::string& val,
32 std::unique_ptr<RequestData> data) {
33 auto wlock = data_.wlock();
34 if (wlock->count(val)) {
35 LOG_FIRST_N(WARNING, 1)
36 << "Called RequestContext::setContextData with data already set";
38 (*wlock)[val] = nullptr;
40 (*wlock)[val] = std::move(data);
44 bool RequestContext::setContextDataIfAbsent(
45 const std::string& val,
46 std::unique_ptr<RequestData> data) {
47 auto ulock = data_.ulock();
48 if (ulock->count(val)) {
52 auto wlock = ulock.moveFromUpgradeToWrite();
53 (*wlock)[val] = std::move(data);
57 bool RequestContext::hasContextData(const std::string& val) const {
58 return data_.rlock()->count(val);
61 RequestData* RequestContext::getContextData(const std::string& val) {
62 const std::unique_ptr<RequestData> dflt{nullptr};
63 return get_ref_default(*data_.rlock(), val, dflt).get();
66 const RequestData* RequestContext::getContextData(
67 const std::string& val) const {
68 const std::unique_ptr<RequestData> dflt{nullptr};
69 return get_ref_default(*data_.rlock(), val, dflt).get();
72 void RequestContext::onSet() {
73 auto rlock = data_.rlock();
74 for (auto const& ent : *rlock) {
75 if (auto& data = ent.second) {
81 void RequestContext::onUnset() {
82 auto rlock = data_.rlock();
83 for (auto const& ent : *rlock) {
84 if (auto& data = ent.second) {
90 void RequestContext::clearContextData(const std::string& val) {
91 std::unique_ptr<RequestData> requestData;
92 // Delete the RequestData after giving up the wlock just in case one of the
93 // RequestData destructors will try to grab the lock again.
95 auto wlock = data_.wlock();
96 auto it = wlock->find(val);
97 if (it != wlock->end()) {
98 requestData = std::move(it->second);
104 std::shared_ptr<RequestContext> RequestContext::setContext(
105 std::shared_ptr<RequestContext> ctx) {
106 auto& curCtx = getStaticContext();
108 FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
121 RequestContext::Provider& RequestContext::requestContextProvider() {
122 class DefaultProvider {
124 constexpr DefaultProvider() = default;
125 DefaultProvider(const DefaultProvider&) = delete;
126 DefaultProvider& operator=(const DefaultProvider&) = delete;
127 DefaultProvider(DefaultProvider&&) = default;
128 DefaultProvider& operator=(DefaultProvider&&) = default;
130 std::shared_ptr<RequestContext>& operator()() {
135 std::shared_ptr<RequestContext> context;
138 static SingletonThreadLocal<Provider> providerSingleton(
139 []() { return new Provider(DefaultProvider()); });
140 return providerSingleton.get();
143 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
144 auto& provider = requestContextProvider();
148 RequestContext* RequestContext::get() {
149 auto& context = getStaticContext();
151 static RequestContext defaultContext;
152 return std::addressof(defaultContext);
154 return context.get();
157 RequestContext::Provider RequestContext::setRequestContextProvider(
158 RequestContext::Provider newProvider) {
160 throw std::runtime_error("RequestContext provider must be non-empty");
163 auto& provider = requestContextProvider();
164 std::swap(provider, newProvider);