Call onSet and onUnset when adding and and deleting RequestData

Summary: Currently, when we add or delete a `RequestData` instance to `RequestContext` for the first time, we do not call the `RequestData`'s `onSet` or `onUnset` callback. This means if the `RequestData` is doing tracking-related work, it may miss the beginning or the request until the first switch of `RequestContext` has happend, and in the end as well. This Diff fixes the issue by calling `onSet` and `onUnset` when adding or deleting.

Reviewed By: anakryiko

Differential Revision: D7978623

fbshipit-source-id: a524286371f890598bc773e957ef609e90086cea

folly/io/async/Request.cpp

@@ -46,6 +46,7 @@ bool RequestContext::doSetContextData(
   if (conflict) {
     if (it->second) {
       if (it->second->hasCallback()) {
+        it->second->onUnset();
         wlock->callbackData_.erase(it->second.get());
       }
       it->second.reset(nullptr);
@@ -55,6 +56,7 @@ bool RequestContext::doSetContextData(
 
   if (data && data->hasCallback()) {
     wlock->callbackData_.insert(data.get());
+    data->onSet();
   }
   wlock->requestData_[val] = std::move(data);
 
@@ -128,6 +130,7 @@ void RequestContext::clearContextData(const std::string& val) {
 
     auto wlock = ulock.moveFromUpgradeToWrite();
     if (it->second && it->second->hasCallback()) {
+      it->second->onUnset();
       wlock->callbackData_.erase(it->second.get());
     }
 

folly/io/async/test/RequestContextTest.cpp

@@ -109,26 +109,30 @@ TEST(RequestContext, testSetUnset) {
   ctx1->setContextData("test", std::make_unique<TestData>(10));
   auto testData1 = dynamic_cast<TestData*>(ctx1->getContextData("test"));
 
+  // onSet called in setContextData
+  EXPECT_EQ(1, testData1->set_);
+
   // Override RequestContext
   RequestContext::create();
   auto ctx2 = RequestContext::saveContext();
   ctx2->setContextData("test", std::make_unique<TestData>(20));
   auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
 
+  // onSet called in setContextData
+  EXPECT_EQ(1, testData2->set_);
+
   // Check ctx1->onUnset was called
-  EXPECT_EQ(0, testData1->set_);
   EXPECT_EQ(1, testData1->unset_);
 
   RequestContext::setContext(ctx1);
-  EXPECT_EQ(1, testData1->set_);
+  EXPECT_EQ(2, testData1->set_);
   EXPECT_EQ(1, testData1->unset_);
-  EXPECT_EQ(0, testData2->set_);
   EXPECT_EQ(1, testData2->unset_);
 
   RequestContext::setContext(ctx2);
-  EXPECT_EQ(1, testData1->set_);
+  EXPECT_EQ(2, testData1->set_);
   EXPECT_EQ(2, testData1->unset_);
-  EXPECT_EQ(1, testData2->set_);
+  EXPECT_EQ(2, testData2->set_);
   EXPECT_EQ(1, testData2->unset_);
 }