impr(CLDSRV-852): Refactor TokenBucket class

- Change constructor to accept resourceClass, resourceId, measure
- Precalculate interval for reuse in constructor
- Add hasCapacity() method
- Add updateLimit() method to encasulate limit change logic
- Remove unused mechanism to record request durations
- Simplify refill logic

Author: tmacro

lib/api/apiUtils/rateLimit/tokenBucket.js

@@ -16,32 +16,44 @@ const { instance: redisClient } = require('./client');
 const { config } = require('../../../Config');
 const { calculateInterval } = require('./gcra');
 
-// Map of bucket name -> WorkerTokenBucket instance
+// Map of `${resourceClass}:${resourceId}:${measure}` -> WorkerTokenBucket instance
 const tokenBuckets = new Map();
 
 /**
- * Per-bucket token bucket for a single worker
+ * Per-resourceClass+resourceID+measure token bucket for a single worker
  */
 class WorkerTokenBucket {
-    constructor(bucketName, limitConfig, log) {
-        this.bucketName = bucketName;
+    constructor(resourceClass, resourceId, measure, limitConfig, log) {
+        this.resourceClass = resourceClass;
+        this.resourceId = resourceId;
+        this.measure = measure;
         this.limitConfig = limitConfig;
         this.log = log;
 
-        // Token buffer configuration
-        this.bufferSize = config.rateLimiting?.tokenBucketBufferSize || 50; // Max tokens to hold
-        this.refillThreshold = config.rateLimiting?.tokenBucketRefillThreshold || 20; // Trigger refill when below this
+        this.bufferSize = config.rateLimiting?.tokenBucketBufferSize; // Max tokens to hold
+        this.refillThreshold = config.rateLimiting?.tokenBucketRefillThreshold; // Trigger refill when below this
         this.tokens = this.bufferSize; // Start with full buffer for fail-open at startup
-
-        // Refill state
-        this.refillInProgress = false;
+        this.interval = calculateInterval(this.limitConfig.limit, config.rateLimiting.nodes);
         this.lastRefillTime = Date.now();
         this.refillCount = 0;
-        this.requestCounter = 0;
-        this.lastRequestTime = 0;
+        this.refillInProgress = false;
+    }
 
-        // Track consumption rate for adaptive refill
-        this.requestTimestamps = new Set();
+    hasCapacity() {
+        return this.tokens > 0;
+    }
+
+    updateLimit(updatedConfig) {
+        if (this.limitConfig.limit !== updatedConfig.limit ||
+            this.limitConfig.burstCapacity !== updatedConfig.burstCapacity
+        ) {
+            const oldConfig = this.limitConfig;
+            this.limitConfig = updatedConfig;
+            this.interval = calculateInterval(updatedConfig.limit, config.rateLimiting.nodes);
+            return { updated: true, oldConfig };
+        }
+
+        return { updated: false };
     }
 
     /**
@@ -50,11 +62,8 @@ class WorkerTokenBucket {
      * @returns {boolean} True if request allowed, false if throttled
      */
     tryConsume() {
-        // Record request for rate calculation
-        this.recordRequest();
-
         if (this.tokens > 0) {
-            this.tokens--;
+            this.tokens -= 1;
             return true; // ALLOWED
         }
 
@@ -78,16 +87,6 @@ class WorkerTokenBucket {
             return;
         }
 
-        // Trigger async refill
-        await this.refill();
-    }
-
-    /**
-     * Request tokens from Redis using GCRA enforcement
-     *
-     * @returns {Promise<void>}
-     */
-    async refill() {
         this.refillInProgress = true;
         const startTime = Date.now();
 
@@ -100,29 +99,28 @@ class WorkerTokenBucket {
             }
 
             // Calculate GCRA parameters
-            const nodes = config.rateLimiting.nodes || 1;
-            const workers = config.clusters || 1;
-            const interval = calculateInterval(this.limitConfig.limit, nodes, workers);
-            const burstCapacitySeconds =
-                config.rateLimiting.bucket?.defaultConfig?.requestsPerSecond?.burstCapacity || 1;
-            const burstCapacity = burstCapacitySeconds * 1000;
-
             let granted = requested;
             if (redisClient.isReady()) {
                 // Request tokens from Redis (atomic GCRA enforcement)
                 granted = await util.promisify(redisClient.grantTokens.bind(redisClient))(
-                    this.bucketName,
+                    this.resourceClass,
+                    this.resourceId,
+                    this.measure,
                     requested,
-                    interval,
-                    burstCapacity,
+                    this.interval,
+                    this.limitConfig.burstCapacity,
                 );
             } else {
                 // Connection to redis has failed in some way.
                 // Client will be reconnecting in the background.
                 // We grant the requested amount of tokens anyway to avoid degrading service availability.
                 this.log.warn(
                     'rate limit redis client not connected. granting tokens anyway to avoid service degradation',
-                    { bucketName: this.bucketName },
+                    {
+                        resourceClass: this.resourceClass,
+                        resourceId: this.resourceId,
+                        measure: this.measure,
+                    },
                 );
             }
 
@@ -134,7 +132,9 @@ class WorkerTokenBucket {
             const duration = this.lastRefillTime - startTime;
 
             this.log.debug('Token refill completed', {
-                bucketName: this.bucketName,
+                resourceClass: this.resourceClass,
+                resourceId: this.resourceId,
+                measure: this.measure,
                 requested,
                 granted,
                 newBalance: this.tokens,
@@ -145,118 +145,67 @@ class WorkerTokenBucket {
             // Warn if refill took too long or granted too few
             if (duration > 100) {
                 this.log.warn('Slow token refill detected', {
-                    bucketName: this.bucketName,
+                    resourceClass: this.resourceClass,
+                    resourceId: this.resourceId,
+                    measure: this.measure,
                     durationMs: duration,
                 });
             }
 
             if (granted === 0 && requested > 0) {
                 this.log.trace('Token refill denied - quota exhausted', {
-                    bucketName: this.bucketName,
+                    resourceClass: this.resourceClass,
+                    resourceId: this.resourceId,
+                    measure: this.measure,
                     requested,
                 });
             }
 
         } catch (err) {
             this.log.error('Token refill failed', {
-                bucketName: this.bucketName,
+                resourceClass: this.resourceClass,
+                resourceId: this.resourceId,
+                measure: this.measure,
                 error: err.message,
                 stack: err.stack,
             });
         } finally {
             this.refillInProgress = false;
         }
     }
-
-    /**
-     * Record request timestamp for rate calculation
-     */
-    recordRequest() {
-        const now = Date.now();
-        if (now == this.lastRequestTime) {
-            this.requestCounter++;
-        } else {
-            this.lastRequestTime = now;
-            this.requestCounter = 0;
-        }
-
-        this.requestTimestamps.add(now * 1000 + this.requestCounter);
-
-        // Keep only last 1 second of timestamps
-        const cutoff = (now - 1000) * 1000;
-        for (const timestamp of this.requestTimestamps.values()) {
-            if (timestamp < cutoff) {
-                this.requestTimestamps.delete(timestamp);
-            }
-        }
-    }
-
-    /**
-     * Get current request rate (requests per second)
-     *
-     * @returns {number}
-     */
-    getCurrentRate() {
-        const now = Date.now();
-        const cutoff = (now - 1000) * 1000;
-
-        let count = 0;
-        for (const timestamp of this.requestTimestamps.values()) {
-            if (timestamp >= cutoff) {
-                count++;
-            }
-        }
-        return count;
-    }
-
-    /**
-     * Get token bucket stats for monitoring
-     *
-     * @returns {object}
-     */
-    getStats() {
-        return {
-            tokens: this.tokens,
-            bufferSize: this.bufferSize,
-            refillThreshold: this.refillThreshold,
-            refillInProgress: this.refillInProgress,
-            lastRefillTime: this.lastRefillTime,
-            refillCount: this.refillCount,
-            currentRate: this.getCurrentRate(),
-        };
-    }
 }
 
 /**
  * Get or create token bucket for a bucket
  *
- * @param {string} bucketName - Bucket name
+ * @param {string} resourceClass - "bucket" or "account"
+ * @param {string} resourceId - bucket name or account canonicalId
+ * @param {string} measure - measure id e.g. "rps"
  * @param {object} limitConfig - Rate limit configuration
  * @param {object} log - Logger instance
  * @returns {WorkerTokenBucket}
  */
-function getTokenBucket(bucketName, limitConfig, log) {
-    let bucket = tokenBuckets.get(bucketName);
-
+function getTokenBucket(resourceClass, resourceId, measure, limitConfig, log) {
+    const cacheKey = `${resourceClass}:${resourceId}:${measure}`;
+    let bucket = tokenBuckets.get(cacheKey);
     if (!bucket) {
-        bucket = new WorkerTokenBucket(bucketName, limitConfig, log);
-        tokenBuckets.set(bucketName, bucket);
+        bucket = new WorkerTokenBucket(resourceClass, resourceId, measure, limitConfig, log);
+        tokenBuckets.set(cacheKey, bucket);
 
         log.debug('Created token bucket', {
-            bucketName,
+            cacheKey,
             bufferSize: bucket.bufferSize,
             refillThreshold: bucket.refillThreshold,
         });
-    } else if (bucket.limitConfig.limit !== limitConfig.limit) {
-        // Update limit config when it changes dynamically
-        const oldLimit = bucket.limitConfig.limit;
-        bucket.limitConfig = limitConfig;
-
-        log.info('Updated token bucket limit config', {
-            bucketName,
-            oldLimit,
-            newLimit: limitConfig.limit,
-        });
+    } else {
+        const { updated, oldConfig } = bucket.updateLimit(limitConfig);
+        if (updated) {
+            log.info('Updated token bucket limit config', {
+                cacheKey,
+                old: oldConfig,
+                new: limitConfig,
+            });
+        }
     }
 
     return bucket;
@@ -282,15 +231,15 @@ function cleanupTokenBuckets(maxIdleMs = 60000) {
     const now = Date.now();
     const toRemove = [];
 
-    for (const [bucketName, bucket] of tokenBuckets.entries()) {
+    for (const [key, bucket] of tokenBuckets.entries()) {
         const idleTime = now - bucket.lastRefillTime;
         if (idleTime > maxIdleMs && bucket.tokens === 0) {
-            toRemove.push(bucketName);
+            toRemove.push(key);
         }
     }
 
-    for (const bucketName of toRemove) {
-        tokenBuckets.delete(bucketName);
+    for (const key of toRemove) {
+        tokenBuckets.delete(key);
     }
 
     return toRemove.length;
@@ -299,11 +248,13 @@ function cleanupTokenBuckets(maxIdleMs = 60000) {
 /**
  * Remove a specific token bucket (used when rate limit config is deleted)
  *
- * @param {string} bucketName - Bucket name
+ * @param {string} resourceClass - "bucket" or "account"
+ * @param {string} resourceId - bucket name or account canonicalId
+ * @param {string} measure - measure id e.g. "rps"
  * @returns {boolean} True if bucket was found and removed
  */
-function removeTokenBucket(bucketName) {
-    return tokenBuckets.delete(bucketName);
+function removeTokenBucket(resourceClass, resourceId, measure) {
+    return tokenBuckets.delete(`${resourceClass}:${resourceId}:${measure}`);
 }
 
 module.exports = {

lib/api/bucketDeleteRateLimit.js

@@ -53,7 +53,7 @@ function bucketDeleteRateLimit(authInfo, request, log, callback) {
             }
             // Invalidate cache and remove token bucket
             cache.deleteCachedConfig(cache.namespace.bucket, bucketName);
-            removeTokenBucket(bucketName);
+            removeTokenBucket('bucket', bucketName, 'rps');
             log.debug('invalidated rate limit cache and token bucket for bucket', { bucketName });
             // TODO: implement Utapi metric support
             return callback(null, corsHeaders);