[WIP] Add partition based even scaling kafka scaler

pkg/scalers/kafka_scaler.go

@@ -23,7 +23,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"os"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -68,8 +70,9 @@ type kafkaMetadata struct {
 
 	// If an invalid offset is found, whether to scale to 1 (false - the default) so consumption can
 	// occur or scale to 0 (true). See discussion in https://github.com/kedacore/keda/issues/2612
-	scaleToZeroOnInvalidOffset bool
-	limitToPartitionsWithLag   bool
+	scaleToZeroOnInvalidOffset         bool
+	limitToPartitionsWithLag           bool
+	ensureEvenDistributionOfPartitions bool
 
 	// SASL
 	saslType kafkaSaslType
@@ -974,8 +977,14 @@ func (s *kafkaScaler) getTotalLag() (int64, int64, error) {
 	s.logger.V(1).Info(fmt.Sprintf("Kafka scaler: Providing metrics based on totalLag %v, topicPartitions %v, threshold %v", totalLag, len(topicPartitions), s.metadata.lagThreshold))
 
 	if !s.metadata.allowIdleConsumers || s.metadata.limitToPartitionsWithLag {
-		// don't scale out beyond the number of topicPartitions or partitionsWithLag depending on settings
 		upperBound := totalTopicPartitions
+		// Ensure that the number of partitions is evenly distributed across the number of consumers
+		if s.metadata.ensureEvenDistributionOfPartitions {
+			// How do we find the current number of consumers?
+			nextFactor := GetNextFactor(1, totalTopicPartitions)
+			totalLag = nextFactor * s.metadata.lagThreshold
+		}
+		// don't scale out beyond the number of topicPartitions or partitionsWithLag depending on settings
 		if s.metadata.limitToPartitionsWithLag {
 			upperBound = partitionsWithLag
 		}
@@ -987,6 +996,16 @@ func (s *kafkaScaler) getTotalLag() (int64, int64, error) {
 	return totalLag, totalLagWithPersistent, nil
 }
 
+func GetNextFactor(current int64, totalTopicPartitions int64) int64 {
+	factors := FindFactors(totalTopicPartitions)
+	for _, factor := range factors {
+		if factor > current {
+			return factor
+		}
+	}
+	return totalTopicPartitions
+}
+
 type brokerOffsetResult struct {
 	offsetResp *sarama.OffsetResponse
 	err        error
@@ -1052,3 +1071,27 @@ func (s *kafkaScaler) getProducerOffsets(topicPartitions map[string][]int32) (ma
 
 	return topicPartitionsOffsets, nil
 }
+
+func FindFactors(n int64) []int64 {
+	if n < 1 {
+		return nil
+	}
+
+	var factors []int64
+	sqrtN := int64(math.Sqrt(float64(n)))
+
+	for i := int64(1); i <= sqrtN; i++ {
+		if n%i == 0 {
+			factors = append(factors, i)
+			if i != n/i {
+				factors = append(factors, n/i)
+			}
+		}
+	}
+
+	sort.Slice(factors, func(i, j int) bool {
+		return factors[i] < factors[j]
+	})
+
+	return factors
+}

pkg/scalers/kafka_scaler_test.go

@@ -650,6 +650,65 @@ func TestGetTopicPartitions(t *testing.T) {
 	}
 }
 
+func TestFindFactors(t *testing.T) {
+	factors := FindFactors(10)
+	if !reflect.DeepEqual(factors, []int64{1, 2, 5, 10}) {
+		t.Errorf("Expected factors to be %v but got %v", []int64{1, 2, 5, 10}, factors)
+	}
+
+	factors = FindFactors(13)
+	if !reflect.DeepEqual(factors, []int64{1, 13}) {
+		t.Errorf("Expected factors to be %v but got %v", []int64{1, 13}, factors)
+	}
+
+	factors = FindFactors(17)
+	if !reflect.DeepEqual(factors, []int64{1, 17}) {
+		t.Errorf("Expected factors to be %v but got %v", []int64{1, 17}, factors)
+	}
+
+	factors = FindFactors(25)
+	if !reflect.DeepEqual(factors, []int64{1, 5, 25}) {
+		t.Errorf("Expected factors to be %v but got %v", []int64{1, 5, 25}, factors)
+	}
+
+	factors = FindFactors(50)
+	if !reflect.DeepEqual(factors, []int64{1, 2, 5, 10, 25, 50}) {
+		t.Errorf("Expected factors to be %v but got %v", []int64{1, 2, 5, 10, 25, 50}, factors)
+	}
+
+	factors = FindFactors(100)
+	if !reflect.DeepEqual(factors, []int64{1, 2, 4, 5, 10, 20, 25, 50, 100}) {
+		t.Errorf("Expected factors to be %v but got %v", []int64{1, 2, 4, 5, 10, 20, 25, 50, 100}, factors)
+	}
+}
+
+func TestGetNextFactor(t *testing.T) {
+	factor := GetNextFactor(1, 100)
+	if factor != 2 {
+		t.Errorf("Expected factor to be %v but got %v", 2, factor)
+	}
+
+	factor = GetNextFactor(2, 100)
+	if factor != 4 {
+		t.Errorf("Expected factor to be %v but got %v", 4, factor)
+	}
+
+	factor = GetNextFactor(4, 100)
+	if factor != 5 {
+		t.Errorf("Expected factor to be %v but got %v", 5, factor)
+	}
+
+	factor = GetNextFactor(5, 100)
+	if factor != 10 {
+		t.Errorf("Expected factor to be %v but got %v", 10, factor)
+	}
+
+	factor = GetNextFactor(100, 100)
+	if factor != 100 {
+		t.Errorf("Expected factor to be %v but got %v", 10, factor)
+	}
+}
+
 type MockClusterAdmin struct {
 	partitionIds []int32
 }