SERVER-29009 Add ability for AsyncTimers to be reset

5 files changed, 73 insertions, 35 deletions

diff --git a/src/mongo/executor/async_timer_asio.cpp b/src/mongo/executor/async_timer_asio.cpp
index 1e418d1ded0..85b1a02cf7a 100644
--- a/src/mongo/executor/async_timer_asio.cpp
+++ b/src/mongo/executor/async_timer_asio.cpp
@@ -32,6 +32,7 @@
 
 #include "mongo/executor/async_timer_asio.h"
 
+#include "mongo/stdx/chrono.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/util/log.h"
 
@@ -53,6 +54,10 @@ void AsyncTimerASIO::asyncWait(AsyncTimerInterface::Handler handler) {
     _timer.async_wait(_strand->wrap(std::move(handler)));
 }
 
+void AsyncTimerASIO::expireAfter(Milliseconds expiration) {
+    _timer.expires_after(std::chrono::milliseconds(expiration.count()));
+}
+
 std::unique_ptr<AsyncTimerInterface> AsyncTimerFactoryASIO::make(asio::io_service::strand* strand,
                                                                  Milliseconds expiration) {
     return stdx::make_unique<AsyncTimerASIO>(strand, expiration);
diff --git a/src/mongo/executor/async_timer_asio.h b/src/mongo/executor/async_timer_asio.h
index 3262ad63c12..01e1fbe18b8 100644
--- a/src/mongo/executor/async_timer_asio.h
+++ b/src/mongo/executor/async_timer_asio.h
@@ -44,6 +44,8 @@ public:
 
     void asyncWait(AsyncTimerInterface::Handler handler) override;
 
+    void expireAfter(Milliseconds expiration) override;
+
 private:
     asio::io_service::strand* const _strand;
     asio::system_timer _timer;
diff --git a/src/mongo/executor/async_timer_interface.h b/src/mongo/executor/async_timer_interface.h
index ffe7d79a0d5..fdfce5b39d8 100644
--- a/src/mongo/executor/async_timer_interface.h
+++ b/src/mongo/executor/async_timer_interface.h
@@ -68,6 +68,12 @@ public:
      */
     virtual void asyncWait(Handler handler) = 0;
 
+    /**
+     * Reset this timer's expiry time relative to now. Any pending asyncWait operations
+     * will be canceled, and their handlers will be invoked with an error code.
+     */
+    virtual void expireAfter(Milliseconds expiration) = 0;
+
 protected:
     AsyncTimerInterface() = default;
 };
