From 5ee1cea0d795721de3d0e6226d74e112c914dc6f Mon Sep 17 00:00:00 2001
From: yoge <yogecoder@gmail.com>
Date: Fri, 26 Dec 2025 17:11:59 +0800
Subject: [PATCH] feat: add gls

---
 go.mod                      |  2 ++
 go.sum                      |  2 ++
 pkg/log/logger.go           | 12 +++++++++---
 pkg/middleware/requestid.go | 18 +++++++++++-------
 pkg/requestid/requestid.go  | 27 +++++++++++++++++++++++++++
 5 files changed, 51 insertions(+), 10 deletions(-)
 create mode 100644 pkg/requestid/requestid.go

diff --git a/go.mod b/go.mod
index b027928..6fb2350 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/gin-gonic/gin v1.10.0
 	github.com/google/uuid v1.6.0
+	github.com/jtolds/gls v4.20.0+incompatible
 	github.com/redis/go-redis/v9 v9.7.0
 	github.com/rs/zerolog v1.33.0
 	github.com/spf13/viper v1.19.0
@@ -43,6 +44,7 @@ require (
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/go-playground/validator/v10 v10.20.0 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/gopherjs/gopherjs v0.0.0-20200217142428-fce0ec30dd00 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
diff --git a/go.sum b/go.sum
index 2af57c4..57e0b0b 100644
--- a/go.sum
+++ b/go.sum
@@ -79,6 +79,7 @@ github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
+github.com/gopherjs/gopherjs v0.0.0-20200217142428-fce0ec30dd00 h1:l5lAOZEym3oK3SQ2HBHWsJUfbNBiTXJDeW2QDxw9AQ0=
 github.com/gopherjs/gopherjs v0.0.0-20200217142428-fce0ec30dd00/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
@@ -89,6 +90,7 @@ github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/
 github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
+github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM=
diff --git a/pkg/log/logger.go b/pkg/log/logger.go
index fba9a3a..6bfc3f3 100644
--- a/pkg/log/logger.go
+++ b/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...)
 }
-
diff --git a/pkg/middleware/requestid.go b/pkg/middleware/requestid.go
index 242ddf6..0f060b1 100644
--- a/pkg/middleware/requestid.go
+++ b/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()
+		})
 	}
 }
-
diff --git a/pkg/requestid/requestid.go b/pkg/requestid/requestid.go
new file mode 100644
index 0000000..f195eae
--- /dev/null
+++ b/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
+}
+