更新脚手架

cmd/internal/grpc/README.tmpl

@@ -9,7 +9,8 @@ Go gRPC项目快速开发脚手架
 - 命令行使用 [cobra](https://github.com/spf13/cobra)
 - 工具包使用 [yiigo](https://github.com/shenghui0779/yiigo)
 - 使用 [grpc-gateway](https://github.com/grpc-ecosystem/grpc-gateway) 同时支持 grpc 和 http 服务
+- http服务支持跨域
 - 支持 proto 参数验证
 - 支持 swagger.json 生成
-- 包含 TraceId、请求日志 等中间价
+- 包含 TraceId、请求日志 等中间件
 - 简单好用的 Result Status 统一输出方式

cmd/internal/grpc/app/README.tmpl

@@ -23,9 +23,9 @@ go install github.com/go-swagger/go-swagger/cmd/swagger@latest
 
 ### 配置运行
 
-1. 配置文件： `config.toml`
+1. 配置文件 `config.toml`
 2. 执行 `buf generate` 编译proto文件
-3. 表 `t_demo` 对应 `ent/schema/demo.go`
+3. 执行 `go mod tidy` 下载依赖
 4. 执行 `ent/generate.go` 生成ORM代码 (只要 `ent/schema` 目录下有变动都需要执行)
 5. 执行 `go run main.go` 运行
 6. 执行 `go run main.go -h` 查看命令

cmd/internal/grpc/app/pkg_app_cmd_root.tmpl

@@ -3,6 +3,8 @@ package cmd
 import (
 	"context"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -50,9 +52,30 @@ func preInit(ctx context.Context) {
 	ent.Init(ctx)
 }
 
+// CleanUp 清理资源
+func CleanUp() {
+	// 关闭数据库连接
+	ent.Close()
+}
+
 func serving(ctx context.Context) {
+	go watchExit()
 	// serve grpc
 	go server.ServeGrpc(ctx)
 	// serve http
 	server.ServeHttp(ctx)
 }
+
+func watchExit() {
+	// 创建一个通道来监听信号
+	ch := make(chan os.Signal, 1)
+	// 监听特定的系统信号
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGTERM)
+	// 处理信号
+	sig := <-ch
+	log.Warn(context.TODO(), "Received Signal", zap.String("signal", sig.String()))
+	// 执行清理操作
+	CleanUp()
+	// 退出程序
+	os.Exit(0)
+}

cmd/internal/grpc/app/pkg_app_main.tmpl

@@ -1,17 +1,8 @@
 package main
 
-import (
-	"{{.Module}}/pkg/{{.AppPkg}}/cmd"
-	"{{.Module}}/pkg/{{.AppPkg}}/ent"
-)
+import "{{.Module}}/pkg/{{.AppPkg}}/cmd"
 
 func main() {
-	defer clean()
+	defer cmd.CleanUp()
 	cmd.Init()
 }
-
-// clean 清理资源
-func clean() {
-	// 关闭数据库连接
-	ent.Close()
-}

cmd/internal/grpc/app/pkg_app_server_http.tmpl

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
@@ -27,7 +28,21 @@ func ServeHttp(ctx context.Context) {
 	}
 	defer conn.Close()
 	// Create http mux with health check
-	mux := runtime.NewServeMux(runtime.WithHealthzEndpoint(grpc_health_v1.NewHealthClient(conn)))
+	mux := runtime.NewServeMux(
+		runtime.WithHealthzEndpoint(grpc_health_v1.NewHealthClient(conn)),
+		runtime.WithIncomingHeaderMatcher(func(s string) (string, bool) {
+			if v, ok := runtime.DefaultHeaderMatcher(s); ok {
+				return v, true
+			}
+			return strings.ToLower(s), true
+		}),
+		runtime.WithOutgoingHeaderMatcher(func(s string) (string, bool) {
+			if s == log.TraceId {
+				return s, true
+			}
+			return runtime.MetadataHeaderPrefix + s, true
+		}),
+	)
 	// Register http handler
 	if err = registerHttp(ctx, mux, conn); err != nil {
 		log.Fatal(ctx, "Error register http", zap.Error(err))
@@ -40,7 +55,7 @@ func ServeHttp(ctx context.Context) {
 		},
 		AllowedMethods:   []string{http.MethodGet, http.MethodPost, http.MethodPatch, http.MethodPut, http.MethodDelete, http.MethodOptions},
 		AllowedHeaders:   []string{"Authorization", "Content-Type", "withCredentials"},
-		ExposedHeaders:   []string{}, // 服务器暴露一些自定义的头信息，允许客户端访问
+		ExposedHeaders:   []string{log.TraceId}, // 服务器暴露一些自定义的头信息，允许客户端访问
 		AllowCredentials: true,
 	}).Handler(mux)
 	// Serve HTTP server

cmd/internal/grpc/pkg_lib_identity_identity.tmpl

@@ -14,10 +14,6 @@ import (
 	"{{.Module}}/pkg/lib/log"
 )
 
-type CtxKeyAuth int
-
-const IdentityKey CtxKeyAuth = 0
-
 // Identity 授权身份
 type Identity interface {
 	// ID 授权ID
@@ -77,17 +73,24 @@ func New(id int64, token string) Identity {
 	}
 }
 
+type identityKey struct{}
+
+// NewContext 将Identity注入context
+func NewContext(ctx context.Context, token string) context.Context {
+	return context.WithValue(ctx, identityKey{}, FromAuthToken(ctx, token))
+}
+
 // FromContext 获取授权信息
 func FromContext(ctx context.Context) Identity {
 	if ctx == nil {
 		return NewEmpty()
 	}
 
-	_identity, ok := ctx.Value(IdentityKey).(Identity)
+	id, ok := ctx.Value(identityKey{}).(Identity)
 	if !ok {
 		return NewEmpty()
 	}
-	return _identity
+	return id
 }
 
 // FromAuthToken 解析授权Token
@@ -105,10 +108,10 @@ func FromAuthToken(ctx context.Context, token string) Identity {
 		return NewEmpty()
 	}
 
-	_identity := NewEmpty()
-	if err = json.Unmarshal(plainText, _identity); err != nil {
+	id := NewEmpty()
+	if err = json.Unmarshal(plainText, id); err != nil {
 		log.Error(ctx, "Error json.Unmarshal AuthToken", zap.Error(err))
 		return NewEmpty()
 	}
-	return _identity
+	return id
 }

cmd/internal/grpc/pkg_lib_log_log.tmpl

@@ -4,24 +4,48 @@ import (
 	"context"
 
 	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 )
 
 func Info(ctx context.Context, msg string, fields ...zap.Field) {
-	logger.Info(msg, append(fields, zap.String("trace_id", GetTraceId(ctx)), zap.String("method", GetFullMethod(ctx)))...)
+	Log(ctx, zapcore.InfoLevel, msg, fields...)
 }
 
 func Warn(ctx context.Context, msg string, fields ...zap.Field) {
-	logger.Warn(msg, append(fields, zap.String("trace_id", GetTraceId(ctx)), zap.String("method", GetFullMethod(ctx)))...)
+	Log(ctx, zapcore.WarnLevel, msg, fields...)
 }
 
 func Error(ctx context.Context, msg string, fields ...zap.Field) {
-	logger.Error(msg, append(fields, zap.String("trace_id", GetTraceId(ctx)), zap.String("method", GetFullMethod(ctx)))...)
+	Log(ctx, zapcore.ErrorLevel, msg, fields...)
 }
 
 func Panic(ctx context.Context, msg string, fields ...zap.Field) {
-	logger.Panic(msg, append(fields, zap.String("trace_id", GetTraceId(ctx)), zap.String("method", GetFullMethod(ctx)))...)
+	Log(ctx, zapcore.PanicLevel, msg, fields...)
 }
 
 func Fatal(ctx context.Context, msg string, fields ...zap.Field) {
-	logger.Fatal(msg, append(fields, zap.String("trace_id", GetTraceId(ctx)), zap.String("method", GetFullMethod(ctx)))...)
+	Log(ctx, zapcore.FatalLevel, msg, fields...)
+}
+
+func Log(ctx context.Context, level zapcore.Level, msg string, fields ...zap.Field) {
+	traceId, fullMethod := GetTraceInfo(ctx)
+	fields = append(fields,
+		zap.String("hostname", hostname),
+		zap.String("trace_id", traceId),
+		zap.String("method", fullMethod),
+	)
+	switch level {
+	case zapcore.InfoLevel:
+		logger.Info(msg, fields...)
+	case zapcore.WarnLevel:
+		logger.Warn(msg, fields...)
+	case zapcore.ErrorLevel:
+		logger.Error(msg, fields...)
+	case zapcore.PanicLevel:
+		logger.Panic(msg, fields...)
+	case zapcore.FatalLevel:
+		logger.Fatal(msg, fields...)
+	default:
+		logger.Debug(msg, fields...)
+	}
 }

cmd/internal/grpc/pkg_lib_log_traceid.tmpl

@@ -8,41 +8,25 @@ import (
 	"os"
 	"strings"
 	"sync/atomic"
-)
 
-// Key to use when setting the trace ID.
-type CtxKeyTraceId int
+	"github.com/shenghui0779/yiigo/xhash"
+	"google.golang.org/grpc/metadata"
+)
 
-// TraceIdKey is the key that holds the unique trace ID in a trace context.
 const (
-	TraceIdKey  CtxKeyTraceId = 0
-	FullMethodKey CtxKeyTraceId = 1
+	TraceId     = "x-trace-id"
+	TraceMethod = "x-trace-method"
 )
 
 var (
-	traceId     uint64
-	TracePrefix string
+	hostname    string
+	tracePrefix string
+	traceSeq    uint64
 )
 
-// A quick note on the statistics here: we're trying to calculate the chance that
-// two randomly generated base62 prefixes will collide. We use the formula from
-// http://en.wikipedia.org/wiki/Birthday_problem
-//
-// P[m, n] \approx 1 - e^{-m^2/2n}
-//
-// We ballpark an upper bound for $m$ by imagining (for whatever reason) a server
-// that restarts every second over 10 years, for $m = 86400 * 365 * 10 = 315360000$
-//
-// For a $k$ character base-62 identifier, we have $n(k) = 62^k$
-//
-// Plugging this in, we find $P[m, n(10)] \approx 5.75%$, which is good enough for
-// our purposes, and is surely more than anyone would ever need in practice -- a
-// process that is rebooted a handful of times a day for a hundred years has less
-// than a millionth of a percent chance of generating two colliding IDs.
-
 func init() {
-	hostname, err := os.Hostname()
-	if hostname == "" || err != nil {
+	hostname, _ = os.Hostname()
+	if len(hostname) == 0 {
 		hostname = "localhost"
 	}
 
@@ -51,38 +35,31 @@ func init() {
 		b64 string
 	)
 	for len(b64) < 10 {
-		rand.Read(buf[:])
+		_, _ = rand.Read(buf[:])
 		b64 = base64.StdEncoding.EncodeToString(buf[:])
 		b64 = strings.NewReplacer("+", "", "/", "").Replace(b64)
 	}
-	TracePrefix = fmt.Sprintf("%s/%s", hostname, b64[0:10])
+	tracePrefix = fmt.Sprintf("%s/%s", hostname, b64)
 }
 
-// GetTraceId returns a trace ID from the given context if one is present.
-// Returns the empty string if a trace ID cannot be found.
-func GetTraceId(ctx context.Context) string {
-	if ctx == nil {
-		return "-"
-	}
-	if v, ok := ctx.Value(TraceIdKey).(string); ok {
-		return v
-	}
-	return "-"
+// NewTraceId generates a new trace ID in the sequence.
+func NewTraceId() string {
+	seq := atomic.AddUint64(&traceSeq, 1)
+	return xhash.MD5(fmt.Sprintf("%s-%d", tracePrefix, seq))
 }
 
-// GetFullMethod returns a full method from the given context if one is present.
-// Returns the empty string if a full method cannot be found.
-func GetFullMethod(ctx context.Context) string {
-	if ctx == nil {
-		return "-"
+func GetTraceInfo(ctx context.Context) (traceId, fullMethod string) {
+	traceId = "-"
+	fullMethod = "-"
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return
 	}
-	if method, ok := ctx.Value(FullMethodKey).(string); ok {
-		return method
+	if v := md.Get(TraceId); len(v) != 0 {
+		traceId = v[0]
 	}
-	return "-"
-}
-
-// NextTraceId generates the next trace ID in the sequence.
-func NextTraceId() uint64 {
-	return atomic.AddUint64(&traceId, 1)
+	if v := md.Get(TraceMethod); len(v) != 0 {
+		fullMethod = v[0]
+	}
+	return
 }

cmd/internal/grpc/pkg_lib_middleware_log.tmpl

@@ -10,9 +10,14 @@ import (
 	"{{.Module}}/pkg/lib/log"
 )
 
-func Log(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
+const HealthCheckMethod = "/grpc.health.v1.Health/Check"
+
+func Log(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
 	now := time.Now().Local()
 	defer func() {
+		if info.FullMethod == HealthCheckMethod {
+			return
+		}
 		log.Info(ctx, "Request info",
 			zap.Any("request", req),
 			zap.Any("response", resp),

cmd/internal/grpc/pkg_lib_middleware_traceid.tmpl

@@ -2,18 +2,24 @@ package middleware
 
 import (
 	"context"
-	"fmt"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	"{{.Module}}/pkg/lib/log"
 )
 
 // TraceId is a middleware that injects a trace ID into the context of each request.
 func TraceId(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-	if v, ok := ctx.Value(log.TraceIdKey).(string); !ok || len(v) == 0 {
-		ctx = context.WithValue(ctx, log.TraceIdKey, fmt.Sprintf("%s-%06d", log.TracePrefix, log.NextTraceId()))
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		md = metadata.Pairs()
 	}
-	ctx = context.WithValue(ctx, log.FullMethodKey, info.FullMethod)
-	return handler(ctx, req)
+	if v := md.Get(log.TraceId); len(v) == 0 {
+		md.Set(log.TraceId, log.NewTraceId())
+	}
+	md.Set(log.TraceMethod, info.FullMethod)
+	// set the response header
+	_ = grpc.SetHeader(ctx, metadata.Pairs(log.TraceId, md.Get(log.TraceId)[0]))
+	return handler(metadata.NewIncomingContext(ctx, md), req)
 }