From f5f7878237b485ac544a80fe0cf5e6536e040940 Mon Sep 17 00:00:00 2001 From: 98001yash Date: Thu, 15 Oct 2026 20:29:14 +0530 Subject: [PATCH] Fix distributed token bucket refill across gateways --- load-tests/distributed.js | 87 ++++++++-- .../config/TrafficControlConfiguration.java | 8 +- .../lease/DistributedTokenBucket.java | 12 ++ .../lease/RedisLeaseCoordinator.java | 154 +++++++++++++++++- .../traffic/TrafficDecisionEngine.java | 57 ++++++- .../RedisLeaseCoordinatorIntegrationTest.java | 32 +++- .../InMemoryDistributedTokenBucket.java | 80 +++++++++ .../traffic/TrafficDecisionEngineTest.java | 78 ++++++++- 8 files changed, 476 insertions(+), 32 deletions(-) create mode 100644 src/main/java/com/yashdotdev/distributed_traffic_control/lease/DistributedTokenBucket.java create mode 100644 src/test/java/com/yashdotdev/distributed_traffic_control/traffic/InMemoryDistributedTokenBucket.java diff --git a/load-tests/distributed.js b/load-tests/distributed.js index 4c976f2..a9dbc02 100644 --- a/load-tests/distributed.js +++ b/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', }); -} \ No newline at end of file +} diff --git a/src/main/java/com/yashdotdev/distributed_traffic_control/config/TrafficControlConfiguration.java b/src/main/java/com/yashdotdev/distributed_traffic_control/config/TrafficControlConfiguration.java index d0328b5..6e3394f 100644 --- a/src/main/java/com/yashdotdev/distributed_traffic_control/config/TrafficControlConfiguration.java +++ b/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 diff --git a/src/main/java/com/yashdotdev/distributed_traffic_control/lease/DistributedTokenBucket.java b/src/main/java/com/yashdotdev/distributed_traffic_control/lease/DistributedTokenBucket.java new file mode 100644 index 0000000..f8e75d3 --- /dev/null +++ b/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); +} diff --git a/src/main/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinator.java b/src/main/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinator.java index 0c99fec..d8523d8 100644 --- a/src/main/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinator.java +++ b/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 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 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: *

