Fixed a data race in the Hash balancer when using custom hasher (segm…

…entio#335)
ravi4422 · Aug 9, 2019 · 433ef07 · 433ef07
1 parent 11670d5
commit 433ef07
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/balancer.go b/balancer.go
@@ -134,6 +134,11 @@ var (
 type Hash struct {
 	rr     RoundRobin
 	Hasher hash.Hash32
+
+	// lock protects Hasher while calculating the hash code.  It is assumed that
+	// the Hasher field is read-only once the Balancer is created, so as a
+	// performance optimization, reads of the field are not protected.
+	lock sync.Mutex
 }
 
 func (h *Hash) Balance(msg Message, partitions ...int) (partition int) {
@@ -142,7 +147,10 @@ func (h *Hash) Balance(msg Message, partitions ...int) (partition int) {
 	}
 
 	hasher := h.Hasher
-	if hasher == nil {
+	if hasher != nil {
+		h.lock.Lock()
+		defer h.lock.Unlock()
+	} else {
 		hasher = fnv1aPool.Get().(hash.Hash32)
 		defer fnv1aPool.Put(hasher)
 	}