async_producer_test.go

//go:build !functional

package sarama

import (
	"errors"
	"log"
	"math"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/goleak"
)

func expectResultsWithTimeout(t *testing.T, p AsyncProducer, successCount, errorCount int, timeout time.Duration) {
	t.Helper()
	expect := successCount + errorCount
	defer func() {
		if successCount != 0 || errorCount != 0 {
			t.Error("Unexpected successes", successCount, "or errors", errorCount)
		}
	}()
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	for expect > 0 {
		select {
		case <-timer.C:
			return
		case msg := <-p.Errors():
			if msg.Msg.flags != 0 {
				t.Error("Message had flags set")
			}
			errorCount--
			expect--
			if errorCount < 0 {
				t.Error(msg.Err)
			}
		case msg := <-p.Successes():
			if msg.flags != 0 {
				t.Error("Message had flags set")
			}
			successCount--
			expect--
			if successCount < 0 {
				t.Error("Too many successes")
			}
		}
	}
}

func expectResults(t *testing.T, p AsyncProducer, successCount, errorCount int) {
	expectResultsWithTimeout(t, p, successCount, errorCount, 5*time.Minute)
}

func TestPartitionProducerFlushRetryBuffersAssignsSequence(t *testing.T) {
	cfg := NewTestConfig()
	cfg.Producer.Idempotent = true

	txnmgr := &transactionManager{
		producerID:      1,
		producerEpoch:   0,
		sequenceNumbers: map[string]int32{"topic-0": 1},
	}

	parent := &asyncProducer{
		conf:   cfg,
		txnmgr: txnmgr,
	}

	bp := &brokerProducer{
		input: make(chan *ProducerMessage, 1),
	}

	pp := &partitionProducer{
		parent:         parent,
		topic:          "topic",
		partition:      0,
		brokerProducer: bp,
		retryState:     make([]partitionRetryState, 1),
		highWatermark:  1,
	}

	msg := &ProducerMessage{Topic: "topic", Partition: 0}
	pp.retryState[0].buf = []*ProducerMessage{msg}

	pp.flushRetryBuffers()

	select {
	case flushed := <-bp.input:
		require.True(t, flushed.hasSequence, "message should have a sequence assigned")
		require.Equal(t, int32(1), flushed.sequenceNumber, "sequence number should have increased")
		require.Equal(t, txnmgr.producerEpoch, flushed.producerEpoch, "producer epoch should be the same")
	default:
		t.Fatal("expected buffered message to flush")
	}
}

type testPartitioner chan *int32

func (p testPartitioner) Partition(msg *ProducerMessage, numPartitions int32) (int32, error) {
	part := <-p
	if part == nil {
		return 0, errors.New("BOOM")
	}

	return *part, nil
}

func (p testPartitioner) RequiresConsistency() bool {
	return true
}

func (p testPartitioner) feed(partition int32) {
	p <- &partition
}

type flakyEncoder bool

func (f flakyEncoder) Length() int {
	return len(TestMessage)
}

func (f flakyEncoder) Encode() ([]byte, error) {
	if !f {
		return nil, errors.New("flaky encoding error")
	}
	return []byte(TestMessage), nil
}

func TestAsyncProducer(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for i := range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Metadata: i}
	}
	for i := range 10 {
		select {
		case msg := <-producer.Errors():
			t.Error(msg.Err)
			if msg.Msg.flags != 0 {
				t.Error("Message had flags set")
			}
		case msg := <-producer.Successes():
			if msg.flags != 0 {
				t.Error("Message had flags set")
			}
			if msg.Metadata.(int) != i {
				t.Error("Message metadata did not match")
			}
		case <-time.After(time.Second):
			t.Errorf("Timeout waiting for msg #%d", i)
			goto done
		}
	}
done:
	closeProducer(t, producer)
	leader.Close()
	seedBroker.Close()
}

func TestAsyncProducerMultipleFlushes(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)
	leader.Returns(prodSuccess)
	leader.Returns(prodSuccess)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 5
	config.Producer.Return.Successes = true
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 3 {
		for range 5 {
			producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
		}
		expectResults(t, producer, 5, 0)
	}

	closeProducer(t, producer)
	leader.Close()
	seedBroker.Close()
}

func TestAsyncProducerMultipleBrokers(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader0 := NewMockBroker(t, 2)
	leader1 := NewMockBroker(t, 3)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader0.Addr(), leader0.BrokerID())
	metadataResponse.AddBroker(leader1.Addr(), leader1.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader0.BrokerID(), nil, nil, nil, ErrNoError)
	metadataResponse.AddTopicPartition("my_topic", 1, leader1.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodResponse0 := new(ProduceResponse)
	prodResponse0.AddTopicPartition("my_topic", 0, ErrNoError)
	leader0.Returns(prodResponse0)

	prodResponse1 := new(ProduceResponse)
	prodResponse1.AddTopicPartition("my_topic", 1, ErrNoError)
	leader1.Returns(prodResponse1)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 5
	config.Producer.Return.Successes = true
	config.Producer.Partitioner = NewRoundRobinPartitioner
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	expectResults(t, producer, 10, 0)

	closeProducer(t, producer)
	leader1.Close()
	leader0.Close()
	seedBroker.Close()
}

func TestAsyncProducerCustomPartitioner(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodResponse := new(ProduceResponse)
	prodResponse.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodResponse)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 2
	config.Producer.Return.Successes = true
	config.Producer.Partitioner = func(topic string) Partitioner {
		p := make(testPartitioner)
		go func() {
			p.feed(0)
			p <- nil
			p <- nil
			p <- nil
			p.feed(0)
		}()
		return p
	}
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 5 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	expectResults(t, producer, 2, 3)

	closeProducer(t, producer)
	leader.Close()
	seedBroker.Close()
}

func TestAsyncProducerFailureRetry(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader1 := NewMockBroker(t, 2)
	leader2 := NewMockBroker(t, 3)

	metadataLeader1 := new(MetadataResponse)
	metadataLeader1.AddBroker(leader1.Addr(), leader1.BrokerID())
	metadataLeader1.AddTopicPartition("my_topic", 0, leader1.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader1)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}
	seedBroker.Close()

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	prodNotLeader := new(ProduceResponse)
	prodNotLeader.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)
	leader1.Returns(prodNotLeader)

	metadataLeader2 := new(MetadataResponse)
	metadataLeader2.AddBroker(leader2.Addr(), leader2.BrokerID())
	metadataLeader2.AddTopicPartition("my_topic", 0, leader2.BrokerID(), nil, nil, nil, ErrNoError)
	leader1.Returns(metadataLeader2)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader2.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)
	leader1.Close()

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	leader2.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)

	leader2.Close()
	closeProducer(t, producer)
}

func TestAsyncProducerRecoveryWithRetriesDisabled(t *testing.T) {
	tt := func(t *testing.T, kErr KError) {
		seedBroker := NewMockBroker(t, 0)
		broker1 := NewMockBroker(t, 1)
		broker2 := NewMockBroker(t, 2)

		mockLeader := func(leaderID int32) *MockMetadataResponse {
			return NewMockMetadataResponse(t).
				SetController(seedBroker.BrokerID()).
				SetBroker(seedBroker.Addr(), seedBroker.BrokerID()).
				SetBroker(broker1.Addr(), broker1.BrokerID()).
				SetBroker(broker2.Addr(), broker2.BrokerID()).
				SetLeader("my_topic", 0, leaderID).
				SetLeader("my_topic", 1, leaderID)
		}

		seedBroker.SetHandlerByMap(
			map[string]MockResponse{
				"MetadataRequest": mockLeader(broker1.BrokerID()),
			},
		)

		config := NewTestConfig()
		config.ClientID = "TestAsyncProducerRecoveryWithRetriesDisabled"
		config.Producer.Flush.Messages = 2
		config.Producer.Flush.Frequency = 100 * time.Millisecond
		config.Producer.Return.Successes = true
		config.Producer.Retry.Max = 0 // disable!
		config.Producer.Retry.Backoff = 0
		config.Producer.Partitioner = NewManualPartitioner
		producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
		if err != nil {
			t.Fatal(err)
		}

		broker1.SetHandlerByMap(
			map[string]MockResponse{
				"MetadataRequest": mockLeader(broker1.BrokerID()),
				"ProduceRequest": NewMockProduceResponse(t).
					SetError("my_topic", 0, kErr).
					SetError("my_topic", 1, kErr),
			},
		)

		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: 0}
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: 1}
		expectResults(t, producer, 0, 2)

		seedBroker.SetHandlerByMap(
			map[string]MockResponse{
				"MetadataRequest": mockLeader(broker2.BrokerID()),
			},
		)
		broker1.SetHandlerByMap(
			map[string]MockResponse{
				"MetadataRequest": mockLeader(broker2.BrokerID()),
			},
		)
		broker2.SetHandlerByMap(
			map[string]MockResponse{
				"MetadataRequest": mockLeader(broker2.BrokerID()),
				"ProduceRequest": NewMockProduceResponse(t).
					SetError("my_topic", 0, ErrNoError).
					SetError("my_topic", 1, ErrNoError),
			},
		)

		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: 0}
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: 1}
		expectResults(t, producer, 2, 0)

		closeProducer(t, producer)
		seedBroker.Close()
		broker1.Close()
		broker2.Close()
	}

	t.Run("retriable error", func(t *testing.T) {
		tt(t, ErrNotLeaderForPartition)
	})

	t.Run("non-retriable error", func(t *testing.T) {
		tt(t, ErrNotController)
	})
}

func TestAsyncProducerEncoderFailures(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)
	leader.Returns(prodSuccess)
	leader.Returns(prodSuccess)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 1
	config.Producer.Return.Successes = true
	config.Producer.Partitioner = NewManualPartitioner
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 3 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: flakyEncoder(true), Value: flakyEncoder(false)}
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: flakyEncoder(false), Value: flakyEncoder(true)}
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: flakyEncoder(true), Value: flakyEncoder(true)}
		expectResults(t, producer, 1, 2)
	}

	closeProducer(t, producer)
	leader.Close()
	seedBroker.Close()
}

// If a Kafka broker becomes unavailable and then returns back in service, then
// producer reconnects to it and continues sending messages.
func TestAsyncProducerBrokerBounce(t *testing.T) {
	// Given
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)
	leaderAddr := leader.Addr()

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leaderAddr, leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 1
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	leader.Returns(prodSuccess)
	expectResults(t, producer, 1, 0)

	// When: a broker connection gets reset by a broker (network glitch, restart, you name it).
	leader.Close()                               // producer should get EOF
	leader = NewMockBrokerAddr(t, 2, leaderAddr) // start it up again right away for giggles
	leader.Returns(metadataResponse)             // tell it to go to broker 2 again

	// Then: a produced message goes through the new broker connection.
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	leader.Returns(prodSuccess)
	expectResults(t, producer, 1, 0)

	closeProducer(t, producer)
	seedBroker.Close()
	leader.Close()
}

func TestAsyncProducerBrokerBounceWithStaleMetadata(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader1 := NewMockBroker(t, 2)
	leader2 := NewMockBroker(t, 3)

	metadataLeader1 := new(MetadataResponse)
	metadataLeader1.AddBroker(leader1.Addr(), leader1.BrokerID())
	metadataLeader1.AddTopicPartition("my_topic", 0, leader1.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader1)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 3
	config.Producer.Retry.Backoff = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	leader1.Close()                     // producer should get EOF
	seedBroker.Returns(metadataLeader1) // tell it to go to leader1 again even though it's still down
	seedBroker.Returns(metadataLeader1) // tell it to go to leader1 again even though it's still down

	// ok fine, tell it to go to leader2 finally
	metadataLeader2 := new(MetadataResponse)
	metadataLeader2.AddBroker(leader2.Addr(), leader2.BrokerID())
	metadataLeader2.AddTopicPartition("my_topic", 0, leader2.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader2)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader2.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)
	seedBroker.Close()
	leader2.Close()

	closeProducer(t, producer)
}

