3 files changed, 45 insertions, 51 deletions

diff --git a/src/mongo/db/s/chunk_split_state_driver.cpp b/src/mongo/db/s/chunk_split_state_driver.cpp
index 24009a312b3..dd80060e642 100644
--- a/src/mongo/db/s/chunk_split_state_driver.cpp
+++ b/src/mongo/db/s/chunk_split_state_driver.cpp
@@ -47,14 +47,21 @@ boost::optional<ChunkSplitStateDriver> ChunkSplitStateDriver::tryInitiateSplit(
 ChunkSplitStateDriver::ChunkSplitStateDriver(ChunkSplitStateDriver&& source) {
     _writesTracker = source._writesTracker;
     source._writesTracker.reset();
+
+    _stashedBytesWritten = source._stashedBytesWritten;
+    source._stashedBytesWritten = 0;
+
     _splitState = source._splitState;
-    // Ensure that the source driver does not cancel the ongoing split
     source._splitState = SplitState::kNotSplitting;
 }
 
 ChunkSplitStateDriver::~ChunkSplitStateDriver() {
-    if (_splitState == SplitState::kSplitInProgress || _splitState == SplitState::kSplitPrepared) {
-        cancelSplit();
+    if (_splitState != SplitState::kSplitCommitted) {
+        auto wt = _writesTracker.lock();
+        if (wt) {
+            wt->releaseSplitLock();
+            wt->addBytesWritten(_stashedBytesWritten);
+        }
     }
 }
 
@@ -62,8 +69,10 @@ void ChunkSplitStateDriver::prepareSplit() {
     invariant(_splitState == SplitState::kSplitInProgress);
     _splitState = SplitState::kSplitPrepared;
 
+    auto wt = _writesTracker.lock();
+    uassert(50873, "Split interrupted due to chunk metadata change.", wt);
     // Clear bytes written and get the previous bytes written.
-    _stashedBytesWritten = _getWritesTracker()->clearBytesWritten();
+    _stashedBytesWritten = wt->clearBytesWritten();
 }
 
 void ChunkSplitStateDriver::commitSplit() {
@@ -71,19 +80,4 @@ void ChunkSplitStateDriver::commitSplit() {
     _splitState = SplitState::kSplitCommitted;
 }
 
-void ChunkSplitStateDriver::cancelSplit() {
-    invariant(_splitState == SplitState::kSplitInProgress ||
-              _splitState == SplitState::kSplitPrepared);
-    _splitState = SplitState::kNotSplitting;
-
-    _getWritesTracker()->releaseSplitLock();
-    _getWritesTracker()->addBytesWritten(_stashedBytesWritten);
-}
-
-std::shared_ptr<ChunkWritesTracker> ChunkSplitStateDriver::_getWritesTracker() {
-    auto wt = _writesTracker.lock();
-    invariant(wt, "ChunkWritesTracker was destructed before ChunkSplitStateDriver");
-    return wt;
-}
-
 }  // namespace mongo
diff --git a/src/mongo/db/s/chunk_split_state_driver.h b/src/mongo/db/s/chunk_split_state_driver.h
index 427f3cda056..b1954a7f9c9 100644
--- a/src/mongo/db/s/chunk_split_state_driver.h
+++ b/src/mongo/db/s/chunk_split_state_driver.h
@@ -86,13 +86,6 @@ public:
      */
     void commitSplit();
 
-    /**
-     * Marks the split state of an in-progress split back to kNotSplitting.
-     * If a split has already been prepared, resets the byte counter to what it
-     * was prior to prepare plus any new bytes that have been written.
-     */
-    void cancelSplit();
-
 private:
     /**
      * Should only be used by tryInitiateSplit
@@ -119,13 +112,6 @@ private:
         kSplitPrepared,
         kSplitCommitted,
     } _splitState{SplitState::kNotSplitting};
-
-
-    /**
-     * Returns the ChunkWritesTracker whose state this driver is controlling,
-     * checking to make sure it has not yet been destroyed.
-     */
-    std::shared_ptr<ChunkWritesTracker> _getWritesTracker();
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/s/chunk_split_state_driver_test.cpp b/src/mongo/db/s/chunk_split_state_driver_test.cpp
index 1619d035f81..6de8ae1ecaa 100644
--- a/src/mongo/db/s/chunk_split_state_driver_test.cpp
+++ b/src/mongo/db/s/chunk_split_state_driver_test.cpp
@@ -36,7 +36,7 @@
 
 namespace mongo {
 
-class ChunkSplitStateDriverTest : public unittest::Test {
+class ChunkSplitStateDriverTestNoTeardown : public unittest::Test {
 public:
     // Add bytes to write tracker and create the ChunkSplitStateDriver object
     // to test, which starts the split on the writes tracker
@@ -48,24 +48,29 @@ public:
             ChunkSplitStateDriver::tryInitiateSplit(_writesTracker));
     }
 
-    void tearDown() override {
-        _splitDriver.reset();
-        _writesTracker.reset();
-    }
+    void tearDown() override {}
 
-    ChunkWritesTracker& writesTracker() {
+    virtual ChunkWritesTracker& writesTracker() {
         return *_writesTracker;
     }
 
-    boost::optional<ChunkSplitStateDriver>& splitDriver() {
+    virtual boost::optional<ChunkSplitStateDriver>& splitDriver() {
         return *_splitDriver;
     }
 
-private:
+protected:
     std::shared_ptr<ChunkWritesTracker> _writesTracker;
     std::unique_ptr<boost::optional<ChunkSplitStateDriver>> _splitDriver;
 };
 
+class ChunkSplitStateDriverTest : public ChunkSplitStateDriverTestNoTeardown {
+public:
+    void tearDown() override {
+        _splitDriver.reset();
+        _writesTracker.reset();
+    }
+};
+
 TEST(ChunkSplitStateDriverTest, InitiateSplitLeavesBytesWrittenUnchanged) {
     auto writesTracker = std::make_shared<ChunkWritesTracker>();
     uint64_t bytesInTrackerBeforeSplit{4};
@@ -82,20 +87,21 @@ TEST_F(ChunkSplitStateDriverTest, PrepareSplitClearsBytesWritten) {
     ASSERT_EQ(writesTracker().getBytesWritten(), 0ull);
 }
 
-TEST_F(ChunkSplitStateDriverTest, PrepareSplitFollowedByCancelSplitRestoresBytesWritten) {
+TEST_F(ChunkSplitStateDriverTestNoTeardown,
+       PrepareSplitFollowedByDestructorWithoutCommitRestoresBytesWritten) {
     auto bytesInTracker = writesTracker().getBytesWritten();
     splitDriver()->prepareSplit();
-    splitDriver()->cancelSplit();
+    splitDriver().reset();
     ASSERT_EQ(writesTracker().getBytesWritten(), bytesInTracker);
 }
 
-TEST_F(ChunkSplitStateDriverTest,
-       PrepareSplitThenAddBytesThenCancelSplitRestoresOldBytesWrittenPlusNewBytesWritten) {
+TEST_F(ChunkSplitStateDriverTestNoTeardown,
+       PrepareSplitFollowedByDestructorWithoutCommitRestoresOldBytesWrittenPlusNewBytesWritten) {
     auto bytesInTracker = writesTracker().getBytesWritten();
     splitDriver()->prepareSplit();
     uint64_t extraBytesToAdd{4};
     writesTracker().addBytesWritten(extraBytesToAdd);
-    splitDriver()->cancelSplit();
+    splitDriver().reset();
 
     ASSERT_EQ(writesTracker().getBytesWritten(), bytesInTracker + extraBytesToAdd);
 }
@@ -125,9 +131,10 @@ TEST_F(ChunkSplitStateDriverTest, ShouldSplitReturnsFalseEvenAfterCommit) {
     ASSERT_FALSE(writesTracker().shouldSplit(maxChunkSize));
 }
 
-TEST_F(ChunkSplitStateDriverTest, ShouldSplitReturnsTrueAgainAfterCancel) {
+TEST_F(ChunkSplitStateDriverTestNoTeardown,
+       ShouldSplitReturnsTrueAfterPrepareSplitThenDestruction) {
     splitDriver()->prepareSplit();
-    splitDriver()->cancelSplit();
+    splitDriver().reset();
 
     uint64_t maxChunkSize{0};
     ASSERT_TRUE(writesTracker().shouldSplit(maxChunkSize));
@@ -139,9 +146,16 @@ DEATH_TEST_F(ChunkSplitStateDriverTest,
     splitDriver()->commitSplit();
 }
 
-DEATH_TEST_F(ChunkSplitStateDriverTest, CancelSplitAfterCommitErrors, "Invariant failure") {
-    splitDriver()->commitSplit();
-    splitDriver()->cancelSplit();
+TEST(ChunkSplitStateDriverTest, PrepareErrorsWhenChunkWritesTrackerNoLongerExists) {
+    std::unique_ptr<boost::optional<ChunkSplitStateDriver>> splitDriver;
+    {
+        auto writesTracker = std::make_shared<ChunkWritesTracker>();
+        uint64_t bytesToAdd{4};
+        writesTracker->addBytesWritten(bytesToAdd);
+        splitDriver = std::make_unique<boost::optional<ChunkSplitStateDriver>>(
+            ChunkSplitStateDriver::tryInitiateSplit(writesTracker));
+    }
+    ASSERT_THROWS(splitDriver->get().prepareSplit(), AssertionException);
 }
 
 }  // namespace mongo