SERVER-58502 mongos hello loadBalanced option

(cherry picked from commit 4a9c2e6f412619a50770ba867c11a6526840558f)

10 files changed, 411 insertions, 19 deletions

diff --git a/jstests/sharding/load_balancer_support/load_balancer_hello.js b/jstests/sharding/load_balancer_support/load_balancer_hello.js
new file mode 100644
index 00000000000..cede53cd8a8
--- /dev/null
+++ b/jstests/sharding/load_balancer_support/load_balancer_hello.js
@@ -0,0 +1,95 @@
+/**
+ * @tags: [featureFlagLoadBalancer]
+ *
+ * Test the extension to the mongos `hello` command by which clients
+ * that have arrived through a load balancer affirm that they are
+ * compatible with the way mongos handles load-balanced clients.
+ * See `src/mongo/s/load_balancing_support.h`.
+ */
+(() => {
+    "use strict";
+
+    load('jstests/libs/fail_point_util.js');
+
+    /**
+     * The whole ShardingTest is restarted just to get a fresh connection.
+     * Obviously this could be accomplished much more efficiently.
+     */
+    var runInShardingTest = (func) => {
+        var st = new ShardingTest({shards: 1, mongos: 1});
+        try {
+            func(st.s0.getDB("admin"));
+        } finally {
+            st.stop();
+        }
+    };
+
+    var doHello = (admin, {lbConnection, lbHello, expectServiceId}) => {
+        if (lbConnection)
+            assert.commandWorked(admin.adminCommand(
+                {configureFailPoint: 'clientIsFromLoadBalancer', mode: 'alwaysOn'}));
+        try {
+            var helloDoc = {};
+            if (lbHello)
+                helloDoc['loadBalanced'] = true;
+            return admin.runCommand("hello", helloDoc);
+        } finally {
+            assert.commandWorked(
+                admin.adminCommand({configureFailPoint: 'clientIsFromLoadBalancer', mode: 'off'}));
+        }
+    };
+
+    var assertServiceId = (res) => {
+        assert.commandWorked(res);
+        assert(res.hasOwnProperty("serviceId"),
+               "serviceId missing from hello response:" + tojson(res));
+        assert(res.serviceId.isObjectId, "res.serviceId = " + tojson(res.serviceId));
+    };
+
+    var assertNoServiceId = (res) => {
+        assert.commandWorked(res);
+        assert(!res.hasOwnProperty("serviceId"), "res.serviceId = " + tojson(res.serviceId));
+    };
+
+    /*
+     * The ordinary baseline non-load-balanced case.
+     */
+    runInShardingTest((admin) => {
+        // Before getting started, confirm that the feature is enabled.
+        assert(admin.adminCommand({getParameter: 1, featureFlagLoadBalancer: 1})
+                   .featureFlagLoadBalancer.value,
+               'featureFlagLoadBalancer should be enabled for this test');
+
+        jsTestLog("Initial hello command");
+        assertNoServiceId(doHello(admin, {}));
+        jsTestLog("Non-initial hello command");
+        assertNoServiceId(doHello(admin, {}));
+    });
+
+    /*
+     * Good case: client arrived via load balancer, and load balancing declared by client.
+     * The load balancing `serviceId` reporting applies only to the initial hello.
+     * The `loadBalanced` option is ignored in subsequent `hello` commands.
+     */
+    runInShardingTest((admin) => {
+        assertServiceId(doHello(admin, {lbConnection: true, lbHello: true}));
+        assertNoServiceId(doHello(admin, {lbConnection: true, lbHello: true}));
+    });
+
+    /*
+     * Client did not arrive via load-balancer, but load balancing support declared by client.
+     * We tolerate the `isLoadBalanced` option but ignore it.
+     */
+    runInShardingTest((admin) => {
+        assertNoServiceId(doHello(admin, {lbHello: true}));
+    });
+
+    /*
+     * Client arrived via load balancer, but load balancing support was not declared by client.
+     * This is an error that should result in a disconnection.
+     */
+    runInShardingTest((admin) => {
+        var res = doHello(admin, {lbConnection: true});
+        assert.commandFailedWithCode(res, ErrorCodes.LoadBalancerSupportMismatch);
+    });
+})();
diff --git a/src/mongo/db/client.cpp b/src/mongo/db/client.cpp
index 7e162fb4335..473150cfaa8 100644
--- a/src/mongo/db/client.cpp
+++ b/src/mongo/db/client.cpp
@@ -48,7 +48,6 @@
 #include "mongo/stdx/thread.h"
 #include "mongo/util/concurrency/thread_name.h"
 #include "mongo/util/exit.h"
