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

     1  // Copyright 2018 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 ycsb
    12  
    13  import (
    14  	"math/rand"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    17  )
    18  
    19  // UniformGenerator is a random number generator that generates draws from a
    20  // uniform distribution.
    21  type UniformGenerator struct {
    22  	iMin uint64
    23  	mu   struct {
    24  		syncutil.Mutex
    25  		r    *rand.Rand
    26  		iMax uint64
    27  	}
    28  }
    29  
    30  // NewUniformGenerator constructs a new UniformGenerator with the given parameters.
    31  // It returns an error if the parameters are outside the accepted range.
    32  func NewUniformGenerator(rng *rand.Rand, iMin, iMax uint64) (*UniformGenerator, error) {
    33  
    34  	z := UniformGenerator{}
    35  	z.iMin = iMin
    36  	z.mu.r = rng
    37  	z.mu.iMax = iMax
    38  
    39  	return &z, nil
    40  }
    41  
    42  // IncrementIMax increments iMax by count.
    43  func (z *UniformGenerator) IncrementIMax(count uint64) error {
    44  	z.mu.Lock()
    45  	defer z.mu.Unlock()
    46  	z.mu.iMax += count
    47  	return nil
    48  }
    49  
    50  // Uint64 returns a random Uint64 between iMin and iMax, drawn from a uniform
    51  // distribution.
    52  func (z *UniformGenerator) Uint64() uint64 {
    53  	z.mu.Lock()
    54  	defer z.mu.Unlock()
    55  	return (uint64)(z.mu.r.Int63n((int64)(z.mu.iMax-z.iMin+1))) + z.iMin
    56  }