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

最近在做负载均衡，需要制作一个可以并发递增的计数器，用来选取worker，并且在特定的数值需要归零，用代码就是：

代码语言：javascript复制

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

第一版用锁可以做到。但想了一下，其实原子变量也能解决这个问题。

代码语言：javascript复制

// 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)
}

测试用例

代码语言：javascript复制

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。感觉应该会比加锁好一点点。

负载均衡 worker 变量 并发 性能

0 人点赞