From 8b51694b56db79f30fe20cb6c019d5d28c030ec4 Mon Sep 17 00:00:00 2001
From: Jody Ho <jodyho@fb.com>
Date: Wed, 15 Nov 2017 23:23:02 -0800
Subject: [PATCH] Expose the time remaining in HHWheelTimer::Callback

Summary:
We would like to know the time remaining for a scheduled timeout to decide
whether a new event should override the scheduled timeout.

Reviewed By: yfeldblum

Differential Revision: D6334067

fbshipit-source-id: f172d5cd7fc804db5fd53a42d06cadfddf857e22
---
 folly/io/async/HHWheelTimer.h            |  9 +++++++++
 folly/io/async/test/HHWheelTimerTest.cpp | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/folly/io/async/HHWheelTimer.h b/folly/io/async/HHWheelTimer.h
index a8978a1b..5544c2d6 100644
--- a/folly/io/async/HHWheelTimer.h
+++ b/folly/io/async/HHWheelTimer.h
@@ -104,6 +104,15 @@ class HHWheelTimer : private folly::AsyncTimeout,
       return wheel_ != nullptr;
     }
 
+    /**
+     * Get the time remaining until this timeout expires. Return 0 if this
+     * timeout is not scheduled or expired. Otherwise, return expiration time
+     * minus getCurTime().
+     */
+    std::chrono::milliseconds getTimeRemaining() {
+      return getTimeRemaining(getCurTime());
+    }
+
    protected:
     /**
      * Don't override this unless you're doing a test. This is mainly here so
diff --git a/folly/io/async/test/HHWheelTimerTest.cpp b/folly/io/async/test/HHWheelTimerTest.cpp
index a05bb698..f7744608 100644
--- a/folly/io/async/test/HHWheelTimerTest.cpp
+++ b/folly/io/async/test/HHWheelTimerTest.cpp
@@ -430,3 +430,27 @@ TEST_F(HHWheelTimerTest, IntrusivePtr) {
   T_CHECK_TIMEOUT(start, t3.timestamps[0], milliseconds(10));
   T_CHECK_TIMEOUT(start, end, milliseconds(10));
 }
+
+TEST_F(HHWheelTimerTest, GetTimeRemaining) {
+  StackWheelTimer t(&eventBase, milliseconds(1));
+  TestTimeout t1;
+
+  // Not scheduled yet, time remaining should be zero
+  ASSERT_EQ(t1.getTimeRemaining(), milliseconds(0));
+  ASSERT_EQ(t.count(), 0);
+
+  // Scheduled, time remaining should be less than or equal to the scheduled
+  // timeout
+  t.scheduleTimeout(&t1, milliseconds(10));
+  ASSERT_LE(t1.getTimeRemaining(), milliseconds(10));
+
+  TimePoint start;
+  eventBase.loop();
+  TimePoint end;
+
+  // Expired and time remaining should be zero
+  ASSERT_EQ(t1.getTimeRemaining(), milliseconds(0));
+
+  ASSERT_EQ(t.count(), 0);
+  T_CHECK_TIMEOUT(start, end, milliseconds(10));
+}
-- 
2.34.1