Create & Init Project...

library/net/rpc/warden/balancer/wrr/test/BUILD

@@ -0,0 +1,51 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["base_test.go"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = [
+        "//library/conf/env:go_default_library",
+        "//library/naming:go_default_library",
+        "//library/net/netutil/breaker:go_default_library",
+        "//library/net/rpc/warden:go_default_library",
+        "//library/net/rpc/warden/balancer/wrr:go_default_library",
+        "//library/net/rpc/warden/proto/testproto:go_default_library",
+        "//library/net/rpc/warden/resolver:go_default_library",
+        "//library/time:go_default_library",
+        "@org_golang_google_grpc//:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [
+        ":package-srcs",
+        "//library/net/rpc/warden/balancer/wrr/test/client:all-srcs",
+        "//library/net/rpc/warden/balancer/wrr/test/server:all-srcs",
+    ],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["base.go"],
+    importpath = "go-common/library/net/rpc/warden/balancer/wrr/test",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/net/rpc/warden/balancer/wrr/test/base.go

@@ -0,0 +1 @@
+package test

library/net/rpc/warden/balancer/wrr/test/base_test.go

@@ -0,0 +1,298 @@
+package test
+
+import (
+	"context"
+	"io"
+	"log"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	"go-common/library/conf/env"
+	"go-common/library/naming"
+	"go-common/library/net/netutil/breaker"
+	"go-common/library/net/rpc/warden"
+	"go-common/library/net/rpc/warden/balancer/wrr"
+	pb "go-common/library/net/rpc/warden/proto/testproto"
+	"go-common/library/net/rpc/warden/resolver"
+	xtime "go-common/library/time"
+
+	"google.golang.org/grpc"
+)
+
+type testBuilder struct {
+	addrs []*naming.Instance
+}
+type testDiscovery struct {
+	mu sync.Mutex
+	b  *testBuilder
+	id string
+	ch chan struct{}
+}
+
+func (b *testBuilder) Build(id string) naming.Resolver {
+	return &testDiscovery{id: id, b: b}
+}
+
+func (b *testBuilder) Scheme() string {
+	return "testbuilder"
+}
+func (d *testDiscovery) Fetch(ctx context.Context) (map[string][]*naming.Instance, bool) {
+	d.mu.Lock()
+	addrs := d.b.addrs
+	d.mu.Unlock()
+	if len(addrs) == 0 {
+		return nil, false
+	}
+	return map[string][]*naming.Instance{env.Zone: addrs}, true
+}
+
+func (d *testDiscovery) Watch() <-chan struct{} {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if d.ch == nil {
+		d.ch = make(chan struct{}, 1)
+	}
+	return d.ch
+}
+
+func (d *testDiscovery) Close() error {
+	return nil
+}
+
+func (d *testDiscovery) Scheme() string {
+	return "discovery"
+}
+
+func (d *testDiscovery) set(addrs []*naming.Instance) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.b.addrs = addrs
+	select {
+	case d.ch <- struct{}{}:
+	default:
+		return
+	}
+}
+
+func TestMain(m *testing.M) {
+	s1 := runServer(":18080")
+	s2 := runServer(":18081")
+	s3 := runServer(":18082")
+	b = &testBuilder{}
+	resolver.Register(b)
+	dis = b.Build("test_app").(*testDiscovery)
+	go func() {
+		time.Sleep(time.Millisecond * 10)
+		dis.set([]*naming.Instance{{
+			Addrs:    []string{"grpc://127.0.0.1:18080"},
+			AppID:    "test_app",
+			Metadata: map[string]string{"weight": "100"},
+		}, {
+			Addrs:    []string{"grpc://127.0.0.1:18081"},
+			AppID:    "test_app",
+			Metadata: map[string]string{"color": "red"},
+		}, {
+			Addrs: []string{"grpc://127.0.0.1:18082"},
+			AppID: "test_app",
+		}})
+	}()
+	c = newClient()
+	time.Sleep(time.Millisecond * 30)
+	ret := m.Run()
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
+	defer cancel()
+	s1.Shutdown(ctx)
+	s2.Shutdown(ctx)
+	s3.Shutdown(ctx)
+	os.Exit(ret)
+}
+
+type helloServer struct {
+	addr string
+}
+
+func (s *helloServer) SayHello(ctx context.Context, in *pb.HelloRequest) (*pb.HelloReply, error) {
+	return &pb.HelloReply{Message: s.addr}, nil
+}
+
+func (s *helloServer) StreamHello(ss pb.Greeter_StreamHelloServer) error {
+	for i := 0; i < 3; i++ {
+		in, err := ss.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		ret := &pb.HelloReply{Message: "Hello " + in.Name, Success: true}
+		err = ss.Send(ret)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func runServer(addr string) *warden.Server {
+	server := warden.NewServer(&warden.ServerConfig{Timeout: xtime.Duration(time.Second)})
+	pb.RegisterGreeterServer(server.Server(), &helloServer{addr: addr})
+	go func() {
+		err := server.Run(addr)
+		if err != nil {
+			panic("run server failed!" + err.Error())
+		}
+	}()
+	return server
+}
+
+// NewClient returns a new blank Client instance with a default client interceptor.
+// opt can be used to add grpc dial options.
+func newClient() (client pb.GreeterClient) {
+	c := warden.NewClient(&warden.ClientConfig{
+		Dial:    xtime.Duration(time.Second * 10),
+		Timeout: xtime.Duration(time.Second * 10),
+		Breaker: &breaker.Config{
+			Window:  xtime.Duration(3 * time.Second),
+			Sleep:   xtime.Duration(3 * time.Second),
+			Bucket:  10,
+			Ratio:   0.3,
+			Request: 20,
+		},
+	},
+		grpc.WithBalancerName(wrr.Name),
+	)
+	conn, err := c.Dial(context.Background(), "discovery://authority/111")
+	if err != nil {
+		log.Fatalf("can't not connect: %v", err)
+	}
+	client = pb.NewGreeterClient(conn)
+	return
+}
+
+var b *testBuilder
+var dis *testDiscovery
+var c pb.GreeterClient
+
+func TestBalancer(t *testing.T) {
+	testBalancerBasic(t)
+	testBalancerFailover(t)
+	testBalancerUpdateColor(t)
+	testBalancerUpdateScore(t)
+}
+
+func testBalancerBasic(t *testing.T) {
+	time.Sleep(time.Millisecond * 10)
+	var idx8080 int
+	var idx8082 int
+	for i := 0; i < 6; i++ {
+		resp, err := c.SayHello(context.Background(), &pb.HelloRequest{Age: 123, Name: "asdasd"})
+		if err != nil {
+			t.Fatalf("testBalancerBasic: say hello failed!err:=%v", err)
+		}
+		if resp.Message == ":18082" {
+			idx8082++
+		} else if resp.Message == ":18080" {
+			idx8080++
+		}
+	}
+	if idx8080 != 3 {
+		t.Fatalf("testBalancerBasic: server 18080 response times should be 3")
+	}
+	if idx8082 != 3 {
+		t.Fatalf("testBalancerBasic: server 18082 response times should be 3")
+	}
+}
+
+func testBalancerFailover(t *testing.T) {
+	dis.set([]*naming.Instance{{
+		Addrs:    []string{"grpc://127.0.0.1:18080"},
+		AppID:    "test_app",
+		Metadata: map[string]string{"weight": "100"},
+	}, {
+		Addrs:    []string{"grpc://127.0.0.1:18081"},
+		AppID:    "test_app",
+		Metadata: map[string]string{"color": "red"},
+	}})
+	time.Sleep(time.Millisecond * 20)
+	var idx8080 int
+	var idx8082 int
+	for i := 0; i < 4; i++ {
+		resp, err := c.SayHello(context.Background(), &pb.HelloRequest{Age: 123, Name: "asdasd"})
+		if err != nil {
+			t.Fatalf("testBalancerFailover: say hello failed!err:=%v", err)
+		}
+		if resp.Message == ":18082" {
+			idx8082++
+		} else if resp.Message == ":18080" {
+			idx8080++
+		}
+	}
+	if idx8080 != 4 {
+		t.Fatalf("testBalancerFailover: server 8080  response should be 4")
+	}
+}
+
+func testBalancerUpdateColor(t *testing.T) {
+	dis.set([]*naming.Instance{{
+		Addrs:    []string{"grpc://127.0.0.1:18080"},
+		AppID:    "test_app",
+		Metadata: map[string]string{"weight": "100"},
+	}, {
+		Addrs: []string{"grpc://127.0.0.1:18081"},
+		AppID: "test_app",
+	}})
+	time.Sleep(time.Millisecond * 30)
+	var idx8080 int
+	var idx8081 int
+	for i := 0; i < 4; i++ {
+		resp, err := c.SayHello(context.Background(), &pb.HelloRequest{Age: 123, Name: "asdasd"})
+		if err != nil {
+			t.Fatalf("testBalancerUpdateColor: say hello failed!err:=%v", err)
+		}
+		if resp.Message == ":18081" {
+			idx8081++
+		} else if resp.Message == ":18080" {
+			idx8080++
+		}
+	}
+	if idx8080 != 2 {
+		t.Fatalf("testBalancerUpdateColor: server 8080 response should be 2")
+	}
+	if idx8081 != 2 {
+		t.Fatalf("testBalancerUpdateColor: server 8081 response should be 2")
+	}
+}
+
+func testBalancerUpdateScore(t *testing.T) {
+	dis.set([]*naming.Instance{{
+		Addrs:    []string{"grpc://127.0.0.1:18080"},
+		AppID:    "test_app",
+		Metadata: map[string]string{"weight": "100"},
+	}, {
+		Addrs:    []string{"grpc://127.0.0.1:18081"},
+		AppID:    "test_app",
+		Metadata: map[string]string{"weight": "300"},
+	}})
+	time.Sleep(time.Millisecond * 10)
+	var idx8080 int
+	var idx8081 int
+	for i := 0; i < 4; i++ {
+		resp, err := c.SayHello(context.Background(), &pb.HelloRequest{Age: 123, Name: "asdasd"})
+		if err != nil {
+			t.Fatalf("testBalancerUpdateScore: say hello failed!err:=%v", err)
+		}
+		if resp.Message == ":18081" {
+			idx8081++
+		} else if resp.Message == ":18080" {
+			idx8080++
+		}
+	}
+	if idx8080 != 1 {
+		t.Fatalf("testBalancerUpdateScore: server 8080 response should be 2")
+	}
+	if idx8081 != 3 {
+		t.Fatalf("testBalancerUpdateScore: server 8081 response should be 2")
+	}
+}

library/net/rpc/warden/balancer/wrr/test/client/BUILD

@@ -0,0 +1,43 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_binary",
+    "go_library",
+)
+
+go_binary(
+    name = "client",
+    embed = [":go_default_library"],
+    tags = ["automanaged"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["client.go"],
+    importpath = "go-common/library/net/rpc/warden/balancer/wrr/test/client",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//library/exp/feature:go_default_library",
+        "//library/log:go_default_library",
+        "//library/net/rpc/warden:go_default_library",
+        "//library/net/rpc/warden/proto/testproto:go_default_library",
+        "//library/net/rpc/warden/resolver:go_default_library",
+        "//library/net/rpc/warden/resolver/direct:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/net/rpc/warden/balancer/wrr/test/client/client.go

@@ -0,0 +1,73 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"sync/atomic"
+	"time"
+
+	"go-common/library/exp/feature"
+	"go-common/library/log"
+	"go-common/library/net/rpc/warden"
+	pb "go-common/library/net/rpc/warden/proto/testproto"
+	"go-common/library/net/rpc/warden/resolver"
+	"go-common/library/net/rpc/warden/resolver/direct"
+)
+
+var addrs string
+var cli pb.GreeterClient
+var concurrency int
+var name string
+var req int64
+var qps int64
+
+func init() {
+	log.Init(&log.Config{Stdout: false})
+	flag.StringVar(&addrs, "addr", "127.0.0.1:8000,127.0.0.1:8001", "-addr 127.0.0.1:8080,127.0.0.1:8081")
+	flag.IntVar(&concurrency, "c", 3, "-c 5")
+	flag.StringVar(&name, "name", "test", "-name test")
+}
+
+func main() {
+	go calcuQPS()
+	feature.DefaultGate.AddFlag(flag.CommandLine)
+	flag.Parse()
+	feature.DefaultGate.SetFromMap(map[string]bool{"dwrr": true})
+	resolver.Register(direct.New())
+	c := warden.NewClient(nil)
+	conn, err := c.Dial(context.Background(), fmt.Sprintf("direct://d/%s", addrs))
+	if err != nil {
+		panic(err)
+	}
+	cli = pb.NewGreeterClient(conn)
+	for i := 0; i < concurrency; i++ {
+		go func() {
+			for {
+				say()
+				time.Sleep(time.Millisecond * 5)
+			}
+		}()
+	}
+	time.Sleep(time.Hour)
+}
+
+func calcuQPS() {
+	var creq, breq int64
+	for {
+		time.Sleep(time.Second * 5)
+		creq = atomic.LoadInt64(&req)
+		delta := creq - breq
+		atomic.StoreInt64(&qps, delta/5)
+		breq = creq
+		fmt.Println("HTTP QPS: ", atomic.LoadInt64(&qps))
+	}
+}
+func say() {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	reply, err := cli.SayHello(ctx, &pb.HelloRequest{Name: name, Age: 10})
+	if err == nil && reply.Success {
+		atomic.AddInt64(&req, 1)
+	}
+}

library/net/rpc/warden/balancer/wrr/test/server/BUILD

@@ -0,0 +1,41 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_binary",
+    "go_library",
+)
+
+go_binary(
+    name = "server",
+    embed = [":go_default_library"],
+    tags = ["automanaged"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["server.go"],
+    importpath = "go-common/library/net/rpc/warden/balancer/wrr/test/server",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//library/ecode:go_default_library",
+        "//library/log:go_default_library",
+        "//library/net/rpc/warden:go_default_library",
+        "//library/net/rpc/warden/proto/testproto:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/net/rpc/warden/balancer/wrr/test/server/server.go

@@ -0,0 +1,91 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"hash/crc32"
+	"io"
+	"math/rand"
+	"sync/atomic"
+	"time"
+
+	"go-common/library/ecode"
+	"go-common/library/log"
+	"go-common/library/net/rpc/warden"
+	pb "go-common/library/net/rpc/warden/proto/testproto"
+)
+
+var (
+	req     int64
+	qps     int64
+	cpu     int
+	errRate int
+	sleep   time.Duration
+)
+
+func init() {
+	log.Init(&log.Config{Stdout: false})
+	flag.IntVar(&cpu, "cpu", 3000, "cpu time")
+	flag.IntVar(&errRate, "err", 0, "error rate")
+	flag.DurationVar(&sleep, "sleep", 0, "sleep time")
+}
+
+func calcuQPS() {
+	var creq, breq int64
+	for {
+		time.Sleep(time.Second * 5)
+		creq = atomic.LoadInt64(&req)
+		delta := creq - breq
+		atomic.StoreInt64(&qps, delta/5)
+		breq = creq
+		fmt.Println("HTTP QPS: ", atomic.LoadInt64(&qps))
+	}
+}
+
+func main() {
+	flag.Parse()
+	server := warden.NewServer(nil)
+	pb.RegisterGreeterServer(server.Server(), &helloServer{})
+	_, err := server.Start()
+	if err != nil {
+		panic(err)
+	}
+	calcuQPS()
+}
+
+type helloServer struct {
+}
+
+func (s *helloServer) SayHello(ctx context.Context, in *pb.HelloRequest) (*pb.HelloReply, error) {
+	atomic.AddInt64(&req, 1)
+	if in.Name == "err" {
+		if rand.Intn(100) < errRate {
+			return nil, ecode.ServiceUnavailable
+		}
+	}
+	time.Sleep(time.Millisecond * time.Duration(in.Age))
+	time.Sleep(sleep)
+	for i := 0; i < cpu+rand.Intn(cpu); i++ {
+		crc32.Checksum([]byte(`testasdwfwfsddsfgwddcscschttp://git.bilibili.co/platform/go-common/merge_requests/new?merge_request%5Bsource_branch%5D=stress%2Fcodel`), crc32.IEEETable)
+	}
+	return &pb.HelloReply{Message: "Hello " + in.Name, Success: true}, nil
+}
+
+func (s *helloServer) StreamHello(ss pb.Greeter_StreamHelloServer) error {
+	for i := 0; i < 3; i++ {
+		in, err := ss.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		ret := &pb.HelloReply{Message: "Hello " + in.Name, Success: true}
+		err = ss.Send(ret)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}