Fix distributed token bucket refill across gateways

load-tests/distributed.js

@@ -1,31 +1,88 @@
+
 import http from 'k6/http';
-import {check} from 'k6';
+import { check } from 'k6';
+import { Counter, Rate } from 'k6/metrics';
 
 export const options = {
-    vus: 10, duration: '30s',
-
+    vus: 1,
+    duration: '30s',
     thresholds: {
-        http_req_failed: ['rate<0.05'], http_req_duration: ['p(95)<500'],
+        http_req_failed: ['rate<0.001'],
+        http_req_duration: ['p(95)<500'],
     },
 };
 
-const BASE_URL = __ENV.BASE_URL || 'http://host.docker.internal:8081';
+const GATEWAY_URLS = (
+    __ENV.GATEWAY_URLS ||
+    'http://host.docker.internal:8081,http://host.docker.internal:8082,http://host.docker.internal:8083'
+)
+    .split(',')
+    .map((url) => url.trim())
+    .filter(Boolean);
+
+if (GATEWAY_URLS.length === 0) {
+    throw new Error('At least one gateway URL is required');
+}
+
+const gatewayRequests = GATEWAY_URLS.map(
+    (_, index) => new Counter(`gateway_${index + 1}_requests`)
+);
+
+const allowedDecisions = new Rate('traffic_decision_allowed');
+const rejectedDecisions = new Rate('traffic_decision_rejected');
+const unknownDecisions = new Rate('traffic_decision_unknown');
 
 export default function () {
+    const gatewayIndex = (__VU + __ITER) % GATEWAY_URLS.length;
+    const baseUrl = GATEWAY_URLS[gatewayIndex];
+
+    gatewayRequests[gatewayIndex].add(1);
+
     const payload = JSON.stringify({
-        requestId: `k6-distributed-${__VU}-${__ITER}-${Date.now()}`, subject: {
-            subjectId: 'load-test-user', type: 'USER',
-        }, resource: '/api/orders', requestedAt: new Date().toISOString(),
+        requestId: `k6-${__VU}-${__ITER}-${Date.now()}`,
+        subject: {
+            subjectId: `load-test-user-${__VU}`,
+            type: 'USER',
+        },
+        resource: '/api/orders',
+        requestedAt: new Date().toISOString(),
     });
 
-    const response = http.post(`${BASE_URL}/api/v1/traffic/evaluate`, payload, {
-        headers: {
-            'Content-Type': 'application/json',
-        }, timeout: '5s',
+    const response = http.post(
+        `${baseUrl}/api/v1/traffic/evaluate`,
+        payload,
+        {
+            headers: { 'Content-Type': 'application/json' },
+            timeout: '5s',
+        }
+    );
+
+    const validHttpResponse = check(response, {
+        'HTTP status is 200': (r) => r.status === 200,
     });
 
+    let decision = 'unknown';
+
+    if (response.status === 200 && response.body) {
+        try {
+            const body = JSON.parse(response.body);
+
+            if (body.status === 'ALLOWED') {
+                decision = 'allowed';
+            } else if (body.status === 'REJECTED') {
+                decision = 'rejected';
+            }
+        } catch (_) {
+            decision = 'unknown';
+        }
+    }
+
+    allowedDecisions.add(decision === 'allowed');
+    rejectedDecisions.add(decision === 'rejected');
+    unknownDecisions.add(decision === 'unknown');
+
     check(response, {
-        'status is 200': (r) => r.status === 200,
-        'response contains decision': (r) => r.status === 200 && r.body != null && (r.body.includes('ALLOWED') || r.body.includes('REJECTED')),
+        'business decision is recognized': () =>
+            validHttpResponse && decision !== 'unknown',
     });
-}
+}

src/main/java/com/yashdotdev/distributed_traffic_control/config/TrafficControlConfiguration.java

@@ -6,6 +6,7 @@
 import com.yashdotdev.distributed_traffic_control.allocation.FixedAllocationStrategy;
 import com.yashdotdev.distributed_traffic_control.allocation.InMemoryCapacityAllocator;
 import com.yashdotdev.distributed_traffic_control.lease.*;
+import com.yashdotdev.distributed_traffic_control.lease.DistributedTokenBucket;
 import com.yashdotdev.distributed_traffic_control.monitoring.TrafficControlMetrics;
 import com.yashdotdev.distributed_traffic_control.policy.*;
 import com.yashdotdev.distributed_traffic_control.quota.InMemoryQuotaCoordinator;
@@ -85,13 +86,16 @@ public TrafficDecisionEngine trafficDecisionEngine(
             QuotaCoordinator quotaCoordinator,
             CapacityAllocator capacityAllocator,
             TrafficControlMetrics metrics,
-            ObservationRegistry observationRegistry) {
+            ObservationRegistry observationRegistry,
+            DistributedTokenBucket distributedTokenBucket) {
+
         return new TrafficDecisionEngine(
                 policyProvider,
                 quotaCoordinator,
                 capacityAllocator,
                 metrics,
-                observationRegistry);
+                observationRegistry,
+                distributedTokenBucket);
     }
 
     @Bean

