SERVER-31484 separate the operation deadline from awaitData deadline in sharded queries.

The deadline has been been already separated for non-sharded queries.

4 files changed, 51 insertions, 18 deletions

diff --git a/src/mongo/executor/task_executor.h b/src/mongo/executor/task_executor.h
index 43a3ee81d9e..a36c9544ac3 100644
--- a/src/mongo/executor/task_executor.h
+++ b/src/mongo/executor/task_executor.h
@@ -39,6 +39,7 @@
 #include "mongo/executor/remote_command_request.h"
 #include "mongo/executor/remote_command_response.h"
 #include "mongo/platform/hash_namespace.h"
+#include "mongo/stdx/condition_variable.h"
 #include "mongo/stdx/functional.h"
 #include "mongo/util/time_support.h"
 
@@ -185,11 +186,12 @@ public:
     virtual void waitForEvent(const EventHandle& event) = 0;
 
     /**
-     * Same as waitForEvent without an OperationContext, but returns an error if the event was not
-     * triggered but the operation was killed - see OperationContext::checkForInterruptNoAssert()
-     * for expected error codes.
+     * Same as waitForEvent without an OperationContext, but returns Status::OK with
+     * cv_status::timeout if the event was not triggered within deadline.
      */
-    virtual Status waitForEvent(OperationContext* opCtx, const EventHandle& event) = 0;
+    virtual StatusWith<stdx::cv_status> waitForEvent(OperationContext* opCtx,
+                                                     const EventHandle& event,
+                                                     Date_t deadline) = 0;
 
     /**
      * Schedules "work" to be run by the executor ASAP.
diff --git a/src/mongo/executor/task_executor_test_common.cpp b/src/mongo/executor/task_executor_test_common.cpp
index 2185c221cf0..acc95e57754 100644
--- a/src/mongo/executor/task_executor_test_common.cpp
+++ b/src/mongo/executor/task_executor_test_common.cpp
@@ -350,10 +350,35 @@ COMMON_EXECUTOR_TEST(EventWaitingWithTimeoutTest) {
     auto client = serviceContext->makeClient("for testing");
     auto opCtx = client->makeOperationContext();
 
-    opCtx->setDeadlineAfterNowBy(Milliseconds{1});
+    auto deadline = mockClock->now() + Milliseconds{1};
     mockClock->advance(Milliseconds(2));
-    ASSERT_EQ(ErrorCodes::ExceededTimeLimit,
-              executor.waitForEvent(opCtx.get(), eventThatWillNeverBeTriggered));
+    ASSERT(stdx::cv_status::timeout ==
+           executor.waitForEvent(opCtx.get(), eventThatWillNeverBeTriggered, deadline));
+    executor.shutdown();
+    joinExecutorThread();
+}
+
+COMMON_EXECUTOR_TEST(EventSignalWithTimeoutTest) {
+    TaskExecutor& executor = getExecutor();
+    launchExecutorThread();
+
+    auto eventSignalled = unittest::assertGet(executor.makeEvent());
+
+    auto serviceContext = getGlobalServiceContext();
+
+    serviceContext->setFastClockSource(stdx::make_unique<ClockSourceMock>());
+    auto mockClock = static_cast<ClockSourceMock*>(serviceContext->getFastClockSource());
+
+    auto client = serviceContext->makeClient("for testing");
+    auto opCtx = client->makeOperationContext();
+
+    auto deadline = mockClock->now() + Milliseconds{1};
+    mockClock->advance(Milliseconds(1));
+
+    executor.signalEvent(eventSignalled);
+
+    ASSERT(stdx::cv_status::no_timeout ==
+           executor.waitForEvent(opCtx.get(), eventSignalled, deadline));
     executor.shutdown();
     joinExecutorThread();
 }
diff --git a/src/mongo/executor/thread_pool_task_executor.cpp b/src/mongo/executor/thread_pool_task_executor.cpp
index 4e6dc5f45fb..202de888b5a 100644
--- a/src/mongo/executor/thread_pool_task_executor.cpp
+++ b/src/mongo/executor/thread_pool_task_executor.cpp
@@ -285,22 +285,26 @@ StatusWith<TaskExecutor::CallbackHandle> ThreadPoolTaskExecutor::onEvent(const E
     return cbHandle;
 }
 
-Status ThreadPoolTaskExecutor::waitForEvent(OperationContext* opCtx, const EventHandle& event) {
+StatusWith<stdx::cv_status> ThreadPoolTaskExecutor::waitForEvent(OperationContext* opCtx,
+                                                                 const EventHandle& event,
+                                                                 Date_t deadline) {
     invariant(opCtx);
     invariant(event.isValid());
     auto eventState = checked_cast<EventState*>(getEventFromHandle(event));
     stdx::unique_lock<stdx::mutex> lk(_mutex);
 
-    try {
-        // std::condition_variable::wait() can wake up spuriously, so provide a callback to detect
-        // when that happens and go back to waiting.
-        opCtx->waitForConditionOrInterrupt(eventState->isSignaledCondition, lk, [&eventState]() {
-            return eventState->isSignaledFlag;
-        });
-    } catch (const DBException& e) {
-        return e.toStatus();
+    // std::condition_variable::wait() can wake up spuriously, so we have to loop until the event
+    // is signalled or we time out.
+    while (!eventState->isSignaledFlag) {
+        auto status = opCtx->waitForConditionOrInterruptNoAssertUntil(
+            eventState->isSignaledCondition, lk, deadline);
+
+        if (!status.isOK() || stdx::cv_status::timeout == status) {
+            return status;
+        }
     }
-    return Status::OK();
+
+    return stdx::cv_status::no_timeout;
 }
 
 void ThreadPoolTaskExecutor::waitForEvent(const EventHandle& event) {
diff --git a/src/mongo/executor/thread_pool_task_executor.h b/src/mongo/executor/thread_pool_task_executor.h
index b443061945e..8e81e3a7f07 100644
--- a/src/mongo/executor/thread_pool_task_executor.h
+++ b/src/mongo/executor/thread_pool_task_executor.h
@@ -73,7 +73,9 @@ public:
     StatusWith<EventHandle> makeEvent() override;
     void signalEvent(const EventHandle& event) override;
     StatusWith<CallbackHandle> onEvent(const EventHandle& event, const CallbackFn& work) override;
-    Status waitForEvent(OperationContext* opCtx, const EventHandle& event) override;
+    StatusWith<stdx::cv_status> waitForEvent(OperationContext* opCtx,
+                                             const EventHandle& event,
+                                             Date_t deadline) override;
     void waitForEvent(const EventHandle& event) override;
     StatusWith<CallbackHandle> scheduleWork(const CallbackFn& work) override;
     StatusWith<CallbackHandle> scheduleWorkAt(Date_t when, const CallbackFn& work) override;