go: 使用原子变量做一个可自动归零的计数器 最近在做负载均衡，需要制作一个可以并发递增的计数器，用来选取worker，并且在特定的数值需要归零，用代码就是： ``` counter.SetMax(len(worker)) ..... // 并发，均衡的选取worker index := counter.Add() workers[index].Run(args) ``` 第一版用锁可以做到。但想了一下，其实原子变量也能解决这个问题。注意，SetMax和Add不能并发。不过max一般和workers的长度绑定，会加读写锁保护。 ``` // 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() } } ``` 性能每次6ns。感觉应该会比加锁好一点点。 来自 大脸猪 写于 2021-06-29 20:24 -- 更新于2021-06-30 10:38 -- 0 条评论