80 lines
1.7 KiB
Go
80 lines
1.7 KiB
Go
package selector
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
mdutil "github.com/go-gost/core/metadata/util"
|
|
"github.com/go-gost/core/selector"
|
|
)
|
|
|
|
type roundRobinStrategy[T selector.Selectable] struct {
|
|
counter uint64
|
|
}
|
|
|
|
// RoundRobinStrategy is a strategy for node selector.
|
|
// The node will be selected by round-robin algorithm.
|
|
func RoundRobinStrategy[T selector.Selectable]() selector.Strategy[T] {
|
|
return &roundRobinStrategy[T]{}
|
|
}
|
|
|
|
func (s *roundRobinStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
|
|
if len(vs) == 0 {
|
|
return
|
|
}
|
|
|
|
n := atomic.AddUint64(&s.counter, 1) - 1
|
|
return vs[int(n%uint64(len(vs)))]
|
|
}
|
|
|
|
type randomStrategy[T selector.Selectable] struct {
|
|
rw *randomWeighted[T]
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// RandomStrategy is a strategy for node selector.
|
|
// The node will be selected randomly.
|
|
func RandomStrategy[T selector.Selectable]() selector.Strategy[T] {
|
|
return &randomStrategy[T]{
|
|
rw: newRandomWeighted[T](),
|
|
}
|
|
}
|
|
|
|
func (s *randomStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
|
|
if len(vs) == 0 {
|
|
return
|
|
}
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
s.rw.Reset()
|
|
for i := range vs {
|
|
weight := mdutil.GetInt(vs[i].Metadata(), labelWeight)
|
|
if weight <= 0 {
|
|
weight = 1
|
|
}
|
|
s.rw.Add(vs[i], weight)
|
|
}
|
|
|
|
return s.rw.Next()
|
|
}
|
|
|
|
type fifoStrategy[T selector.Selectable] struct{}
|
|
|
|
// FIFOStrategy is a strategy for node selector.
|
|
// The node will be selected from first to last,
|
|
// and will stick to the selected node until it is failed.
|
|
func FIFOStrategy[T selector.Selectable]() selector.Strategy[T] {
|
|
return &fifoStrategy[T]{}
|
|
}
|
|
|
|
// Apply applies the fifo strategy for the nodes.
|
|
func (s *fifoStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
|
|
if len(vs) == 0 {
|
|
return
|
|
}
|
|
return vs[0]
|
|
}
|