diff --git a/src/mongo/executor/async_timer_mock.cpp b/src/mongo/executor/async_timer_mock.cpp
index 4148499ffaf..020991d8fe8 100644
--- a/src/mongo/executor/async_timer_mock.cpp
+++ b/src/mongo/executor/async_timer_mock.cpp
@@ -44,45 +44,71 @@ void AsyncTimerMockImpl::cancel() {
 }
 
 void AsyncTimerMockImpl::asyncWait(AsyncTimerInterface::Handler handler) {
-    // If we have expired, run handler now instead of storing.
-    if (_timeLeft == kZeroMilliseconds) {
-        // Callbacks scheduled for now will fire immediately, synchronously.
-        handler(std::error_code());
-    } else {
-        stdx::lock_guard<stdx::mutex> lk(_handlersMutex);
-        _handlers.push_back(handler);
+    {
+        stdx::lock_guard<stdx::mutex> lk(_mutex);
+        if (_timeLeft != kZeroMilliseconds) {
+            _handlers.push_back(handler);
+            return;
+        }
     }
+
+    // If we have expired, run handler now instead of storing.
+    // Callbacks scheduled for now will fire immediately, synchronously.
+    handler(std::error_code());
 }
 
-bool AsyncTimerMockImpl::fastForward(Milliseconds time) {
-    if (time >= _timeLeft) {
-        _timeLeft = kZeroMilliseconds;
-        _expire();
-    } else {
-        _timeLeft -= time;
+void AsyncTimerMockImpl::fastForward(Milliseconds time) {
+    std::vector<AsyncTimerInterface::Handler> tmp;
+
+    // While holding the lock, change the time and remove
+    // handlers that have expired
+    {
+        stdx::lock_guard<stdx::mutex> lk(_mutex);
+        if (time >= _timeLeft) {
+            _timeLeft = kZeroMilliseconds;
+            tmp.swap(_handlers);
+        } else {
+            _timeLeft -= time;
+        }
     }
 
-    return _timeLeft > kZeroMilliseconds;
+    // If handlers expired, call them outside of the lock
+    for (const auto& handler : tmp) {
+        handler(std::error_code());
+    }
 }
 
 Milliseconds AsyncTimerMockImpl::timeLeft() {
+    stdx::lock_guard<stdx::mutex> lk(_mutex);
     return _timeLeft;
 }
 
-void AsyncTimerMockImpl::_callAllHandlers(std::error_code ec) {
+void AsyncTimerMockImpl::expireAfter(Milliseconds expiration) {
     std::vector<AsyncTimerInterface::Handler> tmp;
+
+    // While holding the lock, reset the time and remove all handlers
     {
-        stdx::lock_guard<stdx::mutex> lk(_handlersMutex);
+        stdx::lock_guard<stdx::mutex> lk(_mutex);
+        _timeLeft = expiration;
         tmp.swap(_handlers);
     }
 
+    // Call handlers with a "canceled" error code
     for (const auto& handler : tmp) {
-        handler(ec);
+        handler(asio::error::operation_aborted);
     }
 }
 
-void AsyncTimerMockImpl::_expire() {
-    _callAllHandlers(std::error_code());
+void AsyncTimerMockImpl::_callAllHandlers(std::error_code ec) {
+    std::vector<AsyncTimerInterface::Handler> tmp;
+    {
+        stdx::lock_guard<stdx::mutex> lk(_mutex);
+        tmp.swap(_handlers);
+    }
+
+    for (const auto& handler : tmp) {
+        handler(ec);
+    }
 }
 
 AsyncTimerMock::AsyncTimerMock(std::shared_ptr<AsyncTimerMockImpl> timer) : _timer(timer) {}
@@ -95,6 +121,10 @@ void AsyncTimerMock::asyncWait(AsyncTimerInterface::Handler handler) {
     _timer->asyncWait(handler);
 }
 
+void AsyncTimerMock::expireAfter(Milliseconds expiration) {
+    _timer->expireAfter(expiration);
+}
+
 std::unique_ptr<AsyncTimerInterface> AsyncTimerFactoryMock::make(Milliseconds expiration) {
     return make(nullptr, expiration);
 }
@@ -111,19 +141,10 @@ void AsyncTimerFactoryMock::fastForward(Milliseconds time) {
 
     _curTime += time;
 
-    // erase after iterating to be safe
-    stdx::unordered_set<std::shared_ptr<AsyncTimerMockImpl>> expired;
+    // Timers may be reset, so keep them in our set even if they have expired.
     for (auto elem = _timers.begin(); elem != _timers.end(); elem++) {
         auto timer = *elem;
-
-        // If timer has expired, register it for removal from our set.
-        if (!timer->fastForward(time)) {
-            expired.insert(timer);
-        }
-    }
-
-    for (auto elem = expired.begin(); elem != expired.end(); elem++) {
-        _timers.erase(*elem);
+        timer->fastForward(time);
     }
 }
 
diff --git a/src/mongo/executor/async_timer_mock.h b/src/mongo/executor/async_timer_mock.h
index 073bdc935ac..862d9a93863 100644
--- a/src/mongo/executor/async_timer_mock.h
+++ b/src/mongo/executor/async_timer_mock.h
@@ -62,22 +62,24 @@ public:
     /**
      * Advance current time. If the given interval is greater than or equal to the
      * time left on the timer, expire and call callbacks now.
-     *
-     * Returns true if the timer is still active, false if it has now expired.
      */
-    bool fastForward(Milliseconds time);
+    void fastForward(Milliseconds time);
 
     /**
      * Return the amount of time left on this timer.
      */
     Milliseconds timeLeft();
 
+    /**
+     * Reset the timer.
+     */
+    void expireAfter(Milliseconds expiration);
+
 private:
     void _callAllHandlers(std::error_code ec);
-    void _expire();
 
+    stdx::mutex _mutex;
     Milliseconds _timeLeft;
-    stdx::mutex _handlersMutex;
     std::vector<AsyncTimerInterface::Handler> _handlers;
 };
 
@@ -97,6 +99,8 @@ public:
 
     void asyncWait(AsyncTimerInterface::Handler handler) override;
 
+    void expireAfter(Milliseconds expiration) override;
+
 private:
     // Unfortunate, but it makes the ownership model sane.
     std::shared_ptr<AsyncTimerMockImpl> _timer;