feat: introduced defaultBalancer for Kafka producers that dynamically selects the balancing strategy

Introduced defaultBalancer for Kafka producers that dynamically selects the balancing strategy:
Uses RoundRobin for messages without keys.
Uses Murmur2 for messages with keys.

Author: pinarrkok

balancer.go

@@ -2,8 +2,22 @@ package kafka
 
 import "github.com/segmentio/kafka-go"
 
+var (
+	balancerRoundRobin = GetBalancerRoundRobin()
+	balancerMurmur     = GetBalancerMurmur2Balancer()
+)
+
 type Balancer kafka.Balancer
 
+type defaultBalancer struct{}
+
+func (s *defaultBalancer) Balance(msg kafka.Message, partitions ...int) (partition int) {
+	if msg.Key == nil {
+		return balancerRoundRobin.Balance(msg, partitions...)
+	}
+	return balancerMurmur.Balance(msg, partitions...)
+}
+
 func GetBalancerCRC32() Balancer {
 	return &kafka.CRC32Balancer{}
 }

balancer_test.go

@@ -3,6 +3,8 @@ package kafka
 import (
 	"reflect"
 	"testing"
+
+	"github.com/segmentio/kafka-go"
 )
 
 func TestGetBalancerCRC32(t *testing.T) {
@@ -115,3 +117,103 @@ func TestGetBalancerString(t *testing.T) {
 		})
 	}
 }
+
+func TestDefaultBalancer_Balance(t *testing.T) {
+	partitions := []int{0, 1, 2, 3}
+
+	t.Run("Should_Use_RoundRobin_When_Key_Is_Nil", func(t *testing.T) {
+		// Given
+		msg := kafka.Message{Key: nil}
+		balancer := &defaultBalancer{}
+		expected := GetBalancerRoundRobin().Balance(msg, partitions...)
+
+		// When
+		result := balancer.Balance(msg, partitions...)
+
+		// Then
+		if result != expected {
+			t.Errorf("Expected RoundRobin partition %d, got %d", expected, result)
+		}
+	})
+
+	t.Run("Should_Use_Murmur2_When_Key_Is_Not_Nil", func(t *testing.T) {
+		// Given
+		msg := kafka.Message{Key: []byte("key")}
+		balancer := &defaultBalancer{}
+		expected := GetBalancerMurmur2Balancer().Balance(msg, partitions...)
+
+		// When
+		result := balancer.Balance(msg, partitions...)
+
+		// Then
+		if result != expected {
+			t.Errorf("Expected Murmur2Balancer partition %d, got %d", expected, result)
+		}
+	})
+}
+
+type optimizedBalancer struct{}
+
+func (s *optimizedBalancer) Balance(msg kafka.Message, partitions ...int) int {
+	var balancer kafka.Balancer
+	if msg.Key == nil {
+		balancer = balancerRoundRobin
+	} else {
+		balancer = balancerMurmur
+	}
+	return balancer.Balance(msg, partitions...)
+}
+
+func BenchmarkDefaultBalancer_WithAlloc(b *testing.B) {
+	partitions := []int{0, 1, 2, 3}
+	msgWithKey := kafka.Message{Key: []byte("key")}
+	msgWithoutKey := kafka.Message{Key: nil}
+	balancer := &defaultBalancer{}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		if i%2 == 0 {
+			balancer.Balance(msgWithKey, partitions...)
+		} else {
+			balancer.Balance(msgWithoutKey, partitions...)
+		}
+	}
+}
+
+func BenchmarkDefaultBalancer_Optimized(b *testing.B) {
+	partitions := []int{0, 1, 2, 3}
+	msgWithKey := kafka.Message{Key: []byte("key")}
+	msgWithoutKey := kafka.Message{Key: nil}
+	balancer := &optimizedBalancer{}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		if i%2 == 0 {
+			balancer.Balance(msgWithKey, partitions...)
+		} else {
+			balancer.Balance(msgWithoutKey, partitions...)
+		}
+	}
+}
+
+func BenchmarkDefaultBalancer_Direct(b *testing.B) {
+	partitions := []int{0, 1, 2, 3}
+	msgWithKey := kafka.Message{Key: []byte("key")}
+	msgWithoutKey := kafka.Message{Key: nil}
+	balancer := &optimizedBalancer{}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		if i%2 == 0 {
+			balancer.Balance(msgWithKey, partitions...)
+		} else {
+			balancer.Balance(msgWithoutKey, partitions...)
+		}
+	}
+}

producer.go

@@ -27,7 +27,7 @@ func NewProducer(cfg *ProducerConfig, interceptors ...ProducerInterceptor) (Prod
 	kafkaWriter := &kafka.Writer{
 		Addr:                   kafka.TCP(cfg.Writer.Brokers...),
 		Topic:                  cfg.Writer.Topic,
-		Balancer:               cfg.Writer.Balancer,
+		Balancer:               &defaultBalancer{},
 		MaxAttempts:            cfg.Writer.MaxAttempts,
 		WriteBackoffMin:        cfg.Writer.WriteBackoffMin,
 		WriteBackoffMax:        cfg.Writer.WriteBackoffMax,

producer_config.go

@@ -111,4 +111,8 @@ func (cfg *ProducerConfig) setDefaults() {
 			cfg.DistributedTracingConfiguration.Propagator = otel.GetTextMapPropagator()
 		}
 	}
+
+	if cfg.Writer.Balancer == nil {
+		cfg.Writer.Balancer = &defaultBalancer{}
+	}
 }

producer_config_test.go

@@ -8,19 +8,43 @@ import (
 )
 
 func TestProducerConfig_setDefaults(t *testing.T) {
-	// Given
-	cfg := ProducerConfig{DistributedTracingEnabled: true}
+	t.Run("Should_Set_Default_DistributedTracing_Fields", func(t *testing.T) {
+		// Given
+		cfg := ProducerConfig{DistributedTracingEnabled: true}
 
-	// When
-	cfg.setDefaults()
+		// When
+		cfg.setDefaults()
 
-	// Then
-	if cfg.DistributedTracingConfiguration.TracerProvider == nil {
-		t.Fatal("Traceprovider cannot be null")
-	}
-	if cfg.DistributedTracingConfiguration.Propagator == nil {
-		t.Fatal("Propagator cannot be null")
-	}
+		// Then
+		if cfg.DistributedTracingConfiguration.TracerProvider == nil {
+			t.Fatal("TracerProvider cannot be nil")
+		}
+		if cfg.DistributedTracingConfiguration.Propagator == nil {
+			t.Fatal("Propagator cannot be nil")
+		}
+	})
+
+	t.Run("Should_Set_Default_Balancer_When_Nil", func(t *testing.T) {
+		// Given
+		cfg := ProducerConfig{
+			Writer: WriterConfig{
+				Balancer: nil,
+			},
+		}
+
+		// When
+		cfg.setDefaults()
+
+		// Then
+		if cfg.Writer.Balancer == nil {
+			t.Fatal("Balancer should not be nil")
+		}
+
+		_, ok := cfg.Writer.Balancer.(*defaultBalancer)
+		if !ok {
+			t.Fatalf("Expected balancer to be of type *DefaultBalancer, but got %T", cfg.Writer.Balancer)
+		}
+	})
 }
 
 func TestProducerConfig_Json(t *testing.T) {