SERVER-44951 Expand TaskExecutor to implement exhaust functionality

2 files changed, 17 insertions, 0 deletions

diff --git a/src/mongo/unittest/task_executor_proxy.cpp b/src/mongo/unittest/task_executor_proxy.cpp
index dd583c17056..911123d375c 100644
--- a/src/mongo/unittest/task_executor_proxy.cpp
+++ b/src/mongo/unittest/task_executor_proxy.cpp
@@ -106,6 +106,18 @@ StatusWith<executor::TaskExecutor::CallbackHandle> TaskExecutorProxy::scheduleRe
     return _executor->scheduleRemoteCommandOnAny(request, cb, baton);
 }
 
+StatusWith<executor::TaskExecutor::CallbackHandle>
+TaskExecutorProxy::scheduleExhaustRemoteCommandOnAny(
+    const executor::RemoteCommandRequestOnAny& request,
+    const RemoteCommandOnAnyCallbackFn& cb,
+    const BatonHandle& baton) {
+    return _executor->scheduleExhaustRemoteCommandOnAny(request, cb, baton);
+}
+
+bool TaskExecutorProxy::hasTasks() {
+    return _executor->hasTasks();
+}
+
 void TaskExecutorProxy::cancel(const CallbackHandle& cbHandle) {
     _executor->cancel(cbHandle);
 }
diff --git a/src/mongo/unittest/task_executor_proxy.h b/src/mongo/unittest/task_executor_proxy.h
index 3d37f837e81..fa27e0cde05 100644
--- a/src/mongo/unittest/task_executor_proxy.h
+++ b/src/mongo/unittest/task_executor_proxy.h
@@ -69,6 +69,11 @@ public:
         const executor::RemoteCommandRequestOnAny& request,
         const RemoteCommandOnAnyCallbackFn& cb,
         const BatonHandle& baton = nullptr) override;
+    StatusWith<CallbackHandle> scheduleExhaustRemoteCommandOnAny(
+        const executor::RemoteCommandRequestOnAny& request,
+        const RemoteCommandOnAnyCallbackFn& cb,
+        const BatonHandle& baton = nullptr) override;
+    bool hasTasks() override;
     void cancel(const CallbackHandle& cbHandle) override;
     void wait(const CallbackHandle& cbHandle,
               Interruptible* interruptible = Interruptible::notInterruptible()) override;