diff --git a/async_producer_test.go b/async_producer_test.go index 2855b25f3..e87381ace 100644 --- a/async_producer_test.go +++ b/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) } } diff --git a/client_test.go b/client_test.go index 72e825b22..301d97e49 100644 --- a/client_test.go +++ b/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) } diff --git a/consumer_test.go b/consumer_test.go index 5f796eeb0..4ffb4d9eb 100644 --- a/consumer_test.go +++ b/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() } } diff --git a/functional_consumer_group_test.go b/functional_consumer_group_test.go index 249e1ebf3..be4045e44 100644 --- a/functional_consumer_group_test.go +++ b/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,7 +481,7 @@ 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 }) } @@ -488,7 +489,7 @@ 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) } diff --git a/functional_consumer_staticmembership_test.go b/functional_consumer_staticmembership_test.go index 2cab56cc3..28c1f287a 100644 --- a/functional_consumer_staticmembership_test.go +++ b/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) } diff --git a/mocks/consumer.go b/mocks/consumer.go index 0f6f1ab2a..91074e2e1 100644 --- a/mocks/consumer.go +++ b/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 diff --git a/offset_manager_test.go b/offset_manager_test.go index a1f287637..a981b2adf 100644 --- a/offset_manager_test.go +++ b/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") } }