SERVER-37329 Make ShardRemote interruptable while waiting for a response

3 files changed, 8 insertions, 4 deletions

diff --git a/src/mongo/executor/task_executor.h b/src/mongo/executor/task_executor.h
index d23069ba3ac..404065f3d27 100644
--- a/src/mongo/executor/task_executor.h
+++ b/src/mongo/executor/task_executor.h
@@ -257,8 +257,11 @@ public:
      * callback ran.
      *
      * NOTE: Do not call from a callback running in the executor.
+     *
+     * Prefer the version that takes an OperationContext* to this version.
      */
-    virtual void wait(const CallbackHandle& cbHandle) = 0;
+    virtual void wait(const CallbackHandle& cbHandle,
+                      Interruptible* interruptible = Interruptible::notInterruptible()) = 0;
 
     /**
      * Appends information about the underlying network interface's connections to the given
diff --git a/src/mongo/executor/thread_pool_task_executor.cpp b/src/mongo/executor/thread_pool_task_executor.cpp
index 37181bca6a9..c6015df0150 100644
--- a/src/mongo/executor/thread_pool_task_executor.cpp
+++ b/src/mongo/executor/thread_pool_task_executor.cpp
@@ -473,7 +473,7 @@ void ThreadPoolTaskExecutor::cancel(const CallbackHandle& cbHandle) {
     }
 }
 
-void ThreadPoolTaskExecutor::wait(const CallbackHandle& cbHandle) {
+void ThreadPoolTaskExecutor::wait(const CallbackHandle& cbHandle, Interruptible* interruptible) {
     invariant(cbHandle.isValid());
     auto cbState = checked_cast<CallbackState*>(getCallbackFromHandle(cbHandle));
     if (cbState->isFinished.load()) {
@@ -484,7 +484,7 @@ void ThreadPoolTaskExecutor::wait(const CallbackHandle& cbHandle) {
         cbState->finishedCondition.emplace();
     }
     while (!cbState->isFinished.load()) {
-        cbState->finishedCondition->wait(lk);
+        interruptible->waitForConditionOrInterrupt(*cbState->finishedCondition, lk);
     }
 }
 
diff --git a/src/mongo/executor/thread_pool_task_executor.h b/src/mongo/executor/thread_pool_task_executor.h
index 92b809bc0db..7beea67f044 100644
--- a/src/mongo/executor/thread_pool_task_executor.h
+++ b/src/mongo/executor/thread_pool_task_executor.h
@@ -85,7 +85,8 @@ public:
         const RemoteCommandCallbackFn& cb,
         const transport::BatonHandle& baton = nullptr) override;
     void cancel(const CallbackHandle& cbHandle) override;
-    void wait(const CallbackHandle& cbHandle) override;
+    void wait(const CallbackHandle& cbHandle,
+              Interruptible* interruptible = Interruptible::notInterruptible()) override;
 
     void appendConnectionStats(ConnectionPoolStats* stats) const override;