go: 使用原子变量做一个可自动归零的计数器

counter.SetMax(len(worker))
.....
// 并发，均衡的选取worker
index := counter.Add()
workers[index].Run(args)

// NewCounterIndex ...
func NewCounterIndex(max int64) *CounterIndex {
    c := new(CounterIndex)
    c.max = max
    return c
}

// CounterIndex ...
type CounterIndex struct {
    current int64
    max     int64
}

// Add ...
func (c *CounterIndex) Add() int64 {
    for {
        ret := atomic.AddInt64(&c.current, 1)
        if ret >= atomic.LoadInt64(&c.max) {
            isChange := atomic.CompareAndSwapInt64(&c.current, ret, 0)
            if isChange {
                return 0
            } else {
                continue
            }
        }
        return ret
    }
}

// SetMax ...
func (c *CounterIndex) SetMax(max int64) {
    atomic.StoreInt64(&c.max, max)
}

// Max ...
func (c *CounterIndex) Max() int64 {
    return atomic.LoadInt64(&c.max)
}

func TestCounterIndex_Add(t *testing.T) {
    c := NewCounterIndex(100)
    counts := make([]int, 100)
    ret := make(chan int64, 1000)
    go func() {
        for item := range ret {
            counts[int(item)]++
        }
    }()
    for i := 0; i < 1000; i++ {
        go func() {
            tp := c.Add()
            ret <- tp
        }()
    }
    time.Sleep(time.Millisecond * 10)
    close(ret)
    logrus.Infof("counts:%v", counts)
    for _, c := range counts {
        if c != 10 {
            t.Fail()
        }
    }
}

func BenchmarkCounterIndexAdd(b *testing.B) {
    c := NewCounterIndex(100)
    for i := 0; i < b.N; i++ {
        c.Add()
    }
}