SERVER-29531 Handle rollbacks in SessionTransactionTable

10 files changed, 172 insertions, 0 deletions

diff --git a/src/mongo/db/repl/SConscript b/src/mongo/db/repl/SConscript
index 043aa8599a4..6e4fb1d466b 100644
--- a/src/mongo/db/repl/SConscript
+++ b/src/mongo/db/repl/SConscript
@@ -411,6 +411,7 @@ env.Library(
         'roll_back_local_operations',
         '$BUILD_DIR/mongo/db/catalog/database_holder',
         '$BUILD_DIR/mongo/db/s/sharding',
+        '$BUILD_DIR/mongo/db/sessions',
         '$BUILD_DIR/mongo/util/fail_point',
         '$BUILD_DIR/mongo/db/dbhelpers',
     ],
@@ -453,6 +454,7 @@ env.Library(
         'storage_interface_impl',
         '$BUILD_DIR/mongo/db/service_context',
         '$BUILD_DIR/mongo/db/service_context_d_test_fixture',
+        '$BUILD_DIR/mongo/db/sessions',
         '$BUILD_DIR/mongo/executor/thread_pool_task_executor_test_fixture',
     ],
 )
@@ -497,6 +499,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/concurrency/lock_manager',
         '$BUILD_DIR/mongo/db/s/sharding',
         '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/db/sessions',
         '$BUILD_DIR/mongo/executor/task_executor_interface',
         '$BUILD_DIR/mongo/util/net/network',
     ],
diff --git a/src/mongo/db/repl/rollback_impl.cpp b/src/mongo/db/repl/rollback_impl.cpp
index 84d03035654..b888558a833 100644
--- a/src/mongo/db/repl/rollback_impl.cpp
+++ b/src/mongo/db/repl/rollback_impl.cpp
@@ -40,6 +40,7 @@
 #include "mongo/db/repl/replication_coordinator.h"
 #include "mongo/db/repl/rollback_impl_listener.h"
 #include "mongo/db/s/shard_identity_rollback_notifier.h"
+#include "mongo/db/session_catalog.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/log.h"
@@ -182,6 +183,11 @@ void RollbackImpl::_checkShardIdentityRollback(OperationContext* opCtx) {
     }
 }
 
