refactor(all): replace atomic functions with atomic types (#3699)

Co-authored-by: 1911860538 <alxps1911@gmail.com>

contrib/registry/discovery/discovery.go

@@ -24,7 +24,7 @@ type Discovery struct {
 	httpClient *resty.Client
 
 	node    atomic.Value
-	nodeIdx uint64
+	nodeIdx atomic.Uint64
 
 	mutex       sync.RWMutex
 	apps        map[string]*appInfo
@@ -217,11 +217,11 @@ func (d *Discovery) pickNode() string {
 	if !ok || len(nodes) == 0 {
 		return d.config.Nodes[rand.Intn(len(d.config.Nodes))]
 	}
-	return nodes[atomic.LoadUint64(&d.nodeIdx)%uint64(len(nodes))]
+	return nodes[d.nodeIdx.Load()%uint64(len(nodes))]
 }
 
 func (d *Discovery) switchNode() {
-	atomic.AddUint64(&d.nodeIdx, 1)
+	d.nodeIdx.Add(1)
 }
 
 // renew an instance with Discovery

contrib/registry/zookeeper/watcher.go

@@ -21,7 +21,7 @@ type watcher struct {
 	conn   *zk.Conn
 	cancel context.CancelFunc
 
-	first uint32
+	first atomic.Bool
 	// prefix for ZooKeeper paths or keys (used for filtering or identifying watched nodes)
 	prefix string
 	// the name of the service being watched in ZooKeeper
@@ -61,7 +61,7 @@ func (w *watcher) watch(ctx context.Context) {
 
 func (w *watcher) Next() ([]*registry.ServiceInstance, error) {
 	// TODO: multiple calls to Next may lead to inconsistent service instance information
-	if atomic.CompareAndSwapUint32(&w.first, 0, 1) {
+	if w.first.CompareAndSwap(false, true) {
 		return w.getServices()
 	}
 	select {

internal/context/context.go

@@ -11,7 +11,7 @@ type mergeCtx struct {
 	parent1, parent2 context.Context
 
 	done     chan struct{}
-	doneMark uint32
+	doneMark atomic.Bool
 	doneOnce sync.Once
 	doneErr  error
 
@@ -41,7 +41,7 @@ func Merge(parent1, parent2 context.Context) (context.Context, context.CancelFun
 func (mc *mergeCtx) finish(err error) error {
 	mc.doneOnce.Do(func() {
 		mc.doneErr = err
-		atomic.StoreUint32(&mc.doneMark, 1)
+		mc.doneMark.Store(true)
 		close(mc.done)
 	})
 	return mc.doneErr
@@ -73,7 +73,7 @@ func (mc *mergeCtx) Done() <-chan struct{} {
 
 // Err implements context.Context.
 func (mc *mergeCtx) Err() error {
-	if atomic.LoadUint32(&mc.doneMark) != 0 {
+	if mc.doneMark.Load() {
 		return mc.doneErr
 	}
 	var err error

internal/context/context_test.go

@@ -109,8 +109,8 @@ func TestFinish(t *testing.T) {
 	if !errors.Is(err, context.DeadlineExceeded) {
 		t.Errorf("expect %v, got %v", context.DeadlineExceeded, err)
 	}
-	if !reflect.DeepEqual(mc.doneMark, uint32(1)) {
-		t.Errorf("expect %v, got %v", 1, mc.doneMark)
+	if done := mc.doneMark.Load(); done != true {
+		t.Errorf("expect %v, got %v", true, done)
 	}
 	if <-mc.done != struct{}{} {
 		t.Errorf("expect %v, got %v", struct{}{}, <-mc.done)

selector/node/direct/direct.go

@@ -22,7 +22,7 @@ type Node struct {
 	selector.Node
 
 	// last lastPick timestamp
-	lastPick int64
+	lastPick atomic.Int64
 }
 
 // Builder is direct node builder
@@ -30,12 +30,12 @@ type Builder struct{}
 
 // Build create node
 func (*Builder) Build(n selector.Node) selector.WeightedNode {
-	return &Node{Node: n, lastPick: 0}
+	return &Node{Node: n, lastPick: atomic.Int64{}}
 }
 
 func (n *Node) Pick() selector.DoneFunc {
 	now := time.Now().UnixNano()
-	atomic.StoreInt64(&n.lastPick, now)
+	n.lastPick.Store(now)
 	return func(context.Context, selector.DoneInfo) {}
 }
 
@@ -48,7 +48,7 @@ func (n *Node) Weight() float64 {
 }
 
 func (n *Node) PickElapsed() time.Duration {
-	return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick))
+	return time.Duration(time.Now().UnixNano() - n.lastPick.Load())
 }
 
 func (n *Node) Raw() selector.Node {

selector/node/ewma/node.go

@@ -28,16 +28,16 @@ type Node struct {
 	selector.Node
 
 	// client statistic data
-	lag       int64
-	success   uint64
-	inflight  int64
-	inflights [200]int64
+	lag       atomic.Int64
+	success   atomic.Uint64
+	inflight  atomic.Int64
+	inflights [200]atomic.Int64
 	// last collected timestamp
-	stamp int64
+	stamp atomic.Int64
 	// request number in a period time
-	reqs int64
+	reqs atomic.Int64
 	// last lastPick timestamp
-	lastPick int64
+	lastPick atomic.Int64
 
 	errHandler   func(err error) (isErr bool)
 	cachedWeight *atomic.Value
@@ -57,27 +57,27 @@ type Builder struct {
 func (b *Builder) Build(n selector.Node) selector.WeightedNode {
 	s := &Node{
 		Node:         n,
-		lag:          0,
-		success:      1000,
-		inflight:     1,
+		inflights:    [200]atomic.Int64{},
 		errHandler:   b.ErrHandler,
 		cachedWeight: &atomic.Value{},
 	}
+	s.success.Store(1000)
+	s.inflight.Store(1)
 	return s
 }
 
 func (n *Node) health() uint64 {
-	return atomic.LoadUint64(&n.success)
+	return n.success.Load()
 }
 
 func (n *Node) load() (load uint64) {
 	now := time.Now().UnixNano()
-	avgLag := atomic.LoadInt64(&n.lag)
+	avgLag := n.lag.Load()
 	predict := n.predict(avgLag, now)
 
 	if avgLag == 0 {
 		// penalty is the penalty value when there is no data when the node is just started.
-		load = penalty * uint64(atomic.LoadInt64(&n.inflight))
+		load = penalty * uint64(n.inflight.Load())
 		return
 	}
 	if predict > avgLag {
@@ -86,7 +86,7 @@ func (n *Node) load() (load uint64) {
 	// add 5ms to eliminate the latency gap between different zones
 	avgLag += int64(time.Millisecond * 5)
 	avgLag = int64(math.Sqrt(float64(avgLag)))
-	load = uint64(avgLag) * uint64(atomic.LoadInt64(&n.inflight))
+	load = uint64(avgLag) * uint64(n.inflight.Load())
 	return load
 }
 
@@ -97,7 +97,7 @@ func (n *Node) predict(avgLag int64, now int64) (predict int64) {
 		totalNum int
 	)
 	for i := range n.inflights {
-		start := atomic.LoadInt64(&n.inflights[i])
+		start := n.inflights[i].Load()
 		if start != 0 {
 			totalNum++
 			lag := now - start
@@ -116,20 +116,20 @@ func (n *Node) predict(avgLag int64, now int64) (predict int64) {
 // Pick pick a node.
 func (n *Node) Pick() selector.DoneFunc {
 	start := time.Now().UnixNano()
-	atomic.StoreInt64(&n.lastPick, start)
-	atomic.AddInt64(&n.inflight, 1)
-	reqs := atomic.AddInt64(&n.reqs, 1)
+	n.lastPick.Store(start)
+	n.inflight.Add(1)
+	reqs := n.reqs.Add(1)
 	slot := reqs % 200
-	swapped := atomic.CompareAndSwapInt64(&n.inflights[slot], 0, start)
+	swapped := n.inflights[slot].CompareAndSwap(0, start)
 	return func(_ context.Context, di selector.DoneInfo) {
 		if swapped {
-			atomic.CompareAndSwapInt64(&n.inflights[slot], start, 0)
+			n.inflights[slot].CompareAndSwap(start, 0)
 		}
-		atomic.AddInt64(&n.inflight, -1)
+		n.inflight.Add(-1)
 
 		now := time.Now().UnixNano()
 		// get moving average ratio w
-		stamp := atomic.SwapInt64(&n.stamp, now)
+		stamp := n.stamp.Swap(now)
 		td := now - stamp
 		if td < 0 {
 			td = 0
@@ -140,12 +140,12 @@ func (n *Node) Pick() selector.DoneFunc {
 		if lag < 0 {
 			lag = 0
 		}
-		oldLag := atomic.LoadInt64(&n.lag)
+		oldLag := n.lag.Load()
 		if oldLag == 0 {
 			w = 0.0
 		}
 		lag = int64(float64(oldLag)*w + float64(lag)*(1.0-w))
-		atomic.StoreInt64(&n.lag, lag)
+		n.lag.Store(lag)
 
 		success := uint64(1000) // error value ,if error set 1
 		if di.Err != nil {
@@ -158,9 +158,9 @@ func (n *Node) Pick() selector.DoneFunc {
 				success = 0
 			}
 		}
-		oldSuc := atomic.LoadUint64(&n.success)
+		oldSuc := n.success.Load()
 		success = uint64(float64(oldSuc)*w + float64(success)*(1.0-w))
-		atomic.StoreUint64(&n.success, success)
+		n.success.Store(success)
 	}
 }
 
@@ -183,7 +183,7 @@ func (n *Node) Weight() (weight float64) {
 }
 
 func (n *Node) PickElapsed() time.Duration {
-	return time.Duration(time.Now().UnixNano() - atomic.LoadInt64(&n.lastPick))
+	return time.Duration(time.Now().UnixNano() - n.lastPick.Load())
 }
 
 func (n *Node) Raw() selector.Node {

selector/p2c/p2c.go

@@ -34,7 +34,7 @@ func New(opts ...Option) selector.Selector {
 type Balancer struct {
 	mu     sync.Mutex
 	r      *rand.Rand
-	picked int64
+	picked atomic.Bool
 }
 
 // choose two distinct nodes.
@@ -71,9 +71,9 @@ func (s *Balancer) Pick(_ context.Context, nodes []selector.WeightedNode) (selec
 
 	// If the failed node has never been selected once during forceGap, it is forced to be selected once
 	// Take advantage of forced opportunities to trigger updates of success rate and delay
-	if upc.PickElapsed() > forcePick && atomic.CompareAndSwapInt64(&s.picked, 0, 1) {
+	if upc.PickElapsed() > forcePick && s.picked.CompareAndSwap(false, true) {
+		defer s.picked.Store(false)
 		pc = upc
-		atomic.StoreInt64(&s.picked, 0)
 	}
 	done := pc.Pick()
 	return pc, done, nil