SERVER-20407 Sharded findAndModify should retry on shard version mismatch

This is an approximate backport of
4821a88f0f92ea1f4ec3f46e71d5913ebd815fe5.

5 files changed, 140 insertions, 20 deletions

diff --git a/jstests/sharding/find_and_modify_after_multi_write.js b/jstests/sharding/find_and_modify_after_multi_write.js
index abce2e026a1..c8081ce9119 100644
--- a/jstests/sharding/find_and_modify_after_multi_write.js
+++ b/jstests/sharding/find_and_modify_after_multi_write.js
@@ -11,37 +11,74 @@ var runTest = function(writeFunc) {
     var testDB = st.s.getDB('test');
     testDB.dropDatabase();
 
-    testDB.adminCommand({ enableSharding: 'test' });
+    assert.commandWorked(testDB.adminCommand({ enableSharding: 'test' }));
     st.ensurePrimaryShard('test', 'shard0000');
-    testDB.adminCommand({ shardCollection: 'test.user', key: { x: 1 }});
-    testDB.adminCommand({ split: 'test.user', middle: { x: 0 }});
-    testDB.adminCommand({ moveChunk: 'test.user', find: { x: 0 }, to: 'shard0001' });
+
+    assert.commandWorked(testDB.adminCommand({ shardCollection: 'test.user', key: { x: 1 }}));
+    assert.commandWorked(testDB.adminCommand({ split: 'test.user', middle: { x: 0 }}));
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0001',
+                                               _waitForDelete: true }));
 
     var testDB2 = st.s1.getDB('test');
     testDB2.user.insert({ x: 123456 });
 
-    // Force ssv initialization of 'test.user' ns for this mongos.
-    var doc = testDB2.user.findOne({ x: 123456 });
-    assert.neq(null, doc);
-
     // Move chunk to bump version on a different mongos.
-    testDB.adminCommand({ moveChunk: 'test.user', find: { x: 0 }, to: 'shard0000' });
-
-    writeFunc(testDB2);
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0000',
+                                               _waitForDelete: true }));
 
     // Issue a targetted findAndModify and check that it was upserted to the right shard.
-    var res = testDB2.runCommand({
+    assert.commandWorked(testDB2.runCommand({
         findAndModify: 'user',
         query: { x: 100 },
         update: { $set: { y: 1 }},
         upsert: true
-    });
-
-    assert.commandWorked(res);
+    }));
 
     assert.neq(null, st.d0.getDB('test').user.findOne({ x: 100 }));
     assert.eq(null, st.d1.getDB('test').user.findOne({ x: 100 }));
 
+    // At this point, s1 thinks the version of 'test.user' is 2, bounce it again so it gets
+    // incremented to 3
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0001',
+                                               _waitForDelete: true }));
+
+    assert.commandWorked(testDB2.runCommand({
+        findAndModify: 'user',
+        query: { x: 200 },
+        update: { $set: { y: 1 }},
+        upsert: true
+    }));
+
+    assert.eq(null, st.d0.getDB('test').user.findOne({ x: 200 }));
+    assert.neq(null, st.d1.getDB('test').user.findOne({ x: 200 }));
+
+    // At this point, s0 thinks the version of 'test.user' is 3, bounce it again so it gets
+    // incremented to 4
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0000',
+                                               _waitForDelete: true }));
+
+    // Ensure that write commands with multi version do not reset the connection shard version to
+    // ignored.
+    writeFunc(testDB2);
+
+    assert.commandWorked(testDB2.runCommand({
+        findAndModify: 'user',
+        query: { x: 300 },
+        update: { $set: { y: 1 }},
+        upsert: true
+    }));
+
+    assert.neq(null, st.d0.getDB('test').user.findOne({ x: 300 }));
+    assert.eq(null, st.d1.getDB('test').user.findOne({ x: 300 }));
+
     st.stop();
 };
 
