SERVER-26544: Keep DatabaseCloner active until after the completion callback

2 files changed, 38 insertions, 31 deletions

diff --git a/src/mongo/db/repl/databases_cloner.cpp b/src/mongo/db/repl/databases_cloner.cpp
index 0527a82dfd0..4a4e013de0e 100644
--- a/src/mongo/db/repl/databases_cloner.cpp
+++ b/src/mongo/db/repl/databases_cloner.cpp
@@ -194,8 +194,7 @@ Status DatabasesCloner::startup() {
             RemoteCommandRetryScheduler::kAllRetriableErrors));
     auto s = _listDBsScheduler->startup();
     if (!s.isOK()) {
-        _setStatus_inlock(s);
-        _failed_inlock(lk);
+        _fail_inlock(&lk, s);
     }
 
     return _status;
@@ -215,8 +214,7 @@ void DatabasesCloner::_onListDatabaseFinish(const CommandCallbackArgs& cbd) {
     UniqueLock lk(_mutex);
     if (!respStatus.isOK()) {
         LOG(1) << "listDatabases failed: " << respStatus;
-        _setStatus_inlock(respStatus);
-        _failed_inlock(lk);
+        _fail_inlock(&lk, respStatus);
         return;
     }
 
@@ -300,11 +298,9 @@ void DatabasesCloner::_onListDatabaseFinish(const CommandCallbackArgs& cbd) {
 
     if (_databaseCloners.size() == 0) {
         if (_status.isOK()) {
-            _active = false;
-            lk.unlock();
-            _finishFn(_status);
+            _succeed_inlock(&lk);
         } else {
-            _failed_inlock(lk);
+            _fail_inlock(&lk, _status);
         }
     }
 }
@@ -324,8 +320,7 @@ void DatabasesCloner::_onEachDBCloneFinish(const Status& status, const std::stri
     if (!status.isOK()) {
         warning() << "database '" << name << "' (" << (_stats.databasesCloned + 1) << " of "
                   << _databaseCloners.size() << ") clone failed due to " << status.toString();
-        _setStatus_inlock(status);
-        _failed_inlock(lk);
+        _fail_inlock(&lk, status);
         return;
     }
 
@@ -345,8 +340,7 @@ void DatabasesCloner::_onEachDBCloneFinish(const Status& status, const std::stri
         }
         if (!adminStatus.isOK()) {
             LOG(1) << "Validation failed on 'admin' db due to " << adminStatus;
-            _setStatus_inlock(adminStatus);
-            _failed_inlock(lk);
+            _fail_inlock(&lk, adminStatus);
             return;
         }
     }
@@ -354,11 +348,7 @@ void DatabasesCloner::_onEachDBCloneFinish(const Status& status, const std::stri
     _stats.databasesCloned++;
 
     if (_stats.databasesCloned == _databaseCloners.size()) {
-        _active = false;
-        // All cloners are done, trigger event.
-        LOG(2) << "All database clones finished, calling _finishFn.";
-        lk.unlock();
-        _finishFn(_status);
+        _succeed_inlock(&lk);
         return;
     }
 
@@ -369,25 +359,41 @@ void DatabasesCloner::_onEachDBCloneFinish(const Status& status, const std::stri
         warning() << "failed to schedule database '" << name << "' ("
                   << (_stats.databasesCloned + 1) << " of " << _databaseCloners.size()
                   << ") due to " << startStatus.toString();
-        _setStatus_inlock(startStatus);
-        _failed_inlock(lk);
+        _fail_inlock(&lk, startStatus);
         return;
     }
 }
 
-void DatabasesCloner::_failed_inlock(UniqueLock& lk) {
-    LOG(3) << "DatabasesCloner::_failed_inlock";
+void DatabasesCloner::_fail_inlock(UniqueLock* lk, Status status) {
+    LOG(3) << "DatabasesCloner::_fail_inlock called";
     if (!_active) {
         return;
     }
-    _active = false;
 
+    _setStatus_inlock(status);
     // TODO: shutdown outstanding work, like any cloners active
     auto finish = _finishFn;
-    lk.unlock();
+    lk->unlock();
+
+    LOG(3) << "DatabasesCloner - calling _finishFn with status: " << _status;
+    finish(status);
+
+    lk->lock();
+    _active = false;
+}
 
-    LOG(3) << "calling _finishFn with status: " << _status;
-    _finishFn(_status);
+void DatabasesCloner::_succeed_inlock(UniqueLock* lk) {
+    LOG(3) << "DatabasesCloner::_succeed_inlock called";
+    const auto status = Status::OK();
+    _setStatus_inlock(status);
+    auto finish = _finishFn;
+    lk->unlock();
+
+    LOG(3) << "DatabasesCloner - calling _finishFn with status OK";
+    finish(status);
+
+    lk->lock();
+    _active = false;
 }
 
 void DatabasesCloner::_setStatus_inlock(Status s) {
diff --git a/src/mongo/db/repl/databases_cloner.h b/src/mongo/db/repl/databases_cloner.h
index 14d815f9e1a..99748af768f 100644
--- a/src/mongo/db/repl/databases_cloner.h
+++ b/src/mongo/db/repl/databases_cloner.h
@@ -120,10 +120,11 @@ private:
      */
     void _setStatus_inlock(Status s);
 
-    /**
-     * Will fail the cloner, unlock and call the completion function.
-     */
-    void _failed_inlock(UniqueLock& lk);
+    /** Will fail the cloner, call the completion function, and become inactive. */
+    void _fail_inlock(UniqueLock* lk, Status s);
+
+    /** Will call the completion function, and become inactive. */
+    void _succeed_inlock(UniqueLock* lk);
 
     /** Called each time a database clone is finished */
     void _onEachDBCloneFinish(const Status& status, const std::string& name);
@@ -144,8 +145,8 @@ private:
     Status _status{ErrorCodes::NotYetInitialized, ""};  // (M) If it is not OK, we stop everything.
     executor::TaskExecutor* _exec;                      // (R) executor to schedule things with
     OldThreadPool* _dbWorkThreadPool;  // (R) db worker thread pool for collection cloning.
-    HostAndPort _source;               // (R) The source to use, until we get an error
-    bool _active = false;              // (M) false until we start
+    const HostAndPort _source;         // (R) The source to use.
+    bool _active = false;              // (M) false until we start, and true until finished.
     CollectionCloner::ScheduleDbWorkFn _scheduleDbWorkFn;  // (M)
 
     const IncludeDbFilterFn _includeDbFn;  // (R) function which decides which dbs are cloned.