feat(loki.source.kafka): Add rack_id to fetch from the closest replica

docs/sources/reference/components/loki/loki.source.kafka.md

@@ -43,6 +43,7 @@ You can use the following arguments with `loki.source.kafka`:
 | `assignor`               | `string`             | The consumer group rebalancing strategy to use.         | `"range"`             | no       |
 | `group_id`               | `string`             | The Kafka consumer group ID.                            | `"loki.source.kafka"` | no       |
 | `labels`                 | `map(string)`        | The labels to associate with each received Kafka event. | `{}`                  | no       |
+| `rack_id`                | `string`             | The rack identifier for this client.                    | `""`                  | no       |
 | `relabel_rules`          | `RelabelRules`       | Relabeling rules to apply on log entries.               | `{}`                  | no       |
 | `use_incoming_timestamp` | `bool`               | Whether to use the timestamp received from Kafka.       | `false`               | no       |
 | `version`                | `string`             | Kafka version to connect to.                            | `"2.2.1"`             | no       |
@@ -53,6 +54,9 @@ If a topic starts with a '^', it's treated as a regular expression and may match
 
 Labels from the `labels` argument are applied to every message that the component reads.
 
+The `rack_id` setting enables rack-aware replica selection.
+When it's set and the brokers use a rack-aware replica selector, the component fetches from the closest replica instead of the partition leader.
+
 The `relabel_rules` field can make use of the `rules` export value from a [`loki.relabel`][loki.relabel] component to apply one or more relabeling rules to log entries before they're forwarded to the list of receivers in `forward_to`.
 
 In addition to custom labels, the following internal labels prefixed with `__` are available:

internal/component/loki/source/internal/kafkatarget/config.go

