SERVER-38224 Allow sessions to have more than one outstanding kill request

Also reverts commit e13e069902099d601c6cf64def9fc374082a629e (SERVER-38058 Make retryable_writes_direct_write_to_config_transactions.js expect ConflictingOperationInProgress) since it is no longer necessary.

3 files changed, 50 insertions, 30 deletions

diff --git a/src/mongo/db/session.cpp b/src/mongo/db/session.cpp
index adbd165c267..f72457d7f94 100644
--- a/src/mongo/db/session.cpp
+++ b/src/mongo/db/session.cpp
@@ -45,15 +45,12 @@ OperationContext* Session::currentOperation() const {
 
 Session::KillToken Session::kill(WithLock sessionCatalogLock, ErrorCodes::Error reason) {
     stdx::lock_guard<stdx::mutex> lg(_mutex);
-    uassert(ErrorCodes::ConflictingOperationInProgress,
-            str::stream() << "Session " << getSessionId().getId()
-                          << " is already killed and is in the process of being cleaned up",
-            !_killRequested);
-    _killRequested = true;
+    const bool firstKiller = (0 == _killsRequested);
+    ++_killsRequested;
 
     // For currently checked-out sessions, interrupt the operation context so that the current owner
     // can release the session
-    if (_checkoutOpCtx) {
+    if (firstKiller && _checkoutOpCtx) {
         const auto serviceContext = _checkoutOpCtx->getServiceContext();
 
         stdx::lock_guard<Client> clientLock(*_checkoutOpCtx->getClient());
@@ -65,7 +62,7 @@ Session::KillToken Session::kill(WithLock sessionCatalogLock, ErrorCodes::Error
 
 bool Session::killed() const {
     stdx::lock_guard<stdx::mutex> lg(_mutex);
-    return _killRequested;
+    return _killsRequested > 0;
 }
 
 void Session::_markCheckedOut(WithLock sessionCatalogLock, OperationContext* checkoutOpCtx) {
@@ -82,8 +79,8 @@ void Session::_markCheckedIn(WithLock sessionCatalogLock) {
 
 void Session::_markNotKilled(WithLock sessionCatalogLock, KillToken killToken) {
     stdx::lock_guard<stdx::mutex> lg(_mutex);
-    invariant(_killRequested);
-    _killRequested = false;
+    invariant(_killsRequested > 0);
+    --_killsRequested;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/session.h b/src/mongo/db/session.h
index f7349e0d29e..ba01d6458f5 100644
--- a/src/mongo/db/session.h
+++ b/src/mongo/db/session.h
@@ -65,17 +65,17 @@ public:
     OperationContext* currentOperation() const;
 
     /**
-     * Marks the session as killed and returns a 'kill token' to to be passed later on to
-     * 'checkOutSessionForKill' method of the SessionCatalog in order to permit the caller to
-     * execute any kill cleanup tasks and pass further on to '_markNotKilled' in order to reset the
-     * kill state. Marking session as killed is an internal property only that will cause any
-     * further calls to 'checkOutSession' to block until 'checkOutSessionForKill' is called and the
-     * returned scoped object destroyed.
+     * Increments the number of "killers" for this session and returns a 'kill token' to to be
+     * passed later on to 'checkOutSessionForKill' method of the SessionCatalog in order to permit
+     * the caller to execute any kill cleanup tasks. This token is later on passed to
+     * '_markNotKilled' in order to decrement the number of "killers".
      *
-     * If the session is currently checked-out, this method will also interrupt the operation
-     * context which has it checked-out.
+     * Marking session as killed is an internal property only that will cause any further calls to
+     * 'checkOutSession' to block until 'checkOutSessionForKill' is called the same number of times
+     * as 'kill' was called and the returned scoped object destroyed.
      *
-     * If the session is already killed throws ConflictingOperationInProgress exception.
+     * If the first killer finds the session checked-out, this method will also interrupt the
+     * operation context which has it checked-out.
      *
      * Must be called under the owning SessionCatalog's lock.
      */
@@ -123,8 +123,8 @@ private:
     // is no operation currently running for the Session.
     OperationContext* _checkoutOpCtx{nullptr};
 
-    // Set to true if markKilled has been invoked for this session.
-    bool _killRequested{false};
+    // Incremented every time 'kill' is invoked and decremented by '_markNotKilled'.
+    int _killsRequested{0};
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/session_catalog_test.cpp b/src/mongo/db/session_catalog_test.cpp
index 48431675529..29f0395ed8a 100644
--- a/src/mongo/db/session_catalog_test.cpp
+++ b/src/mongo/db/session_catalog_test.cpp
@@ -270,7 +270,7 @@ TEST_F(SessionCatalogTest, KillSessionWhenSessionIsCheckedOut) {
     future.get();
 }
 
-TEST_F(SessionCatalogTest, MarkSessionAsKilledThrowsWhenCalledTwice) {
+TEST_F(SessionCatalogTest, MarkSessionAsKilledCanBeCalledMoreThanOnce) {
     const auto lsid = makeLogicalSessionIdForTest();
 
     // Create the session so there is something to kill
@@ -279,14 +279,11 @@ TEST_F(SessionCatalogTest, MarkSessionAsKilledThrowsWhenCalledTwice) {
         const auto unusedSession(catalog()->checkOutSession(opCtx.get(), lsid));
     }
 
-    auto killToken = catalog()->killSession(lsid);
-
-    // Second mark as killed attempt will throw since the session is already killed
-    ASSERT_THROWS_CODE(catalog()->killSession(lsid),
-                       AssertionException,
-                       ErrorCodes::ConflictingOperationInProgress);
+    auto killToken1 = catalog()->killSession(lsid);
+    auto killToken2 = catalog()->killSession(lsid);
 
-    // Make sure that regular session check-out will fail because the session is marked as killed
+    // Make sure that regular session check-out will fail because there are two killers on the
+    // session
     {
         auto opCtx = makeOperationContext();
         opCtx->setLogicalSessionId(lsid);
@@ -295,10 +292,36 @@ TEST_F(SessionCatalogTest, MarkSessionAsKilledThrowsWhenCalledTwice) {
             OperationContextSession(opCtx.get()), AssertionException, ErrorCodes::MaxTimeMSExpired);
     }
 
-    // Finish "killing" the session so the SessionCatalog destructor doesn't complain
+    boost::optional<Session::KillToken> killTokenWhileSessionIsCheckedOutForKill;
+
+    // Finish the first killer of the session
     {
         auto opCtx = makeOperationContext();
-        auto scopedSession = catalog()->checkOutSessionForKill(opCtx.get(), std::move(killToken));
+        auto scopedSession = catalog()->checkOutSessionForKill(opCtx.get(), std::move(killToken1));
+        ASSERT_EQ(opCtx.get(), scopedSession->currentOperation());
+
+        // Killing a session while checked out for kill should not affect the killers
+        killTokenWhileSessionIsCheckedOutForKill.emplace(catalog()->killSession(lsid));
+    }
+
+    // Regular session check-out should still fail because there are now still two killers on the
+    // session
+    {
+        auto opCtx = makeOperationContext();
+        opCtx->setLogicalSessionId(lsid);
+        opCtx->setDeadlineAfterNowBy(Milliseconds(10), ErrorCodes::MaxTimeMSExpired);
+        ASSERT_THROWS_CODE(
+            OperationContextSession(opCtx.get()), AssertionException, ErrorCodes::MaxTimeMSExpired);
+    }
+    {
+        auto opCtx = makeOperationContext();
+        auto scopedSession = catalog()->checkOutSessionForKill(opCtx.get(), std::move(killToken2));
+        ASSERT_EQ(opCtx.get(), scopedSession->currentOperation());
+    }
+    {
+        auto opCtx = makeOperationContext();
+        auto scopedSession = catalog()->checkOutSessionForKill(
+            opCtx.get(), std::move(*killTokenWhileSessionIsCheckedOutForKill));
         ASSERT_EQ(opCtx.get(), scopedSession->currentOperation());
     }
 }