Add RequestContext::try_get()

Summary:
Returns the `RequestContext` object if it already exists. Uses
`SingletonThreadLocal::try_get()` underneath.

Reviewed By: yfeldblum

Differential Revision: D31006847

fbshipit-source-id: 2ce67a097c07de10b4359e3e70b820b80d885117

folly/io/async/Request.cpp

@@ -608,10 +608,15 @@ void RequestContext::clearContextData(const RequestToken& val) {
   return prevCtx;
 }
 
-RequestContext::StaticContext& RequestContext::getStaticContext() {
+/* static */ RequestContext::StaticContext& RequestContext::getStaticContext() {
   return StaticContextThreadLocal::get();
 }
 
+/* static */ RequestContext::StaticContext*
+RequestContext::tryGetStaticContext() {
+  return StaticContextThreadLocal::try_get();
+}
+
 /* static */ RequestContext::StaticContextAccessor
 RequestContext::accessAllThreads() {
   return StaticContextAccessor{StaticContextThreadLocal::accessAllThreads()};
@@ -641,7 +646,7 @@ RequestContext::setShallowCopyContext() {
   return child;
 }
 
-RequestContext* RequestContext::get() {
+/* static */ RequestContext* RequestContext::get() {
   auto& context = getStaticContext().requestContext;
   if (!context) {
     static RequestContext defaultContext(0);
@@ -649,4 +654,12 @@ RequestContext* RequestContext::get() {
   }
   return context.get();
 }
+
+/* static */ RequestContext* RequestContext::try_get() {
+  if (auto* staticContext = tryGetStaticContext()) {
+    return staticContext->requestContext.get();
+  }
+  return nullptr;
+}
+
 } // namespace folly

folly/io/async/Request.h

@@ -154,6 +154,9 @@ class RequestContext {
   // Get the current context.
   static RequestContext* get();
 
+  // Get the current context, if it has already been created, or nullptr.
+  static RequestContext* try_get();
+
   intptr_t getRootId() const { return rootId_; }
 
   struct RootIdInfo {
@@ -252,6 +255,7 @@ class RequestContext {
 
  private:
   static StaticContext& getStaticContext();
+  static StaticContext* tryGetStaticContext();
   static std::shared_ptr<RequestContext> setContextHelper(
       std::shared_ptr<RequestContext>& newCtx, StaticContext& staticCtx);
 

folly/io/async/test/RequestContextTest.cpp

@@ -566,3 +566,22 @@ TEST_F(RequestContextTest, AccessAllThreadsDestructionGuard) {
     thread.join();
   }
 }
+
+TEST(RequestContextTryGetTest, TryGetTest) {
+  // try_get() should not create a default RequestContext object if none exists.
+  EXPECT_EQ(RequestContext::try_get(), nullptr);
+  // Explicitly create a new instance so that subsequent calls to try_get()
+  // return it.
+  RequestContext::create();
+  EXPECT_NE(RequestContext::saveContext(), nullptr);
+  EXPECT_NE(RequestContext::try_get(), nullptr);
+  // Make sure that the pointers returned by both get() and try_get() point to
+  // the same underlying instance.
+  EXPECT_EQ(RequestContext::try_get(), RequestContext::get());
+  // Set some context data and read it out via try_get() accessor.
+  RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
+  auto rc = RequestContext::try_get();
+  EXPECT_TRUE(rc->hasContextData("test"));
+  auto* dataPtr = dynamic_cast<TestData*>(rc->getContextData("test"));
+  EXPECT_EQ(dataPtr->data_, 10);
+}