feat : init

.gitignore

@@ -0,0 +1,3 @@
+.idea/
+/app/bench/bench
+/app/bench_cli/bench_cli

Makefile

@@ -0,0 +1,16 @@
+pb:
+	goctl  rpc protoc bench.proto --go_out=./pb --go-grpc_out=./pb --zrpc_out=./app/bench
+
+build:
+	cd app/bench && go build
+	cd app/bench_cli && go build
+
+run_bench:
+	nohup app/bench/bench -f app/bench/etc/bench.yaml &
+	nohup app/bench/bench -f app/bench/etc/bench2.yaml &
+
+run_bench_cli:
+	cd app/bench_cli && ./bench_cli
+
+kill:
+	pkill bench

app/bench/bench.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+
+	"pomelo_bench/app/bench/internal/config"
+	"pomelo_bench/app/bench/internal/server"
+	"pomelo_bench/app/bench/internal/svc"
+	"pomelo_bench/pb/bench"
+
+	"github.com/zeromicro/go-zero/core/conf"
+	"github.com/zeromicro/go-zero/core/service"
+	"github.com/zeromicro/go-zero/zrpc"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/reflection"
+)
+
+var configFile = flag.String("f", "etc/bench.yaml", "the config file")
+
+func main() {
+	flag.Parse()
+
+	var c config.Config
+	conf.MustLoad(*configFile, &c)
+	ctx := svc.NewServiceContext(c)
+
+	s := zrpc.MustNewServer(c.RpcServerConf, func(grpcServer *grpc.Server) {
+		bench.RegisterBenchServer(grpcServer, server.NewBenchServer(ctx))
+
+		if c.Mode == service.DevMode || c.Mode == service.TestMode {
+			reflection.Register(grpcServer)
+		}
+	})
+	defer s.Stop()
+
+	fmt.Printf("Starting rpc server at %s...\n", c.ListenOn)
+	s.Start()
+}

app/bench/etc/bench.yaml

@@ -0,0 +1,13 @@
+Name: "bench.rpc"
+ListenOn: "0.0.0.0:8080"
+Mode: "test"
+Verbose: true
+Timeout: 0
+Prometheus:
+  Host: "0.0.0.0"
+  Port: 9101
+  Path: "/metrics"
+Log:
+  Mode: "file"
+  Path: "logs"
+  Level: "info"

app/bench/etc/bench2.yaml

@@ -0,0 +1,13 @@
+Name: "bench.rpc"
+ListenOn: "0.0.0.0:8081"
+Mode: "test"
+Verbose: true
+Timeout: 0
+Prometheus:
+  Host: "0.0.0.0"
+  Port: 9101
+  Path: "/metrics"
+Log:
+  Mode: "file"
+  Path: "logs"
+  Level: "info"

app/bench/internal/config/config.go

@@ -0,0 +1,7 @@
+package config
+
+import "github.com/zeromicro/go-zero/zrpc"
+
+type Config struct {
+	zrpc.RpcServerConf
+}

app/bench/internal/logic/closeplanlogic.go

