SERVER-53566: Protect ServiceContext from issuing duplicate operation IDs

tseip-mongo · Evergreen Agent · commit dbb6b0d496ba · 2021-03-22T19:20:07.000Z
diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
@@ -459,6 +459,7 @@ env.Library(
         'operation_context.cpp',
         'operation_context_group.cpp',
         'operation_cpu_timer.cpp',
+        'operation_id.cpp',
         'operation_key_manager.cpp',
         'service_context.cpp',
         'server_recovery.cpp',
@@ -2338,6 +2339,7 @@ if wiredtiger:
             'op_observer_registry_test.cpp',
             'operation_context_test.cpp',
             'operation_cpu_timer_test.cpp',
+            'operation_id_test.cpp',
             'operation_time_tracker_test.cpp',
             'persistent_task_store_test.cpp',
             'range_arithmetic_test.cpp',
diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
@@ -72,8 +72,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()) {}
 
@@ -387,7 +390,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
@@ -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/query/datetime/date_time_support.h"
 #include "mongo/db/storage/recovery_unit.h"
 #include "mongo/db/storage/storage_options.h"
@@ -56,7 +57,6 @@
 
 namespace mongo {
 
-class Client;
 class CurOp;
 class ProgressMeter;
 class ServiceContext;
@@ -97,7 +97,12 @@ class OperationContext : public Interruptible, public Decorable<OperationContext
 public:
     static constexpr auto kDefaultOperationContextTimeoutError = ErrorCodes::ExceededTimeLimit;
 
+    /**
+     * 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 {
@@ -185,7 +190,7 @@ class OperationContext : public Interruptible, public Decorable<OperationContext
      * Returns the operation ID associated with this operation.
      */
     OperationId getOpID() const {
-        return _opId;
+        return _opId.getId();
     }
 
     /**
@@ -614,7 +619,7 @@ class OperationContext : public Interruptible, public Decorable<OperationContext
 
     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
@@ -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
@@ -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
@@ -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
diff --git a/src/mongo/db/service_context.h b/src/mongo/db/service_context.h
