forked from TrueCloudLab/frostfs-sdk-go
[#221] pool: Make sampler safe for concurrent using
Signed-off-by: Denis Kirillov <d.kirillov@yadro.com>
This commit is contained in:
parent
3de256d05e
commit
c5c6272029
1 changed files with 15 additions and 4 deletions
|
@ -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]
|
||||
|
|
Loading…
Reference in a new issue