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]
 }