[otel] add kafka partitioner extension

go.mod

@@ -199,6 +199,7 @@ require (
 	github.com/microsoft/go-mssqldb v1.9.5
 	github.com/microsoft/wmi v0.38.3
 	github.com/open-telemetry/opentelemetry-collector-contrib/exporter/elasticsearchexporter v0.150.0
+	github.com/open-telemetry/opentelemetry-collector-contrib/exporter/kafkaexporter v0.150.0
 	github.com/pierrec/lz4/v4 v4.1.26
 	github.com/pkg/xattr v0.4.9
 	github.com/prometheus/prometheus v0.309.1
@@ -237,10 +238,10 @@ require (
 	github.com/coder/websocket v1.8.14
 	github.com/elastic/gokrb5/v8 v8.0.0-20251105095404-23cc45e6a102
 	github.com/mattn/go-sqlite3 v1.14.32
-	github.com/open-telemetry/opentelemetry-collector-contrib/exporter/kafkaexporter v0.150.0
 	github.com/parsiya/golnk v0.0.0-20251207220015-443df11fe4fb
 	github.com/richardlehane/mscfb v1.0.6
 	github.com/rogpeppe/go-internal v1.14.1
+	github.com/twmb/franz-go v1.20.7
 	github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78
 	go.opentelemetry.io/collector/client v1.56.0
 	go.opentelemetry.io/collector/component/componenttest v0.150.0
@@ -253,6 +254,7 @@ require (
 	go.opentelemetry.io/collector/exporter/exportertest v0.150.0
 	go.opentelemetry.io/collector/extension v1.56.0
 	go.opentelemetry.io/collector/extension/extensionauth v1.56.0
+	go.opentelemetry.io/collector/extension/extensiontest v0.150.0
 	go.opentelemetry.io/collector/pipeline v1.56.0
 	go.opentelemetry.io/collector/processor v1.56.0
 	go.opentelemetry.io/collector/processor/processorhelper v0.150.0
@@ -473,7 +475,6 @@ require (
 	github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/tklauser/numcpus v0.11.0 // indirect
-	github.com/twmb/franz-go v1.20.7 // indirect
 	github.com/twmb/franz-go/pkg/kadm v1.17.2 // indirect
 	github.com/twmb/franz-go/pkg/kmsg v1.13.1 // indirect
 	github.com/twmb/franz-go/pkg/sasl/kerberos v1.1.0 // indirect
@@ -503,7 +504,6 @@ require (
 	go.opentelemetry.io/collector/exporter/xexporter v0.150.0 // indirect
 	go.opentelemetry.io/collector/extension/extensioncapabilities v0.150.0 // indirect
 	go.opentelemetry.io/collector/extension/extensionmiddleware v0.150.0 // indirect
-	go.opentelemetry.io/collector/extension/extensiontest v0.150.0 // indirect
 	go.opentelemetry.io/collector/extension/xextension v0.150.0 // indirect
 	go.opentelemetry.io/collector/featuregate v1.56.0 // indirect
 	go.opentelemetry.io/collector/internal/componentalias v0.150.0 // indirect

libbeat/outputs/kafka/partition.go

@@ -77,7 +77,7 @@ func initPartitionStrategy(
 	}
 
 	if len(partition) > 1 {
-		return nil, false, errors.New("Too many partitioners")
+		return nil, false, errors.New("too many partitioners")
 	}
 
 	// extract partitioner from config
@@ -235,7 +235,7 @@ func makeHashPartitioner() partitioner {
 		}
 
 		// create positive hash value
-		return hash2Partition(hash, numPartitions)
+		return Hash2Partition(hash, numPartitions), nil
 	}
 }
 
@@ -249,7 +249,7 @@ func makeFieldsHashPartitioner(log *logp.Logger, fields []string, dropFail bool)
 
 			var err error
 			for _, field := range fields {
-				err = hashFieldValue(hasher, msg.data.Content.Fields, field)
+				err = HashFieldValue(hasher, msg.data.Content.Fields, field)
 				if err != nil {
 					break
 				}
@@ -268,16 +268,16 @@ func makeFieldsHashPartitioner(log *logp.Logger, fields []string, dropFail bool)
 			hash = msg.hash
 		}
 
-		return hash2Partition(hash, numPartitions)
+		return Hash2Partition(hash, numPartitions), nil
 	}
 }
 
-func hash2Partition(hash uint32, numPartitions int32) (int32, error) {
-	p := int32(hash) & 0x7FFFFFFF
-	return p % numPartitions, nil
+func Hash2Partition(hash uint32, numPartitions int32) int32 {
+	p := int32(hash) & 0x7FFFFFFF //nolint:gosec // Conversion from int to int32 is safe here.
+	return p % numPartitions
 }
 
-func hashFieldValue(h hash.Hash32, event mapstr.M, field string) error {
+func HashFieldValue(h hash.Hash32, event mapstr.M, field string) error {
 	type stringer interface {
 		String() string
 	}

libbeat/outputs/kafka/partition_test.go

@@ -366,7 +366,7 @@ var hash2PartitionTests = []struct {
 func TestHash2Partition(t *testing.T) {
 	for _, tt := range hash2PartitionTests {
 		t.Run(tt.testName, func(t *testing.T) {
-			var partition, _ = hash2Partition(tt.hash, tt.numPartitions)
+			var partition = Hash2Partition(tt.hash, tt.numPartitions)
 			assert.Equal(t, tt.expectedResult, partition)
 		})
 	}

x-pack/otel/extension/kafkapartitionerextension/config.go

@@ -0,0 +1,21 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+package kafkapartitionerextension
+
+import (
+	"go.opentelemetry.io/collector/component"
+)
+
+type Config struct {
+	PartitionerConfig map[string]any `mapstructure:",remain"`
+}
+
+func createDefaultConfig() component.Config {
+	return &Config{}
+}
+
+func (c *Config) Validate() error {
+	return nil
+}

x-pack/otel/extension/kafkapartitionerextension/extension.go

@@ -0,0 +1,48 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+package kafkapartitionerextension
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/exporter/kafkaexporter"
+	"github.com/twmb/franz-go/pkg/kgo"
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/extension"
+	"go.uber.org/zap"
+
+	"github.com/elastic/elastic-agent-libs/logp"
+)
+
+var _ extension.Extension = (*kafkaPartitioner)(nil)
+var _ kafkaexporter.RecordPartitionerExtension = (*kafkaPartitioner)(nil)
+
+type kafkaPartitioner struct {
+	cfg         *Config
+	logger      *zap.Logger
+	partitioner kgo.Partitioner
+}
+
+func (k *kafkaPartitioner) Start(context.Context, component.Host) error {
+	logger, err := logp.NewZapLogger(k.logger)
+	if err != nil {
+		return fmt.Errorf("error creating logger: %w", err)
+	}
+	partitioner, err := makePartitioner(logger, k.cfg.PartitionerConfig)
+	if err != nil {
+		return fmt.Errorf("error configuring the partitioner: %w", err)
+	}
+	k.partitioner = partitioner
+	return nil
+}
+
+func (*kafkaPartitioner) Shutdown(context.Context) error {
+	return nil
+}
+
+func (k *kafkaPartitioner) GetPartitioner() kgo.Partitioner {
+	return k.partitioner
+}