github.com/cockroachdb/cockroachdb-parser@v0.23.3-0.20240213214944-911057d40c9a/pkg/util/randutil/rand.go (about) 1 // Copyright 2014 The Cockroach Authors. 2 // 3 // Use of this software is governed by the Business Source License 4 // included in the file licenses/BSL.txt. 5 // 6 // As of the Change Date specified in that file, in accordance with 7 // the Business Source License, use of this software will be governed 8 // by the Apache License, Version 2.0, included in the file 9 // licenses/APL.txt. 10 11 package randutil 12 13 import ( 14 crypto_rand "crypto/rand" 15 "encoding/binary" 16 "fmt" 17 "log" // Don't bring cockroach/util/log into this low-level package. 18 "math/rand" 19 "runtime" 20 "strings" 21 "time" 22 _ "unsafe" // required by go:linkname 23 24 "github.com/cockroachdb/cockroachdb-parser/pkg/util/envutil" 25 "github.com/cockroachdb/cockroachdb-parser/pkg/util/syncutil" 26 ) 27 28 // lockedSource is a thread safe math/rand.Source. See math/rand/rand.go. 29 type lockedSource struct { 30 mu syncutil.Mutex 31 src rand.Source64 32 } 33 34 // NewLockedSource creates random source protected by mutex. 35 func NewLockedSource(seed int64) rand.Source { 36 return &lockedSource{ 37 src: rand.NewSource(seed).(rand.Source64), 38 } 39 } 40 41 func (rng *lockedSource) Int63() (n int64) { 42 rng.mu.Lock() 43 defer rng.mu.Unlock() 44 n = rng.src.Int63() 45 return 46 } 47 48 func (rng *lockedSource) Uint64() (n uint64) { 49 rng.mu.Lock() 50 defer rng.mu.Unlock() 51 n = rng.src.Uint64() 52 return 53 } 54 55 func (rng *lockedSource) Seed(seed int64) { 56 rng.mu.Lock() 57 defer rng.mu.Unlock() 58 rng.src.Seed(seed) 59 } 60 61 // globalSeed contains a pseudo random seed that should only be used in tests. 62 var globalSeed int64 63 64 // rng is a random number generator used to generate seeds for test random 65 // number generators. 66 var rng *rand.Rand 67 68 // lastTestName is the function name of the last test we have seen. 69 var lastTestName string 70 71 // mtx protects rng and lastTestName. 72 var mtx syncutil.Mutex 73 74 // Initializes the global random seed. This value can be specified via an 75 // environment variable COCKROACH_RANDOM_SEED=x. 76 func init() { 77 globalSeed = envutil.EnvOrDefaultInt64("COCKROACH_RANDOM_SEED", NewPseudoSeed()) 78 rng = rand.New(rand.NewSource(globalSeed)) 79 } 80 81 // NewPseudoSeed generates a seed from crypto/rand. 82 func NewPseudoSeed() int64 { 83 var seed int64 84 err := binary.Read(crypto_rand.Reader, binary.LittleEndian, &seed) 85 if err != nil { 86 panic(fmt.Sprintf("could not read from crypto/rand: %s", err)) 87 } 88 return seed 89 } 90 91 // NewPseudoRand returns an instance of math/rand.Rand seeded from the 92 // environment variable COCKROACH_RANDOM_SEED. If that variable is not set, 93 // crypto/rand is used to generate a seed. The seed is also returned so we can 94 // easily and cheaply generate unique streams of numbers. The created object is 95 // not safe for concurrent access. 96 func NewPseudoRand() (*rand.Rand, int64) { 97 seed := envutil.EnvOrDefaultInt64("COCKROACH_RANDOM_SEED", NewPseudoSeed()) 98 return rand.New(rand.NewSource(seed)), seed 99 } 100 101 // Same as NewPseudoRand, but the returned Rand is using thread safe underlying source. 102 func NewLockedPseudoRand() (*rand.Rand, int64) { 103 seed := envutil.EnvOrDefaultInt64("COCKROACH_RANDOM_SEED", NewPseudoSeed()) 104 return rand.New(NewLockedSource(seed)), seed 105 } 106 107 // NewTestRand returns an instance of math/rand.Rand seeded from rng, which is 108 // seeded with the global seed. If the caller is a test with a different 109 // path-qualified name than the previous caller, rng is reseeded from the global 110 // seed. This rand.Rand is useful in testing to produce deterministic, 111 // reproducible behavior. 112 func NewTestRand() (*rand.Rand, int64) { 113 return newTestRandImpl(rand.NewSource) 114 } 115 116 // NewLockedTestRand is identical to NewTestRand but returned rand.Rand is using 117 // thread safe underlying source. 118 func NewLockedTestRand() (*rand.Rand, int64) { 119 return newTestRandImpl(NewLockedSource) 120 } 121 122 func newTestRandImpl(f func(int64) rand.Source) (*rand.Rand, int64) { 123 mtx.Lock() 124 defer mtx.Unlock() 125 fxn := getTestName() 126 if fxn != "" && lastTestName != fxn { 127 // Re-seed rng (the source of seeds for test random number generators) with 128 // the global seed so that individual tests are reproducible using the 129 // random seed. 130 lastTestName = fxn 131 rng = rand.New(f(globalSeed)) 132 } 133 seed := rng.Int63() 134 return rand.New(f(seed)), seed 135 } 136 137 // NewTestRandWithSeed returns an instance of math/rand.Rand, similar to 138 // NewTestRand, but with the seed specified. 139 func NewTestRandWithSeed(seed int64) *rand.Rand { 140 mtx.Lock() 141 defer mtx.Unlock() 142 fxn := getTestName() 143 if fxn != "" && lastTestName != fxn { 144 lastTestName = fxn 145 } 146 return rand.New(rand.NewSource(seed)) 147 } 148 149 // RandIntInRange returns a value in [min, max) 150 func RandIntInRange(r *rand.Rand, min, max int) int { 151 return min + r.Intn(max-min) 152 } 153 154 // RandInt63InRange returns a value in [min, max) 155 func RandInt63InRange(r *rand.Rand, min, max int64) int64 { 156 return min + r.Int63n(max-min) 157 } 158 159 // RandUint64n generates a 64-bit random number in [0, n) range. 160 // Note: n == 0 means n is math.MaxUint64 + 1 161 func RandUint64n(r *rand.Rand, n uint64) uint64 { 162 if n == 0 { 163 return r.Uint64() 164 } 165 // If n is less than 64 bits, delegate to 63 bit version. 166 if n < (1 << 63) { 167 return uint64(r.Int63n(int64(n))) 168 } 169 v := r.Uint64() 170 for v > n { 171 v = r.Uint64() 172 } 173 return v 174 } 175 176 // RandDuration returns a random duration in [0, max). 177 func RandDuration(r *rand.Rand, max time.Duration) time.Duration { 178 return time.Duration(r.Int63n(int64(max))) 179 } 180 181 var randLetters = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ") 182 183 // RandBytes returns a byte slice of the given length with random 184 // data. 185 func RandBytes(r *rand.Rand, size int) []byte { 186 if size <= 0 { 187 return nil 188 } 189 190 arr := make([]byte, size) 191 for i := 0; i < len(arr); i++ { 192 arr[i] = randLetters[r.Intn(len(randLetters))] 193 } 194 return arr 195 } 196 197 // FastUint32 returns a lock free uint32 value. Compared to rand.Uint32, this 198 // implementation scales. We're using the go runtime's implementation through a 199 // linker trick. 200 // 201 //go:linkname FastUint32 runtime.fastrand 202 func FastUint32() uint32 203 204 // FastInt63 returns a non-negative pseudo-random 63-bit integer as an int64. 205 // Compared to rand.Int63(), this implementation scales. 206 func FastInt63() int64 { 207 x, y := FastUint32(), FastUint32() // 32-bit halves 208 u := uint64(x)<<32 ^ uint64(y) 209 i := int64(u >> 1) // clear sign bit 210 return i 211 } 212 213 // ReadTestdataBytes reads random bytes, but then nudges them into printable 214 // ASCII, *reducing their randomness* to make them a little friendlier for 215 // humans using them as testdata. 216 func ReadTestdataBytes(r *rand.Rand, arr []byte) { 217 _, _ = r.Read(arr) 218 for i := range arr { 219 arr[i] = arr[i] & 0x7F // mask out non-ascii 220 if arr[i] < ' ' { // Nudge the control chars up, into the letters. 221 arr[i] += 'A' 222 } 223 } 224 } 225 226 // PrintableKeyAlphabet to use with random string generation to produce strings 227 // that doesn't need to be escaped when found as a part of a key and is 228 // generally human printable. 229 const PrintableKeyAlphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" 230 231 // RandString generates a random string of the desired length from the 232 // input alphabet. It is useful when you want to generate keys that would 233 // be printable without further escaping if alphabet is restricted to 234 // alphanumeric chars. 235 func RandString(rng *rand.Rand, length int, alphabet string) string { 236 runes := []rune(alphabet) 237 buf := &strings.Builder{} 238 for i := 0; i < length; i++ { 239 buf.WriteRune(runes[rng.Intn(len(runes))]) 240 } 241 return buf.String() 242 } 243 244 // SeedForTests seeds the random number generator and prints the seed 245 // value used. This function should be called from TestMain; individual tests 246 // should not touch the seed of the global random number generator. 247 func SeedForTests() { 248 //lint:ignore SA1019 deprecated 249 rand.Seed(globalSeed) 250 log.Printf("random seed: %v", globalSeed) 251 } 252 253 // getTestName returns the calling test function name, returning an empty string 254 // if not found. The number of calls up the call stack is limited. 255 func getTestName() string { 256 pcs := make([]uintptr, 10) 257 n := runtime.Callers(2, pcs) 258 frames := runtime.CallersFrames(pcs[:n]) 259 for { 260 frame, more := frames.Next() 261 fxn := frame.Function 262 if strings.Contains(fxn, ".Test") { 263 return fxn 264 } 265 if !more { 266 break 267 } 268 } 269 return "" 270 }