@@ -0,0 +1,54 @@
+package logic
+
+import (
+	"context"
+	"github.com/zeromicro/go-zero/core/logx"
+	"pomelo_bench/app/bench/internal/service/planmanager"
+	"pomelo_bench/app/bench/internal/svc"
+	"pomelo_bench/pb/bench"
+)
+
+type ClosePlanLogic struct {
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+	logx.Logger
+}
+
+func NewClosePlanLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ClosePlanLogic {
+	return &ClosePlanLogic{
+		ctx:    ctx,
+		svcCtx: svcCtx,
+		Logger: logx.WithContext(ctx),
+	}
+}
+
+// ClosePlan 关闭压测计划
+func (l *ClosePlanLogic) ClosePlan(in *bench.ClosePlanRequest) (*bench.ClosePlanResponse, error) {
+
+	var (
+		plans []*planmanager.Plan
+	)
+
+	if in.Uuid != nil { // 说明单发
+
+		p, err := l.svcCtx.PlanManager.GetPlan(in.GetUuid())
+		if err != nil {
+			return nil, err
+		}
+
+		plans = []*planmanager.Plan{p}
+
+	} else {
+		plans = l.svcCtx.PlanManager.GetAllPlan()
+	}
+
+	for _, p := range plans {
+
+		err := l.svcCtx.PlanManager.ClosePlan(p)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return &bench.ClosePlanResponse{}, nil
+}

app/bench/internal/logic/customsendlogic.go

@@ -0,0 +1,53 @@
+package logic
+
+import (
+	"context"
+	"github.com/zeromicro/go-zero/core/logx"
+	"pomelo_bench/app/bench/internal/service/planmanager"
+	"pomelo_bench/app/bench/internal/svc"
+	"pomelo_bench/pb/bench"
+)
+
+type CustomSendLogic struct {
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+	logx.Logger
+}
+
+func NewCustomSendLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CustomSendLogic {
+	return &CustomSendLogic{
+		ctx:    ctx,
+		svcCtx: svcCtx,
+		Logger: logx.WithContext(ctx),
+	}
+}
+
+// CustomSend 自定义消息发送
+func (l *CustomSendLogic) CustomSend(in *bench.CustomSendRequest) (*bench.CustomSendResponse, error) {
+
+	var (
+		plans []*planmanager.Plan
+	)
+
+	if in.Uuid != nil { // 说明单发
+
+		p, err := l.svcCtx.PlanManager.GetPlan(in.GetUuid())
+		if err != nil {
+			return nil, err
+		}
+
+		plans = []*planmanager.Plan{p}
+
+	} else {
+		plans = l.svcCtx.PlanManager.GetAllPlan()
+	}
+
+	for _, p := range plans {
+		err := p.PlanCustomSend(l.ctx, in.Pool, in.Number, in.Limit, in.Duration)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return &bench.CustomSendResponse{}, nil
+}

app/bench/internal/logic/detailplanlogic.go

@@ -0,0 +1,44 @@
+package logic
+
+import (
+	"context"
+	"github.com/zeromicro/go-zero/core/logx"
+	"pomelo_bench/app/bench/internal/logic/transform"
+	"pomelo_bench/app/bench/internal/svc"
+	"pomelo_bench/pb/bench"
+)
+
+type DetailPlanLogic struct {
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+	logx.Logger
+}
+
+func NewDetailPlanLogic(ctx context.Context, svcCtx *svc.ServiceContext) *DetailPlanLogic {
+	return &DetailPlanLogic{
+		ctx:    ctx,
+		svcCtx: svcCtx,
+		Logger: logx.WithContext(ctx),
+	}
+}
+
+// DetailPlan 查询压测计划详情
+func (l *DetailPlanLogic) DetailPlan(in *bench.DetailPlanRequest) (*bench.DetailPlanResponse, error) {
+	p, err := l.svcCtx.PlanManager.GetPlan(in.GetUuid())
+	if err != nil {
+		return nil, err
+	}
+
+	detail := p.PlanDetail(l.ctx)
+
+	res := &bench.DetailPlanResponse{
+		Uuid: in.GetUuid(),
+		Detail: &bench.PlanDetail{
+			Plan:       detail.Cfg,
+			Connectors: transform.Connectors(detail.Connectors),
+		},
+		Status: detail.Status,
+	}
+
+	return res, nil
+}

app/bench/internal/logic/listplanlogic.go

@@ -0,0 +1,45 @@
+package logic
+
+import (
+	"context"
+	"github.com/zeromicro/go-zero/core/logx"
+	"pomelo_bench/app/bench/internal/logic/transform"
+	"pomelo_bench/app/bench/internal/svc"
+	"pomelo_bench/pb/bench"
+)
+
+type ListPlanLogic struct {
+	ctx    context.Context
+	svcCtx *svc.ServiceContext
+	logx.Logger
+}
+
+func NewListPlanLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ListPlanLogic {
+	return &ListPlanLogic{
+		ctx:    ctx,
+		svcCtx: svcCtx,
+		Logger: logx.WithContext(ctx),
+	}
+}
+
+// ListPlan 查询压测计划
+func (l *ListPlanLogic) ListPlan(in *bench.ListPlanRequest) (*bench.ListPlanResponse, error) {
+
+	plans := l.svcCtx.PlanManager.ListPlan()
+
+	p := make([]*bench.PlanMonitor, 0, len(plans))
+
+	for i := 0; i < len(plans); i++ {
+
+		p = append(p, &bench.PlanMonitor{
+			Uuid:      plans[i].Uuid,
+			Plan:      plans[i].Cfg,
+			Status:    plans[i].Status,
+			Connector: plans[i].Connector,
+			Total:     transform.Statistics(plans[i].TotalStatistics),
+			Stat:      transform.Stat(plans[i].Stat),
+		})
+	}
+
+	return &bench.ListPlanResponse{Plans: p}, nil
+}