both proxy and admin log do not need pass in correlation id explicitly

bricks-cloud · May 21, 2024 · f525612 · f525612
1 parent 681574e
commit f525612
Show file tree

Hide file tree

Showing 31 changed files with 438 additions and 536 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,4 @@
 release_notes.md
 target
-.DS_STORE
+.DS_STORE
+.vscode/launch.json
diff --git a/internal/server/web/admin/admin.go b/internal/server/web/admin/admin.go
diff --git a/internal/server/web/admin/event.go b/internal/server/web/admin/event.go
@@ -5,12 +5,13 @@ import (
 	"time"
 
 	"github.com/bricks-cloud/bricksllm/internal/stats"
+	"github.com/bricks-cloud/bricksllm/internal/util"
 	"github.com/gin-gonic/gin"
-	"go.uber.org/zap"
 )
 
-func getGetUserIdsHandler(m KeyReportingManager, log *zap.Logger, prod bool) gin.HandlerFunc {
+func getGetUserIdsHandler(m KeyReportingManager, prod bool) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		log := util.GetLogFromCtx(c)
 		stats.Incr("bricksllm.admin.get_get_user_ids_handler.requests", nil, 1)
 
 		start := time.Now()
@@ -43,12 +44,11 @@ func getGetUserIdsHandler(m KeyReportingManager, log *zap.Logger, prod bool) gin
 			return
 		}
 
-		cid := c.GetString(correlationId)
 		cids, err := m.GetUserIds(kid)
 		if err != nil {
 			stats.Incr("bricksllm.admin.get_get_user_ids_handler.get_user_ids_err", nil, 1)
 
-			logError(log, "error when getting userIds", prod, cid, err)
+			logError(log, "error when getting userIds", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/key-reporting-manager",
 				Title:    "getting user ids error",
@@ -64,8 +64,9 @@ func getGetUserIdsHandler(m KeyReportingManager, log *zap.Logger, prod bool) gin
 	}
 }
 
-func getGetCustomIdsHandler(m KeyReportingManager, log *zap.Logger, prod bool) gin.HandlerFunc {
+func getGetCustomIdsHandler(m KeyReportingManager, prod bool) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		log := util.GetLogFromCtx(c)
 		stats.Incr("bricksllm.admin.get_get_custom_ids_handler.requests", nil, 1)
 
 		start := time.Now()
@@ -99,13 +100,11 @@ func getGetCustomIdsHandler(m KeyReportingManager, log *zap.Logger, prod bool) g
 			return
 		}
 
