SERVER-48295 Remove operation key in killAndDelistOperation

3 files changed, 16 insertions, 3 deletions

diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index 8bc40f2d991..d468560a8e9 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -87,9 +87,7 @@ OperationContext::OperationContext(Client* client, OperationId opId)
                           : SystemTickSource::get()) {}
 
 OperationContext::~OperationContext() {
-    if (_opKey) {
-        OperationKeyManager::get(_client).remove(*_opKey);
-    }
+    releaseOperationKey();
 }
 
 void OperationContext::setDeadlineAndMaxTime(Date_t when,
@@ -400,6 +398,13 @@ void OperationContext::setOperationKey(OperationKey opKey) {
     OperationKeyManager::get(_client).add(*_opKey, _opId);
 }
 
+void OperationContext::releaseOperationKey() {
+    if (_opKey) {
+        OperationKeyManager::get(_client).remove(*_opKey);
+    }
+    _opKey = boost::none;
+}
+
 void OperationContext::setTxnNumber(TxnNumber txnNumber) {
     invariant(_lsid);
     _txnNumber = txnNumber;
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index ea565afd52c..c5116e805b2 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -185,6 +185,12 @@ public:
     void setOperationKey(OperationKey opKey);
 
     /**
+     * Removes the operation UUID associated with this operation.
+     * DO NOT call this function outside `~OperationContext()` and `killAndDelistOperation()`.
+     */
+    void releaseOperationKey();
+
+    /**
      * Returns the session ID associated with this operation, if there is one.
      */
     const boost::optional<LogicalSessionId>& getLogicalSessionId() const {
diff --git a/src/mongo/db/service_context.cpp b/src/mongo/db/service_context.cpp
index 58f440af121..9d9f907aa1b 100644
--- a/src/mongo/db/service_context.cpp
+++ b/src/mongo/db/service_context.cpp
@@ -385,6 +385,8 @@ void ServiceContext::_delistOperation(OperationContext* opCtx) noexcept {
     if (client->session()) {
         _numCurrentOps.subtractAndFetch(1);
     }
+
+    opCtx->releaseOperationKey();
 }
 
 void ServiceContext::killAndDelistOperation(OperationContext* opCtx,