feat: refactor.

2024-05-06 10:26:59 +08:00
parent 7ce746e181
commit 72daf43151
13 changed files with 100 additions and 81 deletions
--- a/rpc/grpc.go
+++ b/rpc/grpc.go
@@ -0,0 +1,109 @@
+package rpc
+
+import (
+	"context"
+	"time"
+
+	"google.golang.org/grpc"
+
+	"github.com/go-kratos/aegis/ratelimit"
+	"github.com/go-kratos/aegis/ratelimit/bbr"
+
+	"github.com/go-kratos/kratos/v2/log"
+	"github.com/go-kratos/kratos/v2/registry"
+
+	"github.com/go-kratos/kratos/v2/middleware"
+	midRateLimit "github.com/go-kratos/kratos/v2/middleware/ratelimit"
+	"github.com/go-kratos/kratos/v2/middleware/recovery"
+	"github.com/go-kratos/kratos/v2/middleware/tracing"
+	"github.com/go-kratos/kratos/v2/middleware/validate"
+
+	kratosGrpc "github.com/go-kratos/kratos/v2/transport/grpc"
+
+	conf "github.com/tx7do/kratos-bootstrap/api/gen/go/conf/v1"
+)
+
+const defaultTimeout = 5 * time.Second
+
+// CreateGrpcClient 创建GRPC客户端
+func CreateGrpcClient(ctx context.Context, r registry.Discovery, serviceName string, cfg *conf.Bootstrap, m ...middleware.Middleware) grpc.ClientConnInterface {
+	endpoint := "discovery:///" + serviceName
+
+	var ms []middleware.Middleware
+	timeout := defaultTimeout
+	if cfg.Client != nil && cfg.Client.Grpc != nil {
+		if cfg.Client.Grpc.Timeout != nil {
+			timeout = cfg.Client.Grpc.Timeout.AsDuration()
+		}
+
+		if cfg.Client.Grpc.Middleware != nil {
+			if cfg.Client.Grpc.Middleware.GetEnableRecovery() {
+				ms = append(ms, recovery.Recovery())
+			}
+			if cfg.Client.Grpc.Middleware.GetEnableTracing() {
+				ms = append(ms, tracing.Client())
+			}
+			if cfg.Client.Grpc.Middleware.GetEnableValidate() {
+				ms = append(ms, validate.Validator())
+			}
+		}
+	}
+	ms = append(ms, m...)
+
+	conn, err := kratosGrpc.DialInsecure(
+		ctx,
+		kratosGrpc.WithEndpoint(endpoint),
+		kratosGrpc.WithDiscovery(r),
+		kratosGrpc.WithTimeout(timeout),
+		kratosGrpc.WithMiddleware(ms...),
+	)
+	if err != nil {
+		log.Fatalf("dial grpc client [%s] failed: %s", serviceName, err.Error())
+	}
+
+	return conn
+}
+
+// CreateGrpcServer 创建GRPC服务端
+func CreateGrpcServer(cfg *conf.Bootstrap, m ...middleware.Middleware) *kratosGrpc.Server {
+	var opts []kratosGrpc.ServerOption
+
+	var ms []middleware.Middleware
+	if cfg.Server != nil && cfg.Server.Grpc != nil && cfg.Server.Grpc.Middleware != nil {
+		if cfg.Server.Grpc.Middleware.GetEnableRecovery() {
+			ms = append(ms, recovery.Recovery())
+		}
+		if cfg.Server.Grpc.Middleware.GetEnableTracing() {
+			ms = append(ms, tracing.Server())
+		}
+		if cfg.Server.Grpc.Middleware.GetEnableValidate() {
+			ms = append(ms, validate.Validator())
+		}
+		if cfg.Server.Grpc.Middleware.GetEnableCircuitBreaker() {
+		}
+		if cfg.Server.Grpc.Middleware.Limiter != nil {
+			var limiter ratelimit.Limiter
+			switch cfg.Server.Grpc.Middleware.Limiter.GetName() {
+			case "bbr":
+				limiter = bbr.NewLimiter()
+			}
+			ms = append(ms, midRateLimit.Server(midRateLimit.WithLimiter(limiter)))
+		}
+	}
+	ms = append(ms, m...)
+	opts = append(opts, kratosGrpc.Middleware(ms...))
+
+	if cfg.Server.Grpc.Network != "" {
+		opts = append(opts, kratosGrpc.Network(cfg.Server.Grpc.Network))
+	}
+	if cfg.Server.Grpc.Addr != "" {
+		opts = append(opts, kratosGrpc.Address(cfg.Server.Grpc.Addr))
+	}
+	if cfg.Server.Grpc.Timeout != nil {
+		opts = append(opts, kratosGrpc.Timeout(cfg.Server.Grpc.Timeout.AsDuration()))
+	}
+
+	srv := kratosGrpc.NewServer(opts...)
+
+	return srv
+}
--- a/rpc/rest.go
+++ b/rpc/rest.go
@@ -0,0 +1,89 @@
+package rpc
+
+import (
+	"net/http/pprof"
+
+	"github.com/go-kratos/aegis/ratelimit"
+	"github.com/go-kratos/aegis/ratelimit/bbr"
+
+	"github.com/go-kratos/kratos/v2/middleware"
+	midRateLimit "github.com/go-kratos/kratos/v2/middleware/ratelimit"
+	"github.com/go-kratos/kratos/v2/middleware/recovery"
+	"github.com/go-kratos/kratos/v2/middleware/tracing"
+	"github.com/go-kratos/kratos/v2/middleware/validate"
+
+	kratosRest "github.com/go-kratos/kratos/v2/transport/http"
+
+	"github.com/gorilla/handlers"
+
+	conf "github.com/tx7do/kratos-bootstrap/api/gen/go/conf/v1"
+)
+
+// CreateRestServer 创建REST服务端
+func CreateRestServer(cfg *conf.Bootstrap, m ...middleware.Middleware) *kratosRest.Server {
+	var opts = []kratosRest.ServerOption{
+		kratosRest.Filter(handlers.CORS(
+			handlers.AllowedHeaders(cfg.Server.Rest.Cors.Headers),
+			handlers.AllowedMethods(cfg.Server.Rest.Cors.Methods),
+			handlers.AllowedOrigins(cfg.Server.Rest.Cors.Origins),
+		)),
+	}
+
+	var ms []middleware.Middleware
+	if cfg.Server != nil && cfg.Server.Rest != nil && cfg.Server.Rest.Middleware != nil {
+		if cfg.Server.Rest.Middleware.GetEnableRecovery() {
+			ms = append(ms, recovery.Recovery())
+		}
+		if cfg.Server.Rest.Middleware.GetEnableTracing() {
+			ms = append(ms, tracing.Server())
+		}
+		if cfg.Server.Rest.Middleware.GetEnableValidate() {
+			ms = append(ms, validate.Validator())
+		}
+		if cfg.Server.Rest.Middleware.GetEnableCircuitBreaker() {
+		}
+		if cfg.Server.Rest.Middleware.Limiter != nil {
+			var limiter ratelimit.Limiter
+			switch cfg.Server.Rest.Middleware.Limiter.GetName() {
+			case "bbr":
+				limiter = bbr.NewLimiter()
+			}
+			ms = append(ms, midRateLimit.Server(midRateLimit.WithLimiter(limiter)))
+		}
+	}
+	ms = append(ms, m...)
+	opts = append(opts, kratosRest.Middleware(ms...))
+
+	if cfg.Server.Rest.Network != "" {
+		opts = append(opts, kratosRest.Network(cfg.Server.Rest.Network))
+	}
+	if cfg.Server.Rest.Addr != "" {
+		opts = append(opts, kratosRest.Address(cfg.Server.Rest.Addr))
+	}
+	if cfg.Server.Rest.Timeout != nil {
+		opts = append(opts, kratosRest.Timeout(cfg.Server.Rest.Timeout.AsDuration()))
+	}
+
+	srv := kratosRest.NewServer(opts...)
+
+	if cfg.Server.Rest.GetEnablePprof() {
+		registerHttpPprof(srv)
+	}
+
+	return srv
+}
+
+func registerHttpPprof(s *kratosRest.Server) {
+	s.HandleFunc("/debug/pprof", pprof.Index)
+	s.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	s.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	s.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	s.HandleFunc("/debug/pprof/trace", pprof.Trace)
+
+	s.HandleFunc("/debug/pprof/allocs", pprof.Handler("allocs").ServeHTTP)
+	s.HandleFunc("/debug/pprof/block", pprof.Handler("block").ServeHTTP)
+	s.HandleFunc("/debug/pprof/goroutine", pprof.Handler("goroutine").ServeHTTP)
+	s.HandleFunc("/debug/pprof/heap", pprof.Handler("heap").ServeHTTP)
+	s.HandleFunc("/debug/pprof/mutex", pprof.Handler("mutex").ServeHTTP)
+	s.HandleFunc("/debug/pprof/threadcreate", pprof.Handler("threadcreate").ServeHTTP)
+}