From a8a6c806d6871928e91a35429fd36db245f5f08e Mon Sep 17 00:00:00 2001 From: 98001yash Date: Sat, 24 Oct 2026 14:04:33 +0530 Subject: [PATCH] Make lease allocation strategy pluggable --- .../AllocationStrategyConfiguration.java | 18 ++++++ .../config/TrafficControlConfiguration.java | 31 +++++----- .../AllocationStrategyConfigurationTest.java | 57 +++++++++++++++++++ 3 files changed, 91 insertions(+), 15 deletions(-) create mode 100644 src/main/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfiguration.java create mode 100644 src/test/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfigurationTest.java diff --git a/src/main/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfiguration.java b/src/main/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfiguration.java new file mode 100644 index 0000000..6934221 --- /dev/null +++ b/src/main/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfiguration.java @@ -0,0 +1,18 @@ +package com.yashdotdev.distributed_traffic_control.allocation; + +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +@AutoConfiguration +@EnableConfigurationProperties(AllocationProperties.class) +public class AllocationStrategyConfiguration { + + @Bean + @ConditionalOnMissingBean(AllocationStrategy.class) + public AllocationStrategy allocationStrategy( + AllocationProperties allocationProperties) { + return new FixedAllocationStrategy(allocationProperties); + } +} \ 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 6e3394f..c679a8f 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 @@ -2,13 +2,22 @@ import com.yashdotdev.distributed_traffic_control.allocation.AllocationProperties; import com.yashdotdev.distributed_traffic_control.allocation.AllocationStrategy; +import com.yashdotdev.distributed_traffic_control.allocation.AllocationStrategyConfiguration; import com.yashdotdev.distributed_traffic_control.allocation.CapacityAllocator; -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.lease.InMemoryLeaseStore; +import com.yashdotdev.distributed_traffic_control.lease.LeaseCoordinator; +import com.yashdotdev.distributed_traffic_control.lease.LeaseStore; +import com.yashdotdev.distributed_traffic_control.lease.RedisLeaseCoordinator; import com.yashdotdev.distributed_traffic_control.monitoring.TrafficControlMetrics; -import com.yashdotdev.distributed_traffic_control.policy.*; +import com.yashdotdev.distributed_traffic_control.policy.DefaultPolicyManagementService; +import com.yashdotdev.distributed_traffic_control.policy.InMemoryPolicyProvider; +import com.yashdotdev.distributed_traffic_control.policy.PolicyManagementService; +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.policy.PolicyStatus; import com.yashdotdev.distributed_traffic_control.quota.InMemoryQuotaCoordinator; import com.yashdotdev.distributed_traffic_control.quota.QuotaCoordinator; import com.yashdotdev.distributed_traffic_control.traffic.DefaultTrafficControlService; @@ -18,14 +27,13 @@ import io.micrometer.observation.ObservationRegistry; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.Import; import org.springframework.data.redis.core.StringRedisTemplate; import java.time.Clock; -import org.springframework.boot.context.properties.EnableConfigurationProperties; - @Configuration -@EnableConfigurationProperties(AllocationProperties.class) +@Import(AllocationStrategyConfiguration.class) public class TrafficControlConfiguration { @Bean @@ -70,11 +78,6 @@ public LeaseStore leaseStore() { return new InMemoryLeaseStore(); } - @Bean - public FixedAllocationStrategy allocationStrategy(AllocationProperties allocationProperties) { - return new FixedAllocationStrategy(allocationProperties); - } - @Bean public CapacityAllocator capacityAllocator(LeaseCoordinator leaseCoordinator, AllocationStrategy allocationStrategy, AllocationProperties allocationProperties, LeaseStore leaseStore, Clock clock) { return new InMemoryCapacityAllocator(leaseCoordinator, allocationStrategy, allocationProperties, leaseStore, clock); @@ -85,19 +88,17 @@ public TrafficDecisionEngine trafficDecisionEngine( PolicyProvider policyProvider, QuotaCoordinator quotaCoordinator, CapacityAllocator capacityAllocator, - TrafficControlMetrics metrics, + TrafficControlMetrics trafficControlMetrics, ObservationRegistry observationRegistry, DistributedTokenBucket distributedTokenBucket) { - return new TrafficDecisionEngine( policyProvider, quotaCoordinator, capacityAllocator, - metrics, + trafficControlMetrics, observationRegistry, distributedTokenBucket); } - @Bean public TrafficControlService trafficControlService(TrafficDecisionEngine trafficDecisionEngine) { return new DefaultTrafficControlService(trafficDecisionEngine); diff --git a/src/test/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfigurationTest.java b/src/test/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfigurationTest.java new file mode 100644 index 0000000..0d51800 --- /dev/null +++ b/src/test/java/com/yashdotdev/distributed_traffic_control/allocation/AllocationStrategyConfigurationTest.java @@ -0,0 +1,57 @@ +package com.yashdotdev.distributed_traffic_control.allocation; + +import com.yashdotdev.distributed_traffic_control.policy.TrafficPolicy; +import org.junit.jupiter.api.Test; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +import static org.assertj.core.api.Assertions.assertThat; + +class AllocationStrategyConfigurationTest { + + private final ApplicationContextRunner contextRunner = + new ApplicationContextRunner() + .withConfiguration( + org.springframework.boot.autoconfigure.AutoConfigurations.of( + AllocationStrategyConfiguration.class)); + + @Test + void shouldRegisterFixedAllocationStrategyByDefault() { + contextRunner.run(context -> { + assertThat(context).hasNotFailed(); + assertThat(context).hasSingleBean(AllocationStrategy.class); + assertThat(context.getBean(AllocationStrategy.class)) + .isInstanceOf(FixedAllocationStrategy.class); + }); + } + + @Test + void shouldUseCustomAllocationStrategyWhenProvided() { + contextRunner + .withUserConfiguration(CustomStrategyConfiguration.class) + .run(context -> { + assertThat(context).hasNotFailed(); + assertThat(context).hasSingleBean(AllocationStrategy.class); + assertThat(context.getBean(AllocationStrategy.class)) + .isInstanceOf(CustomAllocationStrategy.class); + }); + } + + @Configuration(proxyBeanMethods = false) + static class CustomStrategyConfiguration { + + @Bean + AllocationStrategy customAllocationStrategy() { + return new CustomAllocationStrategy(); + } + } + + static class CustomAllocationStrategy implements AllocationStrategy { + + @Override + public long determineCapacity(TrafficPolicy policy) { + return 1; + } + } +} \ No newline at end of file