SERVER-53566: Protect ServiceContext from issuing duplicate operation IDs

8 files changed, 354 insertions, 18 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 1156249c3f0..cf429e583dc 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -434,6 +434,7 @@ env.Library(
         'default_baton.cpp',
         'operation_context.cpp',
         'operation_context_group.cpp',
+        'operation_id.cpp',
         'operation_key_manager.cpp',
         'service_context.cpp',
         'server_recovery.cpp',
@@ -1853,6 +1854,7 @@ envWithAsio.CppUnitTest(
         'op_observer_impl_test.cpp',
         'op_observer_registry_test.cpp',
         'operation_context_test.cpp',
+        'operation_id_test.cpp',
         'operation_time_tracker_test.cpp',
         'range_arithmetic_test.cpp',
         'read_write_concern_defaults_test.cpp',
diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index d468560a8e9..8817793d618 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -81,8 +81,11 @@ const auto kNoWaiterThread = stdx::thread::id();
 }  // namespace
 
 OperationContext::OperationContext(Client* client, OperationId opId)
+    : OperationContext(client, OperationIdSlot(opId)) {}
+
+OperationContext::OperationContext(Client* client, OperationIdSlot&& opIdSlot)
     : _client(client),
-      _opId(opId),
+      _opId(std::move(opIdSlot)),
       _elapsedTime(client ? client->getServiceContext()->getTickSource()
                           : SystemTickSource::get()) {}
 
@@ -395,7 +398,7 @@ void OperationContext::setOperationKey(OperationKey opKey) {
     invariant(!_opKey);
 
     _opKey.emplace(std::move(opKey));
-    OperationKeyManager::get(_client).add(*_opKey, _opId);
+    OperationKeyManager::get(_client).add(*_opKey, _opId.getId());
 }
 
 void OperationContext::releaseOperationKey() {
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index d8f92af60cf..1ba407704b1 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -36,6 +36,7 @@
 #include "mongo/db/client.h"
 #include "mongo/db/concurrency/locker.h"
 #include "mongo/db/logical_session_id.h"
+#include "mongo/db/operation_id.h"
 #include "mongo/db/storage/recovery_unit.h"
 #include "mongo/db/storage/storage_options.h"
 #include "mongo/db/storage/write_unit_of_work.h"
@@ -53,7 +54,6 @@
 
 namespace mongo {
 
-class Client;
 class CurOp;
 class ProgressMeter;
 class ServiceContext;
@@ -79,7 +79,12 @@ class OperationContext : public Interruptible, public Decorable<OperationContext
     OperationContext& operator=(const OperationContext&) = delete;
 
 public:
+    /**
+     * Creates an op context with no unique operation ID tracking - prefer using the OperationIdSlot
+     * CTOR if possible to avoid OperationId collisions.
+     */
     OperationContext(Client* client, OperationId opId);
+    OperationContext(Client* client, OperationIdSlot&& opIdSlot);
     virtual ~OperationContext();
 
     bool shouldParticipateInFlowControl() const {
@@ -167,7 +172,7 @@ public:
      * Returns the operation ID associated with this operation.
      */
     OperationId getOpID() const {
-        return _opId;
+        return _opId.getId();
     }
 
     /**
@@ -546,7 +551,7 @@ private:
 
     Client* const _client;
 
-    const OperationId _opId;
+    const OperationIdSlot _opId;
     boost::optional<OperationKey> _opKey;
 
     boost::optional<LogicalSessionId> _lsid;
diff --git a/src/mongo/db/operation_id.cpp b/src/mongo/db/operation_id.cpp
new file mode 100644
index 00000000000..5a4ca82ed13
--- /dev/null
+++ b/src/mongo/db/operation_id.cpp
@@ -0,0 +1,56 @@
+/**
+ *    Copyright (C) 2018-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/operation_id.h"
+
+#include "mongo/util/assert_util.h"
+
+namespace mongo {
+
+OperationIdSlot UniqueOperationIdRegistry::acquireSlot() {
+    stdx::lock_guard lk(_mutex);
+
+    // Make sure the set isn't absolutely enormous. If it is, something else is wrong,
+    // and the loop below could fail.
+    invariant(_activeIds.size() < (1 << 20));
+
+    while (true) {
+        const auto&& [it, ok] = _activeIds.insert(_nextOpId++);
+        if (ok) {
+            return OperationIdSlot(*it, shared_from_this());
+        }
+    }
+}
+
+void UniqueOperationIdRegistry::_releaseSlot(OperationId id) {
+    stdx::lock_guard lk(_mutex);
+    invariant(_activeIds.erase(id));
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_id.h b/src/mongo/db/operation_id.h
new file mode 100644
index 00000000000..44c2a28663d
--- /dev/null
+++ b/src/mongo/db/operation_id.h
@@ -0,0 +1,145 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/platform/mutex.h"
+#include "mongo/stdx/unordered_set.h"
+
+namespace mongo {
+
+/**
+ * Every OperationContext is expected to have a unique OperationId within the domain of its
+ * ServiceContext. Generally speaking, OperationId is used for forming maps of OperationContexts and
+ * directing metaoperations like killop.
+ */
+using OperationId = uint32_t;
+
+/**
+ * This class issues guaranteed unique OperationIds for a given instance of this class.
+ */
+class UniqueOperationIdRegistry : public std::enable_shared_from_this<UniqueOperationIdRegistry> {
+public:
+    /**
+     * This class represents a slot issued by a UniqueOperationIdRegistry.
+     * It functions as an RAII wrapper for a unique OperationId.
+     */
+    class OperationIdSlot {
+    public:
+        explicit OperationIdSlot(OperationId id) : _id(id), _registry() {}
+
+        OperationIdSlot(OperationId id, std::weak_ptr<UniqueOperationIdRegistry> registry)
+            : _id(id), _registry(std::move(registry)) {}
+
+        OperationIdSlot(OperationIdSlot&& other) = default;
+
+        OperationIdSlot& operator=(OperationIdSlot&& other) noexcept {
+            if (&other == this) {
+                return *this;
+            }
+            _releaseSlot();
+            _id = std::exchange(other._id, {});
+            _registry = std::exchange(other._registry, {});
+            return *this;
+        }
+
+        // Disable copies.
+        OperationIdSlot(const OperationIdSlot&) = delete;
+        OperationIdSlot& operator=(const OperationIdSlot&) = delete;
+
+        ~OperationIdSlot() {
+            _releaseSlot();
+        }
+
+        /**
+         * Get the contained ID.
+         */
+        OperationId getId() const {
+            return _id;
+        }
+
+    private:
+        void _releaseSlot() {
+            if (auto registry = _registry.lock()) {
+                registry->_releaseSlot(_id);
+            }
+        }
+
+        OperationId _id;
+        std::weak_ptr<UniqueOperationIdRegistry> _registry;
+    };
+
+    /**
+     * Public factory function.
+     */
+    static std::shared_ptr<UniqueOperationIdRegistry> create() {
+        return std::shared_ptr<UniqueOperationIdRegistry>(new UniqueOperationIdRegistry());
+    }
+
+    /**
+     * Gets a unique OperationIdSlot which will clear itself from the map when destroyed.
+     */
+    OperationIdSlot acquireSlot();
+
+    /**
+     * A helper class for exposing test functions.
+     */
+    class UniqueOperationIdRegistryTestHarness {
+    public:
+        /**
+         * Returns true if the given operation ID exists.
+         */
+        static bool isActive(UniqueOperationIdRegistry& registry, OperationId id) {
+            stdx::lock_guard lk(registry._mutex);
+            return registry._activeIds.find(id) != registry._activeIds.end();
+        }
+
+        static void setNextOpId(UniqueOperationIdRegistry& registry, OperationId id) {
+            stdx::lock_guard lk(registry._mutex);
+            registry._nextOpId = id;
+        }
+    };
+
+private:
+    UniqueOperationIdRegistry() = default;
+
+    /**
+     * Clears a unique ID from the set.
+     */
+    void _releaseSlot(OperationId id);
+
+    Mutex _mutex = MONGO_MAKE_LATCH("UniqueOperationIdRegistry::_mutex");
+    stdx::unordered_set<OperationId> _activeIds;
+
+    OperationId _nextOpId = 1U;
+};
+
+using OperationIdSlot = UniqueOperationIdRegistry::OperationIdSlot;
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_id_test.cpp b/src/mongo/db/operation_id_test.cpp
new file mode 100644
index 00000000000..825891d90ae
--- /dev/null
+++ b/src/mongo/db/operation_id_test.cpp
@@ -0,0 +1,110 @@
+/**
+ *    Copyright (C) 2018-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/operation_id.h"
+#include "mongo/logv2/log.h"
+#include "mongo/platform/mutex.h"
+#include "mongo/unittest/death_test.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo {
+namespace {
+
+using UniqueOperationIdRegistryTestHarness =
+    UniqueOperationIdRegistry::UniqueOperationIdRegistryTestHarness;
+
+TEST(OperationIdTest, OperationIdIncrementsProperly) {
+    auto registry = UniqueOperationIdRegistry::create();
+    auto slot = registry->acquireSlot();
+    ASSERT_EQ(slot.getId(), 1);
+
+    {
+        auto slot2 = registry->acquireSlot();
+        ASSERT_EQ(slot2.getId(), 2);
+    }
+
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 1));
+    ASSERT(!UniqueOperationIdRegistryTestHarness::isActive(*registry, 2));
+
+    slot = registry->acquireSlot();
+    ASSERT_EQ(slot.getId(), 3);
+
+    ASSERT(!UniqueOperationIdRegistryTestHarness::isActive(*registry, 1));
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 3));
+}
+
+TEST(OperationIdTest, OperationIdsAreUnique) {
+    auto registry = UniqueOperationIdRegistry::create();
+    auto slot = registry->acquireSlot();
+    ASSERT_EQ(slot.getId(), 1);
+
+    auto slot2 = registry->acquireSlot();
+    ASSERT_EQ(slot2.getId(), 2);
+
+    // If the registry's state points to a currently active operation ID, we will
+    // find the next hole and issue that instead.
+    UniqueOperationIdRegistryTestHarness::setNextOpId(*registry, 1);
+    auto slot3 = registry->acquireSlot();
+    ASSERT_EQ(slot3.getId(), 3);
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 1));
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 2));
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 3));
+}
+
+TEST(OperationIdTest, OperationIdsAreMovable) {
+    auto registry = UniqueOperationIdRegistry::create();
+    auto slot = registry->acquireSlot();
+    ASSERT_EQ(slot.getId(), 1);
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 1));
+
+    auto moveSlot(std::move(slot));
+    ASSERT_EQ(moveSlot.getId(), 1);
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 1));
+
+    auto assignSlot = std::move(moveSlot);
+    ASSERT_EQ(assignSlot.getId(), 1);
+    ASSERT(UniqueOperationIdRegistryTestHarness::isActive(*registry, 1));
+}
+
+DEATH_TEST(OperationIdTest, TooManyTransactionsShouldCrash, "invariant") {
+    auto registry = UniqueOperationIdRegistry::create();
+    std::vector<OperationIdSlot> slots;
+
+    while (true) {
+        slots.push_back(registry->acquireSlot());
+    }
+}
+
+}  // namespace
+
+}  // namespace mongo
diff --git a/src/mongo/db/service_context.cpp b/src/mongo/db/service_context.cpp
index 45cacbb5786..2f3b3399c89 100644
--- a/src/mongo/db/service_context.cpp
+++ b/src/mongo/db/service_context.cpp
@@ -103,10 +103,12 @@ bool supportsDocLocking() {
 }
 
 ServiceContext::ServiceContext()
-    : _tickSource(std::make_unique<SystemTickSource>()),
+    : _opIdRegistry(UniqueOperationIdRegistry::create()),
+      _tickSource(std::make_unique<SystemTickSource>()),
       _fastClockSource(std::make_unique<SystemClockSource>()),
       _preciseClockSource(std::make_unique<SystemClockSource>()) {}
 
+
 ServiceContext::~ServiceContext() {
     stdx::lock_guard<Latch> lk(_mutex);
     for (const auto& client : _clients) {
@@ -248,12 +250,21 @@ void ServiceContext::ClientDeleter::operator()(Client* client) const {
 }
 
 ServiceContext::UniqueOperationContext ServiceContext::makeOperationContext(Client* client) {
-    auto opCtx = std::make_unique<OperationContext>(client, _nextOpId.fetchAndAdd(1));
+    auto opCtx = std::make_unique<OperationContext>(client, _opIdRegistry->acquireSlot());
+
     if (client->session()) {
         _numCurrentOps.addAndFetch(1);
     }
 
+    auto numOpsGuard = makeGuard([&] {
+        if (client->session()) {
+            _numCurrentOps.subtractAndFetch(1);
+        }
+    });
+
     onCreate(opCtx.get(), _clientObservers);
+    auto onCreateGuard = makeGuard([&] { onDestroy(opCtx.get(), _clientObservers); });
+
     if (!opCtx->lockState()) {
         opCtx->setLockState(std::make_unique<LockerNoop>());
     }
@@ -267,14 +278,23 @@ ServiceContext::UniqueOperationContext ServiceContext::makeOperationContext(Clie
     } else {
         makeBaton(opCtx.get());
     }
+
+    auto batonGuard = makeGuard([&] { opCtx->getBaton()->detach(); });
+
     {
         stdx::lock_guard<Client> lk(*client);
         client->setOperationContext(opCtx.get());
     }
 
+    numOpsGuard.dismiss();
+    onCreateGuard.dismiss();
+    batonGuard.dismiss();
+
     {
         stdx::lock_guard lk(_mutex);
-        _clientByOperationId.emplace(opCtx->getOpID(), client);
+        bool clientByOperationContextInsertionSuccessful =
+            _clientByOperationId.insert({opCtx->getOpID(), client}).second;
+        invariant(clientByOperationContextInsertionSuccessful);
     }
 
     return UniqueOperationContext(opCtx.release());
diff --git a/src/mongo/db/service_context.h b/src/mongo/db/service_context.h
index b93e97e4c2b..57302f199d5 100644
--- a/src/mongo/db/service_context.h
+++ b/src/mongo/db/service_context.h
@@ -37,6 +37,7 @@
 
 #include "mongo/base/global_initializer_registerer.h"
 #include "mongo/db/logical_session_id.h"
+#include "mongo/db/operation_id.h"
 #include "mongo/db/storage/storage_engine.h"
 #include "mongo/platform/atomic_word.h"
 #include "mongo/platform/mutex.h"
@@ -124,13 +125,6 @@ private:
 };
 
 /**
- * Every OperationContext is expected to have a unique OperationId within the domain of its
- * ServiceContext. Generally speaking, OperationId is used for forming maps of OperationContexts and
- * directing metaoperations like killop.
- */
-using OperationId = uint32_t;
-
-/**
  * Users may provide an OperationKey when sending a command request as a stable token by which to
  * refer to an operation (and thus an OperationContext). An OperationContext is not required to have
  * an OperationKey. The presence of an OperationKey implies that the client is either closely
@@ -663,6 +657,10 @@ private:
     std::vector<ClientObserverHolder> _clientObservers;
     ClientSet _clients;
 
+    /**
+     * Managing classes for our issued operation IDs.
+     */
+    std::shared_ptr<UniqueOperationIdRegistry> _opIdRegistry;
     stdx::unordered_map<OperationId, Client*> _clientByOperationId;
 
     /**
@@ -689,9 +687,6 @@ private:
     // protected by _mutex
     std::vector<KillOpListenerInterface*> _killOpListeners;
 
-    // Counter for assigning operation ids.
-    AtomicWord<OperationId> _nextOpId{1};
-
     // When the catalog is restarted, the generation goes up by one each time.
     AtomicWord<uint64_t> _catalogGeneration{0};