github.com/yandex/pandora@v0.5.32/lib/mp/iterator.go (about)

     1  package mp
     2  
     3  import (
     4  	"math/rand"
     5  	"sync"
     6  	"sync/atomic"
     7  )
     8  
     9  type Iterator interface {
    10  	Next(segment string) int
    11  	Rand(length int) int
    12  }
    13  
    14  func NewNextIterator(seed int64) *NextIterator {
    15  	return &NextIterator{
    16  		gs:  make(map[string]*atomic.Uint64),
    17  		rnd: rand.New(rand.NewSource(seed)),
    18  	}
    19  }
    20  
    21  type NextIterator struct {
    22  	mx  sync.Mutex
    23  	gs  map[string]*atomic.Uint64
    24  	rnd *rand.Rand
    25  }
    26  
    27  func (n *NextIterator) Rand(length int) int {
    28  	return n.rnd.Intn(length)
    29  }
    30  
    31  func (n *NextIterator) Next(segment string) int {
    32  	n.mx.Lock()
    33  	defer n.mx.Unlock()
    34  	a, ok := n.gs[segment]
    35  	if !ok {
    36  		n.gs[segment] = &atomic.Uint64{}
    37  		return 0
    38  	}
    39  	add := a.Add(1)
    40  	return int(add)
    41  }