diff --git a/jstests/sharding/query_after_multi_write.js b/jstests/sharding/query_after_multi_write.js
new file mode 100644
index 00000000000..74867dfd970
--- /dev/null
+++ b/jstests/sharding/query_after_multi_write.js
@@ -0,0 +1,70 @@
+(function() {
+"use strict";
+
+/**
+ * Test that queries will be properly routed after executing a write that does not
+ * perform any shard version checks.
+ */
+var runTest = function(writeFunc) {
+    var st = new ShardingTest({ shards: 2, mongos: 2 });
+
+    var testDB = st.s.getDB('test');
+    testDB.dropDatabase();
+
+    assert.commandWorked(testDB.adminCommand({ enableSharding: 'test' }));
+    st.ensurePrimaryShard('test', 'shard0000');
+
+    assert.commandWorked(testDB.adminCommand({ shardCollection: 'test.user', key: { x: 1 }}));
+    assert.commandWorked(testDB.adminCommand({ split: 'test.user', middle: { x: 0 }}));
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0001',
+                                               _waitForDelete: true }));
+
+    var testDB2 = st.s1.getDB('test');
+    testDB2.user.insert({ x: 123456 });
+
+    // Move chunk to bump version on a different mongos.
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0000',
+                                               _waitForDelete: true }));
+
+    // Issue a query and make sure it gets routed to the right shard.
+    assert.neq(null, testDB2.user.findOne({ x: 123456 }));
+
+    // At this point, s1 thinks the version of 'test.user' is 2, bounce it again so it gets
+    // incremented to 3
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0001',
+                                               _waitForDelete: true }));
+
+    // Issue a query and make sure it gets routed to the right shard again.
+    assert.neq(null, testDB2.user.findOne({ x: 123456 }));
+
+    // At this point, s0 thinks the version of 'test.user' is 3, bounce it again so it gets
+    // incremented to 4
+    assert.commandWorked(testDB.adminCommand({ moveChunk: 'test.user',
+                                               find: { x: 0 },
+                                               to: 'shard0000',
+                                               _waitForDelete: true }));
+
+    // Ensure that write commands with multi version do not reset the connection shard version to
+    // ignored.
+    writeFunc(testDB2);
+
+    assert.neq(null, testDB2.user.findOne({ x: 123456 }));
+
+    st.stop();
+};
+
+runTest(function(db) {
+    db.user.update({}, { $inc: { y: 987654 }}, false, true);
+});
+
+runTest(function(db) {
+    db.user.remove({ y: 'noMatch' }, false);
+});
+
+})();
diff --git a/src/mongo/s/commands_public.cpp b/src/mongo/s/commands_public.cpp
index 0d9d0d3b80e..f2d3789c2a6 100644
--- a/src/mongo/s/commands_public.cpp
+++ b/src/mongo/s/commands_public.cpp
@@ -1315,7 +1315,7 @@ public:
         uassert(13343, "query for sharded findAndModify must have shardkey", !shardKey.isEmpty());
 
         ChunkPtr chunk = cm->findIntersectingChunk(shardKey);
-        ShardConnection conn(chunk->getShard(), fullns);
+        ShardConnection conn(chunk->getShard(), fullns, cm);
         BSONObj res;
         bool ok = conn->runCommand(conf->getName(), cmdObj, res);
         conn.done();
diff --git a/src/mongo/s/s_only.cpp b/src/mongo/s/s_only.cpp
index b529d3c8c94..c4bf137c386 100644
--- a/src/mongo/s/s_only.cpp
+++ b/src/mongo/s/s_only.cpp
@@ -148,13 +148,14 @@ void Command::execCommandClientBasic(OperationContext* txn,
     c->_commandsExecuted.increment();
 
     std::string errmsg;
-    bool ok;
+    bool ok = false;
     try {
         ok = c->run(txn, dbname, cmdObj, queryOptions, errmsg, result, false);
-    } catch (DBException& e) {
-        ok = false;
+    } catch (const DBException& e) {
         int code = e.getCode();
-        if (code == RecvStaleConfigCode) {  // code for StaleConfigException
+
+        // Codes for StaleConfigException
+        if (code == RecvStaleConfigCode || code == SendStaleConfigCode) {
             throw;
         }
 
diff --git a/src/mongo/s/shard.h b/src/mongo/s/shard.h
index f8672d6536e..6f69d3351f6 100644
--- a/src/mongo/s/shard.h
+++ b/src/mongo/s/shard.h
@@ -238,6 +238,18 @@ typedef boost::shared_ptr<const ChunkManager> ChunkManagerPtr;
 
 class ShardConnection : public AScopedConnection {
 public:
+    /**
+     * Instantiates a new sharded connection, which will be associated with the specified chunk
+     * manager for versioning purposes.
+     *
+     * @param shard/addr Shard to use.
+     * @param ns Namespace to associate the version with.
+     * @param manager The chunk manager, which should be used for obtaining shard version
+     *  information to be set on the connection. This value can be nullptr, if the connection will
+     *  not be versioned and not associated with a namespace, in which case only the
+     *  setShardVersion command will be invoked to initialize the remote shard. Otherwise, the
+     *  chunk manager will be used to obtain the shard version to set on the connection.
+     */
     ShardConnection(const Shard* s,
                     const std::string& ns,
                     ChunkManagerPtr manager = ChunkManagerPtr());