+void RollbackImpl::_clearSessionTransactionTable(OperationContext* opCtx) {
+    invariant(opCtx);
+    SessionCatalog::get(opCtx)->clearTransactionTable();
+}
+
 void RollbackImpl::_transitionFromRollbackToSecondary(OperationContext* opCtx) {
     invariant(opCtx);
 
@@ -208,6 +214,7 @@ void RollbackImpl::_tearDown(OperationContext* opCtx) {
     invariant(opCtx);
 
     _checkShardIdentityRollback(opCtx);
+    _clearSessionTransactionTable(opCtx);
     _transitionFromRollbackToSecondary(opCtx);
 }
 
diff --git a/src/mongo/db/repl/rollback_impl.h b/src/mongo/db/repl/rollback_impl.h
index cbddbca41e7..1142ad73f6a 100644
--- a/src/mongo/db/repl/rollback_impl.h
+++ b/src/mongo/db/repl/rollback_impl.h
@@ -193,6 +193,16 @@ private:
     void _checkShardIdentityRollback(OperationContext* opCtx);
 
     /**
+     * The in-memory session transaction table needs to be cleared after rollback, so it is forced
+     * to refetch from storage.
+     *
+     * 'opCtx' cannot be null.
+     *
+     * Called by _tearDown().
+     */
+    void _clearSessionTransactionTable(OperationContext* opCtx);
+
+    /**
      * Transitions the current member state from ROLLBACK to SECONDARY.
      * This operation must succeed. Otherwise, we will shut down the server.
      *
diff --git a/src/mongo/db/repl/rollback_test_fixture.cpp b/src/mongo/db/repl/rollback_test_fixture.cpp
index 0eae16c83b9..c3097819656 100644
--- a/src/mongo/db/repl/rollback_test_fixture.cpp
+++ b/src/mongo/db/repl/rollback_test_fixture.cpp
@@ -38,6 +38,7 @@
 #include "mongo/db/repl/replication_coordinator.h"
 #include "mongo/db/repl/replication_coordinator_mock.h"
 #include "mongo/db/repl/replication_process.h"
+#include "mongo/db/session_catalog.h"
 #include "mongo/stdx/memory.h"
 #include "mongo/util/mongoutils/str.h"
 
@@ -83,6 +84,8 @@ void RollbackTest::setUp() {
                                 std::unique_ptr<ReplicationCoordinator>(_coordinator));
     setOplogCollectionName();
 
+    SessionCatalog::create(serviceContext);
+
     _opCtx = cc().makeOperationContext();
     _replicationProcess->getConsistencyMarkers()->setAppliedThrough(_opCtx.get(), OpTime{});
     _replicationProcess->getConsistencyMarkers()->setMinValid(_opCtx.get(), OpTime{});
@@ -95,6 +98,8 @@ void RollbackTest::tearDown() {
     _coordinator = nullptr;
     _opCtx.reset();
 
+    SessionCatalog::reset_forTest(_serviceContextMongoDTest.getServiceContext());
+
     // We cannot unset the global replication coordinator because ServiceContextMongoD::tearDown()
     // calls dropAllDatabasesExceptLocal() which requires the replication coordinator to clear all
     // snapshots.
diff --git a/src/mongo/db/repl/rs_rollback.cpp b/src/mongo/db/repl/rs_rollback.cpp
index 8054c31c5e3..62c5a190e44 100644
--- a/src/mongo/db/repl/rs_rollback.cpp
+++ b/src/mongo/db/repl/rs_rollback.cpp
@@ -50,6 +50,7 @@
 #include "mongo/db/db_raii.h"
 #include "mongo/db/dbhelpers.h"
 #include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/logical_session_id.h"
 #include "mongo/db/ops/delete.h"
 #include "mongo/db/ops/update.h"
 #include "mongo/db/ops/update_lifecycle_impl.h"
@@ -66,6 +67,7 @@
 #include "mongo/db/repl/rollback_source.h"
 #include "mongo/db/repl/rslog.h"
 #include "mongo/db/s/shard_identity_rollback_notifier.h"
+#include "mongo/db/session_catalog.h"
 #include "mongo/util/exit.h"
 #include "mongo/util/fail_point_service.h"
 #include "mongo/util/log.h"
@@ -160,6 +162,29 @@ Status rollback_internal::updateFixUpInfoFromLocalOplogEntry(FixUpInfo& fixUpInf
                                              << redact(oplogEntry.toBSON()));
     }
 
+    // If the operation being rolled back has a txnNumber, then the corresponding entry in the
+    // session transaction table needs to be refetched.
+    auto operationSessionInfo = oplogEntry.getOperationSessionInfo();
+    auto txnNumber = operationSessionInfo.getTxnNumber();
+    if (txnNumber) {
+        auto sessionId = operationSessionInfo.getSessionId();
+        invariant(sessionId);
+        invariant(oplogEntry.getStatementId());
+
+        DocID txnDoc;
+        BSONObjBuilder txnBob;
+        txnBob.append("_id", sessionId->toBSON());
+        txnDoc.ownedObj = txnBob.obj();
+        txnDoc._id = txnDoc.ownedObj.firstElement();
+        // TODO: SERVER-29667
+        // Once collection uuids replace namespace strings for rollback, this will need to be
+        // changed to the uuid of the session transaction table collection.
+        txnDoc.ns = NamespaceString::kSessionTransactionsTableNamespace.ns().c_str();
+
+        fixUpInfo.docsToRefetch.insert(txnDoc);
+        fixUpInfo.refetchTransactionDocs = true;
+    }
+
     if (oplogEntry.getOpType() == OpTypeEnum::kCommand) {
 
         // The first element of the object is the name of the command
@@ -858,6 +883,11 @@ void syncFixUp(OperationContext* opCtx,
         fassertFailedNoTrace(40496);
     }
 
+    // If necessary, clear the in-memory session transaction table.
+    if (fixUpInfo.refetchTransactionDocs) {
+        SessionCatalog::get(opCtx)->clearTransactionTable();
+    }
+
     // Reload the lastAppliedOpTime and lastDurableOpTime value in the replcoord and the
     // lastAppliedHash value in bgsync to reflect our new last op.
     replCoord->resetLastOpTimesFromOplog(opCtx);
diff --git a/src/mongo/db/repl/rs_rollback.h b/src/mongo/db/repl/rs_rollback.h
index 7fa272dc22e..f2583919e5e 100644
--- a/src/mongo/db/repl/rs_rollback.h
+++ b/src/mongo/db/repl/rs_rollback.h
@@ -250,6 +250,10 @@ struct FixUpInfo {
     stdx::unordered_map<UUID, std::pair<OpTime, std::string>, UUID::Hash>
         collectionsToRollBackPendingDrop;
 
+    // True if rollback requires re-fetching documents in the session transaction table. If true,
+    // after rollback the in-memory transaction table is cleared.
+    bool refetchTransactionDocs = false;
+
     OpTime commonPoint;
     RecordId commonPointOurDiskloc;
 
diff --git a/src/mongo/db/repl/rs_rollback_test.cpp b/src/mongo/db/repl/rs_rollback_test.cpp
index 3d2bd032c64..b36dbb08789 100644
--- a/src/mongo/db/repl/rs_rollback_test.cpp
+++ b/src/mongo/db/repl/rs_rollback_test.cpp
@@ -1396,6 +1396,93 @@ TEST(RSRollbackTest, LocalEntryWithoutO2IsFatal) {
                   RSFatalException);
 }
 
+DEATH_TEST_F(RSRollbackTest, LocalEntryWithTxnNumberWithoutSessionIdIsFatal, "invariant") {
+    auto validOplogEntry = BSON("ts" << Timestamp(Seconds(1), 0) << "t" << 1LL << "h" << 1LL << "op"
+                                     << "i"
+                                     << "ui"
+                                     << UUID::gen()
+                                     << "ns"
+                                     << "test.t"
+                                     << "o"
+                                     << BSON("_id" << 1 << "a" << 1));
+    FixUpInfo fui;
+    ASSERT_OK(updateFixUpInfoFromLocalOplogEntry(fui, validOplogEntry));
+
+    const auto txnNumber = BSON("txnNumber" << 1LL);
+    const auto noSessionIdOrStmtId = validOplogEntry.addField(txnNumber.firstElement());
+
+    const auto stmtId = BSON("stmtId" << 1);
+    const auto noSessionId = noSessionIdOrStmtId.addField(stmtId.firstElement());
+    ASSERT_THROWS(updateFixUpInfoFromLocalOplogEntry(fui, noSessionId).transitional_ignore(),
+                  RSFatalException);
+}
+
+DEATH_TEST_F(RSRollbackTest, LocalEntryWithTxnNumberWithoutStmtIdIsFatal, "invariant") {
+    auto validOplogEntry = BSON("ts" << Timestamp(Seconds(1), 0) << "t" << 1LL << "h" << 1LL << "op"
+                                     << "i"
+                                     << "ui"
+                                     << UUID::gen()
+                                     << "ns"
+                                     << "test.t"
+                                     << "o"
+                                     << BSON("_id" << 1 << "a" << 1));
+    FixUpInfo fui;
+    ASSERT_OK(updateFixUpInfoFromLocalOplogEntry(fui, validOplogEntry));
+
+    const auto txnNumber = BSON("txnNumber" << 1LL);
+    const auto noSessionIdOrStmtId = validOplogEntry.addField(txnNumber.firstElement());
+
+    const auto uuid = UUID::gen();
+    const auto sessionId = BSON("lsid" << BSON("id" << uuid.toBSON()));
+    const auto noStmtId = noSessionIdOrStmtId.addField(sessionId.firstElement());
+    ASSERT_THROWS(updateFixUpInfoFromLocalOplogEntry(fui, noStmtId).transitional_ignore(),
+                  RSFatalException);
+}
+
+TEST(RSRollbackTest, LocalEntryWithTxnNumberAddsTransactionTableDocToBeRefetched) {
+    FixUpInfo fui;
+    auto entryWithoutTxnNumber =
+        BSON("ts" << Timestamp(Seconds(1), 0) << "t" << 1LL << "h" << 1LL << "op"
+                  << "i"
+                  << "ui"
+                  << UUID::gen()
+                  << "ns"
+                  << "test.t2"
+                  << "o"
+                  << BSON("_id" << 2 << "a" << 2));
+    ASSERT_OK(updateFixUpInfoFromLocalOplogEntry(fui, entryWithoutTxnNumber));
+
+    // With no txnNumber present, no extra documents need to be refetched.
+    ASSERT_EQ(fui.docsToRefetch.size(), 1U);
+
+    auto entryWithTxnNumber =
+        BSON("ts" << Timestamp(Seconds(1), 0) << "t" << 1LL << "h" << 1LL << "op"
+                  << "i"
+                  << "ui"
+                  << UUID::gen()
+                  << "ns"
+                  << "test.t"
+                  << "o"
+                  << BSON("_id" << 1 << "a" << 1)
+                  << "txnNumber"
+                  << 1LL
+                  << "stmtId"
+                  << 1
+                  << "lsid"
+                  << BSON("id" << UUID::gen().toBSON()));
+    ASSERT_OK(updateFixUpInfoFromLocalOplogEntry(fui, entryWithTxnNumber));
+
+    // If txnNumber is present, the session transactions table document corresponding to the oplog
+    // entry's sessionId also needs to be refetched.
+    ASSERT_EQ(fui.docsToRefetch.size(), 3U);
+
+    DocID expectedTxnDoc;
+    expectedTxnDoc.ownedObj = BSON("_id" << entryWithTxnNumber["lsid"]);
+    expectedTxnDoc._id = expectedTxnDoc.ownedObj.firstElement();
+    expectedTxnDoc.ns = NamespaceString::kSessionTransactionsTableNamespace.ns().c_str();
+    ASSERT_TRUE(fui.docsToRefetch.find(expectedTxnDoc) != fui.docsToRefetch.end());
+}
+
 TEST_F(RSRollbackTest, RollbackReturnsImmediatelyOnFailureToTransitionToRollback) {
     // On failing to transition to ROLLBACK, rollback() should return immediately and not call
     // syncRollback(). We provide an empty oplog so that if syncRollback() is called erroneously,
diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index 932794cb9b8..37dd16cccee 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -64,6 +64,11 @@ void SessionCatalog::create(ServiceContext* service) {
     sessionTransactionTable.emplace(service);
 }
 
+void SessionCatalog::reset_forTest(ServiceContext* service) {
+    auto& sessionTransactionTable = sessionTransactionTableDecoration(service);
+    sessionTransactionTable.reset();
+}
+
 SessionCatalog* SessionCatalog::get(OperationContext* opCtx) {
     return get(opCtx->getServiceContext());
 }
@@ -130,6 +135,11 @@ ScopedSession SessionCatalog::checkOutSession(OperationContext* opCtx) {
     return ScopedSession(opCtx, std::move(sri));
 }
 
+void SessionCatalog::clearTransactionTable() {
+    stdx::lock_guard<stdx::mutex> lg(_mutex);
+    _txnTable.clear();
+}
+
 void SessionCatalog::_releaseSession(const LogicalSessionId& lsid) {
     stdx::lock_guard<stdx::mutex> lg(_mutex);
 
diff --git a/src/mongo/db/session_catalog.h b/src/mongo/db/session_catalog.h
index 3d77f64a249..c6a67f6c44c 100644
--- a/src/mongo/db/session_catalog.h
+++ b/src/mongo/db/session_catalog.h
@@ -60,6 +60,12 @@ public:
     static void create(ServiceContext* service);
 
     /**
+     * Resets the transaction table on the specified service context to an uninitialized state.
+     * Meant only for testing.
+     */
+    static void reset_forTest(ServiceContext* service);
+
+    /**
      * Retrieves the session transaction table associated with the service or operation context.
      * Must only be called after 'create' has been called.
      */
@@ -86,6 +92,11 @@ public:
      */
     ScopedSession checkOutSession(OperationContext* opCtx);
 
+    /**
+     * Clears the entire transaction table. Invoked after rollback.
+     */
+    void clearTransactionTable();
+
 private:
     struct SessionRuntimeInfo {
         SessionRuntimeInfo(LogicalSessionId lsid) : txnState(std::move(lsid)) {}
diff --git a/src/mongo/db/session_catalog_test.cpp b/src/mongo/db/session_catalog_test.cpp
index cc8ee4ba9b9..0d5f4e9122a 100644
--- a/src/mongo/db/session_catalog_test.cpp
+++ b/src/mongo/db/session_catalog_test.cpp
@@ -45,6 +45,11 @@ protected:
         ServiceContextMongoDTest::setUp();
         SessionCatalog::create(getServiceContext());
     }
+
+    void tearDown() final {
+        SessionCatalog::reset_forTest(getServiceContext());
+        ServiceContextMongoDTest::tearDown();
+    }
 };
 
 TEST_F(SessionCatalogTest, CheckoutAndReleaseSession) {