SERVER-37496 Register balancer shutdown task when MongoD is single-threaded

(cherry picked from commit 871f1234c72f3f390f81dc31fd8a95d8bc69cf91)

2 files changed, 6 insertions, 9 deletions

diff --git a/src/mongo/db/db.cpp b/src/mongo/db/db.cpp
index 05fc368e3b6..e05ed922183 100644
--- a/src/mongo/db/db.cpp
+++ b/src/mongo/db/db.cpp
@@ -1062,6 +1062,12 @@ static void shutdownTask() {
     Client::initThreadIfNotAlready();
     Client& client = cc();
 
+    // Terminate the balancer thread so it doesn't leak memory.
+    if (auto balancer = Balancer::get(serviceContext)) {
+        balancer->interruptBalancer();
+        balancer->waitForBalancerToStop();
+    }
+
     ServiceContext::UniqueOperationContext uniqueTxn;
     OperationContext* txn = client.getOperationContext();
     if (!txn && serviceContext->getGlobalStorageEngine()) {
diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index d6291503cc4..fb635d5fa3c 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -168,15 +168,6 @@ Balancer::~Balancer() {
 void Balancer::create(ServiceContext* serviceContext) {
     invariant(!getBalancer(serviceContext));
     getBalancer(serviceContext) = stdx::make_unique<Balancer>(serviceContext);
-
-    // Register a shutdown task to terminate the balancer thread so that it doesn't leak memory.
-    registerShutdownTask([serviceContext] {
-        auto balancer = Balancer::get(serviceContext);
-        // Make sure that the balancer thread has been interrupted.
-        balancer->interruptBalancer();
-        // Make sure the balancer thread has terminated.
-        balancer->waitForBalancerToStop();
-    });
 }
 
 Balancer* Balancer::get(ServiceContext* serviceContext) {