add weight for selector

selector/filter.go

@@ -1,6 +1,7 @@
 package selector
 
 import (
+	"context"
 	"time"
 
 	mdutil "github.com/go-gost/core/metadata/util"
@@ -22,7 +23,7 @@ func FailFilter[T selector.Selectable](maxFails int, timeout time.Duration) sele
 }
 
 // Filter filters dead objects.
-func (f *failFilter[T]) Filter(vs ...T) []T {
+func (f *failFilter[T]) Filter(ctx context.Context, vs ...T) []T {
 	if len(vs) <= 1 {
 		return vs
 	}
@@ -66,7 +67,7 @@ func BackupFilter[T selector.Selectable]() selector.Filter[T] {
 }
 
 // Filter filters backup objects.
-func (f *backupFilter[T]) Filter(vs ...T) []T {
+func (f *backupFilter[T]) Filter(ctx context.Context, vs ...T) []T {
 	if len(vs) <= 1 {
 		return vs
 	}

selector/selector.go

@@ -1,9 +1,9 @@
 package selector
 
 import (
+	"context"
 	"time"
 
-	"github.com/go-gost/core/chain"
 	"github.com/go-gost/core/selector"
 )
 
@@ -20,17 +20,6 @@ const (
 	labelFailTimeout = "failTimeout"
 )
 
-var (
-	DefaultNodeSelector = NewSelector(
-		RoundRobinStrategy[*chain.Node](),
-		// FailFilter[*Node](1, DefaultFailTimeout),
-	)
-	DefaultChainSelector = NewSelector(
-		RoundRobinStrategy[chain.SelectableChainer](),
-		// FailFilter[SelectableChainer](1, DefaultFailTimeout),
-	)
-)
-
 type defaultSelector[T selector.Selectable] struct {
 	strategy selector.Strategy[T]
 	filters  []selector.Filter[T]
@@ -43,12 +32,12 @@ func NewSelector[T selector.Selectable](strategy selector.Strategy[T], filters .
 	}
 }
 
-func (s *defaultSelector[T]) Select(vs ...T) (v T) {
+func (s *defaultSelector[T]) Select(ctx context.Context, vs ...T) (v T) {
 	for _, filter := range s.filters {
-		vs = filter.Filter(vs...)
+		vs = filter.Filter(ctx, vs...)
 	}
 	if len(vs) == 0 {
 		return
 	}
-	return s.strategy.Apply(vs...)
+	return s.strategy.Apply(ctx, vs...)
 }

selector/strategy.go

@@ -1,11 +1,11 @@
 package selector
 
 import (
-	"math/rand"
+	"context"
 	"sync"
 	"sync/atomic"
-	"time"
 
+	mdutil "github.com/go-gost/core/metadata/util"
 	"github.com/go-gost/core/selector"
 )
 
@@ -19,7 +19,7 @@ func RoundRobinStrategy[T selector.Selectable]() selector.Strategy[T] {
 	return &roundRobinStrategy[T]{}
 }
 
-func (s *roundRobinStrategy[T]) Apply(vs ...T) (v T) {
+func (s *roundRobinStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
 	if len(vs) == 0 {
 		return
 	}
@@ -29,29 +29,36 @@ func (s *roundRobinStrategy[T]) Apply(vs ...T) (v T) {
 }
 
 type randomStrategy[T selector.Selectable] struct {
-	rand *rand.Rand
-	mux  sync.Mutex
+	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]{
-		rand: rand.New(rand.NewSource(time.Now().UnixNano())),
+		rw: newRandomWeighted[T](),
 	}
 }
 
-func (s *randomStrategy[T]) Apply(vs ...T) (v T) {
+func (s *randomStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
 	if len(vs) == 0 {
 		return
 	}
 
-	s.mux.Lock()
-	defer s.mux.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
-	r := s.rand.Int()
+	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 vs[r%len(vs)]
+	return s.rw.Next()
 }
 
 type fifoStrategy[T selector.Selectable] struct{}
@@ -64,7 +71,7 @@ func FIFOStrategy[T selector.Selectable]() selector.Strategy[T] {
 }
 
 // Apply applies the fifo strategy for the nodes.
-func (s *fifoStrategy[T]) Apply(vs ...T) (v T) {
+func (s *fifoStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
 	if len(vs) == 0 {
 		return
 	}

selector/weighted.go

@@ -0,0 +1,54 @@
+package selector
+
+import (
+	"math/rand"
+	"time"
+
+	"github.com/go-gost/core/selector"
+)
+
+type randomWeightedItem[T selector.Selectable] struct {
+	item   T
+	weight int
+}
+
+type randomWeighted[T selector.Selectable] struct {
+	items []*randomWeightedItem[T]
+	sum   int
+	r     *rand.Rand
+}
+
+func newRandomWeighted[T selector.Selectable]() *randomWeighted[T] {
+	return &randomWeighted[T]{
+		r: rand.New(rand.NewSource(time.Now().UnixNano())),
+	}
+}
+
+func (rw *randomWeighted[T]) Add(item T, weight int) {
+	ri := &randomWeightedItem[T]{item: item, weight: weight}
+	rw.items = append(rw.items, ri)
+	rw.sum += weight
+}
+
+func (rw *randomWeighted[T]) Next() (v T) {
+	if len(rw.items) == 0 {
+		return
+	}
+	if rw.sum <= 0 {
+		return
+	}
+	weight := rw.r.Intn(rw.sum) + 1
+	for _, item := range rw.items {
+		weight -= item.weight
+		if weight <= 0 {
+			return item.item
+		}
+	}
+
+	return rw.items[len(rw.items)-1].item
+}
+
+func (rw *randomWeighted[T]) Reset() {
+	rw.items = nil
+	rw.sum = 0
+}