SERVER-104413: admission rate limiter metrics (#36574)

GitOrigin-RevId: 13fd645d22cda2d4854f22e6697f60efdd6d0f5b

Author: dgoffredo

jstests/noPassthrough/connection_establishment_rate_limiting_client_disconnect.js

@@ -43,10 +43,10 @@ const testKillOnClientDisconnect = (conn) => {
     if (isLinux()) {
         assert.soon(() => checkLog.checkContainsOnceJson(
                         conn, 20883));  // Interrupted operation as its client disconnected
-        assert.soon(
-            () => (1 ==
-                   getConnectionStats(
-                       conn)["establishmentRateLimit"]["totalInterruptedDueToClientDisconnect"]));
+        assert.soon(() =>
+                        (1 ==
+                         getConnectionStats(
+                             conn)["establishmentRateLimit"]["interruptedDueToClientDisconnect"]));
     }
 };
 

jstests/noPassthrough/connection_establishment_rate_limiting_exemptions.js

@@ -40,11 +40,11 @@ const testExemptIPsFromRateLimit = (conn) => {
         return false;
     });
 
-    assert.soon(() => 1 == getConnectionStats(conn)["establishmentRateLimit"]["totalRejected"]);
+    assert.soon(() => 1 == getConnectionStats(conn)["establishmentRateLimit"]["rejected"]);
 
     // Connections over an exempted ip will succeed.
     assert(new Mongo(`mongodb://${exemptIP}:${conn.port}`));
-    assert.soon(() => getConnectionStats(conn)["establishmentRateLimit"]["totalExempted"] >= 1);
+    assert.soon(() => getConnectionStats(conn)["establishmentRateLimit"]["exempted"] >= 1);
 };
 
 const testExemptIPsFromRateLimitOpts = {

jstests/noPassthrough/connection_establishment_rate_limiting_proxy_protocol.js

@@ -58,9 +58,8 @@ rs.initiate();
         assert.neq(null, conn, 'Client was unable to connect');
     }
 
-    assert.soon(
-        () => getConnectionStats(rs.getPrimary())["establishmentRateLimit"]["totalExempted"] >=
-            numConnections);
+    assert.soon(() => getConnectionStats(rs.getPrimary())["establishmentRateLimit"]["exempted"] >=
+                    numConnections);
 
     proxy_server.stop();
 }
@@ -100,8 +99,8 @@ rs.getPrimary().adminCommand({
         return false;
     });
 
-    assert.soon(() => 1 ==
-                    getConnectionStats(rs.getPrimary())["establishmentRateLimit"]["totalRejected"]);
+    assert.soon(() =>
+                    1 == getConnectionStats(rs.getPrimary())["establishmentRateLimit"]["rejected"]);
 
     proxy_server.stop();
 }

jstests/noPassthrough/connection_establishment_rate_limiting_stats.js

@@ -10,20 +10,28 @@
 import {configureFailPoint} from "jstests/libs/fail_point_util.js";
 import {Thread} from "jstests/libs/parallelTester.js";
 import {
-    getConnectionStats,
+    getLimiterStats,
     runTestReplSet,
     runTestShardedCluster,
     runTestStandaloneParamsSetAtRuntime,
     runTestStandaloneParamsSetAtStartup
 } from "jstests/noPassthrough/libs/conn_establishment_rate_limiter_helpers.js";
 
