SERVER-21004 Add setAlarm capability and virtualization flag to ClockSourceMock.

1 files changed, 34 insertions, 1 deletions

diff --git a/src/mongo/util/clock_source_mock.cpp b/src/mongo/util/clock_source_mock.cpp
index f2257bf2604..66879ca5438 100644
--- a/src/mongo/util/clock_source_mock.cpp
+++ b/src/mongo/util/clock_source_mock.cpp
@@ -29,7 +29,8 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/util/clock_source_mock.h"
-#include "mongo/util/time_support.h"
+
+#include <algorithm>
 
 namespace mongo {
 
@@ -38,14 +39,46 @@ Milliseconds ClockSourceMock::getPrecision() {
 }
 
 Date_t ClockSourceMock::now() {
+    stdx::lock_guard<stdx::mutex> lk(_mutex);
     return _now;
 }
 
 void ClockSourceMock::advance(Milliseconds ms) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
     _now += ms;
+    _processAlarms(std::move(lk));
 }
 
 void ClockSourceMock::reset(Date_t newNow) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
     _now = newNow;
+    _processAlarms(std::move(lk));
+}
+
+Status ClockSourceMock::setAlarm(Date_t when, stdx::function<void()> action) {
+    stdx::unique_lock<stdx::mutex> lk(_mutex);
+    if (when <= _now) {
+        lk.unlock();
+        action();
+        return Status::OK();
+    }
+    _alarms.emplace_back(std::make_pair(when, std::move(action)));
+    return Status::OK();
 }
+
+void ClockSourceMock::_processAlarms(stdx::unique_lock<stdx::mutex> lk) {
+    using std::swap;
+    invariant(lk.owns_lock());
+    std::vector<Alarm> readyAlarms;
+    std::vector<Alarm>::iterator iter;
+    auto alarmIsNotExpired = [&](const Alarm& alarm) { return alarm.first > _now; };
+    auto expiredAlarmsBegin = std::partition(_alarms.begin(), _alarms.end(), alarmIsNotExpired);
+    std::move(expiredAlarmsBegin, _alarms.end(), std::back_inserter(readyAlarms));
+    _alarms.erase(expiredAlarmsBegin, _alarms.end());
+    lk.unlock();
+    for (const auto& alarm : readyAlarms) {
+        alarm.second();
+    }
+}
+
 }  // namespace mongo