feat: exponential backoff for clients

Signed-off-by: Wenli Wan <[email protected]>

Author: wanwenli

async_producer_test.go

@@ -638,6 +638,70 @@ func TestAsyncProducerMultipleRetriesWithBackoffFunc(t *testing.T) {
 	}
 }
 
+func TestAsyncProducerWithExponentialBackoffDurations(t *testing.T) {
+	Logger = log.New(os.Stdout, "[sarama] ", log.LstdFlags)
+	var backoffDurations []time.Duration
+	var mu sync.Mutex
+
+	topic := "my_topic"
+	maxBackoff := 2 * time.Second
+	config := NewTestConfig()
+
+	innerBackoffFunc := NewExponentialBackoff(100*time.Millisecond, maxBackoff)
+	backoffFunc := func(retries, maxRetries int) time.Duration {
+		duration := innerBackoffFunc(retries, maxRetries)
+		mu.Lock()
+		backoffDurations = append(backoffDurations, duration)
+		mu.Unlock()
+		return duration
+	}
+
+	config.Producer.Flush.Messages = 5
+	config.Producer.Return.Successes = true
+	config.Producer.Retry.Max = 3
+	config.Producer.Retry.BackoffFunc = backoffFunc
+
+	broker := NewMockBroker(t, 1)
+
+	metadataResponse := new(MetadataResponse)
+	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
+	metadataResponse.AddTopicPartition(topic, 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
+	broker.Returns(metadataResponse)
+
+	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	failResponse := new(ProduceResponse)
+	failResponse.AddTopicPartition(topic, 0, ErrNotLeaderForPartition)
+	successResponse := new(ProduceResponse)
+	successResponse.AddTopicPartition(topic, 0, ErrNoError)
+
+	broker.Returns(failResponse)
+	broker.Returns(metadataResponse)
+	broker.Returns(failResponse)
+	broker.Returns(metadataResponse)
+	broker.Returns(successResponse)
+
+	for i := 0; i < 5; i++ {
+		producer.Input() <- &ProducerMessage{Topic: topic, Value: StringEncoder("test")}
+	}
+
+	expectResults(t, producer, 5, 0)
+	closeProducer(t, producer)
+	broker.Close()
+
+	for i := 1; i < len(backoffDurations); i++ {
+		if backoffDurations[i] < backoffDurations[i-1] {
+			t.Errorf("expected backoff[%d] >= backoff[%d], got %v < %v", i, i-1, backoffDurations[i], backoffDurations[i-1])
+		}
+		if backoffDurations[i] > maxBackoff {
+			t.Errorf("backoff exceeded max: %v", backoffDurations[i])
+		}
+	}
+}
+
 // https://github.com/IBM/sarama/issues/2129
 func TestAsyncProducerMultipleRetriesWithConcurrentRequests(t *testing.T) {
 	// Logger = log.New(os.Stdout, "[sarama] ", log.LstdFlags)

utils.go

@@ -3,8 +3,10 @@ package sarama
 import (
 	"bufio"
 	"fmt"
+	"math/rand"
 	"net"
 	"regexp"
+	"time"
 )
 
 type none struct{}
@@ -344,3 +346,34 @@ func (v KafkaVersion) String() string {
 
 	return fmt.Sprintf("%d.%d.%d", v.version[0], v.version[1], v.version[2])
 }
+
+// NewExponentialBackoff returns a function that implements an exponential backoff strategy with jitter.
+// It follows KIP-580, implementing the formula:
+// MIN(retry.backoff.max.ms, (retry.backoff.ms * 2**(failures - 1)) * random(0.8, 1.2))
+// This ensures retries start with `backoff` and exponentially increase until `maxBackoff`, with added jitter.
+//
+// Example usage:
+//
+//	backoffFunc := sarama.NewExponentialBackoff(config.Producer.Retry.Backoff, 2*time.Second)
+//	config.Producer.Retry.BackoffFunc = backoffFunc
+func NewExponentialBackoff(backoff time.Duration, maxBackoff time.Duration) func(retries, maxRetries int) time.Duration {
+	backoff = max(backoff, 0)
+	maxBackoff = max(maxBackoff, 0)
+
+	if backoff > maxBackoff {
+		Logger.Println("Warning: backoff is greater than maxBackoff, using maxBackoff instead.")
+		backoff = maxBackoff
+	}
+
+	return func(retries, maxRetries int) time.Duration {
+		if retries <= 0 {
+			return backoff
+		}
+
+		calculatedBackoff := backoff * time.Duration(1<<(retries-1))
+		jitter := 0.8 + 0.4*rand.Float64()
+		calculatedBackoff = time.Duration(float64(calculatedBackoff) * jitter)
+
+		return min(calculatedBackoff, maxBackoff)
+	}
+}

utils_test.go

@@ -2,7 +2,11 @@
 
 package sarama
 
-import "testing"
+import (
+	"sync"
+	"testing"
+	"time"
+)
 
 func TestVersionCompare(t *testing.T) {
 	if V0_8_2_0.IsAtLeast(V0_8_2_1) {
@@ -95,3 +99,48 @@ func TestVersionParsing(t *testing.T) {
 		}
 	}
 }
+
+func TestExponentialBackoffCorrectness(t *testing.T) {
+	testCases := []struct {
+		backoff            time.Duration
+		maxBackoff         time.Duration
+		retries            int
+		maxRetries         int
+		minBackoff         time.Duration
+		maxBackoffExpected time.Duration
+	}{
+		{100 * time.Millisecond, 2 * time.Second, 0, 5, 100 * time.Millisecond, 100 * time.Millisecond},
+		{100 * time.Millisecond, 2 * time.Second, 1, 5, 80 * time.Millisecond, 120 * time.Millisecond},
+		{100 * time.Millisecond, 2 * time.Second, 3, 5, 320 * time.Millisecond, 480 * time.Millisecond},
+		{100 * time.Millisecond, 2 * time.Second, 5, 5, 1280 * time.Millisecond, 1920 * time.Millisecond},
+		{-100 * time.Millisecond, 2 * time.Second, 3, 5, 0, 480 * time.Millisecond},
+		{100 * time.Millisecond, -2 * time.Second, 3, 5, 0, 0},
+		{-100 * time.Millisecond, -2 * time.Second, 3, 5, 0, 0},
+		{0 * time.Millisecond, 2 * time.Second, 3, 5, 0, 480 * time.Millisecond},
+		{100 * time.Millisecond, 0 * time.Second, 3, 5, 0, 0},
+		{0 * time.Millisecond, 0 * time.Second, 3, 5, 0, 0},
+	}
+
+	for _, tc := range testCases {
+		backoffFunc := NewExponentialBackoff(tc.backoff, tc.maxBackoff)
+		backoff := backoffFunc(tc.retries, tc.maxRetries)
+		if backoff < tc.minBackoff || backoff > tc.maxBackoffExpected {
+			t.Errorf("backoff(%d, %d): expected between %v and %v, got %v", tc.retries, tc.maxRetries, tc.minBackoff, tc.maxBackoffExpected, backoff)
+		}
+	}
+}
+
+func TestExponentialBackoffRaceDetection(t *testing.T) {
+	backoffFunc := NewExponentialBackoff(100*time.Millisecond, 2*time.Second)
+	var wg sync.WaitGroup
+	concurrency := 1000
+
+	wg.Add(concurrency)
+	for i := 0; i < concurrency; i++ {
+		go func(i int) {
+			defer wg.Done()
+			_ = backoffFunc(i%10, 5)
+		}(i)
+	}
+	wg.Wait()
+}