func TestAsyncProducerMultipleRetries(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader1 := NewMockBroker(t, 2)
	leader2 := NewMockBroker(t, 3)

	metadataLeader1 := new(MetadataResponse)
	metadataLeader1.AddBroker(leader1.Addr(), leader1.BrokerID())
	metadataLeader1.AddTopicPartition("my_topic", 0, leader1.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader1)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 4
	config.Producer.Retry.Backoff = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	prodNotLeader := new(ProduceResponse)
	prodNotLeader.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)
	leader1.Returns(prodNotLeader)

	metadataLeader2 := new(MetadataResponse)
	metadataLeader2.AddBroker(leader2.Addr(), leader2.BrokerID())
	metadataLeader2.AddTopicPartition("my_topic", 0, leader2.BrokerID(), nil, nil, nil, ErrNoError)

	leader1.Returns(metadataLeader2)
	leader2.Returns(prodNotLeader)
	leader2.Returns(metadataLeader1)
	leader1.Returns(prodNotLeader)
	leader1.Returns(metadataLeader1)
	leader1.Returns(prodNotLeader)
	leader1.Returns(metadataLeader2)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader2.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	leader2.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)

	seedBroker.Close()
	leader1.Close()
	leader2.Close()
	closeProducer(t, producer)
}

func TestAsyncProducerMultipleRetriesWithBackoffFunc(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader1 := NewMockBroker(t, 2)
	leader2 := NewMockBroker(t, 3)

	metadataLeader1 := new(MetadataResponse)
	metadataLeader1.AddBroker(leader1.Addr(), leader1.BrokerID())
	metadataLeader1.AddTopicPartition("my_topic", 0, leader1.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader1)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 1
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 4

	// We use a pointer to atomic to prevent the possibility of a reallocation causing a copy.
	backoffCalled := make([]*atomic.Int32, config.Producer.Retry.Max+1)
	for i := range backoffCalled {
		backoffCalled[i] = new(atomic.Int32)
	}

	config.Producer.Retry.BackoffFunc = func(retries, maxRetries int) time.Duration {
		backoffCalled[retries-1].Add(1)
		return 0
	}
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	prodNotLeader := new(ProduceResponse)
	prodNotLeader.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)

	metadataLeader2 := new(MetadataResponse)
	metadataLeader2.AddBroker(leader2.Addr(), leader2.BrokerID())
	metadataLeader2.AddTopicPartition("my_topic", 0, leader2.BrokerID(), nil, nil, nil, ErrNoError)

	leader1.Returns(prodNotLeader)
	leader1.Returns(metadataLeader2)
	leader2.Returns(prodNotLeader)
	leader2.Returns(metadataLeader1)
	leader1.Returns(prodNotLeader)
	leader1.Returns(metadataLeader1)
	leader1.Returns(prodNotLeader)
	leader1.Returns(metadataLeader2)
	leader2.Returns(prodSuccess)

	expectResults(t, producer, 1, 0)

	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	leader2.Returns(prodSuccess)
	expectResults(t, producer, 1, 0)

	seedBroker.Close()
	leader1.Close()
	leader2.Close()
	closeProducer(t, producer)

	for i := 0; i < config.Producer.Retry.Max; i++ {
		if backoffCalled[i].Load() != 1 {
			t.Errorf("expected one retry attempt #%d", i)
		}
	}
	if backoffCalled[config.Producer.Retry.Max].Load() != 0 {
		t.Errorf("expected no retry attempt #%d", config.Producer.Retry.Max)
	}
}

func TestAsyncProducerWithExponentialBackoffDurations(t *testing.T) {
	var backoffDurations []time.Duration
	var mu sync.Mutex

	topic := "my_topic"
	maxBackoff := 2 * time.Second
	config := NewTestConfig()

	innerBackoffFunc := NewExponentialBackoff(defaultRetryBackoff, 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 range 5 {
		producer.Input() <- &ProducerMessage{Topic: topic, Value: StringEncoder("test")}
	}

	expectResults(t, producer, 5, 0)
	closeProducer(t, producer)
	broker.Close()

	assert.Greater(t, backoffDurations[0], time.Duration(0),
		"Expected first backoff duration to be greater than 0")
	for i := 1; i < len(backoffDurations); i++ {
		assert.Greater(t, backoffDurations[i], time.Duration(0))
		assert.GreaterOrEqual(t, backoffDurations[i], backoffDurations[i-1])
		assert.LessOrEqual(t, backoffDurations[i], maxBackoff)
	}
}

// https://github.com/IBM/sarama/issues/2129
func TestAsyncProducerMultipleRetriesWithConcurrentRequests(t *testing.T) {
	// Logger = log.New(os.Stdout, "[sarama] ", log.LstdFlags)
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	// The seed broker only handles Metadata request
	seedBroker.setHandler(func(req *request) (res encoderWithHeader) {
		metadataLeader := new(MetadataResponse)
		metadataLeader.AddBroker(leader.Addr(), leader.BrokerID())
		metadataLeader.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
		return metadataLeader
	})

	// Simulate a slow broker by taking ~200ms to handle requests
	// therefore triggering the read timeout and the retry logic
	leader.setHandler(func(req *request) (res encoderWithHeader) {
		time.Sleep(200 * time.Millisecond)
		// Will likely not be read by the producer (read timeout)
		prodSuccess := new(ProduceResponse)
		prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
		return prodSuccess
	})

	config := NewTestConfig()
	// Use very short read to simulate read error on unresponsive broker
	config.Net.ReadTimeout = 50 * time.Millisecond
	// Flush every record to generate up to 5 in-flight Produce requests
	// because config.Net.MaxOpenRequests defaults to 5
	config.Producer.Flush.MaxMessages = 1
	config.Producer.Return.Successes = true
	// Reduce retries to speed up the test while keeping the default backoff
	config.Producer.Retry.Max = 1
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	expectResults(t, producer, 0, 10)

	seedBroker.Close()
	leader.Close()
	closeProducer(t, producer)
}

func TestAsyncProducerBrokerRestart(t *testing.T) {
	// Logger = log.New(os.Stdout, "[sarama] ", log.LstdFlags)

	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	var leaderLock sync.Mutex
	metadataRequestHandlerFunc := func(req *request) (res encoderWithHeader) {
		leaderLock.Lock()
		defer leaderLock.Unlock()
		metadataLeader := new(MetadataResponse)
		metadataLeader.AddBroker(leader.Addr(), leader.BrokerID())
		metadataLeader.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
		return metadataLeader
	}

	// The seed broker only handles Metadata request in bootstrap
	seedBroker.setHandler(metadataRequestHandlerFunc)

	var emptyValues atomic.Int32

	countRecordsWithEmptyValue := func(req *request) {
		preq := req.body.(*ProduceRequest)
		if batch := preq.records["my_topic"][0].RecordBatch; batch != nil {
			for _, record := range batch.Records {
				if len(record.Value) == 0 {
					emptyValues.Add(1)
				}
			}
		}
		if batch := preq.records["my_topic"][0].MsgSet; batch != nil {
			for _, record := range batch.Messages {
				if len(record.Msg.Value) == 0 {
					emptyValues.Add(1)
				}
			}
		}
	}

	failedProduceRequestHandlerFunc := func(req *request) (res encoderWithHeader) {
		countRecordsWithEmptyValue(req)

		time.Sleep(50 * time.Millisecond)

		prodSuccess := new(ProduceResponse)
		prodSuccess.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)
		return prodSuccess
	}

	succeededProduceRequestHandlerFunc := func(req *request) (res encoderWithHeader) {
		countRecordsWithEmptyValue(req)

		prodSuccess := new(ProduceResponse)
		prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
		return prodSuccess
	}

	leader.SetHandlerFuncByMap(map[string]requestHandlerFunc{
		"ProduceRequest":  failedProduceRequestHandlerFunc,
		"MetadataRequest": metadataRequestHandlerFunc,
	})

	config := NewTestConfig()
	config.Producer.Retry.Backoff = 250 * time.Millisecond
	config.Producer.Flush.MaxMessages = 1
	config.Producer.Return.Errors = true
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 10

	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	var wg sync.WaitGroup

	pushMsg := func() {
		defer wg.Done()
		for range 10 {
			producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
			time.Sleep(50 * time.Millisecond)
		}
	}

	wg.Add(1)
	go pushMsg()

	for range 3 {
		time.Sleep(100 * time.Millisecond)

		wg.Add(1)
		go pushMsg()
	}

	leader.Close()
	leaderLock.Lock()
	leader = NewMockBroker(t, 2)
	leaderLock.Unlock()
	leader.SetHandlerFuncByMap(map[string]requestHandlerFunc{
		"ProduceRequest":  succeededProduceRequestHandlerFunc,
		"MetadataRequest": metadataRequestHandlerFunc,
	})

	wg.Wait()

	expectResultsWithTimeout(t, producer, 40, 0, 10*time.Second)

	seedBroker.Close()
	leader.Close()

	closeProducerWithTimeout(t, producer, 5*time.Second)

	if emptyValues := emptyValues.Load(); emptyValues > 0 {
		t.Fatalf("%d empty values", emptyValues)
	}
}

func TestAsyncProducerOutOfRetries(t *testing.T) {
	t.Skip("Enable once bug #294 is fixed.")

	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	config.Producer.Retry.Max = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	prodNotLeader := new(ProduceResponse)
	prodNotLeader.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)
	leader.Returns(prodNotLeader)

	for range 10 {
		select {
		case msg := <-producer.Errors():
			if !errors.Is(msg.Err, ErrNotLeaderForPartition) {
				t.Error(msg.Err)
			}
		case <-producer.Successes():
			t.Error("Unexpected success")
		}
	}

	seedBroker.Returns(metadataResponse)

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)

	expectResults(t, producer, 10, 0)

	leader.Close()
	seedBroker.Close()
	safeClose(t, producer)
}

func TestAsyncProducerRetryWithReferenceOpen(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)
	leaderAddr := leader.Addr()

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leaderAddr, leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	metadataResponse.AddTopicPartition("my_topic", 1, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	config := NewTestConfig()
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	config.Producer.Retry.Max = 1
	config.Producer.Partitioner = NewRoundRobinPartitioner
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	// prime partition 0
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)
	expectResults(t, producer, 1, 0)

	// prime partition 1
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	prodSuccess = new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 1, ErrNoError)
	leader.Returns(prodSuccess)
	expectResults(t, producer, 1, 0)

	// reboot the broker (the producer will get EOF on its existing connection)
	leader.Close()
	leader = NewMockBrokerAddr(t, 2, leaderAddr)

	// send another message on partition 0 to trigger the EOF and retry
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}

	// tell partition 0 to go to that broker again
	leader.Returns(metadataResponse)

	// succeed this time
	prodSuccess = new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)
	expectResults(t, producer, 1, 0)

	// shutdown
	closeProducer(t, producer)
	seedBroker.Close()
	leader.Close()
}

