SERVER-58505 Add metric to serverStatus for number of load balanced connections

6 files changed, 99 insertions, 1 deletions

diff --git a/jstests/sharding/load_balancer_support/load_balanced_server_status_metrics.js b/jstests/sharding/load_balancer_support/load_balanced_server_status_metrics.js
new file mode 100644
index 00000000000..1e0f60ccfa5
--- /dev/null
+++ b/jstests/sharding/load_balancer_support/load_balanced_server_status_metrics.js
@@ -0,0 +1,68 @@
+/**
+ * @tags: [featureFlagLoadBalancer]
+ *
+ * Tests that load-balanced connections are reported correctly in server status metrics.
+ */
+
+(() => {
+    "use strict";
+
+    load('jstests/libs/fail_point_util.js');
+
+    const numConnections = 10;
+
+    function createTemporaryConnection(uri, dbName, collectionName) {
+        // Retry connecting until you are successful
+        var pollString = "var conn = null;" +
+            "assert.soon(function() {" +
+            "try { conn = new Mongo(\"" + uri + "\"); return conn" +
+            "} catch (x) {return false;}}, " +
+            "\"Timed out waiting for temporary connection to connect\", 30000, 5000);";
+        // Poll the signal collection until it is told to terminate.
+        pollString += "assert.soon(function() {" +
+            "return conn.getDB('" + dbName + "').getCollection('" + collectionName + "')" +
+            ".findOne().stop;}, \"Parallel shell never told to terminate\", 10 * 60000);";
+        return startParallelShell(pollString, null, true);
+    }
+
+    function waitForConnections(db, expected) {
+        assert.soon(() => admin.serverStatus().connections.loadBalanced == expected,
+                    () => "Incorrect number of load-balanced connections: expected " + expected +
+                        ", but serverStatus() reports " +
+                        admin.serverStatus().connections.loadBalanced,
+                    5 * 60000);
+    }
+
+    var st = new ShardingTest({shards: 1, mongos: 1});
+    let admin = st.s.getDB("admin");
+
+    assert(admin.adminCommand({getParameter: 1, featureFlagLoadBalancer: 1})
+               .featureFlagLoadBalancer.value,
+           'featureFlagLoadBalancer should be enabled for this test');
+    assert.commandWorked(
+        admin.adminCommand({configureFailPoint: 'clientIsFromLoadBalancer', mode: 'alwaysOn'}));
+
+    var uri = "mongodb://" + admin.getMongo().host + "/?loadBalanced=true";
+
+    var testDB = 'connectionsOpenedTest';
+    var signalCollection = 'keepRunning';
+
+    admin.getSiblingDB(testDB).dropDatabase();
+    admin.getSiblingDB(testDB).getCollection(signalCollection).insert({stop: false});
+
+    var connections = [];
+    for (var i = 0; i < numConnections; i++) {
+        connections.push(createTemporaryConnection(uri, testDB, signalCollection));
+        waitForConnections(admin, i + 1);
+    }
+
+    admin.getSiblingDB(testDB).getCollection(signalCollection).update({}, {$set: {stop: true}});
+    for (var i = 0; i < numConnections; i++) {
+        connections[i]();
+    }
+    waitForConnections(admin, 0);
+
+    assert.commandWorked(
+        admin.adminCommand({configureFailPoint: 'clientIsFromLoadBalancer', mode: 'off'}));
+    st.stop();
+})();
diff --git a/src/mongo/s/service_entry_point_mongos.cpp b/src/mongo/s/service_entry_point_mongos.cpp
index dc7a9319aa6..89e71ea5e05 100644
--- a/src/mongo/s/service_entry_point_mongos.cpp
+++ b/src/mongo/s/service_entry_point_mongos.cpp
@@ -208,8 +208,16 @@ Future<DbResponse> ServiceEntryPointMongos::handleRequest(OperationContext* opCt
     return hr->run();
 }
 
