diff --git a/internal/infra/logs/encoder.go b/internal/infra/logs/encoder.go new file mode 100644 index 0000000..b14a771 --- /dev/null +++ b/internal/infra/logs/encoder.go @@ -0,0 +1,102 @@ +package logs + +import ( + "bytes" + + "go.uber.org/zap/buffer" + "go.uber.org/zap/zapcore" +) + +// traceKey 是日志中用于调用链的字段名，必须与 GetLogger 中使用的保持一致。 +const traceKey = "逻辑调用链" + +// coloredConsoleEncoder 是一个自定义的 zapcore.Encoder，它实现了两个核心功能： +// 1. 为控制台输出的日志行根据级别添加 ANSI 颜色。 +// 2. 确保代表“调用链”的结构化字段总是出现在日志行的末尾。 +type coloredConsoleEncoder struct { + zapcore.Encoder // 嵌入一个标准的 ConsoleEncoder，复用其大部分能力。 +} + +// NewColoredConsoleEncoder 创建一个新的 coloredConsoleEncoder 实例。 +func NewColoredConsoleEncoder(cfg zapcore.EncoderConfig) zapcore.Encoder { + return &coloredConsoleEncoder{ + Encoder: zapcore.NewConsoleEncoder(cfg), + } +} + +// EncodeEntry 重写了核心的编码方法。 +func (c *coloredConsoleEncoder) EncodeEntry(entry zapcore.Entry, fields []zapcore.Field) (*buffer.Buffer, error) { + // 1. 从所有字段中分离出“调用链”字段和其他字段 + var traceField zapcore.Field + otherFields := make([]zapcore.Field, 0, len(fields)) + for _, field := range fields { + if field.Key == traceKey { + traceField = field + } else { + otherFields = append(otherFields, field) + } + } + + // 2. 使用内嵌的 ConsoleEncoder 先编码日志条目和“其他”字段 + // 这会生成不包含调用链的日志主体部分。 + line, err := c.Encoder.EncodeEntry(entry, otherFields) + if err != nil { + return nil, err + } + + // 3. 如果存在“调用链”字段，则将其手动编码并追加到末尾 + if traceField.Key != "" { + // 为了保持格式一致，我们创建一个临时的 Encoder 来只编码这一个字段 + // 注意：这里我们不能直接使用 c.Encoder，因为它会添加不必要的前缀（如时间、级别等） + tempEncoder := zapcore.NewConsoleEncoder(zapcore.EncoderConfig{}) + traceLine, err := tempEncoder.EncodeEntry(zapcore.Entry{}, []zapcore.Field{traceField}) + if err != nil { + return nil, err // 理论上不应失败 + } + + // 将编码后的调用链字段（例如 "逻辑调用链=xxx->yyy"）追加到日志主体后 + // TrimSpace 用于移除 tempEncoder 可能产生的额外换行符 + line.AppendString("\t") // 使用制表符分隔 + line.AppendString(string(bytes.TrimSpace(traceLine.Bytes()))) + } + + // 4. 为最终的日志行（不包括调用链）添加颜色 + var color string + switch entry.Level { + case zapcore.DebugLevel: + color = blue + case zapcore.InfoLevel: + color = green + case zapcore.WarnLevel: + color = yellow + case zapcore.ErrorLevel: + color = red + case zapcore.DPanicLevel, zapcore.PanicLevel, zapcore.FatalLevel: + color = bold + red + default: + color = reset + } + + // 创建一个新的 buffer，将颜色、日志内容和重置代码包裹起来 + finalBuf := buffer.NewPool().Get() + finalBuf.AppendString(color) + finalBuf.Write(line.Bytes()) // 写入已包含调用链的日志行 + finalBuf.AppendString(reset) + + // 如果原始日志行末尾没有换行符，则添加一个 + if line.Bytes()[line.Len()-1] != '\n' { + finalBuf.AppendByte('\n') + } + + line.Free() // 释放由 c.Encoder.EncodeEntry 创建的 buffer + + return finalBuf, nil +} + +// Clone 必须被重写，以确保克隆时返回的是我们自定义的 encoder 类型。 +func (c *coloredConsoleEncoder) Clone() zapcore.Encoder { + // 克隆内嵌的 Encoder，并用它来创建一个新的 coloredConsoleEncoder + return &coloredConsoleEncoder{ + Encoder: c.Encoder.Clone(), + } +} diff --git a/internal/infra/logs/logs.go b/internal/infra/logs/logs.go index c206bec..d60b05d 100644 --- a/internal/infra/logs/logs.go +++ b/internal/infra/logs/logs.go @@ -14,7 +14,6 @@ import ( "git.huangwc.com/pig/pig-farm-controller/internal/infra/config" "go.uber.org/zap" - "go.uber.org/zap/buffer" "go.uber.org/zap/zapcore" "gopkg.in/natefinch/lumberjack.v2" gormlogger "gorm.io/gorm/logger" @@ -97,8 +96,8 @@ func GetLogger(ctx context.Context) *Logger { return defaultLogger } - // 使用 With 方法创建带有 "trace" 字段的 Logger 副本 - newSugaredLogger := defaultLogger.With("trace", strings.Join(chain, "->")) + // 使用 With 方法创建带有 traceKey 字段的 Logger 副本 + newSugaredLogger := defaultLogger.With(traceKey, strings.Join(chain, "->")) return &Logger{newSugaredLogger} } @@ -111,6 +110,14 @@ func Trace(upstreamCtx context.Context, selfCtx context.Context, funcName string return newCtx, logger } +// TraceLogger 类似于 Trace，但只返回一个包含完整调用链的 Logger 实例， +// 而不返回更新后的 Context。适用于调用链的末端或不需要向下传递 Context 的场景。 +func TraceLogger(upstreamCtx context.Context, selfCtx context.Context, funcName string) *Logger { + newCtx := AddFuncName(upstreamCtx, selfCtx, funcName) + logger := GetLogger(newCtx) + return logger +} + // GetEncoder 根据指定的格式返回一个 zapcore.Encoder。 func GetEncoder(format string) zapcore.Encoder { encoderConfig := zap.NewProductionEncoderConfig() @@ -124,63 +131,6 @@ func GetEncoder(format string) zapcore.Encoder { return NewColoredConsoleEncoder(encoderConfig) } -// coloredConsoleEncoder 是一个自定义的 zapcore.Encoder，用于为整个日志行添加颜色。 -type coloredConsoleEncoder struct { - zapcore.Encoder // 嵌入默认的 ConsoleEncoder，以便委托其大部分功能 - cfg zapcore.EncoderConfig -} - -// NewColoredConsoleEncoder 创建一个新的 coloredConsoleEncoder 实例。 -func NewColoredConsoleEncoder(cfg zapcore.EncoderConfig) zapcore.Encoder { - // 创建一个标准的 ConsoleEncoder 作为基础 - baseEncoder := zapcore.NewConsoleEncoder(cfg) - return &coloredConsoleEncoder{ - Encoder: baseEncoder, - cfg: cfg, - } -} - -// EncodeEntry 重写 EncodeEntry 方法，在原始日志输出前后添加颜色代码。 -func (c *coloredConsoleEncoder) EncodeEntry(entry zapcore.Entry, fields []zapcore.Field) (*buffer.Buffer, error) { - // 首先，让嵌入的默认编码器生成原始的日志行 - buf, err := c.Encoder.EncodeEntry(entry, fields) - if err != nil { - return nil, err - } - - // 根据日志级别确定颜色 - var color string - switch entry.Level { - case zapcore.DebugLevel: - color = blue - case zapcore.InfoLevel: - color = green - case zapcore.WarnLevel: - color = yellow - case zapcore.ErrorLevel: - color = red - case zapcore.DPanicLevel, zapcore.PanicLevel, zapcore.FatalLevel: - color = bold + red - default: - color = reset - } - - // 创建一个新的 buffer 来存储带颜色的日志行 - coloredBuf := buffer.NewPool().Get() - coloredBuf.AppendString(color) - coloredBuf.Write(buf.Bytes()) - coloredBuf.AppendString(reset) - - buf.Free() // 释放原始 buffer 回池中 - - return coloredBuf, nil -} - -// Clone 方法也需要重写，以确保返回一个新的 coloredConsoleEncoder 实例。 -func (c *coloredConsoleEncoder) Clone() zapcore.Encoder { - return NewColoredConsoleEncoder(c.cfg) -} - // getWriteSyncer 根据配置创建日志写入目标。 func getWriteSyncer(cfg config.LogConfig) zapcore.WriteSyncer { writers := []zapcore.WriteSyncer{os.Stdout}