2017

[folly.git] / folly / io / async / Request.cpp

diff --git a/folly/io/async/Request.cpp b/folly/io/async/Request.cpp
index 97e2dda4f5f43b4db5dcc3c636e975c83871d011..8310033b2020529c1a9f16453f44c0a521a27455 100644 (file)
--- a/folly/io/async/Request.cpp
+++ b/folly/io/async/Request.cpp
@@ -1,4 +1,6 @@
 /*
+ * Copyright 2017 Facebook, Inc.
+ *
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements. See the NOTICE file
  * distributed with this work for additional information
@@ -16,19 +18,109 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-#include "folly/io/async/Request.h"
 
-#ifndef NO_LIB_GFLAGS
-  DEFINE_bool(enable_request_context, true,
-              "Enable collection of per-request queueing stats for thrift");
-#endif
+#include <folly/io/async/Request.h>
+#include <folly/tracing/StaticTracepoint.h>
+
+#include <glog/logging.h>
+
+#include <folly/MapUtil.h>
+#include <folly/SingletonThreadLocal.h>
 
 namespace folly {
 
-#ifdef NO_LIB_GFLAGS
-  bool FLAGS_enable_thrift_request_context = true;
-#endif
+void RequestContext::setContextData(
+    const std::string& val,
+    std::unique_ptr<RequestData> data) {
+  auto wlock = data_.wlock();
+  if (wlock->count(val)) {
+    LOG_FIRST_N(WARNING, 1)
+        << "Called RequestContext::setContextData with data already set";
 
-RequestContext* defaultContext;
+    (*wlock)[val] = nullptr;
+  } else {
+    (*wlock)[val] = std::move(data);
+  }
+}
+
+bool RequestContext::setContextDataIfAbsent(
+    const std::string& val,
+    std::unique_ptr<RequestData> data) {
+  auto ulock = data_.ulock();
+  if (ulock->count(val)) {
+    return false;
+  }
+
+  auto wlock = ulock.moveFromUpgradeToWrite();
+  (*wlock)[val] = std::move(data);
+  return true;
+}
+
+bool RequestContext::hasContextData(const std::string& val) const {
+  return data_.rlock()->count(val);
+}
+
+RequestData* RequestContext::getContextData(const std::string& val) {
+  return get_ref_default(*data_.rlock(), val, nullptr).get();
+}
+
+const RequestData* RequestContext::getContextData(
+    const std::string& val) const {
+  return get_ref_default(*data_.rlock(), val, nullptr).get();
+}
 
+void RequestContext::onSet() {
+  auto rlock = data_.rlock();
+  for (auto const& ent : *rlock) {
+    if (auto& data = ent.second) {
+      data->onSet();
+    }
+  }
+}
+
+void RequestContext::onUnset() {
+  auto rlock = data_.rlock();
+  for (auto const& ent : *rlock) {
+    if (auto& data = ent.second) {
+      data->onUnset();
+    }
+  }
+}
+
+void RequestContext::clearContextData(const std::string& val) {
+  data_.wlock()->erase(val);
+}
+
+std::shared_ptr<RequestContext> RequestContext::setContext(
+    std::shared_ptr<RequestContext> ctx) {
+  auto& curCtx = getStaticContext();
+  if (ctx != curCtx) {
+    FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
+    using std::swap;
+    if (curCtx) {
+      curCtx->onUnset();
+    }
+    swap(ctx, curCtx);
+    if (curCtx) {
+      curCtx->onSet();
+    }
+  }
+  return ctx;
+}
+
+std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
+  using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
+  static SingletonT singleton;
+
+  return singleton.get();
+}
+
+RequestContext* RequestContext::get() {
+  auto context = getStaticContext();
+  if (!context) {
+    static RequestContext defaultContext;
+    return std::addressof(defaultContext);
+  }
+  return context.get();
+}
 }