src/main/java/com/yashdotdev/distributed_traffic_control/lease/DistributedTokenBucket.java

@@ -0,0 +1,12 @@
+package com.yashdotdev.distributed_traffic_control.lease;
+
+import java.time.Instant;
+
+public interface DistributedTokenBucket {
+
+    long tryConsume(
+            GlobalCapacityKey capacityKey,
+            long capacity,
+            long refillRate,
+            Instant currentTime);
+}

src/main/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinator.java

@@ -13,12 +13,110 @@
 import java.util.Optional;
 import java.util.UUID;
 
-public class RedisLeaseCoordinator implements LeaseCoordinator {
+public class RedisLeaseCoordinator
+        implements LeaseCoordinator, DistributedTokenBucket {
 
     private static final String QUOTA_KEY_PREFIX = "traffic-control:quota:";
     private static final String LEASE_KEY_PREFIX = "traffic-control:lease:";
     private static final String LEASE_REGISTRY_SUFFIX = ":leases";
 
+
+    private static final String TOKEN_BUCKET_KEY_PREFIX =
+            "traffic-control:token-bucket:";
+
+    private static final DefaultRedisScript<Long> TOKEN_BUCKET_SCRIPT =
+            new DefaultRedisScript<>("""
+                local capacity = tonumber(ARGV[1])
+                local refillRate = tonumber(ARGV[2])
+
+                -- Use Redis server time so all gateways share one clock.
+                local redisTime = redis.call('TIME')
+                local now = tonumber(redisTime[1]) * 1000
+                    + math.floor(tonumber(redisTime[2]) / 1000)
+
+                if not capacity or capacity <= 0
+                   or not refillRate or refillRate <= 0 then
+                    return -1
+                end
+
+                -- Read the current shared bucket state.
+                local state = redis.call(
+                    'HMGET',
+                    KEYS[1],
+                    'tokens',
+                    'lastRefillMillis'
+                )
+
+                local tokens = tonumber(state[1])
+                local lastRefill = tonumber(state[2])
+
+                -- Initialize a missing bucket with full capacity.
+                if not tokens or not lastRefill then
+                    tokens = capacity
+                    lastRefill = now
+                end
+
+                -- Refill proportionally to elapsed time, capped at capacity.
+                if now > lastRefill then
+                    local elapsed = now - lastRefill
+
+                    tokens = math.min(
+                        capacity,
+                        tokens + elapsed * refillRate / 1000
+                    )
+
+                    lastRefill = now
+                end
+
+                -- Reject without consuming when fewer than one token remains.
+                if tokens < 1 then
+                    redis.call(
+                        'HSET',
+                        KEYS[1],
+                        'tokens', tostring(tokens),
+                        'lastRefillMillis', tostring(lastRefill)
+                    )
+
+                    return -1
+                end
+
+                -- Atomically consume one token.
+                tokens = tokens - 1
+
+                redis.call(
+                    'HSET',
+                    KEYS[1],
+                    'tokens', tostring(tokens),
+                    'lastRefillMillis', tostring(lastRefill)
+                )
+
+                return math.floor(tokens)
+                """, Long.class);
+
+
+
+    private static final DefaultRedisScript<Long> INITIALIZE_TOKEN_BUCKET_SCRIPT =
+            new DefaultRedisScript<>("""
+                if redis.call('EXISTS', KEYS[1]) == 1 then
+                    return 0
+                end
+
+                local redisTime = redis.call('TIME')
+                local now = tonumber(redisTime[1]) * 1000
+                    + math.floor(tonumber(redisTime[2]) / 1000)
+
+                redis.call(
+                    'HSET',
+                    KEYS[1],
+                    'tokens', ARGV[1],
+                    'lastRefillMillis', tostring(now)
+                )
+
+                return 1
+                """, Long.class);
+
+
+
     /**
      * Atomically:
      * <p>
@@ -380,20 +478,38 @@ public void registerCapacity(GlobalCapacityKey capacityKey, long capacity) {
         }
 
         if (capacity <= 0) {
-            throw new IllegalArgumentException("capacity must be greater than zero");
+            throw new IllegalArgumentException(
+                    "capacity must be greater than zero");
         }
-        redisTemplate.opsForValue().setIfAbsent(buildGlobalCapacityRedisKey(capacityKey), String.valueOf(capacity));
+
+        redisTemplate.opsForValue().setIfAbsent(
+                buildGlobalCapacityRedisKey(capacityKey),
+                String.valueOf(capacity));
+
+        redisTemplate.execute(
+                INITIALIZE_TOKEN_BUCKET_SCRIPT,
+                List.of(buildTokenBucketRedisKey(capacityKey)),
+                String.valueOf(capacity));
     }
 
     @Override
     public boolean removeCapacity(GlobalCapacityKey capacityKey) {
         if (capacityKey == null) {
-            throw new IllegalArgumentException("capacityKey must not be null");
+            throw new IllegalArgumentException(
+                    "capacityKey must not be null");
         }
-        Boolean deleted = redisTemplate.delete(buildGlobalCapacityRedisKey(capacityKey));
-        return Boolean.TRUE.equals(deleted);
+
+        Boolean capacityDeleted = redisTemplate.delete(
+                buildGlobalCapacityRedisKey(capacityKey));
+
+        Boolean tokenBucketDeleted = redisTemplate.delete(
+                buildTokenBucketRedisKey(capacityKey));
+
+        return Boolean.TRUE.equals(capacityDeleted)
+                || Boolean.TRUE.equals(tokenBucketDeleted);
     }
 
+
     @Override
     public Optional<QuotaLease> acquireLease(QuotaKey quotaKey, String nodeId, long requestedCapacity, Duration leaseDuration) {
         validateQuotaKey(quotaKey);
@@ -531,9 +647,35 @@ public boolean releaseLease(QuotaLease lease) {
         return Long.valueOf(1L).equals(result);
     }
 
+    @Override
+    public long tryConsume(GlobalCapacityKey capacityKey, long capacity, long refillRate, Instant currentTime) {
+
+        if (capacityKey == null) {
+            throw new IllegalArgumentException("capacityKey must not be null");
+        }
+        if (capacity <= 0 || refillRate <= 0) {
+            throw new IllegalArgumentException("capacity and refillRate must be greater than zero");
+        }
+        if (currentTime == null) {
+            throw new IllegalArgumentException("currentTime must not be null");
+        }
+        String key = TOKEN_BUCKET_KEY_PREFIX + capacityKey.policyId() + ":" + capacityKey.resource();
+        Long result = observe("redis.token-bucket.consume",
+                () -> redisTemplate.execute(TOKEN_BUCKET_SCRIPT, List.of(key),
+                        String.valueOf(capacity), String.valueOf(refillRate)));
+
+        return result == null ? -1L : result;
+    }
+
     private String buildGlobalCapacityRedisKey(GlobalCapacityKey capacityKey) {
         return QUOTA_KEY_PREFIX + capacityKey.policyId() + ":" + capacityKey.resource();
     }
+
+    private String buildTokenBucketRedisKey(GlobalCapacityKey capacityKey) {
+        return TOKEN_BUCKET_KEY_PREFIX
+                + capacityKey.policyId() + ":" + capacityKey.resource();
+    }
+
     private String buildLeaseRegistryRedisKey(GlobalCapacityKey capacityKey) {
         return buildGlobalCapacityRedisKey(capacityKey) + LEASE_REGISTRY_SUFFIX;
     }