func TestAsyncProducerFlusherRetryCondition(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	metadataResponse.AddTopicPartition("my_topic", 1, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 5
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	config.Producer.Retry.Max = 1
	config.Producer.Partitioner = NewManualPartitioner
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	// prime partitions
	for p := range int32(2) {
		for range 5 {
			producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: p}
		}
		prodSuccess := new(ProduceResponse)
		prodSuccess.AddTopicPartition("my_topic", p, ErrNoError)
		leader.Returns(prodSuccess)
		expectResults(t, producer, 5, 0)
	}

	// send more messages on partition 0
	for range 5 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: 0}
	}
	prodNotLeader := new(ProduceResponse)
	prodNotLeader.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)
	leader.Returns(prodNotLeader)

	time.Sleep(50 * time.Millisecond)

	// tell partition 0 to go to that broker again
	leader.Returns(metadataResponse)
	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)

	// succeed this time
	expectResults(t, producer, 5, 0)

	// put five more through
	for range 5 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage), Partition: 0}
	}
	prodSuccess = new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)
	expectResults(t, producer, 5, 0)

	// shutdown
	closeProducer(t, producer)
	seedBroker.Close()
	leader.Close()
}

func TestAsyncProducerRetryShutdown(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataLeader := new(MetadataResponse)
	metadataLeader.AddBroker(leader.Addr(), leader.BrokerID())
	metadataLeader.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}
	producer.AsyncClose()
	time.Sleep(5 * time.Millisecond) // let the shutdown goroutine kick in

	producer.Input() <- &ProducerMessage{Topic: "FOO"}
	if err := <-producer.Errors(); !errors.Is(err.Err, ErrShuttingDown) {
		t.Error(err)
	}

	prodNotLeader := new(ProduceResponse)
	prodNotLeader.AddTopicPartition("my_topic", 0, ErrNotLeaderForPartition)
	leader.Returns(prodNotLeader)

	leader.Returns(metadataLeader)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)

	seedBroker.Close()
	leader.Close()

	// wait for the async-closed producer to shut down fully
	for err := range producer.Errors() {
		t.Error(err)
	}
}

func TestAsyncProducerNoReturns(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataLeader := new(MetadataResponse)
	metadataLeader.AddBroker(leader.Addr(), leader.BrokerID())
	metadataLeader.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = false
	config.Producer.Return.Errors = false
	config.Producer.Retry.Backoff = 0
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	wait := make(chan bool)
	go func() {
		if err := producer.Close(); err != nil {
			t.Error(err)
		}
		close(wait)
	}()

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)

	<-wait
	seedBroker.Close()
	leader.Close()
}

func TestAsyncProducerIdempotentGoldenPath(t *testing.T) {
	broker := NewMockBroker(t, 1)

	metadataResponse := &MetadataResponse{
		Version:      4,
		ControllerID: 1,
	}
	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataResponse)

	initProducerID := &InitProducerIDResponse{
		ThrottleTime:  0,
		ProducerID:    1000,
		ProducerEpoch: 1,
	}
	broker.Returns(initProducerID)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 4
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Retry.Backoff = 0
	config.Producer.Idempotent = true
	config.Net.MaxOpenRequests = 1
	config.Version = V0_11_0_0
	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	prodSuccess := &ProduceResponse{
		Version:      3,
		ThrottleTime: 0,
	}
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	broker.Returns(prodSuccess)
	expectResults(t, producer, 10, 0)

	broker.Close()
	closeProducer(t, producer)
}

func TestAsyncProducerIdempotentRetryCheckBatch(t *testing.T) {
	// Logger = log.New(os.Stderr, "", log.LstdFlags)
	tests := []struct {
		name           string
		failAfterWrite bool
	}{
		{"FailAfterWrite", true},
		{"FailBeforeWrite", false},
	}

	for _, test := range tests {
		broker := NewMockBroker(t, 1)

		metadataResponse := &MetadataResponse{
			Version:      4,
			ControllerID: 1,
		}
		metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
		metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)

		initProducerIDResponse := &InitProducerIDResponse{
			ThrottleTime:  0,
			ProducerID:    1000,
			ProducerEpoch: 1,
		}

		prodNotLeaderResponse := &ProduceResponse{
			Version:      3,
			ThrottleTime: 0,
		}
		prodNotLeaderResponse.AddTopicPartition("my_topic", 0, ErrNotEnoughReplicas)

		prodDuplicate := &ProduceResponse{
			Version:      3,
			ThrottleTime: 0,
		}
		prodDuplicate.AddTopicPartition("my_topic", 0, ErrDuplicateSequenceNumber)

		prodOutOfSeq := &ProduceResponse{
			Version:      3,
			ThrottleTime: 0,
		}
		prodOutOfSeq.AddTopicPartition("my_topic", 0, ErrOutOfOrderSequenceNumber)

		prodSuccessResponse := &ProduceResponse{
			Version:      3,
			ThrottleTime: 0,
		}
		prodSuccessResponse.AddTopicPartition("my_topic", 0, ErrNoError)

		prodCounter := 0
		lastBatchFirstSeq := -1
		lastBatchSize := -1
		lastSequenceWrittenToDisk := -1
		handlerFailBeforeWrite := func(req *request) (res encoderWithHeader) {
			switch req.body.key() {
			case 3:
				return metadataResponse
			case 22:
				return initProducerIDResponse
			case 0:
				prodCounter++

				preq := req.body.(*ProduceRequest)
				batch := preq.records["my_topic"][0].RecordBatch
				batchFirstSeq := int(batch.FirstSequence)
				batchSize := len(batch.Records)

				if lastSequenceWrittenToDisk == batchFirstSeq-1 { // in sequence append
					if lastBatchFirstSeq == batchFirstSeq { // is a batch retry
						if lastBatchSize == batchSize { // good retry
							// mock write to disk
							lastSequenceWrittenToDisk = batchFirstSeq + batchSize - 1
							return prodSuccessResponse
						}
						t.Errorf("[%s] Retried Batch firstSeq=%d with different size old=%d new=%d", test.name, batchFirstSeq, lastBatchSize, batchSize)
						return prodOutOfSeq
					} // not a retry
					// save batch just received for future check
					lastBatchFirstSeq = batchFirstSeq
					lastBatchSize = batchSize

					if prodCounter%2 == 1 {
						if test.failAfterWrite {
							// mock write to disk
							lastSequenceWrittenToDisk = batchFirstSeq + batchSize - 1
						}
						return prodNotLeaderResponse
					}
					// mock write to disk
					lastSequenceWrittenToDisk = batchFirstSeq + batchSize - 1
					return prodSuccessResponse
				}
				if lastBatchFirstSeq == batchFirstSeq && lastBatchSize == batchSize { // is a good batch retry
					if lastSequenceWrittenToDisk == (batchFirstSeq + batchSize - 1) { // we already have the messages
						return prodDuplicate
					}
					// mock write to disk
					lastSequenceWrittenToDisk = batchFirstSeq + batchSize - 1
					return prodSuccessResponse
				} // out of sequence / bad retried batch
				if lastBatchFirstSeq == batchFirstSeq && lastBatchSize != batchSize {
					t.Errorf("[%s] Retried Batch firstSeq=%d with different size old=%d new=%d", test.name, batchFirstSeq, lastBatchSize, batchSize)
				} else if lastSequenceWrittenToDisk+1 != batchFirstSeq {
					t.Errorf("[%s] Out of sequence message lastSequence=%d new batch starts at=%d", test.name, lastSequenceWrittenToDisk, batchFirstSeq)
				} else {
					t.Errorf("[%s] Unexpected error", test.name)
				}

				return prodOutOfSeq
			}
			return nil
		}

		config := NewTestConfig()
		config.Version = V0_11_0_0
		config.Producer.Idempotent = true
		config.Net.MaxOpenRequests = 1
		config.Producer.RequiredAcks = WaitForAll
		config.Producer.Return.Successes = true
		config.Producer.Flush.Frequency = 50 * time.Millisecond
		config.Producer.Retry.Backoff = 100 * time.Millisecond

		broker.setHandler(handlerFailBeforeWrite)
		producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
		if err != nil {
			t.Fatal(err)
		}

		for range 3 {
			producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
		}

		go func() {
			for range 7 {
				producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder("goroutine")}
				time.Sleep(100 * time.Millisecond)
			}
		}()

		expectResults(t, producer, 10, 0)

		broker.Close()
		closeProducer(t, producer)
	}
}

// test case for https://github.com/IBM/sarama/pull/2378
func TestAsyncProducerIdempotentRetryCheckBatch_2378(t *testing.T) {
	broker := NewMockBroker(t, 1)

	metadataResponse := &MetadataResponse{
		Version:      4,
		ControllerID: 1,
	}
	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)

	initProducerIDResponse := &InitProducerIDResponse{
		ThrottleTime:  0,
		ProducerID:    1000,
		ProducerEpoch: 1,
	}

	prodNotLeaderResponse := &ProduceResponse{
		Version:      3,
		ThrottleTime: 0,
	}
	prodNotLeaderResponse.AddTopicPartition("my_topic", 0, ErrNotEnoughReplicas)

	handlerFailBeforeWrite := func(req *request) (res encoderWithHeader) {
		switch req.body.key() {
		case 3:
			return metadataResponse
		case 22:
			return initProducerIDResponse
		case 0: // for msg, always return error to trigger retryBatch
			return prodNotLeaderResponse
		}
		return nil
	}

	config := NewTestConfig()
	config.Version = V0_11_0_0
	config.Producer.Idempotent = true
	config.Net.MaxOpenRequests = 1
	config.Producer.Retry.Max = 1 // set max retry to 1
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Return.Successes = true
	config.Producer.Flush.Frequency = 50 * time.Millisecond
	config.Producer.Retry.Backoff = 100 * time.Millisecond

	broker.setHandler(handlerFailBeforeWrite)
	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 3 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	go func() {
		for range 7 {
			producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder("goroutine")}
			time.Sleep(100 * time.Millisecond)
		}
	}()

	// this will block until 5 minutes timeout before pr 2378 merge
	expectResults(t, producer, 0, 10)

	broker.Close()
	closeProducer(t, producer)
}

// test case for https://github.com/IBM/sarama/issues/2469: idempotent producer
// retries via retryBatch must honor Retry.Backoff / Retry.BackoffFunc when the
// broker repeatedly returns a retriable error.
func TestAsyncProducerIdempotentRetryBatchBackoff(t *testing.T) {
	broker := NewMockBroker(t, 1)

	metadataResponse := &MetadataResponse{
		Version:      4,
		ControllerID: 1,
	}
	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)

	initProducerIDResponse := &InitProducerIDResponse{
		ThrottleTime:  0,
		ProducerID:    1000,
		ProducerEpoch: 1,
	}

	prodNotLeaderResponse := &ProduceResponse{
		Version:      3,
		ThrottleTime: 0,
	}
	prodNotLeaderResponse.AddTopicPartition("my_topic", 0, ErrNotEnoughReplicas)

	handler := func(req *request) (res encoderWithHeader) {
		switch req.body.key() {
		case 3:
			return metadataResponse
		case 22:
			return initProducerIDResponse
		case 0:
			return prodNotLeaderResponse
		}
		return nil
	}

	config := NewTestConfig()
	config.Version = V0_11_0_0
	config.Producer.Idempotent = true
	config.Net.MaxOpenRequests = 1
	config.Producer.Retry.Max = 3
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Return.Successes = true
	config.Producer.Flush.Frequency = 50 * time.Millisecond

	backoffCalls := new(atomic.Int32)
	config.Producer.Retry.BackoffFunc = func(retries, maxRetries int) time.Duration {
		backoffCalls.Add(1)
		return 0
	}

	broker.setHandler(handler)
	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
	require.NoError(t, err)

	producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}

	expectResults(t, producer, 0, 1)

	broker.Close()
	closeProducer(t, producer)

	assert.GreaterOrEqual(t, int(backoffCalls.Load()), config.Producer.Retry.Max,
		"BackoffFunc should be called at least Retry.Max times during idempotent retryBatch")
}

