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
|
package pool
|
||||||
|
|
||||||
import "math/rand"
|
import (
|
||||||
|
"math/rand"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
// sampler implements weighted random number generation using Vose's Alias
|
// sampler implements weighted random number generation using Vose's Alias
|
||||||
// Method (https://www.keithschwarz.com/darts-dice-coins/).
|
// Method (https://www.keithschwarz.com/darts-dice-coins/).
|
||||||
type sampler struct {
|
type sampler struct {
|
||||||
|
mu sync.Mutex
|
||||||
randomGenerator *rand.Rand
|
randomGenerator *rand.Rand
|
||||||
probabilities []float64
|
|
||||||
alias []int
|
probabilities []float64
|
||||||
|
alias []int
|
||||||
}
|
}
|
||||||
|
|
||||||
// newSampler creates new sampler with a given set of probabilities using
|
// 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.
|
// 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 {
|
func (g *sampler) Next() int {
|
||||||
n := len(g.alias)
|
n := len(g.alias)
|
||||||
|
|
||||||
|
g.mu.Lock()
|
||||||
i := g.randomGenerator.Intn(n)
|
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 i
|
||||||
}
|
}
|
||||||
return g.alias[i]
|
return g.alias[i]
|
||||||
|
|
Loading…
Reference in a new issue