SERVER-106149 RateLimiter should return tokens to bucket upon rejection (#37246)

GitOrigin-RevId: 53e95641575e220ec4b3ea2525ec9cf9cb3f48a5

Author: erin2722

src/mongo/db/admission/rate_limiter.cpp

@@ -32,6 +32,7 @@
 #include <folly/TokenBucket.h>
 
 #include "mongo/logv2/log.h"
+#include "mongo/logv2/log_severity_suppressor.h"
 #include "mongo/util/duration.h"
 #include "mongo/util/scopeguard.h"
 
@@ -123,6 +124,10 @@ Status RateLimiter::acquireToken(OperationContext* opCtx) {
         // don't advance the mock clock before the sleep deadline is calculated.
         Date_t deadline = opCtx->getServiceContext()->getPreciseClockSource()->now() + napTime;
         if (auto status = _impl->enqueue(); !status.isOK()) {
+            {
+                auto lk = _impl->rwMutex.readLock();
+                _impl->tokenBucket.returnTokens(1.0);
+            }
             _impl->stats.rejectedAdmissions.incrementRelaxed();
             return status;
         }
@@ -157,6 +162,7 @@ Status RateLimiter::acquireToken(OperationContext* opCtx) {
 
     _impl->stats.successfulAdmissions.incrementRelaxed();
     _impl->stats.averageTimeQueuedMicros.addSample(waitForTokenSecs * 1'000'000);
+
     return Status::OK();
 }
 
@@ -206,6 +212,11 @@ double RateLimiter::tokensAvailable() const {
     return _impl->tokenBucket.available();
 }
 
+double RateLimiter::tokenBalance() const {
+    auto lk = _impl->rwMutex.readLock();
+    return _impl->tokenBucket.balance();
+}
+
 int64_t RateLimiter::queued() const {
     return _impl->queued.load();
 }

src/mongo/db/admission/rate_limiter.h

@@ -129,6 +129,12 @@ class RateLimiter {
     /** Returns the number of tokens available in the underlying token bucket. **/
     double tokensAvailable() const;
 
+    /**
+     * Returns the balance of tokens in the bucket, which may be negative if requests have
+     * "borrowed" tokens.
+     * */
+    double tokenBalance() const;
+
     /** Returns the number of sessions that are sleeping in acquireToken(...). **/
     int64_t queued() const;