pool/sampler.go

@@ -1,13 +1,18 @@
 package pool
 
-import "math/rand"
+import (
+	"math/rand"
+	"sync"
+)
 
 // sampler implements weighted random number generation using Vose's Alias
 // Method (https://www.keithschwarz.com/darts-dice-coins/).
 type sampler struct {
+	mu              sync.Mutex
 	randomGenerator *rand.Rand
-	probabilities   []float64
+
-	alias           []int
+	probabilities []float64
+	alias         []int
 }
 
 // newSampler creates new sampler with a given set of probabilities using
@@ -58,10 +63,16 @@ func newSampler(probabilities []float64, source rand.Source) *sampler {
 }
 
 // Next returns the next (not so) random number from sampler.
+// This method is safe for concurrent use by multiple goroutines.
 func (g *sampler) Next() int {
 	n := len(g.alias)
+
+	g.mu.Lock()
 	i := g.randomGenerator.Intn(n)
-	if g.randomGenerator.Float64() < g.probabilities[i] {
+	f := g.randomGenerator.Float64()
+	g.mu.Unlock()
+
+	if f < g.probabilities[i] {
 		return i
 	}
 	return g.alias[i]