-#include "mongo/util/fail_point.h"
 #include "mongo/util/str.h"
 
 namespace mongo {
@@ -56,8 +55,6 @@ namespace mongo {
 namespace {
 thread_local ServiceContext::UniqueClient currentClient;
 
-MONGO_FAIL_POINT_DEFINE(clientIsFromLoadBalancer);
-
 void invariantNoCurrentClient() {
     invariant(!haveClient(),
               str::stream() << "Already have client on this thread: "  //
@@ -204,11 +201,4 @@ Client* ThreadClient::get() const {
     return &cc();
 }
 
-bool Client::isFromLoadBalancer() const {
-    if (MONGO_unlikely(clientIsFromLoadBalancer.shouldFail())) {
-        return true;
-    }
-    return _isFromLoadBalancer;
-}
-
 }  // namespace mongo
diff --git a/src/mongo/db/client.h b/src/mongo/db/client.h
index dd7cf1107a3..67b8ed959d5 100644
--- a/src/mongo/db/client.h
+++ b/src/mongo/db/client.h
@@ -260,11 +260,6 @@ public:
         _supportsHello = newVal;
     }
 
-    /**
-     * Returns whether the connection was established through a load balancer.
-     */
-    bool isFromLoadBalancer() const;
-
 private:
     friend class ServiceContext;
     friend class ThreadClient;
@@ -309,9 +304,6 @@ private:
     bool _supportsHello = false;
 
     UUID _uuid;
-
-    // Whether the connection was established through a load balancer
-    const bool _isFromLoadBalancer = false;
 };
 
 /**
diff --git a/src/mongo/db/repl/hello.idl b/src/mongo/db/repl/hello.idl
index a3e7392585b..90e7c329bc5 100644
--- a/src/mongo/db/repl/hello.idl
+++ b/src/mongo/db/repl/hello.idl
@@ -121,6 +121,9 @@ structs:
             msg:
                 type: string
                 optional: true
+            serviceId:
+                type: objectid
+                optional: true
             ##
             ## ReplicationInfo
             ##
@@ -252,3 +255,6 @@ commands:
             speculativeAuthenticate:
                 type: object
                 optional: true
+            loadBalanced:
+                type: bool
+                optional: true
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index cb79bf17c1f..98e81ffb1be 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -131,6 +131,30 @@ env.Library(
 )
 
 env.Library(
+    target="load_balancer_feature_flag",
+    source=[
+        'load_balancer_feature_flag.idl'
+    ],
+    LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/idl/feature_flag',
+    ]
+)
+
+env.Library(
+    target="load_balancer_support",
+    source=[
+        'load_balancer_support.cpp',
+    ],
+    LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/db/server_options_core',
+        '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/idl/feature_flag',
+        "load_balancer_feature_flag",
+    ],
+)
+
+# Sharding code needed by both mongos and mongod.
+env.Library(
     target='common_s',
     source=[
         'cannot_implicitly_create_collection_info.cpp',
@@ -149,7 +173,6 @@ env.Library(
         'chunk_version.idl',
         'database_version.cpp',
         'database_version.idl',
-        'load_balancer_feature_flag.idl',
         'mongod_and_mongos_server_parameters.idl',
         'request_types/abort_reshard_collection.idl',
         'request_types/add_shard_request_type.cpp',
@@ -208,6 +231,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/timeseries/timeseries_options',
         '$BUILD_DIR/mongo/idl/feature_flag',
         '$BUILD_DIR/mongo/idl/idl_parser',
+        "load_balancer_feature_flag",
     ]
 )
 
@@ -614,6 +638,7 @@ env.CppUnitTest(
         'comparable_chunk_version_test.cpp',
         'comparable_database_version_test.cpp',
         'hedge_options_util_test.cpp',
+        'load_balancer_support_test.cpp',
         'mock_ns_targeter.cpp',
         'mongos_topology_coordinator_test.cpp',
         'request_types/add_shard_request_test.cpp',
@@ -653,6 +678,7 @@ env.CppUnitTest(
         'chunk_writes_tracker',
         'cluster_last_error_info',
         'coreshard',
+        'load_balancer_support',
         'mongos_topology_coordinator',
         'sessions_collection_sharded',
         'sharding_api',
diff --git a/src/mongo/s/commands/SConscript b/src/mongo/s/commands/SConscript
index 0fb51bc21c6..3cbd37bdf44 100644
--- a/src/mongo/s/commands/SConscript
+++ b/src/mongo/s/commands/SConscript
@@ -156,6 +156,7 @@ env.Library(
         '$BUILD_DIR/mongo/rpc/client_metadata',
         '$BUILD_DIR/mongo/rpc/rewrite_state_change_errors',
         '$BUILD_DIR/mongo/s/cluster_last_error_info',
+        '$BUILD_DIR/mongo/s/load_balancer_support',
         '$BUILD_DIR/mongo/s/mongos_topology_coordinator',
         '$BUILD_DIR/mongo/s/query/cluster_aggregate',
         '$BUILD_DIR/mongo/s/query/cluster_client_cursor',
diff --git a/src/mongo/s/commands/cluster_hello_cmd.cpp b/src/mongo/s/commands/cluster_hello_cmd.cpp
index a2b404b1513..4e2014f1058 100644
--- a/src/mongo/s/commands/cluster_hello_cmd.cpp
+++ b/src/mongo/s/commands/cluster_hello_cmd.cpp
@@ -47,6 +47,7 @@
 #include "mongo/rpc/metadata/client_metadata.h"
 #include "mongo/rpc/rewrite_state_change_errors.h"
 #include "mongo/rpc/topology_version_gen.h"
+#include "mongo/s/load_balancer_support.h"
 #include "mongo/s/mongos_topology_coordinator.h"
 #include "mongo/transport/message_compressor_manager.h"
 #include "mongo/util/net/socket_utils.h"
@@ -168,6 +169,8 @@ public:
         // The hello response always includes a topologyVersion.
         auto currentMongosTopologyVersion = mongosHelloResponse->getTopologyVersion();
 
+        load_balancer_support::handleHello(opCtx, &result, cmd.getLoadBalanced().value_or(false));
+
         // Try to parse the optional 'helloOk' field. On mongos, if we see this field, we will
         // respond with helloOk: true so the client knows that it can continue to send the hello
         // command to mongos.
diff --git a/src/mongo/s/load_balancer_support.cpp b/src/mongo/s/load_balancer_support.cpp
new file mode 100644
index 00000000000..b1ebf6bb6d3
--- /dev/null
+++ b/src/mongo/s/load_balancer_support.cpp
@@ -0,0 +1,118 @@
+/**
+ *    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/load_balancer_support.h"
+
+#include <memory>
+
+#include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/bson/oid.h"
+#include "mongo/db/client.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/repl/hello_gen.h"
+#include "mongo/db/service_context.h"
+#include "mongo/s/load_balancer_feature_flag_gen.h"
+#include "mongo/util/assert_util.h"
+#include "mongo/util/fail_point.h"
+
+namespace mongo::load_balancer_support {
+namespace {
+
+MONGO_FAIL_POINT_DEFINE(clientIsFromLoadBalancer);
+
+bool featureEnabled() {
+    return feature_flags::gFeatureFlagLoadBalancer.isEnabled(
+        serverGlobalParams.featureCompatibility);
+}
+
+struct PerService {
+    /**
+     * When a client reaches a mongos through a load balancer, the `serviceId`
+     * identifies the mongos to which it is connected. It persists through the
+     * lifespan of the service context.
+     */
+    OID serviceId = OID::gen();
+};
+
+class PerClient {
+public:
+    bool isFromLoadBalancer() const {
+        if (MONGO_unlikely(clientIsFromLoadBalancer.shouldFail())) {
+            return true;
+        }
+        return _isFromLoadBalancer;
+    }
+
+    void setIsFromLoadBalancer() {
+        _isFromLoadBalancer = true;
+    }
+
+    bool didHello() const {
+        return _didHello;
+    }
+
+    void setDidHello() {
+        _didHello = true;
+    }
+
+private:
+    /** True if the connection was established through a load balancer. */
+    bool _isFromLoadBalancer = false;
+
+    /** True after we send this client a hello reply. */
+    bool _didHello = false;
+};
+
+const auto getPerServiceState = ServiceContext::declareDecoration<PerService>();
+const auto getPerClientState = Client::declareDecoration<PerClient>();
+}  // namespace
+
+void setClientIsFromLoadBalancer(Client* client) {
+    if (!featureEnabled())
+        return;
+    auto& perClient = getPerClientState(client);
+    perClient.setIsFromLoadBalancer();
+}
+
+void handleHello(OperationContext* opCtx, BSONObjBuilder* result, bool helloHasLoadBalancedOption) {
+    if (!featureEnabled())
+        return;
+    auto& perClient = getPerClientState(opCtx->getClient());
+    if (perClient.didHello() || !perClient.isFromLoadBalancer())
+        return;
+
+    uassert(ErrorCodes::LoadBalancerSupportMismatch,
+            "The server is being accessed through a load balancer, but "
+            "this driver does not have load balancing enabled",
+            helloHasLoadBalancedOption);
+    result->append(HelloCommandReply::kServiceIdFieldName,
+                   getPerServiceState(opCtx->getServiceContext()).serviceId);
+    perClient.setDidHello();
+}
+
+}  // namespace mongo::load_balancer_support
diff --git a/src/mongo/s/load_balancer_support.h b/src/mongo/s/load_balancer_support.h
new file mode 100644
index 00000000000..c3d46d65823
--- /dev/null
+++ b/src/mongo/s/load_balancer_support.h
@@ -0,0 +1,57 @@
+/**
+ *    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/bson/bsonobjbuilder.h"
+#include "mongo/db/client.h"
+#include "mongo/db/operation_context.h"
+
+namespace mongo::load_balancer_support {
+
+/**
+ * When a connection is made, we identify whether it came in through a load
+ * balancer. We associate this information with the `client`.
+ */
+void setClientIsFromLoadBalancer(Client* client);
+
+/**
+ * Helper for handling the `hello` command on mongos.
+ * `helloHasLoadBalancedOption` must be true if the hello command had the
+ * `loadBalanced` option present and set to `true`.
+ *
+ * For only the initial `hello` command, we respond to a `loadBalanced: true`
+ * option by including the `serviceId` of this mongos in the hello reply.
+ *
+ * A connection marked as having come in through a load balancer must confirm
+ * that it is using a load-balancer-aware driver by setting the `loadBalanced:
+ * true` option in its first `hello` command. Otherwise, this function will
+ * throw with `ErrorCodes::LoadBalancerSupportMismatch`.
+ */
+void handleHello(OperationContext* opCtx, BSONObjBuilder* result, bool helloHasLoadBalancedOption);
+
+}  // namespace mongo::load_balancer_support
diff --git a/src/mongo/s/load_balancer_support_test.cpp b/src/mongo/s/load_balancer_support_test.cpp
new file mode 100644
index 00000000000..cb3e359c612
--- /dev/null
+++ b/src/mongo/s/load_balancer_support_test.cpp
@@ -0,0 +1,104 @@
+/**
+ *    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.
+ */
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest
+
+#include "mongo/s/load_balancer_support.h"
+
+#include "mongo/bson/json.h"
+#include "mongo/db/service_context_test_fixture.h"
+#include "mongo/idl/server_parameter_test_util.h"
+#include "mongo/logv2/log.h"
+#include "mongo/s/load_balancer_feature_flag_gen.h"
+#include "mongo/unittest/assert_that.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/fail_point.h"
+
+namespace mongo {
+namespace {
+
+using namespace unittest::match;
+
+class LoadBalancerSupportTest : public ServiceContextTest {
+public:
+    using ServiceContextTest::ServiceContextTest;
+
+    BSONObj doHello(bool lbOption) {
+        BSONObjBuilder bob;
+        load_balancer_support::handleHello(&*makeOperationContext(), &bob, lbOption);
+        return bob.obj();
+    }
+
+    struct HasServiceId : Matcher {
+        std::string describe() const {
+            return "HasServiceId";
+        }
+        MatchResult match(BSONObj obj) const {
+            bool ok = obj.hasElement("serviceId");
+            std::string msg;
+            if (!ok)
+                msg = tojson(obj);
+            return {ok, msg};
+        }
+    };
+
+    FailPointEnableBlock simulateLoadBalancerConnection() const {
+        return FailPointEnableBlock("clientIsFromLoadBalancer");
+    }
+
+    RAIIServerParameterControllerForTest featureEnabler{"featureFlagLoadBalancer", true};
+};
+
+TEST_F(LoadBalancerSupportTest, HelloNormalClientNoOption) {
+    ASSERT_THAT(doHello(false), Not(HasServiceId()));
+}
+
+TEST_F(LoadBalancerSupportTest, HelloNormalClientGivesOption) {
+    ASSERT_THAT(doHello(true), Not(HasServiceId()));
+}
+
+TEST_F(LoadBalancerSupportTest, HelloLoadBalancedClientNoOption) {
+    auto simLB = simulateLoadBalancerConnection();
+    try {
+        doHello(false);
+        FAIL("Expected to throw");
+    } catch (const DBException& ex) {
+        ASSERT_THAT(ex.toStatus(),
+                    StatusIs(Eq(ErrorCodes::LoadBalancerSupportMismatch),
+                             ContainsRegex("load balancer.*but.*driver")));
+    }
+}
+
+TEST_F(LoadBalancerSupportTest, HelloLoadBalancedClientGivesOption) {
+    auto simLB = simulateLoadBalancerConnection();
+    ASSERT_THAT(doHello(true), HasServiceId());
+    ASSERT_THAT(doHello(true), Not(HasServiceId())) << "only first hello is special";
+}
+
+}  // namespace
+}  // namespace mongo