feat: refresh cluster metadata asynchronously on stale errors

metadata_refresh.go

@@ -0,0 +1,89 @@
+package kafka
+
+import (
+	"errors"
+
+	"github.com/segmentio/kafka-go/protocol"
+	fetchAPI "github.com/segmentio/kafka-go/protocol/fetch"
+	produceAPI "github.com/segmentio/kafka-go/protocol/produce"
+)
+
+// isStaleMetadataError reports whether a kafka error code indicates that the
+// cached cluster metadata is likely out of date and should be refreshed.
+func isStaleMetadataError(err Error) bool {
+	switch err {
+	case UnknownTopicOrPartition,
+		LeaderNotAvailable,
+		NotLeaderForPartition,
+		ReplicaNotAvailable,
+		BrokerNotAvailable,
+		KafkaStorageError,
+		FencedLeaderEpoch,
+		UnknownLeaderEpoch,
+		UnknownTopicID:
+		return true
+	default:
+		return false
+	}
+}
+
+// errorRequiresMetadataRefresh reports whether a transport-level error returned
+// by a failed round trip likely indicates stale cluster metadata: transient
+// network errors, routing failures over a stale layout (unknown topic, partition
+// or leader), or kafka errors that signal the metadata is out of date.
+func errorRequiresMetadataRefresh(err error) bool {
+	if isTransientNetworkError(err) {
+		return true
+	}
+	// Routing failures raised while resolving the target broker from the cached
+	// cluster layout indicate the layout is stale and should be refreshed.
+	if errors.Is(err, protocol.ErrNoTopic) ||
+		errors.Is(err, protocol.ErrNoPartition) ||
+		errors.Is(err, protocol.ErrNoLeader) {
+		return true
+	}
+	var kafkaErr Error
+	if errors.As(err, &kafkaErr) {
+		return isStaleMetadataError(kafkaErr)
+	}
+	return false
+}
+
+// responseRequiresMetadataRefresh inspects a successful response body and
+// reports whether any partition (or, for Fetch, the session-level error code)
+// reported a stale-metadata error. Only response types routed to specific
+// partition leaders are inspected, since those are impacted by leader changes.
+func responseRequiresMetadataRefresh(r Response) bool {
+	switch resp := r.(type) {
+	case *produceAPI.Response:
+		if resp == nil {
+			return false
+		}
+		for i := range resp.Topics {
+			partitions := resp.Topics[i].Partitions
+			for j := range partitions {
+				if isStaleMetadataError(Error(partitions[j].ErrorCode)) {
+					return true
+				}
+			}
+		}
+	case *fetchAPI.Response:
+		if resp == nil {
+			return false
+		}
+		// Fetch v7+ may report a session-level error in addition to the
+		// per-partition error codes.
+		if isStaleMetadataError(Error(resp.ErrorCode)) {
+			return true
+		}
+		for i := range resp.Topics {
+			partitions := resp.Topics[i].Partitions
+			for j := range partitions {
+				if isStaleMetadataError(Error(partitions[j].ErrorCode)) {
+					return true
+				}
+			}
+		}
+	}
+	return false
+}

metadata_refresh_test.go

