github.com/polarismesh/polaris@v1.17.8/common/srand/scalable_rand.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package srand
    19  
    20  import (
    21  	"math/rand"
    22  	"runtime"
    23  	"sync"
    24  	"sync/atomic"
    25  	"time"
    26  )
    27  
    28  // ScalableRand 可水平扩展的随机数发生器
    29  type ScalableRand struct {
    30  	initSeed int64
    31  	randPool *sync.Pool
    32  }
    33  
    34  // NewScalableRand 初始化随机数发生器
    35  func NewScalableRand() *ScalableRand {
    36  	scalableRand := &ScalableRand{
    37  		randPool: &sync.Pool{},
    38  	}
    39  	cpuCount := runtime.NumCPU()
    40  	for i := 0; i < cpuCount; i++ {
    41  		scalableRand.randPool.Put(rand.New(rand.NewSource(scalableRand.getRandSeed())))
    42  	}
    43  	return scalableRand
    44  }
    45  
    46  // getRandSeed 循环并获取唯一的随机数种子
    47  func (s *ScalableRand) getRandSeed() int64 {
    48  	var seed int64
    49  	for {
    50  		seed = time.Now().UnixNano()
    51  		if s.getAndSetInitSeed(seed) {
    52  			break
    53  		}
    54  		time.Sleep(1)
    55  	}
    56  	return seed
    57  }
    58  
    59  // getAndSetInitSeed 获取并比较种子数
    60  func (s *ScalableRand) getAndSetInitSeed(seed int64) bool {
    61  	initSeed := atomic.LoadInt64(&s.initSeed)
    62  	if initSeed == seed {
    63  		return false
    64  	}
    65  	return atomic.CompareAndSwapInt64(&s.initSeed, initSeed, seed)
    66  }
    67  
    68  // Intn 获取随机数
    69  func (s *ScalableRand) Intn(n int) int {
    70  	var randSeed *rand.Rand
    71  	value := s.randPool.Get()
    72  	if value != nil {
    73  		randSeed = value.(*rand.Rand)
    74  	} else {
    75  		randSeed = rand.New(rand.NewSource(s.getRandSeed()))
    76  	}
    77  	randValue := randSeed.Intn(n)
    78  	s.randPool.Put(randSeed)
    79  	return randValue
    80  }
    81  
    82  // 全局随机种子
    83  var globalRand *ScalableRand
    84  
    85  // Intn 返回全局随机数
    86  func Intn(n int) int {
    87  	return globalRand.Intn(n)
    88  }
    89  
    90  // init 初始化全局随机种子
    91  func init() {
    92  	globalRand = NewScalableRand()
    93  }