@@ -39,6 +39,11 @@ type TargetConfig struct {
 	// Rebalancing strategy to use. (e.g. sticky, roundrobin or range)
 	Assignor string `yaml:"assignor"`
 
+	// RackID is the rack identifier for this client. When set and the brokers
+	// use a rack-aware replica selector, the consumer fetches from the closest
+	// replica instead of the partition leader (KIP-392).
+	RackID string `yaml:"rack_id"`
+
 	// Authentication strategy with Kafka brokers
 	Authentication Authentication `yaml:"authentication"`
 

internal/component/loki/source/internal/kafkatarget/target_syncer.go

@@ -52,28 +52,10 @@ func NewSyncer(
 	if err := validateConfig(&cfg); err != nil {
 		return nil, err
 	}
-	version, err := sarama.ParseKafkaVersion(cfg.KafkaConfig.Version)
+	config, err := newSaramaConfig(cfg)
 	if err != nil {
 		return nil, err
 	}
-	config := sarama.NewConfig()
-	config.Version = version
-	config.Consumer.Offsets.Initial = sarama.OffsetOldest
-
-	switch cfg.KafkaConfig.Assignor {
-	case sarama.StickyBalanceStrategyName:
-		config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategySticky()
-	case sarama.RoundRobinBalanceStrategyName:
-		config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategyRoundRobin()
-	case sarama.RangeBalanceStrategyName, "":
-		config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategyRange()
-	default:
-		return nil, fmt.Errorf("unrecognized consumer group partition assignor: %s", cfg.KafkaConfig.Assignor)
-	}
-	config, err = withAuthentication(*config, cfg.KafkaConfig.Authentication)
-	if err != nil {
-		return nil, fmt.Errorf("error setting up kafka authentication: %w", err)
-	}
 	client, err := sarama.NewClient(cfg.KafkaConfig.Brokers, config)
 	if err != nil {
 		return nil, fmt.Errorf("error creating kafka client: %w", err)
@@ -113,6 +95,36 @@ func NewSyncer(
 	return t, nil
 }
 
+// newSaramaConfig builds the Sarama client configuration used by both the
+// client and the consumer group.
+func newSaramaConfig(cfg Config) (*sarama.Config, error) {
+	version, err := sarama.ParseKafkaVersion(cfg.KafkaConfig.Version)
+	if err != nil {
+		return nil, err
+	}
+	config := sarama.NewConfig()
+	config.Version = version
+	config.Consumer.Offsets.Initial = sarama.OffsetOldest
+	config.RackID = cfg.KafkaConfig.RackID
+
+	switch cfg.KafkaConfig.Assignor {
+	case sarama.StickyBalanceStrategyName:
+		config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategySticky()
+	case sarama.RoundRobinBalanceStrategyName:
+		config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategyRoundRobin()
+	case sarama.RangeBalanceStrategyName, "":
+		config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategyRange()
+	default:
+		return nil, fmt.Errorf("unrecognized consumer group partition assignor: %s", cfg.KafkaConfig.Assignor)
+	}
+
+	config, err = withAuthentication(*config, cfg.KafkaConfig.Authentication)
+	if err != nil {
+		return nil, fmt.Errorf("error setting up kafka authentication: %w", err)
+	}
+	return config, nil
+}
+
 func withAuthentication(cfg sarama.Config, authCfg Authentication) (*sarama.Config, error) {
 	if len(authCfg.Type) == 0 || authCfg.Type == AuthenticationTypeNone {
 		return &cfg, nil

internal/component/loki/source/internal/kafkatarget/target_syncer_test.go

@@ -330,3 +330,52 @@ func Test_withAuthentication(t *testing.T) {
 	assert.NotNil(t, saslCfg.Net.TLS.Config.RootCAs)
 	assert.NoError(t, saslCfg.Validate())
 }
+
+func Test_newSaramaConfig_RackID(t *testing.T) {
+	cfg := Config{
+		KafkaConfig: TargetConfig{
+			Brokers: []string{"localhost:9092"},
+			Topics:  []string{"topic1"},
+			Version: "2.2.1",
+			RackID:  "eu-west-1a",
+		},
+	}
+
+	saramaCfg, err := newSaramaConfig(cfg)
+	require.NoError(t, err)
+	require.Equal(t, "eu-west-1a", saramaCfg.RackID)
+}
+
+func Test_newSaramaConfig_Assignor(t *testing.T) {
+	tests := map[string]struct {
+		assignor string
+		expected string
+		wantErr  bool
+	}{
+		"empty defaults to range": {assignor: "", expected: sarama.RangeBalanceStrategyName},
+		"range":                   {assignor: "range", expected: sarama.RangeBalanceStrategyName},
+		"roundrobin":              {assignor: "roundrobin", expected: sarama.RoundRobinBalanceStrategyName},
+		"sticky":                  {assignor: "sticky", expected: sarama.StickyBalanceStrategyName},
+		"unknown":                 {assignor: "nonexistent", wantErr: true},
+	}
+
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			saramaCfg, err := newSaramaConfig(Config{
+				KafkaConfig: TargetConfig{
+					Brokers:  []string{"localhost:9092"},
+					Topics:   []string{"topic1"},
+					Version:  "2.2.1",
+					Assignor: tt.assignor,
+				},
+			})
+
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, saramaCfg.Consumer.Group.Rebalance.Strategy.Name())
+		})
+	}
+}

internal/component/loki/source/kafka/kafka.go

@@ -37,6 +37,7 @@ type Arguments struct {
 	GroupID              string              `alloy:"group_id,attr,optional"`
 	Assignor             string              `alloy:"assignor,attr,optional"`
 	Version              string              `alloy:"version,attr,optional"`
+	RackID               string              `alloy:"rack_id,attr,optional"`
 	Authentication       KafkaAuthentication `alloy:"authentication,block,optional"`
 	UseIncomingTimestamp bool                `alloy:"use_incoming_timestamp,attr,optional"`
 	Labels               map[string]string   `alloy:"labels,attr,optional"`
@@ -177,6 +178,7 @@ func (args *Arguments) Convert() kt.Config {
 			Topics:               args.Topics,
 			Version:              args.Version,
 			Assignor:             args.Assignor,
+			RackID:               args.RackID,
 			Authentication:       args.Authentication.Convert(),
 		},
 		RelabelConfigs: alloy_relabel.ComponentToPromRelabelConfigs(args.RelabelRules),

internal/component/loki/source/kafka/kafka_test.go

@@ -84,3 +84,18 @@ func TestSASLOAuthAlloyConfig(t *testing.T) {
 	err := syntax.Unmarshal([]byte(exampleAlloyConfig), &args)
 	require.NoError(t, err)
 }
+
+func TestRackIDAlloyConfig(t *testing.T) {
+	var exampleAlloyConfig = `
+	brokers    = ["localhost:9092", "localhost:23456"]
+	topics     = ["quickstart-events"]
+	rack_id    = "eu-west-1a"
+	forward_to = []
+`
+
+	var args Arguments
+	err := syntax.Unmarshal([]byte(exampleAlloyConfig), &args)
+	require.NoError(t, err)
+	require.Equal(t, "eu-west-1a", args.RackID)
+	require.Equal(t, "eu-west-1a", args.Convert().KafkaConfig.RackID)
+}