+void ServiceEntryPointMongos::onClientConnect(Client* client) {
+    if (load_balancer_support::isFromLoadBalancer(client)) {
+        _loadBalancedConnections.increment();
+    }
+}
+
 void ServiceEntryPointMongos::onClientDisconnect(Client* client) {
     if (load_balancer_support::isFromLoadBalancer(client)) {
+        _loadBalancedConnections.decrement();
+
         auto killerOperationContext = client->makeOperationContext();
 
         // Kill any cursors opened by the given Client.
@@ -247,4 +255,11 @@ void ServiceEntryPointMongos::onClientDisconnect(Client* client) {
         }
     }
 }
+
+void ServiceEntryPointMongos::appendStats(BSONObjBuilder* bob) const {
+    ServiceEntryPointImpl::appendStats(bob);
+    if (load_balancer_support::isEnabled()) {
+        bob->append("loadBalanced", _loadBalancedConnections);
+    }
+}
 }  // namespace mongo
diff --git a/src/mongo/s/service_entry_point_mongos.h b/src/mongo/s/service_entry_point_mongos.h
index 444f89eef4c..c821ab184ef 100644
--- a/src/mongo/s/service_entry_point_mongos.h
+++ b/src/mongo/s/service_entry_point_mongos.h
@@ -47,7 +47,13 @@ public:
     Future<DbResponse> handleRequest(OperationContext* opCtx,
                                      const Message& request) noexcept override;
 
+    void appendStats(BSONObjBuilder* bob) const override;
+
+    void onClientConnect(Client* client) override;
     void onClientDisconnect(Client* client) override;
+
+private:
+    Counter64 _loadBalancedConnections;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/transport/service_entry_point.h b/src/mongo/transport/service_entry_point.h
index 44786921474..bc67de9cd7b 100644
--- a/src/mongo/transport/service_entry_point.h
+++ b/src/mongo/transport/service_entry_point.h
@@ -104,6 +104,11 @@ public:
     virtual void onEndSession(const transport::SessionHandle&) {}
 
     /**
+     * Optional handler which is invoked after a client connects.
+     */
+    virtual void onClientConnect(Client* client) {}
+
+    /**
      * Optional handler which is invoked after a client disconnect. A client disconnect occurs when
      * the connection between the mongo process and client is closed for any reason, and is defined
      * by the destruction and cleanup of the ServiceStateMachine that manages the client.
diff --git a/src/mongo/transport/service_entry_point_impl.cpp b/src/mongo/transport/service_entry_point_impl.cpp
index 92d80da802e..1821eae0cfe 100644
--- a/src/mongo/transport/service_entry_point_impl.cpp
+++ b/src/mongo/transport/service_entry_point_impl.cpp
@@ -175,10 +175,14 @@ void ServiceEntryPointImpl::startSession(transport::SessionHandle session) {
             return boost::none;
         }
 
+        auto clientPtr = client.get();
         auto it = _sessions.emplace(_sessions.begin(), std::move(client));
+
         connectionCount = _sessions.size();
         _currentConnections.store(connectionCount);
         _createdConnections.addAndFetch(1);
+        onClientConnect(clientPtr);
+
         return it;
     }();
 
diff --git a/src/mongo/transport/service_entry_point_impl.h b/src/mongo/transport/service_entry_point_impl.h
index a8c1b080c7c..72f798908d5 100644
--- a/src/mongo/transport/service_entry_point_impl.h
+++ b/src/mongo/transport/service_entry_point_impl.h
@@ -74,7 +74,7 @@ public:
     bool shutdownAndWait(Milliseconds timeout);
     bool waitForNoSessions(Milliseconds timeout);
 
-    void appendStats(BSONObjBuilder* bob) const final;
+    void appendStats(BSONObjBuilder* bob) const override;
 
     size_t numOpenSessions() const final {
         return _currentConnections.load();