1 files changed, 25 insertions, 4 deletions

diff --git a/lib/cpp/src/thrift/concurrency/TimerManager.h b/lib/cpp/src/thrift/concurrency/TimerManager.h
index f6643483d..2bfc6a755 100644
--- a/lib/cpp/src/thrift/concurrency/TimerManager.h
+++ b/lib/cpp/src/thrift/concurrency/TimerManager.h
@@ -42,6 +42,9 @@ namespace concurrency {
 class TimerManager {
 
 public:
+  class Task;
+  typedef stdcxx::weak_ptr<Task> Timer;
+
   TimerManager();
 
   virtual ~TimerManager();
@@ -69,28 +72,33 @@ public:
    *
    * @param task The task to execute
    * @param timeout Time in milliseconds to delay before executing task
+   * @return Handle of the timer, which can be used to remove the timer.
    */
-  virtual void add(stdcxx::shared_ptr<Runnable> task, int64_t timeout);
+  virtual Timer add(stdcxx::shared_ptr<Runnable> task, int64_t timeout);
 
   /**
    * Adds a task to be executed at some time in the future by a worker thread.
    *
    * @param task The task to execute
    * @param timeout Absolute time in the future to execute task.
+   * @return Handle of the timer, which can be used to remove the timer.
    */
-  virtual void add(stdcxx::shared_ptr<Runnable> task, const struct THRIFT_TIMESPEC& timeout);
+  virtual Timer add(stdcxx::shared_ptr<Runnable> task, const struct THRIFT_TIMESPEC& timeout);
 
   /**
    * Adds a task to be executed at some time in the future by a worker thread.
    *
    * @param task The task to execute
    * @param timeout Absolute time in the future to execute task.
+   * @return Handle of the timer, which can be used to remove the timer.
    */
-  virtual void add(stdcxx::shared_ptr<Runnable> task, const struct timeval& timeout);
+  virtual Timer add(stdcxx::shared_ptr<Runnable> task, const struct timeval& timeout);
 
   /**
    * Removes a pending task
    *
+   * @param task The task to remove. All timers which execute this task will
+   * be removed.
    * @throws NoSuchTaskException Specified task doesn't exist. It was either
    *                             processed already or this call was made for a
    *                             task that was never added to this timer
@@ -100,13 +108,26 @@ public:
    */
   virtual void remove(stdcxx::shared_ptr<Runnable> task);
 
+  /**
+   * Removes a single pending task
+   *
+   * @param timer The timer to remove. The timer is returned when calling the
+   * add() method.
+   * @throws NoSuchTaskException Specified task doesn't exist. It was either
+   *                             processed already or this call was made for a
+   *                             task that was never added to this timer
+   *
+   * @throws UncancellableTaskException Specified task is already being
+   *                                    executed or has completed execution.
+   */
+  virtual void remove(Timer timer);
+
   enum STATE { UNINITIALIZED, STARTING, STARTED, STOPPING, STOPPED };
 
   virtual STATE state() const;
 
 private:
   stdcxx::shared_ptr<const ThreadFactory> threadFactory_;
-  class Task;
   friend class Task;
   std::multimap<int64_t, stdcxx::shared_ptr<Task> > taskMap_;
   size_t taskCount_;