From 6596261c307c4824c9ddd6de01faaf27fba2c9b6 Mon Sep 17 00:00:00 2001
From: Bailey Pearson <bailey.pearson@mongodb.com>
Date: Tue, 30 Aug 2022 11:38:50 -0400
Subject: [PATCH] fix: shuffle servers when there are only two servers

---
 src/sdam/topology.ts                          |  4 +---
 ...er_selection.prose.operation_count.test.ts | 23 +++++++++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/sdam/topology.ts b/src/sdam/topology.ts
index 69b82befac..03a0627343 100644
--- a/src/sdam/topology.ts
+++ b/src/sdam/topology.ts
@@ -904,9 +904,7 @@ function processWaitQueue(topology: Topology) {
     } else if (selectedDescriptions.length === 1) {
       selectedServer = topology.s.servers.get(selectedDescriptions[0].address);
     } else {
-      // don't shuffle the array if there are only two elements
-      const descriptions =
-        selectedDescriptions.length === 2 ? selectedDescriptions : shuffle(selectedDescriptions, 2);
+      const descriptions = shuffle(selectedDescriptions, 2);
       const server1 = topology.s.servers.get(descriptions[0].address);
       const server2 = topology.s.servers.get(descriptions[1].address);
 
diff --git a/test/integration/server-selection/server_selection.prose.operation_count.test.ts b/test/integration/server-selection/server_selection.prose.operation_count.test.ts
index 48e8c23cb8..7865ce507a 100644
--- a/test/integration/server-selection/server_selection.prose.operation_count.test.ts
+++ b/test/integration/server-selection/server_selection.prose.operation_count.test.ts
@@ -165,4 +165,27 @@ describe('operationCount-based Selection Within Latency Window - Prose Test', fu
     expect(percentageToHost1).to.be.greaterThan(40).and.lessThan(60);
     expect(percentageToHost2).to.be.greaterThan(40).and.lessThan(60);
   });
+
+  it(
+    'equally distributes operations with both hosts when requests are in parallel',
+    TEST_METADATA,
+    async function () {
+      const collection = client.db('test-db').collection('collection0');
+
+      const { insertedId } = await collection.insertOne({ name: 'bumpy' });
+
+      const n = 1000;
+
+      for (let i = 0; i < n; ++i) {
+        await collection.findOne({ _id: insertedId });
+      }
+
+      // Step 9: Using command monitoring events, assert that each mongos was selected roughly 50% of the time (within +/- 10%).
+      const [host1, host2] = seeds.map(seed => seed.split(':')[1]);
+      const percentageToHost1 = (counts[host1] / n) * 100;
+      const percentageToHost2 = (counts[host2] / n) * 100;
+      expect(percentageToHost1).to.be.greaterThan(40).and.lessThan(60);
+      expect(percentageToHost2).to.be.greaterThan(40).and.lessThan(60);
+    }
+  );
 });