func TestAsyncProducerIdempotentErrorOnOutOfSeq(t *testing.T) {
	broker := NewMockBroker(t, 1)

	metadataResponse := &MetadataResponse{
		Version:      4,
		ControllerID: 1,
	}
	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataResponse)

	initProducerID := &InitProducerIDResponse{
		ThrottleTime:  0,
		ProducerID:    1000,
		ProducerEpoch: 1,
	}
	broker.Returns(initProducerID)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 400000
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Retry.Backoff = 0
	config.Producer.Idempotent = true
	config.Net.MaxOpenRequests = 1
	config.Version = V0_11_0_0

	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	prodOutOfSeq := &ProduceResponse{
		Version:      3,
		ThrottleTime: 0,
	}
	prodOutOfSeq.AddTopicPartition("my_topic", 0, ErrOutOfOrderSequenceNumber)
	broker.Returns(prodOutOfSeq)
	expectResults(t, producer, 0, 10)

	broker.Close()
	closeProducer(t, producer)
}

func TestAsyncProducerIdempotentEpochRollover(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	metadataResponse := &MetadataResponse{
		Version:      4,
		ControllerID: 1,
	}
	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataResponse)

	initProducerID := &InitProducerIDResponse{
		ThrottleTime:  0,
		ProducerID:    1000,
		ProducerEpoch: 1,
	}
	broker.Returns(initProducerID)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Flush.Frequency = 10 * time.Millisecond
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 1 // This test needs to exercise what happens when retries exhaust
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Retry.Backoff = 0
	config.Producer.Idempotent = true
	config.Net.MaxOpenRequests = 1
	config.Version = V0_11_0_0

	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}
	defer closeProducer(t, producer)

	producer.Input() <- &ProducerMessage{Topic: "my_topic", Value: StringEncoder("hello")}
	prodError := &ProduceResponse{
		Version:      3,
		ThrottleTime: 0,
	}
	prodError.AddTopicPartition("my_topic", 0, ErrBrokerNotAvailable)
	broker.Returns(prodError)
	<-producer.Errors()

	lastReqRes := broker.history[len(broker.history)-1]
	lastProduceBatch := lastReqRes.Request.(*ProduceRequest).records["my_topic"][0].RecordBatch
	if lastProduceBatch.FirstSequence != 0 {
		t.Error("first sequence not zero")
	}
	if lastProduceBatch.ProducerEpoch != 1 {
		t.Error("first epoch was not one")
	}

	// Now if we produce again, the epoch should have rolled over.
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Value: StringEncoder("hello")}
	broker.Returns(prodError)
	<-producer.Errors()

	lastReqRes = broker.history[len(broker.history)-1]
	lastProduceBatch = lastReqRes.Request.(*ProduceRequest).records["my_topic"][0].RecordBatch
	if lastProduceBatch.FirstSequence != 0 {
		t.Error("second sequence not zero")
	}
	if lastProduceBatch.ProducerEpoch <= 1 {
		t.Error("second epoch was not > 1")
	}
}

// TestAsyncProducerIdempotentEpochExhaustion ensures that producer requests
// a new producerID when producerEpoch is exhausted
func TestAsyncProducerIdempotentEpochExhaustion(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	var (
		initialProducerID = int64(1000)
		newProducerID     = initialProducerID + 1
	)

	metadataResponse := &MetadataResponse{
		Version:      4,
		ControllerID: 1,
	}
	metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataResponse)

	initProducerID := &InitProducerIDResponse{
		ThrottleTime:  0,
		ProducerID:    initialProducerID,
		ProducerEpoch: math.MaxInt16, // Mock ProducerEpoch at the exhaustion point
	}
	broker.Returns(initProducerID)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Flush.Frequency = 10 * time.Millisecond
	config.Producer.Return.Successes = true
	config.Producer.Retry.Max = 1 // This test needs to exercise what happens when retries exhaust
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Retry.Backoff = 0
	config.Producer.Idempotent = true
	config.Net.MaxOpenRequests = 1
	config.Version = V0_11_0_0

	producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}
	defer closeProducer(t, producer)

	producer.Input() <- &ProducerMessage{Topic: "my_topic", Value: StringEncoder("hello")}
	prodError := &ProduceResponse{
		Version:      3,
		ThrottleTime: 0,
	}
	prodError.AddTopicPartition("my_topic", 0, ErrBrokerNotAvailable)
	broker.Returns(prodError)
	broker.Returns(&InitProducerIDResponse{
		ProducerID: newProducerID,
	})

	<-producer.Errors()

	lastProduceReqRes := broker.history[len(broker.history)-2] // last is InitProducerIDRequest
	lastProduceBatch := lastProduceReqRes.Request.(*ProduceRequest).records["my_topic"][0].RecordBatch
	if lastProduceBatch.FirstSequence != 0 {
		t.Error("first sequence not zero")
	}
	if lastProduceBatch.ProducerEpoch <= 1 {
		t.Error("first epoch was not at exhaustion point")
	}

	// Now we should produce with a new ProducerID
	producer.Input() <- &ProducerMessage{Topic: "my_topic", Value: StringEncoder("hello")}
	broker.Returns(prodError)
	<-producer.Errors()

	lastProduceReqRes = broker.history[len(broker.history)-1]
	lastProduceBatch = lastProduceReqRes.Request.(*ProduceRequest).records["my_topic"][0].RecordBatch
	if lastProduceBatch.ProducerID != newProducerID || lastProduceBatch.ProducerEpoch != 0 {
		t.Error("producer did not requested a new producerID")
	}
}

// TestBrokerProducerShutdown ensures that a call to shutdown stops the
// brokerProducer run() loop and doesn't leak any goroutines
//
//nolint:paralleltest
func TestBrokerProducerShutdown(t *testing.T) {
	defer goleak.VerifyNone(t,
		goleak.IgnoreCurrent(),
		goleak.IgnoreTopFunction("github.com/rcrowley/go-metrics.(*meterArbiter).tick"),
	)

	mockBroker := NewMockBroker(t, 1)
	metadataResponse := &MetadataResponse{}
	metadataResponse.AddBroker(mockBroker.Addr(), mockBroker.BrokerID())
	metadataResponse.AddTopicPartition(
		"my_topic", 0, mockBroker.BrokerID(), nil, nil, nil, ErrNoError)
	mockBroker.Returns(metadataResponse)

	producer, err := NewAsyncProducer([]string{mockBroker.Addr()}, NewTestConfig())
	if err != nil {
		t.Fatal(err)
	}
	broker := &Broker{
		addr: mockBroker.Addr(),
		id:   mockBroker.BrokerID(),
	}
	// Starts various goroutines in newBrokerProducer
	bp := producer.(*asyncProducer).getBrokerProducer(broker)
	// Initiate the shutdown of all of them
	producer.(*asyncProducer).unrefBrokerProducer(broker, bp)

	_ = producer.Close()
	mockBroker.Close()
}

// TestBrokerProducerWaitForSpaceEmptyBufferRollover ensures forced rollovers with an empty buffer
// do not deadlock waiting for responses when no partitions are muted.
func TestBrokerProducerWaitForSpaceEmptyBufferRollover(t *testing.T) {
	config := NewTestConfig()
	parent := &asyncProducer{
		conf:   config,
		muter:  newPartitionMuter(),
		txnmgr: &transactionManager{},
	}

	bp := &brokerProducer{
		parent:            parent,
		accumulatingBatch: newProduceSet(parent),
		output:            make(chan *produceSet, 1),
		responses:         make(chan *brokerProducerResponse),
	}

	done := make(chan error, 1)
	go func() {
		done <- bp.waitForSpace(&ProducerMessage{Topic: "topic", Partition: 0}, true)
	}()

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("expected nil error, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("waitForSpace blocked on empty buffer rollover")
	}
}

func awaitMuterBlocked(t *testing.T, m *partitionMuter, set *produceSet) {
	t.Helper()
	require.EventuallyWithT(t, func(c *assert.CollectT) {
		m.mu.Lock()
		defer m.mu.Unlock()
		assert.True(c, m.isAnyMuted(set), "muter is not blocked on set")
	}, 500*time.Millisecond, 5*time.Millisecond)
}

func assertNotDone[T any](t *testing.T, ch <-chan T, wait time.Duration) {
	t.Helper()
	time.Sleep(wait)
	select {
	case <-ch:
		t.Fatal("channel should not be ready")
	default:
	}
}

func assertDoneWithin[T any](t *testing.T, ch <-chan T, timeout time.Duration) T {
	t.Helper()
	select {
	case v := <-ch:
		return v
	case <-time.After(timeout):
		t.Fatal("timed out waiting for channel")
		var zero T
		return zero
	}
}

// TestBrokerProducerWaitForSpaceRespectsExternalUnmute ensures waitForSpace does not
// deadlock when partitions are muted by another producer and are unmuted elsewhere.
func TestBrokerProducerWaitForSpaceRespectsExternalUnmute(t *testing.T) {
	config := NewTestConfig()
	txnMgr := &transactionManager{
		producerID:      0,
		producerEpoch:   0,
		sequenceNumbers: make(map[string]int32),
	}
	parent := &asyncProducer{
		conf:   config,
		muter:  newPartitionMuter(),
		txnmgr: txnMgr,
	}

	externallyMutedSet := newProduceSet(parent)
	safeAddMessage(t, externallyMutedSet, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("held")})
	if !parent.muter.tryMute(externallyMutedSet) {
		t.Fatal("expected to mute partition")
	}

	output := make(chan *produceSet, 1)
	bp := &brokerProducer{
		parent:            parent,
		accumulatingBatch: newProduceSet(parent),
		output:            output,
		responses:         make(chan *brokerProducerResponse),
	}
	msg := &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("wait")}
	safeAddMessage(t, bp.accumulatingBatch, msg)

	done := make(chan error, 1)
	go func() {
		done <- bp.waitForSpace(msg, true)
	}()

	awaitMuterBlocked(t, parent.muter, bp.accumulatingBatch)
	parent.muter.unmute(externallyMutedSet)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("expected nil error, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("waitForSpace blocked waiting for external unmute")
	}
}

func TestBrokerProducerTryBuildFlushingBatch(t *testing.T) {
	newTestBrokerProducer := func() *brokerProducer {
		parent := &asyncProducer{
			conf:   NewTestConfig(),
			muter:  newPartitionMuter(),
			txnmgr: &transactionManager{},
		}
		return &brokerProducer{
			parent:            parent,
			accumulatingBatch: newProduceSet(parent),
		}
	}

	t.Run("signals when a batch stranded behind a muted partition can be retried", func(t *testing.T) {
		bp := newTestBrokerProducer()
		safeAddMessage(t, bp.accumulatingBatch, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("stranded")})

		inFlightBatch := newProduceSet(bp.parent)
		safeAddMessage(t, inFlightBatch, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("retrying")})
		require.True(t, bp.parent.muter.tryMute(inFlightBatch), "expected to mute partition")

		unmuteSignal := bp.tryBuildFlushingBatch()
		require.Nil(t, bp.flushingBatch, "expected the muted partition to remain buffered")
		require.NotNil(t, unmuteSignal, "expected to await the muted partition")
		select {
		case <-unmuteSignal:
			assert.Fail(t, "unmute signaled before the partition was unmuted")
		default:
		}

		bp.parent.muter.unmute(inFlightBatch)
		select {
		case <-unmuteSignal:
		default:
			require.FailNow(t, "unmute was not signaled")
		}

		require.Nil(t, bp.tryBuildFlushingBatch(), "expected no unmute wait after building the batch")
		require.NotNil(t, bp.flushingBatch, "expected to build the flushing batch")
	})

	t.Run("builds a partial batch from unmuted partitions", func(t *testing.T) {
		bp := newTestBrokerProducer()
		safeAddMessage(t, bp.accumulatingBatch, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("p0")})
		safeAddMessage(t, bp.accumulatingBatch, &ProducerMessage{Topic: "topic", Partition: 1, Value: StringEncoder("p1")})

		mutedBatch := newProduceSet(bp.parent)
		safeAddMessage(t, mutedBatch, &ProducerMessage{Topic: "topic", Partition: 1, Value: StringEncoder("held")})
		require.True(t, bp.parent.muter.tryMute(mutedBatch), "expected to mute partition")

		require.Nil(t, bp.tryBuildFlushingBatch(), "expected no unmute wait after building a partial batch")
		require.NotNil(t, bp.flushingBatch, "expected to flush available partitions")
		assert.Contains(t, bp.flushingBatch.msgs["topic"], int32(0), "expected unmuted partition to flush")
		assert.NotContains(t, bp.flushingBatch.msgs["topic"], int32(1), "expected muted partition to stay buffered")
		assert.Contains(t, bp.accumulatingBatch.msgs["topic"], int32(1), "expected muted partition to remain buffered")
	})
}