+// Return whether the left Number or NumberLong is equal to the right Number or NumberLong.
+const equal = (left, right) => {
+    if (left instanceof NumberLong && right instanceof NumberLong) {
+        return left.compare(right) === 0;
+    }
+    return left == right;
+};
+
 const maxQueueSize = 3;
 // Hardcoded so that we can assert on "available" connections in serverStatus.
 const maxIncomingConnections = 1000;
 
 const testRateLimiterStats = (conn) => {
-    // Start maxQueueDepth + 3 threads that will all try to connect to the server. The rate limiter
-    // should allow maxQueueDepth connections to be queued, and the rest should be rejected.
+    // Start maxQueueSize + 3 threads that will all try to connect to the server. The rate limiter
+    // should allow maxQueueSize connections to be queued, and the rest should be rejected.
     let connDelayFailPoint = configureFailPoint(conn, 'hangInRateLimiter');
     const extraConns = 3;
     const threads = [];
@@ -41,15 +49,39 @@ const testRateLimiterStats = (conn) => {
     }
 
     assert.soon(() => {
-        const stats = getConnectionStats(conn);
-        // Queued connections should be counted in "queued", "totalCreated", and "available" stats.
-        return (stats["queued"] == maxQueueSize && stats["totalCreated"] >= maxQueueSize &&
-                stats["available"] <= maxIncomingConnections - maxQueueSize) &&
+        const {connections: cstats, ingressSessionEstablishmentQueues: qstats} =
+            getLimiterStats(conn, {log: false});
+
+        jsTestLog("stats: " + tojson({cstats, qstats}));
+
+        const checks = [
+            // Queued connections should be counted in "queuedForEstablishment", "totalCreated", and
+            // "available" stats.
+            () => equal(cstats["queuedForEstablishment"], maxQueueSize),
+            () => cstats["totalCreated"] >= maxQueueSize,
+            () => cstats["available"] <= maxIncomingConnections - maxQueueSize,
             // "extraConns" connections should be rejected because we create more than the rate
             // limit can handle-- they should be counted as rejected in the establishmentRateLimit
             // subsection and the overall "rejected" count.
-            (stats["rejected"] == extraConns &&
-             stats["establishmentRateLimit"]["totalRejected"] == extraConns);
+            () => equal(cstats["rejected"], extraConns),
+            () => equal(cstats["establishmentRateLimit"]["rejected"], extraConns),
+            // There is a correspondence between the connections stats and the session
+            // establishment queue stats, because they use the same underlying rate limiter.
+            () => equal(cstats["establishmentRateLimit"]["rejected"], qstats["rejectedAdmissions"]),
+            () => equal(cstats["queuedForEstablishment"],
+                        qstats["addedToQueue"] - qstats["removedFromQueue"]),
+            // Somebody either waited or was admitted immediately, so there is an average wait
+            // time.
+            () => qstats["averageTimeQueuedMicros"] >= 0
+        ];
+
+        // results :: {<body of check>: <boolean result of check>}
+        const results = Object.fromEntries(checks.map(check => [check.toString(), check()]));
+
+        jsTestLog("checks: " + tojson(results));
+
+        // "all of the checks passed"
+        return Object.values(results).every(result => result);
     });
 
     connDelayFailPoint.off();

jstests/noPassthrough/libs/conn_establishment_rate_limiter_helpers.js

@@ -1,8 +1,20 @@
+export const getLimiterStats = (conn, {log = false}) => {
+    const serverStatus = conn.adminCommand({serverStatus: 1});
+    assert(serverStatus, "Failed to get server status");
+    const result = {
+        connections: serverStatus.connections,
+        ingressSessionEstablishmentQueues: serverStatus.queues.ingressSessionEstablishment
+    };
+    if (log) {
+        jsTestLog("Limiter stats: " + tojson(result));
+    }
+    return result;
+};
+
 export const getConnectionStats = (conn) => {
-    const connStats = conn.adminCommand({serverStatus: 1})["connections"];
-    assert.neq(null, connStats, "Failed to get connection stats");
-    jsTestLog("Connection stats: " + tojson(connStats));
-    return connStats;
+    const {connections} = getLimiterStats(conn, {log: true});
+    assert(connections, "Failed to get connection stats");
+    return connections;
 };
 
 export const runTestStandaloneParamsSetAtStartup = (setParams, testCase) => {

src/mongo/db/admission/SConscript

@@ -117,6 +117,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/commands/server_status_core',
         '$BUILD_DIR/mongo/db/server_feature_flags',
         '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/transport/transport_layer_common',
         'ingress_admission_control',
         'ticketholder_manager',
     ],

src/mongo/db/admission/queues_server_status_section.cpp

@@ -30,9 +30,9 @@
 #include "mongo/db/admission/ingress_admission_control_gen.h"
 #include "mongo/db/admission/ingress_admission_controller.h"
 #include "mongo/db/admission/ticketholder_manager.h"
-#include "mongo/db/server_feature_flags_gen.h"
-
 #include "mongo/db/commands/server_status.h"
+#include "mongo/transport/session_establishment_rate_limiter.h"
+#include "mongo/transport/transport_layer.h"
 
 namespace mongo {
 namespace admission {
@@ -48,27 +48,37 @@ class AdmissionServerStatusSection : public ServerStatusSection {
 
     BSONObj generateSection(OperationContext* opCtx,
                             const BSONElement& configElement) const override {
+        const ClusterRole role = opCtx->getService()->role();
+
         BSONObjBuilder admissionBuilder;
         auto ticketHolderManager = TicketHolderManager::get(opCtx->getServiceContext());
-        if (ticketHolderManager) {
+        if (ticketHolderManager && role.has(ClusterRole::ShardServer)) {
             BSONObjBuilder executionBuilder(admissionBuilder.subobjStart("execution"));
             ticketHolderManager->appendStats(executionBuilder);
             executionBuilder.done();
         }
 
-        if (gIngressAdmissionControlEnabled.load()) {
+        if (gIngressAdmissionControlEnabled.load() && role.has(ClusterRole::ShardServer)) {
             BSONObjBuilder ingressBuilder(admissionBuilder.subobjStart("ingress"));
             auto& controller = IngressAdmissionController::get(opCtx);
             controller.appendStats(ingressBuilder);
             ingressBuilder.done();
         }
 
+        if (auto* limiter =
+                transport::SessionEstablishmentRateLimiter::get(*opCtx->getServiceContext())) {
+            BSONObjBuilder ingressSessionEstablishmentBuilder(
+                admissionBuilder.subobjStart("ingressSessionEstablishment"));
+            limiter->appendStatsQueues(&ingressSessionEstablishmentBuilder);
+            ingressSessionEstablishmentBuilder.done();
+        }
+
         return admissionBuilder.obj();
     }
 };
 
 auto& admissionSection =
-    *ServerStatusSectionBuilder<AdmissionServerStatusSection>("queues").forShard();
+    *ServerStatusSectionBuilder<AdmissionServerStatusSection>("queues").forShard().forRouter();
 
 }  // namespace
 }  // namespace admission

src/mongo/db/admission/rate_limiter.cpp

@@ -32,6 +32,8 @@
 #include <folly/TokenBucket.h>
 
 #include "mongo/logv2/log.h"
+#include "mongo/util/duration.h"
+#include "mongo/util/scopeguard.h"
 
 #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kDefault
 
@@ -46,37 +48,42 @@ Milliseconds doubleToMillis(double t) {
 }  // namespace
 
 struct RateLimiter::RateLimiterPrivate {
-    RateLimiterPrivate(
-        double r, double b, int64_t m, std::string n, std::unique_ptr<RateLimiter::Stats> s)
-        : tokenBucket{r, b}, maxQueueDepth(m), name(std::move(n)), stats(std::move(s)) {}
+    RateLimiterPrivate(double r, double b, int64_t m, std::string n)
+        // Initialize the token bucket with one "burst" of tokens. The third parameter to
+        // tokenBucket's constructor ("zeroTime") is interpreted as a number of seconds from the
+        // epoch of the clock used by the token bucket. The clock is
+        // `std::chrono::steady_clock`, whose epoch is unspecified but is usually the boot time
+        // of the machine. Rather than have an initial accumulation of tokens based on some
+        // unknown point in the past, set the zero time to a known time in the past: enough time
+        // for burst size (b) tokens to have accumulated.
+        : tokenBucket{r, b, folly::TokenBucket::defaultClockNow() - b / r},
+          maxQueueDepth(m),
+          queued(0),
+          name(std::move(n)) {}
 
     WriteRarelyRWMutex rwMutex;
     folly::TokenBucket tokenBucket;
 
+    Stats stats;
+
     Atomic<int64_t> maxQueueDepth;
-    Atomic<int64_t> numWaiters;
+    Atomic<int64_t> queued;
 
     std::string name;
 
-    /**
-     * Users may inherit from the RateLimiter::Stats class and define their own stats to track.
-     */
-    std::unique_ptr<Stats> stats;
-
-    Status rejectIfOverQueueLimit(double nWaiters) {
-        auto maxDepth = maxQueueDepth.loadRelaxed();
-        if (MONGO_unlikely(nWaiters >= maxDepth)) {
-            {
-                auto lk = rwMutex.readLock();
-                tokenBucket.returnTokens(1.0);
+    Status enqueue() {
+        const auto maxDepth = maxQueueDepth.loadRelaxed();
+        int64_t expected = queued.load();
+        do {
+            if (expected >= maxDepth) {
+                return Status(ErrorCodes::TemporarilyUnavailable,
+                              fmt::format("RateLimiter queue depth has exceeded the maxQueueDepth. "
+                                          "numWaiters={}; maxQueueDepth={}; rateLimiterName={}",
+                                          expected,
+                                          maxDepth,
+                                          name));
             }
-            return Status(ErrorCodes::TemporarilyUnavailable,
-                          fmt::format("RateLimiter queue depth has exceeded the maxQueueDepth. "
-                                      "numWaiters={}; maxQueueDepth={}; rateLimiterName={}",
-                                      nWaiters,
-                                      maxDepth,
-                                      name));
-        }
+        } while (!queued.compareAndSwap(&expected, expected + 1));
 
         return Status::OK();
     }
@@ -85,20 +92,21 @@ struct RateLimiter::RateLimiterPrivate {
 RateLimiter::RateLimiter(double refreshRatePerSec,
                          double burstSize,
                          int64_t maxQueueDepth,
-                         std::string name,
-                         std::unique_ptr<Stats> stats) {
+                         std::string name) {
     uassert(ErrorCodes::InvalidOptions,
             fmt::format("burstSize cannot be less than 1.0. burstSize={}; rateLimiterName={}",
                         burstSize,
                         name),
             burstSize >= 1.0);
     _impl = std::make_unique<RateLimiterPrivate>(
-        refreshRatePerSec, burstSize, maxQueueDepth, std::move(name), std::move(stats));
+        refreshRatePerSec, burstSize, maxQueueDepth, std::move(name));
 }
 
 RateLimiter::~RateLimiter() = default;
 
 Status RateLimiter::acquireToken(OperationContext* opCtx) {
+    _impl->stats.attemptedAdmissions.incrementRelaxed();
+
     // The consumeWithBorrowNonBlocking API consumes a token (possibly leading to a negative
     // bucket balance), and returns how long the consumer should nap until their token
     // reservation becomes valid.
@@ -111,12 +119,15 @@ Status RateLimiter::acquireToken(OperationContext* opCtx) {
     }
 
     if (auto napTime = doubleToMillis(waitForTokenSecs); napTime > Milliseconds{0}) {
-        auto nWaiters = _impl->numWaiters.fetchAndAdd(1);
-        ON_BLOCK_EXIT([&] { _impl->numWaiters.fetchAndSubtract(1); });
-        if (auto s = _impl->rejectIfOverQueueLimit(nWaiters); !s.isOK()) {
-            return s;
+        if (auto status = _impl->enqueue(); !status.isOK()) {
+            _impl->stats.rejectedAdmissions.incrementRelaxed();
+            return status;
         }
-
+        _impl->stats.addedToQueue.incrementRelaxed();
+        ON_BLOCK_EXIT([&] {
+            _impl->stats.removedFromQueue.incrementRelaxed();
+            _impl->queued.fetchAndSubtract(1);
+        });
         try {
             LOGV2_DEBUG(10550200,
                         4,
@@ -125,6 +136,7 @@ Status RateLimiter::acquireToken(OperationContext* opCtx) {
                         "napTimeMillis"_attr = napTime.toString());
             opCtx->sleepFor(napTime);
         } catch (const DBException& e) {
+            _impl->stats.interruptedInQueue.incrementRelaxed();
             LOGV2_DEBUG(10440800,
                         4,
                         "Interrupted while waiting in rate limiter queue",
@@ -139,15 +151,14 @@ Status RateLimiter::acquireToken(OperationContext* opCtx) {
                             _impl->name));
         }
     }
-    return Status::OK();
-}
 
-RateLimiter::Stats* RateLimiter::getRateLimiterStats() {
-    return _impl->stats.get();
+    _impl->stats.successfulAdmissions.incrementRelaxed();
+    _impl->stats.averageTimeQueuedMicros.addSample(waitForTokenSecs * 1'000'000);
+    return Status::OK();
 }
 
-int64_t RateLimiter::getNumWaiters() {
-    return _impl->numWaiters.load();
+void RateLimiter::recordExemption() {
+    _impl->stats.exemptedAdmissions.incrementRelaxed();
 }
 
 void RateLimiter::setRefreshRatePerSec(double refreshRatePerSec) {
@@ -169,4 +180,33 @@ void RateLimiter::setBurstSize(double burstSize) {
 void RateLimiter::setMaxQueueDepth(int64_t maxQueueDepth) {
     _impl->maxQueueDepth.storeRelaxed(maxQueueDepth);
 }
+
+const RateLimiter::Stats& RateLimiter::stats() const {
+    return _impl->stats;
+}
+
+void RateLimiter::appendStats(BSONObjBuilder* bob) const {
+    invariant(bob);
+    bob->append("addedToQueue", stats().addedToQueue.get());
+    bob->append("removedFromQueue", stats().removedFromQueue.get());
+    bob->append("interruptedInQueue", stats().interruptedInQueue.get());
+    bob->append("rejectedAdmissions", stats().rejectedAdmissions.get());
+    bob->append("exemptedAdmissions", stats().exemptedAdmissions.get());
+    bob->append("successfulAdmissions", stats().successfulAdmissions.get());
+    bob->append("attemptedAdmissions", stats().attemptedAdmissions.get());
+    if (const auto avg = stats().averageTimeQueuedMicros.get()) {
+        bob->append("averageTimeQueuedMicros", *avg);
+    }
+    bob->append("totalAvailableTokens", tokensAvailable());
+}
+
+double RateLimiter::tokensAvailable() const {
+    auto lk = _impl->rwMutex.readLock();
+    return _impl->tokenBucket.available();
+}
+
+int64_t RateLimiter::queued() const {
+    return _impl->queued.load();
+}
+
 }  // namespace mongo::admission