github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/balancer/roundrobin.go (about)

     1  package balancer
     2  
     3  import (
     4  	"reflect"
     5  	"sync/atomic"
     6  )
     7  
     8  type RoundRobin struct {
     9  	current uint64
    10  }
    11  
    12  // NextPeer 返回轮盘的下一节点
    13  func (r *RoundRobin) NextPeer(nodes interface{}) (error, interface{}) {
    14  	v := reflect.ValueOf(nodes)
    15  	if v.Kind() != reflect.Slice {
    16  		return ErrInputNotSlice, nil
    17  	}
    18  
    19  	if v.Len() == 0 {
    20  		return ErrNoAvaliableNode, nil
    21  	}
    22  
    23  	next := r.NextIndex(v.Len())
    24  	return nil, v.Index(next).Interface()
    25  }
    26  
    27  // NextIndex 返回下一个节点下标
    28  func (r *RoundRobin) NextIndex(total int) int {
    29  	return int(atomic.AddUint64(&r.current, uint64(1)) % uint64(total))
    30  }