// TestAsyncProducerUnblocksOnExternalUnmute verifies that a batch stranded
// behind a partition muted by another brokerProducer flushes once the mute is
// released (#3689).
func TestAsyncProducerUnblocksOnExternalUnmute(t *testing.T) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)

	metadataResponse := new(MetadataResponse)
	metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
	metadataResponse.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataResponse)

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)

	config := NewTestConfig()
	config.Producer.Return.Successes = true
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	require.NoError(t, err)

	// mute the partition externally, as a retryBatch with a batch still in flight would
	parent := producer.(*asyncProducer)
	retryingBatch := newProduceSet(parent)
	safeAddMessage(t, retryingBatch, &ProducerMessage{Topic: "my_topic", Partition: 0, Value: StringEncoder("in-flight")})
	require.True(t, parent.muter.tryMute(retryingBatch), "expected to mute partition")

	producer.Input() <- &ProducerMessage{Topic: "my_topic", Value: StringEncoder(TestMessage)}
	assertNotDone(t, producer.Successes(), 100*time.Millisecond)

	parent.muter.unmute(retryingBatch)
	expectResultsWithTimeout(t, producer, 1, 0, 5*time.Second)

	closeProducerWithTimeout(t, producer, 5*time.Second)
	leader.Close()
	seedBroker.Close()
}

// TestBrokerProducerWaitForSpaceAllPartitionsMuted verifies that waitForSpace unblocks
// when all partitions in the accumulating batch are externally muted and later unmuted.
func TestBrokerProducerWaitForSpaceAllPartitionsMuted(t *testing.T) {
	config := NewTestConfig()
	parent := &asyncProducer{
		conf:   config,
		muter:  newPartitionMuter(),
		txnmgr: &transactionManager{},
	}

	blockedSet := newProduceSet(parent)
	safeAddMessage(t, blockedSet, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("held")})
	if !parent.muter.tryMute(blockedSet) {
		t.Fatal("expected to mute partition externally")
	}

	bp := &brokerProducer{
		parent:            parent,
		accumulatingBatch: newProduceSet(parent),
		output:            make(chan *produceSet, 1),
		responses:         make(chan *brokerProducerResponse),
		currentRetries:    make(map[string]map[int32]error),
	}
	safeAddMessage(t, bp.accumulatingBatch, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("waiting")})

	done := make(chan error, 1)
	go func() {
		done <- bp.waitForSpace(&ProducerMessage{Topic: "topic", Partition: 0}, true)
	}()

	assertNotDone(t, done, 50*time.Millisecond)
	parent.muter.unmute(blockedSet)
	require.NoError(t, assertDoneWithin(t, done, 2*time.Second))
}

// TestPartitionMuterCloseWakesWaitUntilMuted verifies that closing the muter wakes
// goroutines blocked in waitUntilMuted.
func TestPartitionMuterCloseWakesWaitUntilMuted(t *testing.T) {
	config := NewTestConfig()
	parent := &asyncProducer{
		conf:   config,
		muter:  newPartitionMuter(),
		txnmgr: &transactionManager{},
	}

	blockedSet := newProduceSet(parent)
	safeAddMessage(t, blockedSet, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("held")})
	if !parent.muter.tryMute(blockedSet) {
		t.Fatal("expected to mute partition")
	}

	waitSet := newProduceSet(parent)
	safeAddMessage(t, waitSet, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("waiting")})

	done := make(chan bool, 1)
	go func() {
		done <- parent.muter.waitUntilMuted(waitSet)
	}()

	assertNotDone(t, done, 50*time.Millisecond)
	parent.muter.close()

	select {
	case result := <-done:
		if result {
			t.Fatal("expected waitUntilMuted to return false after close")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out")
	}
}

// TestBrokerProducerRollOverClearsTimer ensures timer events from a previous batch
// do not cause a flush of a fresh empty batch after rollOver.
func TestBrokerProducerRollOverClearsTimer(t *testing.T) {
	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())

	config := NewTestConfig()
	config.Producer.Flush.Frequency = 10 * time.Millisecond
	parent := &asyncProducer{
		conf:   config,
		muter:  newPartitionMuter(),
		txnmgr: &transactionManager{},
	}
	output := make(chan *produceSet, 2)
	responses := make(chan *brokerProducerResponse)
	input := make(chan *ProducerMessage)
	bp := &brokerProducer{
		parent:            parent,
		broker:            &Broker{id: 1},
		input:             input,
		output:            output,
		responses:         responses,
		accumulatingBatch: newProduceSet(parent),
		currentRetries:    make(map[string]map[int32]error),
	}

	done := make(chan struct{})
	go func() {
		bp.run()
		close(done)
	}()

	msg := &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("v")}
	input <- msg

	select {
	case first := <-output:
		if first == nil || first.empty() {
			t.Fatal("expected flushed batch to contain message")
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("expected batch flush after timer fired")
	}

	select {
	case <-output:
		t.Fatal("unexpected flush after rollOver")
	case <-time.After(50 * time.Millisecond):
	}

	close(responses)
	close(input)

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("brokerProducer did not shut down")
	}
}

func TestRetryBatchRespectsPartitionMuter(t *testing.T) {
	config := NewTestConfig()
	config.Producer.Idempotent = true
	txnMgr := &transactionManager{
		producerID:      0,
		producerEpoch:   0,
		sequenceNumbers: make(map[string]int32),
	}

	parent := &asyncProducer{
		conf:       config,
		muter:      newPartitionMuter(),
		brokers:    make(map[*Broker]*brokerProducer),
		brokerRefs: make(map[*brokerProducer]int),
		txnmgr:     txnMgr,
	}
	leader := &Broker{}
	parent.client = &stubLeaderClient{leader: leader, cfg: config}

	output := make(chan *produceSet, 1)
	bp := &brokerProducer{
		parent: parent,
		broker: leader,
		output: output,
		input:  make(chan *ProducerMessage),
	}
	parent.brokers[leader] = bp

	retrySet := newProduceSet(parent)
	safeAddMessage(t, retrySet, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("retry")})
	retryPartitionSet := retrySet.msgs["topic"][0]
	if !parent.muter.tryMute(retrySet) {
		t.Fatal("expected retry set to mute partitions")
	}
	parent.muter.unmute(retrySet)

	parent.retryBatch("topic", 0, retryPartitionSet, ErrNotEnoughReplicas, false)

	select {
	case sent := <-output:
		set := sent.msgs["topic"][0]
		require.Equal(t, retryPartitionSet, set)
	default:
		t.Fatal("expected retry batch to be dispatched")
	}

	contender := newProduceSet(parent)
	safeAddMessage(t, contender, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("next")})
	if parent.muter.tryMute(contender) {
		t.Fatal("expected partition to remain muted by retry batch")
	}
}

func TestRetryBatchReleasesMuteOnShutdown(t *testing.T) {
	config := NewTestConfig()
	config.Producer.Idempotent = false
	config.Producer.Retry.Max = 1
	config.Producer.Retry.Backoff = 0
	config.Producer.Return.Errors = true

	parent := &asyncProducer{
		conf:       config,
		muter:      newPartitionMuter(),
		brokers:    make(map[*Broker]*brokerProducer),
		brokerRefs: make(map[*brokerProducer]int),
		errors:     make(chan *ProducerError, 1),
		done:       make(chan struct{}),
		txnmgr:     &transactionManager{},
	}
	leader := &Broker{id: 1}
	parent.client = &stubLeaderClient{leader: leader, cfg: config}

	// unbuffered and never read, so the handoff can only complete via shutdown
	parent.brokers[leader] = &brokerProducer{
		parent: parent,
		broker: leader,
		output: make(chan *produceSet),
		input:  make(chan *ProducerMessage),
	}

	sent := newProduceSet(parent)
	safeAddMessage(t, sent, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("retry")})
	retryPartitionSet := sent.msgs["topic"][0]
	require.True(t, parent.muter.tryMute(sent), "sent batch should mute its partition")
	parent.inFlight.Add(1)

	done := make(chan struct{})
	go func() {
		parent.retryBatch("topic", 0, retryPartitionSet, ErrOutOfBrokers, true)
		close(done)
	}()

	require.True(t, parent.closed.CompareAndSwap(false, true))
	close(parent.done)

	producerErr := assertDoneWithin(t, parent.errors, 2*time.Second)
	require.Equal(t, ErrShuttingDown, producerErr.Err)
	assertDoneWithin(t, done, 2*time.Second)

	contender := newProduceSet(parent)
	safeAddMessage(t, contender, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("next")})
	require.True(t, parent.muter.tryMute(contender), "partition mute should be released after aborted handoff")
	parent.muter.unmute(contender)
}

func TestBrokerProducerHandleError(t *testing.T) {
	t.Run("keeps non-idempotent connection retries muted", func(t *testing.T) {
		config := NewTestConfig()
		config.Producer.Idempotent = false
		config.Producer.Retry.Max = 1
		config.Producer.Retry.Backoff = 0

		parent := &asyncProducer{
			conf:       config,
			muter:      newPartitionMuter(),
			brokers:    make(map[*Broker]*brokerProducer),
			brokerRefs: make(map[*brokerProducer]int),
			retries:    make(chan *ProducerMessage, 4),
			txnmgr:     &transactionManager{},
		}
		retryLeader := &Broker{id: 2}
		parent.client = &stubLeaderClient{leader: retryLeader, cfg: config}

		output := make(chan *produceSet, 1)
		parent.brokers[retryLeader] = &brokerProducer{
			parent: parent,
			broker: retryLeader,
			output: output,
			input:  make(chan *ProducerMessage),
		}

		bp := &brokerProducer{
			parent:            parent,
			broker:            &Broker{id: 1},
			input:             make(chan *ProducerMessage),
			accumulatingBatch: newProduceSet(parent),
			currentRetries:    make(map[string]map[int32]error),
		}

		sent := newProduceSet(parent)
		safeAddMessage(t, sent, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("retry")})
		retryPartitionSet := sent.msgs["topic"][0]
		require.True(t, parent.muter.tryMute(sent), "sent batch should mute its partition")

		bp.handleError(sent, ErrOutOfBrokers)

		retrySet := assertDoneWithin(t, output, 2*time.Second)
		defer parent.muter.unmute(retrySet)
		require.Equal(t, retryPartitionSet, retrySet.msgs["topic"][0])
		require.Equal(t, 1, retryPartitionSet.msgs[0].retries)

		contender := newProduceSet(parent)
		safeAddMessage(t, contender, &ProducerMessage{Topic: "topic", Partition: 0, Value: StringEncoder("next")})
		require.False(t, parent.muter.tryMute(contender), "partition should stay muted by the retrying batch")
	})
}

type stubLeaderClient struct {
	cfg    *Config
	leader *Broker
}

