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 void RequestContext::setContextData(
28 const std::string& val,
29 std::unique_ptr<RequestData> data) {
30 auto wlock = data_.wlock();
31 if (wlock->count(val)) {
32 LOG_FIRST_N(WARNING, 1)
33 << "Called RequestContext::setContextData with data already set";
35 (*wlock)[val] = nullptr;
37 (*wlock)[val] = std::move(data);
41 bool RequestContext::setContextDataIfAbsent(
42 const std::string& val,
43 std::unique_ptr<RequestData> data) {
44 auto ulock = data_.ulock();
45 if (ulock->count(val)) {
49 auto wlock = ulock.moveFromUpgradeToWrite();
50 (*wlock)[val] = std::move(data);
54 bool RequestContext::hasContextData(const std::string& val) const {
55 return data_.rlock()->count(val);
58 RequestData* RequestContext::getContextData(const std::string& val) {
59 const std::unique_ptr<RequestData> dflt{nullptr};
60 return get_ref_default(*data_.rlock(), val, dflt).get();
63 const RequestData* RequestContext::getContextData(
64 const std::string& val) const {
65 const std::unique_ptr<RequestData> dflt{nullptr};
66 return get_ref_default(*data_.rlock(), val, dflt).get();
69 void RequestContext::onSet() {
70 auto rlock = data_.rlock();
71 for (auto const& ent : *rlock) {
72 if (auto& data = ent.second) {
78 void RequestContext::onUnset() {
79 auto rlock = data_.rlock();
80 for (auto const& ent : *rlock) {
81 if (auto& data = ent.second) {
87 void RequestContext::clearContextData(const std::string& val) {
88 std::unique_ptr<RequestData> requestData;
89 // Delete the RequestData after giving up the wlock just in case one of the
90 // RequestData destructors will try to grab the lock again.
92 auto wlock = data_.wlock();
93 auto it = wlock->find(val);
94 if (it != wlock->end()) {
95 requestData = std::move(it->second);
101 std::shared_ptr<RequestContext> RequestContext::setContext(
102 std::shared_ptr<RequestContext> ctx) {
103 auto& curCtx = getStaticContext();
105 FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
118 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
119 using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
120 static SingletonT singleton;
122 return singleton.get();
125 RequestContext* RequestContext::get() {
126 auto context = getStaticContext();
128 static RequestContext defaultContext;
129 return std::addressof(defaultContext);
131 return context.get();