From aabebd047bbc4d428cf3c48a71a74c5ea9ec401f Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Mon, 25 May 2026 23:38:35 +0800 Subject: [PATCH] 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. --- selector/filter_test.go | 266 +++++++++++++++++++++++++++++++++++++ selector/parallel.go | 9 +- selector/parallel_test.go | 238 +++++++++++++++++++++++++++++++++ selector/selector_test.go | 53 ++++++++ selector/setup_test.go | 14 ++ selector/strategy.go | 9 +- selector/strategy_test.go | 271 ++++++++++++++++++++++++++++++++++++++ selector/weighted.go | 7 + selector/weighted_test.go | 130 ++++++++++++++++++ 9 files changed, 992 insertions(+), 5 deletions(-) create mode 100644 selector/filter_test.go create mode 100644 selector/parallel_test.go create mode 100644 selector/selector_test.go create mode 100644 selector/setup_test.go create mode 100644 selector/strategy_test.go create mode 100644 selector/weighted_test.go diff --git a/selector/filter_test.go b/selector/filter_test.go new file mode 100644 index 00000000..98b3b6b9 --- /dev/null +++ b/selector/filter_test.go @@ -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) + } +} diff --git a/selector/parallel.go b/selector/parallel.go index 2f418926..742931bf 100644 --- a/selector/parallel.go +++ b/selector/parallel.go @@ -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...), diff --git a/selector/parallel_test.go b/selector/parallel_test.go new file mode 100644 index 00000000..f6751b0a --- /dev/null +++ b/selector/parallel_test.go @@ -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() +} diff --git a/selector/selector_test.go b/selector/selector_test.go new file mode 100644 index 00000000..32125f12 --- /dev/null +++ b/selector/selector_test.go @@ -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) + } +} diff --git a/selector/setup_test.go b/selector/setup_test.go new file mode 100644 index 00000000..f6a26efe --- /dev/null +++ b/selector/setup_test.go @@ -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()) +} diff --git a/selector/strategy.go b/selector/strategy.go index 8458d797..352c925a 100644 --- a/selector/strategy.go +++ b/selector/strategy.go @@ -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))] } diff --git a/selector/strategy_test.go b/selector/strategy_test.go new file mode 100644 index 00000000..a25c1fba --- /dev/null +++ b/selector/strategy_test.go @@ -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]()) +} diff --git a/selector/weighted.go b/selector/weighted.go index bb50f2ac..a1826314 100644 --- a/selector/weighted.go +++ b/selector/weighted.go @@ -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 diff --git a/selector/weighted_test.go b/selector/weighted_test.go new file mode 100644 index 00000000..a7115a70 --- /dev/null +++ b/selector/weighted_test.go @@ -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) + } +}