From 93d33c65420c0bf1f1ca1f82a4cd4bc2cf312e77 Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Sun, 14 Jul 2024 16:57:47 -0400
Subject: [PATCH] refactor: use $createModifiedPathsSnapshot for transaction
 wrapper

Re: #14699
---
 lib/connection.js               | 27 +--------------------------
 lib/document.js                 |  4 +++-
 lib/modifiedPathsSnapshot.js    |  3 ++-
 lib/plugins/trackTransaction.js | 17 ++---------------
 test/document.test.js           |  4 ++++
 5 files changed, 12 insertions(+), 43 deletions(-)

diff --git a/lib/connection.js b/lib/connection.js
index 6e52d6ca4a0..0a746ac5074 100644
--- a/lib/connection.js
+++ b/lib/connection.js
@@ -574,32 +574,7 @@ async function _wrapUserTransaction(fn, session, mongoose) {
 function _resetSessionDocuments(session) {
   for (const doc of session[sessionNewDocuments].keys()) {
     const state = session[sessionNewDocuments].get(doc);
-    if (state.hasOwnProperty('isNew')) {
-      doc.$isNew = state.isNew;
-    }
-    if (state.hasOwnProperty('versionKey')) {
-      doc.set(doc.schema.options.versionKey, state.versionKey);
-    }
-
-    if (state.modifiedPaths.length > 0 && doc.$__.activePaths.states.modify == null) {
-      doc.$__.activePaths.states.modify = {};
-    }
-    for (const path of state.modifiedPaths) {
-      const currentState = doc.$__.activePaths.paths[path];
-      if (currentState != null) {
-        delete doc.$__.activePaths[currentState][path];
-      }
-      doc.$__.activePaths.paths[path] = 'modify';
-      doc.$__.activePaths.states.modify[path] = true;
-    }
-
-    for (const path of state.atomics.keys()) {
-      const val = doc.$__getValue(path);
-      if (val == null) {
-        continue;
-      }
-      val[arrayAtomicsSymbol] = state.atomics.get(path);
-    }
+    doc.$restoreModifiedPathsSnapshot(state);
   }
 }
 
diff --git a/lib/document.js b/lib/document.js
index c4afeb763a7..d0e42deb86d 100644
--- a/lib/document.js
+++ b/lib/document.js
@@ -5249,7 +5249,8 @@ Document.prototype.$createModifiedPathsSnapshot = function $createModifiedPathsS
   return new ModifiedPathsSnapshot(
     subdocSnapshot,
     this.$__.activePaths.clone(),
-    this.$__.version
+    this.$__.version,
+    this.$isNew
   );
 };
 
@@ -5281,6 +5282,7 @@ Document.prototype.$createModifiedPathsSnapshot = function $createModifiedPathsS
 Document.prototype.$restoreModifiedPathsSnapshot = function $restoreModifiedPathsSnapshot(snapshot) {
   this.$__.activePaths = snapshot.activePaths.clone();
   this.$__.version = snapshot.version;
+  this.$isNew = snapshot.isNew;
   if (!this.$isSubdocument) {
     const subdocs = this.$getAllSubdocs();
     for (const child of subdocs) {
diff --git a/lib/modifiedPathsSnapshot.js b/lib/modifiedPathsSnapshot.js
index 54d6b30d70b..54fd849383f 100644
--- a/lib/modifiedPathsSnapshot.js
+++ b/lib/modifiedPathsSnapshot.js
@@ -1,9 +1,10 @@
 'use strict';
 
 module.exports = class ModifiedPathsSnapshot {
-  constructor(subdocSnapshot, activePaths, version) {
+  constructor(subdocSnapshot, activePaths, version, isNew) {
     this.subdocSnapshot = subdocSnapshot;
     this.activePaths = activePaths;
     this.version = version;
+    this.isNew = isNew;
   }
 };
diff --git a/lib/plugins/trackTransaction.js b/lib/plugins/trackTransaction.js
index af5c7c84da4..24c3aec2ddf 100644
--- a/lib/plugins/trackTransaction.js
+++ b/lib/plugins/trackTransaction.js
@@ -15,24 +15,11 @@ module.exports = function trackTransaction(schema) {
     }
 
     if (!session[sessionNewDocuments].has(this)) {
-      const initialState = {};
-      if (this.isNew) {
-        initialState.isNew = true;
-      }
-      if (this.$__schema.options.versionKey) {
-        initialState.versionKey = this.get(this.$__schema.options.versionKey);
-      }
-
-      initialState.modifiedPaths = new Set(Object.keys(this.$__.activePaths.getStatePaths('modify')));
-      initialState.atomics = _getAtomics(this);
-
-      session[sessionNewDocuments].set(this, initialState);
+      session[sessionNewDocuments].set(this, this.$createModifiedPathsSnapshot());
     } else {
       const state = session[sessionNewDocuments].get(this);
 
-      for (const path of Object.keys(this.$__.activePaths.getStatePaths('modify'))) {
-        state.modifiedPaths.add(path);
-      }
+      state.activePaths.states.modify = { ...this.$__.activePaths.getStatePaths('modify') };
       state.atomics = _getAtomics(this, state.atomics);
     }
   });
diff --git a/test/document.test.js b/test/document.test.js
index 0e7eae43e52..3a95b5a71e9 100644
--- a/test/document.test.js
+++ b/test/document.test.js
@@ -13631,6 +13631,7 @@ describe('document', function() {
       docArr: [new Schema({ subprop3: String }, { _id: false })]
     });
     const Test = db.model('Test', schema);
+    await Test.deleteMany({});
 
     const doc = new Test({});
     await doc.save();
@@ -13651,14 +13652,17 @@ describe('document', function() {
     assert.deepStrictEqual(doc.subdoc.getChanges(), { $set: { subprop2: 'test3' } });
     assert.deepStrictEqual(doc.docArr[0].getChanges(), { $set: { subprop3: 'test4' } });
 
+    assert.ok(!doc.$isNew);
     const snapshot = doc.$createModifiedPathsSnapshot();
     doc.$clearModifiedPaths();
+    doc.$isNew = true;
 
     assert.deepStrictEqual(doc.getChanges(), {});
     assert.deepStrictEqual(doc.subdoc.getChanges(), {});
     assert.deepStrictEqual(doc.docArr[0].getChanges(), {});
 
     doc.$restoreModifiedPathsSnapshot(snapshot);
+    assert.ok(!doc.$isNew);
     assert.deepStrictEqual(doc.getChanges().$set, {
       name: 'test1',
       nested: { subprop1: 'test2' },