feat: support min and max replica counts on spec (#107)

BrandwatchLtd · Sep 6, 2023 · ae605ca · ae605ca
2 parents b858365 + dd924d0
commit ae605ca
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -23,7 +23,10 @@ metadata:
 spec:
   dryRun: false # set true to see what the autoscaler _would_ do
   cooloffSeconds: 300
-  maxScaleIncrements: 1
+  maxScaleIncrements: 1 # how many increments to scale by, increment of 1 will allow scaling from 1->2, 2->4, 4->8 
+                        #2 will allow 1->4 but not 1->8 
+  minReplicas: 1        # optional
+  maxReplicas: 16       # optional: partition count will be used if missing 
   scaleTargetRef:
     apiVersion: apps/v1
     kind: Deployment

diff --git a/kustomize/crd/kafkapodautoscaler.brandwatch.com-v1alpha1.yaml b/kustomize/crd/kafkapodautoscaler.brandwatch.com-v1alpha1.yaml
@@ -78,6 +78,10 @@ spec:
                       type: object
                   type: object
                 type: array
+              minReplicas:
+                type: integer
+              maxReplicas:
+                type: integer
               maxScaleIncrements:
                 type: integer
             required:

diff --git a/src/main/java/com/brandwatch/kafka_pod_autoscaler/KafkaPodAutoscalerReconciler.java b/src/main/java/com/brandwatch/kafka_pod_autoscaler/KafkaPodAutoscalerReconciler.java
@@ -33,6 +33,7 @@
 import com.brandwatch.kafka_pod_autoscaler.scaledresources.GenericScaledResourceFactory;
 import com.brandwatch.kafka_pod_autoscaler.triggers.TriggerProcessor;
 import com.brandwatch.kafka_pod_autoscaler.v1alpha1.KafkaPodAutoscaler;
+import com.brandwatch.kafka_pod_autoscaler.v1alpha1.KafkaPodAutoscalerSpec;
 import com.brandwatch.kafka_pod_autoscaler.v1alpha1.KafkaPodAutoscalerStatus;
 import com.brandwatch.kafka_pod_autoscaler.v1alpha1.kafkapodautoscalerspec.ScaleTargetRef;
 import com.brandwatch.kafka_pod_autoscaler.v1alpha1.kafkapodautoscalerspec.TriggerDefinition;
@@ -106,14 +107,13 @@ public UpdateControl<KafkaPodAutoscaler> reconcile(KafkaPodAutoscaler kafkaPodAu
                 return replicas;
             })
             .max().orElse(1);
-        var partitionCount = getPartitionCount(kafkaPodAutoscaler);
 
+        var partitionCount = getPartitionCount(kafkaPodAutoscaler);
         if (partitionCount.isPresent()) {
             statusLogger.recordPartitionCount(partitionCount.getAsInt());
         }
-        var maxScaleIncrements = kafkaPodAutoscaler.getSpec().getMaxScaleIncrements();
-        var finalReplicaCount = fitReplicaCount(currentReplicaCount, calculatedReplicaCount, partitionCount.orElse(calculatedReplicaCount),
-                                                maxScaleIncrements);
+        var finalReplicaCount = fitReplicaCount(currentReplicaCount, calculatedReplicaCount, partitionCount,
+                                                kafkaPodAutoscaler.getSpec());
 
         statusLogger.recordCurrentReplicaCount(currentReplicaCount);
         statusLogger.recordCalculatedReplicaCount(calculatedReplicaCount);
@@ -210,11 +210,24 @@ private com.brandwatch.kafka_pod_autoscaler.triggers.TriggerResult calculateTrig
                          .process(client, scaledResource, autoscaler, trigger, replicaCount);
     }
 
