github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/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  
    20  	"github.com/cockroachdb/cockroach/pkg/util/envutil"
    21  )
    22  
    23  // NewPseudoSeed generates a seed from crypto/rand.
    24  func NewPseudoSeed() int64 {
    25  	var seed int64
    26  	err := binary.Read(crypto_rand.Reader, binary.LittleEndian, &seed)
    27  	if err != nil {
    28  		panic(fmt.Sprintf("could not read from crypto/rand: %s", err))
    29  	}
    30  	return seed
    31  }
    32  
    33  // NewPseudoRand returns an instance of math/rand.Rand seeded from the
    34  // environment variable COCKROACH_RANDOM_SEED.  If that variable is not set,
    35  // crypto/rand is used to generate a seed. The seed is also returned so we can
    36  // easily and cheaply generate unique streams of numbers. The created object is
    37  // not safe for concurrent access.
    38  func NewPseudoRand() (*rand.Rand, int64) {
    39  	seed := envutil.EnvOrDefaultInt64("COCKROACH_RANDOM_SEED", NewPseudoSeed())
    40  	return rand.New(rand.NewSource(seed)), seed
    41  }
    42  
    43  // RandIntInRange returns a value in [min, max)
    44  func RandIntInRange(r *rand.Rand, min, max int) int {
    45  	return min + r.Intn(max-min)
    46  }
    47  
    48  var randLetters = []byte("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
    49  
    50  // RandBytes returns a byte slice of the given length with random
    51  // data.
    52  func RandBytes(r *rand.Rand, size int) []byte {
    53  	if size <= 0 {
    54  		return nil
    55  	}
    56  
    57  	arr := make([]byte, size)
    58  	for i := 0; i < len(arr); i++ {
    59  		arr[i] = randLetters[r.Intn(len(randLetters))]
    60  	}
    61  	return arr
    62  }
    63  
    64  // ReadTestdataBytes reads random bytes, but then nudges them into printable
    65  // ASCII, *reducing their randomness* to make them a little friendlier for
    66  // humans using them as testdata.
    67  func ReadTestdataBytes(r *rand.Rand, arr []byte) {
    68  	_, _ = r.Read(arr)
    69  	for i := range arr {
    70  		arr[i] = arr[i] & 0x7F // mask out non-ascii
    71  		if arr[i] < ' ' {      // Nudge the control chars up, into the letters.
    72  			arr[i] += 'A'
    73  		}
    74  	}
    75  }
    76  
    77  // SeedForTests seeds the random number generator and prints the seed
    78  // value used. This value can be specified via an environment variable
    79  // COCKROACH_RANDOM_SEED=x to reuse the same value later. This function should
    80  // be called from TestMain; individual tests should not touch the seed
    81  // of the global random number generator.
    82  func SeedForTests() {
    83  	seed := envutil.EnvOrDefaultInt64("COCKROACH_RANDOM_SEED", NewPseudoSeed())
    84  	rand.Seed(seed)
    85  	log.Printf("Random seed: %v", seed)
    86  }