github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/workloadimpl/precomputedrand.go (about)

     1  // Copyright 2019 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 workloadimpl
    12  
    13  import (
    14  	"sync"
    15  
    16  	"golang.org/x/exp/rand"
    17  )
    18  
    19  // PrecomputedRand is a precomputed sequence of random data in some alphabet.
    20  type PrecomputedRand []byte
    21  
    22  // PrecomputedRandInit returns a init function that lazily initializes and
    23  // returns a PrecomputedRand. This initialization work is done once and the
    24  // result is shared, subsequent calls to return this shared one. The init
    25  // function is concurrency safe.
    26  func PrecomputedRandInit(rng rand.Source, length int, alphabet string) func() PrecomputedRand {
    27  	var prOnce sync.Once
    28  	var pr PrecomputedRand
    29  	return func() PrecomputedRand {
    30  		prOnce.Do(func() {
    31  			pr = make(PrecomputedRand, length)
    32  			RandStringFast(rng, pr, alphabet)
    33  		})
    34  		return pr
    35  	}
    36  }
    37  
    38  // FillBytes fills the given buffer with precomputed random data, starting at
    39  // the given offset (which is like a seed) and returning a new offset to be used
    40  // on the next call. FillBytes is concurrency safe.
    41  func (pr PrecomputedRand) FillBytes(offset int, buf []byte) int {
    42  	if len(pr) == 0 {
    43  		panic(`cannot fill from empty precomputed rand`)
    44  	}
    45  	prIdx := offset
    46  	for bufIdx := 0; bufIdx < len(buf); {
    47  		if prIdx == len(pr) {
    48  			prIdx = 0
    49  		}
    50  		need, remaining := len(buf)-bufIdx, len(pr)-prIdx
    51  		copyLen := need
    52  		if copyLen > remaining {
    53  			copyLen = remaining
    54  		}
    55  		newBufIdx, newPRIdx := bufIdx+copyLen, prIdx+copyLen
    56  		copy(buf[bufIdx:newBufIdx], pr[prIdx:newPRIdx])
    57  		bufIdx = newBufIdx
    58  		prIdx = newPRIdx
    59  	}
    60  	return prIdx
    61  }