io/async: Add support for RequestData::onClear() callback function

Summary:
Add support for onClear() callback function to RequestData. The callback is to be executed on the last clearing of RequestData. This typically happens upon the destruction of the last RequestContext that contains the RequestData or due to explicit calls to RequestContext::clearContextData that remove the RequestData from the RequestContext.

The onClear() callback is executed exactly once. If there are no calls to clearContextData, then onClear is executed upon the destruction of the last RequestContext that refers to the RequestData. If there are calls to clearContextData, then onClear is executed upon the removal of the last reference to the RequestData from a RequestContext.

RequestData::hasCallback() applies only to onSet() and onUnset() which may be called multiple times, whereas onClear() is called exactly once.

If a class derived from RequestData overrides onClear and either or both of onSet and onUnset, then care must be taken that the callbacks are mutually safe, as it is possible for the execution of onClear to be concurrent with executions of onSet and onUnset.

Motivation: Some uses of RequestContext::clearContextData expect the immediate destruction of associated RequestData. The hazard pointer-based implementation of RequestContext guarantees only that all associated RequestData is destroyed before the completion of the destructor of the RequestContext. This callback provides users of RequestContext more control of the timing of executing actions related to clearing the data. The addition of this capability is intended to provide a tool to eliminate differences in timing of executing actions between the two implementations of RequestContext (hazard pointer-based vs. read lock-based).

Also, fixes a minor bug in request_context_test, where the key "test2" was used in two tests (setIfAbsentTest and deadlockTest) without being cleared, which generated a warning when the tests were run in the same process.

Reviewed By: davidtgoldblatt

Differential Revision: D20332008

fbshipit-source-id: c9cefffb28c8da098e97f851f31d29a1bf704061

folly/io/async/Request.cpp

@@ -75,14 +75,40 @@ Synchronized<F14FastMap<std::string, uint32_t>>& RequestToken::getCache() {
 }
 
 void RequestData::acquireRef() {
-  auto rc = keepAliveCounter_.fetch_add(1, std::memory_order_relaxed);
+  auto rc = keepAliveCounter_.fetch_add(
+      kClearCount + kDeleteCount, std::memory_order_relaxed);
   DCHECK_GE(rc, 0);
 }
 
