SERVER-20630 find command on mongod needs to wait until out of critical section

5 files changed, 21 insertions, 9 deletions

diff --git a/jstests/noPassthroughWithMongod/sharding_balance4.js b/jstests/noPassthroughWithMongod/sharding_balance4.js
index 542c0931b78..9ce404d9f95 100644
--- a/jstests/noPassthroughWithMongod/sharding_balance4.js
+++ b/jstests/noPassthroughWithMongod/sharding_balance4.js
@@ -5,9 +5,6 @@
 s = new ShardingTest( "slow_sharding_balance4" , 2 , 1 , 1 , { chunksize : 1 } )
 s.stopBalancer();
 
-// TODO: SERVER-20194. This test forces use of the old mongos query path.
-assert.commandWorked(s.s0.adminCommand({setParameter: 1, useClusterClientCursor: false}));
-
 s.adminCommand( { enablesharding : "test" } );
 s.ensurePrimaryShard('test', 'shard0001');
 s.adminCommand( { shardcollection : "test.foo" , key : { _id : 1 } } );
diff --git a/jstests/sharding/auto2.js b/jstests/sharding/auto2.js
index 0a869884b1d..0c0b58eef1c 100644
--- a/jstests/sharding/auto2.js
+++ b/jstests/sharding/auto2.js
@@ -6,9 +6,6 @@ s.adminCommand( { enablesharding : "test" } );
 s.ensurePrimaryShard('test', 'shard0001');
 s.adminCommand( { shardcollection : "test.foo" , key : { num : 1 } } );
 
-// TODO: SERVER-20194. This test forces use of the old mongos query path.
-assert.commandWorked(s.s0.adminCommand({setParameter: 1, useClusterClientCursor: false}));
-
 bigString = "";
 while ( bigString.length < 1024 * 50 )
     bigString += "asocsancdnsjfnsdnfsjdhfasdfasdfasdfnsadofnsadlkfnsaldknfsad";
diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp
index cd2d663c9be..088ecbcdb08 100644
--- a/src/mongo/db/commands/find_cmd.cpp
+++ b/src/mongo/db/commands/find_cmd.cpp
@@ -234,13 +234,25 @@ public:
         }
         std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
 
-        // If the received version is newer than the version cached in 'shardingState', then we have
-        // to refresh 'shardingState' from the config servers. We do this before acquiring locks so
-        // that we don't hold locks while waiting on the network.
         ShardingState* const shardingState = ShardingState::get(txn);
+
         if (OperationShardVersion::get(txn).hasShardVersion() && shardingState->enabled()) {
             ChunkVersion receivedVersion = OperationShardVersion::get(txn).getShardVersion(nss);
             ChunkVersion latestVersion;
+            // Wait for migration completion to get the correct chunk version.
+            const int maxTimeoutSec = 30;
+            int timeoutSec = cq->getParsed().getMaxTimeMS() / 1000;
+            if (!timeoutSec || timeoutSec > maxTimeoutSec) {
+                timeoutSec = maxTimeoutSec;
+            }
+
+            if (!shardingState->waitTillNotInCriticalSection(timeoutSec)) {
+                uasserted(ErrorCodes::LockTimeout, "Timeout while waiting for migration commit");
+            }
+
+            // If the received version is newer than the version cached in 'shardingState', then we
+            // have to refresh 'shardingState' from the config servers. We do this before acquiring
+            // locks so that we don't hold locks while waiting on the network.
             uassertStatusOK(shardingState->refreshMetadataIfNeeded(
                 txn, nss.ns(), receivedVersion, &latestVersion));
         }
diff --git a/src/mongo/db/s/migration_impl.cpp b/src/mongo/db/s/migration_impl.cpp
index b909a50227b..95a8b3514c5 100644
--- a/src/mongo/db/s/migration_impl.cpp
+++ b/src/mongo/db/s/migration_impl.cpp
@@ -78,6 +78,7 @@ WriteConcernOptions getDefaultWriteConcernForMigration() {
 }
 
 MONGO_FP_DECLARE(failMigrationCommit);
+MONGO_FP_DECLARE(hangBeforeLeavingCriticalSection);
 MONGO_FP_DECLARE(failMigrationConfigWritePrepare);
 MONGO_FP_DECLARE(failMigrationApplyOps);
 
@@ -511,6 +512,8 @@ Status ChunkMoveOperationState::commitMigration(OperationContext* txn) {
         }
     }
 
+    MONGO_FAIL_POINT_PAUSE_WHILE_SET(hangBeforeLeavingCriticalSection);
+
     shardingState->migrationSourceManager()->setInCriticalSection(false);
 
     // Migration is done, just log some diagnostics information
diff --git a/src/mongo/db/s/migration_source_manager.cpp b/src/mongo/db/s/migration_source_manager.cpp
index 0bb2cec7f61..f497c3f4d3d 100644
--- a/src/mongo/db/s/migration_source_manager.cpp
+++ b/src/mongo/db/s/migration_source_manager.cpp
@@ -548,6 +548,9 @@ bool MigrationSourceManager::waitTillNotInCriticalSection(int maxSecondsToWait)
     const auto deadline = stdx::chrono::system_clock::now() + Seconds(maxSecondsToWait);
     stdx::unique_lock<stdx::mutex> lk(_mutex);
     while (_inCriticalSection) {
+        log() << "Waiting for " << maxSecondsToWait
+              << " seconds for the migration critical section to end";
+
         if (stdx::cv_status::timeout == _inCriticalSectionCV.wait_until(lk, deadline)) {
             return false;
         }