1 files changed, 9 insertions, 0 deletions
diff --git a/jstests/concurrency/fsm_workloads/internal_transactions_sharded.js b/jstests/concurrency/fsm_workloads/internal_transactions_sharded.js
index 0e06f84b798..f55fa716950 100644
--- a/jstests/concurrency/fsm_workloads/internal_transactions_sharded.js
+++ b/jstests/concurrency/fsm_workloads/internal_transactions_sharded.js
@@ -17,6 +17,7 @@
 load('jstests/concurrency/fsm_libs/extend_workload.js');
 load('jstests/concurrency/fsm_workloads/random_moveChunk_base.js');
 load('jstests/concurrency/fsm_workloads/internal_transactions_unsharded.js');
+load('jstests/concurrency/fsm_workload_helpers/balancer.js');
 load('jstests/libs/fail_point_util.js');
 
 var $config = extendWorkload($config, function($config, $super) {
@@ -64,6 +65,10 @@ var $config = extendWorkload($config, function($config, $super) {
     $config.setup = function setup(db, collName, cluster) {
         const ns = db.getName() + "." + collName;
 
+        // Disallow balancing 'ns' during $setup so it does not interfere with the splits.
+        BalancerHelper.disableBalancerForCollection(db, ns);
+        BalancerHelper.joinBalancerRound(db);
+
         // Move the initial chunk to shard0.
         const shards = Object.keys(cluster.getSerializedCluster().shards);
         ChunkHelper.moveChunk(
@@ -82,6 +87,7 @@ var $config = extendWorkload($config, function($config, $super) {
                 assert.commandWorked(db.adminCommand(
                     {split: ns, middle: {[this.defaultShardKeyField]: partition.lower}}));
             }
+
             assert.commandWorked(
                 db.adminCommand({split: ns, middle: {[this.defaultShardKeyField]: partition.mid}}));
 
@@ -108,6 +114,9 @@ var $config = extendWorkload($config, function($config, $super) {
             }
         }
 
+        // Allow balancing 'ns' again.
+        BalancerHelper.enableBalancerForCollection(db, ns);
+
         this.overrideInternalTransactionsReapThreshold(cluster);
         if (this.lowerTransactionLifetimeLimitSeconds) {
             this.overrideTransactionLifetimeLimit(cluster);