SERVER-58960 Add benchmarks to measure the overhead from ShardingWriteRouter on the write codepath

8 files changed, 384 insertions, 4 deletions

diff --git a/buildscripts/resmokeconfig/suites/benchmarks.yml b/buildscripts/resmokeconfig/suites/benchmarks.yml
index 2676e09a82f..5ac56882be0 100644
--- a/buildscripts/resmokeconfig/suites/benchmarks.yml
+++ b/buildscripts/resmokeconfig/suites/benchmarks.yml
@@ -9,11 +9,13 @@ selector:
   # The trailing asterisk is for handling the .exe extension on Windows.
   # These benchmarks are being run as part of the benchmarks_sharding.yml test suite.
   - build/install/bin/chunk_manager_refresh_bm*
+  - build/install/bin/sharding_write_router_bm*
   # These benchmarks are being run as part of the benchmarks_cst.yml test suite.
   - build/install/bin/cst_bm*
   # Hash table benchmark is really slow, don't run on evergreen
   - build/install/bin/hash_table_bm*
 
+
 executor:
   config: {}
   hooks:
diff --git a/buildscripts/resmokeconfig/suites/benchmarks_sharding.yml b/buildscripts/resmokeconfig/suites/benchmarks_sharding.yml
index 4ee0fb66fa9..d6dd3fd5d53 100644
--- a/buildscripts/resmokeconfig/suites/benchmarks_sharding.yml
+++ b/buildscripts/resmokeconfig/suites/benchmarks_sharding.yml
@@ -6,6 +6,7 @@ selector:
   # The trailing asterisk is for handling the .exe extension on Windows.
   - build/**/system_resource_canary_bm*
   - build/install/bin/chunk_manager_refresh_bm*
+  - build/install/bin/sharding_write_router_bm*
 
 executor:
   config: {}
diff --git a/src/mongo/db/s/SConscript b/src/mongo/db/s/SConscript
index f8373b1e0b0..8fbed5c7247 100644
--- a/src/mongo/db/s/SConscript
+++ b/src/mongo/db/s/SConscript
@@ -587,3 +587,15 @@ env.CppUnitTest(
         'config_server_test_fixture',
     ],
 )
