From 5629faccb40177a4102408567c085a75b96aee92 Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Wed, 26 Aug 2020 12:21:39 -0400
Subject: [PATCH] fix(model): dont wipe out changes made while `save()` is
 in-flight

Fix #9327
---
 lib/document.js | 39 +++++++++++++++++++++++++++++++++++++++
 lib/model.js    | 24 +++++++++++++++---------
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/lib/document.js b/lib/document.js
index 57060f507d5..7598bdea93f 100644
--- a/lib/document.js
+++ b/lib/document.js
@@ -38,6 +38,7 @@ const clone = utils.clone;
 const deepEqual = utils.deepEqual;
 const isMongooseObject = utils.isMongooseObject;
 
+const arrayAtomicsBackupSymbol = Symbol('mongoose.Array#atomicsBackup');
 const arrayAtomicsSymbol = require('./helpers/symbols').arrayAtomicsSymbol;
 const documentArrayParent = require('./helpers/symbols').documentArrayParent;
 const documentSchemaSymbol = require('./helpers/symbols').documentSchemaSymbol;
@@ -2764,6 +2765,7 @@ Document.prototype.$__reset = function reset() {
 
       _this.$__.activePaths.init(array.$path());
 
+      array[arrayAtomicsBackupSymbol] = array[arrayAtomicsSymbol];
       array[arrayAtomicsSymbol] = {};
     });
 
@@ -2784,10 +2786,19 @@ Document.prototype.$__reset = function reset() {
     const type = dirt.value;
 
     if (type && type[arrayAtomicsSymbol]) {
+      type[arrayAtomicsBackupSymbol] = type[arrayAtomicsSymbol];
       type[arrayAtomicsSymbol] = {};
     }
   });
 
+  this.$__.backup = {};
+  this.$__.backup.activePaths = {
+    modify: Object.assign({}, this.$__.activePaths.states.modify),
+    default: Object.assign({}, this.$__.activePaths.states.default)
+  };
+  this.$__.backup.validationError = this.$__.validationError;
+  this.$__.backup.errors = this.errors;
+
   // Clear 'dirty' cache
   this.$__.activePaths.clear('modify');
   this.$__.activePaths.clear('default');
@@ -2801,6 +2812,34 @@ Document.prototype.$__reset = function reset() {
   return this;
 };
 
+/*!
+ * ignore
+ */
+
+Document.prototype.$__undoReset = function $__undoReset() {
+  if (this.$__.backup == null || this.$__.backup.activePaths == null) {
+    return;
+  }
+
+  this.$__.activePaths.states.modify = this.$__.backup.activePaths.modify;
+  this.$__.activePaths.states.default = this.$__.backup.activePaths.default;
+
+  this.$__.validationError = this.$__.backup.validationError;
+  this.errors = this.$__.backup.errors;
+
+  for (const dirt of this.$__dirty()) {
+    const type = dirt.value;
+
+    if (type && type[arrayAtomicsSymbol] && type[arrayAtomicsBackupSymbol]) {
+      type[arrayAtomicsSymbol] = type[arrayAtomicsBackupSymbol];
+    }
+  }
+
+  for (const subdoc of this.$__getAllSubdocs()) {
+    subdoc.$__undoReset();
+  }
+};
+
 /**
  * Returns this documents dirty paths / vals.
  *
diff --git a/lib/model.js b/lib/model.js
index 0a1f01c0904..4ce3e39e648 100644
--- a/lib/model.js
+++ b/lib/model.js
@@ -282,6 +282,7 @@ Model.prototype.$__handleSave = function(options, callback) {
 
       callback(null, ret);
     });
+
     this.$__reset();
     _setIsNew(this, false);
     // Make it possible to retry the insert
@@ -307,8 +308,10 @@ Model.prototype.$__handleSave = function(options, callback) {
 
       _applyCustomWhere(this, where);
 
-      this[modelCollectionSymbol].updateOne(where, delta[1], saveOptions, function(err, ret) {
+      this[modelCollectionSymbol].updateOne(where, delta[1], saveOptions, (err, ret) => {
         if (err) {
+          this.$__undoReset();
+
           callback(err);
           return;
         }
@@ -319,15 +322,21 @@ Model.prototype.$__handleSave = function(options, callback) {
       const optionsWithCustomValues = Object.assign({}, options, saveOptions);
       this.constructor.exists(this.$__where(), optionsWithCustomValues)
         .then((documentExists) => {
-          if (!documentExists) throw new DocumentNotFoundError(this.$__where(), this.constructor.modelName);
+          if (!documentExists) {
+            throw new DocumentNotFoundError(this.$__where(), this.constructor.modelName);
+          }
 
-          this.$__reset();
           callback();
         })
         .catch(callback);
       return;
     }
 
+    // store the modified paths before the document is reset
+    this.$__.modifiedPaths = this.modifiedPaths();
+
+    this.$__reset();
+
     _setIsNew(this, false);
   }
 };
@@ -345,11 +354,6 @@ Model.prototype.$__save = function(options, callback) {
       });
     }
 
-    // store the modified paths before the document is reset
-    const modifiedPaths = this.modifiedPaths();
-
-    this.$__reset();
-
     let numAffected = 0;
     if (get(options, 'safe.w') !== 0 && get(options, 'w') !== 0) {
       // Skip checking if write succeeded if writeConcern is set to
@@ -376,8 +380,9 @@ Model.prototype.$__save = function(options, callback) {
 
         if (numAffected <= 0) {
           // the update failed. pass an error back
+          this.$__undoReset();
           const err = this.$__.$versionError ||
-            new VersionError(this, version, modifiedPaths);
+            new VersionError(this, version, this.$__.modifiedPaths);
           return callback(err);
         }
 
@@ -388,6 +393,7 @@ Model.prototype.$__save = function(options, callback) {
       }
 
       if (result != null && numAffected <= 0) {
+        this.$__undoReset();
         error = new DocumentNotFoundError(result.$where,
           this.constructor.modelName, numAffected, result);
         return hooks.execPost('save:error', this, [this], { error: error }, (error) => {