SERVER-39692 Make mongos shutdown drain all in-progress transactions

1 files changed, 17 insertions, 0 deletions

diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index d8ac79858a2..55b0c331ad4 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -251,6 +251,23 @@ OperationContextSession::OperationContextSession(OperationContext* opCtx) : _opC
     checkOut(opCtx);
 }
 
+OperationContextSession::OperationContextSession(OperationContext* opCtx,
+                                                 SessionCatalog::KillToken killToken)
+    : _opCtx(opCtx) {
+    auto& checkedOutSession = operationSessionDecoration(opCtx);
+
+    invariant(!checkedOutSession);
+    invariant(!opCtx->getLogicalSessionId());  // lsid is specified by killToken argument.
+
+    const auto catalog = SessionCatalog::get(opCtx);
+    auto scopedSessionForKill = catalog->checkOutSessionForKill(opCtx, std::move(killToken));
+
+    // We acquire a Client lock here to guard the construction of this session so that references to
+    // this session are safe to use while the lock is held
+    stdx::lock_guard lk(*opCtx->getClient());
+    checkedOutSession.emplace(std::move(scopedSessionForKill._scos));
+}
+
 OperationContextSession::~OperationContextSession() {
     // Only release the checked out session at the end of the top-level request from the client, not
     // at the end of a nested DBDirectClient call