From c5c6272029fefe895ae899bb0d1d6eeaa01619b1 Mon Sep 17 00:00:00 2001 From: Denis Kirillov Date: Wed, 8 May 2024 14:54:39 +0300 Subject: [PATCH] [#221] pool: Make sampler safe for concurrent using Signed-off-by: Denis Kirillov --- pool/sampler.go | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/pool/sampler.go b/pool/sampler.go index f5b6635..6f46886 100644 --- a/pool/sampler.go +++ b/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]