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 0f4e332bc6f..635f13bbd19 100644
--- a/src/mongo/db/db.cpp
+++ b/src/mongo/db/db.cpp
@@ -1267,6 +1267,12 @@ void shutdownTask() {
     auto const client = Client::getCurrent();
     auto const serviceContext = client->getServiceContext();
 
+    // Terminate the balancer thread so it doesn't leak memory.
+    if (auto balancer = Balancer::get(serviceContext)) {
+        balancer->interruptBalancer();
+        balancer->waitForBalancerToStop();
+    }
+
     // Shutdown the TransportLayer so that new connections aren't accepted
     if (auto tl = serviceContext->getTransportLayer()) {
         log(LogComponent::kNetwork) << "shutdown: going to close listening sockets...";
diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index 6622a04ad47..30d766726bf 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -170,15 +170,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) {