feat: add gls

pkg/log/logger.go

@@ -8,6 +8,8 @@ import (
 	"runtime"
 	"time"
 
+	"gitea.com/bitwsd/document_ai/pkg/requestid"
+
 	"github.com/rs/zerolog"
 	"gopkg.in/natefinch/lumberjack.v2"
 )
@@ -67,8 +69,13 @@ func log(ctx context.Context, level zerolog.Level, logType LogType, kv ...interf
 	// 添加日志类型
 	event.Str("type", string(logType))
 
-	// 添加请求ID
-	if reqID, exists := ctx.Value("request_id").(string); exists {
+	reqID := requestid.GetRequestID()
+	if reqID == "" {
+		if id, exists := ctx.Value("request_id").(string); exists {
+			reqID = id
+		}
+	}
+	if reqID != "" {
 		event.Str("request_id", reqID)
 	}
 
@@ -149,4 +156,3 @@ func Fatal(ctx context.Context, kv ...interface{}) {
 func Access(ctx context.Context, kv ...interface{}) {
 	log(ctx, zerolog.InfoLevel, TypeAccess, kv...)
 }
-

pkg/middleware/requestid.go

@@ -1,19 +1,23 @@
 package middleware
 
 import (
+	"gitea.com/bitwsd/document_ai/pkg/requestid"
+
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 )
 
 func RequestID() gin.HandlerFunc {
 	return func(c *gin.Context) {
-		requestID := c.Request.Header.Get("X-Request-ID")
-		if requestID == "" {
-			requestID = uuid.New().String()
+		reqID := c.Request.Header.Get("X-Request-ID")
+		if reqID == "" {
+			reqID = uuid.New().String()
 		}
-		c.Request.Header.Set("X-Request-ID", requestID)
-		c.Set("request_id", requestID)
-		c.Next()
+		c.Request.Header.Set("X-Request-ID", reqID)
+		c.Set("request_id", reqID)
+
+		requestid.SetRequestID(reqID, func() {
+			c.Next()
+		})
 	}
 }
-

pkg/requestid/requestid.go

@@ -0,0 +1,27 @@
+package requestid
+
+import (
+	"github.com/jtolds/gls"
+)
+
+// requestIDKey 是 gls 中存储 request_id 的 key
+var requestIDKey = gls.GenSym()
+
+// glsMgr 是 gls 管理器
+var glsMgr = gls.NewContextManager()
+
+// SetRequestID 在 gls 中设置 request_id，并在 fn 执行期间保持有效
+func SetRequestID(requestID string, fn func()) {
+	glsMgr.SetValues(gls.Values{requestIDKey: requestID}, fn)
+}
+
+// GetRequestID 从 gls 中获取当前 goroutine 的 request_id
+func GetRequestID() string {
+	val, ok := glsMgr.GetValue(requestIDKey)
+	if !ok {
+		return ""
+	}
+	reqID, _ := val.(string)
+	return reqID
+}
+