package p2c import ( "context" "fmt" "math/rand" "reflect" "sync" "sync/atomic" "testing" "time" "github.com/go-kratos/kratos/v2/registry" "github.com/go-kratos/kratos/v2/selector" "github.com/go-kratos/kratos/v2/selector/filter" ) func TestWrr3(t *testing.T) { p2c := New() var nodes []selector.Node for i := 0; i < 3; i++ { addr := fmt.Sprintf("127.0.0.%d:8080", i) nodes = append(nodes, selector.NewNode( "http", addr, ®istry.ServiceInstance{ ID: addr, Version: "v2.0.0", Metadata: map[string]string{"weight": "10"}, })) } p2c.Apply(nodes) var count1, count2, count3 int64 group := &sync.WaitGroup{} var lk sync.Mutex for i := 0; i < 9000; i++ { group.Add(1) go func() { defer group.Done() lk.Lock() d := time.Duration(rand.Intn(500)) * time.Millisecond lk.Unlock() time.Sleep(d) n, done, err := p2c.Select(context.Background(), selector.WithNodeFilter(filter.Version("v2.0.0"))) if err != nil { t.Errorf("expect %v, got %v", nil, err) } if n == nil { t.Errorf("expect %v, got %v", nil, n) } if done == nil { t.Errorf("expect %v, got %v", nil, done) } time.Sleep(time.Millisecond * 10) done(context.Background(), selector.DoneInfo{}) if n.Address() == "127.0.0.0:8080" { atomic.AddInt64(&count1, 1) } else if n.Address() == "127.0.0.1:8080" { atomic.AddInt64(&count2, 1) } else if n.Address() == "127.0.0.2:8080" { atomic.AddInt64(&count3, 1) } }() } group.Wait() if count1 <= int64(1500) { t.Errorf("count1(%v) <= int64(1500)", count1) } if count1 >= int64(4500) { t.Errorf("count1(%v) >= int64(4500),", count1) } if count2 <= int64(1500) { t.Errorf("count2(%v) <= int64(1500)", count2) } if count2 >= int64(4500) { t.Errorf("count2(%v) >= int64(4500),", count2) } if count3 <= int64(1500) { t.Errorf("count3(%v) <= int64(1500)", count3) } if count3 >= int64(4500) { t.Errorf("count3(%v) >= int64(4500),", count3) } } func TestEmpty(t *testing.T) { b := &Balancer{} _, _, err := b.Pick(context.Background(), []selector.WeightedNode{}) if err == nil { t.Errorf("expect %v, got %v", nil, err) } } func TestOne(t *testing.T) { p2c := New() var nodes []selector.Node for i := 0; i < 1; i++ { addr := fmt.Sprintf("127.0.0.%d:8080", i) nodes = append(nodes, selector.NewNode( "http", addr, ®istry.ServiceInstance{ ID: addr, Version: "v2.0.0", Metadata: map[string]string{"weight": "10"}, })) } p2c.Apply(nodes) n, done, err := p2c.Select(context.Background(), selector.WithNodeFilter(filter.Version("v2.0.0"))) if err != nil { t.Errorf("expect %v, got %v", nil, err) } if n == nil { t.Errorf("expect %v, got %v", nil, n) } if done == nil { t.Errorf("expect %v, got %v", nil, done) } if !reflect.DeepEqual("127.0.0.0:8080", n.Address()) { t.Errorf("expect %v, got %v", "127.0.0.0:8080", n.Address()) } }