From b5338e82505238ec3f4d4639098127e326344870 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@fb.com>
Date: Sun, 2 Aug 2015 16:15:04 -0700
Subject: [PATCH] Fix ASAN failure in
 folly/io/async/test/NotificationQueueTest.cpp.

Summary: [Folly] Fix ASAN failure in folly/io/async/test/NotificationQueueTest.cpp.

Reviewed By: @djwatson

Differential Revision: D2299112
---
 folly/io/async/NotificationQueue.h            | 21 +++++++++++++------
 folly/io/async/test/NotificationQueueTest.cpp | 16 ++++++++------
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/folly/io/async/NotificationQueue.h b/folly/io/async/NotificationQueue.h
index 8a7219e5..b0a2bd73 100644
--- a/folly/io/async/NotificationQueue.h
+++ b/folly/io/async/NotificationQueue.h
@@ -21,6 +21,7 @@
 
 #include <folly/io/async/EventBase.h>
 #include <folly/io/async/EventHandler.h>
+#include <folly/io/async/DelayedDestruction.h>
 #include <folly/io/async/Request.h>
 #include <folly/Likely.h>
 #include <folly/ScopeGuard.h>
@@ -60,7 +61,7 @@ class NotificationQueue {
   /**
    * A callback interface for consuming messages from the queue as they arrive.
    */
-  class Consumer : private EventHandler {
+  class Consumer : public DelayedDestruction, private EventHandler {
    public:
     enum : uint16_t { kDefaultMaxReadAtOnce = 10 };
 
@@ -69,8 +70,6 @@ class NotificationQueue {
         destroyedFlagPtr_(nullptr),
         maxReadAtOnce_(kDefaultMaxReadAtOnce) {}
 
-    virtual ~Consumer();
-
     /**
      * messageAvailable() will be invoked whenever a new
      * message is available from the pipe.
@@ -152,7 +151,13 @@ class NotificationQueue {
       return base_;
     }
 
-    virtual void handlerReady(uint16_t events) noexcept;
+    void handlerReady(uint16_t events) noexcept override;
+
+   protected:
+
+    void destroy() override;
+
+    virtual ~Consumer() {}
 
    private:
     /**
@@ -577,7 +582,7 @@ class NotificationQueue {
 };
 
 template<typename MessageT>
-NotificationQueue<MessageT>::Consumer::~Consumer() {
+void NotificationQueue<MessageT>::Consumer::destroy() {
   // If we are in the middle of a call to handlerReady(), destroyedFlagPtr_
   // will be non-nullptr.  Mark the value that it points to, so that
   // handlerReady() will know the callback is destroyed, and that it cannot
@@ -585,6 +590,8 @@ NotificationQueue<MessageT>::Consumer::~Consumer() {
   if (destroyedFlagPtr_) {
     *destroyedFlagPtr_ = true;
   }
+  stopConsuming();
+  DelayedDestruction::destroy();
 }
 
 template<typename MessageT>
@@ -596,6 +603,7 @@ void NotificationQueue<MessageT>::Consumer::handlerReady(uint16_t /*events*/)
 template<typename MessageT>
 void NotificationQueue<MessageT>::Consumer::consumeMessages(
     bool isDrain, size_t* numConsumed) noexcept {
+  DestructorGuard dg(this);
   uint32_t numProcessed = 0;
   bool firstRun = true;
   setActive(true);
@@ -661,6 +669,7 @@ void NotificationQueue<MessageT>::Consumer::consumeMessages(
       CHECK(destroyedFlagPtr_ == nullptr);
       destroyedFlagPtr_ = &callbackDestroyed;
       messageAvailable(std::move(msg));
+      destroyedFlagPtr_ = nullptr;
 
       RequestContext::setContext(old_ctx);
 
@@ -668,7 +677,6 @@ void NotificationQueue<MessageT>::Consumer::consumeMessages(
       if (callbackDestroyed) {
         return;
       }
-      destroyedFlagPtr_ = nullptr;
 
       // If the callback is no longer installed, we are done.
       if (queue_ == nullptr) {
@@ -767,6 +775,7 @@ void NotificationQueue<MessageT>::Consumer::stopConsuming() {
 template<typename MessageT>
 bool NotificationQueue<MessageT>::Consumer::consumeUntilDrained(
     size_t* numConsumed) noexcept {
+  DestructorGuard dg(this);
   {
     folly::SpinLockGuard g(queue_->spinlock_);
     if (queue_->draining_) {
diff --git a/folly/io/async/test/NotificationQueueTest.cpp b/folly/io/async/test/NotificationQueueTest.cpp
index c831ee10..2bfc1365 100644
--- a/folly/io/async/test/NotificationQueueTest.cpp
+++ b/folly/io/async/test/NotificationQueueTest.cpp
@@ -360,15 +360,16 @@ void QueueTest::destroyCallback() {
   // avoid destroying the function object.
   class DestroyTestConsumer : public IntQueue::Consumer {
    public:
-    DestroyTestConsumer() {}
-
     void messageAvailable(int&& value) override {
+      DestructorGuard g(this);
       if (fn && *fn) {
         (*fn)(value);
       }
     }
 
     std::function<void(int)> *fn;
+   protected:
+    virtual ~DestroyTestConsumer() = default;
   };
 
   EventBase eventBase;
@@ -381,11 +382,13 @@ void QueueTest::destroyCallback() {
   // This way one consumer will be destroyed from inside its messageAvailable()
   // callback, and one consume will be destroyed when it isn't inside
   // messageAvailable().
-  std::unique_ptr<DestroyTestConsumer> consumer1(new DestroyTestConsumer);
-  std::unique_ptr<DestroyTestConsumer> consumer2(new DestroyTestConsumer);
+  std::unique_ptr<DestroyTestConsumer, DelayedDestruction::Destructor>
+    consumer1(new DestroyTestConsumer);
+  std::unique_ptr<DestroyTestConsumer, DelayedDestruction::Destructor>
+    consumer2(new DestroyTestConsumer);
   std::function<void(int)> fn = [&](int) {
-    consumer1.reset();
-    consumer2.reset();
+    consumer1 = nullptr;
+    consumer2 = nullptr;
   };
   consumer1->fn = &fn;
   consumer2->fn = &fn;
@@ -617,6 +620,7 @@ TEST(NotificationQueueTest, UseAfterFork) {
       // We shouldn't reach here.
       _exit(0);
     }
+    PCHECK(pid > 0);
 
     // Parent.  Wait for the child to exit.
     auto waited = waitpid(pid, &childStatus, 0);
-- 
2.34.1