feat: implement KIP-392 (consumer rack-aware fetching)

batch.go

@@ -28,7 +28,11 @@ type Batch struct {
 	partition     int
 	offset        int64
 	highWaterMark int64
-	err           error
+	// preferredReadReplica is the broker id (>=0) the broker would prefer
+	// the consumer fetch from for this partition (KIP-392). -1 means no
+	// preference (fetch from leader).
+	preferredReadReplica int32
+	err                  error
 	// The last offset in the batch.
 	//
 	// We use lastOffset to skip offsets that have been compacted away.
@@ -51,6 +55,17 @@ func (batch *Batch) HighWaterMark() int64 {
 	return batch.highWaterMark
 }
 
+// PreferredReadReplica returns the broker id of the replica the kafka broker
+// recommends the consumer fetch from for this partition, as introduced in
+// KIP-392 (Allow Consumers to Fetch from Closest Replica). A value of -1
+// means no preference and the consumer should keep fetching from the
+// partition leader. The value is only populated when the connection
+// negotiated Fetch v11 or higher and the broker chose a preferred replica;
+// otherwise it is -1.
+func (batch *Batch) PreferredReadReplica() int32 {
+	return batch.preferredReadReplica
+}
+
 // Partition returns the batch partition.
 func (batch *Batch) Partition() int {
 	return batch.partition

conn.go

@@ -54,6 +54,7 @@ type Conn struct {
 	fetchMinSize  int32
 	broker        int32
 	rack          string
+	clientRack    string
 
 	// correlation ID generator (synchronized on wlock)
 	correlationID int32
@@ -91,6 +92,13 @@ type ConnConfig struct {
 	Broker    int
 	Rack      string
 
+	// ClientRack is the consumer's rack id used by the broker to determine
+	// the closest replica to fetch from (KIP-392). When set and the broker
+	// supports Fetch v11 or higher, the rack id is sent in fetch requests
+	// and the broker may direct the consumer to read from a follower
+	// replica that is closer to the consumer.
+	ClientRack string
+
 	// The transactional id to use for transactional delivery. Idempotent
 	// deliver should be enabled if transactional id is configured.
 	// For more details look at transactional.id description here: http://kafka.apache.org/documentation.html#producerconfigs
@@ -179,6 +187,7 @@ func NewConnWith(conn net.Conn, config ConnConfig) *Conn {
 		partition:       int32(config.Partition),
 		broker:          int32(config.Broker),
 		rack:            config.Rack,
+		clientRack:      config.ClientRack,
 		offset:          FirstOffset,
 		requiredAcks:    -1,
 		transactionalID: emptyToNullable(config.TransactionalID),
@@ -777,7 +786,7 @@ func (c *Conn) ReadBatchWith(cfg ReadBatchConfig) *Batch {
 		return &Batch{err: dontExpectEOF(err)}
 	}
 
-	fetchVersion, err := c.negotiateVersion(fetch, v2, v5, v10)
+	fetchVersion, err := c.negotiateVersion(fetch, v2, v5, v10, v11)
 	if err != nil {
 		return &Batch{err: dontExpectEOF(err)}
 	}
@@ -799,6 +808,19 @@ func (c *Conn) ReadBatchWith(cfg ReadBatchConfig) *Batch {
 		// truncated messages.
 		adjustedDeadline = deadline
 		switch fetchVersion {
+		case v11:
+			return c.wb.writeFetchRequestV11(
+				id,
+				c.clientID,
+				c.topic,
+				c.partition,
+				offset,
+				cfg.MinBytes,
+				cfg.MaxBytes+int(c.fetchMinSize),
+				timeout,
+				int8(cfg.IsolationLevel),
+				c.clientRack,
+			)
 		case v10:
 			return c.wb.writeFetchRequestV10(
 				id,
@@ -848,8 +870,11 @@ func (c *Conn) ReadBatchWith(cfg ReadBatchConfig) *Batch {
 	var throttle int32
 	var highWaterMark int64
 	var remain int
+	var preferredReadReplica int32 = -1
 
 	switch fetchVersion {
+	case v11:
+		throttle, highWaterMark, preferredReadReplica, remain, err = readFetchResponseHeaderV11(&c.rbuf, size)
 	case v10:
 		throttle, highWaterMark, remain, err = readFetchResponseHeaderV10(&c.rbuf, size)
 	case v5:
@@ -874,15 +899,16 @@ func (c *Conn) ReadBatchWith(cfg ReadBatchConfig) *Batch {
 	}
 
 	return &Batch{
-		conn:          c,
-		msgs:          msgs,
-		deadline:      adjustedDeadline,
-		throttle:      makeDuration(throttle),
-		lock:          lock,
-		topic:         c.topic,          // topic is copied to Batch to prevent race with Batch.close
-		partition:     int(c.partition), // partition is copied to Batch to prevent race with Batch.close
-		offset:        offset,
-		highWaterMark: highWaterMark,
+		conn:                 c,
+		msgs:                 msgs,
+		deadline:             adjustedDeadline,
+		throttle:             makeDuration(throttle),
+		lock:                 lock,
+		topic:                c.topic,          // topic is copied to Batch to prevent race with Batch.close
+		partition:            int(c.partition), // partition is copied to Batch to prevent race with Batch.close
+		offset:               offset,
+		highWaterMark:        highWaterMark,
+		preferredReadReplica: preferredReadReplica,
 		// there shouldn't be a short read on initially setting up the batch.
 		// as such, any io.EOF is re-mapped to an io.ErrUnexpectedEOF so that we
 		// don't accidentally signal that we successfully reached the end of the

dialer.go

@@ -89,6 +89,11 @@ type Dialer struct {
 	// For more details look at transactional.id description here: http://kafka.apache.org/documentation.html#producerconfigs
 	// Empty string means that the connection will be non-transactional.
 	TransactionalID string
+
+	// ClientRack is the consumer's rack id (KIP-392). When set and the broker
+	// supports Fetch v11+, the broker may direct the consumer to fetch from
+	// the closest replica instead of the partition leader.
+	ClientRack string
 }
 
 // Dial connects to the address on the named network.
@@ -117,6 +122,7 @@ func (d *Dialer) DialContext(ctx context.Context, network string, address string
 		ConnConfig{
 			ClientID:        d.ClientID,
 			TransactionalID: d.TransactionalID,
+			ClientRack:      d.ClientRack,
 		},
 	)
 }
@@ -148,6 +154,7 @@ func (d *Dialer) DialPartition(ctx context.Context, network string, address stri
 		Broker:          partition.Leader.ID,
 		Rack:            partition.Leader.Rack,
 		TransactionalID: d.TransactionalID,
+		ClientRack:      d.ClientRack,
 	})
 }
 

fetch.go

@@ -37,6 +37,14 @@ type FetchRequest struct {
 	// This field requires the kafka broker to support the Fetch API in version
 	// 4 or above (otherwise the value is ignored).
 	IsolationLevel IsolationLevel
+
+	// RackID is the consumer's rack id (KIP-392). When set, the broker may
+	// direct the consumer to fetch from the closest replica via the
+	// PreferredReadReplica field of the response.
+	//
+	// This field requires the kafka broker to support the Fetch API in
+	// version 11 or above (otherwise the value is ignored).
+	RackID string
 }
 
 // FetchResponse represents a response from a kafka broker to a fetch request.
@@ -60,6 +68,13 @@ type FetchResponse struct {
 	LastStableOffset int64
 	LogStartOffset   int64
 
+	// PreferredReadReplica is the broker id the broker would prefer the
+	// consumer fetch from for this partition (KIP-392). A value of -1 means
+	// no preference. Only populated when the broker supports Fetch v11 or
+	// above AND the caller set RackID on the request; in any other case
+	// the value is forced to -1.
+	PreferredReadReplica int32
+
 	// An error that may have occurred while attempting to fetch the records.
 	//
 	// The error contains both the kafka error code, and an error message
@@ -145,6 +160,7 @@ func (c *Client) Fetch(ctx context.Context, req *FetchRequest) (*FetchResponse,
 				PartitionMaxBytes:  int32(req.MaxBytes),
 			}},
 		}},
+		RackID: req.RackID,
 	})
 
 	if err != nil {
@@ -162,14 +178,26 @@ func (c *Client) Fetch(ctx context.Context, req *FetchRequest) (*FetchResponse,
 	partition := &topic.Partitions[0]
 
 	ret := &FetchResponse{
-		Throttle:         makeDuration(res.ThrottleTimeMs),
-		Topic:            topic.Topic,
-		Partition:        int(partition.Partition),
-		Error:            makeError(res.ErrorCode, ""),
-		HighWatermark:    partition.HighWatermark,
-		LastStableOffset: partition.LastStableOffset,
-		LogStartOffset:   partition.LogStartOffset,
-		Records:          partition.RecordSet.Records,
+		Throttle:             makeDuration(res.ThrottleTimeMs),
+		Topic:                topic.Topic,
+		Partition:            int(partition.Partition),
+		Error:                makeError(res.ErrorCode, ""),
+		HighWatermark:        partition.HighWatermark,
+		LastStableOffset:     partition.LastStableOffset,
+		LogStartOffset:       partition.LogStartOffset,
+		PreferredReadReplica: partition.PreferredReadReplica,
+		Records:              partition.RecordSet.Records,
+	}
+
+	// KIP-392: PreferredReadReplica is only meaningful in Fetch v11+ responses.
+	// When the broker negotiated an older version, the protocol decoder leaves
+	// the field at Go's zero value (0), which is indistinguishable from broker
+	// id 0 being the preferred replica. We can't read the negotiated version
+	// from here, but we know that if the caller did not set RackID then they
+	// did not opt into KIP-392 and we should not surface any preferred replica
+	// id at all. Force -1 so callers can rely on "-1 means no preference".
+	if req.RackID == "" {
+		ret.PreferredReadReplica = -1
 	}
 
 	if partition.ErrorCode != 0 {

fetch_kip392_test.go

@@ -0,0 +1,107 @@
+package kafka
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/segmentio/kafka-go/protocol"
+	fetchAPI "github.com/segmentio/kafka-go/protocol/fetch"
+)
+
+// stubRoundTripper returns a canned response for any request it receives.
+type stubRoundTripper struct {
+	resp protocol.Message
+	err  error
+}
+
+func (s *stubRoundTripper) RoundTrip(_ context.Context, _ net.Addr, _ protocol.Message) (protocol.Message, error) {
+	return s.resp, s.err
+}
+
+// TestClientFetch_PreferredReadReplicaDefaultsMinusOneWithoutRackID asserts
+// that Client.Fetch surfaces PreferredReadReplica = -1 when the caller did
+// not opt into KIP-392 (req.RackID == ""). Without this guard, a v10 broker
+// response (or any response that omits the v11 field) would surface the Go
+// zero value 0, indistinguishable from broker id 0 being preferred.
+func TestClientFetch_PreferredReadReplicaDefaultsMinusOneWithoutRackID(t *testing.T) {
+	stub := &stubRoundTripper{
+		resp: &fetchAPI.Response{
+			ThrottleTimeMs: 0,
+			Topics: []fetchAPI.ResponseTopic{{
+				Topic: "topic-x",
+				Partitions: []fetchAPI.ResponsePartition{{
+					Partition:            0,
+					ErrorCode:            0,
+					HighWatermark:        100,
+					LastStableOffset:     100,
+					LogStartOffset:       0,
+					PreferredReadReplica: 0, // simulates v10 zero-value field
+				}},
+			}},
+		},
+	}
+
+	client := &Client{
+		Addr:      TCP("127.0.0.1:9092"),
+		Timeout:   time.Second,
+		Transport: stub,
+	}
+
+	resp, err := client.Fetch(context.Background(), &FetchRequest{
+		Topic:     "topic-x",
+		Partition: 0,
+		Offset:    0,
+		MinBytes:  1,
+		MaxBytes:  1024,
+		MaxWait:   100 * time.Millisecond,
+		// RackID intentionally left empty -- caller did not opt in.
+	})
+	if err != nil {
+		t.Fatalf("Client.Fetch: %v", err)
+	}
+	if resp.PreferredReadReplica != -1 {
+		t.Fatalf("PreferredReadReplica: got %d, want -1 when RackID is empty", resp.PreferredReadReplica)
+	}
+}
+
+// TestClientFetch_PreferredReadReplicaPassThroughWithRackID asserts that when
+// the caller did opt into KIP-392 by setting RackID, broker id 0 is passed
+// through verbatim (broker 0 may legitimately be the preferred replica).
+func TestClientFetch_PreferredReadReplicaPassThroughWithRackID(t *testing.T) {
+	stub := &stubRoundTripper{
+		resp: &fetchAPI.Response{
+			Topics: []fetchAPI.ResponseTopic{{
+				Topic: "topic-x",
+				Partitions: []fetchAPI.ResponsePartition{{
+					Partition:            0,
+					HighWatermark:        100,
+					PreferredReadReplica: 7,
+				}},
+			}},
+		},
+	}
+
+	client := &Client{
+		Addr:      TCP("127.0.0.1:9092"),
+		Timeout:   time.Second,
+		Transport: stub,
+	}
+
+	resp, err := client.Fetch(context.Background(), &FetchRequest{
+		Topic:     "topic-x",
+		Partition: 0,
+		Offset:    0,
+		MinBytes:  1,
+		MaxBytes:  1024,
+		MaxWait:   100 * time.Millisecond,
+		RackID:    "rack-nl",
+	})
+	if err != nil {
+		t.Fatalf("Client.Fetch: %v", err)
+	}
+	if resp.PreferredReadReplica != 7 {
+		t.Fatalf("PreferredReadReplica: got %d, want 7 (passed through verbatim)", resp.PreferredReadReplica)
+	}
+}