SERVER-32783 make RemoteCommandRetryScheduler single-use - cannot be be restarted once completed

(cherry picked from commit 35b9b4287581fdc9f37d3afeebfb2c9895b2428b)

3 files changed, 61 insertions, 13 deletions

diff --git a/src/mongo/client/remote_command_retry_scheduler.cpp b/src/mongo/client/remote_command_retry_scheduler.cpp
index 8f9f73058c8..ee067daa07d 100644
--- a/src/mongo/client/remote_command_retry_scheduler.cpp
+++ b/src/mongo/client/remote_command_retry_scheduler.cpp
@@ -167,20 +167,34 @@ RemoteCommandRetryScheduler::~RemoteCommandRetryScheduler() {
 
 bool RemoteCommandRetryScheduler::isActive() const {
     stdx::lock_guard<stdx::mutex> lock(_mutex);
-    return _active;
+    return _isActive_inlock();
+}
+
+bool RemoteCommandRetryScheduler::_isActive_inlock() const {
+    return State::kRunning == _state || State::kShuttingDown == _state;
 }
 
 Status RemoteCommandRetryScheduler::startup() {
     stdx::lock_guard<stdx::mutex> lock(_mutex);
 
-    if (_active) {
-        return Status(ErrorCodes::IllegalOperation, "fetcher already scheduled");
+    switch (_state) {
+        case State::kPreStart:
+            _state = State::kRunning;
+            break;
+        case State::kRunning:
+            return Status(ErrorCodes::IllegalOperation, "scheduler already started");
+        case State::kShuttingDown:
+            return Status(ErrorCodes::ShutdownInProgress, "scheduler shutting down");
+        case State::kComplete:
+            return Status(ErrorCodes::ShutdownInProgress, "scheduler completed");
     }
 
     auto scheduleStatus = _schedule_inlock();
     if (!scheduleStatus.isOK()) {
+        _state = State::kComplete;
         return scheduleStatus;
     }
+
     return Status::OK();
 }
 
@@ -188,9 +202,18 @@ void RemoteCommandRetryScheduler::shutdown() {
     executor::TaskExecutor::CallbackHandle remoteCommandCallbackHandle;
     {
         stdx::lock_guard<stdx::mutex> lock(_mutex);
-
-        if (!_active) {
-            return;
+        switch (_state) {
+            case State::kPreStart:
+                // Transition directly from PreStart to Complete if not started yet.
+                _state = State::kComplete;
+                return;
+            case State::kRunning:
+                _state = State::kShuttingDown;
+                break;
+            case State::kShuttingDown:
+            case State::kComplete:
+                // Nothing to do if we are already in ShuttingDown or Complete state.
+                return;
         }
 
         remoteCommandCallbackHandle = _remoteCommandCallbackHandle;
@@ -202,7 +225,7 @@ void RemoteCommandRetryScheduler::shutdown() {
 
 void RemoteCommandRetryScheduler::join() {
     stdx::unique_lock<stdx::mutex> lock(_mutex);
-    _condition.wait(lock, [this]() { return !_active; });
+    _condition.wait(lock, [this]() { return !_isActive_inlock(); });
 }
 
 std::string RemoteCommandRetryScheduler::toString() const {
@@ -210,7 +233,7 @@ std::string RemoteCommandRetryScheduler::toString() const {
     str::stream output;
     output << "RemoteCommandRetryScheduler";
     output << " request: " << _request.toString();
-    output << " active: " << _active;
+    output << " active: " << _isActive_inlock();
     if (_remoteCommandCallbackHandle.isValid()) {
         output << " callbackHandle.valid: " << _remoteCommandCallbackHandle.isValid();
         output << " callbackHandle.cancelled: " << _remoteCommandCallbackHandle.isCanceled();
@@ -232,7 +255,6 @@ Status RemoteCommandRetryScheduler::_schedule_inlock() {
     }
 
     _remoteCommandCallbackHandle = scheduleResult.getValue();
-    _active = true;
     return Status::OK();
 }
 
@@ -274,8 +296,8 @@ void RemoteCommandRetryScheduler::_onComplete(
     _callback = {};
 
     stdx::lock_guard<stdx::mutex> lock(_mutex);
-    invariant(_active);
-    _active = false;
+    invariant(_isActive_inlock());
+    _state = State::kComplete;
     _condition.notify_all();
 }
 
diff --git a/src/mongo/client/remote_command_retry_scheduler.h b/src/mongo/client/remote_command_retry_scheduler.h
index 93c8ca17ad0..bd963d3a2c1 100644
--- a/src/mongo/client/remote_command_retry_scheduler.h
+++ b/src/mongo/client/remote_command_retry_scheduler.h
@@ -104,6 +104,7 @@ public:
      * Returns true if we have scheduled a remote command and are waiting for the response.
      */
     bool isActive() const;
+    bool _isActive_inlock() const;
 
     /**
      * Schedules remote command request.
@@ -159,8 +160,13 @@ private:
 
     mutable stdx::condition_variable _condition;
 
-    // _active is true when remote command is scheduled to be run by the executor.
-    bool _active = false;
+    // State transitions:
+    // PreStart --> Running --> ShuttingDown --> Complete
+    // It is possible to skip intermediate states. For example,
+    // Calling shutdown() when the scheduler has not started will transition from PreStart directly
+    // to Complete.
+    enum class State { kPreStart, kRunning, kShuttingDown, kComplete };
+    State _state = State::kPreStart;  // (M)
 
     // Callback handle to the scheduled remote command.
     executor::TaskExecutor::CallbackHandle _remoteCommandCallbackHandle;
diff --git a/src/mongo/client/remote_command_retry_scheduler_test.cpp b/src/mongo/client/remote_command_retry_scheduler_test.cpp
index 23f6e6ba3be..c6331327960 100644
--- a/src/mongo/client/remote_command_retry_scheduler_test.cpp
+++ b/src/mongo/client/remote_command_retry_scheduler_test.cpp
@@ -299,6 +299,19 @@ TEST_F(RemoteCommandRetrySchedulerTest, StartupFailsWhenExecutorIsShutDown) {
     ASSERT_FALSE(scheduler.isActive());
 }
 
+TEST_F(RemoteCommandRetrySchedulerTest, StartupFailsWhenSchedulerIsShutDown) {
+    auto callback = [](const executor::TaskExecutor::RemoteCommandCallbackArgs&) {};
+    auto policy = RemoteCommandRetryScheduler::makeNoRetryPolicy();
+
+    RemoteCommandRetryScheduler scheduler(&getExecutor(), request, callback, std::move(policy));
+    ASSERT_FALSE(scheduler.isActive());
+
+    scheduler.shutdown();
+
+    ASSERT_EQUALS(ErrorCodes::ShutdownInProgress, scheduler.startup());
+    ASSERT_FALSE(scheduler.isActive());
+}
+
 TEST_F(RemoteCommandRetrySchedulerTest,
        ShuttingDownExecutorAfterSchedulerStartupInvokesCallbackWithCallbackCanceledError) {
     CallbackResponseSaver callback;
@@ -349,6 +362,9 @@ TEST_F(RemoteCommandRetrySchedulerTest, SchedulerInvokesCallbackOnNonRetryableEr
     ResponseStatus rs(ErrorCodes::OperationFailed, "injected error", Milliseconds(0));
     processNetworkResponse(rs);
     checkCompletionStatus(&scheduler, callback, rs);
+
+    // Scheduler cannot be restarted once it has run to completion.
+    ASSERT_EQUALS(ErrorCodes::ShutdownInProgress, scheduler.startup());
 }
 
 TEST_F(RemoteCommandRetrySchedulerTest, SchedulerInvokesCallbackOnFirstSuccessfulResponse) {
@@ -364,6 +380,10 @@ TEST_F(RemoteCommandRetrySchedulerTest, SchedulerInvokesCallbackOnFirstSuccessfu
 
     processNetworkResponse(response);
     checkCompletionStatus(&scheduler, callback, response);
+
+    // Scheduler cannot be restarted once it has run to completion.
+    ASSERT_EQUALS(ErrorCodes::ShutdownInProgress, scheduler.startup());
+    ASSERT_FALSE(scheduler.isActive());
 }
 
 TEST_F(RemoteCommandRetrySchedulerTest, SchedulerIgnoresEmbeddedErrorInSuccessfulResponse) {