@@ -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 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; } diff --git a/src/main/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngine.java b/src/main/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngine.java index 3073b86..3ea28cc 100644 --- a/src/main/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngine.java +++ b/src/main/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngine.java @@ -1,11 +1,14 @@ package com.yashdotdev.distributed_traffic_control.traffic; import com.yashdotdev.distributed_traffic_control.allocation.CapacityAllocator; +import com.yashdotdev.distributed_traffic_control.lease.DistributedTokenBucket; +import com.yashdotdev.distributed_traffic_control.lease.GlobalCapacityKey; import com.yashdotdev.distributed_traffic_control.lease.LeaseConsumptionResult; import com.yashdotdev.distributed_traffic_control.lease.QuotaLease; import com.yashdotdev.distributed_traffic_control.monitoring.TrafficControlMetrics; import com.yashdotdev.distributed_traffic_control.policy.PolicyProvider; import com.yashdotdev.distributed_traffic_control.policy.TrafficPolicy; +import com.yashdotdev.distributed_traffic_control.policy.TrafficPolicyType; import com.yashdotdev.distributed_traffic_control.quota.QuotaConsumptionResult; import com.yashdotdev.distributed_traffic_control.quota.QuotaCoordinator; import com.yashdotdev.distributed_traffic_control.quota.QuotaKey; @@ -22,12 +25,14 @@ public class TrafficDecisionEngine { private final CapacityAllocator capacityAllocator; private final TrafficControlMetrics metrics; private final ObservationRegistry observationRegistry; + private final DistributedTokenBucket distributedTokenBucket; public TrafficDecisionEngine( PolicyProvider policyProvider, QuotaCoordinator quotaCoordinator, CapacityAllocator capacityAllocator) { - this(policyProvider, quotaCoordinator, capacityAllocator, null, null); + this(policyProvider, quotaCoordinator, capacityAllocator, + null, null, null); } public TrafficDecisionEngine( @@ -35,7 +40,8 @@ public TrafficDecisionEngine( QuotaCoordinator quotaCoordinator, CapacityAllocator capacityAllocator, TrafficControlMetrics metrics) { - this(policyProvider, quotaCoordinator, capacityAllocator, metrics, null); + this(policyProvider, quotaCoordinator, capacityAllocator, + metrics, null, null); } public TrafficDecisionEngine( @@ -44,11 +50,23 @@ public TrafficDecisionEngine( CapacityAllocator capacityAllocator, TrafficControlMetrics metrics, ObservationRegistry observationRegistry) { + this(policyProvider, quotaCoordinator, capacityAllocator, + metrics, observationRegistry, null); + } + + public TrafficDecisionEngine( + PolicyProvider policyProvider, + QuotaCoordinator quotaCoordinator, + CapacityAllocator capacityAllocator, + TrafficControlMetrics metrics, + ObservationRegistry observationRegistry, + DistributedTokenBucket distributedTokenBucket) { this.policyProvider = policyProvider; this.quotaCoordinator = quotaCoordinator; this.capacityAllocator = capacityAllocator; this.metrics = metrics; this.observationRegistry = observationRegistry; + this.distributedTokenBucket = distributedTokenBucket; } public TrafficDecision evaluate(TrafficRequest request) { @@ -88,6 +106,41 @@ private TrafficDecision evaluateDecision(TrafficRequest request) { request.getSubject(), request.getResource()); + if (trafficPolicy.getType() == TrafficPolicyType.TOKEN_BUCKET) { + if (distributedTokenBucket == null) { + throw new IllegalStateException( + "Distributed token bucket is required for TOKEN_BUCKET policies"); + } + + GlobalCapacityKey capacityKey = new GlobalCapacityKey( + trafficPolicy.getPolicyId(), + request.getResource()); + + long remainingCapacity = observe( + "traffic.token-bucket.consume", + () -> distributedTokenBucket.tryConsume( + capacityKey, + trafficPolicy.getCapacity(), + trafficPolicy.getRefillRate(), + request.getRequestedAt())); + + if (remainingCapacity >= 0) { + recordAllowed(); + return new TrafficDecision( + TrafficDecisionStatus.ALLOWED, + "Request allowed", + remainingCapacity); + } + + recordQuotaExhausted(); + recordRejected(); + + return new TrafficDecision( + TrafficDecisionStatus.REJECTED, + "Traffic quota exhausted", + 0); + } + QuotaConsumptionResult consumptionResult = observe( "traffic.quota.evaluate", () -> quotaCoordinator.tryConsume(quotaKey, trafficPolicy)); diff --git a/src/test/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinatorIntegrationTest.java b/src/test/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinatorIntegrationTest.java index 682815c..e58257b 100644 --- a/src/test/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinatorIntegrationTest.java +++ b/src/test/java/com/yashdotdev/distributed_traffic_control/lease/RedisLeaseCoordinatorIntegrationTest.java @@ -207,6 +207,37 @@ void shouldReclaimCapacityFromExpiredLease() { assertEquals("node-b", replacementLease.get().getNodeId()); } + @Test + void shouldRefillGlobalCapacityOverTime() throws InterruptedException { + RedisLeaseCoordinator coordinator = new RedisLeaseCoordinator(createRedisTemplate()); + GlobalCapacityKey capacityKey = new GlobalCapacityKey("refill-policy", "/api/orders"); + for (int i = 0; i < 5; i++) { + long remaining = coordinator.tryConsume(capacityKey, 5, 1, Instant.now()); + assertTrue(remaining >= 0, "Expected the initial bucket to allow 5 requests"); + } + assertEquals(-1L, coordinator.tryConsume(capacityKey, 5, 1, + Instant.now()), "Expected the bucket to be exhausted"); + Thread.sleep(1100); + assertTrue(coordinator.tryConsume(capacityKey, 5, 1, + Instant.now()) >= 0, "Expected a refilled token to be available"); + } + + @Test + void shouldShareTokenBucketAcrossCoordinatorInstances() { + RedisLeaseCoordinator firstCoordinator = new RedisLeaseCoordinator(createRedisTemplate()); + RedisLeaseCoordinator secondCoordinator = new RedisLeaseCoordinator(createRedisTemplate()); + + GlobalCapacityKey capacityKey = new GlobalCapacityKey("shared-token-bucket-policy", "/api/shared-test"); + for (int i = 0; i < 3; i++) { + assertTrue(firstCoordinator.tryConsume(capacityKey, 5, 1, Instant.now()) >= 0, + "First coordinator should consume available tokens"); + } + for (int i = 0; i < 2; i++) { + assertTrue(secondCoordinator.tryConsume(capacityKey, 5, 1, Instant.now()) >= 0, "Second coordinator should consume the shared tokens"); + } + assertEquals(-1L, firstCoordinator.tryConsume(capacityKey, 5, 1, Instant.now())); + assertEquals(-1L, secondCoordinator.tryConsume(capacityKey, 5, 1, Instant.now())); + } @BeforeEach void cleanRedis() { @@ -214,7 +245,6 @@ void cleanRedis() { Objects.requireNonNull(redisTemplate.getConnectionFactory()).getConnection().serverCommands().flushDb(); } - private static class MutableClock extends Clock { private Instant currentTime; private MutableClock(Instant currentTime) { diff --git a/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/InMemoryDistributedTokenBucket.java b/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/InMemoryDistributedTokenBucket.java new file mode 100644 index 0000000..c10581f --- /dev/null +++ b/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/InMemoryDistributedTokenBucket.java @@ -0,0 +1,80 @@ +package com.yashdotdev.distributed_traffic_control.traffic; + +import com.yashdotdev.distributed_traffic_control.lease.DistributedTokenBucket; +import com.yashdotdev.distributed_traffic_control.lease.GlobalCapacityKey; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +final class InMemoryDistributedTokenBucket implements DistributedTokenBucket { + + private static final class Bucket { + private double tokens; + private Instant lastRefill; + + private Bucket(long capacity, Instant now) { + this.tokens = capacity; + this.lastRefill = now; + } + } + + private final Map buckets = new HashMap<>(); + private final Clock clock; + + InMemoryDistributedTokenBucket() { + this(Clock.systemUTC()); + } + + InMemoryDistributedTokenBucket(Clock clock) { + this.clock = Objects.requireNonNull(clock, "clock must not be null"); + } + + @Override + public synchronized long tryConsume( + GlobalCapacityKey capacityKey, + long capacity, + long refillRate, + Instant currentTime) { + + Objects.requireNonNull(capacityKey, "capacityKey must not be null"); + Objects.requireNonNull(currentTime, "currentTime must not be null"); + + if (capacity <= 0 || refillRate <= 0) { + throw new IllegalArgumentException( + "Capacity and refill rate must be positive"); + } + + Instant now = clock.instant(); + + Bucket bucket = buckets.computeIfAbsent( + capacityKey, + key -> new Bucket(capacity, now)); + + long elapsedMillis = Math.max( + 0L, + Duration.between(bucket.lastRefill, now).toMillis()); + + double refilledTokens = + elapsedMillis * (double) refillRate / 1000.0; + + bucket.tokens = Math.min( + capacity, + bucket.tokens + refilledTokens); + + if (now.isAfter(bucket.lastRefill)) { + bucket.lastRefill = now; + } + + if (bucket.tokens < 1.0) { + return -1L; + } + + bucket.tokens -= 1.0; + + return (long) Math.floor(bucket.tokens); + } +} diff --git a/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngineTest.java b/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngineTest.java index 71c51c9..9dceb97 100644 --- a/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngineTest.java +++ b/src/test/java/com/yashdotdev/distributed_traffic_control/traffic/TrafficDecisionEngineTest.java @@ -97,7 +97,11 @@ void shouldAllowRequestAfterQuotaIsRefilled() { MutableClock clock = new MutableClock(Instant.parse("2026-08-28T10:00:00Z")); TrafficPolicy policy = new TrafficPolicy("refill-policy", "Refill Traffic Policy", TrafficPolicyType.TOKEN_BUCKET, PolicyStatus.ACTIVE, 1, 1, Instant.now(clock)); - TrafficDecisionEngine trafficDecisionEngine = createTrafficDecisionEngine(policy, new InMemoryQuotaCoordinator(clock)); + TrafficDecisionEngine trafficDecisionEngine = + createTrafficDecisionEngine( + policy, + new InMemoryQuotaCoordinator(clock), + new InMemoryDistributedTokenBucket(clock)); TrafficRequest request = new TrafficRequest("request-1", new TrafficSubject("user-123", TrafficSubjectType.USER), "/api/orders", Instant.now(clock)); TrafficDecision firstDecision = trafficDecisionEngine.evaluate(request); TrafficDecision secondDecision = trafficDecisionEngine.evaluate(request); @@ -114,7 +118,11 @@ void shouldNotRefillQuotaBeyondMaximumCapacity() { MutableClock clock = new MutableClock(Instant.parse("2026-08-28T10:00:00Z")); int capacity = 5; TrafficPolicy policy = new TrafficPolicy("capacity-policy", "Capacity Traffic Policy", TrafficPolicyType.TOKEN_BUCKET, PolicyStatus.ACTIVE, capacity, 2, Instant.now(clock)); - TrafficDecisionEngine trafficDecisionEngine = createTrafficDecisionEngine(policy, new InMemoryQuotaCoordinator(clock)); + TrafficDecisionEngine trafficDecisionEngine = + createTrafficDecisionEngine( + policy, + new InMemoryQuotaCoordinator(clock), + new InMemoryDistributedTokenBucket(clock)); TrafficRequest request = new TrafficRequest("request-1", new TrafficSubject("user-123", TrafficSubjectType.USER), "/api/orders", Instant.now(clock)); for (int i = 0; i < capacity; i++) { @@ -139,7 +147,15 @@ void shouldNotRefillQuotaBeyondMaximumCapacity() { @Test void shouldAllowRequestWhenLocalQuotaIsExhaustedAndCapacityIsAllocated() { - TrafficPolicy policy = new TrafficPolicy("allocation-policy", "Allocation Traffic Policy", TrafficPolicyType.TOKEN_BUCKET, PolicyStatus.ACTIVE, 1, 1, Instant.parse("2026-08-28T10:00:00Z")); + TrafficPolicy policy = new TrafficPolicy( + "allocation-policy", + "Allocation Traffic Policy", + TrafficPolicyType.FIXED_WINDOW, + PolicyStatus.ACTIVE, + 1, + 1, + Duration.ofSeconds(60), + Instant.parse("2026-08-28T10:00:00Z")); InMemoryQuotaCoordinator quotaCoordinator = new InMemoryQuotaCoordinator(); QuotaKey quotaKey = new QuotaKey(policy.getPolicyId(), new TrafficSubject("user-123", TrafficSubjectType.USER), "/api/orders"); InMemoryLeaseCoordinator leaseCoordinator = new InMemoryLeaseCoordinator(); @@ -151,7 +167,14 @@ void shouldAllowRequestWhenLocalQuotaIsExhaustedAndCapacityIsAllocated() { Clock clock = Clock.fixed(Instant.parse("2026-08-28T10:00:00Z"), ZoneOffset.UTC); LeaseStore leaseStore = new InMemoryLeaseStore(); InMemoryCapacityAllocator capacityAllocator = new InMemoryCapacityAllocator(leaseCoordinator, new FixedAllocationStrategy(), allocationProperties, leaseStore, clock); - TrafficDecisionEngine trafficDecisionEngine = new TrafficDecisionEngine(new InMemoryPolicyProvider(policy), quotaCoordinator, capacityAllocator); + TrafficDecisionEngine trafficDecisionEngine = + new TrafficDecisionEngine( + new InMemoryPolicyProvider(policy), + quotaCoordinator, + capacityAllocator, + null, + null, + new InMemoryDistributedTokenBucket()); TrafficRequest request = new TrafficRequest("request-1", new TrafficSubject("user-123", TrafficSubjectType.USER), "/api/orders", Instant.parse("2026-08-28T10:00:00Z")); TrafficDecision firstDecision = trafficDecisionEngine.evaluate(request); TrafficDecision secondDecision = trafficDecisionEngine.evaluate(request); @@ -181,7 +204,13 @@ public LeaseConsumptionResult tryConsume(QuotaLease lease, Instant currentTime) } }; - return new TrafficDecisionEngine(new InMemoryPolicyProvider(policy), quotaCoordinator, noOpCapacityAllocator); + return new TrafficDecisionEngine( + new InMemoryPolicyProvider(policy), + quotaCoordinator, + noOpCapacityAllocator, + null, + null, + new InMemoryDistributedTokenBucket()); } @Test @@ -206,7 +235,14 @@ public LeaseConsumptionResult tryConsume(QuotaLease quotaLease, Instant currentT } }; - TrafficDecisionEngine trafficDecisionEngine = new TrafficDecisionEngine(new InMemoryPolicyProvider(policy), quotaCoordinator, capacityAllocator); + TrafficDecisionEngine trafficDecisionEngine = + new TrafficDecisionEngine( + new InMemoryPolicyProvider(policy), + quotaCoordinator, + capacityAllocator, + null, + null, + new InMemoryDistributedTokenBucket()); TrafficRequest request = new TrafficRequest("request-1", new TrafficSubject("user-123", TrafficSubjectType.USER), "/api/orders", Instant.parse("2026-08-28T10:00:00Z")); @@ -223,6 +259,36 @@ public LeaseConsumptionResult tryConsume(QuotaLease quotaLease, Instant currentT assertEquals("Traffic quota exhausted", secondDecision.getReason()); } + + private TrafficDecisionEngine createTrafficDecisionEngine( + TrafficPolicy policy, + InMemoryQuotaCoordinator quotaCoordinator, + InMemoryDistributedTokenBucket distributedTokenBucket) { + + CapacityAllocator noOpCapacityAllocator = new CapacityAllocator() { + @Override + public Optional allocate( + TrafficPolicy trafficPolicy, QuotaKey quotaKey) { + return Optional.empty(); + } + + @Override + public LeaseConsumptionResult tryConsume( + QuotaLease lease, Instant currentTime) { + return new LeaseConsumptionResult( + false, lease.getRemainingCapacity()); + } + }; + + return new TrafficDecisionEngine( + new InMemoryPolicyProvider(policy), + quotaCoordinator, + noOpCapacityAllocator, + null, + null, + distributedTokenBucket); + } + private static class MutableClock extends Clock { private Instant currentTime;