From 3945a65b62316c92066bd499a0e5d780b7106b0e Mon Sep 17 00:00:00 2001
From: Francois Ferrand <thetypz@gmail.com>
Date: Fri, 25 Aug 2017 09:01:26 +0200
Subject: [PATCH] THRIFT-4292: Implement TimerManager::remove()

---
 .../src/thrift/concurrency/TimerManager.cpp   | 16 ++++-
 lib/cpp/test/concurrency/Tests.cpp            | 14 ++++
 lib/cpp/test/concurrency/TimerManagerTests.h  | 66 +++++++++++++++++++
 3 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/lib/cpp/src/thrift/concurrency/TimerManager.cpp b/lib/cpp/src/thrift/concurrency/TimerManager.cpp
index b03ff42af67..9ae1f941927 100644
--- a/lib/cpp/src/thrift/concurrency/TimerManager.cpp
+++ b/lib/cpp/src/thrift/concurrency/TimerManager.cpp
@@ -52,6 +52,8 @@ class TimerManager::Task : public Runnable {
     }
   }
 
+  bool operator==(const shared_ptr<Runnable> & runnable) const { return runnable_ == runnable; }
+
 private:
   shared_ptr<Runnable> runnable_;
   friend class TimerManager::Dispatcher;
@@ -290,11 +292,23 @@ void TimerManager::add(shared_ptr<Runnable> task, const struct timeval& value) {
 }
 
 void TimerManager::remove(shared_ptr<Runnable> task) {
-  (void)task;
   Synchronized s(monitor_);
   if (state_ != TimerManager::STARTED) {
     throw IllegalStateException();
   }
+  bool found = false;
+  for (task_iterator ix = taskMap_.begin(); ix != taskMap_.end();) {
+    if (*ix->second == task) {
+      found = true;
+      taskCount_--;
+      taskMap_.erase(ix++);
+    } else {
+      ++ix;
+    }
+  }
+  if (!found) {
+    throw NoSuchTaskException();
+  }
 }
 
 TimerManager::STATE TimerManager::state() const {
diff --git a/lib/cpp/test/concurrency/Tests.cpp b/lib/cpp/test/concurrency/Tests.cpp
index f49bb9feaa8..d09d438d659 100644
--- a/lib/cpp/test/concurrency/Tests.cpp
+++ b/lib/cpp/test/concurrency/Tests.cpp
@@ -123,6 +123,20 @@ int main(int argc, char** argv) {
       std::cerr << "\t\tTimerManager tests FAILED" << std::endl;
       return 1;
     }
+
+    std::cout << "\t\tTimerManager test01" << std::endl;
+
+    if (!timerManagerTests.test01()) {
+      std::cerr << "\t\tTimerManager tests FAILED" << std::endl;
+      return 1;
+    }
+
+    std::cout << "\t\tTimerManager test02" << std::endl;
+
+    if (!timerManagerTests.test02()) {
+      std::cerr << "\t\tTimerManager tests FAILED" << std::endl;
+      return 1;
+    }
   }
 
   if (runAll || args[0].compare("thread-manager") == 0) {
diff --git a/lib/cpp/test/concurrency/TimerManagerTests.h b/lib/cpp/test/concurrency/TimerManagerTests.h
index 32d39355f88..80d373bef78 100644
--- a/lib/cpp/test/concurrency/TimerManagerTests.h
+++ b/lib/cpp/test/concurrency/TimerManagerTests.h
@@ -126,6 +126,72 @@ class TimerManagerTests {
     return true;
   }
 
+  /**
+   * This test creates two tasks, removes the first one then waits for the second one. It then
+   * verifies that the timer manager properly clean up itself and the remaining orphaned timeout
+   * task when the manager goes out of scope and its destructor is called.
+   */
+  bool test01(int64_t timeout = 1000LL) {
+    TimerManager timerManager;
+    timerManager.threadFactory(shared_ptr<PlatformThreadFactory>(new PlatformThreadFactory()));
+    timerManager.start();
+    assert(timerManager.state() == TimerManager::STARTED);
+
+    Synchronized s(_monitor);
+
+    // Setup the two tasks
+    shared_ptr<TimerManagerTests::Task> taskToRemove
+      = shared_ptr<TimerManagerTests::Task>(new TimerManagerTests::Task(_monitor, timeout / 2));
+    timerManager.add(taskToRemove, taskToRemove->_timeout);
+
+    shared_ptr<TimerManagerTests::Task> task
+      = shared_ptr<TimerManagerTests::Task>(new TimerManagerTests::Task(_monitor, timeout));
+    timerManager.add(task, task->_timeout);
+
+    // Remove one task and wait until the other has completed
+    timerManager.remove(taskToRemove);
+    _monitor.wait(timeout * 2);
+
+    assert(!taskToRemove->_done);
+    assert(task->_done);
+
+    return true;
+  }
+
+  /**
+   * This test creates two tasks with the same callback and another one, then removes the two
+   * duplicated then waits for the last one. It then verifies that the timer manager properly
+   * clean up itself and the remaining orphaned timeout task when the manager goes out of scope
+   * and its destructor is called.
+   */
+  bool test02(int64_t timeout = 1000LL) {
+    TimerManager timerManager;
+    timerManager.threadFactory(shared_ptr<PlatformThreadFactory>(new PlatformThreadFactory()));
+    timerManager.start();
+    assert(timerManager.state() == TimerManager::STARTED);
+
+    Synchronized s(_monitor);
+
+    // Setup the one tasks and add it twice
+    shared_ptr<TimerManagerTests::Task> taskToRemove
+      = shared_ptr<TimerManagerTests::Task>(new TimerManagerTests::Task(_monitor, timeout / 3));
+    timerManager.add(taskToRemove, taskToRemove->_timeout);
+    timerManager.add(taskToRemove, taskToRemove->_timeout * 2);
+
+    shared_ptr<TimerManagerTests::Task> task
+      = shared_ptr<TimerManagerTests::Task>(new TimerManagerTests::Task(_monitor, timeout));
+    timerManager.add(task, task->_timeout);
+
+    // Remove the first task (e.g. two timers) and wait until the other has completed
+    timerManager.remove(taskToRemove);
+    _monitor.wait(timeout * 2);
+
+    assert(!taskToRemove->_done);
+    assert(task->_done);
+
+    return true;
+  }
+
   friend class TestTask;
 
   Monitor _monitor;