1 files changed, 8 insertions, 7 deletions

diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index 6236e45a6a0..4da1d88b854 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -662,8 +662,8 @@ PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
                 BSONObj newObj = _specificStats.objInserted;
                 *out = _ws->allocate();
                 WorkingSetMember* member = _ws->get(*out);
-                member->obj = Snapshotted<BSONObj>(getOpCtx()->recoveryUnit()->getSnapshotId(),
-                                                   newObj.getOwned());
+                member->resetDocument(getOpCtx()->recoveryUnit()->getSnapshotId(),
+                                      newObj.getOwned());
                 member->transitionToOwnedObj();
                 return PlanStage::ADVANCED;
             }
@@ -759,14 +759,15 @@ PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
 
         // If we care about the pre-updated version of the doc, save it out here.
         BSONObj oldObj;
+        SnapshotId oldSnapshot = member->doc.snapshotId();
         if (_params.request->shouldReturnOldDocs()) {
-            oldObj = member->obj.value().getOwned();
+            oldObj = member->doc.value().toBson().getOwned();
         }
 
         BSONObj newObj;
         try {
             // Do the update, get us the new version of the doc.
-            newObj = transformAndUpdate(member->obj, recordId);
+            newObj = transformAndUpdate({oldSnapshot, member->doc.value().toBson()}, recordId);
         } catch (const WriteConflictException&) {
             memberFreer.dismiss();  // Keep this member around so we can retry updating it.
             return prepareToRetryWSM(id, out);
@@ -775,11 +776,11 @@ PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
         // Set member's obj to be the doc we want to return.
         if (_params.request->shouldReturnAnyDocs()) {
             if (_params.request->shouldReturnNewDocs()) {
-                member->obj = Snapshotted<BSONObj>(getOpCtx()->recoveryUnit()->getSnapshotId(),
-                                                   newObj.getOwned());
+                member->resetDocument(getOpCtx()->recoveryUnit()->getSnapshotId(),
+                                      newObj.getOwned());
             } else {
                 invariant(_params.request->shouldReturnOldDocs());
-                member->obj.setValue(oldObj);
+                member->resetDocument(oldSnapshot, oldObj);
             }
             member->recordId = RecordId();
             member->transitionToOwnedObj();