SERVER-39630 Allow updates to the shard key value only when the document will not change shards

9 files changed, 919 insertions, 976 deletions

diff --git a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
index 6b57a4c6940..3433f173efc 100644
--- a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
@@ -63,6 +63,9 @@ selector:
   - jstests/sharding/shard7.js
   - jstests/sharding/shard_config_db_collections.js
   - jstests/sharding/unsharded_collection_targetting.js
+  - jstests/sharding/array_shard_key.js
+  - jstests/sharding/update_immutable_fields.js
+  - jstests/sharding/update_shard_key_doc_on_same_shard.js
   # TODO: SERVER-38541 remove from blacklist
   - jstests/sharding/shard_collection_existing_zones.js
   - jstests/sharding/single_shard_transaction_with_arbiter.js
diff --git a/buildscripts/resmokeconfig/suites/sharding_misc.yml b/buildscripts/resmokeconfig/suites/sharding_misc.yml
index 75a7d0e9c7a..2ef54fe8f49 100644
--- a/buildscripts/resmokeconfig/suites/sharding_misc.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_misc.yml
@@ -256,7 +256,6 @@ selector:
   - jstests/sharding/autosplit.js
   - jstests/sharding/limit_push.js
   - jstests/sharding/shard_keycount.js
-  - jstests/sharding/shard_key_immutable.js
   - jstests/sharding/sharding_system_namespaces.js
   - jstests/sharding/logical_time_metadata.js
   - jstests/sharding/time_zone_info_mongos.js
@@ -290,7 +289,6 @@ selector:
   - jstests/sharding/authConnectionHook.js
   - jstests/sharding/update_immutable_fields.js
   - jstests/sharding/min_optime_recovery_on_failed_move_chunk_commit.js
-  - jstests/sharding/SERVER-7379.js
   - jstests/sharding/create_database.js
   - jstests/sharding/shard_aware_on_add_shard.js
   - jstests/sharding/killop.js
diff --git a/jstests/multiVersion/update_shard_key_disallowed_fcv40.js b/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
index fc1cf7e8571..65c23d0bf1a 100644
--- a/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
+++ b/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
@@ -46,8 +46,11 @@
 
         // Assert that updating the shard key when the doc would remain on the same shard fails for
         // both modify and replacement updates
-        assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 5}}));
-        assert.writeError(sessionDB.foo.update({x: 30}, {x: 5}));
+
+        // TODO SERVER-40225: Uncomment the assertions below
+        // assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 5}}));
+        // assert.writeError(sessionDB.foo.update({x: 30}, {x: 5}));
+
         assert.throws(function() {
             sessionDB.foo.findAndModify({query: {x: 80}, update: {$set: {x: 100}}});
         });
@@ -81,8 +84,11 @@
 
         // Assert that updating the shard key when the doc would remain on the same shard fails for
         // both modify and replacement updates
-        assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 5}}));
-        assert.writeError(sessionDB.foo.update({x: 30}, {x: 5}));
+
+        // TODO SERVER-40225: Uncomment the assertions below
+        // assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 5}}));
+        // assert.writeError(sessionDB.foo.update({x: 30}, {x: 5}));
+
         assert.throws(function() {
             sessionDB.foo.findAndModify({query: {x: 80}, update: {$set: {x: 100}}});
         });
@@ -113,7 +119,9 @@
         // Assert that updating the shard key when the doc would remain on the same shard fails for
         // both modify and replacement updates
         session.startTransaction();
-        assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 5}}));
+        // TODO SERVER-40225: Uncomment the assertions below
+        // assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 5}}));
+        assert.writeError(sessionDB.foo.update({x: 80}, {$set: {x: 100}}));
         session.abortTransaction();
 
         session.startTransaction();
