SERVER-37367 Make Notification::waitFor interruptible

2 files changed, 6 insertions, 8 deletions

diff --git a/src/mongo/db/s/operation_sharding_state.cpp b/src/mongo/db/s/operation_sharding_state.cpp
index 0f92bbd5492..2f0911bc9ac 100644
--- a/src/mongo/db/s/operation_sharding_state.cpp
+++ b/src/mongo/db/s/operation_sharding_state.cpp
@@ -40,7 +40,7 @@ const OperationContext::Decoration<OperationShardingState> shardingMetadataDecor
     OperationContext::declareDecoration<OperationShardingState>();
 
 // Max time to wait for the migration critical section to complete
-const Microseconds kMaxWaitForMigrationCriticalSection = Minutes(5);
+const Milliseconds kMaxWaitForMigrationCriticalSection = Minutes(5);
 
 }  // namespace
 
@@ -109,7 +109,7 @@ bool OperationShardingState::waitForMigrationCriticalSectionSignal(OperationCont
         _migrationCriticalSectionSignal->waitFor(
             txn,
             txn->hasDeadline()
-                ? std::min(txn->getRemainingMaxTimeMicros(), kMaxWaitForMigrationCriticalSection)
+                ? std::min(txn->getRemainingMaxTimeMillis(), kMaxWaitForMigrationCriticalSection)
                 : kMaxWaitForMigrationCriticalSection);
         _migrationCriticalSectionSignal = nullptr;
         return true;
diff --git a/src/mongo/util/concurrency/notification.h b/src/mongo/util/concurrency/notification.h
index d24fc84e5f9..25f0c65f187 100644
--- a/src/mongo/util/concurrency/notification.h
+++ b/src/mongo/util/concurrency/notification.h
@@ -102,12 +102,10 @@ public:
      * set (in which case a subsequent call to get is guaranteed to not block) or false otherwise.
      * If the wait is interrupted, throws an exception.
      */
-    bool waitFor(OperationContext* txn, Microseconds waitTimeout) {
-        const auto waitDeadline = Date_t::now() + waitTimeout;
-
+    bool waitFor(OperationContext* txn, Milliseconds waitTimeout) {
         stdx::unique_lock<stdx::mutex> lock(_mutex);
-        return _condVar.wait_until(
-            lock, waitDeadline.toSystemTimePoint(), [&]() { return !!_value; });
+        return txn->waitForConditionOrInterruptFor(
+            _condVar, lock, waitTimeout, [&]() { return !!_value; });
     }
 
 private:
@@ -137,7 +135,7 @@ public:
         _notification.set(true);
     }
 
-    bool waitFor(OperationContext* txn, Microseconds waitTimeout) {
+    bool waitFor(OperationContext* txn, Milliseconds waitTimeout) {
         return _notification.waitFor(txn, waitTimeout);
     }