fix(selector): safe type assertion, clean up debug logging, add doc comments and 64 tests
Replace bare type assertion in ParallelStrategy with comma-ok check to prevent panic on non-Node inputs. Downgrade noisy Infof hash-selection log to Tracef and remove Chinese-language debug output. Add doc comments to all exported symbols in weighted.go.
This commit is contained in:
ginuerzh
@@ -0,0 +1,266 @@
|
||||
package selector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/metadata"
|
||||
"github.com/go-gost/core/selector"
|
||||
xmd "github.com/go-gost/x/metadata"
|
||||
mdutil "github.com/go-gost/x/metadata/util"
|
||||
)
|
||||
|
||||
// testMarkable is a test item that implements both Metadatable and Markable.
|
||||
type testMarkable struct {
|
||||
md metadata.Metadata
|
||||
marker selector.Marker
|
||||
}
|
||||
|
||||
func newTestMarkable(md map[string]any) *testMarkable {
|
||||
return &testMarkable{
|
||||
md: xmd.NewMetadata(md),
|
||||
marker: selector.NewFailMarker(),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *testMarkable) Metadata() metadata.Metadata {
|
||||
return t.md
|
||||
}
|
||||
|
||||
func (t *testMarkable) Marker() selector.Marker {
|
||||
return t.marker
|
||||
}
|
||||
|
||||
// testMetadatable only implements Metadatable, not Markable.
|
||||
type testMetadatable struct {
|
||||
md metadata.Metadata
|
||||
}
|
||||
|
||||
func newTestMetadatable(md map[string]any) *testMetadatable {
|
||||
return &testMetadatable{md: xmd.NewMetadata(md)}
|
||||
}
|
||||
|
||||
func (t *testMetadatable) Metadata() metadata.Metadata {
|
||||
return t.md
|
||||
}
|
||||
|
||||
// --- FailFilter tests ---
|
||||
|
||||
func TestFailFilter_Empty(t *testing.T) {
|
||||
f := FailFilter[int](1, time.Second)
|
||||
if vs := f.Filter(context.Background()); len(vs) != 0 {
|
||||
t.Fatalf("expected empty, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_Single(t *testing.T) {
|
||||
f := FailFilter[int](1, time.Second)
|
||||
vs := f.Filter(context.Background(), 42)
|
||||
if len(vs) != 1 || vs[0] != 42 {
|
||||
t.Fatalf("expected [42], got %v", vs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_HealthyItems(t *testing.T) {
|
||||
f := FailFilter[*testMarkable](2, 10*time.Second)
|
||||
|
||||
items := []*testMarkable{
|
||||
newTestMarkable(nil),
|
||||
newTestMarkable(nil),
|
||||
}
|
||||
vs := f.Filter(context.Background(), items[0], items[1])
|
||||
if len(vs) != 2 {
|
||||
t.Fatalf("expected 2 items, got %d", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_DeadItem(t *testing.T) {
|
||||
f := FailFilter[*testMarkable](1, 10*time.Second)
|
||||
|
||||
alive := newTestMarkable(nil)
|
||||
dead := newTestMarkable(nil)
|
||||
dead.Marker().Mark() // count=1, within failTimeout
|
||||
|
||||
vs := f.Filter(context.Background(), alive, dead)
|
||||
if len(vs) != 1 || vs[0] != alive {
|
||||
t.Fatalf("expected only alive item, got %v", vs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_ExpiredFailTimeout(t *testing.T) {
|
||||
f := FailFilter[*testMarkable](1, 50*time.Millisecond)
|
||||
|
||||
item := newTestMarkable(nil)
|
||||
item.Marker().Mark()
|
||||
|
||||
// Wait for fail timeout to expire
|
||||
time.Sleep(80 * time.Millisecond)
|
||||
|
||||
vs := f.Filter(context.Background(), item)
|
||||
if len(vs) != 1 {
|
||||
t.Fatalf("expected item to be included after timeout, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_PerItemOverrides(t *testing.T) {
|
||||
f := FailFilter[*testMarkable](1, 10*time.Second)
|
||||
|
||||
// Item with custom maxFails=5
|
||||
item := newTestMarkable(map[string]any{
|
||||
"maxFails": 5,
|
||||
"failTimeout": "1h",
|
||||
})
|
||||
item.Marker().Mark() // count=1, but maxFails=5 so still alive
|
||||
|
||||
vs := f.Filter(context.Background(), item)
|
||||
if len(vs) != 1 {
|
||||
t.Fatalf("expected item with high maxFails to pass, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_NilMarker(t *testing.T) {
|
||||
f := FailFilter[*testMetadatable](1, time.Second)
|
||||
|
||||
item := newTestMetadatable(nil)
|
||||
vs := f.Filter(context.Background(), item)
|
||||
if len(vs) != 1 {
|
||||
t.Fatalf("expected non-markable item to pass, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_AllDead(t *testing.T) {
|
||||
f := FailFilter[*testMarkable](1, 10*time.Second)
|
||||
|
||||
dead1 := newTestMarkable(nil)
|
||||
dead1.Marker().Mark()
|
||||
dead2 := newTestMarkable(nil)
|
||||
dead2.Marker().Mark()
|
||||
|
||||
vs := f.Filter(context.Background(), dead1, dead2)
|
||||
if len(vs) != 0 {
|
||||
t.Fatalf("expected all dead items filtered, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_ZeroMaxFails(t *testing.T) {
|
||||
// Zero maxFails should default to 1, so a single mark should filter the item
|
||||
f := FailFilter[*testMarkable](0, 10*time.Second)
|
||||
|
||||
item := newTestMarkable(nil)
|
||||
item.Marker().Mark() // count=1, which >= maxFails(defaulted to 1)
|
||||
|
||||
vs := f.Filter(context.Background(), item, newTestMarkable(nil))
|
||||
if len(vs) != 1 {
|
||||
t.Fatalf("expected zero maxFails to default to 1, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailFilter_ZeroFailTimeout(t *testing.T) {
|
||||
// Zero failTimeout should default to DefaultFailTimeout (10s)
|
||||
f := FailFilter[*testMarkable](1, 0)
|
||||
|
||||
item := newTestMarkable(nil)
|
||||
item.Marker().Mark() // count=1, within DefaultFailTimeout → dead
|
||||
|
||||
vs := f.Filter(context.Background(), item, newTestMarkable(nil))
|
||||
if len(vs) != 1 {
|
||||
t.Fatalf("expected zero failTimeout to use default, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
// --- BackupFilter tests ---
|
||||
|
||||
func TestBackupFilter_Empty(t *testing.T) {
|
||||
f := BackupFilter[int]()
|
||||
if vs := f.Filter(context.Background()); len(vs) != 0 {
|
||||
t.Fatalf("expected empty, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupFilter_Single(t *testing.T) {
|
||||
f := BackupFilter[int]()
|
||||
vs := f.Filter(context.Background(), 42)
|
||||
if len(vs) != 1 || vs[0] != 42 {
|
||||
t.Fatalf("expected [42], got %v", vs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupFilter_PrimaryPreferred(t *testing.T) {
|
||||
f := BackupFilter[*testMetadatable]()
|
||||
|
||||
primary := newTestMetadatable(nil)
|
||||
backup := newTestMetadatable(map[string]any{"backup": true})
|
||||
|
||||
vs := f.Filter(context.Background(), primary, backup)
|
||||
if len(vs) != 1 || vs[0] != primary {
|
||||
t.Fatalf("expected only primary, got %v", vs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupFilter_AllBackup(t *testing.T) {
|
||||
f := BackupFilter[*testMetadatable]()
|
||||
|
||||
b1 := newTestMetadatable(map[string]any{"backup": true})
|
||||
b2 := newTestMetadatable(map[string]any{"backup": true})
|
||||
|
||||
vs := f.Filter(context.Background(), b1, b2)
|
||||
if len(vs) != 2 {
|
||||
t.Fatalf("expected all backups returned when no primaries, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupFilter_NoBackup(t *testing.T) {
|
||||
f := BackupFilter[*testMetadatable]()
|
||||
|
||||
p1 := newTestMetadatable(nil)
|
||||
p2 := newTestMetadatable(nil)
|
||||
|
||||
vs := f.Filter(context.Background(), p1, p2)
|
||||
if len(vs) != 2 {
|
||||
t.Fatalf("expected all primaries, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupFilter_BackupFalse(t *testing.T) {
|
||||
f := BackupFilter[*testMetadatable]()
|
||||
|
||||
item := newTestMetadatable(map[string]any{"backup": false})
|
||||
vs := f.Filter(context.Background(), item)
|
||||
if len(vs) != 1 {
|
||||
t.Fatalf("expected item with backup=false to be treated as primary, got %d items", len(vs))
|
||||
}
|
||||
}
|
||||
|
||||
// Verify the constants are accessible
|
||||
func TestFilterConstants(t *testing.T) {
|
||||
if DefaultMaxFails != 1 {
|
||||
t.Fatalf("DefaultMaxFails expected 1, got %d", DefaultMaxFails)
|
||||
}
|
||||
if DefaultFailTimeout != 10*time.Second {
|
||||
t.Fatalf("DefaultFailTimeout expected 10s, got %v", DefaultFailTimeout)
|
||||
}
|
||||
}
|
||||
|
||||
// Verify metadata label constants via mdutil
|
||||
func TestFilterMetadataLabels(t *testing.T) {
|
||||
md := xmd.NewMetadata(map[string]any{
|
||||
"weight": 5,
|
||||
"backup": true,
|
||||
"maxFails": 3,
|
||||
"failTimeout": "30s",
|
||||
})
|
||||
|
||||
if w := mdutil.GetInt(md, labelWeight); w != 5 {
|
||||
t.Fatalf("weight: expected 5, got %d", w)
|
||||
}
|
||||
if !mdutil.GetBool(md, labelBackup) {
|
||||
t.Fatal("backup: expected true")
|
||||
}
|
||||
if mf := mdutil.GetInt(md, labelMaxFails); mf != 3 {
|
||||
t.Fatalf("maxFails: expected 3, got %d", mf)
|
||||
}
|
||||
if ft := mdutil.GetDuration(md, labelFailTimeout); ft != 30*time.Second {
|
||||
t.Fatalf("failTimeout: expected 30s, got %v", ft)
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,8 @@ import (
|
||||
|
||||
type parallelStrategy[T any] struct{}
|
||||
|
||||
// ParallelStrategy is a strategy for node selector.
|
||||
// It dials all nodes concurrently and uses the first successful connection.
|
||||
func ParallelStrategy[T any]() selector.Strategy[T] {
|
||||
return ¶llelStrategy[T]{}
|
||||
}
|
||||
@@ -29,7 +31,11 @@ func (s *parallelStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
|
||||
|
||||
nodes := make([]*chain.Node, 0, len(vs))
|
||||
for _, node := range vs {
|
||||
nodes = append(nodes, any(node).(*chain.Node))
|
||||
n, ok := any(node).(*chain.Node)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
nodes = append(nodes, n)
|
||||
}
|
||||
|
||||
vn := chain.NewNode("parallel", "", chain.TransportNodeOption(¶llelTransporter{nodes: nodes}))
|
||||
@@ -160,6 +166,7 @@ func (tr *parallelTransporter) Options() *chain.TransportOptions {
|
||||
return &chain.TransportOptions{}
|
||||
}
|
||||
|
||||
// Copy implements chain.Transporter.Copy.
|
||||
func (tr *parallelTransporter) Copy() chain.Transporter {
|
||||
return ¶llelTransporter{
|
||||
nodes: append([]*chain.Node(nil), tr.nodes...),
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
package selector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"testing"
|
||||
|
||||
"github.com/go-gost/core/chain"
|
||||
"github.com/go-gost/core/connector"
|
||||
)
|
||||
|
||||
// mockTransport implements chain.Transporter for testing.
|
||||
type mockTransport struct {
|
||||
dialResult net.Conn
|
||||
dialErr error
|
||||
handshakeErr error
|
||||
connectErr error
|
||||
bindErr error
|
||||
}
|
||||
|
||||
func (m *mockTransport) Dial(ctx context.Context, addr string) (net.Conn, error) {
|
||||
if m.dialErr != nil {
|
||||
return nil, m.dialErr
|
||||
}
|
||||
return m.dialResult, nil
|
||||
}
|
||||
|
||||
func (m *mockTransport) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
|
||||
if m.handshakeErr != nil {
|
||||
return nil, m.handshakeErr
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (m *mockTransport) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
|
||||
if m.connectErr != nil {
|
||||
return nil, m.connectErr
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (m *mockTransport) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
|
||||
if m.bindErr != nil {
|
||||
return nil, m.bindErr
|
||||
}
|
||||
return net.Listen("tcp", "127.0.0.1:0")
|
||||
}
|
||||
|
||||
func (m *mockTransport) Multiplex() bool { return false }
|
||||
func (m *mockTransport) Options() *chain.TransportOptions { return &chain.TransportOptions{} }
|
||||
func (m *mockTransport) Copy() chain.Transporter { return m }
|
||||
|
||||
func makeNode(tr chain.Transporter) *chain.Node {
|
||||
return chain.NewNode("test", "127.0.0.1:0", chain.TransportNodeOption(tr))
|
||||
}
|
||||
|
||||
// --- ParallelStrategy ---
|
||||
|
||||
func TestParallelStrategy_Empty(t *testing.T) {
|
||||
s := ParallelStrategy[*chain.Node]()
|
||||
if v := s.Apply(context.Background()); v != nil {
|
||||
t.Fatalf("expected nil, got %v", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelStrategy_Single(t *testing.T) {
|
||||
s := ParallelStrategy[*chain.Node]()
|
||||
tr := &mockTransport{}
|
||||
node := makeNode(tr)
|
||||
|
||||
result := s.Apply(context.Background(), node)
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result for single node")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelStrategy_MultipleNodes(t *testing.T) {
|
||||
s := ParallelStrategy[*chain.Node]()
|
||||
nodes := []*chain.Node{
|
||||
makeNode(&mockTransport{}),
|
||||
makeNode(&mockTransport{}),
|
||||
makeNode(&mockTransport{}),
|
||||
}
|
||||
|
||||
result := s.Apply(context.Background(), nodes[0], nodes[1], nodes[2])
|
||||
if result == nil {
|
||||
t.Fatal("expected non-nil result")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelStrategy_NonNodeInput(t *testing.T) {
|
||||
s := ParallelStrategy[string]()
|
||||
// Single non-node input: len==1 returns vs[0] directly
|
||||
if v := s.Apply(context.Background(), "hello"); v != "hello" {
|
||||
t.Fatalf("expected 'hello' for single input, got %q", v)
|
||||
}
|
||||
|
||||
// Multiple non-node inputs: type assertion fails, returns zero value
|
||||
if v := s.Apply(context.Background(), "a", "b"); v != "" {
|
||||
t.Fatalf("expected zero value for multiple non-node inputs, got %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
// --- parallelTransporter unit tests ---
|
||||
|
||||
func TestParallelTransporter_Dial_NoNodes(t *testing.T) {
|
||||
tr := ¶llelTransporter{nodes: nil}
|
||||
_, err := tr.Dial(context.Background(), "")
|
||||
if err == nil {
|
||||
t.Fatal("expected error with no nodes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelTransporter_Copy(t *testing.T) {
|
||||
nodes := []*chain.Node{makeNode(&mockTransport{}), makeNode(&mockTransport{})}
|
||||
tr := ¶llelTransporter{nodes: nodes}
|
||||
|
||||
cp := tr.Copy()
|
||||
cpTr := cp.(*parallelTransporter)
|
||||
|
||||
if len(cpTr.nodes) != len(tr.nodes) {
|
||||
t.Fatalf("copy should have same number of nodes: %d vs %d", len(cpTr.nodes), len(tr.nodes))
|
||||
}
|
||||
|
||||
// Should be a shallow copy of the slice (different backing array)
|
||||
cpTr.nodes[0] = nil
|
||||
if tr.nodes[0] == nil {
|
||||
t.Fatal("copy should not share backing array")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelTransporter_Multiplex(t *testing.T) {
|
||||
tr := ¶llelTransporter{}
|
||||
if tr.Multiplex() {
|
||||
t.Fatal("parallel transporter should not multiplex")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelTransporter_Options(t *testing.T) {
|
||||
tr := ¶llelTransporter{}
|
||||
if opts := tr.Options(); opts == nil {
|
||||
t.Fatal("expected non-nil options")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelTransporter_Handshake_NonParallelConn(t *testing.T) {
|
||||
tr := ¶llelTransporter{}
|
||||
c1, c2 := net.Pipe()
|
||||
defer c1.Close()
|
||||
defer c2.Close()
|
||||
|
||||
conn, err := tr.Handshake(context.Background(), c1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
// Non-parallelConn should be returned as-is
|
||||
if conn != c1 {
|
||||
t.Fatal("expected original conn for non-parallelConn")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelTransporter_Connect_NonParallelConn(t *testing.T) {
|
||||
tr := ¶llelTransporter{}
|
||||
c1, _ := net.Pipe()
|
||||
defer c1.Close()
|
||||
|
||||
_, err := tr.Connect(context.Background(), c1, "tcp", "addr")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-parallelConn")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelTransporter_Bind_NonParallelConn(t *testing.T) {
|
||||
tr := ¶llelTransporter{}
|
||||
c1, _ := net.Pipe()
|
||||
defer c1.Close()
|
||||
|
||||
_, err := tr.Bind(context.Background(), c1, "tcp", "addr")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for non-parallelConn")
|
||||
}
|
||||
}
|
||||
|
||||
// --- parallelConn tests ---
|
||||
|
||||
func TestParallelConn_Handshake_Delegates(t *testing.T) {
|
||||
tr := ¶llelTransporter{
|
||||
nodes: []*chain.Node{makeNode(&mockTransport{})},
|
||||
}
|
||||
|
||||
c1, c2 := net.Pipe()
|
||||
defer c2.Close()
|
||||
|
||||
pc := ¶llelConn{Conn: c1, node: makeNode(&mockTransport{})}
|
||||
conn, err := tr.Handshake(context.Background(), pc)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if conn == nil {
|
||||
t.Fatal("expected non-nil conn after handshake")
|
||||
}
|
||||
conn.Close()
|
||||
}
|
||||
|
||||
func TestParallelConn_Connect_Delegates(t *testing.T) {
|
||||
tr := ¶llelTransporter{
|
||||
nodes: []*chain.Node{makeNode(&mockTransport{})},
|
||||
}
|
||||
|
||||
c1, c2 := net.Pipe()
|
||||
defer c2.Close()
|
||||
|
||||
pc := ¶llelConn{Conn: c1, node: makeNode(&mockTransport{})}
|
||||
conn, err := tr.Connect(context.Background(), pc, "tcp", "127.0.0.1:80")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if conn == nil {
|
||||
t.Fatal("expected non-nil conn after connect")
|
||||
}
|
||||
conn.Close()
|
||||
}
|
||||
|
||||
func TestParallelConn_Bind_Delegates(t *testing.T) {
|
||||
tr := ¶llelTransporter{
|
||||
nodes: []*chain.Node{makeNode(&mockTransport{})},
|
||||
}
|
||||
|
||||
c1, c2 := net.Pipe()
|
||||
defer c2.Close()
|
||||
|
||||
pc := ¶llelConn{Conn: c1, node: makeNode(&mockTransport{})}
|
||||
ln, err := tr.Bind(context.Background(), pc, "tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
ln.Close()
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
package selector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/go-gost/core/selector"
|
||||
)
|
||||
|
||||
func TestNewSelector(t *testing.T) {
|
||||
s := NewSelector[string](RoundRobinStrategy[string]())
|
||||
if s == nil {
|
||||
t.Fatal("expected non-nil selector")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultSelector_Select_Empty(t *testing.T) {
|
||||
s := NewSelector[int](RoundRobinStrategy[int]())
|
||||
if v := s.Select(context.Background()); v != 0 {
|
||||
t.Fatalf("expected zero value, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultSelector_Select_Single(t *testing.T) {
|
||||
s := NewSelector[int](RoundRobinStrategy[int]())
|
||||
if v := s.Select(context.Background(), 42); v != 42 {
|
||||
t.Fatalf("expected 42, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultSelector_Select_WithFilters(t *testing.T) {
|
||||
strategy := FIFOStrategy[int]()
|
||||
filter := selector.Filter[int](BackupFilter[int]())
|
||||
s := NewSelector[int](strategy, filter)
|
||||
|
||||
// Without backup flag, all items pass through; FIFO picks first
|
||||
v := s.Select(context.Background(), 1, 2, 3)
|
||||
if v != 1 {
|
||||
t.Fatalf("expected 1, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultSelector_FiltersAppliedBeforeStrategy(t *testing.T) {
|
||||
// Verify that filters narrow the list before strategy picks
|
||||
strategy := FIFOStrategy[string]()
|
||||
s := NewSelector[string](strategy, FailFilter[string](1, 0))
|
||||
|
||||
// All items are non-markable, so failFilter passes them all through
|
||||
v := s.Select(context.Background(), "a", "b", "c")
|
||||
if v != "a" {
|
||||
t.Fatalf("expected 'a', got %q", v)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package selector
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
corelogger "github.com/go-gost/core/logger"
|
||||
xlogger "github.com/go-gost/x/logger"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
corelogger.SetDefault(xlogger.Nop())
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
@@ -93,6 +93,9 @@ type hashStrategy[T any] struct {
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// HashStrategy is a strategy for node selector.
|
||||
// The node will be selected by hashing the client ID or hash source
|
||||
// from the context. Falls back to random selection when neither is available.
|
||||
func HashStrategy[T any]() selector.Strategy[T] {
|
||||
return &hashStrategy[T]{
|
||||
r: rand.New(rand.NewSource(time.Now().UnixNano())),
|
||||
@@ -104,11 +107,9 @@ func (s *hashStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
|
||||
return
|
||||
}
|
||||
|
||||
// 打印 基于用户 id 进行 hash 选择最终出口
|
||||
if clientID := xctx.ClientIDFromContext(ctx); clientID != "" {
|
||||
idStr := string(clientID)
|
||||
value := uint64(crc32.ChecksumIEEE([]byte(idStr)))
|
||||
logger.Default().Infof("【命中】从 xctx 拿到 ID: %s, Hash: %d", idStr, value)
|
||||
value := uint64(crc32.ChecksumIEEE([]byte(string(clientID))))
|
||||
logger.Default().Tracef("hash client %s %d", clientID, value)
|
||||
return vs[value%uint64(len(vs))]
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
package selector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"hash/crc32"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/go-gost/core/metadata"
|
||||
"github.com/go-gost/core/selector"
|
||||
xctx "github.com/go-gost/x/ctx"
|
||||
xmd "github.com/go-gost/x/metadata"
|
||||
)
|
||||
|
||||
// --- RoundRobinStrategy ---
|
||||
|
||||
func TestRoundRobinStrategy_Empty(t *testing.T) {
|
||||
s := RoundRobinStrategy[int]()
|
||||
if v := s.Apply(context.Background()); v != 0 {
|
||||
t.Fatalf("expected zero value, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinStrategy_Single(t *testing.T) {
|
||||
s := RoundRobinStrategy[int]()
|
||||
for i := 0; i < 5; i++ {
|
||||
if v := s.Apply(context.Background(), 42); v != 42 {
|
||||
t.Fatalf("iteration %d: expected 42, got %d", i, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinStrategy_Sequence(t *testing.T) {
|
||||
s := RoundRobinStrategy[int]()
|
||||
items := []int{10, 20, 30}
|
||||
|
||||
// Should cycle through 10, 20, 30, 10, 20, 30, ...
|
||||
for i := 0; i < 9; i++ {
|
||||
v := s.Apply(context.Background(), items...)
|
||||
expected := items[i%3]
|
||||
if v != expected {
|
||||
t.Fatalf("iteration %d: expected %d, got %d", i, expected, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinStrategy_Concurrent(t *testing.T) {
|
||||
s := RoundRobinStrategy[int]()
|
||||
items := []int{0, 1, 2, 3, 4}
|
||||
const goroutines = 100
|
||||
|
||||
var wg sync.WaitGroup
|
||||
results := make(chan int, goroutines)
|
||||
for i := 0; i < goroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
results <- s.Apply(context.Background(), items...)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(results)
|
||||
|
||||
// Every result should be a valid item
|
||||
for v := range results {
|
||||
if v < 0 || v > 4 {
|
||||
t.Fatalf("unexpected value %d", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- FIFOStrategy ---
|
||||
|
||||
func TestFIFOStrategy_Empty(t *testing.T) {
|
||||
s := FIFOStrategy[string]()
|
||||
if v := s.Apply(context.Background()); v != "" {
|
||||
t.Fatalf("expected zero value, got %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFIFOStrategy_AlwaysFirst(t *testing.T) {
|
||||
s := FIFOStrategy[int]()
|
||||
for i := 0; i < 10; i++ {
|
||||
v := s.Apply(context.Background(), 100, 200, 300)
|
||||
if v != 100 {
|
||||
t.Fatalf("iteration %d: expected 100, got %d", i, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- RandomStrategy ---
|
||||
|
||||
func TestRandomStrategy_Empty(t *testing.T) {
|
||||
s := RandomStrategy[int]()
|
||||
if v := s.Apply(context.Background()); v != 0 {
|
||||
t.Fatalf("expected zero value, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomStrategy_Single(t *testing.T) {
|
||||
s := RandomStrategy[int]()
|
||||
for i := 0; i < 10; i++ {
|
||||
if v := s.Apply(context.Background(), 7); v != 7 {
|
||||
t.Fatalf("expected 7, got %d", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomStrategy_Distribution(t *testing.T) {
|
||||
s := RandomStrategy[int]()
|
||||
items := []int{0, 1, 2}
|
||||
counts := make(map[int]int)
|
||||
const n = 3000
|
||||
|
||||
for i := 0; i < n; i++ {
|
||||
v := s.Apply(context.Background(), items...)
|
||||
counts[v]++
|
||||
}
|
||||
|
||||
// With uniform weights, each should get roughly n/3
|
||||
for _, item := range items {
|
||||
if counts[item] < n/len(items)/2 {
|
||||
t.Fatalf("item %d underrepresented: %d/%d", item, counts[item], n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// weightedItem implements metadata.Metadatable for weight testing.
|
||||
type weightedItem struct {
|
||||
md metadata.Metadata
|
||||
}
|
||||
|
||||
func (w *weightedItem) Metadata() metadata.Metadata { return w.md }
|
||||
|
||||
func TestRandomStrategy_Weighted(t *testing.T) {
|
||||
s := RandomStrategy[*weightedItem]()
|
||||
heavy := &weightedItem{md: xmd.NewMetadata(map[string]any{"weight": 100})}
|
||||
light := &weightedItem{md: xmd.NewMetadata(map[string]any{"weight": 1})}
|
||||
|
||||
heavyCount := 0
|
||||
const n = 2000
|
||||
for i := 0; i < n; i++ {
|
||||
v := s.Apply(context.Background(), heavy, light)
|
||||
if v == heavy {
|
||||
heavyCount++
|
||||
}
|
||||
}
|
||||
|
||||
// Heavy should win ~99% of the time with weight 100 vs 1
|
||||
if heavyCount < int(float64(n)*0.9) {
|
||||
t.Fatalf("heavy item underrepresented: %d/%d", heavyCount, n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomStrategy_ZeroWeight(t *testing.T) {
|
||||
s := RandomStrategy[int]()
|
||||
// Zero/negative weights default to 1, so all items should be selectable
|
||||
v := s.Apply(context.Background(), 1, 2, 3)
|
||||
if v < 1 || v > 3 {
|
||||
t.Fatalf("expected 1-3, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
// --- HashStrategy ---
|
||||
|
||||
func TestHashStrategy_Empty(t *testing.T) {
|
||||
s := HashStrategy[int]()
|
||||
if v := s.Apply(context.Background()); v != 0 {
|
||||
t.Fatalf("expected zero value, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHashStrategy_ClientID(t *testing.T) {
|
||||
s := HashStrategy[int]()
|
||||
items := []int{10, 20, 30, 40, 50}
|
||||
ctx := xctx.ContextWithClientID(context.Background(), "test-client")
|
||||
|
||||
// Same client ID should always map to the same item
|
||||
v1 := s.Apply(ctx, items...)
|
||||
v2 := s.Apply(ctx, items...)
|
||||
if v1 != v2 {
|
||||
t.Fatalf("hash strategy should be deterministic for same client ID: %d != %d", v1, v2)
|
||||
}
|
||||
|
||||
// Verify it matches manual CRC32 calculation
|
||||
expectedIdx := uint64(crc32.ChecksumIEEE([]byte("test-client"))) % uint64(len(items))
|
||||
if v1 != items[expectedIdx] {
|
||||
t.Fatalf("expected items[%d]=%d, got %d", expectedIdx, items[expectedIdx], v1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHashStrategy_HashSource(t *testing.T) {
|
||||
s := HashStrategy[int]()
|
||||
items := []int{10, 20, 30}
|
||||
ctx := xctx.ContextWithHash(context.Background(), &xctx.Hash{Source: "my-hash-key"})
|
||||
|
||||
v1 := s.Apply(ctx, items...)
|
||||
v2 := s.Apply(ctx, items...)
|
||||
if v1 != v2 {
|
||||
t.Fatalf("hash strategy should be deterministic for same hash source: %d != %d", v1, v2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHashStrategy_ClientIDPriorityOverHash(t *testing.T) {
|
||||
s := HashStrategy[int]()
|
||||
items := []int{10, 20, 30, 40}
|
||||
|
||||
ctxClient := xctx.ContextWithClientID(context.Background(), "clientA")
|
||||
ctxBoth := xctx.ContextWithHash(
|
||||
xctx.ContextWithClientID(context.Background(), "clientA"),
|
||||
&xctx.Hash{Source: "different-hash"},
|
||||
)
|
||||
|
||||
vClient := s.Apply(ctxClient, items...)
|
||||
vBoth := s.Apply(ctxBoth, items...)
|
||||
|
||||
// Client ID should take priority over hash source
|
||||
if vClient != vBoth {
|
||||
t.Fatalf("client ID should take priority: %d != %d", vClient, vBoth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHashStrategy_FallbackRandom(t *testing.T) {
|
||||
s := HashStrategy[int]()
|
||||
items := []int{0, 1, 2}
|
||||
counts := make(map[int]int)
|
||||
const n = 3000
|
||||
|
||||
for i := 0; i < n; i++ {
|
||||
v := s.Apply(context.Background(), items...)
|
||||
counts[v]++
|
||||
}
|
||||
|
||||
// Without client ID or hash, falls back to random
|
||||
for _, item := range items {
|
||||
if counts[item] < n/len(items)/2 {
|
||||
t.Fatalf("item %d underrepresented in random fallback: %d/%d", item, counts[item], n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHashStrategy_DifferentClientIDs(t *testing.T) {
|
||||
s := HashStrategy[int]()
|
||||
items := make([]int, 100)
|
||||
for i := range items {
|
||||
items[i] = i
|
||||
}
|
||||
|
||||
selected := make(map[int]bool)
|
||||
for i := 0; i < 100; i++ {
|
||||
ctx := xctx.ContextWithClientID(context.Background(), xctx.ClientID(string(rune('A'+i))))
|
||||
v := s.Apply(ctx, items...)
|
||||
selected[v] = true
|
||||
}
|
||||
|
||||
// Different client IDs should spread across multiple items
|
||||
if len(selected) < 10 {
|
||||
t.Fatalf("expected distribution across many items, got %d unique selections", len(selected))
|
||||
}
|
||||
}
|
||||
|
||||
// --- Interface compliance ---
|
||||
|
||||
func TestStrategyInterfaceCompliance(t *testing.T) {
|
||||
// Verify all strategies implement selector.Strategy
|
||||
_ = selector.Strategy[int](RoundRobinStrategy[int]())
|
||||
_ = selector.Strategy[int](RandomStrategy[int]())
|
||||
_ = selector.Strategy[int](FIFOStrategy[int]())
|
||||
_ = selector.Strategy[int](HashStrategy[int]())
|
||||
_ = selector.Strategy[int](ParallelStrategy[int]())
|
||||
}
|
||||
@@ -10,24 +10,30 @@ type randomWeightedItem[T any] struct {
|
||||
weight int
|
||||
}
|
||||
|
||||
// RandomWeighted is a weighted random selector.
|
||||
// Items are selected randomly with probability proportional to their weight.
|
||||
type RandomWeighted[T any] struct {
|
||||
items []*randomWeightedItem[T]
|
||||
sum int
|
||||
r *rand.Rand
|
||||
}
|
||||
|
||||
// NewRandomWeighted creates a new RandomWeighted selector.
|
||||
func NewRandomWeighted[T any]() *RandomWeighted[T] {
|
||||
return &RandomWeighted[T]{
|
||||
r: rand.New(rand.NewSource(time.Now().UnixNano())),
|
||||
}
|
||||
}
|
||||
|
||||
// Add adds an item with the given weight to the selector.
|
||||
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
|
||||
}
|
||||
|
||||
// Next selects an item randomly based on weights.
|
||||
// Returns the zero value if no items are added.
|
||||
func (rw *RandomWeighted[T]) Next() (v T) {
|
||||
if len(rw.items) == 0 {
|
||||
return
|
||||
@@ -46,6 +52,7 @@ func (rw *RandomWeighted[T]) Next() (v T) {
|
||||
return rw.items[len(rw.items)-1].item
|
||||
}
|
||||
|
||||
// Reset clears all items and weights.
|
||||
func (rw *RandomWeighted[T]) Reset() {
|
||||
rw.items = nil
|
||||
rw.sum = 0
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
package selector
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRandomWeighted_Empty(t *testing.T) {
|
||||
rw := NewRandomWeighted[int]()
|
||||
if v := rw.Next(); v != 0 {
|
||||
t.Fatalf("expected zero value, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_Single(t *testing.T) {
|
||||
rw := NewRandomWeighted[int]()
|
||||
rw.Add(42, 10)
|
||||
for i := 0; i < 100; i++ {
|
||||
if v := rw.Next(); v != 42 {
|
||||
t.Fatalf("expected 42, got %d", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_EqualWeights(t *testing.T) {
|
||||
rw := NewRandomWeighted[int]()
|
||||
rw.Add(1, 1)
|
||||
rw.Add(2, 1)
|
||||
rw.Add(3, 1)
|
||||
|
||||
counts := make(map[int]int)
|
||||
const n = 3000
|
||||
for i := 0; i < n; i++ {
|
||||
counts[rw.Next()]++
|
||||
}
|
||||
|
||||
for _, item := range []int{1, 2, 3} {
|
||||
if counts[item] < n/3/2 {
|
||||
t.Fatalf("item %d underrepresented: %d/%d", item, counts[item], n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_WeightedDistribution(t *testing.T) {
|
||||
rw := NewRandomWeighted[string]()
|
||||
rw.Add("heavy", 99)
|
||||
rw.Add("light", 1)
|
||||
|
||||
heavyCount := 0
|
||||
const n = 5000
|
||||
for i := 0; i < n; i++ {
|
||||
if rw.Next() == "heavy" {
|
||||
heavyCount++
|
||||
}
|
||||
}
|
||||
|
||||
// ~99% should be heavy
|
||||
if heavyCount < int(float64(n)*0.95) {
|
||||
t.Fatalf("heavy underrepresented: %d/%d", heavyCount, n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_ZeroWeight(t *testing.T) {
|
||||
rw := NewRandomWeighted[int]()
|
||||
rw.Add(1, 0)
|
||||
rw.Add(2, 0)
|
||||
// sum=0, should return zero value
|
||||
if v := rw.Next(); v != 0 {
|
||||
t.Fatalf("expected zero value with sum=0, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_Reset(t *testing.T) {
|
||||
rw := NewRandomWeighted[int]()
|
||||
rw.Add(1, 10)
|
||||
rw.Add(2, 20)
|
||||
|
||||
rw.Reset()
|
||||
if v := rw.Next(); v != 0 {
|
||||
t.Fatalf("expected zero value after reset, got %d", v)
|
||||
}
|
||||
|
||||
// Can add again after reset
|
||||
rw.Add(3, 1)
|
||||
if v := rw.Next(); v != 3 {
|
||||
t.Fatalf("expected 3 after re-add, got %d", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_MultipleWeights(t *testing.T) {
|
||||
rw := NewRandomWeighted[int]()
|
||||
rw.Add(0, 1)
|
||||
rw.Add(1, 2)
|
||||
rw.Add(2, 3)
|
||||
// Total weight = 6
|
||||
|
||||
counts := make(map[int]int)
|
||||
const n = 6000
|
||||
for i := 0; i < n; i++ {
|
||||
counts[rw.Next()]++
|
||||
}
|
||||
|
||||
// Item 2 should appear roughly 2x more than item 0
|
||||
ratio := float64(counts[2]) / float64(counts[0])
|
||||
if ratio < 1.5 || ratio > 3.5 {
|
||||
t.Fatalf("expected ratio ~2.0 for item2/item0, got %.2f (counts: %v)", ratio, counts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_StringItems(t *testing.T) {
|
||||
rw := NewRandomWeighted[string]()
|
||||
rw.Add("a", 1)
|
||||
rw.Add("b", 1)
|
||||
|
||||
v := rw.Next()
|
||||
if v != "a" && v != "b" {
|
||||
t.Fatalf("expected 'a' or 'b', got %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRandomWeighted_StructItems(t *testing.T) {
|
||||
type item struct{ Name string }
|
||||
rw := NewRandomWeighted[item]()
|
||||
rw.Add(item{Name: "x"}, 1)
|
||||
rw.Add(item{Name: "y"}, 1)
|
||||
|
||||
v := rw.Next()
|
||||
if v.Name != "x" && v.Name != "y" {
|
||||
t.Fatalf("unexpected item: %+v", v)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user