SERVER-76046 Suspend the balancer while running test teardown hooks

4 files changed, 18 insertions, 24 deletions

diff --git a/buildscripts/resmokelib/testing/fixtures/shardedcluster.py b/buildscripts/resmokelib/testing/fixtures/shardedcluster.py
index eb7b3f894c6..193c3768448 100644
--- a/buildscripts/resmokelib/testing/fixtures/shardedcluster.py
+++ b/buildscripts/resmokelib/testing/fixtures/shardedcluster.py
@@ -139,7 +139,7 @@ class ShardedClusterFixture(interface.Fixture):
 
         # Turn off the balancer if it is not meant to be enabled.
         if not self.enable_balancer:
-            self.stop_balancer()
+            self.stop_balancer(join_migrations=False)
 
         # Inform mongos about each of the shards
         for idx, shard in enumerate(self.shards):
@@ -192,10 +192,15 @@ class ShardedClusterFixture(interface.Fixture):
                                 .format(port, interface.Fixture.AWAIT_READY_TIMEOUT_SECS))
                         time.sleep(0.1)
 
-    def stop_balancer(self, timeout_ms=300000):
+    # TODO SERVER-76343 remove the join_migrations parameter and the if clause depending on it.
+    def stop_balancer(self, timeout_ms=300000, join_migrations=True):
         """Stop the balancer."""
         client = interface.build_client(self, self.auth_options)
         client.admin.command({"balancerStop": 1}, maxTimeMS=timeout_ms)
+        if join_migrations:
+            for shard in self.shards:
+                shard_client = interface.build_client(shard.get_primary(), self.auth_options)
+                shard_client.admin.command({"_shardsvrJoinMigrations": 1})
         self.logger.info("Stopped the balancer")
 
     def start_balancer(self, timeout_ms=300000):
diff --git a/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py b/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py
index e391f8e5664..1c1caaf3cb5 100644
--- a/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py
+++ b/buildscripts/resmokelib/testing/hooks/cleanup_concurrency_workloads.py
@@ -68,11 +68,6 @@ 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:
@@ -93,7 +88,3 @@ 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()
diff --git a/buildscripts/resmokelib/testing/hooks/orphans.py b/buildscripts/resmokelib/testing/hooks/orphans.py
index 38d62f51010..2d2d3947aec 100644
--- a/buildscripts/resmokelib/testing/hooks/orphans.py
+++ b/buildscripts/resmokelib/testing/hooks/orphans.py
@@ -22,16 +22,3 @@ class CheckOrphansDeleted(jsfile.DataConsistencyHook):
         js_filename = os.path.join("jstests", "hooks", "run_check_orphans_are_deleted.js")
         super().__init__(hook_logger, fixture, js_filename, description,
                          shell_options=shell_options)
-
-    def after_test(self, test, test_report):
-        """Run the run_check_orphans_are_deleted.js hook."""
-
-        # We temporarily disable the balancer so more work isn't generated for the range deleter
-        # while the hook is running.
-        if self.fixture.enable_balancer:
-            self.fixture.stop_balancer()
-
-        super().after_test(test, test_report)
-
-        if self.fixture.enable_balancer:
-            self.fixture.start_balancer()
diff --git a/buildscripts/resmokelib/testing/job.py b/buildscripts/resmokelib/testing/job.py
index db29bb06a87..3379dd51777 100644
--- a/buildscripts/resmokelib/testing/job.py
+++ b/buildscripts/resmokelib/testing/job.py
@@ -7,6 +7,7 @@ from collections import namedtuple
 from buildscripts.resmokelib import config
 from buildscripts.resmokelib import errors
 from buildscripts.resmokelib.testing import testcases
+from buildscripts.resmokelib.testing.fixtures import shardedcluster
 from buildscripts.resmokelib.testing.fixtures.interface import create_fixture_table
 from buildscripts.resmokelib.testing.testcases import fixture as _fixture
 from buildscripts.resmokelib.utils import queue as _queue
@@ -282,7 +283,13 @@ class Job(object):
         @param test: the test after which we run the hooks.
         @param background: whether to run background hooks.
         """
+        suite_with_balancer = isinstance(
+            self.fixture, shardedcluster.ShardedClusterFixture) and self.fixture.enable_balancer
+
         try:
+            if not background and suite_with_balancer:
+                self.logger.info("Stopping the balancer before running end-test hooks")
+                self.fixture.stop_balancer()
             for hook in self.hooks:
                 if hook.IS_BACKGROUND == background:
                     self._run_hook(hook, hook.after_test, test, hook_failure_flag)
@@ -307,6 +314,10 @@ class Job(object):
             self.report.setError(test, sys.exc_info())
             raise
 
+        if not background and suite_with_balancer:
+            self.logger.info("Resuming the balancer after running end-test hooks")
+            self.fixture.start_balancer()
+
     def _fail_test(self, test, exc_info, return_code=1):
         """Provide helper to record a test as a failure with the provided return code.