func (c *stubLeaderClient) Config() *Config                            { return c.cfg }
func (c *stubLeaderClient) Controller() (*Broker, error)               { return nil, nil }
func (c *stubLeaderClient) RefreshController() (*Broker, error)        { return nil, nil }
func (c *stubLeaderClient) Brokers() []*Broker                         { return nil }
func (c *stubLeaderClient) Broker(int32) (*Broker, error)              { return nil, nil }
func (c *stubLeaderClient) Topics() ([]string, error)                  { return nil, nil }
func (c *stubLeaderClient) Partitions(string) ([]int32, error)         { return nil, nil }
func (c *stubLeaderClient) WritablePartitions(string) ([]int32, error) { return nil, nil }
func (c *stubLeaderClient) Leader(topic string, partitionID int32) (*Broker, error) {
	return c.leader, nil
}
func (c *stubLeaderClient) LeaderAndEpoch(string, int32) (*Broker, int32, error) {
	return c.leader, 0, nil
}
func (c *stubLeaderClient) Replicas(string, int32) ([]int32, error)          { return nil, nil }
func (c *stubLeaderClient) InSyncReplicas(string, int32) ([]int32, error)    { return nil, nil }
func (c *stubLeaderClient) OfflineReplicas(string, int32) ([]int32, error)   { return nil, nil }
func (c *stubLeaderClient) RefreshBrokers([]string) error                    { return nil }
func (c *stubLeaderClient) RefreshMetadata(...string) error                  { return nil }
func (c *stubLeaderClient) GetOffset(string, int32, int64) (int64, error)    { return 0, nil }
func (c *stubLeaderClient) Coordinator(string) (*Broker, error)              { return nil, nil }
func (c *stubLeaderClient) RefreshCoordinator(string) error                  { return nil }
func (c *stubLeaderClient) TransactionCoordinator(string) (*Broker, error)   { return nil, nil }
func (c *stubLeaderClient) RefreshTransactionCoordinator(string) error       { return nil }
func (c *stubLeaderClient) InitProducerID() (*InitProducerIDResponse, error) { return nil, nil }
func (c *stubLeaderClient) LeastLoadedBroker() *Broker                       { return c.leader }
func (c *stubLeaderClient) PartitionNotReadable(string, int32) bool          { return false }
func (c *stubLeaderClient) Close() error                                     { return nil }
func (c *stubLeaderClient) Closed() bool                                     { return false }

func testProducerInterceptor(
	t *testing.T,
	interceptors []ProducerInterceptor,
	expectationFn func(*testing.T, int, *ProducerMessage),
) {
	seedBroker := NewMockBroker(t, 1)
	leader := NewMockBroker(t, 2)
	metadataLeader := new(MetadataResponse)
	metadataLeader.AddBroker(leader.Addr(), leader.BrokerID())
	metadataLeader.AddTopicPartition("my_topic", 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
	seedBroker.Returns(metadataLeader)

	config := NewTestConfig()
	config.Producer.Flush.Messages = 10
	config.Producer.Return.Successes = true
	config.Producer.Interceptors = interceptors
	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
	if err != nil {
		t.Fatal(err)
	}

	for range 10 {
		producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder(TestMessage)}
	}

	prodSuccess := new(ProduceResponse)
	prodSuccess.AddTopicPartition("my_topic", 0, ErrNoError)
	leader.Returns(prodSuccess)

	for i := range 10 {
		select {
		case msg := <-producer.Errors():
			t.Error(msg.Err)
		case msg := <-producer.Successes():
			expectationFn(t, i, msg)
		}
	}

	closeProducer(t, producer)
	leader.Close()
	seedBroker.Close()
}

func TestAsyncProducerInterceptors(t *testing.T) {
	tests := []struct {
		name          string
		interceptors  []ProducerInterceptor
		expectationFn func(*testing.T, int, *ProducerMessage)
	}{
		{
			name:         "intercept messages",
			interceptors: []ProducerInterceptor{&appendInterceptor{i: 0}},
			expectationFn: func(t *testing.T, i int, msg *ProducerMessage) {
				v, _ := msg.Value.Encode()
				expected := TestMessage + strconv.Itoa(i)
				if string(v) != expected {
					t.Errorf("Interceptor should have incremented the value, got %s, expected %s", v, expected)
				}
			},
		},
		{
			name:         "interceptor chain",
			interceptors: []ProducerInterceptor{&appendInterceptor{i: 0}, &appendInterceptor{i: 1000}},
			expectationFn: func(t *testing.T, i int, msg *ProducerMessage) {
				v, _ := msg.Value.Encode()
				expected := TestMessage + strconv.Itoa(i) + strconv.Itoa(i+1000)
				if string(v) != expected {
					t.Errorf("Interceptor should have incremented the value, got %s, expected %s", v, expected)
				}
			},
		},
		{
			name:         "interceptor chain with one interceptor failing",
			interceptors: []ProducerInterceptor{&appendInterceptor{i: -1}, &appendInterceptor{i: 1000}},
			expectationFn: func(t *testing.T, i int, msg *ProducerMessage) {
				v, _ := msg.Value.Encode()
				expected := TestMessage + strconv.Itoa(i+1000)
				if string(v) != expected {
					t.Errorf("Interceptor should have incremented the value, got %s, expected %s", v, expected)
				}
			},
		},
		{
			name:         "interceptor chain with all interceptors failing",
			interceptors: []ProducerInterceptor{&appendInterceptor{i: -1}, &appendInterceptor{i: -1}},
			expectationFn: func(t *testing.T, i int, msg *ProducerMessage) {
				v, _ := msg.Value.Encode()
				expected := TestMessage
				if string(v) != expected {
					t.Errorf("Interceptor should have not changed the value, got %s, expected %s", v, expected)
				}
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			testProducerInterceptor(t, tt.interceptors, tt.expectationFn)
		})
	}
}

func TestProducerError(t *testing.T) {
	t.Parallel()
	err := ProducerError{Err: ErrOutOfBrokers}
	if !errors.Is(err, ErrOutOfBrokers) {
		t.Error("unexpected errors.Is")
	}
}

func TestTxmngInitProducerId(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	metadataLeader := new(MetadataResponse)
	metadataLeader.Version = 4
	metadataLeader.AddBroker(broker.Addr(), broker.BrokerID())
	broker.Returns(metadataLeader)

	config := NewTestConfig()
	config.Producer.Idempotent = true
	config.Version = V0_11_0_0
	config.Producer.RequiredAcks = WaitForAll
	config.Net.MaxOpenRequests = 1

	client, err := NewClient([]string{broker.Addr()}, config)
	require.NoError(t, err)
	defer client.Close()

	producerIdResponse := &InitProducerIDResponse{
		Err:           ErrNoError,
		ProducerID:    1,
		ProducerEpoch: 0,
	}
	broker.Returns(producerIdResponse)

	txmng, err := newTransactionManager(config, client)
	require.NoError(t, err)

	require.Equal(t, int64(1), txmng.producerID)
	require.Equal(t, int16(0), txmng.producerEpoch)
}

func TestTxnProduceBumpEpoch(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	config := NewTestConfig()
	config.Producer.Idempotent = true
	config.Producer.Transaction.ID = "test"
	config.Version = V2_6_0_0
	config.Producer.RequiredAcks = WaitForAll
	config.Net.MaxOpenRequests = 1
	config.Producer.Return.Errors = false

	config.ApiVersionsRequest = false

	metadataLeader := new(MetadataResponse)
	metadataLeader.Version = 9
	metadataLeader.ControllerID = broker.brokerID
	metadataLeader.AddBroker(broker.Addr(), broker.BrokerID())
	metadataLeader.AddTopic("test-topic", ErrNoError)
	metadataLeader.AddTopicPartition("test-topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataLeader)

	client, err := NewClient([]string{broker.Addr()}, config)
	require.NoError(t, err)
	defer client.Close()

	findCoordinatorResponse := FindCoordinatorResponse{
		Coordinator: client.Brokers()[0],
		Err:         ErrNoError,
		Version:     1,
	}
	broker.Returns(&findCoordinatorResponse)

	producerIdResponse := &InitProducerIDResponse{
		Err:           ErrNoError,
		ProducerID:    1000,
		ProducerEpoch: 0,
		Version:       3,
	}
	broker.Returns(producerIdResponse)

	ap, err := NewAsyncProducerFromClient(client)
	producer := ap.(*asyncProducer)
	require.NoError(t, err)
	defer ap.Close()
	require.Equal(t, int64(1000), producer.txnmgr.producerID)
	require.Equal(t, int16(0), producer.txnmgr.producerEpoch)

	addPartitionsToTxnResponse := &AddPartitionsToTxnResponse{
		Errors: map[string][]*PartitionError{
			"test-topic": {
				{
					Partition: 0,
				},
			},
		},
	}
	broker.Returns(addPartitionsToTxnResponse)

	produceResponse := new(ProduceResponse)
	produceResponse.AddTopicPartition("test-topic", 0, ErrOutOfOrderSequenceNumber)
	broker.Returns(produceResponse)

	err = producer.BeginTxn()
	require.NoError(t, err)

	producer.Input() <- &ProducerMessage{Topic: "test-topic", Key: nil, Value: StringEncoder(TestMessage)}

	// Force send
	producer.inFlight.Add(1)
	producer.Input() <- &ProducerMessage{flags: shutdown}
	producer.inFlight.Wait()

	err = producer.CommitTxn()
	require.Error(t, err)
	require.Equal(t, ProducerTxnFlagInError|ProducerTxnFlagAbortableError, producer.txnmgr.status)

	err = producer.CommitTxn()
	require.Error(t, err)
	require.Equal(t, ProducerTxnFlagInError|ProducerTxnFlagAbortableError, producer.txnmgr.status)

	endTxnResponse := &EndTxnResponse{
		Err: ErrNoError,
	}
	broker.Returns(endTxnResponse)

	producerBumpIdResponse := &InitProducerIDResponse{
		Err:           ErrNoError,
		ProducerID:    1000,
		ProducerEpoch: 1,
		Version:       3,
	}
	broker.Returns(producerBumpIdResponse)

	err = producer.AbortTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)
	require.Equal(t, int64(1000), producer.txnmgr.producerID)
	require.Equal(t, int16(1), producer.txnmgr.producerEpoch)
}

func TestTxnProduceRecordWithCommit(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	config := NewTestConfig()
	config.Producer.Idempotent = true
	config.Producer.Transaction.ID = "test"
	config.Version = V0_11_0_0
	config.Producer.RequiredAcks = WaitForAll
	config.Net.MaxOpenRequests = 1

	metadataLeader := new(MetadataResponse)
	metadataLeader.Version = 4
	metadataLeader.ControllerID = broker.brokerID
	metadataLeader.AddBroker(broker.Addr(), broker.BrokerID())
	metadataLeader.AddTopic("test-topic", ErrNoError)
	metadataLeader.AddTopicPartition("test-topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataLeader)

	client, err := NewClient([]string{broker.Addr()}, config)
	require.NoError(t, err)
	defer client.Close()

	findCoordinatorResponse := FindCoordinatorResponse{
		Coordinator: client.Brokers()[0],
		Err:         ErrNoError,
		Version:     1,
	}
	broker.Returns(&findCoordinatorResponse)

	producerIdResponse := &InitProducerIDResponse{
		Err:           ErrNoError,
		ProducerID:    1,
		ProducerEpoch: 0,
	}
	broker.Returns(producerIdResponse)

	ap, err := NewAsyncProducerFromClient(client)
	producer := ap.(*asyncProducer)
	require.NoError(t, err)
	defer ap.Close()

	addPartitionsToTxnResponse := &AddPartitionsToTxnResponse{
		Errors: map[string][]*PartitionError{
			"test-topic": {
				{
					Partition: 0,
				},
			},
		},
	}
	broker.Returns(addPartitionsToTxnResponse)

	produceResponse := new(ProduceResponse)
	produceResponse.Version = 3
	produceResponse.AddTopicPartition("test-topic", 0, ErrNoError)
	broker.Returns(produceResponse)

	endTxnResponse := &EndTxnResponse{
		Err: ErrNoError,
	}
	broker.Returns(endTxnResponse)

	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)

	err = producer.BeginTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagInTransaction, producer.txnmgr.status)

	producer.Input() <- &ProducerMessage{Topic: "test-topic", Key: nil, Value: StringEncoder(TestMessage)}
	err = producer.CommitTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)
}

