SERVER-32783 fix race in RemoteCommandRetryScheduler between shutdown() and resending command

(cherry picked from commit acf7bec77edde339ed6fb1bb89f7f03888144476)

2 files changed, 59 insertions, 1 deletions

diff --git a/src/mongo/client/remote_command_retry_scheduler.cpp b/src/mongo/client/remote_command_retry_scheduler.cpp
index ee067daa07d..e894e708f3d 100644
--- a/src/mongo/client/remote_command_retry_scheduler.cpp
+++ b/src/mongo/client/remote_command_retry_scheduler.cpp
@@ -277,7 +277,15 @@ void RemoteCommandRetryScheduler::_remoteCommandCallback(
 
     // TODO(benety): Check cumulative elapsed time of failed responses received against retry
     // policy. Requires SERVER-24067.
-    auto scheduleStatus = _schedule_inlock();
+    auto scheduleStatus = [this]() {
+        stdx::lock_guard<stdx::mutex> lock(_mutex);
+        if (State::kShuttingDown == _state) {
+            return Status(ErrorCodes::CallbackCanceled,
+                          "scheduler was shut down before retrying command");
+        }
+        return _schedule_inlock();
+    }();
+
     if (!scheduleStatus.isOK()) {
         _onComplete({rcba.executor, rcba.myHandle, rcba.request, scheduleStatus});
         return;
diff --git a/src/mongo/client/remote_command_retry_scheduler_test.cpp b/src/mongo/client/remote_command_retry_scheduler_test.cpp
index c6331327960..dca50732e93 100644
--- a/src/mongo/client/remote_command_retry_scheduler_test.cpp
+++ b/src/mongo/client/remote_command_retry_scheduler_test.cpp
@@ -461,6 +461,56 @@ TEST_F(RemoteCommandRetrySchedulerTest, SchedulerShouldRetryUntilSuccessfulRespo
     checkCompletionStatus(&scheduler, callback, response);
 }
 
+/**
+ * Retry policy that shuts down the scheduler whenever it is consulted by the scheduler.
+ * Results from getMaximumAttempts() and shouldRetryOnError() must cause the scheduler
+ * to resend the request.
+ */
+class ShutdownSchedulerRetryPolicy : public RemoteCommandRetryScheduler::RetryPolicy {
+public:
+    std::size_t getMaximumAttempts() const override {
+        if (scheduler) {
+            scheduler->shutdown();
+        }
+        return 2U;
+    }
+    Milliseconds getMaximumResponseElapsedTotal() const override {
+        return executor::RemoteCommandRequest::kNoTimeout;
+    }
+    bool shouldRetryOnError(ErrorCodes::Error) const override {
+        if (scheduler) {
+            scheduler->shutdown();
+        }
+        return true;
+    }
+    std::string toString() const override {
+        return "";
+    }
+
+    // This must be set before starting the scheduler.
+    RemoteCommandRetryScheduler* scheduler = nullptr;
+};
+
+TEST_F(RemoteCommandRetrySchedulerTest,
+       SchedulerReturnsCallbackCanceledIfShutdownBeforeSendingRetryCommand) {
+    CallbackResponseSaver callback;
+    auto policy = stdx::make_unique<ShutdownSchedulerRetryPolicy>();
+    auto policyPtr = policy.get();
+    TaskExecutorWithFailureInScheduleRemoteCommand badExecutor(&getExecutor());
+    RemoteCommandRetryScheduler scheduler(
+        &badExecutor, request, stdx::ref(callback), std::move(policy));
+    policyPtr->scheduler = &scheduler;
+    start(&scheduler);
+
+    processNetworkResponse({ErrorCodes::HostNotFound, "first", Milliseconds(0)});
+
+    checkCompletionStatus(&scheduler,
+                          callback,
+                          {ErrorCodes::CallbackCanceled,
+                           "scheduler was shut down before retrying command",
+                           Milliseconds(0)});
+}
+
 bool sharedCallbackStateDestroyed = false;
 class SharedCallbackState {
     MONGO_DISALLOW_COPYING(SharedCallbackState);