chore: finish up the move to atomic types

async_producer_test.go

@@ -629,9 +629,14 @@ func TestAsyncProducerMultipleRetriesWithBackoffFunc(t *testing.T) {
 	config.Producer.Return.Successes = true
 	config.Producer.Retry.Max = 4
 
-	backoffCalled := make([]int32, config.Producer.Retry.Max+1)
+	// 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 {
-		atomic.AddInt32(&backoffCalled[retries-1], 1)
+		backoffCalled[retries-1].Add(1)
 		return 0
 	}
 	producer, err := NewAsyncProducer([]string{seedBroker.Addr()}, config)
@@ -672,11 +677,11 @@ func TestAsyncProducerMultipleRetriesWithBackoffFunc(t *testing.T) {
 	closeProducer(t, producer)
 
 	for i := 0; i < config.Producer.Retry.Max; i++ {
-		if atomic.LoadInt32(&backoffCalled[i]) != 1 {
+		if backoffCalled[i].Load() != 1 {
 			t.Errorf("expected one retry attempt #%d", i)
 		}
 	}
-	if atomic.LoadInt32(&backoffCalled[config.Producer.Retry.Max]) != 0 {
+	if backoffCalled[config.Producer.Retry.Max].Load() != 0 {
 		t.Errorf("expected no retry attempt #%d", config.Producer.Retry.Max)
 	}
 }
@@ -811,21 +816,21 @@ func TestAsyncProducerBrokerRestart(t *testing.T) {
 	// The seed broker only handles Metadata request in bootstrap
 	seedBroker.setHandler(metadataRequestHandlerFunc)
 
-	var emptyValues int32 = 0
+	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 {
-					atomic.AddInt32(&emptyValues, 1)
+					emptyValues.Add(1)
 				}
 			}
 		}
 		if batch := preq.records["my_topic"][0].MsgSet; batch != nil {
 			for _, record := range batch.Messages {
 				if len(record.Msg.Value) == 0 {
-					atomic.AddInt32(&emptyValues, 1)
+					emptyValues.Add(1)
 				}
 			}
 		}
@@ -904,7 +909,7 @@ func TestAsyncProducerBrokerRestart(t *testing.T) {
 
 	closeProducerWithTimeout(t, producer, 5*time.Second)
 
-	if emptyValues := atomic.LoadInt32(&emptyValues); emptyValues > 0 {
+	if emptyValues := emptyValues.Load(); emptyValues > 0 {
 		t.Fatalf("%d empty values", emptyValues)
 	}
 }

client_test.go

@@ -362,12 +362,12 @@ func TestClientReceivingUnknownTopicWithBackoffFunc(t *testing.T) {
 	metadataResponse1.AddBroker(seedBroker.Addr(), seedBroker.BrokerID())
 	seedBroker.Returns(metadataResponse1)
 
-	retryCount := int32(0)
+	var retryCount atomic.Int32
 
 	config := NewTestConfig()
 	config.Metadata.Retry.Max = 1
 	config.Metadata.Retry.BackoffFunc = func(retries, maxRetries int) time.Duration {
-		atomic.AddInt32(&retryCount, 1)
+		retryCount.Add(1)
 		return 0
 	}
 	client, err := NewClient([]string{seedBroker.Addr()}, config)
@@ -388,7 +388,7 @@ func TestClientReceivingUnknownTopicWithBackoffFunc(t *testing.T) {
 	safeClose(t, client)
 	seedBroker.Close()
 
-	actualRetryCount := atomic.LoadInt32(&retryCount)
+	actualRetryCount := retryCount.Load()
 	if actualRetryCount != 1 {
 		t.Fatalf("Expected BackoffFunc to be called exactly once, but saw %d", actualRetryCount)
 	}

consumer_test.go

@@ -491,12 +491,12 @@ func TestConsumerLeaderRefreshError(t *testing.T) {
 }
 
 func TestConsumerLeaderRefreshErrorWithBackoffFunc(t *testing.T) {
-	var calls int32 = 0
+	var calls atomic.Int32
 
 	config := NewTestConfig()
 	config.Net.ReadTimeout = 100 * time.Millisecond
 	config.Consumer.Retry.BackoffFunc = func(retries int) time.Duration {
-		atomic.AddInt32(&calls, 1)
+		calls.Add(1)
 		return 200 * time.Millisecond
 	}
 	config.Consumer.Return.Errors = true
@@ -505,7 +505,7 @@ func TestConsumerLeaderRefreshErrorWithBackoffFunc(t *testing.T) {
 	runConsumerLeaderRefreshErrorTestWithConfig(t, config)
 
 	// we expect at least one call to our backoff function
-	if calls == 0 {
+	if calls.Load() == 0 {
 		t.Fail()
 	}
 }

functional_consumer_group_test.go

@@ -348,20 +348,20 @@ type testFuncConsumerGroupMessage struct {
 
 type testFuncConsumerGroupSink struct {
 	msgs  chan testFuncConsumerGroupMessage
-	count int32
+	count atomic.Int32
 }
 
 func (s *testFuncConsumerGroupSink) Len() int {
 	if s == nil {
 		return -1
 	}
-	return int(atomic.LoadInt32(&s.count))
+	return int(s.count.Load())
 }
 
 func (s *testFuncConsumerGroupSink) Push(clientID string, m *ConsumerMessage) {
 	if s != nil {
 		s.msgs <- testFuncConsumerGroupMessage{ClientID: clientID, ConsumerMessage: m}
-		atomic.AddInt32(&s.count, 1)
+		s.count.Add(1)
 	}
 }
 
@@ -378,18 +378,19 @@ func (s *testFuncConsumerGroupSink) Close() map[string][]string {
 
 type testFuncConsumerGroupMember struct {
 	ConsumerGroup
-	clientID     string
-	claims       map[string]int
-	generationId int32
-	state        int32
-	handlers     int32
-	errs         []error
-	maxMessages  int32
-	isCapped     bool
-	sink         *testFuncConsumerGroupSink
+	t        *testing.T
+	clientID string
+	isCapped bool
+	sink     *testFuncConsumerGroupSink
 
-	t  *testing.T
-	mu sync.RWMutex
+	generationId atomic.Int32
+	state        atomic.Int32
+	handlers     atomic.Int32
+	maxMessages  atomic.Int32
+
+	mu     sync.RWMutex
+	claims map[string]int
+	errs   []error
 }
 
 func defaultConfig(clientID string) *Config {
@@ -443,11 +444,11 @@ func runTestFuncConsumerGroupMemberWithConfig(
 		ConsumerGroup: group,
 		clientID:      config.ClientID,
 		claims:        make(map[string]int),
-		maxMessages:   maxMessages,
 		isCapped:      maxMessages != 0,
 		sink:          sink,
 		t:             t,
 	}
+	member.maxMessages.Store(maxMessages)
 	go member.loop(topics)
 	return member
 }
@@ -480,15 +481,15 @@ func (m *testFuncConsumerGroupMember) WaitForState(expected int32) {
 	m.t.Helper()
 
 	m.waitFor("state", expected, func() (interface{}, error) {
-		return atomic.LoadInt32(&m.state), nil
+		return m.state.Load(), nil
 	})
 }
 
 func (m *testFuncConsumerGroupMember) WaitForHandlers(expected int) {
 	m.t.Helper()
 
 	m.waitFor("handlers", expected, func() (interface{}, error) {
-		return int(atomic.LoadInt32(&m.handlers)), nil
+		return int(m.handlers.Load()), nil
 	})
 }
 
@@ -518,17 +519,17 @@ func (m *testFuncConsumerGroupMember) Setup(s ConsumerGroupSession) error {
 	m.mu.Unlock()
 
 	// store generationID
-	atomic.StoreInt32(&m.generationId, s.GenerationID())
+	m.generationId.Store(s.GenerationID())
 
 	// enter post-setup state
-	atomic.StoreInt32(&m.state, 2)
+	m.state.Store(2)
 	return nil
 }
 
 func (m *testFuncConsumerGroupMember) Cleanup(s ConsumerGroupSession) error {
 	m.t.Logf("Consumer %s: session ended %s in generation %d", m.clientID, s.MemberID(), s.GenerationID())
 	// enter post-cleanup state
-	atomic.StoreInt32(&m.state, 3)
+	m.state.Store(3)
 	return nil
 }
 
@@ -538,8 +539,8 @@ func (m *testFuncConsumerGroupMember) ConsumeClaim(s ConsumerGroupSession, c Con
 			m.t.Errorf("panic in ConsumeClaim: %v", r)
 		}
 	}()
-	atomic.AddInt32(&m.handlers, 1)
-	defer atomic.AddInt32(&m.handlers, -1)
+	m.handlers.Add(1)
+	defer m.handlers.Add(-1)
 
 	consumed := 0
 	for {
@@ -549,7 +550,7 @@ func (m *testFuncConsumerGroupMember) ConsumeClaim(s ConsumerGroupSession, c Con
 				m.t.Logf("Consumer %s: message channel closed, consumed %d messages", m.clientID, consumed)
 				return nil
 			}
-			if n := atomic.AddInt32(&m.maxMessages, -1); m.isCapped && n < 0 {
+			if n := m.maxMessages.Add(-1); m.isCapped && n < 0 {
 				m.t.Logf("Consumer %s: reached max messages, consumed %d messages", m.clientID, consumed)
 				return nil
 			}
@@ -597,7 +598,7 @@ func (m *testFuncConsumerGroupMember) loop(topics []string) {
 		if r := recover(); r != nil {
 			m.t.Errorf("panic in loop for %s: %v", m.clientID, r)
 		}
-		atomic.StoreInt32(&m.state, 4)
+		m.state.Store(4)
 	}()
 
 	go func() {
@@ -614,7 +615,7 @@ func (m *testFuncConsumerGroupMember) loop(topics []string) {
 	ctx := context.Background()
 	for {
 		// set state to pre-consume
-		atomic.StoreInt32(&m.state, 1)
+		m.state.Store(1)
 
 		if err := m.Consume(ctx, topics, m); errors.Is(err, ErrClosedConsumerGroup) {
 			m.t.Logf("Consumer %s: closed consumer group", m.clientID)
@@ -634,7 +635,7 @@ func (m *testFuncConsumerGroupMember) loop(topics []string) {
 		}
 
 		// return if capped
-		if n := atomic.LoadInt32(&m.maxMessages); m.isCapped && n < 0 {
+		if n := m.maxMessages.Load(); m.isCapped && n < 0 {
 			m.t.Logf("Consumer %s: reached max messages, returning from loop", m.clientID)
 			return
 		}
@@ -651,7 +652,7 @@ func newTestStatefulStrategy(t *testing.T) *testStatefulStrategy {
 type testStatefulStrategy struct {
 	BalanceStrategy
 	t       *testing.T
-	initial int32
+	initial atomic.Int32
 	state   sync.Map
 }
 
@@ -664,7 +665,7 @@ func (h *testStatefulStrategy) Plan(members map[string]ConsumerGroupMemberMetada
 	for memberID, metadata := range members {
 		if !strings.HasSuffix(string(metadata.UserData), "-stateful") {
 			metadata.UserData = []byte(string(metadata.UserData) + "-stateful")
-			atomic.AddInt32(&h.initial, 1)
+			h.initial.Add(1)
 		}
 		h.state.Store(memberID, metadata.UserData)
 	}
@@ -680,7 +681,7 @@ func (h *testStatefulStrategy) AssignmentData(memberID string, topics map[string
 
 func (h *testStatefulStrategy) AssertInitialValues(count int32) {
 	h.t.Helper()
-	actual := atomic.LoadInt32(&h.initial)
+	actual := h.initial.Load()
 	if actual != count {
 		h.t.Fatalf("unexpected count of initial values: %d, expected: %d", actual, count)
 	}

functional_consumer_staticmembership_test.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"math"
 	"reflect"
-	"sync/atomic"
 	"testing"
 )
 
@@ -107,7 +106,7 @@ func TestFuncConsumerGroupStaticMembership_RejoinAndLeave(t *testing.T) {
 		t.Errorf("should have 2 members in group , got %v\n", len(res1[0].Members))
 	}
 
-	generationId1 := m1.generationId
+	generationId1 := m1.generationId.Load()
 
 	// shut down m2, membership should not change (we didn't leave group when close)
 	m2.AssertCleanShutdown()
@@ -123,7 +122,7 @@ func TestFuncConsumerGroupStaticMembership_RejoinAndLeave(t *testing.T) {
 		t.Errorf("group description be the same before %s, after %s", res1Bytes, res2Bytes)
 	}
 
-	generationId2 := atomic.LoadInt32(&m1.generationId)
+	generationId2 := m1.generationId.Load()
 	if generationId2 != generationId1 {
 		t.Errorf("m1 generation should not increase expect %v, actual %v", generationId1, generationId2)
 	}
@@ -147,7 +146,7 @@ func TestFuncConsumerGroupStaticMembership_RejoinAndLeave(t *testing.T) {
 		t.Errorf("should have 2 members in group , got %v\n", len(res3[0].Members))
 	}
 
-	generationId3 := atomic.LoadInt32(&m1.generationId)
+	generationId3 := m1.generationId.Load()
 	if generationId3 != generationId1 {
 		t.Errorf("m1 generation should not increase expect %v, actual %v", generationId1, generationId2)
 	}
@@ -162,7 +161,7 @@ func TestFuncConsumerGroupStaticMembership_RejoinAndLeave(t *testing.T) {
 	}
 	m1.WaitForHandlers(4)
 
-	generationId4 := atomic.LoadInt32(&m1.generationId)
+	generationId4 := m1.generationId.Load()
 	if generationId4 == generationId1 {
 		t.Errorf("m1 generation should increase expect %v, actual %v", generationId1, generationId2)
 	}

mocks/consumer.go

@@ -398,7 +398,8 @@ func (pc *PartitionConsumer) YieldMessage(msg *sarama.ConsumerMessage) *Partitio
 	msg.Partition = pc.partition
 
 	if pc.paused {
-		msg.Offset = atomic.AddInt64(&pc.suppressedHighWaterMarkOffset, 1) - 1
+		msg.Offset = pc.suppressedHighWaterMarkOffset
+		pc.suppressedHighWaterMarkOffset++
 		pc.suppressedMessages <- msg
 	} else {
 		msg.Offset = pc.highWaterMarkOffset.Add(1) - 1

offset_manager_test.go

@@ -112,7 +112,7 @@ func TestNewOffsetManager(t *testing.T) {
 // Test that the correct sequence of offset commit messages is sent to a broker when
 // multiple goroutines for a group are committing offsets at the same time
 func TestOffsetManagerCommitSequence(t *testing.T) {
-	lastOffset := map[int32]int64{}
+	lastOffset := make(map[int32]int64)
 	var outOfOrder atomic.Pointer[string]
 	seedBroker := NewMockBroker(t, 1)
 	defer seedBroker.Close()
@@ -384,9 +384,9 @@ func TestOffsetManagerFetchInitialFail(t *testing.T) {
 
 // Test fetchInitialOffset retry on ErrOffsetsLoadInProgress
 func TestOffsetManagerFetchInitialLoadInProgress(t *testing.T) {
-	retryCount := int32(0)
+	var retryCount atomic.Int32
 	backoff := func(retries, maxRetries int) time.Duration {
-		atomic.AddInt32(&retryCount, 1)
+		retryCount.Add(1)
 		return 0
 	}
 	om, testClient, broker, coordinator := initOffsetManagerWithBackoffFunc(t, 0, backoff, NewTestConfig())
@@ -420,7 +420,7 @@ func TestOffsetManagerFetchInitialLoadInProgress(t *testing.T) {
 	safeClose(t, om)
 	safeClose(t, testClient)
 
-	if atomic.LoadInt32(&retryCount) == 0 {
+	if retryCount.Load() == 0 {
 		t.Fatal("Expected at least one retry")
 	}
 }