func TestTxnProduceBatchAddPartition(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	config := NewTestConfig()
	config.Producer.Idempotent = true
	config.Producer.Transaction.ID = "test"
	config.Version = V0_11_0_0
	config.Producer.RequiredAcks = WaitForAll
	config.Net.MaxOpenRequests = 1

	config.Producer.Retry.Max = 1
	config.Producer.Flush.Messages = 3
	config.Producer.Flush.Frequency = 30 * time.Second
	config.Producer.Flush.Bytes = 1 << 12
	config.Producer.Partitioner = NewManualPartitioner

	metadataLeader := new(MetadataResponse)
	metadataLeader.Version = 4
	metadataLeader.ControllerID = broker.brokerID
	metadataLeader.AddBroker(broker.Addr(), broker.BrokerID())
	metadataLeader.AddTopic("test-topic", ErrNoError)
	metadataLeader.AddTopicPartition("test-topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	metadataLeader.AddTopicPartition("test-topic", 1, broker.BrokerID(), nil, nil, nil, ErrNoError)
	metadataLeader.AddTopicPartition("test-topic", 2, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataLeader)

	client, err := NewClient([]string{broker.Addr()}, config)
	require.NoError(t, err)
	defer client.Close()

	findCoordinatorResponse := FindCoordinatorResponse{
		Coordinator: client.Brokers()[0],
		Err:         ErrNoError,
		Version:     1,
	}
	broker.Returns(&findCoordinatorResponse)

	producerIdResponse := &InitProducerIDResponse{
		Err:           ErrNoError,
		ProducerID:    1,
		ProducerEpoch: 0,
	}
	broker.Returns(producerIdResponse)

	ap, err := NewAsyncProducerFromClient(client)
	producer := ap.(*asyncProducer)
	require.NoError(t, err)
	defer ap.Close()

	go func() {
		for err := range producer.Errors() {
			assert.NoError(t, err)
		}
	}()

	broker.Returns(&AddPartitionsToTxnResponse{
		Errors: map[string][]*PartitionError{
			"test-topic": {
				{
					Partition: 0,
					Err:       ErrNoError,
				},
				{
					Partition: 1,
					Err:       ErrNoError,
				},
				{
					Partition: 2,
					Err:       ErrNoError,
				},
			},
		},
	})

	produceResponse := new(ProduceResponse)
	produceResponse.Version = 3
	produceResponse.AddTopicPartition("test-topic", 0, ErrNoError)
	produceResponse.AddTopicPartition("test-topic", 1, ErrNoError)
	produceResponse.AddTopicPartition("test-topic", 2, ErrNoError)
	broker.Returns(produceResponse)

	endTxnResponse := &EndTxnResponse{
		Err: ErrNoError,
	}
	broker.Returns(endTxnResponse)

	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)

	err = producer.BeginTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagInTransaction, producer.txnmgr.status)

	producer.Input() <- &ProducerMessage{Topic: "test-topic", Partition: 0, Key: nil, Value: StringEncoder("partition-0")}
	producer.Input() <- &ProducerMessage{Topic: "test-topic", Partition: 1, Key: nil, Value: StringEncoder("partition-1")}
	producer.Input() <- &ProducerMessage{Topic: "test-topic", Partition: 2, Key: nil, Value: StringEncoder("partition-2")}

	err = producer.CommitTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)

	produceExchange := broker.History()[len(broker.History())-2]
	produceRequest := produceExchange.Request.(*ProduceRequest)
	require.Len(t, produceRequest.records["test-topic"], 3)

	addPartitionExchange := broker.History()[len(broker.History())-3]
	addpartitionRequest := addPartitionExchange.Request.(*AddPartitionsToTxnRequest)
	require.Len(t, addpartitionRequest.TopicPartitions["test-topic"], 3)
	require.Contains(t, addpartitionRequest.TopicPartitions["test-topic"], int32(0))
	require.Contains(t, addpartitionRequest.TopicPartitions["test-topic"], int32(1))
	require.Contains(t, addpartitionRequest.TopicPartitions["test-topic"], int32(2))
}

func TestTxnProduceRecordWithAbort(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	config := NewTestConfig()
	config.Producer.Idempotent = true
	config.Producer.Transaction.ID = "test"
	config.Version = V0_11_0_0
	config.Producer.RequiredAcks = WaitForAll
	config.Net.MaxOpenRequests = 1

	metadataLeader := new(MetadataResponse)
	metadataLeader.Version = 4
	metadataLeader.ControllerID = broker.brokerID
	metadataLeader.AddBroker(broker.Addr(), broker.BrokerID())
	metadataLeader.AddTopic("test-topic", ErrNoError)
	metadataLeader.AddTopicPartition("test-topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
	broker.Returns(metadataLeader)

	client, err := NewClient([]string{broker.Addr()}, config)
	require.NoError(t, err)
	defer client.Close()

	findCoordinatorResponse := FindCoordinatorResponse{
		Coordinator: client.Brokers()[0],
		Err:         ErrNoError,
		Version:     1,
	}
	broker.Returns(&findCoordinatorResponse)

	producerIdResponse := &InitProducerIDResponse{
		Err:           ErrNoError,
		ProducerID:    1,
		ProducerEpoch: 0,
	}
	broker.Returns(producerIdResponse)

	ap, err := NewAsyncProducerFromClient(client)
	producer := ap.(*asyncProducer)
	require.NoError(t, err)
	defer ap.Close()

	broker.Returns(&AddPartitionsToTxnResponse{
		Errors: map[string][]*PartitionError{
			"test-topic": {
				{
					Partition: 0,
					Err:       ErrNoError,
				},
			},
		},
	})

	produceResponse := new(ProduceResponse)
	produceResponse.Version = 3
	produceResponse.AddTopicPartition("test-topic", 0, ErrNoError)
	broker.Returns(produceResponse)

	endTxnResponse := &EndTxnResponse{
		Err: ErrNoError,
	}
	broker.Returns(endTxnResponse)

	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)

	err = producer.BeginTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagInTransaction, producer.txnmgr.status)

	producer.Input() <- &ProducerMessage{Topic: "test-topic", Key: nil, Value: StringEncoder(TestMessage)}
	err = producer.AbortTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)
}

func TestTxnCanAbort(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()

	config := NewTestConfig()
	config.Producer.Idempotent = true
	config.Producer.Transaction.ID = "test"
	config.Version = V0_11_0_0
	config.Producer.RequiredAcks = WaitForAll
	config.Producer.Return.Errors = false
	config.Producer.Return.Successes = true
	config.Producer.Retry.Backoff = 0
	config.Producer.Flush.Messages = 1
	config.Producer.Retry.Max = 1
	config.Net.MaxOpenRequests = 1

	var (
		mu                   sync.Mutex
		addPartitionsCount   int
		produceRequestsCount int
	)

	broker.SetHandlerFuncByMap(map[string]requestHandlerFunc{
		"MetadataRequest": func(req *request) encoderWithHeader {
			resp := new(MetadataResponse)
			resp.Version = 4
			resp.ControllerID = broker.BrokerID()
			resp.AddBroker(broker.Addr(), broker.BrokerID())
			resp.AddTopic("test-topic", ErrNoError)
			resp.AddTopic("test-topic-2", ErrNoError)
			resp.AddTopicPartition("test-topic", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
			resp.AddTopicPartition("test-topic-2", 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
			return resp
		},
		"FindCoordinatorRequest": func(req *request) encoderWithHeader {
			resp := new(FindCoordinatorResponse)
			resp.Version = 1
			resp.Coordinator = &Broker{id: broker.BrokerID(), addr: broker.Addr()}
			resp.Err = ErrNoError
			return resp
		},
		"InitProducerIDRequest": func(req *request) encoderWithHeader {
			return &InitProducerIDResponse{
				Err:           ErrNoError,
				ProducerID:    1,
				ProducerEpoch: 0,
			}
		},
		"AddPartitionsToTxnRequest": func(req *request) encoderWithHeader {
			mu.Lock()
			addPartitionsCount++
			count := addPartitionsCount
			mu.Unlock()

			if count == 1 {
				return &AddPartitionsToTxnResponse{
					Errors: map[string][]*PartitionError{
						"test-topic-2": {{Partition: 0, Err: ErrNoError}},
					},
				}
			}
			return &AddPartitionsToTxnResponse{
				Errors: map[string][]*PartitionError{
					"test-topic": {{Partition: 0, Err: ErrTopicAuthorizationFailed}},
				},
			}
		},
		"ProduceRequest": func(req *request) encoderWithHeader {
			mu.Lock()
			produceRequestsCount++
			mu.Unlock()

			resp := new(ProduceResponse)
			resp.Version = 3
			resp.AddTopicPartition("test-topic-2", 0, ErrNoError)
			return resp
		},
		"EndTxnRequest": func(req *request) encoderWithHeader {
			return &EndTxnResponse{Err: ErrNoError}
		},
	})

	client, err := NewClient([]string{broker.Addr()}, config)
	require.NoError(t, err)
	defer client.Close()

	ap, err := NewAsyncProducerFromClient(client)
	producer := ap.(*asyncProducer)
	require.NoError(t, err)
	defer ap.Close()

	require.Equal(t, ProducerTxnFlagReady, producer.txnmgr.status)

	err = producer.BeginTxn()
	require.NoError(t, err)
	require.Equal(t, ProducerTxnFlagInTransaction, producer.txnmgr.status)

	producer.Input() <- &ProducerMessage{Topic: "test-topic-2", Partition: 0, Key: nil, Value: StringEncoder(TestMessage)}
	<-producer.Successes()

	producer.Input() <- &ProducerMessage{Topic: "test-topic", Partition: 0, Key: nil, Value: StringEncoder(TestMessage)}

	err = producer.CommitTxn()
	require.Error(t, err)
	require.NotEqual(t, 0, producer.txnmgr.status&ProducerTxnFlagAbortableError)

	err = producer.AbortTxn()
	require.NoError(t, err)
}

func TestProducerRetryBufferLimits(t *testing.T) {
	broker := NewMockBroker(t, 1)
	defer broker.Close()
	topic := "test-topic"

	metadataRequestHandlerFunc := func(req *request) (res encoderWithHeader) {
		r := new(MetadataResponse)
		r.AddBroker(broker.Addr(), broker.BrokerID())
		r.AddTopicPartition(topic, 0, broker.BrokerID(), nil, nil, nil, ErrNoError)
		return r
	}

	produceRequestHandlerFunc := func(req *request) (res encoderWithHeader) {
		r := new(ProduceResponse)
		r.AddTopicPartition(topic, 0, ErrNotLeaderForPartition)
		return r
	}

	broker.SetHandlerFuncByMap(map[string]requestHandlerFunc{
		"ProduceRequest":  produceRequestHandlerFunc,
		"MetadataRequest": metadataRequestHandlerFunc,
	})

	tests := []struct {
		name            string
		configureBuffer func(*Config)
		messageSize     int
		numMessages     int
	}{
		{
			name: "MaxBufferLength",
			configureBuffer: func(config *Config) {
				config.Producer.Flush.MaxMessages = 1
				config.Producer.Retry.MaxBufferLength = minFunctionalRetryBufferLength
			},
			messageSize: 1, // Small message size
			numMessages: 10000,
		},
		{
			name: "MaxBufferBytes",
			configureBuffer: func(config *Config) {
				config.Producer.Flush.MaxMessages = 1
				config.Producer.Retry.MaxBufferBytes = minFunctionalRetryBufferBytes
			},
			messageSize: 950 * 1024, // 950 KB
			numMessages: 1000,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			config := NewTestConfig()
			config.Producer.Return.Successes = true
			tt.configureBuffer(config)

			producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
			if err != nil {
				t.Fatal(err)
			}

			var (
				wg                        sync.WaitGroup
				successes, producerErrors int
				errorFound                bool
			)

			wg.Go(func() {
				for range producer.Successes() {
					successes++
				}
			})

			wg.Go(func() {
				for errMsg := range producer.Errors() {
					if errors.Is(errMsg.Err, ErrProducerRetryBufferOverflow) {
						errorFound = true
					}
					producerErrors++
				}
			})

			longString := strings.Repeat("a", tt.messageSize)
			val := StringEncoder(longString)

			for i := 0; i < tt.numMessages; i++ {
				msg := &ProducerMessage{
					Topic: topic,
					Value: val,
				}
				producer.Input() <- msg
			}

			producer.AsyncClose()
			wg.Wait()

			assert.Equal(t, successes+producerErrors, tt.numMessages, "Expected all messages to be processed")
			assert.True(t, errorFound, "Expected at least one error matching ErrProducerRetryBufferOverflow")
		})
	}
}

// This example shows how to use the producer while simultaneously
// reading the Errors channel to know about any failures.
func ExampleAsyncProducer_select() {
	producer, err := NewAsyncProducer([]string{"localhost:9092"}, nil)
	if err != nil {
		panic(err)
	}

	defer func() {
		if err := producer.Close(); err != nil {
			log.Fatalln(err)
		}
	}()

	// Trap SIGINT to trigger a shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	var enqueued, producerErrors int
ProducerLoop:
	for {
		select {
		case producer.Input() <- &ProducerMessage{Topic: "my_topic", Key: nil, Value: StringEncoder("testing 123")}:
			enqueued++
		case err := <-producer.Errors():
			log.Println("Failed to produce message", err)
			producerErrors++
		case <-signals:
			break ProducerLoop
		}
	}

	log.Printf("Enqueued: %d; errors: %d\n", enqueued, producerErrors)
}

// This example shows how to use the producer with separate goroutines
// reading from the Successes and Errors channels. Note that in order
// for the Successes channel to be populated, you have to set
// config.Producer.Return.Successes to true.
func ExampleAsyncProducer_goroutines() {
	config := NewTestConfig()
	config.Producer.Return.Successes = true
	producer, err := NewAsyncProducer([]string{"localhost:9092"}, config)
	if err != nil {
		panic(err)
	}

	// Trap SIGINT to trigger a graceful shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	var (
		wg                                  sync.WaitGroup
		enqueued, successes, producerErrors int
	)

	wg.Go(func() {
		for range producer.Successes() {
			successes++
		}
	})

	wg.Go(func() {
		for err := range producer.Errors() {
			log.Println(err)
			producerErrors++
		}
	})

ProducerLoop:
	for {
		message := &ProducerMessage{Topic: "my_topic", Value: StringEncoder("testing 123")}
		select {
		case producer.Input() <- message:
			enqueued++

		case <-signals:
			producer.AsyncClose() // Trigger a shutdown of the producer.
			break ProducerLoop
		}
	}

	wg.Wait()

	log.Printf("Successfully produced: %d; errors: %d\n", successes, producerErrors)
}

