SERVER-35313 CleanupConcurrencyWorkloads resmoke hook needs to handle the balancer

(cherry picked from commit 844c1aac8ad3a3f596c06a0654e36a73262b0d03)

2 files changed, 23 insertions, 5 deletions

diff --git a/buildscripts/resmokelib/testing/fixtures/shardedcluster.py b/buildscripts/resmokelib/testing/fixtures/shardedcluster.py
index 70c1eaa4329..394323630bd 100644
--- a/buildscripts/resmokelib/testing/fixtures/shardedcluster.py
+++ b/buildscripts/resmokelib/testing/fixtures/shardedcluster.py
@@ -110,7 +110,7 @@ class ShardedClusterFixture(interface.Fixture):  # pylint: disable=too-many-inst
 
         # Turn off the balancer if it is not meant to be enabled.
         if not self.enable_balancer:
-            self._stop_balancer()
+            self.stop_balancer()
 
         # Turn off autosplit if it is not meant to be enabled.
         if not self.enable_autosplit:
@@ -147,11 +147,19 @@ class ShardedClusterFixture(interface.Fixture):  # pylint: disable=too-many-inst
                                  password=self.auth_options["password"],
                                  mechanism=self.auth_options["authenticationMechanism"])
 
-    def _stop_balancer(self, timeout_ms=60000):
+    def stop_balancer(self, timeout_ms=60000):
         """Stop the balancer."""
         client = self.mongo_client()
         self._auth_to_db(client)
         client.admin.command({"balancerStop": 1}, maxTimeMS=timeout_ms)
+        self.logger.info("Stopped the balancer")
+
+    def start_balancer(self, timeout_ms=60000):
+        """Start the balancer."""
+        client = self.mongo_client()
+        self._auth_to_db(client)
+        client.admin.command({"balancerStart": 1}, maxTimeMS=timeout_ms)
+        self.logger.info("Started the balancer")
 
     def _do_teardown(self):
         """Shut down the sharded cluster."""
@@ -163,7 +171,7 @@ class ShardedClusterFixture(interface.Fixture):  # pylint: disable=too-many-inst
                                 "but weren't.")
 
         if self.enable_balancer:
-            self._stop_balancer()
+            self.stop_balancer()
 
         teardown_handler = interface.FixtureTeardownHandler(self.logger)
 
diff --git a/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py b/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py
index c7954ca187b..4719b806cf2 100644
--- a/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py
+++ b/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py
@@ -4,8 +4,9 @@ from __future__ import absolute_import
 
 import copy
 
-from . import interface
-from ... import utils
+from buildscripts.resmokelib import utils
+from buildscripts.resmokelib.testing.hooks import interface
+from buildscripts.resmokelib.testing.fixtures import shardedcluster
 
 
 class CleanupConcurrencyWorkloads(interface.Hook):
@@ -67,6 +68,11 @@ class CleanupConcurrencyWorkloadsTestCase(interface.DynamicTestCase):
                 exclude_dbs.append(same_db_name)
         self.logger.info("Dropping all databases except for %s", exclude_dbs)
 
+        is_sharded_fixture = isinstance(self._hook.fixture, shardedcluster.ShardedClusterFixture)
+        # Stop the balancer.
+        if is_sharded_fixture and self._hook.fixture.enable_balancer:
+            self._hook.fixture.stop_balancer()
+
         for db_name in [db for db in db_names if db not in exclude_dbs]:
             self.logger.info("Dropping database %s", db_name)
             try:
@@ -87,3 +93,7 @@ class CleanupConcurrencyWorkloadsTestCase(interface.DynamicTestCase):
                     self.logger.exception("Encountered an error while dropping db % collection %s.",
                                           same_db_name, coll)
                     raise
+
+        # Start the balancer.
+        if is_sharded_fixture and self._hook.fixture.enable_balancer:
+            self._hook.fixture.start_balancer()