1package rand
2
3import (
4	"math/rand"
5	"sync/atomic"
6)
7
8// Pool represents a pool of random number generators.
9// To generate a random id, round robin through the source pool with atomic increment.
10// With more and more goroutines, Pool improves the performance of Random vs naive global random
11// mutex exponentially.
12// Try tests with 20000 goroutines and 500 calls to observe the difference
13type Pool struct {
14	sources []NumberGenerator
15	counter uint64 // used for round robin
16	size    uint64
17}
18
19// see bit twiddling hacks: https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
20func nextNearestPow2uint64(v uint64) uint64 {
21	v--
22	v |= v >> 1
23	v |= v >> 2
24	v |= v >> 4
25	v |= v >> 8
26	v |= v >> 16
27	v |= v >> 32
28	v++
29	return v
30}
31
32// NewPool takes in a size and creates a pool of random id generators with size equal to next closest power of 2.
33// eg: NewPool(10) returns a pool with 2^4 = 16 random sources.
34func NewPool(seed int64, size uint64) *Pool {
35	groupsize := nextNearestPow2uint64(size)
36	pool := &Pool{
37		size:    groupsize,
38		sources: make([]NumberGenerator, groupsize),
39	}
40	// seed the pool
41	pool.seed(seed)
42	return pool
43}
44
45// seed initializes the pool using a randomized sequence with given seed.
46func (r *Pool) seed(seed int64) {
47	// init a random sequence to seed all sources
48	seedRan := rand.NewSource(seed)
49	for i := uint64(0); i < r.size; i++ {
50		r.sources[i] = NewLockedRand(seedRan.Int63())
51	}
52}
53
54// Pick returns a NumberGenerator from a pool of NumberGenerators
55func (r *Pool) Pick() NumberGenerator {
56	// use round robin with fast modulus of pow2 numbers
57	selection := atomic.AddUint64(&r.counter, 1) & (r.size - 1)
58	return r.sources[selection]
59}
60