github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/xrand/xrand.go (about)

     1  package xrand
     2  
     3  import (
     4  	"math/rand"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  type Rand interface {
    10  	Int64(max int64) int64
    11  	Int(max int) int
    12  	Shuffle(n int, swap func(i, j int))
    13  }
    14  
    15  type r struct {
    16  	m *sync.Mutex
    17  	r *rand.Rand
    18  }
    19  
    20  type option func(r *r)
    21  
    22  func WithLock() option {
    23  	return func(r *r) {
    24  		r.m = &sync.Mutex{}
    25  	}
    26  }
    27  
    28  func WithSeed(seed int64) option {
    29  	return func(r *r) {
    30  		r.r = rand.New(rand.NewSource(seed)) //nolint:gosec
    31  	}
    32  }
    33  
    34  func New(opts ...option) Rand {
    35  	r := &r{
    36  		r: rand.New(rand.NewSource(time.Now().Unix())), //nolint:gosec
    37  	}
    38  	for _, o := range opts {
    39  		if o != nil {
    40  			o(r)
    41  		}
    42  	}
    43  
    44  	return r
    45  }
    46  
    47  func (r *r) int64n(max int64) int64 {
    48  	if r.m != nil {
    49  		r.m.Lock()
    50  		defer r.m.Unlock()
    51  	}
    52  
    53  	return r.r.Int63n(max)
    54  }
    55  
    56  func (r *r) Int64(max int64) int64 {
    57  	return r.int64n(max)
    58  }
    59  
    60  func (r *r) Int(max int) int {
    61  	return int(r.int64n(int64(max)))
    62  }
    63  
    64  func (r *r) Shuffle(n int, swap func(i, j int)) {
    65  	if r.m != nil {
    66  		r.m.Lock()
    67  		defer r.m.Unlock()
    68  	}
    69  
    70  	r.r.Shuffle(n, swap)
    71  }