github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/selector/strategy.go (about)

     1  package selector
     2  
     3  import (
     4  	"math/rand"
     5  	"strings"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/volts-dev/volts/registry"
    10  )
    11  
    12  var strategyMap = map[string]func(services []*registry.Service) Next{
    13  	"random":     Random,
    14  	"roundrobin": RoundRobin,
    15  }
    16  
    17  func init() {
    18  	rand.Seed(time.Now().UnixNano())
    19  }
    20  
    21  func Register(name string, fn func(services []*registry.Service) Next) {
    22  	strategyMap[strings.ToLower(name)] = fn
    23  }
    24  
    25  func Use(name string) func(services []*registry.Service) Next {
    26  	if fn, has := strategyMap[strings.ToLower(name)]; has {
    27  		return fn
    28  	}
    29  	return nil
    30  }
    31  
    32  // Random is a random strategy algorithm for node selection
    33  func Random(services []*registry.Service) Next {
    34  	nodes := make([]*registry.Node, 0, len(services))
    35  
    36  	for _, service := range services {
    37  		nodes = append(nodes, service.Nodes...)
    38  	}
    39  
    40  	return func() (*registry.Node, error) {
    41  		if len(nodes) == 0 {
    42  			return nil, ErrNoneAvailable
    43  		}
    44  
    45  		i := rand.Int() % len(nodes)
    46  		return nodes[i], nil
    47  	}
    48  }
    49  
    50  // RoundRobin is a roundrobin strategy algorithm for node selection
    51  func RoundRobin(services []*registry.Service) Next {
    52  	nodes := make([]*registry.Node, 0, len(services))
    53  
    54  	for _, service := range services {
    55  		nodes = append(nodes, service.Nodes...)
    56  	}
    57  
    58  	var i = rand.Int()
    59  	var mtx sync.Mutex
    60  
    61  	return func() (*registry.Node, error) {
    62  		if len(nodes) == 0 {
    63  			return nil, ErrNoneAvailable
    64  		}
    65  
    66  		mtx.Lock()
    67  		node := nodes[i%len(nodes)]
    68  		i++
    69  		mtx.Unlock()
    70  
    71  		return node, nil
    72  	}
    73  }