frostfs-sdk-go/pool/sampler.go
Denis Kirillov c5c6272029
All checks were successful
DCO / DCO (pull_request) Successful in 1m4s
Tests and linters / Tests (1.21) (pull_request) Successful in 1m20s
Tests and linters / Tests (1.20) (pull_request) Successful in 1m31s
Tests and linters / Lint (pull_request) Successful in 2m25s
[#221] pool: Make sampler safe for concurrent using
Signed-off-by: Denis Kirillov <d.kirillov@yadro.com>
2024-05-20 14:14:14 +03:00

92 lines
1.9 KiB
Go

package pool
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
}
// newSampler creates new sampler with a given set of probabilities using
// given source of randomness. Created sampler will produce numbers from
// 0 to len(probabilities).
func newSampler(probabilities []float64, source rand.Source) *sampler {
sampler := &sampler{}
var (
small workList
large workList
)
n := len(probabilities)
sampler.randomGenerator = rand.New(source)
sampler.probabilities = make([]float64, n)
sampler.alias = make([]int, n)
// Compute scaled probabilities.
p := make([]float64, n)
for i := 0; i < n; i++ {
p[i] = probabilities[i] * float64(n)
}
for i, pi := range p {
if pi < 1 {
small.add(i)
} else {
large.add(i)
}
}
for len(small) > 0 && len(large) > 0 {
l, g := small.remove(), large.remove()
sampler.probabilities[l] = p[l]
sampler.alias[l] = g
p[g] = p[g] + p[l] - 1
if p[g] < 1 {
small.add(g)
} else {
large.add(g)
}
}
for len(large) > 0 {
g := large.remove()
sampler.probabilities[g] = 1
}
for len(small) > 0 {
l := small.remove()
sampler.probabilities[l] = 1
}
return 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)
f := g.randomGenerator.Float64()
g.mu.Unlock()
if f < g.probabilities[i] {
return i
}
return g.alias[i]
}
type workList []int
func (wl *workList) add(e int) {
*wl = append(*wl, e)
}
func (wl *workList) remove() int {
l := len(*wl) - 1
n := (*wl)[l]
*wl = (*wl)[:l]
return n
}