-void RequestData::releaseRefDeleteIfNoRefs() {
-  auto rc = keepAliveCounter_.fetch_sub(1, std::memory_order_acq_rel);
+void RequestData::releaseRefClearOnly() {
+  auto rc =
+      keepAliveCounter_.fetch_sub(kClearCount, std::memory_order_acq_rel) -
+      kClearCount;
   DCHECK_GT(rc, 0);
-  if (rc == 1) {
+  if (rc < kClearCount) {
+    this->onClear();
+  }
+}
+
+void RequestData::releaseRefDeleteOnly() {
+  auto rc =
+      keepAliveCounter_.fetch_sub(kDeleteCount, std::memory_order_acq_rel) -
+      kDeleteCount;
+  DCHECK_GE(rc, 0);
+  if (rc == 0) {
+    delete this;
+  }
+}
+
+void RequestData::releaseRefClearDelete() {
+  auto rc = keepAliveCounter_.fetch_sub(
+                kClearCount + kDeleteCount, std::memory_order_acq_rel) -
+      (kClearCount + kDeleteCount);
+  DCHECK_GE(rc, 0);
+  if (rc < kClearCount) {
+    this->onClear();
+  }
+  if (rc == 0) {
     delete this;
   }
 }
@@ -94,6 +120,7 @@ void RequestData::DestructPtr::operator()(RequestData* ptr) {
     // Note: this is the value before decrement, hence == 1 check
     DCHECK(keepAliveCounter > 0);
     if (keepAliveCounter == 1) {
+      ptr->onClear();
       delete ptr;
     }
   }
@@ -121,6 +148,12 @@ struct RequestContext::StateHazptr::Combined : hazptr_obj_base<Combined> {
   SingleWriterFixedHashMap<RequestToken, RequestData*> requestData_;
   // This must be optimized for iteration, its hot path is setContext
   SingleWriterFixedHashMap<RequestData*, bool> callbackData_;
+  // Hash map to keep track of Clear and Delete counts. The presence
+  // of a key indicates holding a Delete count for the request data
+  // (i.e., delete when the Delete count goes to zero). A value of
+  // true indicates holding a Clear counts for the request data (i.e.,
+  // call onClear() when the Clear count goes to zero).
+  F14FastMap<RequestData*, bool> refs_;
 
   Combined()
       : requestData_(kInitialCapacity), callbackData_(kInitialCapacity) {}
@@ -145,6 +178,7 @@ struct RequestContext::StateHazptr::Combined : hazptr_obj_base<Combined> {
     for (auto it = requestData_.begin(); it != requestData_.end(); ++it) {
       auto p = it.value();
       if (p) {
+        refs_.insert({p, true});
         p->acquireRef();
       }
     }
@@ -152,10 +186,11 @@ struct RequestContext::StateHazptr::Combined : hazptr_obj_base<Combined> {
 
   /* releaseDataRefs - Called only once from ~Combined */
   void releaseDataRefs() {
-    for (auto it = requestData_.begin(); it != requestData_.end(); ++it) {
-      auto p = it.value();
-      if (p) {
-        p->releaseRefDeleteIfNoRefs();
+    for (auto pair : refs_) {
+      if (pair.second) {
+        pair.first->releaseRefClearDelete();
+      } else {
+        pair.first->releaseRefDeleteOnly();
       }
     }
   }
@@ -315,7 +350,8 @@ RequestContext::StateHazptr::eraseOldData(
     // If the caller guarantees thread-safety, then erase the
     // entry in the current version.
     cur->requestData_.erase(token);
-    olddata->releaseRefDeleteIfNoRefs();
+    cur->refs_.erase(olddata);
+    olddata->releaseRefClearDelete();
   } else {
     // If there may be concurrent readers, then copy-on-erase.
     // Update the data reference counts to account for the
@@ -347,6 +383,7 @@ RequestContext::StateHazptr::insertNewData(
     data->onSet();
   }
   if (data) {
+    cur->refs_.insert({data.get(), true});
     data->acquireRef();
   }
   cur->requestData_.insert(token, data.release());
@@ -416,6 +453,7 @@ void RequestContext::StateHazptr::onUnset() {
 }
 
 void RequestContext::StateHazptr::clearContextData(const RequestToken& token) {
+  RequestData* data;
   Combined* replaced = nullptr;
   { // Lock mutex_
     std::lock_guard<std::mutex> g(mutex_);
@@ -427,7 +465,7 @@ void RequestContext::StateHazptr::clearContextData(const RequestToken& token) {
     if (it == cur->requestData_.end()) {
       return;
     }
-    RequestData* data = it.value();
+    data = it.value();
     if (!data) {
       cur->requestData_.erase(token);
       return;
@@ -442,6 +480,9 @@ void RequestContext::StateHazptr::clearContextData(const RequestToken& token) {
     cur->acquireDataRefs();
     setCombined(cur);
   } // Unlock mutex_
+  DCHECK(data);
+  data->releaseRefClearOnly();
+  replaced->refs_[data] = false; // Clear reference already released
   DCHECK(replaced);
   replaced->retire();
 }

folly/io/async/Request.h

@@ -93,6 +93,8 @@ class RequestData {
   // could fix these deadlocks, but only at significant performance penalty, so
   // just don't do it!
 
+  // hasCallback() applies only to onSet() and onUnset().
+  // onClear() is always executed exactly once.
   virtual bool hasCallback() = 0;
   // Callback executed when setting RequestContext. Make sure your RequestData
   // instance overrides the hasCallback method to return true otherwise
@@ -102,6 +104,13 @@ class RequestData {
   // instance overrides the hasCallback method to return true otherwise
   // the callback will not be executed
   virtual void onUnset() {}
+  // Callback executed exactly once upon the release of the last
+  // reference to the request data (as a result of either a call to
+  // clearContextData or the destruction of a request context that
+  // contains a reference to the data). It can be overridden in
+  // derived classes. There may be concurrent executions of onSet()
+  // and onUnset() with that of onClear().
+  virtual void onClear() {}
   // For debugging
   int refCount() {
     return keepAliveCounter_.load(std::memory_order_acquire);
@@ -116,11 +125,18 @@ class RequestData {
 
   friend class RequestContext;
 
+  static constexpr int kDeleteCount = 0x1;
+  static constexpr int kClearCount = 0x1000;
+
   // Reference-counting functions used by the hazptr-based implementation.
   // Increment the reference count
   void acquireRef();
-  // Decrement the reference count and delete if zero
-  void releaseRefDeleteIfNoRefs();
+  // Decrement the reference count. Clear only if last.
+  void releaseRefClearOnly();
+  // Decrement the reference count. Delete only if last.
+  void releaseRefDeleteOnly();
+  // Decrement the reference count. Clear and delete if last.
+  void releaseRefClearDelete();
 
   // Unique ptr with custom destructor, decrement the counter
   // and only free if 0
@@ -379,11 +395,19 @@ class RequestContext {
     StateHazptr& operator=(StateHazptr&&) = delete;
     ~StateHazptr();
 
+   private:
+    friend class RequestContext;
+
+    struct SetContextDataResult {
+      bool changed; // Changes were made
+      bool unexpected; // Update was unexpected
+      Combined* replaced; // The combined structure was replaced
+    };
+
     Combined* combined() const;
     Combined* ensureCombined(); // Lazy allocation if needed
     void setCombined(Combined* p);
     Combined* expand(Combined* combined);
-
     bool doSetContextData(
         const RequestToken& token,
         std::unique_ptr<RequestData>& data,
@@ -395,14 +419,6 @@ class RequestContext {
     void onSet();
     void onUnset();
     void clearContextData(const RequestToken& token);
-
-   private:
-    struct SetContextDataResult {
-      bool changed; // Changes were made
-      bool unexpected; // Update was unexpected
-      Combined* replaced; // The combined structure was replaced
-    };
-
     SetContextDataResult doSetContextDataHelper(
         const RequestToken& token,
         std::unique_ptr<RequestData>& data,

folly/io/async/test/RequestContextTest.cpp

@@ -233,7 +233,7 @@ TEST_F(RequestContextTest, deadlockTest) {
   };
 
   RequestContext::get()->setContextData(
-      "test", std::make_unique<DeadlockTestData>("test2"));
+      "test", std::make_unique<DeadlockTestData>("test1"));
   RequestContext::get()->clearContextData(testtoken);
 }
 
@@ -399,3 +399,50 @@ TEST_F(RequestContextTest, ThreadId) {
   EXPECT_EQ(*folly::getThreadName(rootids[0].tid), "DummyThread");
   EXPECT_FALSE(folly::getThreadName(rootids[1].tid));
 }
+
+TEST_F(RequestContextTest, Clear) {
+  struct Foo : public RequestData {
+    bool& cleared;
+    bool& deleted;
+    Foo(bool& c, bool& d) : cleared(c), deleted(d) {}
+    ~Foo() override {
+      EXPECT_TRUE(cleared);
+      deleted = true;
+    }
+    bool hasCallback() override {
+      return false;
+    }
+    void onClear() override {
+      EXPECT_FALSE(cleared);
+      cleared = true;
+    }
+  };
+
+  std::string key = "clear";
+  {
+    bool cleared = false;
+    bool deleted = false;
+    {
+      RequestContextScopeGuard g;
+      RequestContext::get()->setContextData(
+          key, std::make_unique<Foo>(cleared, deleted));
+      EXPECT_FALSE(cleared);
+      RequestContext::get()->clearContextData(key);
+      EXPECT_TRUE(cleared);
+    }
+    EXPECT_TRUE(deleted);
+  }
+  {
+    bool cleared = false;
+    bool deleted = false;
+    {
+      RequestContextScopeGuard g;
+      RequestContext::get()->setContextData(
+          key, std::make_unique<Foo>(cleared, deleted));
+      EXPECT_FALSE(cleared);
+      EXPECT_FALSE(deleted);
+    }
+    EXPECT_TRUE(cleared);
+    EXPECT_TRUE(deleted);
+  }
+}