@@ -0,0 +1,119 @@
+package kafka
+
+import (
+	"io"
+	"testing"
+
+	fetchAPI "github.com/segmentio/kafka-go/protocol/fetch"
+	meta "github.com/segmentio/kafka-go/protocol/metadata"
+	produceAPI "github.com/segmentio/kafka-go/protocol/produce"
+)
+
+func TestIsStaleMetadataError(t *testing.T) {
+	staleErrors := []Error{
+		UnknownTopicOrPartition,
+		LeaderNotAvailable,
+		NotLeaderForPartition,
+		ReplicaNotAvailable,
+		BrokerNotAvailable,
+		KafkaStorageError,
+		FencedLeaderEpoch,
+		UnknownLeaderEpoch,
+		UnknownTopicID,
+	}
+
+	for _, err := range staleErrors {
+		if !isStaleMetadataError(err) {
+			t.Errorf("expected %v (%d) to be classified as a stale metadata error", err, int(err))
+		}
+	}
+
+	nonStaleErrors := []Error{
+		Unknown,
+		InvalidMessage,
+		RequestTimedOut,
+		MessageSizeTooLarge,
+		TopicAuthorizationFailed,
+		InvalidRequiredAcks,
+	}
+
+	for _, err := range nonStaleErrors {
+		if isStaleMetadataError(err) {
+			t.Errorf("did not expect %v (%d) to be classified as a stale metadata error", err, int(err))
+		}
+	}
+}
+
+func TestErrorRequiresMetadataRefresh(t *testing.T) {
+	tests := []struct {
+		scenario string
+		err      error
+		want     bool
+	}{
+		{scenario: "nil", err: nil, want: false},
+		{scenario: "transient network", err: io.ErrUnexpectedEOF, want: true},
+		{scenario: "broker not available", err: BrokerNotAvailable, want: true},
+		{scenario: "not leader for partition", err: NotLeaderForPartition, want: true},
+		{scenario: "non retriable kafka error", err: TopicAuthorizationFailed, want: false},
+	}
+
+	for _, test := range tests {
+		t.Run(test.scenario, func(t *testing.T) {
+			if got := errorRequiresMetadataRefresh(test.err); got != test.want {
+				t.Errorf("errorRequiresMetadataRefresh(%v) = %v, want %v", test.err, got, test.want)
+			}
+		})
+	}
+}
+
+func TestResponseRequiresMetadataRefresh(t *testing.T) {
+	produceWith := func(code Error) *produceAPI.Response {
+		return &produceAPI.Response{
+			Topics: []produceAPI.ResponseTopic{{
+				Topic: "topic",
+				Partitions: []produceAPI.ResponsePartition{{
+					Partition: 0,
+					ErrorCode: int16(code),
+				}},
+			}},
+		}
+	}
+
+	fetchWith := func(topLevel, partition Error) *fetchAPI.Response {
+		return &fetchAPI.Response{
+			ErrorCode: int16(topLevel),
+			Topics: []fetchAPI.ResponseTopic{{
+				Topic: "topic",
+				Partitions: []fetchAPI.ResponsePartition{{
+					Partition: 0,
+					ErrorCode: int16(partition),
+				}},
+			}},
+		}
+	}
+
+	tests := []struct {
+		scenario string
+		resp     Response
+		want     bool
+	}{
+		{scenario: "produce stale", resp: produceWith(NotLeaderForPartition), want: true},
+		{scenario: "produce ok", resp: produceWith(0), want: false},
+		{scenario: "produce non-stale error", resp: produceWith(InvalidMessageSize), want: false},
+		{scenario: "fetch partition stale", resp: fetchWith(0, LeaderNotAvailable), want: true},
+		{scenario: "fetch session stale", resp: fetchWith(FencedLeaderEpoch, 0), want: true},
+		{scenario: "fetch ok", resp: fetchWith(0, 0), want: false},
+		{scenario: "unrelated response", resp: &meta.Response{}, want: false},
+		{scenario: "nil interface", resp: nil, want: false},
+		{scenario: "typed nil produce", resp: (*produceAPI.Response)(nil), want: false},
+		{scenario: "typed nil fetch", resp: (*fetchAPI.Response)(nil), want: false},
+	}
+
+	for _, test := range tests {
+		t.Run(test.scenario, func(t *testing.T) {
+			if got := responseRequiresMetadataRefresh(test.resp); got != test.want {
+				t.Errorf("responseRequiresMetadataRefresh(%T) = %v, want %v", test.resp, got, test.want)
+			}
+		})
+	}
+}

transport.go

@@ -296,6 +296,9 @@ type connPool struct {
 	ready  event      // triggered after the first metadata update
 	wake   chan event // used to force metadata updates
 	cancel context.CancelFunc
+	// Unix-nanos timestamp of the last error-triggered metadata refresh
+	// request, used to throttle refreshes under cascading failures.
+	lastMetadataRefresh int64
 	// Mutable fields of the connection pool, access must be synchronized.
 	mutex sync.RWMutex
 	conns map[int32]*connGroup // data connections used for produce/fetch/etc...
@@ -398,9 +401,20 @@ func (p *connPool) roundTrip(ctx context.Context, req Request) (Response, error)
 
 	r, err := response.await(ctx)
 	if err != nil {
+		// A communication or routing error likely means the cached cluster
+		// view is stale; refresh it so retries reach the right brokers.
+		if errorRequiresMetadataRefresh(err) {
+			p.requestMetadataUpdate()
+		}
 		return r, err
 	}
 
+	// A successful response may still report per-partition errors that
+	// indicate stale metadata (e.g. the partition leader moved).
+	if responseRequiresMetadataRefresh(r) {
+		p.requestMetadataUpdate()
+	}
+
 	switch resp := r.(type) {
 	case *createtopics.Response:
 		// Force an update of the metadata when adding topics,
@@ -442,6 +456,31 @@ func (p *connPool) roundTrip(ctx context.Context, req Request) (Response, error)
 	return r, nil
 }
 
+// metadataRefreshThrottle is the minimum interval between two error-triggered
+// metadata refreshes. It bounds the load on the cluster when many round trips
+// fail in cascade.
+const metadataRefreshThrottle = time.Second
+
+// requestMetadataUpdate triggers an asynchronous refresh of the cached cluster
+// metadata without blocking the caller. The discover goroutine performs the
+// refresh on its next iteration. Requests are coalesced and throttled to at
+// most one per metadataRefreshThrottle to avoid storming the brokers under
+// cascading failures.
+func (p *connPool) requestMetadataUpdate() {
+	now := time.Now().UnixNano()
+	last := atomic.LoadInt64(&p.lastMetadataRefresh)
+	if now-last < int64(metadataRefreshThrottle) {
+		return
+	}
+	if !atomic.CompareAndSwapInt64(&p.lastMetadataRefresh, last, now) {
+		return
+	}
+	select {
+	case p.wake <- make(event):
+	default:
+	}
+}
+
 // refreshMetadata forces an update of the cached cluster metadata, and waits
 // for the given list of topics to appear. This waiting mechanism is necessary
 // to account for the fact that topic creation is asynchronous in kafka, and