SERVER-49468 Kill and throw when OperationContexts are overwritten

4 files changed, 36 insertions, 39 deletions

diff --git a/src/mongo/db/client.cpp b/src/mongo/db/client.cpp
index 67d957d7559..cd331631002 100644
--- a/src/mongo/db/client.cpp
+++ b/src/mongo/db/client.cpp
@@ -115,17 +115,6 @@ ServiceContext::UniqueOperationContext Client::makeOperationContext() {
     return getServiceContext()->makeOperationContext(this);
 }
 
-void Client::setOperationContext(OperationContext* opCtx) {
-    // We can only set the OperationContext once before resetting it.
-    invariant(opCtx != nullptr && _opCtx == nullptr);
-    _opCtx = opCtx;
-}
-
-void Client::resetOperationContext() {
-    invariant(_opCtx != nullptr);
-    _opCtx = nullptr;
-}
-
 std::string Client::clientAddress(bool includePort) const {
     if (!hasRemote()) {
         return "";
diff --git a/src/mongo/db/client.h b/src/mongo/db/client.h
index 55724f1a90e..65e181fc847 100644
--- a/src/mongo/db/client.h
+++ b/src/mongo/db/client.h
@@ -160,22 +160,6 @@ public:
     ServiceContext::UniqueOperationContext makeOperationContext();
 
     /**
-     * Sets the active operation context on this client to "opCtx", which must be non-NULL.
-     *
-     * It is an error to call this method if there is already an operation context on Client.
-     * It is an error to call this on an unlocked client.
-     */
-    void setOperationContext(OperationContext* opCtx);
-
-    /**
-     * Clears the active operation context on this client.
-     *
-     * There must already be such a context set on this client.
-     * It is an error to call this on an unlocked client.
-     */
-    void resetOperationContext();
-
-    /**
      * Gets the operation context active on this client, or nullptr if there is no such context.
      *
      * It is an error to call this method on an unlocked client, or to use the value returned
@@ -278,6 +262,13 @@ private:
                     ServiceContext* serviceContext,
                     transport::SessionHandle session);
 
+    /**
+     * Sets the active operation context on this client to "opCtx".
+     */
+    void _setOperationContext(OperationContext* opCtx) {
+        _opCtx = opCtx;
+    }
+
     ServiceContext* const _serviceContext;
     const transport::SessionHandle _session;
 
diff --git a/src/mongo/db/service_context.cpp b/src/mongo/db/service_context.cpp
index 8796d7fdae6..d69416b33b0 100644
--- a/src/mongo/db/service_context.cpp
+++ b/src/mongo/db/service_context.cpp
@@ -253,9 +253,21 @@ ServiceContext::UniqueOperationContext ServiceContext::makeOperationContext(Clie
     } else {
         makeBaton(opCtx.get());
     }
+
     {
         stdx::lock_guard<Client> lk(*client);
-        client->setOperationContext(opCtx.get());
+
+        // If we have a previous operation context, it's not worth crashing the process in
+        // production. However, we do want to prevent it from doing more work and complain loudly.
+        auto lastOpCtx = client->getOperationContext();
+        if (lastOpCtx) {
+            killOperation(lk, lastOpCtx, ErrorCodes::Error(4946800));
+            tasserted(
+                4946801,
+                "Client has attempted to create a new OperationContext, but it already has one");
+        }
+
+        client->_setOperationContext(opCtx.get());
     }
 
     {
@@ -372,7 +384,7 @@ void ServiceContext::_delistOperation(OperationContext* opCtx) noexcept {
     // Assigning a new opCtx to the client must never precede the destruction of any existing opCtx
     // that references the client.
     invariant(client->getOperationContext() == opCtx);
-    client->resetOperationContext();
+    client->_setOperationContext({});
 
     if (client->session()) {
         _numCurrentOps.subtractAndFetch(1);
diff --git a/src/mongo/dbtests/cursor_manager_test.cpp b/src/mongo/dbtests/cursor_manager_test.cpp
index a5065dc13f7..9916407cd48 100644
--- a/src/mongo/dbtests/cursor_manager_test.cpp
+++ b/src/mongo/dbtests/cursor_manager_test.cpp
@@ -56,12 +56,17 @@ const NamespaceString kTestNss{"test.collection"};
 
 class CursorManagerTest : public unittest::Test {
 public:
-    CursorManagerTest()
-        : _queryServiceContext(std::make_unique<QueryTestServiceContext>()),
-          _opCtx(_queryServiceContext->makeOperationContext()) {
-        _opCtx->getServiceContext()->setPreciseClockSource(std::make_unique<ClockSourceMock>());
-        _clock =
-            static_cast<ClockSourceMock*>(_opCtx->getServiceContext()->getPreciseClockSource());
+    CursorManagerTest() : _queryServiceContext(std::make_unique<QueryTestServiceContext>()) {
+        _queryServiceContext->getServiceContext()->setPreciseClockSource(
+            std::make_unique<ClockSourceMock>());
+    }
+
+    void setUp() override {
+        _opCtx = _queryServiceContext->makeOperationContext();
+    }
+
+    void tearDown() override {
+        // Do nothing.
     }
 
     std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> makeFakePlanExecutor() {
@@ -102,7 +107,8 @@ public:
     }
 
     ClockSourceMock* useClock() {
-        return _clock;
+        auto svcCtx = _queryServiceContext->getServiceContext();
+        return static_cast<ClockSourceMock*>(svcCtx->getPreciseClockSource());
     }
 
     CursorManager* useCursorManager() {
@@ -114,17 +120,16 @@ protected:
     ServiceContext::UniqueOperationContext _opCtx;
 
 private:
-    ClockSourceMock* _clock;
     CursorManager _cursorManager;
 };
 
 class CursorManagerTestCustomOpCtx : public CursorManagerTest {
     void setUp() override {
-        _queryServiceContext->getClient()->resetOperationContext();
+        // Do nothing.
     }
 
     void tearDown() override {
-        _queryServiceContext->getClient()->setOperationContext(_opCtx.get());
+        // Do nothing.
     }
 };