[#221] pool: Make sampler safe for concurrent using

Signed-off-by: Denis Kirillov <d.kirillov@yadro.com>
This commit is contained in:
Denis Kirillov 2024-05-08 14:54:39 +03:00
parent 3de256d05e
commit c5c6272029

View file

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