-    static int fitReplicaCount(int currentReplicaCount, int idealReplicaCount, int partitionCount, int maxScaleIncrements) {
+    static int fitReplicaCount(int currentReplicaCount, int idealReplicaCount, OptionalInt partitionCount, KafkaPodAutoscalerSpec spec) {
+        int minReplicas = Optional.ofNullable(spec.getMinReplicas()).orElse(1);
+        int maxReplicas = Optional.ofNullable(spec.getMaxReplicas()).orElse(partitionCount.orElse(Math.max(currentReplicaCount, idealReplicaCount)));
+        if (minReplicas < 1) {
+            throw new IllegalArgumentException("minReplicas cannot be less than 1: " + minReplicas);
+        }
+        if (partitionCount.isPresent() && maxReplicas > partitionCount.getAsInt()) {
+            throw new IllegalArgumentException("maxReplicas cannot be greater than the number of partitions on the topic: "
+                                                   + minReplicas + " (partitions=" + partitionCount + ")");
+        }
         if (idealReplicaCount == 0) {
-            return 1;
+            return Math.max(currentReplicaCount, minReplicas);
         }
-        var allowedReplicaCounts = partitionCountCache.get(partitionCount);
+        idealReplicaCount = Math.max(idealReplicaCount, minReplicas);
+        idealReplicaCount = Math.min(idealReplicaCount, maxReplicas);
+
+        var maxScaleIncrements = spec.getMaxScaleIncrements();
+        var allowedReplicaCounts = partitionCountCache.get(partitionCount.orElse(maxReplicas));
         var currentReplicaCountIndex = findReplicaCountIndex(currentReplicaCount, allowedReplicaCounts);
         var newReplicaCountIndex = findReplicaCountIndex(idealReplicaCount, allowedReplicaCounts);
 
@@ -223,7 +236,7 @@ static int fitReplicaCount(int currentReplicaCount, int idealReplicaCount, int p
             if (currentReplicaCountIndex > newReplicaCountIndex) {
                 newReplicaCountIndex = Math.max(currentReplicaCountIndex - maxScaleIncrements, 0);
             } else {
-                newReplicaCountIndex = Math.min(currentReplicaCountIndex + maxScaleIncrements, partitionCount);
+                newReplicaCountIndex = Math.min(currentReplicaCountIndex + maxScaleIncrements, partitionCount.orElse(maxReplicas));
             }
         }
 

diff --git a/src/main/java/com/brandwatch/kafka_pod_autoscaler/v1alpha1/KafkaPodAutoscalerSpec.java b/src/main/java/com/brandwatch/kafka_pod_autoscaler/v1alpha1/KafkaPodAutoscalerSpec.java
@@ -57,6 +57,16 @@ public class KafkaPodAutoscalerSpec implements KubernetesResource {
     private List<TriggerDefinition> triggers = Serialization.unmarshal("[]", List.class);
     @Getter
     @Setter
+    @JsonProperty("minReplicas")
+    @JsonSetter(nulls = Nulls.SKIP)
+    private Integer minReplicas = null;
+    @Getter
+    @Setter
+    @JsonProperty("maxReplicas")
+    @JsonSetter(nulls = Nulls.SKIP)
+    private Integer maxReplicas = null;
+    @Getter
+    @Setter
     @JsonProperty("maxScaleIncrements")
     @JsonSetter(nulls = Nulls.SKIP)
     private int maxScaleIncrements = 1;

diff --git a/src/test/java/com/brandwatch/kafka_pod_autoscaler/KafkaPodAutoscalerReconcilerTest.java b/src/test/java/com/brandwatch/kafka_pod_autoscaler/KafkaPodAutoscalerReconcilerTest.java
@@ -14,6 +14,7 @@
 import java.time.ZoneOffset;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
 import java.util.stream.Stream;
 
 import org.junit.jupiter.api.BeforeEach;
@@ -409,22 +410,31 @@ public void canScale_withKafkaConfig_multipleStaticTriggers() {
 
     @ParameterizedTest
     @MethodSource
-    public void fitReplicaCount(int currentReplicaCount, int idealReplicaCount, int partitionCount, int maxScaleIncrements,
+    public void fitReplicaCount(int currentReplicaCount, int idealReplicaCount, OptionalInt partitionCount,
+                                Integer minReplicas, Integer maxReplicas, int maxScaleIncrements,
                                 int expectedResult) {
-        assertThat(KafkaPodAutoscalerReconciler.fitReplicaCount(currentReplicaCount, idealReplicaCount, partitionCount, maxScaleIncrements))
+        var spec = mock(KafkaPodAutoscalerSpec.class);
+        when(spec.getMinReplicas()).thenReturn(minReplicas);
+        when(spec.getMaxReplicas()).thenReturn(maxReplicas);
+        when(spec.getMaxScaleIncrements()).thenReturn(maxScaleIncrements);
+        assertThat(KafkaPodAutoscalerReconciler.fitReplicaCount(currentReplicaCount, idealReplicaCount, partitionCount, spec))
             .isEqualTo(expectedResult);
     }
 
     public static Stream<Arguments> fitReplicaCount() {
         return Stream.of(
             // Scaling up
-            Arguments.of(1, 1, 16, 1000, 1),
-            Arguments.of(1, 3, 16, 1000, 4),
-            Arguments.of(4, 200, 16, 1000, 16),
-            Arguments.of(1, 3, 16, 1, 2),
+            Arguments.of(1, 1, OptionalInt.of(16), null, null, 1000, 1),
+            Arguments.of(1, 3, OptionalInt.of(16), null, null, 1000, 4),
+            Arguments.of(4, 200, OptionalInt.of(16), null, null, 1000, 16),
+            Arguments.of(4, 200, OptionalInt.of(16), null, 8, 1000, 8),
+            Arguments.of(1, 3, OptionalInt.of(16), null, null, 1, 2),
+            Arguments.of(1, 3, OptionalInt.empty(), null, null, 1, 3),
             // Scaling down
-            Arguments.of(16, 3, 16, 1000, 4),
-            Arguments.of(16, 3, 16, 1, 8)
+            Arguments.of(16, 3, OptionalInt.of(16), null, null, 1000, 4),
+            Arguments.of(16, 3, OptionalInt.of(16), 8, null, 1000, 8),
+            Arguments.of(16, 3, OptionalInt.of(16), null, null, 1, 8),
+            Arguments.of(16, 3, OptionalInt.empty(), null, null, 1, 8)
         );
     }