+
+env.Benchmark(
+    target="sharding_write_router_bm",
+    source=["sharding_write_router_bm.cpp"],
+    LIBDEPS=[
+        "$BUILD_DIR/mongo/db/auth/authmocks",
+        "$BUILD_DIR/mongo/s/grid",
+        "$BUILD_DIR/mongo/s/sharding_test_fixture_common",
+        "sharding_api_d",
+        "sharding_runtime_d",
+    ],
+)
diff --git a/src/mongo/db/s/sharding_write_router_bm.cpp b/src/mongo/db/s/sharding_write_router_bm.cpp
new file mode 100644
index 00000000000..a9420c64c3a
--- /dev/null
+++ b/src/mongo/db/s/sharding_write_router_bm.cpp
@@ -0,0 +1,239 @@
+/**
+ *    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.
+ */
+
+#include "mongo/db/keypattern.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/s/collection_sharding_state_factory_shard.h"
+#include "mongo/db/service_context.h"
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/s/collection_sharding_runtime.h"
+#include "mongo/s/catalog_cache.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
+#include "mongo/s/catalog_cache_mock.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/processinfo.h"
+#include "mongo/util/uuid.h"
+#include <benchmark/benchmark.h>
+
+#include "mongo/base/init.h"
+#include "mongo/db/s/collection_metadata.h"
+#include "mongo/db/s/operation_sharding_state.h"
+#include "mongo/db/s/sharding_state.h"
+#include "mongo/db/s/sharding_write_router.h"
+#include "mongo/platform/random.h"
+#include "mongo/s/chunk_manager.h"
+#include "mongo/util/assert_util.h"
+#include "mongo/util/str.h"
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+
+namespace mongo {
+namespace {
+
+const NamespaceString kNss("test", "foo");
+ShardId pessimalShardSelector(int i, int nShards, int nChunks) {
+    return ShardId(str::stream() << "shard" << (i % nShards));
+}
+
+ChunkRange getRangeForChunk(int i, int nChunks) {
+    invariant(i >= 0);
+    invariant(nChunks > 0);
+    invariant(i < nChunks);
+    if (i == 0) {
+        return {BSON("_id" << MINKEY), BSON("_id" << 0)};
+    }
+    if (i + 1 == nChunks) {
+        return {BSON("_id" << (i - 1) * 100), BSON("_id" << MAXKEY)};
+    }
+    return {BSON("_id" << (i - 1) * 100), BSON("_id" << i * 100)};
+}
+
+RoutingTableHistoryValueHandle makeStandaloneRoutingTableHistory(RoutingTableHistory rt) {
+    const auto version = rt.getVersion();
+    return RoutingTableHistoryValueHandle(
+        std::move(rt), ComparableChunkVersion::makeComparableChunkVersion(version));
+}
+
+std::pair<std::vector<mongo::ChunkType>, mongo::ChunkManager> createChunks(
+    size_t nShards, uint32_t nChunks, std::vector<ShardId> shards) {
+    invariant(shards.size() == nShards);
+
+    const auto collIdentifier = UUID::gen();
+    const auto shardKeyPattern = KeyPattern(BSON("_id" << 1));
+    const auto reshardKeyPattern = KeyPattern(BSON("y" << 1));
+    const auto collEpoch = OID::gen();
+    const auto tempNss =
+        NamespaceString(kNss.db(),
+                        fmt::format("{}{}",
+                                    NamespaceString::kTemporaryReshardingCollectionPrefix,
+                                    collIdentifier.toString()));
+
+    std::vector<ChunkType> chunks;
+    chunks.reserve(nChunks);
+
+    for (uint32_t i = 0; i < nChunks; ++i) {
+        chunks.emplace_back(collIdentifier,
+                            getRangeForChunk(i, nChunks),
+                            ChunkVersion{i + 1, 0, collEpoch, Timestamp() /* timestamp */},
+                            pessimalShardSelector(i, nShards, nChunks));
+    }
+
+    TypeCollectionReshardingFields reshardingFields{UUID::gen()};
+    reshardingFields.setState(CoordinatorStateEnum::kPreparingToDonate);
+    // ShardingWriteRouter is only meant to be used by the donor.
+    reshardingFields.setDonorFields(TypeCollectionDonorFields{tempNss, reshardKeyPattern, shards});
+
+    ChunkManager cm(shards[0],
+                    DatabaseVersion(UUID::gen(), Timestamp()),
+                    makeStandaloneRoutingTableHistory(
+                        RoutingTableHistory::makeNew(kNss,
+                                                     collIdentifier,
+                                                     shardKeyPattern,
+                                                     nullptr,
+                                                     false,
+                                                     collEpoch,
+                                                     Timestamp(),
+                                                     boost::none /* timeseriesFields */,
+                                                     reshardingFields, /* reshardingFields */
+                                                     true,
+                                                     chunks)),
+                    boost::none);
+
+
+    return std::make_pair(chunks, cm);
+}
+
+
+std::unique_ptr<CatalogCacheMock> createCatalogCacheMock(OperationContext* opCtx) {
+
+    const size_t nShards = 1;
+    const uint32_t nChunks = 60;
+    const auto clusterId = OID::gen();
+    const auto shards = std::vector<ShardId>{ShardId("shard0")};
+    const auto originatorShard = shards[0];
+
+
+    const auto [chunks, chunkManager] = createChunks(nShards, nChunks, shards);
+
+    // Necessary to make usages of 'CollectionShardingState::get(opCtx, nss);' work
+    CollectionShardingStateFactory::set(
+        opCtx->getServiceContext(),
+        std::make_unique<CollectionShardingStateFactoryShard>(opCtx->getServiceContext()));
+
+    // necessary to make ShardingState::get(_serviceContext)->enabled() return true
+    ShardingState::get(opCtx->getServiceContext())->setInitialized(originatorShard, clusterId);
+
+    BSONObjBuilder builder;
+    chunkManager.getVersion(originatorShard).appendToCommand(&builder);
+    // necessary to set the _shardVersions and _databaseVersions to true. Which is needed to get
+    // `getCollectionDescription` to work
+    OperationShardingState::get(opCtx).initializeClientRoutingVersionsFromCommand(kNss,
+                                                                                  builder.obj());
+
+    // Configuring the filtering metadata such that calls to getCollectionDescription return what we
+    // want. Specifically the reshardingFields are what we use. Its specified by the chunkManager.
+    CollectionShardingRuntime::get(opCtx, kNss)
+        ->setFilteringMetadata(opCtx, CollectionMetadata(chunkManager, originatorShard));
+
+    auto catalogCache = CatalogCacheMock::make();
+    catalogCache->setChunkManagerReturnValue(chunkManager);
+
+    return catalogCache;
+}
+
+static void BM_InsertGetDestinedRecipient(benchmark::State& state) {
+    // ShardingWriteRouter currently requires the ShardServer cluster role.
+    serverGlobalParams.clusterRole = ClusterRole::ShardServer;
+
+    auto serviceContext = ServiceContext::make();
+    const auto client = serviceContext->makeClient("test");
+    const auto opCtx = client->makeOperationContext();
+
+    const auto catalogCache = createCatalogCacheMock(opCtx.get());
+
+    ShardingWriteRouter writeRouter(opCtx.get(), kNss, catalogCache.get());
+
+    for (auto keepRunning : state) {
+        benchmark::ClobberMemory();
+        auto shardId = writeRouter.getReshardingDestinedRecipient(BSON("_id" << 0));
+        ASSERT(shardId != boost::none);
+    }
+}
+
+static void BM_UpdateGetDestinedRecipient(benchmark::State& state) {
+    // ShardingWriteRouter currently requires the ShardServer cluster role.
+    serverGlobalParams.clusterRole = ClusterRole::ShardServer;
+
+    auto serviceContext = ServiceContext::make();
+    const auto client = serviceContext->makeClient("test");
+    const auto opCtx = client->makeOperationContext();
+
+    const auto catalogCache = createCatalogCacheMock(opCtx.get());
+
+    for (auto keepRunning : state) {
+        benchmark::ClobberMemory();
+        ShardingWriteRouter writeRouter(opCtx.get(), kNss, catalogCache.get());
+        auto shardId = writeRouter.getReshardingDestinedRecipient(BSON("_id" << 0));
+        ASSERT(shardId != boost::none);
+    }
+}
+
+static void BM_UnshardedDestinedRecipient(benchmark::State& state) {
+    serverGlobalParams.clusterRole = ClusterRole::None;
+
+    auto serviceContext = ServiceContext::make();
+    const auto client = serviceContext->makeClient("test");
+    const auto opCtx = client->makeOperationContext();
+
+    const auto catalogCache = CatalogCacheMock::make();
+
+    for (auto keepRunning : state) {
+        benchmark::ClobberMemory();
+        ShardingWriteRouter writeRouter(opCtx.get(), kNss, catalogCache.get());
+        auto shardId = writeRouter.getReshardingDestinedRecipient(BSON("_id" << 0));
+        ASSERT(shardId == boost::none);
+    }
+}
+
+BENCHMARK(BM_InsertGetDestinedRecipient)
+    ->Range(1, 1 << 4)
+    ->ThreadRange(1, ProcessInfo::getNumAvailableCores());
+
+BENCHMARK(BM_UpdateGetDestinedRecipient)
+    ->Range(1, 1 << 4)
+    ->ThreadRange(1, ProcessInfo::getNumAvailableCores());
+
+BENCHMARK(BM_UnshardedDestinedRecipient)
+    ->Range(1, 1 << 4)
+    ->ThreadRange(1, ProcessInfo::getNumAvailableCores());
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index db171385b21..51a250bd53d 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -215,6 +215,7 @@ env.Library(
     source=[
         'sharding_test_fixture_common.cpp',
         'catalog_cache_loader_mock.cpp',
+        'catalog_cache_mock.cpp'
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/client/remote_command_targeter_mock',
diff --git a/src/mongo/s/catalog_cache.h b/src/mongo/s/catalog_cache.h
index 863ee2ec8cf..17e3b04d4c6 100644
--- a/src/mongo/s/catalog_cache.h
+++ b/src/mongo/s/catalog_cache.h
@@ -79,7 +79,7 @@ class CatalogCache {
 
 public:
     CatalogCache(ServiceContext* service, CatalogCacheLoader& cacheLoader);
-    ~CatalogCache();
+    virtual ~CatalogCache();
 
     /**
      * Blocking method that ensures the specified database is in the cache, loading it if necessary,
@@ -113,9 +113,9 @@ public:
      * guaranteed to never return StaleClusterTime, because the latest routing information should
      * always be available.
      */
-    StatusWith<ChunkManager> getCollectionRoutingInfo(OperationContext* opCtx,
-                                                      const NamespaceString& nss,
-                                                      bool allowLocks = false);
+    virtual StatusWith<ChunkManager> getCollectionRoutingInfo(OperationContext* opCtx,
+                                                              const NamespaceString& nss,
+                                                              bool allowLocks = false);
 
     /**
      * Same as getDatbase above, but in addition forces the database entry to be refreshed.
diff --git a/src/mongo/s/catalog_cache_mock.cpp b/src/mongo/s/catalog_cache_mock.cpp
new file mode 100644
index 00000000000..706b6a836d8
--- /dev/null
+++ b/src/mongo/s/catalog_cache_mock.cpp
@@ -0,0 +1,61 @@
+/**
+ *    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.
+ */
+
+#include "mongo/s/catalog_cache_mock.h"
+#include "mongo/db/service_context.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
+
+namespace mongo {
+
+const Status CatalogCacheMock::kChunkManagerInternalErrorStatus = {
+    ErrorCodes::InternalError, "Mocked catalog cache received unexpected chunks manager"};
+
+
+CatalogCacheMock::CatalogCacheMock(ServiceContext* serviceContext, CatalogCacheLoaderMock& loader)
+    : CatalogCache(serviceContext, loader) {}
+
+StatusWith<ChunkManager> CatalogCacheMock::getCollectionRoutingInfo(OperationContext* opCtx,
+                                                                    const NamespaceString& nss,
+                                                                    bool allowLocks) {
+    return _swChunkManagerReturnValue;
+}
+
+void CatalogCacheMock::setChunkManagerReturnValue(StatusWith<ChunkManager> statusWithChunks) {
+    _swChunkManagerReturnValue = statusWithChunks;
+}
+void CatalogCacheMock::clearChunkManagerReturnValue() {
+    _swChunkManagerReturnValue = kChunkManagerInternalErrorStatus;
+}
+
+std::unique_ptr<CatalogCacheMock> CatalogCacheMock::make() {
+    auto catalogCacheLoader = std::make_unique<CatalogCacheLoaderMock>();
+    auto serviceContext = ServiceContext::make();
+    return std::make_unique<CatalogCacheMock>(serviceContext.get(), *catalogCacheLoader);
+}
+}  // namespace mongo
diff --git a/src/mongo/s/catalog_cache_mock.h b/src/mongo/s/catalog_cache_mock.h
new file mode 100644
index 00000000000..306529b60d3
--- /dev/null
+++ b/src/mongo/s/catalog_cache_mock.h
@@ -0,0 +1,64 @@
+/**
+ *    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/s/catalog_cache.h"
+#include "mongo/s/catalog_cache_loader_mock.h"
+
+namespace mongo {
+
+/**
+ * Mocks the metadata refresh results with settable return values. The purpose of this class is to
+ * facilitate testing of classes that use a CatalogCacheLoader.
+ */
+class CatalogCacheMock final : public CatalogCache {
+    CatalogCacheMock(const CatalogCacheMock&) = delete;
+    CatalogCacheMock& operator=(const CatalogCacheMock&) = delete;
+
+public:
+    CatalogCacheMock(ServiceContext* context, CatalogCacheLoaderMock& loader);
+    ~CatalogCacheMock() = default;
+
+    StatusWith<ChunkManager> getCollectionRoutingInfo(OperationContext* opCtx,
+                                                      const NamespaceString& nss,
+                                                      bool allowLocks) override;
+
+    void setChunkManagerReturnValue(StatusWith<ChunkManager> statusWithChunks);
+    void clearChunkManagerReturnValue();
+
+    static std::unique_ptr<CatalogCacheMock> make();
+
+    static const Status kChunkManagerInternalErrorStatus;
+
+private:
+    StatusWith<ChunkManager> _swChunkManagerReturnValue{kChunkManagerInternalErrorStatus};
+};
+
+}  // namespace mongo