From f430679a0d826faf77c999803556dec21ed7b829 Mon Sep 17 00:00:00 2001
From: Zizhong Zhang <zizhang@linkedin.com>
Date: Wed, 22 Aug 2018 16:58:48 -0700
Subject: [PATCH] revoke promise when provider is gone

---
 include/tscpp/api/Async.h | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/include/tscpp/api/Async.h b/include/tscpp/api/Async.h
index d5227f0b0fb..b392f494eda 100644
--- a/include/tscpp/api/Async.h
+++ b/include/tscpp/api/Async.h
@@ -23,9 +23,9 @@
 
 #pragma once
 
-#include <list>
 #include <memory>
 #include <mutex>
+#include <unordered_map>
 
 #include "tscpp/api/noncopyable.h"
 
@@ -91,7 +91,7 @@ class AsyncProvider
     }
   }
 
-  virtual ~AsyncProvider() {}
+  virtual ~AsyncProvider() { this->cancel(); }
 
 protected:
   std::shared_ptr<AsyncDispatchControllerBase>
@@ -137,7 +137,10 @@ class AsyncDispatchController : public AsyncDispatchControllerBase
   disable() override
   {
     std::lock_guard<Mutex> scopedLock(*dispatch_mutex_);
-    event_receiver_ = nullptr;
+    if (event_receiver_ != nullptr) {
+      event_receiver_->revokePromise(this);
+      event_receiver_ = nullptr;
+    }
   }
 
   bool
@@ -209,13 +212,20 @@ template <typename AsyncProviderType> class AsyncReceiver : noncopyable
    */
   virtual void handleAsyncComplete(AsyncProviderType &provider) = 0;
   virtual ~AsyncReceiver() {}
+  void
+  revokePromise(AsyncDispatchController<AsyncReceiver<AsyncProviderType>, AsyncProviderType> *dispatch_controller_ptr)
+  {
+    receiver_promises_.erase(dispatch_controller_ptr);
+  }
 
 protected:
   AsyncReceiver() {}
   friend class Async;
 
 private:
-  mutable std::list<std::shared_ptr<AsyncReceiverPromise<AsyncReceiver<AsyncProviderType>, AsyncProviderType>>> receiver_promises_;
+  mutable std::unordered_map<AsyncDispatchController<AsyncReceiver<AsyncProviderType>, AsyncProviderType> *,
+                             std::shared_ptr<AsyncReceiverPromise<AsyncReceiver<AsyncProviderType>, AsyncProviderType>>>
+    receiver_promises_;
 };
 
 /**
@@ -246,7 +256,7 @@ class Async : noncopyable
       new AsyncDispatchController<AsyncReceiver<AsyncProviderType>, AsyncProviderType>(event_receiver, provider, mutex));
     std::shared_ptr<AsyncReceiverPromise<AsyncReceiver<AsyncProviderType>, AsyncProviderType>> receiver_promise(
       new AsyncReceiverPromise<AsyncReceiver<AsyncProviderType>, AsyncProviderType>(dispatcher));
-    event_receiver->receiver_promises_.push_back(receiver_promise); // now if the event receiver dies, we're safe.
+    event_receiver->receiver_promises_[dispatcher.get()] = receiver_promise;
     provider->doRun(dispatcher);
   }
 };