From 2af3101bdba0698ff97117cd2b0051510d996df7 Mon Sep 17 00:00:00 2001 From: Peter Dannemann <28637185+petedannemann@users.noreply.github.com> Date: Wed, 13 Dec 2023 09:30:14 -0500 Subject: [PATCH] fix: data race in roundrobin balancer (#1251) --- balancer.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/balancer.go b/balancer.go index 4136fce7..ee3a2588 100644 --- a/balancer.go +++ b/balancer.go @@ -7,7 +7,6 @@ import ( "math/rand" "sort" "sync" - "sync/atomic" ) // The Balancer interface provides an abstraction of the message distribution @@ -42,8 +41,10 @@ func (f BalancerFunc) Balance(msg Message, partitions ...int) int { type RoundRobin struct { ChunkSize int // Use a 32 bits integer so RoundRobin values don't need to be aligned to - // apply atomic increments. + // apply increments. counter uint32 + + mutex sync.Mutex } // Balance satisfies the Balancer interface. @@ -52,14 +53,17 @@ func (rr *RoundRobin) Balance(msg Message, partitions ...int) int { } func (rr *RoundRobin) balance(partitions []int) int { + rr.mutex.Lock() + defer rr.mutex.Unlock() + if rr.ChunkSize < 1 { rr.ChunkSize = 1 } length := len(partitions) - counterNow := atomic.LoadUint32(&rr.counter) + counterNow := rr.counter offset := int(counterNow / uint32(rr.ChunkSize)) - atomic.AddUint32(&rr.counter, 1) + rr.counter++ return partitions[offset%length] }