-
Notifications
You must be signed in to change notification settings - Fork 70
/
rand.go
76 lines (66 loc) · 1.59 KB
/
rand.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
package pkg
import (
"math/rand"
"time"
)
// randIDGenerator generates random ID that combines round-robin and zipf distribution and make sure not duplicated.
type randIDGenerator struct {
allocated map[uint64]struct{}
zipf *rand.Zipf
uniform *rand.Rand
tableNum int
max uint64
partitionSize uint64
partitionIdx uint64
}
func (r *randIDGenerator) nextTableAndRowPairs() (fromTableID, toTableID int, fromRowID, toRowID uint64) {
fromTableID = r.nextTableID()
toTableID = r.nextTableID()
fromRowID = r.nextRowID()
toRowID = r.nextRowID()
return
}
func (r *randIDGenerator) nextTableID() int {
return r.uniform.Intn(r.tableNum)
}
func (r *randIDGenerator) nextRowID() uint64 {
for {
v := r.zipf.Uint64()
if _, ok := r.allocated[v]; ok {
continue
}
r.allocated[v] = struct{}{}
return v
}
}
func (r *randIDGenerator) nextNonExistRowID() uint64 {
for {
v := r.zipf.Uint64() + r.max
if _, ok := r.allocated[v]; ok {
continue
}
r.allocated[v] = struct{}{}
return v
}
}
func (r *randIDGenerator) reset() {
r.allocated = map[uint64]struct{}{}
}
func (r *randIDGenerator) nextNumStatements(n int) int {
return r.uniform.Intn(n)
}
func (r *randIDGenerator) nextUniqueIndex() uint64 {
return r.uniform.Uint64()
}
func newRandIDGenerator(tableNum int, maxSize uint64) *randIDGenerator {
src := rand.NewSource(time.Now().UnixNano())
ran := rand.New(src)
zipf := rand.NewZipf(ran, 1.01, 1, maxSize-1)
return &randIDGenerator{
allocated: map[uint64]struct{}{},
zipf: zipf,
uniform: ran,
tableNum: tableNum,
max: maxSize,
}
}