SERVER-52974 Avoid extra calls to _doc.getObject() in update codepath.

2 files changed, 39 insertions, 15 deletions

diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index f561d3d28b5..3d4eb8a5ac1 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -273,9 +273,7 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
                 Snapshotted<RecordData> snap(oldObj.snapshotId(), oldRec);
 
                 if (_isUserInitiatedWrite &&
-                    (checkUpdateChangesShardKeyFields(oldObj) ||
-                     wasReshardingKeyUpdated(_doc.getObject(), oldObj)) &&
-                    !args.preImageDoc) {
+                    checkUpdateChangesShardKeyFields(boost::none, oldObj) && !args.preImageDoc) {
                     args.preImageDoc = oldObj.value().getOwned();
                 }
 
@@ -299,9 +297,7 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
                     newObj.objsize() <= BSONObjMaxUserSize);
 
             if (!request->explain()) {
-                if (_isUserInitiatedWrite &&
-                    (checkUpdateChangesShardKeyFields(oldObj) ||
-                     wasReshardingKeyUpdated(_doc.getObject(), oldObj)) &&
+                if (_isUserInitiatedWrite && checkUpdateChangesShardKeyFields(newObj, oldObj) &&
                     !args.preImageDoc) {
                     args.preImageDoc = oldObj.value().getOwned();
                 }
@@ -573,12 +569,10 @@ PlanStage::StageState UpdateStage::prepareToRetryWSM(WorkingSetID idToRetry, Wor
     return NEED_YIELD;
 }
 
-bool UpdateStage::wasReshardingKeyUpdated(const BSONObj& newObj,
+bool UpdateStage::wasReshardingKeyUpdated(const ScopedCollectionDescription& collDesc,
+                                          const BSONObj& newObj,
                                           const Snapshotted<BSONObj>& oldObj) {
-    auto css = CollectionShardingState::get(opCtx(), collection()->ns());
-
-    auto reshardingKeyPattern =
-        css->getCollectionDescription(opCtx()).getReshardingKeyIfShouldForwardOps();
+    auto reshardingKeyPattern = collDesc.getReshardingKeyIfShouldForwardOps();
     if (!reshardingKeyPattern)
         return false;
 
@@ -606,14 +600,27 @@ bool UpdateStage::wasReshardingKeyUpdated(const BSONObj& newObj,
     return true;
 }
 
-bool UpdateStage::checkUpdateChangesShardKeyFields(const Snapshotted<BSONObj>& oldObj) {
+bool UpdateStage::checkUpdateChangesShardKeyFields(const boost::optional<BSONObj>& newObjCopy,
+                                                   const Snapshotted<BSONObj>& oldObj) {
     auto* const css = CollectionShardingState::get(opCtx(), collection()->ns());
     const auto collDesc = css->getCollectionDescription(opCtx());
+
+    // Calling mutablebson::Document::getObject() renders a full copy of the updated document. This
+    // can be expensive for larger documents, so we skip calling it when the collection isn't even
+    // sharded.
     if (!collDesc.isSharded()) {
         return false;
     }
 
-    auto newObj = _doc.getObject();
+    const auto& newObj = newObjCopy ? *newObjCopy : _doc.getObject();
+    return wasExistingShardKeyUpdated(css, collDesc, newObj, oldObj) ||
+        wasReshardingKeyUpdated(collDesc, newObj, oldObj);
+}
+
+bool UpdateStage::wasExistingShardKeyUpdated(CollectionShardingState* css,
+                                             const ScopedCollectionDescription& collDesc,
+                                             const BSONObj& newObj,
+                                             const Snapshotted<BSONObj>& oldObj) {
     const ShardKeyPattern shardKeyPattern(collDesc.getKeyPattern());
     auto oldShardKey = shardKeyPattern.extractShardKeyFromDoc(oldObj.value());
     auto newShardKey = shardKeyPattern.extractShardKeyFromDoc(newObj);
diff --git a/src/mongo/db/exec/update_stage.h b/src/mongo/db/exec/update_stage.h
index 90ac0cf0625..80e7162e515 100644
--- a/src/mongo/db/exec/update_stage.h
+++ b/src/mongo/db/exec/update_stage.h
@@ -145,6 +145,17 @@ private:
     StageState prepareToRetryWSM(WorkingSetID idToRetry, WorkingSetID* out);
 
     /**
+     * Returns true if the owning shard under the current key pattern would change as a result of
+     * the update, or if the destined recipient under the new shard key pattern from resharding
+     * would change as a result of the update, and returns false otherwise.
+     *
+     * Accepting a 'newObjCopy' parameter is a performance enhancement for updates which weren't
+     * performed in-place to avoid rendering a full copy of the updated document multiple times.
+     */
+    bool checkUpdateChangesShardKeyFields(const boost::optional<BSONObj>& newObjCopy,
+                                          const Snapshotted<BSONObj>& oldObj);
+
+    /**
      * 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
@@ -155,8 +166,14 @@ private:
      * If the update changes shard key fields but the new shard key remains on the same node,
      * returns true. If the update does not change shard key fields, returns false.
      */
-    bool checkUpdateChangesShardKeyFields(const Snapshotted<BSONObj>& oldObj);
-    bool wasReshardingKeyUpdated(const BSONObj& newObj, const Snapshotted<BSONObj>& oldObj);
+    bool wasExistingShardKeyUpdated(CollectionShardingState* css,
+                                    const ScopedCollectionDescription& collDesc,
+                                    const BSONObj& newObj,
+                                    const Snapshotted<BSONObj>& oldObj);
+
+    bool wasReshardingKeyUpdated(const ScopedCollectionDescription& collDesc,
+                                 const BSONObj& newObj,
+                                 const Snapshotted<BSONObj>& oldObj);
 
     // If not WorkingSet::INVALID_ID, we use this rather than asking our child what to do next.
     WorkingSetID _idRetrying;