SERVER-74303 Fix WiredTigerRecordStore not marking sizeInfo dirty on rollback

(cherry picked from commit 2d07f31853fac1e9496be12523c8e7a72dccd81f)

2 files changed, 51 insertions, 8 deletions

diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp
index e80ba142aec..e0f3938539a 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp
@@ -2006,20 +2006,23 @@ void WiredTigerRecordStore::_changeNumRecordsAndDataSize(OperationContext* opCtx
         return;
     }
 
-    opCtx->recoveryUnit()->onRollback([this, numRecordDiff, dataSizeDiff]() {
+    const auto updateAndStoreSizeInfo = [this](int64_t numRecordDiff, int64_t dataSizeDiff) {
+        _sizeInfo->numRecords.addAndFetch(numRecordDiff);
+        _sizeInfo->dataSize.addAndFetch(dataSizeDiff);
+
+        if (_sizeStorer)
+            _sizeStorer->store(_uri, _sizeInfo);
+    };
+
+    opCtx->recoveryUnit()->onRollback([updateAndStoreSizeInfo, numRecordDiff, dataSizeDiff]() {
         LOGV2_DEBUG(7105300,
                     3,
                     "WiredTigerRecordStore: rolling back change to numRecords and dataSize",
                     "numRecordDiff"_attr = -numRecordDiff,
                     "dataSizeDiff"_attr = -dataSizeDiff);
-        _sizeInfo->numRecords.addAndFetch(-numRecordDiff);
-        _sizeInfo->dataSize.addAndFetch(-dataSizeDiff);
+        updateAndStoreSizeInfo(-numRecordDiff, -dataSizeDiff);
     });
-    _sizeInfo->numRecords.addAndFetch(numRecordDiff);
-    _sizeInfo->dataSize.addAndFetch(dataSizeDiff);
-
-    if (_sizeStorer)
-        _sizeStorer->store(_uri, _sizeInfo);
+    updateAndStoreSizeInfo(numRecordDiff, dataSizeDiff);
 }
 
 void WiredTigerRecordStore::setNumRecords(long long numRecords) {
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_standard_record_store_test.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_standard_record_store_test.cpp
index 16704935431..86c6abf08fa 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_standard_record_store_test.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_standard_record_store_test.cpp
@@ -247,5 +247,45 @@ TEST_F(SizeStorerUpdateTest, DataSizeModification) {
     }
 }
 
+// Verify that the size storer contains accurate data after a transaction rollback just before a
+// flush (simulating a shutdown). That is, that the rollback marks the size info as dirty, and is
+// properly flushed to disk.
+TEST_F(SizeStorerUpdateTest, ReloadAfterRollbackAndFlush) {
+    ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+    // Do an op for which the sizeInfo is persisted, for safety so we don't check against 0.
+    {
+        WriteUnitOfWork uow(opCtx.get());
+        auto rId = rs->insertRecord(opCtx.get(), "12345", 5, Timestamp{1});
+        ASSERT_TRUE(rId.isOK());
+
+        uow.commit();
+    }
+
+    // An operation to rollback, with a flush between the original modification and the rollback.
+    {
+        WriteUnitOfWork uow(opCtx.get());
+        auto rId = rs->insertRecord(opCtx.get(), "12345", 5, Timestamp{2});
+        ASSERT_TRUE(rId.isOK());
+
+        ASSERT_EQ(getNumRecords(opCtx.get()), 2);
+        ASSERT_EQ(getDataSize(opCtx.get()), 10);
+        // Mark size info as clean, before rollback is done.
+        sizeStorer->flush(false);
+    }
+
+    // Simulate a shutdown and restart, which loads the size storer from disk.
+    sizeStorer->flush(true);
+    sizeStorer.reset(new WiredTigerSizeStorer(harnessHelper->conn(),
+                                              WiredTigerKVEngine::kTableUriPrefix + "sizeStorer"));
+    WiredTigerRecordStore* wtrs = checked_cast<WiredTigerRecordStore*>(rs.get());
+    wtrs->setSizeStorer(sizeStorer.get());
+
+    // As the operation was rolled back, numRecords and dataSize should be for the first op only. If
+    // rollback does not properly mark the sizeInfo as dirty, on load sizeInfo will account for the
+    // two operations, as the rollback sizeInfo update has not been flushed.
+    ASSERT_EQ(getNumRecords(opCtx.get()), 1);
+    ASSERT_EQ(getDataSize(opCtx.get()), 5);
+};
+
 }  // namespace
 }  // namespace mongo