SERVER-44152: Pre-warm connection pools in mongosr4.2.17-rc0r4.2.17

9 files changed, 318 insertions, 0 deletions

diff --git a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
index 7f21f6aea90..0fca651f4a0 100644
--- a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
@@ -182,6 +182,8 @@ selector:
   - jstests/sharding/pipeline_length_limit.js
   # Remove once SERVER-55648 is backported to v4.0
   - jstests/sharding/retryable_mongos_write_errors.js
+  - jstests/sharding/warm_up_connection_pool.js
+
 executor:
   config:
     shell_options:
diff --git a/jstests/sharding/warm_up_connection_pool.js b/jstests/sharding/warm_up_connection_pool.js
new file mode 100644
index 00000000000..486fe4b1825
--- /dev/null
+++ b/jstests/sharding/warm_up_connection_pool.js
@@ -0,0 +1,124 @@
+/**
+ * This test checks that Mongos is pre-warming connections before accepting connections.
+ *
+ * This test requires users to persist across a restart.
+ * @tags: [requires_persistence]
+ */
+
+(function() {
+'use strict';
+load("jstests/replsets/rslib.js");
+
+function runTest(setParams, assertCheck, extraOptionsAction) {
+    jsTestLog('Starting next iteration');
+    const test = new ShardingTest({shards: 2, mongosOptions: setParams});
+    var db = test.getDB("test");
+
+    const shardCommand = {shardcollection: "test.foo", key: {num: 1}};
+
+    // shard
+    assert.commandWorked(test.s0.adminCommand({enablesharding: "test"}));
+    assert.commandWorked(test.s0.adminCommand(shardCommand));
+
+    var primary;
+    var mId;
+    if (extraOptionsAction !== undefined) {
+        const resp = extraOptionsAction(test);
+        primary = resp.connString;
+        mId = resp.nodeId;
+    }
+
+    test.restartMongos(0);
+
+    db = test.getDB("admin");
+    var connPoolStats = db.runCommand({connPoolStats: 1});
+    var shardList = db.runCommand({listShards: 1});
+
+    for (var shard in shardList["shards"]) {
+        var currentShard = getShardHostName(shardList, shard);
+        assertCheck(connPoolStats, currentShard, primary);
+    }
+
+    if (extraOptionsAction !== undefined) {
+        jsTestLog(`Restart ${mId}`);
+        test.rs0.restart(mId);
+        // Wait for mongos to recognize that the host is up.
+        awaitRSClientHosts(test.s0, test.rs0.nodes[mId], {ok: true});
+    }
+
+    test.stop();
+}
+
+function getShardHostName(shardlist, shard) {
+    return shardlist["shards"][shard]["host"].split("/")[1];
+}
+
+// Tests basic warm up of connection pool
+var testWarmUpParams = {};
+var testWarmUpAssertCheck = function(connPoolStats, currentShard) {
+    assert.soon(() => connPoolStats["hosts"][currentShard]["inUse"] +
+                        connPoolStats["hosts"][currentShard]["available"] +
+                        connPoolStats["hosts"][currentShard]["refreshing"] >=
+                    1,
+                `Connections ${tojson(connPoolStats)}`,
+                5 * 60 * 1000,
+                1000);
+};
+
+runTest(testWarmUpParams, testWarmUpAssertCheck);
+
+// Tests does not warm up connection pool when parameter is disabled
+var warmUpDisabledParams = {
+    setParameter: {warmMinConnectionsInShardingTaskExecutorPoolOnStartup: false}
+};
+var warmUpDisabledAssertCheck = function(connPoolStats, currentShard) {
+    assert.eq(null, connPoolStats["hosts"][currentShard]);
+};
+
+runTest(warmUpDisabledParams, warmUpDisabledAssertCheck);
+
+// Tests establishes more connections when parameter is set
+var biggerPoolSizeParams = {setParameter: {ShardingTaskExecutorPoolMinSize: 3}};
+var biggerPoolSizeAssertCheck = function(connPoolStats, currentShard) {
+    assert.soon(() => connPoolStats["hosts"][currentShard]["inUse"] +
+                        connPoolStats["hosts"][currentShard]["available"] +
+                        connPoolStats["hosts"][currentShard]["refreshing"] >=
+                    3,
+                `Connections ${tojson(connPoolStats)}`,
+                5 * 60 * 1000,
+                1000);
+};
+
+runTest(biggerPoolSizeParams, biggerPoolSizeAssertCheck);
+
+// Tests establishes connections it can and ignores the ones it can't
+var shutdownNodeParams = {};
+var shutdownNodeAssertCheck = function(connPoolStats, currentShard, primary) {
+    if (currentShard == primary) {
+        assert.soon(() => connPoolStats["hosts"][currentShard]["inUse"] +
+                            connPoolStats["hosts"][currentShard]["available"] ===
+                        0,
+                    `Connections ${tojson(connPoolStats)}`,
+                    5 * 60 * 1000,
+                    1000);
+    } else {
+        assert.soon(() => connPoolStats["hosts"][currentShard]["inUse"] +
+                            connPoolStats["hosts"][currentShard]["available"] +
+                            connPoolStats["hosts"][currentShard]["refreshing"] >=
+                        1,
+                    `Connections ${tojson(connPoolStats)}`,
+                    5 * 60 * 1000,
+                    1000);
+    }
+};
+var shutdownNodeExtraOptions = function(test) {
+    const nodeList = test.rs0.nodeList();
+    const master = test.rs0.getPrimary();
+    var mId = test.rs0.getNodeId(master);
+
+    test.rs0.stop(mId);
+    return {connString: nodeList[mId], nodeId: mId};
+};
+
+runTest(shutdownNodeParams, shutdownNodeAssertCheck, shutdownNodeExtraOptions);
+})();
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index 0376952c50e..8e71a31c4db 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -105,11 +105,13 @@ env.Library(
     source=[
         'sharding_initialization.cpp',
         'sharding_task_executor_pool_controller.cpp',
+        'pre_warm_connection_pool_impl.cpp',
         env.Idlc('sharding_task_executor_pool.idl')[0],
         'client/sharding_connection_hook.cpp',
         'client/sharding_network_connection_hook.cpp',
     ],
     LIBDEPS=[
+        '$BUILD_DIR/mongo/executor/async_multicaster',
         '$BUILD_DIR/mongo/s/catalog/dist_lock_catalog_impl',
         '$BUILD_DIR/mongo/s/catalog/replset_dist_lock_manager',
         '$BUILD_DIR/mongo/s/catalog/sharding_catalog_client_impl',
diff --git a/src/mongo/s/mongos_options.idl b/src/mongo/s/mongos_options.idl
index f92d2474fad..5c8033a24b0 100644
--- a/src/mongo/s/mongos_options.idl
+++ b/src/mongo/s/mongos_options.idl
@@ -31,10 +31,26 @@ global:
   cpp_namespace: "mongo"
   cpp_includes:
     - "mongo/s/mongos_options.h"
+    - "mongo/s/pre_warm_connection_pool_impl.h"
   configs:
     section: "Sharding options"
     source: [ yaml, cli, ini ]
 
+server_parameters: 
+  warmMinConnectionsInShardingTaskExecutorPoolOnStartup:
+    description: <-
+        Enables prewarming of the connection pool.
+    set_at: [ startup ]
+    cpp_varname: "gWarmMinConnectionsInShardingTaskExecutorPoolOnStartup"
+    default: true
+
+  warmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS:
+    description: <-
+        How long to wait for all hosts to have at least one connection.
+    set_at: [ startup ]
+    cpp_varname: "gWarmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS"
+    default: 2000 # 2secs
+
 configs:
   "sharding.configDB":
     description: >-
diff --git a/src/mongo/s/pre_warm_connection_pool_impl.cpp b/src/mongo/s/pre_warm_connection_pool_impl.cpp
new file mode 100644
index 00000000000..4721498be63
--- /dev/null
+++ b/src/mongo/s/pre_warm_connection_pool_impl.cpp
@@ -0,0 +1,37 @@
+/**
+ *    Copyright (C) 2019-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/pre_warm_connection_pool_impl.h"
+
+namespace mongo {
+
+bool gWarmMinConnectionsInShardingTaskExecutorPoolOnStartup = true;
+int gWarmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS = 2000;
+
+}  // namespace mongo
diff --git a/src/mongo/s/pre_warm_connection_pool_impl.h b/src/mongo/s/pre_warm_connection_pool_impl.h
new file mode 100644
index 00000000000..f5b4a5dfbe0
--- /dev/null
+++ b/src/mongo/s/pre_warm_connection_pool_impl.h
@@ -0,0 +1,41 @@
+/**
+ *    Copyright (C) 2019-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
+
+namespace mongo {
+/**
+ * Set parameters used to control whether or not the mongos attempts to warm up the connection
+ * pool on start up and for how long it should try.
+ */
+
+extern bool gWarmMinConnectionsInShardingTaskExecutorPoolOnStartup;
+extern int gWarmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS;
+
+}  // namespace mongo
diff --git a/src/mongo/s/server.cpp b/src/mongo/s/server.cpp
index c43f6528798..3ac128f4794 100644
--- a/src/mongo/s/server.cpp
+++ b/src/mongo/s/server.cpp
@@ -426,6 +426,11 @@ Status initializeSharding(OperationContext* opCtx) {
         return status;
     }
 
+    status = preWarmConnectionPool(opCtx);
+    if (!status.isOK()) {
+        return status;
+    }
+
     Grid::get(opCtx)->setShardingInitialized();
 
     return Status::OK();
diff --git a/src/mongo/s/sharding_initialization.cpp b/src/mongo/s/sharding_initialization.cpp
index 29da54a9a30..4dc7372960e 100644
--- a/src/mongo/s/sharding_initialization.cpp
+++ b/src/mongo/s/sharding_initialization.cpp
@@ -44,9 +44,12 @@
 #include "mongo/db/server_options.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/time_proof_service.h"
+#include "mongo/executor/async_multicaster.h"
 #include "mongo/executor/connection_pool.h"
+#include "mongo/executor/connection_pool_stats.h"
 #include "mongo/executor/network_interface_factory.h"
 #include "mongo/executor/network_interface_thread_pool.h"
+#include "mongo/executor/scoped_task_executor.h"
 #include "mongo/executor/task_executor.h"
 #include "mongo/executor/task_executor_pool.h"
 #include "mongo/executor/thread_pool_task_executor.h"
@@ -56,12 +59,14 @@
 #include "mongo/s/catalog/dist_lock_catalog_impl.h"
 #include "mongo/s/catalog/replset_dist_lock_manager.h"
 #include "mongo/s/catalog/sharding_catalog_client_impl.h"
+#include "mongo/s/catalog/type_shard.h"
 #include "mongo/s/catalog_cache.h"
 #include "mongo/s/client/num_hosts_targeted_metrics.h"
 #include "mongo/s/client/shard_factory.h"
 #include "mongo/s/client/sharding_network_connection_hook.h"
 #include "mongo/s/cluster_identity_loader.h"
 #include "mongo/s/grid.h"
+#include "mongo/s/pre_warm_connection_pool_impl.h"
 #include "mongo/s/query/cluster_cursor_manager.h"
 #include "mongo/s/sharding_task_executor.h"
 #include "mongo/s/sharding_task_executor_pool_controller.h"
@@ -141,6 +146,26 @@ std::unique_ptr<TaskExecutorPool> makeShardingTaskExecutorPool(
     return executorPool;
 }
 
+/**
+ * Uses an AsyncMulticaster to ping all of the hosts in order to establish
+ * ShardingTaskExecutorPoolMinSize connections. This does not wait
+ * for the connections to be established nor does it check how many were established.
+ */
+void preWarmConnections(OperationContext* opCtx, std::vector<HostAndPort> allHosts) {
+    auto const grid = Grid::get(opCtx);
+    auto arbi = grid->getExecutorPool()->getArbitraryExecutor();
+    auto executor = executor::ScopedTaskExecutor(arbi);
+    executor::AsyncMulticaster::Options options;
+
+    auto results =
+        executor::AsyncMulticaster(*executor, options)
+            .multicast(allHosts,
+                       "admin",
+                       BSON("ping" << 1),
+                       opCtx,
+                       Milliseconds(gWarmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS));
+}
+
 }  // namespace
 
 std::unique_ptr<executor::TaskExecutor> makeShardingTaskExecutor(
@@ -245,4 +270,64 @@ Status waitForShardRegistryReload(OperationContext* opCtx) {
     return {ErrorCodes::ShutdownInProgress, "aborting shard loading attempt"};
 }
 
+Status preWarmConnectionPool(OperationContext* opCtx) {
+    if (!gWarmMinConnectionsInShardingTaskExecutorPoolOnStartup) {
+        return Status::OK();
+    }
+
+    if (serverGlobalParams.clusterRole == ClusterRole::ConfigServer) {
+        log() << "Pre-warming pooled connections for config server is disabled";
+        return Status::OK();
+    }
+
+    // Should not be called by mongod
+    invariant(serverGlobalParams.clusterRole != ClusterRole::ShardServer);
+
+    Timer timer;
+    std::vector<HostAndPort> allHosts;
+    auto const grid = Grid::get(opCtx);
+    auto allShardsStatus =
+        grid->catalogClient()->getAllShards(opCtx, repl::ReadConcernLevel::kMajorityReadConcern);
+    if (!allShardsStatus.isOK()) {
+        return allShardsStatus.getStatus();
+    }
+    auto allShards = allShardsStatus.getValue().value;
+
+    for (auto& shard : allShards) {
+        auto connStrStatus = ConnectionString::parse(shard.getHost());
+        if (!connStrStatus.isOK()) {
+            return connStrStatus.getStatus();
+        }
+        auto connStr = connStrStatus.getValue();
+        for (const auto& hostEntry : connStr.getServers()) {
+            allHosts.push_back(hostEntry);
+        }
+    }
+
+    if (allHosts.empty()) {
+        log() << "No hosts found to pre-warm connections to";
+        return Status::OK();
+    }
+    log() << "Pre-warming connections to " << allHosts.size() << " hosts";
+
+    try {
+        opCtx->runWithDeadline(
+            opCtx->getServiceContext()->getPreciseClockSource()->now() +
+                Milliseconds(gWarmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS),
+            ErrorCodes::ExceededTimeLimit,
+            [&] { preWarmConnections(opCtx, allHosts); });
+    } catch (const ExceptionFor<ErrorCodes::ExceededTimeLimit>&) {
+        // if we've timed out, eat the exception and continue
+    } catch (const DBException& ex) {
+        warning() << "Connection pool pre-warm failure " << causedBy(ex.toStatus()) << ", timeMs "
+                  << timer.millis();
+        return ex.toStatus();
+    }
+
+    if (timer.millis() > gWarmMinConnectionsInShardingTaskExecutorPoolOnStartupWaitMS / 2) {
+        log() << "Slow connection pool pre-warm, timeMs " << timer.millis();
+    }
+    return Status::OK();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/sharding_initialization.h b/src/mongo/s/sharding_initialization.h
index 8f306dbde9a..56c5f652ddd 100644
--- a/src/mongo/s/sharding_initialization.h
+++ b/src/mongo/s/sharding_initialization.h
@@ -89,4 +89,10 @@ Status initializeGlobalShardingState(OperationContext* opCtx,
 
 Status waitForShardRegistryReload(OperationContext* opCtx);
 
+/**
+ * Warms up connections to shards with best effort strategy.
+ */
+
+Status preWarmConnectionPool(OperationContext* opCtx);
+
 }  // namespace mongo