github.com/xtls/xray-core@v1.8.12-0.20240518155711-3168d27b0bdb/common/protocol/server_picker.go (about)

     1  package protocol
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type ServerList struct {
     8  	sync.RWMutex
     9  	servers []*ServerSpec
    10  }
    11  
    12  func NewServerList() *ServerList {
    13  	return &ServerList{}
    14  }
    15  
    16  func (sl *ServerList) AddServer(server *ServerSpec) {
    17  	sl.Lock()
    18  	defer sl.Unlock()
    19  
    20  	sl.servers = append(sl.servers, server)
    21  }
    22  
    23  func (sl *ServerList) Size() uint32 {
    24  	sl.RLock()
    25  	defer sl.RUnlock()
    26  
    27  	return uint32(len(sl.servers))
    28  }
    29  
    30  func (sl *ServerList) GetServer(idx uint32) *ServerSpec {
    31  	sl.Lock()
    32  	defer sl.Unlock()
    33  
    34  	for {
    35  		if idx >= uint32(len(sl.servers)) {
    36  			return nil
    37  		}
    38  
    39  		server := sl.servers[idx]
    40  		if !server.IsValid() {
    41  			sl.removeServer(idx)
    42  			continue
    43  		}
    44  
    45  		return server
    46  	}
    47  }
    48  
    49  func (sl *ServerList) removeServer(idx uint32) {
    50  	n := len(sl.servers)
    51  	sl.servers[idx] = sl.servers[n-1]
    52  	sl.servers = sl.servers[:n-1]
    53  }
    54  
    55  type ServerPicker interface {
    56  	PickServer() *ServerSpec
    57  }
    58  
    59  type RoundRobinServerPicker struct {
    60  	sync.Mutex
    61  	serverlist *ServerList
    62  	nextIndex  uint32
    63  }
    64  
    65  func NewRoundRobinServerPicker(serverlist *ServerList) *RoundRobinServerPicker {
    66  	return &RoundRobinServerPicker{
    67  		serverlist: serverlist,
    68  		nextIndex:  0,
    69  	}
    70  }
    71  
    72  func (p *RoundRobinServerPicker) PickServer() *ServerSpec {
    73  	p.Lock()
    74  	defer p.Unlock()
    75  
    76  	next := p.nextIndex
    77  	server := p.serverlist.GetServer(next)
    78  	if server == nil {
    79  		next = 0
    80  		server = p.serverlist.GetServer(0)
    81  	}
    82  	next++
    83  	if next >= p.serverlist.Size() {
    84  		next = 0
    85  	}
    86  	p.nextIndex = next
    87  
    88  	return server
    89  }