// TestAsyncProducerRetryOrdering verifies that message ordering is preserved during retries,
// both with and without request pipelining (MaxOpenRequests=1 vs >1).
func TestAsyncProducerRetryOrdering(t *testing.T) {
	const topic = "my_topic"

	extractValue := func(pr *ProduceRequest) string {
		recordsByPartition := pr.records[topic]
		if recordsByPartition == nil {
			return ""
		}
		records := recordsByPartition[0]
		if rb := records.RecordBatch; rb != nil && len(rb.Records) > 0 {
			return string(rb.Records[0].Value)
		}
		if ms := records.MsgSet; ms != nil && len(ms.Messages) > 0 {
			return string(ms.Messages[0].Msg.Value)
		}
		return ""
	}

	tests := []struct {
		name            string
		maxOpenRequests int
		retryBackoff    time.Duration
	}{
		{
			name:            "no pipelining (MaxOpenRequests=1)",
			maxOpenRequests: 1,
			retryBackoff:    0,
		},
		{
			name:            "with pipelining (MaxOpenRequests=5)",
			maxOpenRequests: 5,
			retryBackoff:    50 * time.Millisecond,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			seedBroker := NewMockBroker(t, 1)
			leader := NewMockBroker(t, 2)

			metadataResponse := new(MetadataResponse)
			metadataResponse.AddBroker(leader.Addr(), leader.BrokerID())
			metadataResponse.AddTopicPartition(topic, 0, leader.BrokerID(), nil, nil, nil, ErrNoError)
			seedBroker.Returns(metadataResponse)

			var (
				mu              sync.Mutex
				produceAttempts int
				valuesSeen      []string
			)

			leader.setHandler(func(req *request) (res encoderWithHeader) {
				switch typed := req.body.(type) {
				case *MetadataRequest:
					return metadataResponse
				case *ProduceRequest:
					mu.Lock()
					defer mu.Unlock()

					produceAttempts++
					value := extractValue(typed)
					valuesSeen = append(valuesSeen, value)

					resp := new(ProduceResponse)
					resp.Version = typed.Version
					switch produceAttempts {
					case 1:
						resp.AddTopicPartition(topic, 0, ErrNotLeaderForPartition)
					case 2, 3:
						resp.AddTopicPartition(topic, 0, ErrNoError)
					default:
						t.Errorf("unexpected attempt %d", produceAttempts)
						resp.AddTopicPartition(topic, 0, ErrNoError)
					}
					return resp
				default:
					return nil
				}
			})

			config := NewTestConfig()
			config.Producer.Return.Successes = true
			config.Producer.Retry.Max = 3
			config.Producer.Retry.Backoff = tt.retryBackoff
			config.Producer.Flush.Messages = 1
			config.Producer.Flush.MaxMessages = 1
			config.Producer.Partitioner = NewManualPartitioner
			config.Net.MaxOpenRequests = tt.maxOpenRequests

			producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
			if err != nil {
				t.Fatal(err)
			}

			msgValues := []string{"msg-0", "msg-1"}
			for _, val := range msgValues {
				producer.Input() <- &ProducerMessage{Topic: topic, Partition: 0, Value: StringEncoder(val)}
			}

			expectResultsWithTimeout(t, producer, len(msgValues), 0, 10*time.Second)

			mu.Lock()
			attempts := produceAttempts
			seen := make([]string, len(valuesSeen))
			copy(seen, valuesSeen)
			mu.Unlock()

			closeProducer(t, producer)
			seedBroker.Close()
			leader.Close()

			if attempts != 3 {
				t.Errorf("expected 3 produce attempts, got %d", attempts)
			}

			// Both configurations should maintain ordering: msg-0 (fail), msg-0 (retry), msg-1
			expectedOrder := []string{"msg-0", "msg-0", "msg-1"}
			if !assert.Equal(t, expectedOrder, seen) {
				t.Errorf("messages out of order: got %v, want %v", seen, expectedOrder)
			}
		})
	}
}

// TestAsyncProducerPartitionUnmuting verifies that partitions are properly unmuted
// in all error paths: send errors, NoResponse acks, etc. Without proper unmuting,
// partitions remain muted and subsequent messages would block indefinitely.
func TestAsyncProducerPartitionUnmuting(t *testing.T) {
	const topic = "test_topic"

	t.Run("NoResponse acks unmute partitions", func(t *testing.T) {
		broker := NewMockBroker(t, 1)
		defer broker.Close()

		metadataResponse := NewMockMetadataResponse(t).
			SetBroker(broker.Addr(), broker.BrokerID()).
			SetLeader(topic, 0, broker.BrokerID())
		broker.SetHandlerByMap(map[string]MockResponse{
			"MetadataRequest": metadataResponse,
		})

		config := NewTestConfig()
		config.Producer.RequiredAcks = NoResponse
		config.Producer.Return.Successes = true
		config.Producer.Flush.Messages = 1
		config.Net.MaxOpenRequests = 5

		producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
		if err != nil {
			t.Fatal(err)
		}

		for range 3 {
			producer.Input() <- &ProducerMessage{
				Topic:     topic,
				Partition: 0,
				Value:     StringEncoder("msg"),
			}
		}

		successCount := 0
		for i := range 3 {
			select {
			case <-producer.Successes():
				successCount++
			case err := <-producer.Errors():
				t.Fatalf("unexpected error: %v", err)
			case <-time.After(5 * time.Second):
				t.Fatalf("timeout waiting for success %d (got %d) - partitions may not be unmuted", i+1, successCount)
			}
		}

		if successCount != 3 {
			t.Errorf("expected 3 successes, got %d", successCount)
		}

		closeProducer(t, producer)
	})

	t.Run("retry keeps partition muted until queued", func(t *testing.T) {
		broker := NewMockBroker(t, 1)
		defer broker.Close()

		metadataResponse := new(MetadataResponse)
		metadataResponse.AddBroker(broker.Addr(), broker.BrokerID())
		metadataResponse.AddTopicPartition(topic, 0, broker.BrokerID(), nil, nil, nil, ErrNoError)

		var attemptCount int
		var mu sync.Mutex
		firstRequestReceived := make(chan struct{})

		broker.setHandler(func(req *request) (res encoderWithHeader) {
			switch req.body.(type) {
			case *MetadataRequest:
				return metadataResponse
			case *ProduceRequest:
				mu.Lock()
				attemptCount++
				attempt := attemptCount
				mu.Unlock()

				if attempt == 1 {
					close(firstRequestReceived)
				}

				resp := new(ProduceResponse)
				if attempt == 1 {
					resp.AddTopicPartition(topic, 0, ErrNotLeaderForPartition)
				} else {
					resp.AddTopicPartition(topic, 0, ErrNoError)
				}
				return resp
			}
			return nil
		})

		config := NewTestConfig()
		config.Producer.Return.Successes = true
		config.Producer.Retry.Max = 1
		config.Producer.Retry.Backoff = 10 * time.Millisecond
		config.Producer.Flush.Messages = 1
		config.Net.MaxOpenRequests = 5

		producer, err := NewAsyncProducer([]string{broker.Addr()}, config)
		if err != nil {
			t.Fatal(err)
		}

		producer.Input() <- &ProducerMessage{
			Topic:     topic,
			Partition: 0,
			Value:     StringEncoder("msg-0"),
		}

		<-firstRequestReceived

		producer.Input() <- &ProducerMessage{
			Topic:     topic,
			Partition: 0,
			Value:     StringEncoder("msg-1"),
		}

		var successCount int
		for range 2 {
			select {
			case <-producer.Successes():
				successCount++
			case err := <-producer.Errors():
				t.Fatalf("unexpected error: %v", err)
			case <-time.After(5 * time.Second):
				t.Fatalf("timeout waiting for success %d/%d - partition may be deadlocked", successCount, 2)
			}
		}

		if successCount != 2 {
			t.Errorf("expected 2 successes, got %d", successCount)
		}

		mu.Lock()
		attempts := attemptCount
		mu.Unlock()
		if attempts != 3 {
			t.Errorf("expected 3 produce attempts, got %d", attempts)
		}

		closeProducer(t, producer)
	})
}