Line data Source code
1 : // Copyright 2019 The LevelDB-Go and Pebble Authors. All rights reserved. Use 2 : // of this source code is governed by a BSD-style license that can be found in 3 : // the LICENSE file. 4 : 5 : package randvar 6 : 7 : import ( 8 : "encoding/binary" 9 : "regexp" 10 : "strconv" 11 : "strings" 12 : 13 : "github.com/cockroachdb/errors" 14 : "golang.org/x/exp/rand" 15 : ) 16 : 17 : var randVarRE = regexp.MustCompile(`^(?:(latest|uniform|zipf):)?(\d+)(?:-(\d+))?$`) 18 : 19 : // Flag provides a command line flag interface for specifying static random 20 : // variables. 21 : type Flag struct { 22 : Static 23 : spec string 24 : } 25 : 26 : // NewFlag creates a new Flag initialized with the specified spec. 27 1 : func NewFlag(spec string) *Flag { 28 1 : f := &Flag{} 29 1 : if err := f.Set(spec); err != nil { 30 0 : panic(err) 31 : } 32 1 : return f 33 : } 34 : 35 1 : func (f *Flag) String() string { 36 1 : return f.spec 37 1 : } 38 : 39 : // Type implements the Flag.Value interface. 40 0 : func (f *Flag) Type() string { 41 0 : return "randvar" 42 0 : } 43 : 44 : // Set implements the Flag.Value interface. 45 1 : func (f *Flag) Set(spec string) error { 46 1 : m := randVarRE.FindStringSubmatch(spec) 47 1 : if m == nil { 48 0 : return errors.Errorf("invalid random var spec: %s", errors.Safe(spec)) 49 0 : } 50 : 51 1 : min, err := strconv.Atoi(m[2]) 52 1 : if err != nil { 53 0 : return err 54 0 : } 55 1 : max := min 56 1 : if m[3] != "" { 57 1 : max, err = strconv.Atoi(m[3]) 58 1 : if err != nil { 59 0 : return err 60 0 : } 61 : } 62 : 63 1 : switch strings.ToLower(m[1]) { 64 1 : case "", "uniform": 65 1 : f.Static = NewUniform(uint64(min), uint64(max)) 66 0 : case "latest": 67 0 : f.Static, err = NewSkewedLatest(uint64(min), uint64(max), 0.99) 68 0 : if err != nil { 69 0 : return err 70 0 : } 71 1 : case "zipf": 72 1 : var err error 73 1 : f.Static, err = NewZipf(uint64(min), uint64(max), 0.99) 74 1 : if err != nil { 75 0 : return err 76 0 : } 77 0 : default: 78 0 : return errors.Errorf("unknown random var distribution: %s", errors.Safe(m[1])) 79 : } 80 1 : f.spec = spec 81 1 : return nil 82 : } 83 : 84 : // BytesFlag provides a command line flag interface for specifying random 85 : // bytes. The specification provides for both the length of the random bytes 86 : // and a target compression ratio. 87 : type BytesFlag struct { 88 : sizeFlag Flag 89 : targetCompression float64 90 : spec string 91 : } 92 : 93 : // NewBytesFlag creates a new BytesFlag initialized with the specified spec. 94 1 : func NewBytesFlag(spec string) *BytesFlag { 95 1 : f := &BytesFlag{} 96 1 : if err := f.Set(spec); err != nil { 97 0 : panic(err) 98 : } 99 1 : return f 100 : } 101 : 102 1 : func (f *BytesFlag) String() string { 103 1 : return f.spec 104 1 : } 105 : 106 : // Type implements the Flag.Value interface. 107 0 : func (f *BytesFlag) Type() string { 108 0 : return "randbytes" 109 0 : } 110 : 111 : // Set implements the Flag.Value interface. 112 1 : func (f *BytesFlag) Set(spec string) error { 113 1 : parts := strings.Split(spec, "/") 114 1 : if len(parts) == 0 || len(parts) > 2 { 115 0 : return errors.Errorf("invalid randbytes spec: %s", errors.Safe(spec)) 116 0 : } 117 1 : if err := f.sizeFlag.Set(parts[0]); err != nil { 118 0 : return err 119 0 : } 120 1 : f.targetCompression = 1.0 121 1 : if len(parts) == 2 { 122 1 : var err error 123 1 : f.targetCompression, err = strconv.ParseFloat(parts[1], 64) 124 1 : if err != nil { 125 0 : return err 126 0 : } 127 : } 128 1 : f.spec = spec 129 1 : return nil 130 : } 131 : 132 : // Bytes returns random bytes. The length of the random bytes comes from the 133 : // internal sizeFlag. 134 0 : func (f *BytesFlag) Bytes(r *rand.Rand, buf []byte) []byte { 135 0 : size := int(f.sizeFlag.Uint64(r)) 136 0 : uniqueSize := int(float64(size) / f.targetCompression) 137 0 : if uniqueSize < 1 { 138 0 : uniqueSize = 1 139 0 : } 140 0 : if cap(buf) < size { 141 0 : buf = make([]byte, size) 142 0 : } 143 0 : data := buf[:size] 144 0 : offset := 0 145 0 : for offset+8 <= uniqueSize { 146 0 : binary.LittleEndian.PutUint64(data[offset:], r.Uint64()) 147 0 : offset += 8 148 0 : } 149 0 : word := r.Uint64() 150 0 : for offset < uniqueSize { 151 0 : data[offset] = byte(word) 152 0 : word >>= 8 153 0 : offset++ 154 0 : } 155 0 : for offset < size { 156 0 : data[offset] = data[offset-uniqueSize] 157 0 : offset++ 158 0 : } 159 0 : return data 160 : }