diff --git a/jstests/sharding/SERVER-7379.js b/jstests/sharding/SERVER-7379.js
deleted file mode 100644
index ed30be53bc6..00000000000
--- a/jstests/sharding/SERVER-7379.js
+++ /dev/null
@@ -1,45 +0,0 @@
-var st = new ShardingTest({shards: 2});
-
-st.adminCommand({enablesharding: "test"});
-st.ensurePrimaryShard('test', st.shard1.shardName);
-st.adminCommand(
-    {shardcollection: "test.offerChange", key: {"categoryId": 1, "store": 1, "_id": 1}});
-
-var db = st.s.getDB('test');
-var offerChange = db.getCollection('offerChange');
-var testDoc = {"_id": 123, "categoryId": 9881, "store": "NEW"};
-
-offerChange.remove({}, false);
-offerChange.insert(testDoc);
-assert.writeError(offerChange.update({_id: 123}, {$set: {store: "NEWEST"}}, true, false));
-var doc = offerChange.findOne();
-assert(friendlyEqual(doc, testDoc), 'doc changed: ' + tojson(doc));
-
-offerChange.remove({}, false);
-offerChange.insert(testDoc);
-assert.writeError(
-    offerChange.update({_id: 123}, {_id: 123, categoryId: 9881, store: "NEWEST"}, true, false));
-doc = offerChange.findOne();
-assert(friendlyEqual(doc, testDoc), 'doc changed: ' + tojson(doc));
-
-offerChange.remove({}, false);
-offerChange.insert(testDoc);
-assert.writeError(offerChange.save({"_id": 123, "categoryId": 9881, "store": "NEWEST"}));
-doc = offerChange.findOne();
-assert(friendlyEqual(doc, testDoc), 'doc changed: ' + tojson(doc));
-
-offerChange.remove({}, false);
-offerChange.insert(testDoc);
-assert.writeError(offerChange.update(
-    {_id: 123, store: "NEW"}, {_id: 123, categoryId: 9881, store: "NEWEST"}, true, false));
-doc = offerChange.findOne();
-assert(friendlyEqual(doc, testDoc), 'doc changed: ' + tojson(doc));
-
-offerChange.remove({}, false);
-offerChange.insert(testDoc);
-assert.writeError(offerChange.update(
-    {_id: 123, categoryId: 9881}, {_id: 123, categoryId: 9881, store: "NEWEST"}, true, false));
-doc = offerChange.findOne();
-assert(friendlyEqual(doc, testDoc), 'doc changed: ' + tojson(doc));
-
-st.stop();
diff --git a/jstests/sharding/shard_key_immutable.js b/jstests/sharding/shard_key_immutable.js
deleted file mode 100644
index 6e8c787393b..00000000000
--- a/jstests/sharding/shard_key_immutable.js
+++ /dev/null
@@ -1,869 +0,0 @@
-/**
- * Shard key invariant:
- *
- * A document must be created with a full non-array-or-regex shard key, and the value of that shard
- * key can never change.
- *
- * To enforce this invariant, we have the following mongos rule:
- *
- * - Upserts must always contain the full shard key and must only be targeted* to the applicable
- *shard.
- *
- * and the following mongod rules:
- *
- * - Upserted shard key values must not be arrays (or regexes).
- * - If a shard key value is present in the update query, upserts must only insert documents which
- *   match this value.
- * - Updates must not modify shard keys.
- *
- * *Updates are targeted by the update query if $op-style, or the update document if
- *replacement-style.
- *
- * NOTE: The above is enough to ensure that shard keys do not change.  It is not enough to ensure
- * uniqueness of an upserted document based on the upsert query.  This is necessary due to the
- *save()
- * style operation:
- * db.coll.update({ _id : xxx }, { _id : xxx, shard : xxx, key : xxx, other : xxx }, { upsert : true
- *})
- *
- * TODO: Minimize the impact of this hole by disallowing anything but save-style upserts of this
- *form.
- * Save-style upserts of this form are not safe (duplicate _ids can be created) but the user is
- * explicitly responsible for this for the _id field.
- *
- * In addition, there is an rule where non-multi updates can only affect 0 or 1 documents.
- *
- * To enforce this, we have the following mongos rule:
- *
- * - Non-multi updates must be targeted based on an exact _id query or the full shard key.
- *
- * Test setup:
- * - replacement style updates have the multiUpdate flag set to false.
- * - $ op updates have multiUpdate flag set to true.
- */
-
-var st = new ShardingTest({shards: 2});
-
-st.adminCommand({enablesharding: "test"});
-st.ensurePrimaryShard('test', st.shard1.shardName);
-st.adminCommand({shardcollection: "test.col0", key: {a: 1, b: 1}});
-st.adminCommand({shardcollection: "test.col1", key: {'x.a': 1}});
-
-var db = st.s.getDB('test');
-var compoundColl = db.getCollection('col0');
-var dotColl = db.getCollection('col1');
-
-//
-// Empty query update
-//
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({}, {a: 1}, false));
-var doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-doc = compoundColl.findOne();
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({}, {a: 1, b: 1}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({}, {a: 100, b: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({}, {a: 100, b: 100, _id: 1}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({}, {$set: {a: 1, b: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({}, {$set: {a: 100, b: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Cannot modify _id
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({}, {$set: {a: 1, b: 1, _id: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({}, {$set: {c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'doc did not change: ' + tojson(doc));
-
-//
-// Empty query upsert
-//
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({}, {a: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({}, {a: 1, b: 1}, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 1, b: 1}), 'doc not upserted properly: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({}, {a: 1, b: 1, _id: 1}, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 1, b: 1}), 'doc not upserted properly: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({}, {$set: {a: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({}, {$set: {a: 1, b: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({}, {$set: {a: 1, b: 1, _id: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({}, {$set: {c: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-//
-// Partial skey query update
-//
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {a: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {a: 2}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {a: 100, b: 1}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Inspecting query and update alone is not enough to tell whether a shard key will change.
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100}, {a: 100, b: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {a: 100, b: 100, _id: 1}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100}, {$set: {a: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {$set: {b: 200}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Inspecting query and update alone is not enough to tell whether a shard key will change.
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100}, {$set: {b: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {$set: {a: 100, b: 200}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Inspecting query and update alone is not enough to tell whether a shard key will change.
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100}, {$set: {a: 100, b: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100}, {$set: {a: 100, b: 100, _id: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100}, {$set: {c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'doc did not change: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100}, {$rename: {c: 'a'}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-//
-// Partial skey query upsert
-//
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {a: 100}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {a: 2}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {a: 1, b: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {a: 1, b: 1, _id: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {$set: {a: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {$set: {b: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {$set: {a: 100, b: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {$set: {a: 100, b: 1, _id: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {$set: {c: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100}, {$rename: {c: 'a'}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-//
-// Not prefix of skey query update
-//
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {b: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {b: 2}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {a: 1}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Inspecting query and update alone is not enough to tell whether a shard key will change.
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {a: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {a: 1, b: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Inspecting query and update alone is not enough to tell whether a shard key will change.
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({b: 100}, {a: 100, b: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {a: 1, b: 1, _id: 1}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {$set: {b: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {$set: {a: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({b: 100}, {$set: {a: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {$set: {a: 1, b: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Inspecting query and update alone is not enough to tell whether a shard key will change.
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({b: 100}, {$set: {a: 100, b: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({b: 100}, {$set: {a: 100, b: 100, _id: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({b: 100}, {$set: {c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'doc did not change: ' + tojson(doc));
-
-//
-// Not prefix of skey query upsert
-//
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {b: 100}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {b: 2}, true));
-
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {a: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {a: 1, b: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {a: 1, b: 1, _id: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {$set: {b: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {$set: {a: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {$set: {a: 1, b: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {$set: {a: 1, b: 1, _id: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({b: 100}, {$set: {c: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc upserted: ' + tojson(doc));
-
-//
-// Full skey query update
-//
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100, b: 100}, {a: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {a: 100, b: 100, c: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 100}), 'doc did not change: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100, b: 100}, {a: 100, b: 100, _id: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100, b: 100}, {b: 100}, false));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {b: 100, c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'doc did not change: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {a: 100, b: 100, c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'doc did not change: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(
-    compoundColl.update({a: 100, b: 100}, {$set: {a: 100, b: 100, _id: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeError(compoundColl.update({a: 100, b: 100}, {$set: {a: 100, b: 2, c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({a: 100, b: 100});
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {c: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'doc did not change: ' + tojson(doc));
-
-//
-// Full skey query upsert
-//
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100, b: 100}, {a: 100}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {a: 100, b: 100, c: 1}, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'wrong doc: ' + tojson(doc));
-
-// Cannot modify _id!
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {a: 100, b: 100, _id: 100}, true));
-doc = compoundColl.findOne();
-assert(friendlyEqual(doc, {_id: 100, a: 100, b: 100}), 'wrong doc: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100, b: 100}, {b: 100}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {b: 100, c: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc != null, 'doc was not upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {a: 100, b: 100, c: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc != null, 'doc was not upserted: ' + tojson(doc));
-
-// Can upsert with new _id
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {a: 100, b: 100, _id: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc != null, 'doc was not upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({a: 100, b: 100}, {$set: {a: 100, b: 2, c: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({a: 100, b: 100}, {$set: {c: 1}}, true, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100, c: 1}), 'wrong doc: ' + tojson(doc));
-
-//
-// _id query update
-//
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeError(compoundColl.update({_id: 1}, {a: 1}));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-// Special case for _id. This is for making save method work.
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeOK(compoundColl.update({_id: 1}, {a: 100, b: 100}));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeError(compoundColl.update({_id: 1}, {a: 1, b: 1}));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeError(compoundColl.update({_id: 1}, {$set: {a: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeOK(compoundColl.update({_id: 1}, {$set: {a: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeError(compoundColl.update({_id: 1}, {$set: {b: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeOK(compoundColl.update({_id: 1}, {$set: {b: 100}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-compoundColl.insert({_id: 1, a: 100, b: 100});
-assert.writeError(compoundColl.update({_id: 1}, {$set: {a: 1, b: 1}}, false, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 100, b: 100}), 'doc changed: ' + tojson(doc));
-
-//
-// _id query upsert
-//
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({_id: 1}, {a: 1}, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeOK(compoundColl.update({_id: 1}, {a: 1, b: 1}, true));
-doc = compoundColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {a: 1, b: 1}), 'bad doc: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({_id: 1}, {$set: {a: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({_id: 1}, {$set: {b: 1}}, true, true));
-doc = compoundColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-compoundColl.remove({}, false);
-assert.writeError(compoundColl.update({_id: 1}, {$set: {a: 1, b: 1}}, true, true));
-assert.eq(0, compoundColl.count(), 'doc should not be inserted');
-
-//
-// Dotted query update
-//
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeOK(dotColl.update({'x.a': 100}, {x: {a: 100, b: 2}}));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100, b: 2}}), 'doc did not change: ' + tojson(doc));
-
-// Dotted field names in the update object do not represent paths. The path must be expanded
-// since single document updates require a match on the shard key.
-// (e.g. 'x.a' will not work, but x: {a: ...} will work).
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeErrorWithCode(dotColl.update({'x.a': 100}, {x: {'a.z': 100}}),
-                          ErrorCodes.ShardKeyNotFound);
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeErrorWithCode(dotColl.update({'x.a': 100}, {'x.a': 100}), ErrorCodes.ShardKeyNotFound);
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeOK(dotColl.update({'x.a': 100}, {x: {a: 100}, b: 1}));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {x: 100}));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {x: {b: 100}}));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeOK(dotColl.update({'x.a': 100}, {$set: {x: {a: 100, b: 2}}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100, b: 2}}), 'doc did not change: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: {a: 2}}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: {b: 100}}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-// Dotted field names within $set are treated as paths.
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeOK(dotColl.update({'x.a': 100}, {$set: {'x.a': 100, b: 2}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}, b: 2}), 'doc did not change: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: {'a.z': 100}}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {'x.a.z': 100}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: 100}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}}), 'doc changed: ' + tojson(doc));
-
-dotColl.remove({}, false);
-dotColl.insert({x: {a: 100}});
-assert.writeOK(dotColl.update({'x.a': 100}, {$set: {'x.b': 200}}, false, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100, b: 200}}), 'doc did not change: ' + tojson(doc));
-
-//
-// Dotted query upsert
-//
-
-dotColl.remove({}, false);
-assert.writeOK(dotColl.update({'x.a': 100}, {x: {a: 100, b: 2}}, true));
-doc = dotColl.findOne();
-assert(doc != null, 'doc was not upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {x: {z: 100}}, true));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {'x.a': 100}, true));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {'x.a.z': 100}, true));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {x: 100}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {x: {b: 100}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeOK(dotColl.update({'x.a': 100}, {$set: {x: {a: 100, b: 2}}}, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100, b: 2}}), 'bad doc: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: {a: 2}}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: {b: 100}}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeOK(dotColl.update({'x.a': 100}, {$set: {'x.a': 100, b: 3}}, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100}, b: 3}), 'bad doc: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {'x.a': 2}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: {'a.z': 100}}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {'x.a.z': 100}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeError(dotColl.update({'x.a': 100}, {$set: {x: 100}}, true));
-doc = dotColl.findOne();
-assert(doc == null, 'doc was upserted: ' + tojson(doc));
-
-dotColl.remove({}, false);
-assert.writeOK(dotColl.update({'x.a': 100}, {$set: {'x.b': 2}}, true));
-doc = dotColl.findOne();
-delete doc._id;
-assert(friendlyEqual(doc, {x: {a: 100, b: 2}}), 'bad doc: ' + tojson(doc));
-
-st.stop();
diff --git a/jstests/sharding/update_immutable_fields.js b/jstests/sharding/update_immutable_fields.js
index b2a422a98cb..5e81d075c37 100644
--- a/jstests/sharding/update_immutable_fields.js
+++ b/jstests/sharding/update_immutable_fields.js
@@ -1,4 +1,4 @@
-// Tests that updates can't change immutable fields (used in sharded system)
+// Tests that save style updates correctly change immutable fields
 (function() {
     'use strict';
 
diff --git a/jstests/sharding/update_shard_key_doc_on_same_shard.js b/jstests/sharding/update_shard_key_doc_on_same_shard.js
new file mode 100644
index 00000000000..6bc22918bc5
--- /dev/null
+++ b/jstests/sharding/update_shard_key_doc_on_same_shard.js
@@ -0,0 +1,833 @@
+/*
+ * Tests that changing the shard key value of a document using update and findAndModify works
+ * correctly when the new shard key value belongs to the same shard.
+ * @tags: [uses_transactions, uses_multi_shard_transaction]
+ */
+
+(function() {
+    'use strict';
+
+    load("jstests/aggregation/extras/utils.js");
+
+    var st = new ShardingTest({mongos: 1, shards: 2});
+    var kDbName = 'db';
+    var mongos = st.s0;
+    var shard0 = st.shard0.shardName;
+    var shard1 = st.shard1.shardName;
+
+    assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
+    st.ensurePrimaryShard(kDbName, shard0);
+
+    function shardCollectionMoveChunks(shardKey, docsToInsert, splitDoc, moveChunkDoc) {
+        assert.commandWorked(mongos.getDB(kDbName).foo.createIndex(shardKey));
+
+        var ns = kDbName + '.foo';
+        assert.eq(mongos.getDB('config').collections.count({_id: ns, dropped: false}), 0);
+
+        for (let i = 0; i < docsToInsert.length; i++) {
+            assert.commandWorked(mongos.getDB(kDbName).foo.insert(docsToInsert[i]));
+        }
+
+        assert.commandWorked(mongos.adminCommand({shardCollection: ns, key: shardKey}));
+
+        if (docsToInsert.length > 0) {
+            assert.commandWorked(mongos.adminCommand({split: ns, find: splitDoc}));
+            assert.commandWorked(
+                mongos.adminCommand({moveChunk: ns, find: moveChunkDoc, to: shard1}));
+        }
+
+        assert.commandWorked(st.shard0.adminCommand({_flushDatabaseCacheUpdates: kDbName}));
+        assert.commandWorked(st.shard0.adminCommand({_flushRoutingTableCacheUpdates: ns}));
+        assert.commandWorked(st.shard1.adminCommand({_flushDatabaseCacheUpdates: kDbName}));
+        assert.commandWorked(st.shard1.adminCommand({_flushRoutingTableCacheUpdates: ns}));
+    }
+
+    function runUpdateCmdSuccess(inTxn, queries, updates, upsert) {
+        let res;
+        for (let i = 0; i < queries.length; i++) {
+            if (inTxn) {
+                session.startTransaction();
+                res = sessionDB.foo.update(queries[i], updates[i], {"upsert": upsert});
+                assert.commandWorked(res);
+                session.commitTransaction();
+            } else {
+                res = sessionDB.foo.update(queries[i], updates[i], {"upsert": upsert});
+                assert.commandWorked(res);
+            }
+
+            let updatedVal = updates[i]["$set"] ? updates[i]["$set"] : updates[i];
+            assert.eq(0, mongos.getDB(kDbName).foo.find(queries[i]).itcount());
+            assert.eq(1, mongos.getDB(kDbName).foo.find(updatedVal).itcount());
+            if (upsert) {
+                assert.eq(1, res.nUpserted);
+                assert.eq(0, res.nMatched);
+                assert.eq(0, res.nModified);
+            } else {
+                assert.eq(0, res.nUpserted);
+                assert.eq(1, res.nMatched);
+                assert.eq(1, res.nModified);
+            }
+        }
+    }
+
+    function runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, returnNew) {
+        let res;
+        for (let i = 0; i < queries.length; i++) {
+            let oldDoc;
+            if (!returnNew && !upsert) {
+                oldDoc = mongos.getDB(kDbName).foo.find(queries[i]).toArray();
+            }
+
+            if (inTxn) {
+                session.startTransaction();
+                res = sessionDB.foo.findAndModify(
+                    {query: queries[i], update: updates[i], "upsert": upsert, "new": returnNew});
+                session.commitTransaction();
+            } else {
+                res = sessionDB.foo.findAndModify(
+                    {query: queries[i], update: updates[i], "upsert": upsert, "new": returnNew});
+            }
+
+            let updatedVal = updates[i]["$set"] ? updates[i]["$set"] : updates[i];
+            let newDoc = mongos.getDB(kDbName).foo.find(updatedVal).toArray();
+            assert.eq(0, mongos.getDB(kDbName).foo.find(queries[i]).itcount());
+            assert.eq(1, newDoc.length);
+            if (returnNew) {
+                assert(resultsEq([res], newDoc));
+            } else {
+                if (upsert) {
+                    assert.eq(null, res);
+                } else {
+                    assert(resultsEq([res], oldDoc));
+                }
+            }
+        }
+    }
+
+    function runUpdateCmdFail(inTxn, query, update, multiParamSet) {
+        let res;
+        if (inTxn) {
+            session.startTransaction();
+            res = sessionDB.foo.update(query, update, {multi: multiParamSet});
+            assert.writeError(res);
+            session.abortTransaction();
+        } else {
+            res = sessionDB.foo.update(query, update, {multi: multiParamSet});
+            assert.writeError(res);
+        }
+
+        let updatedVal = update["$set"] ? update["$set"] : update;
+        assert.eq(1, mongos.getDB(kDbName).foo.find(query).itcount());
+        if (!update["$unset"]) {
+            assert.eq(0, mongos.getDB(kDbName).foo.find(updatedVal).itcount());
+        }
+    }
+
+    function runFindAndModifyCmdFail(inTxn, query, update, upsert) {
+        if (inTxn) {
+            session.startTransaction();
+            assert.throws(function() {
+                sessionDB.foo.findAndModify({query: query, update: update, "upsert": upsert});
+            });
+            session.abortTransaction();
+        } else {
+            assert.throws(function() {
+                sessionDB.foo.findAndModify({query: query, update: update, "upsert": upsert});
+            });
+        }
+        let updatedVal = update["$set"] ? update["$set"] : update;
+        assert.eq(1, mongos.getDB(kDbName).foo.find(query).itcount());
+        if (!update["$unset"]) {
+            assert.eq(0, mongos.getDB(kDbName).foo.find(updatedVal).itcount());
+        }
+    }
+
+    function assertCanUpdatePrimitiveShardKey(inTxn, isFindAndModify, queries, updates, upsert) {
+        let docsToInsert = [];
+        if (!upsert) {
+            docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+        }
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505}));
+        if (isFindAndModify) {
+            // Run once with {new: false} and once with {new: true} to make sure findAndModify
+            // returns pre and post images correctly
+            runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, false);
+            mongos.getDB(kDbName).foo.drop();
+
+            shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+            // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort
+            // txn
+            assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505}));
+            runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, true);
+        } else {
+            runUpdateCmdSuccess(inTxn, queries, updates, upsert);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCanUpdateDottedPath(inTxn, isFindAndModify, queries, updates, upsert) {
+        let docsToInsert = [];
+        if (!upsert) {
+            docsToInsert = [
+                {"x": {"a": 4, "y": 1}, "a": 3},
+                {"x": {"a": 100, "y": 1}},
+                {"x": {"a": 300, "y": 1}, "a": 3},
+                {"x": {"a": 500, "y": 1}, "a": 6}
+            ];
+        }
+        shardCollectionMoveChunks({"x.a": 1}, docsToInsert, {"x.a": 100}, {"x.a": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": {"a": 505}}));
+        if (isFindAndModify) {
+            // Run once with {new: false} and once with {new: true} to make sure findAndModify
+            // returns pre and post images correctly
+            runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, false);
+            mongos.getDB(kDbName).foo.drop();
+
+            shardCollectionMoveChunks({"x.a": 1}, docsToInsert, {"x.a": 100}, {"x.a": 300});
+            // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort
+            // txn
+            assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": {"a": 505}}));
+            runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, true);
+        } else {
+            runUpdateCmdSuccess(inTxn, queries, updates, upsert);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCanUpdatePartialShardKey(inTxn, isFindAndModify, queries, updates, upsert) {
+        let docsToInsert = [];
+        if (!upsert) {
+            docsToInsert =
+                [{"x": 4, "y": 3}, {"x": 100, "y": 50}, {"x": 300, "y": 80}, {"x": 500, "y": 600}];
+        }
+        shardCollectionMoveChunks(
+            {"x": 1, "y": 1}, docsToInsert, {"x": 100, "y": 50}, {"x": 300, "y": 80});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505, "y": 90}));
+        if (isFindAndModify) {
+            // Run once with {new: false} and once with {new: true} to make sure findAndModify
+            // returns pre and post images correctly
+            runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, false);
+            mongos.getDB(kDbName).foo.drop();
+
+            shardCollectionMoveChunks(
+                {"x": 1, "y": 1}, docsToInsert, {"x": 100, "y": 50}, {"x": 300, "y": 80});
+            // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort
+            // txn
+            assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505, "y": 90}));
+            runFindAndModifyCmdSuccess(inTxn, queries, updates, upsert, true);
+        } else {
+            runUpdateCmdSuccess(inTxn, queries, updates, upsert);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCannotUpdate_id(inTxn, isFindAndModify, query, update) {
+        let docsToInsert =
+            [{"_id": 4, "a": 3}, {"_id": 100}, {"_id": 300, "a": 3}, {"_id": 500, "a": 6}];
+        shardCollectionMoveChunks({"_id": 1}, docsToInsert, {"_id": 100}, {"_id": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"_id": 505}));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(inTxn, query, update);
+        } else {
+            runUpdateCmdFail(inTxn, query, update, false);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCannotUpdate_idDottedPath(inTxn, isFindAndModify, query, update) {
+        let docsToInsert = [
+            {"_id": {"a": 4, "y": 1}, "a": 3},
+            {"_id": {"a": 100, "y": 1}},
+            {"_id": {"a": 300, "y": 1}, "a": 3},
+            {"_id": {"a": 500, "y": 1}, "a": 6}
+        ];
+        shardCollectionMoveChunks({"_id.a": 1}, docsToInsert, {"_id.a": 100}, {"_id.a": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"_id": {"a": 505}}));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(inTxn, query, update);
+        } else {
+            runUpdateCmdFail(inTxn, query, update, false);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        inTxn, isFindAndModify, query, update) {
+        let docsToInsert =
+            [{"x": 4, "y": 3}, {"x": 100, "y": 50}, {"x": 300, "y": 80}, {"x": 500, "y": 600}];
+        shardCollectionMoveChunks(
+            {"x": 1, "y": 1}, docsToInsert, {"x": 100, "y": 50}, {"x": 300, "y": 80});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505, "y": 90}));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(inTxn, query, update);
+        } else {
+            runUpdateCmdFail(inTxn, query, update, false);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCannotUpdateWithMultiTrue(inTxn, query, update) {
+        let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505}));
+        runUpdateCmdFail(inTxn, query, update, true);
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCannotUpdateSKToArray(inTxn, isFindAndModify, query, update) {
+        docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505}));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(inTxn, query, update);
+        } else {
+            runUpdateCmdFail(inTxn, query, update, false);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    function assertCannotUnsetSKField(inTxn, isFindAndModify, query, update) {
+        // Updates to the shard key cannot $unset a shard key field from a doc
+        docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        assert.commandWorked(mongos.getDB(kDbName).foo.insert({"x": 505}));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(inTxn, query, update);
+        } else {
+            runUpdateCmdFail(inTxn, query, update, false);
+        }
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    // Shard key updates are allowed in bulk ops if the update doesn't cause the doc to move shards
+    function assertCanUpdateInBulkOp(inTxn, ordered) {
+        let bulkOp;
+        let bulkRes;
+
+        // Update multiple documents on different shards
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        mongos.getDB(kDbName).foo.insert({"x": 505});
+        if (inTxn) {
+            session.startTransaction();
+        }
+        if (ordered) {
+            bulkOp = sessionDB.foo.initializeOrderedBulkOp();
+        } else {
+            bulkOp = sessionDB.foo.initializeUnorderedBulkOp();
+        }
+        bulkOp.find({"x": 300}).updateOne({"$set": {"x": 600}});
+        bulkOp.find({"x": 4}).updateOne({"$set": {"x": 30}});
+        bulkRes = bulkOp.execute();
+        assert.commandWorked(bulkRes);
+        if (inTxn) {
+            session.commitTransaction();
+        }
+        assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
+        assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 4}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 30}).itcount());
+
+        assert.eq(0, bulkRes.nUpserted);
+        assert.eq(2, bulkRes.nMatched);
+        assert.eq(2, bulkRes.nModified);
+
+        mongos.getDB(kDbName).foo.drop();
+
+        // Check that final doc is correct after doing $inc on doc A and then updating the shard key
+        // for doc A. The outcome should be the same for both ordered and unordered bulk ops because
+        // the doc will not change shards, so both udpates will be targeted to the same shard.
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        mongos.getDB(kDbName).foo.insert({"x": 505});
+        if (inTxn) {
+            session.startTransaction();
+        }
+        if (ordered) {
+            bulkOp = sessionDB.foo.initializeOrderedBulkOp();
+        } else {
+            bulkOp = sessionDB.foo.initializeUnorderedBulkOp();
+        }
+        bulkOp.find({"x": 500}).updateOne({"$inc": {"a": 1}});
+        bulkOp.find({"x": 500}).updateOne({"$set": {"x": 400}});
+        bulkRes = bulkOp.execute();
+        assert.commandWorked(bulkRes);
+        if (inTxn) {
+            session.commitTransaction();
+        }
+        assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 400}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 400, "a": 7}).itcount());
+
+        assert.eq(0, bulkRes.nUpserted);
+        assert.eq(2, bulkRes.nMatched);
+        assert.eq(2, bulkRes.nModified);
+
+        mongos.getDB(kDbName).foo.drop();
+
+        // Check that updating the shard key for doc A, then doing $inc on the old doc A does not
+        // inc the field on the final doc. The outcome should be the same for both ordered and
+        // unordered bulk ops because the doc will not change shards, so both udpates will be
+        // targeted to the same shard.
+        shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+        // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+        mongos.getDB(kDbName).foo.insert({"x": 505});
+        if (inTxn) {
+            session.startTransaction();
+        }
+        if (ordered) {
+            bulkOp = sessionDB.foo.initializeOrderedBulkOp();
+        } else {
+            bulkOp = sessionDB.foo.initializeUnorderedBulkOp();
+        }
+        bulkOp.find({"x": 500}).updateOne({"x": 400, "a": 6});
+        bulkOp.find({"x": 500}).updateOne({"$inc": {"a": 1}});
+        bulkOp.insert({"x": 1, "a": 1});
+        bulkOp.find({"x": 400}).updateOne({"$set": {"x": 600}});
+        bulkRes = bulkOp.execute();
+        assert.commandWorked(bulkRes);
+        if (inTxn) {
+            session.commitTransaction();
+        }
+        assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
+        assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 400}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 600, "a": 6}).itcount());
+        assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 1}).itcount());
+
+        assert.eq(0, bulkRes.nUpserted);
+        assert.eq(2, bulkRes.nMatched);
+        assert.eq(2, bulkRes.nModified);
+        assert.eq(1, bulkRes.nInserted);
+
+        mongos.getDB(kDbName).foo.drop();
+    }
+    // -----------------------------------------
+    // Updates to the shard key are not allowed if write is not retryable and not in a multi-stmt
+    // txn
+    // -----------------------------------------
+
+    let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+    shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+    assert.writeError(mongos.getDB(kDbName).foo.update({"x": 300}, {"x": 600}));
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
+
+    assert.throws(function() {
+        mongos.getDB(kDbName).foo.findAndModify({query: {"x": 300}, update: {$set: {"x": 600}}});
+    });
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
+
+    mongos.getDB(kDbName).foo.drop();
+
+    // ---------------------------------
+    // Update shard key retryable write
+    // ---------------------------------
+
+    let session = st.s.startSession({retryWrites: true});
+    let sessionDB = session.getDatabase(kDbName);
+
+    // Modify updates
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        false, false, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], false);
+    assertCanUpdateDottedPath(false,
+                              false,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              false);
+    assertCanUpdatePartialShardKey(false,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        false, false, [{"x": 900}, {"x": 3}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], true);
+    assertCanUpdateDottedPath(false,
+                              false,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              true);
+    assertCanUpdatePartialShardKey(false,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(false, false, {"_id": 300}, {"$set": {"_id": 600}});
+    assertCannotUpdate_idDottedPath(false, false, {"_id.a": 300}, {"$set": {"_id": {"a": 600}}});
+    assertCannotUpdateWithMultiTrue(false, {"x": 300}, {"$set": {"x": 600}});
+    assertCannotUpdateSKToArray(false, false, {"x": 300}, {"$set": {"x": [300]}});
+    assertCannotUnsetSKField(false, false, {"x": 300}, {"$unset": {"x": 1}});
+
+    // Replacement updates
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        false, false, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], false);
+    assertCanUpdateDottedPath(
+        false, false, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], false);
+    assertCanUpdatePartialShardKey(false,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        false, false, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], true);
+    assertCanUpdateDottedPath(
+        false, false, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], true);
+    assertCanUpdatePartialShardKey(false,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(false, false, {"_id": 300}, {"_id": 600});
+    assertCannotUpdate_idDottedPath(false, false, {"_id.a": 300}, {"_id": {"a": 600}});
+    assertCannotUpdateWithMultiTrue(false, {"x": 300}, {"x": 600});
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        false, false, {"x": 300, "y": 80}, {"x": 600});
+    assertCannotUpdateSKToArray(false, false, {"x": 300}, {"x": [300]});
+
+    // Modify style findAndModify
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        false, true, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], false);
+    assertCanUpdateDottedPath(false,
+                              true,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              false);
+    assertCanUpdatePartialShardKey(false,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        false, false, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], true);
+    assertCanUpdateDottedPath(false,
+                              false,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              true);
+    assertCanUpdatePartialShardKey(false,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(false, true, {"_id": 300}, {"$set": {"_id": 600}});
+    assertCannotUpdate_idDottedPath(false, true, {"_id.a": 300}, {"$set": {"_id": {"a": 600}}});
+    assertCannotUpdateSKToArray(false, true, {"x": 300}, {"$set": {"x": [300]}});
+    assertCannotUnsetSKField(false, true, {"x": 300}, {"$unset": {"x": 1}});
+
+    // Replacement style findAndModify
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        false, true, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], false);
+    assertCanUpdateDottedPath(
+        false, true, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], false);
+    assertCanUpdatePartialShardKey(false,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   false);
+
+    // upsert: true
+    assertCanUpdatePrimitiveShardKey(
+        false, false, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], true);
+    assertCanUpdateDottedPath(
+        false, false, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], true);
+    assertCanUpdatePartialShardKey(false,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(false, true, {"_id": 300}, {"_id": 600});
+    assertCannotUpdate_idDottedPath(false, true, {"_id.a": 300}, {"_id": {"a": 600}});
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        false, true, {"x": 300, "y": 80}, {"x": 600});
+    assertCannotUpdateSKToArray(false, true, {"x": 300}, {"x": [300]});
+
+    // Bulk writes retryable writes
+    assertCanUpdateInBulkOp(false, false);
+    assertCanUpdateInBulkOp(false, true);
+
+    // ---------------------------------------
+    // Update shard key in multi statement txn
+    // ---------------------------------------
+
+    session = st.s.startSession();
+    sessionDB = session.getDatabase(kDbName);
+
+    // ----Single writes in txn----
+
+    // Modify updates
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        true, false, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], false);
+    assertCanUpdateDottedPath(true,
+                              false,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              false);
+    assertCanUpdatePartialShardKey(true,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        true, false, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], true);
+    assertCanUpdateDottedPath(true,
+                              false,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              true);
+    assertCanUpdatePartialShardKey(true,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(true, false, {"_id": 300}, {"$set": {"_id": 600}});
+    assertCannotUpdate_idDottedPath(true, false, {"_id.a": 300}, {"$set": {"_id": {"a": 600}}});
+    assertCannotUpdateWithMultiTrue(true, {"x": 300}, {"$set": {"x": 600}});
+    assertCannotUpdateSKToArray(true, false, {"x": 300}, {"$set": {"x": [300]}});
+    assertCannotUnsetSKField(true, false, {"x": 300}, {"$unset": {"x": 1}});
+
+    // Replacement updates
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        true, false, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], false);
+    assertCanUpdateDottedPath(
+        true, false, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], false);
+    assertCanUpdatePartialShardKey(true,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        true, false, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], true);
+    assertCanUpdateDottedPath(
+        true, false, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], true);
+    assertCanUpdatePartialShardKey(true,
+                                   false,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(true, false, {"_id": 300}, {"_id": 600});
+    assertCannotUpdate_idDottedPath(true, false, {"_id.a": 300}, {"_id": {"a": 600}});
+    assertCannotUpdateWithMultiTrue(true, {"x": 300}, {"x": 600});
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        true, false, {"x": 300, "y": 80}, {"x": 600});
+    assertCannotUpdateSKToArray(true, false, {"x": 300}, {"x": [300]});
+
+    // Modify style findAndModify
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        true, true, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], false);
+    assertCanUpdateDottedPath(true,
+                              true,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              false);
+    assertCanUpdatePartialShardKey(true,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        true, true, [{"x": 300}, {"x": 4}], [{"$set": {"x": 600}}, {"$set": {"x": 30}}], true);
+    assertCanUpdateDottedPath(true,
+                              true,
+                              [{"x.a": 300}, {"x.a": 4}],
+                              [{"$set": {"x": {"a": 600}}}, {"$set": {"x": {"a": 30}}}],
+                              true);
+    assertCanUpdatePartialShardKey(true,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"$set": {"x": 600}}, {"$set": {"x": 30}}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(true, true, {"_id": 300}, {"$set": {"_id": 600}});
+    assertCannotUpdate_idDottedPath(true, true, {"_id.a": 300}, {"$set": {"_id": {"a": 600}}});
+    assertCannotUpdateSKToArray(true, true, {"x": 300}, {"$set": {"x": [300]}});
+    assertCannotUnsetSKField(true, true, {"x": 300}, {"$unset": {"x": 1}});
+
+    // Replacement style findAndModify
+
+    // upsert : false
+    assertCanUpdatePrimitiveShardKey(
+        true, true, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], false);
+    assertCanUpdateDottedPath(
+        true, true, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], false);
+    assertCanUpdatePartialShardKey(true,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   false);
+
+    // upsert : true
+    assertCanUpdatePrimitiveShardKey(
+        true, true, [{"x": 300}, {"x": 4}], [{"x": 600}, {"x": 30}], true);
+    assertCanUpdateDottedPath(
+        true, true, [{"x.a": 300}, {"x.a": 4}], [{"x": {"a": 600}}, {"x": {"a": 30}}], true);
+    assertCanUpdatePartialShardKey(true,
+                                   true,
+                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                   [{"x": 600, "y": 80}, {"x": 30, "y": 3}],
+                                   true);
+
+    // failing cases
+    assertCannotUpdate_id(true, true, {"_id": 300}, {"_id": 600});
+    assertCannotUpdate_idDottedPath(true, true, {"_id.a": 300}, {"_id": {"a": 600}});
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        true, false, {"x": 300, "y": 80}, {"x": 600});
+    assertCannotUpdateSKToArray(true, true, {"x": 300}, {"x": [300]});
+
+    // ----Multiple writes in txn-----
+
+    // Bulk writes in txn
+    assertCanUpdateInBulkOp(true, false);
+    assertCanUpdateInBulkOp(true, true);
+
+    // Update two docs, updating one twice
+    docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+    shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    mongos.getDB(kDbName).foo.insert({"x": 505});
+
+    session.startTransaction();
+    let id = mongos.getDB(kDbName).foo.find({"x": 500}).toArray()[0]._id;
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$set": {"x": 400}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 400}, {"x": 600, "_id": id}));
+    assert.commandWorked(sessionDB.foo.update({"x": 4}, {"$set": {"x": 30}}));
+    session.commitTransaction();
+
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 400}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
+    assert.eq(id, mongos.getDB(kDbName).foo.find({"x": 600}).toArray()[0]._id);
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 4}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 30}).itcount());
+
+    mongos.getDB(kDbName).foo.drop();
+
+    // Check that doing $inc on doc A, then updating shard key for doc A, then $inc again only incs
+    // once
+    shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    mongos.getDB(kDbName).foo.insert({"x": 505});
+
+    session.startTransaction();
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$set": {"x": 400}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
+    session.commitTransaction();
+
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 400}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"a": 7}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 400, "a": 7}).itcount());
+
+    mongos.getDB(kDbName).foo.drop();
+
+    // Check that doing findAndModify to update shard key followed by $inc works correctly
+    shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    mongos.getDB(kDbName).foo.insert({"x": 505});
+
+    session.startTransaction();
+    sessionDB.foo.findAndModify({query: {"x": 500}, update: {$set: {"x": 600}}});
+    assert.commandWorked(sessionDB.foo.update({"x": 600}, {"$inc": {"a": 1}}));
+    session.commitTransaction();
+
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"a": 7}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 600, "a": 7}).itcount());
+
+    mongos.getDB(kDbName).foo.drop();
+
+    // Check that doing findAndModify followed by and update on a shard key works correctly
+    shardCollectionMoveChunks({"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+
+    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    mongos.getDB(kDbName).foo.insert({"x": 505});
+
+    id = mongos.getDB(kDbName).foo.find({"x": 4}).toArray()[0]._id;
+    session.startTransaction();
+    sessionDB.foo.findAndModify({query: {"x": 4}, update: {$set: {"x": 20}}});
+    assert.commandWorked(sessionDB.foo.update({"x": 20}, {$set: {"x": 1}}));
+    session.commitTransaction();
+
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 4}).itcount());
+    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 20}).itcount());
+    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 1}).itcount());
+    assert.eq(id, mongos.getDB(kDbName).foo.find({"x": 1}).toArray()[0]._id);
+
+    mongos.getDB(kDbName).foo.drop();
+
+    st.stop();
+
+})();
\ No newline at end of file
diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index d2a6213d496..54119affd1d 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -103,24 +103,24 @@ void addObjectIDIdField(mb::Document* doc) {
 }
 
 /**
- * Uasserts if any of the paths in 'immutablePaths' are not present in 'document', or if they are
+ * Uasserts if any of the paths in 'requiredPaths' are not present in 'document', or if they are
  * arrays or array descendants.
  */
-void checkImmutablePathsPresent(const mb::Document& document, const FieldRefSet& immutablePaths) {
-    for (auto path = immutablePaths.begin(); path != immutablePaths.end(); ++path) {
+void assertRequiredPathsPresent(const mb::Document& document, const FieldRefSet& requiredPaths) {
+    for (const auto& path : requiredPaths) {
         auto elem = document.root();
-        for (size_t i = 0; i < (*path)->numParts(); ++i) {
-            elem = elem[(*path)->getPart(i)];
+        for (size_t i = 0; i < (*path).numParts(); ++i) {
+            elem = elem[(*path).getPart(i)];
             uassert(ErrorCodes::NoSuchKey,
                     str::stream() << "After applying the update, the new document was missing the "
                                      "required field '"
-                                  << (*path)->dottedField()
+                                  << (*path).dottedField()
                                   << "'",
                     elem.ok());
             uassert(
                 ErrorCodes::NotSingleValueField,
                 str::stream() << "After applying the update to the document, the required field '"
-                              << (*path)->dottedField()
+                              << (*path).dottedField()
                               << "' was found to be an array or array descendant.",
                 elem.getType() != BSONType::Array);
         }
@@ -139,17 +139,6 @@ bool shouldRestartUpdateIfNoLongerMatches(const UpdateStageParams& params) {
     return params.request->shouldReturnAnyDocs() && !params.request->getSort().isEmpty();
 };
 
-const std::vector<std::unique_ptr<FieldRef>>* getImmutableFields(OperationContext* opCtx,
-                                                                 const NamespaceString& ns) {
-    auto metadata = CollectionShardingState::get(opCtx, ns)->getCurrentMetadata();
-    if (metadata->isSharded()) {
-        const std::vector<std::unique_ptr<FieldRef>>& fields = metadata->getKeyPatternFields();
-        // Return shard-keys as immutable for the update system.
-        return &fields;
-    }
-    return NULL;
-}
-
 CollectionUpdateArgs::StoreDocOption getStoreDocMode(const UpdateRequest& updateRequest) {
     if (updateRequest.shouldReturnNewDocs()) {
         return CollectionUpdateArgs::StoreDocOption::PostImage;
@@ -162,7 +151,6 @@ CollectionUpdateArgs::StoreDocOption getStoreDocMode(const UpdateRequest& update
     invariant(!updateRequest.shouldReturnAnyDocs());
     return CollectionUpdateArgs::StoreDocOption::None;
 }
-
 }  // namespace
 
 const char* UpdateStage::kStageType = "UPDATE";
@@ -219,15 +207,21 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
 
     bool docWasModified = false;
 
+    const auto isFCV42 = serverGlobalParams.featureCompatibility.isVersionInitialized() &&
+        serverGlobalParams.featureCompatibility.getVersion() ==
+            ServerGlobalParams::FeatureCompatibility::Version::kFullyUpgradedTo42;
+
+    auto* const css = CollectionShardingState::get(getOpCtx(), collection()->ns());
+    auto metadata = css->getCurrentMetadata();
     Status status = Status::OK();
     const bool validateForStorage = getOpCtx()->writesAreReplicated() && _enforceOkForStorage;
     const bool isInsert = false;
     FieldRefSet immutablePaths;
     if (getOpCtx()->writesAreReplicated() && !request->isFromMigration()) {
-        auto immutablePathsVector = getImmutableFields(getOpCtx(), request->getNamespaceString());
-        if (immutablePathsVector) {
+        if (!isFCV42 && metadata->isSharded()) {
+            auto& immutablePathsVector = metadata->getKeyPatternFields();
             immutablePaths.fillFrom(
-                transitional_tools_do_not_use::unspool_vector(*immutablePathsVector));
+                transitional_tools_do_not_use::unspool_vector(immutablePathsVector));
         }
         immutablePaths.keepShortest(&idFieldRef);
     }
@@ -301,8 +295,6 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
         RecordId newRecordId;
         CollectionUpdateArgs args;
 
-        auto* const css = CollectionShardingState::get(getOpCtx(), collection()->ns());
-        auto metadata = css->getCurrentMetadata();
         if (!request->isExplain()) {
             args.stmtId = request->getStmtId();
             args.update = logObj;
@@ -317,9 +309,6 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
             }
         }
 
-        const auto isFCV42 = serverGlobalParams.featureCompatibility.isVersionInitialized() &&
-            serverGlobalParams.featureCompatibility.getVersion() ==
-                ServerGlobalParams::FeatureCompatibility::Version::kFullyUpgradedTo42;
         if (inPlace) {
             if (!request->isExplain()) {
                 newObj = oldObj.value();
@@ -328,7 +317,7 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
                 Snapshotted<RecordData> snap(oldObj.snapshotId(), oldRec);
 
                 if (isFCV42 && metadata->isSharded()) {
-                    assertDocStillBelongsToNode(metadata, oldObj);
+                    assertUpdateToShardKeyFieldsIsValidAndDocStillBelongsToNode(metadata, oldObj);
                 }
 
                 WriteUnitOfWork wunit(getOpCtx());
@@ -351,7 +340,7 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
                     newObj.objsize() <= BSONObjMaxUserSize);
 
             if (isFCV42 && metadata->isSharded()) {
-                assertDocStillBelongsToNode(metadata, oldObj);
+                assertUpdateToShardKeyFieldsIsValidAndDocStillBelongsToNode(metadata, oldObj);
             }
 
             if (!request->isExplain()) {
@@ -368,7 +357,6 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
             }
         }
 
-
         // If the document moved, we might see it again in a collection scan (maybe it's
         // a document after our current document).
         //
@@ -406,18 +394,30 @@ BSONObj UpdateStage::applyUpdateOpsForInsert(OperationContext* opCtx,
     // Some mods may only work in that context (e.g. $setOnInsert).
     driver->setLogOp(false);
 
+    auto* const css = CollectionShardingState::get(opCtx, ns);
+    auto metadata = css->getCurrentMetadata();
+
+    const auto isFCV42 = serverGlobalParams.featureCompatibility.isVersionInitialized() &&
+        serverGlobalParams.featureCompatibility.getVersion() ==
+            ServerGlobalParams::FeatureCompatibility::Version::kFullyUpgradedTo42;
+
     FieldRefSet immutablePaths;
-    if (!isInternalRequest) {
-        auto immutablePathsVector = getImmutableFields(opCtx, ns);
-        if (immutablePathsVector) {
-            immutablePaths.fillFrom(
-                transitional_tools_do_not_use::unspool_vector(*immutablePathsVector));
-        }
+    if (!isFCV42 && metadata->isSharded()) {
+        auto& immutablePathsVector = metadata->getKeyPatternFields();
+        immutablePaths.fillFrom(
+            transitional_tools_do_not_use::unspool_vector(immutablePathsVector));
     }
     immutablePaths.keepShortest(&idFieldRef);
 
     if (cq) {
-        uassertStatusOK(driver->populateDocumentWithQueryFields(*cq, immutablePaths, *doc));
+        FieldRefSet requiredPaths;
+        if (metadata->isSharded()) {
+            const auto& shardKeyPathsVector = metadata->getKeyPatternFields();
+            requiredPaths.fillFrom(
+                transitional_tools_do_not_use::unspool_vector(shardKeyPathsVector));
+        }
+        requiredPaths.keepShortest(&idFieldRef);
+        uassertStatusOK(driver->populateDocumentWithQueryFields(*cq, requiredPaths, *doc));
         if (driver->isDocReplacement())
             stats->fastmodinsert = true;
     } else {
@@ -448,13 +448,20 @@ BSONObj UpdateStage::applyUpdateOpsForInsert(OperationContext* opCtx,
     }
 
     // Validate that the object replacement or modifiers resulted in a document
-    // that contains all the immutable keys and can be stored if it isn't coming
+    // that contains all the required keys and can be stored if it isn't coming
     // from a migration or via replication.
     if (!isInternalRequest) {
         if (enforceOkForStorage) {
             storage_validation::storageValid(*doc);
         }
-        checkImmutablePathsPresent(*doc, immutablePaths);
+        FieldRefSet requiredPaths;
+        if (metadata->isSharded()) {
+            const auto& shardKeyPathsVector = metadata->getKeyPatternFields();
+            requiredPaths.fillFrom(
+                transitional_tools_do_not_use::unspool_vector(shardKeyPathsVector));
+        }
+        requiredPaths.keepShortest(&idFieldRef);
+        assertRequiredPathsPresent(*doc, requiredPaths);
     }
 
     BSONObj newObj = doc->getObject();
@@ -883,17 +890,25 @@ PlanStage::StageState UpdateStage::prepareToRetryWSM(WorkingSetID idToRetry, Wor
     return NEED_YIELD;
 }
 
-void UpdateStage::assertDocStillBelongsToNode(ScopedCollectionMetadata metadata,
-                                              const Snapshotted<BSONObj>& oldObj) {
+void UpdateStage::assertUpdateToShardKeyFieldsIsValidAndDocStillBelongsToNode(
+    ScopedCollectionMetadata metadata, const Snapshotted<BSONObj>& oldObj) {
+    auto newObj = _doc.getObject();
     auto oldShardKey = metadata->extractDocumentKey(oldObj.value());
-    auto newShardKey = metadata->extractDocumentKey(_doc.getObject());
+    auto newShardKey = metadata->extractDocumentKey(newObj);
 
-    // If this document is an orphan and so does not belong to this shard or if the shard key fields
-    // remain unchanged by this update, we can skip the rest of the checks.
-    if (!metadata->keyBelongsToMe(oldShardKey) || (newShardKey.woCompare(oldShardKey) == 0)) {
+    // If the shard key fields remain unchanged by this update or if this document is an orphan and
+    // so does not belong to this shard, we can skip the rest of the checks.
+    if ((newShardKey.woCompare(oldShardKey) == 0) || !metadata->keyBelongsToMe(oldShardKey)) {
         return;
     }
 
+    // Assert that the updated doc has all shard key fields and none are arrays or array
+    // descendants.
+    FieldRefSet shardKeyPaths;
+    const auto& shardKeyPathsVector = metadata->getKeyPatternFields();
+    shardKeyPaths.fillFrom(transitional_tools_do_not_use::unspool_vector(shardKeyPathsVector));
+    assertRequiredPathsPresent(_doc, shardKeyPaths);
+
     // We do not allow updates to the shard key when 'multi' is true.
     uassert(ErrorCodes::InvalidOptions,
             str::stream() << "Multi-update operations are not allowed when updating the"
@@ -909,12 +924,9 @@ void UpdateStage::assertDocStillBelongsToNode(ScopedCollectionMetadata metadata,
             txnParticipant);
 
     if (!metadata->keyBelongsToMe(newShardKey)) {
-        // If this update is in a multi-stmt txn, attach the post image to the error. Otherwise,
-        // attach the original update field.
         boost::optional<BSONObj> originalQuery{txnParticipant.inMultiDocumentTransaction(),
                                                _params.request->getQuery()};
-        boost::optional<BSONObj> postImg{txnParticipant.inMultiDocumentTransaction(),
-                                         _doc.getObject()};
+        boost::optional<BSONObj> postImg{txnParticipant.inMultiDocumentTransaction(), newObj};
 
         uasserted(WouldChangeOwningShardInfo(originalQuery, postImg),
                   str::stream() << "This update would cause the doc to change owning shards");
diff --git a/src/mongo/db/exec/update_stage.h b/src/mongo/db/exec/update_stage.h
index 2b989812007..1a37554bee2 100644
--- a/src/mongo/db/exec/update_stage.h
+++ b/src/mongo/db/exec/update_stage.h
@@ -198,12 +198,15 @@ private:
     StageState prepareToRetryWSM(WorkingSetID idToRetry, WorkingSetID* out);
 
     /**
-     * Checks if the updated doc still belongs to this node and throws if it does not. This means
-     * that one or more shard key field values have been updated causing this doc to belong to
-     * a chunk that is not owned by this shard. We cannot apply this update atomically.
+     * Checks that the updated doc has all required shard key fields and throws if it does not.
+     *
+     * Also checks if the updated doc still belongs to this node and throws if it does not. If the
+     * doc no longer belongs to this shard, this means that one or more shard key field values have
+     * been updated to a value belonging to a chunk that is not owned by this shard. We cannot apply
+     * this update atomically.
      */
-    void assertDocStillBelongsToNode(ScopedCollectionMetadata metadata,
-                                     const Snapshotted<BSONObj>& oldObj);
+    void assertUpdateToShardKeyFieldsIsValidAndDocStillBelongsToNode(
+        ScopedCollectionMetadata metadata, const Snapshotted<BSONObj>& oldObj);
 
     UpdateStageParams _params;