RequestContext::StaticContextAccessor

Summary: Redo of D29536635 (https://github.com/facebook/folly/commit/bc0818f89e7ab3802cdac95a25e29f4061021c3c) but rebased on top of `UniqueInstance.h` fixes.

Reviewed By: yfeldblum

Differential Revision: D29630206

fbshipit-source-id: ed103849edd6edd3f9c4ef75368c2fa1b038182b

folly/io/async/Request.cpp

@@ -18,18 +18,12 @@
 
 #include <folly/GLog.h>
 #include <folly/MapUtil.h>
-#include <folly/SingletonThreadLocal.h>
 #include <folly/experimental/SingleWriterFixedHashMap.h>
 #include <folly/synchronization/Hazptr.h>
 #include <folly/tracing/StaticTracepoint.h>
 
 namespace folly {
 
-namespace {
-using SingletonT =
-    SingletonThreadLocal<RequestContext::StaticContext, RequestContext>;
-}
-
 RequestToken::RequestToken(const std::string& str) {
   auto& cache = getCache();
   {
@@ -615,18 +609,20 @@ void RequestContext::clearContextData(const RequestToken& val) {
 }
 
 RequestContext::StaticContext& RequestContext::getStaticContext() {
-  return SingletonT::get();
+  return StaticContextThreadLocal::get();
+}
+
+/* static */ RequestContext::StaticContextAccessor
+RequestContext::accessAllThreads() {
+  return StaticContextAccessor{StaticContextThreadLocal::accessAllThreads()};
 }
 
 /* static */ std::vector<RequestContext::RootIdInfo>
 RequestContext::getRootIdsFromAllThreads() {
   std::vector<RootIdInfo> result;
-  auto accessor = SingletonT::accessAllThreads();
+  auto accessor = RequestContext::accessAllThreads();
   for (auto it = accessor.begin(); it != accessor.end(); ++it) {
-    result.push_back(
-        {it->rootId.load(std::memory_order_relaxed),
-         it.getThreadId(),
-         it.getOSThreadId()});
+    result.push_back(it.getRootIdInfo());
   }
   return result;
 }

folly/io/async/Request.h

@@ -22,8 +22,10 @@
 #include <string>
 #include <utility>
 
+#include <folly/SingletonThreadLocal.h>
 #include <folly/Synchronized.h>
 #include <folly/container/F14Map.h>
+#include <folly/detail/Iterators.h>
 #include <folly/synchronization/Hazptr.h>
 
 namespace folly {
@@ -253,6 +255,60 @@ class RequestContext {
   static std::shared_ptr<RequestContext> setContextHelper(
       std::shared_ptr<RequestContext>& newCtx, StaticContext& staticCtx);
 
+  using StaticContextThreadLocal = SingletonThreadLocal<
+      RequestContext::StaticContext,
+      RequestContext /* Tag */>;
+
+ public:
+  class StaticContextAccessor {
+   private:
+    using Inner = StaticContextThreadLocal::Accessor;
+    using IteratorBase = Inner::Iterator;
+    using IteratorTag = std::bidirectional_iterator_tag;
+
+    Inner inner_;
+
+    explicit StaticContextAccessor(Inner&& inner) noexcept
+        : inner_(std::move(inner)) {}
+
+   public:
+    friend class RequestContext;
+
+    class Iterator : public detail::IteratorAdaptor<
+                         Iterator,
+                         IteratorBase,
+                         StaticContext,
+                         IteratorTag> {
+      using Super = detail::
+          IteratorAdaptor<Iterator, IteratorBase, StaticContext, IteratorTag>;
+
+     public:
+      using Super::Super;
+
+      StaticContext& dereference() const { return *base(); }
+
+      RootIdInfo getRootIdInfo() const {
+        return {
+            base()->rootId.load(std::memory_order_relaxed),
+            base().getThreadId(),
+            base().getOSThreadId()};
+      }
+    };
+
+    StaticContextAccessor(const StaticContextAccessor&) = delete;
+    StaticContextAccessor& operator=(const StaticContextAccessor&) = delete;
+    StaticContextAccessor(StaticContextAccessor&&) = default;
+    StaticContextAccessor& operator=(StaticContextAccessor&&) = default;
+
+    Iterator begin() const { return Iterator(inner_.begin()); }
+    Iterator end() const { return Iterator(inner_.end()); }
+  };
+  // Returns an accessor object that blocks the construction and destruction of
+  // StaticContext objects on all threads. This is useful to quickly introspect
+  // the context from all threads while ensuring that their thread-local
+  // StaticContext object is not destroyed.
+  static StaticContextAccessor accessAllThreads();
+
   // Start shallow copy guard implementation details:
   // All methods are private to encourage proper use
   friend struct ShallowCopyRequestContextScopeGuard;

folly/io/async/test/RequestContextTest.cpp

@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include <atomic>
+#include <cstdint>
 #include <thread>
 
 #include <folly/Memory.h>
@@ -23,6 +25,8 @@
 #include <folly/portability/GTest.h>
 #include <folly/system/ThreadName.h>
 
+#include <boost/thread/barrier.hpp>
+
 using namespace folly;
 
 RequestToken testtoken("test");
@@ -523,3 +527,42 @@ TEST_F(RequestContextTest, ConcurrentDataRefRelease) {
     th2.join();
   }
 }
+
+TEST_F(RequestContextTest, AccessAllThreadsDestructionGuard) {
+  constexpr auto kNumThreads = 128;
+
+  std::vector<std::thread> threads{kNumThreads};
+  boost::barrier barrier{kNumThreads + 1};
+
+  std::atomic<std::size_t> count{0};
+  for (auto& thread : threads) {
+    thread = std::thread([&] {
+      // Force creation of thread local
+      RequestContext::get();
+      ++count;
+      // Wait for all other threads to do the same
+      barrier.wait();
+      // Wait until signaled to die
+      barrier.wait();
+    });
+  }
+
+  barrier.wait();
+  // Sanity check
+  EXPECT_EQ(count.load(), kNumThreads);
+
+  {
+    auto accessor = RequestContext::accessAllThreads();
+    // Allow threads to die (but they should not as long as we hold accessor!)
+    barrier.wait();
+    auto accessorsCount = std::distance(accessor.begin(), accessor.end());
+    EXPECT_EQ(accessorsCount, kNumThreads + 1);
+    for (RequestContext::StaticContext& staticContext : accessor) {
+      EXPECT_EQ(staticContext.requestContext, nullptr);
+    }
+  }
+
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}