Add support for partition based scaling on the kafka scaler (#6558)

Scaling kafka consumers should be done in factors of partition count on the topic. This is to ensure that the partitions are evenly spread across all consumers. If the paritions are not evenly spread we run the risk of some partitions being consumed faster than other. This PR adds a new property on the kafka scaler `ensureEvenDistributionOfPartitions`. When this property is set to true the scaler ensure that the number of pods are always evenly spread across the number of topics on the partition.

Signed-off-by: Vishal Patel <38330754+patelvp@users.noreply.github.com>

Author: patelvp

CHANGELOG.md

@@ -85,12 +85,14 @@ To learn more about active deprecations, we recommend checking [GitHub Discussio
 - **General**: Allow excluding labels from being propagated from ScaledObject and ScaledJob to generated HPA and Job objects ([#6849](https://github.com/kedacore/keda/issues/6849))
 - **General**: chore: only add webhook DNS names when webhook patching is enabled ([#7002](https://github.com/kedacore/keda/issues/7002))
 - **General**: Improve Events emitted from ScaledObject controller ([#6802](https://github.com/kedacore/keda/issues/6802))
+- **Apache Kafka Scaler**: Add support for even distribution of partitions to consumers ([#2581](https://github.com/kedacore/keda/issues/2581))
 - **Azure Pipelines Scaler**: Add new `fetchUnfinishedJobsOnly` property to fetch only unfinished pipeline jobs for a pool ([#6819](https://github.com/kedacore/keda/issues/6819))
 - **Datadog Scaler**: Add a specific timeout configuration parameter for the Datadog trigger ([#6999](https://github.com/kedacore/keda/pull/6999))
 - **Datadog Scaler**: Fix bug with datadogNamespace config ([#6828](https://github.com/kedacore/keda/pull/6828))
 - **Datadog Scaler**: Improve Datadog scaler error messages ([#6999](https://github.com/kedacore/keda/pull/6999))
 - **Github Scaler**: Add support to control unlabeled job/runner matching ([#6900](https://github.com/kedacore/keda/issues/6900))
 - **InfluxDB Scaler**: Add support for InfluxDB v3 ([#6981](https://github.com/kedacore/keda/issues/6981))
+- **Kafka Scaler**: Add support for even distribution of partitions to consumers ([#2581](https://github.com/kedacore/keda/issues/2581))
 - **Metrics API Scaler**: Support AuthParams for authMode ([#6939](https://github.com/kedacore/keda/issues/6939))
 - **Metrics API Scaler**: Support multiple auth methods simultaneously ([#6642](https://github.com/kedacore/keda/issues/6642))
 - **RabbitMQ Scaler**: add `DeliverGetRate`, `PublishedToDeliveredRatio` and `ExpectedQueueConsumptionTime` trigger modes to RabbitMQ scaler ([#7071](https://github.com/kedacore/keda/issues/7071))

pkg/scalers/apache_kafka_scaler.go

@@ -60,8 +60,9 @@ type apacheKafkaMetadata 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 `keda:"name=scaleToZeroOnInvalidOffset, order=triggerMetadata, optional"`
-	LimitToPartitionsWithLag   bool `keda:"name=limitToPartitionsWithLag,   order=triggerMetadata, optional"`
+	ScaleToZeroOnInvalidOffset         bool `keda:"name=scaleToZeroOnInvalidOffset,         order=triggerMetadata, optional"`
+	LimitToPartitionsWithLag           bool `keda:"name=limitToPartitionsWithLag,           order=triggerMetadata, optional"`
+	EnsureEvenDistributionOfPartitions bool `keda:"name=ensureEvenDistributionOfPartitions, order=triggerMetadata, optional"`
 
 	// SASL
 	SASLType kafkaSaslType `keda:"name=sasl,     order=triggerMetadata;authParams, enum=none;plaintext;scram_sha256;scram_sha512;gssapi;aws_msk_iam, default=none"`
@@ -505,9 +506,15 @@ func (s *apacheKafkaScaler) getTotalLag(ctx context.Context) (int64, int64, erro
 
 	s.logger.V(1).Info(fmt.Sprintf("Kafka scaler: Consumer offsets %v, producer offsets %v", consumerOffsets, producerOffsets))
 
-	if !s.metadata.AllowIdleConsumers || s.metadata.LimitToPartitionsWithLag {
+	if !s.metadata.AllowIdleConsumers || s.metadata.LimitToPartitionsWithLag || s.metadata.EnsureEvenDistributionOfPartitions {
 		// 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 {
+			nextFactor := getNextFactorThatBalancesConsumersToTopicPartitions(totalLag, totalTopicPartitions, s.metadata.LagThreshold)
+			s.logger.V(1).Info(fmt.Sprintf("Kafka scaler: Providing metrics to ensure even distribution of partitions on totalLag %v, topicPartitions %v, evenPartitions %v", totalLag, totalTopicPartitions, nextFactor))
+			totalLag = nextFactor * s.metadata.LagThreshold
+		}
 		if s.metadata.LimitToPartitionsWithLag {
 			upperBound = partitionsWithLag
 		}

pkg/scalers/apache_kafka_scaler_test.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
@@ -594,6 +597,22 @@ func parseKafkaMetadata(config *scalersconfig.ScalerConfig, logger logr.Logger)
 		}
 	}
 
+	meta.ensureEvenDistributionOfPartitions = false
+	if val, ok := config.TriggerMetadata["ensureEvenDistributionOfPartitions"]; ok {
+		t, err := strconv.ParseBool(val)
+		if err != nil {
+			return meta, fmt.Errorf("error parsing ensureEvenDistributionOfPartitions: %w", err)
+		}
+		meta.ensureEvenDistributionOfPartitions = t
+
+		if meta.limitToPartitionsWithLag && meta.ensureEvenDistributionOfPartitions {
+			return meta, fmt.Errorf("limitToPartitionsWithLag and ensureEvenDistributionOfPartitions cannot be set simultaneously")
+		}
+		if len(meta.topic) == 0 && meta.ensureEvenDistributionOfPartitions {
+			return meta, fmt.Errorf("topic must be specified when using ensureEvenDistributionOfPartitions")
+		}
+	}
+
 	meta.version = sarama.V1_0_0_0
 	if val, ok := config.TriggerMetadata["version"]; ok {
 		val = strings.TrimSpace(val)
@@ -972,9 +991,15 @@ 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 {
+	if !s.metadata.allowIdleConsumers || s.metadata.limitToPartitionsWithLag || s.metadata.ensureEvenDistributionOfPartitions {
 		// 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 {
+			nextFactor := getNextFactorThatBalancesConsumersToTopicPartitions(totalLag, totalTopicPartitions, s.metadata.lagThreshold)
+			s.logger.V(1).Info(fmt.Sprintf("Kafka scaler: Providing metrics to ensure even distribution of partitions on totalLag %v, topicPartitions %v, evenPartitions %v", totalLag, totalTopicPartitions, nextFactor))
+			totalLag = nextFactor * s.metadata.lagThreshold
+		}
 		if s.metadata.limitToPartitionsWithLag {
 			upperBound = partitionsWithLag
 		}
@@ -986,6 +1011,16 @@ func (s *kafkaScaler) getTotalLag() (int64, int64, error) {
 	return totalLag, totalLagWithPersistent, nil
 }
 
+func getNextFactorThatBalancesConsumersToTopicPartitions(totalLag int64, totalTopicPartitions int64, lagThreshold int64) int64 {
+	factors := FindFactors(totalTopicPartitions)
+	for _, factor := range factors {
+		if factor*lagThreshold >= totalLag {
+			return factor
+		}
+	}
+	return totalTopicPartitions
+}
+
 type brokerOffsetResult struct {
 	offsetResp *sarama.OffsetResponse
 	err        error
@@ -1051,3 +1086,28 @@ func (s *kafkaScaler) getProducerOffsets(topicPartitions map[string][]int32) (ma
 
 	return topicPartitionsOffsets, nil
 }
+
+// FindFactors returns all factors of a given number
+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
+}