-		cid := c.GetString(correlationId)
-
 		cids, err := m.GetCustomIds(kid)
 		if err != nil {
 			stats.Incr("bricksllm.admin.get_get_user_ids_handler.get_custom_ids_err", nil, 1)
 
-			logError(log, "error when getting custom ids", prod, cid, err)
+			logError(log, "error when getting custom ids", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/key-reporting-manager",
 				Title:    "getting custom ids error",

diff --git a/internal/server/web/admin/events.go b/internal/server/web/admin/events.go
@@ -8,12 +8,13 @@ import (
 
 	"github.com/bricks-cloud/bricksllm/internal/event"
 	"github.com/bricks-cloud/bricksllm/internal/stats"
+	"github.com/bricks-cloud/bricksllm/internal/util"
 	"github.com/gin-gonic/gin"
-	"go.uber.org/zap"
 )
 
-func getGetEventsV2Handler(m KeyReportingManager, log *zap.Logger, prod bool) gin.HandlerFunc {
+func getGetEventsV2Handler(m KeyReportingManager, prod bool) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		log := util.GetLogFromCtx(c)
 		stats.Incr("bricksllm.admin.get_get_events_v2_handler.requests", nil, 1)
 
 		start := time.Now()
@@ -34,10 +35,9 @@ func getGetEventsV2Handler(m KeyReportingManager, log *zap.Logger, prod bool) gi
 			return
 		}
 
-		cid := c.GetString(correlationId)
 		data, err := io.ReadAll(c.Request.Body)
 		if err != nil {
-			logError(log, "error when reading get events request body", prod, cid, err)
+			logError(log, "error when reading get events request body", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/request-body-read",
 				Title:    "get events request body reader error",
@@ -51,7 +51,7 @@ func getGetEventsV2Handler(m KeyReportingManager, log *zap.Logger, prod bool) gi
 		request := &event.EventRequest{}
 		err = json.Unmarshal(data, request)
 		if err != nil {
-			logError(log, "error when unmarshalling get events request body", prod, cid, err)
+			logError(log, "error when unmarshalling get events request body", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/json-unmarshal",
 				Title:    "json unmarshaller error",
@@ -84,7 +84,7 @@ func getGetEventsV2Handler(m KeyReportingManager, log *zap.Logger, prod bool) gi
 				return
 			}
 
-			logError(log, "error when getting events", prod, cid, err)
+			logError(log, "error when getting events", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/event-manager",
 				Title:    "getting events errored out",

diff --git a/internal/server/web/admin/middleware.go b/internal/server/web/admin/middleware.go
@@ -16,17 +16,20 @@ func getAdminLoggerMiddleware(log *zap.Logger, prefix string, prod bool, adminPa
 			return
 		}
 
-		c.Set(correlationId, util.NewUuid())
+		cid := util.NewUuid()
+		c.Set(util.STRING_CORRELATION_ID, cid)
+		logWithCid := log.With(zap.String(util.STRING_CORRELATION_ID, cid))
+		util.SetLogToCtx(c, logWithCid)
+
 		start := time.Now()
 		c.Next()
 		latency := time.Since(start).Milliseconds()
 		if !prod {
-			log.Sugar().Infof("%s | %d | %s | %s | %dms", prefix, c.Writer.Status(), c.Request.Method, c.FullPath(), latency)
+			logWithCid.Sugar().Infof("%s | %d | %s | %s | %dms", prefix, c.Writer.Status(), c.Request.Method, c.FullPath(), latency)
 		}
 
 		if prod {
-			log.Info("request to admin management api",
-				zap.String(correlationId, c.GetString(correlationId)),
+			logWithCid.Info("request to admin management api",
 				zap.Int("code", c.Writer.Status()),
 				zap.String("method", c.Request.Method),
 				zap.String("path", c.FullPath()),

diff --git a/internal/server/web/admin/policy.go b/internal/server/web/admin/policy.go
@@ -8,12 +8,13 @@ import (
 
 	"github.com/bricks-cloud/bricksllm/internal/policy"
 	"github.com/bricks-cloud/bricksllm/internal/stats"
+	"github.com/bricks-cloud/bricksllm/internal/util"
 	"github.com/gin-gonic/gin"
-	"go.uber.org/zap"
 )
 
-func getCreatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.HandlerFunc {
+func getCreatePolicyHandler(pm PoliciesManager, prod bool) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		log := util.GetLogFromCtx(c)
 		stats.Incr("bricksllm.admin.get_get_create_policy_handler.requests", nil, 1)
 
 		start := time.Now()
@@ -34,10 +35,9 @@ func getCreatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 			return
 		}
 
-		cid := c.GetString(correlationId)
 		data, err := io.ReadAll(c.Request.Body)
 		if err != nil {
-			logError(log, "error when reading policy creation request body", prod, cid, err)
+			logError(log, "error when reading policy creation request body", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/request-body-read",
 				Title:    "request body reader error",
@@ -51,7 +51,7 @@ func getCreatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 		p := &policy.Policy{}
 		err = json.Unmarshal(data, p)
 		if err != nil {
-			logError(log, "error when unmarshalling policy creation request body", prod, cid, err)
+			logError(log, "error when unmarshalling policy creation request body", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/json-unmarshal",
 				Title:    "json unmarshaller error",
@@ -66,7 +66,7 @@ func getCreatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 		if err != nil {
 			stats.Incr("bricksllm.admin.get_get_create_policy_handler.creat_policy_error", nil, 1)
 
-			logError(log, "error when creating a policy", prod, cid, err)
+			logError(log, "error when creating a policy", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/policies/creation",
 				Title:    "policy creation failed",
@@ -83,8 +83,9 @@ func getCreatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 	}
 }
 
-func getUpdatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.HandlerFunc {
+func getUpdatePolicyHandler(pm PoliciesManager, prod bool) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		log := util.GetLogFromCtx(c)
 		stats.Incr("bricksllm.admin.get_update_policy_handler.requests", nil, 1)
 
 		start := time.Now()
@@ -118,10 +119,9 @@ func getUpdatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 			return
 		}
 
-		cid := c.GetString(correlationId)
 		data, err := io.ReadAll(c.Request.Body)
 		if err != nil {
-			logError(log, "error when reading policy creation request body", prod, cid, err)
+			logError(log, "error when reading policy creation request body", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/request-body-read",
 				Title:    "request body reader error",
@@ -135,7 +135,7 @@ func getUpdatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 		p := &policy.UpdatePolicy{}
 		err = json.Unmarshal(data, p)
 		if err != nil {
-			logError(log, "error when unmarshalling policy creation request body", prod, cid, err)
+			logError(log, "error when unmarshalling policy creation request body", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/json-unmarshal",
 				Title:    "json unmarshaller error",
@@ -150,7 +150,7 @@ func getUpdatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 		if err != nil {
 			stats.Incr("bricksllm.admin.get_update_policy_handler.update_policy_error", nil, 1)
 
-			logError(log, "error when updating a policy by id", prod, cid, err)
+			logError(log, "error when updating a policy by id", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/policie/updates",
 				Title:    "update a policy failed",
@@ -167,8 +167,9 @@ func getUpdatePolicyHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.
 	}
 }
 
-func getGetPoliciesByTagsHandler(pm PoliciesManager, log *zap.Logger, prod bool) gin.HandlerFunc {
+func getGetPoliciesByTagsHandler(pm PoliciesManager, prod bool) gin.HandlerFunc {
 	return func(c *gin.Context) {
+		log := util.GetLogFromCtx(c)
 		stats.Incr("bricksllm.admin.get_get_policies_by_tags_handler.requests", nil, 1)
 
 		start := time.Now()
@@ -202,7 +203,6 @@ func getGetPoliciesByTagsHandler(pm PoliciesManager, log *zap.Logger, prod bool)
 			return
 		}
 
-		cid := c.GetString(correlationId)
 		policies, err := pm.GetPoliciesByTags(c.QueryArray("tags"))
 		if err != nil {
 			errType := "internal"
@@ -213,7 +213,7 @@ func getGetPoliciesByTagsHandler(pm PoliciesManager, log *zap.Logger, prod bool)
 				}, 1)
 			}()
 
-			logError(log, "error when getting policies by tags", prod, cid, err)
+			logError(log, "error when getting policies by tags", prod, err)
 			c.JSON(http.StatusInternalServerError, &ErrorResponse{
 				Type:     "/errors/policies",
 				Title:    "get policies by tags failed",