package log import ( "bytes" "context" "io" "log/slog" "os" "strings" "sync" "sync/atomic" "github.com/fatih/color" "github.com/tidwall/pretty" "github.com/tigorlazuardi/bluemage/go/pkg/caller" ) type PrettyHandler struct { opts *slog.HandlerOptions output WriteSyncer replaceAttr func(groups []string, attr slog.Attr) slog.Attr withAttrs []slog.Attr withGroup []string isFlushing atomic.Bool flushingChan chan struct{} wg sync.WaitGroup } // NewPrettyHandler creates a human friendly readable logs. func NewPrettyHandler(writer WriteSyncer, opts *slog.HandlerOptions) *PrettyHandler { if opts == nil { opts = &slog.HandlerOptions{Level: slog.LevelDebug} } if opts.ReplaceAttr == nil { opts.ReplaceAttr = func(groups []string, attr slog.Attr) slog.Attr { return attr } } return &PrettyHandler{ opts: opts, output: writer, replaceAttr: func(groups []string, attr slog.Attr) slog.Attr { if len(groups) > 0 { return opts.ReplaceAttr(groups, attr) } switch attr.Key { case slog.TimeKey, slog.LevelKey, slog.SourceKey, slog.MessageKey: return slog.Attr{} default: return opts.ReplaceAttr(groups, attr) } }, flushingChan: make(chan struct{}), } } // Flush waits until all logs are written to the underlying writer // and then flush the buffer. // // Flush blocks other log writes until it's done. func (pr *PrettyHandler) Flush() error { pr.isFlushing.Store(true) defer func() { pr.isFlushing.Store(false) }() pr.wg.Wait() err := pr.output.Sync() for { select { // signal to handlers that flushing is done. case pr.flushingChan <- struct{}{}: default: return err } } } // Sync attemps to flush the buffer of the underlying writer. func (pr *PrettyHandler) Sync() error { return pr.output.Sync() } // Enabled implements slog.Handler interface. func (pr *PrettyHandler) Enabled(ctx context.Context, lvl slog.Level) bool { return pr.opts.Level.Level() <= lvl } var bufferPool = sync.Pool{ New: func() interface{} { buf := &bytes.Buffer{} buf.Grow(1024) return buf }, } func putBuffer(buf *bytes.Buffer) { const limit = 1024 * 512 // 512KB if buf.Cap() < limit { buf.Reset() bufferPool.Put(buf) } } var jsonColorStyle = &pretty.Style{ Key: [2]string{"\x1B[95m", "\x1B[0m"}, String: [2]string{"\x1B[32m", "\x1B[0m"}, Number: [2]string{"\x1B[33m", "\x1B[0m"}, True: [2]string{"\x1B[36m", "\x1B[0m"}, False: [2]string{"\x1B[36m", "\x1B[0m"}, Null: [2]string{"\x1B[2m", "\x1B[0m"}, Escape: [2]string{"\x1B[35m", "\x1B[0m"}, Brackets: [2]string{"\x1B[0m", "\x1B[0m"}, Append: pretty.TerminalStyle.Append, } var jsonPrettyOpts = &pretty.Options{ Width: 80, Prefix: "", Indent: " ", SortKeys: false, } // Handle implements slog.Handler interface. func (pr *PrettyHandler) Handle(ctx context.Context, record slog.Record) error { if pr.isFlushing.Load() { // pr is currently flushing. Wait until the flushing process is done. <-pr.flushingChan } pr.wg.Add(1) defer pr.wg.Done() var levelColor *color.Color switch { case record.Level >= slog.LevelError: levelColor = color.New(color.FgRed) case record.Level >= slog.LevelWarn: levelColor = color.New(color.FgYellow) case record.Level >= slog.LevelInfo: levelColor = color.New(color.FgGreen) default: levelColor = color.New(color.FgWhite) } buf := bufferPool.Get().(*bytes.Buffer) jsonBuf := bufferPool.Get().(*bytes.Buffer) defer putBuffer(buf) defer putBuffer(jsonBuf) if record.PC != 0 && pr.opts.AddSource { frame := caller.From(record.PC).Frame levelColor.Fprint(buf, frame.File) levelColor.Fprint(buf, ":") levelColor.Fprint(buf, frame.Line) levelColor.Fprint(buf, " -- ") split := strings.Split(frame.Function, string(os.PathSeparator)) fnName := split[len(split)-1] levelColor.Fprint(buf, fnName) buf.WriteByte('\n') } if !record.Time.IsZero() { buf.WriteString(record.Time.Format("[2006-01-02 15:04:05] ")) } buf.WriteByte('[') levelColor.Add(color.Bold).Fprint(buf, record.Level.String()) buf.WriteString("] ") if record.Message != "" { buf.WriteString(record.Message) } buf.WriteByte('\n') serializer := pr.createSerializer(jsonBuf) _ = serializer.Handle(ctx, record) if jsonBuf.Len() > 3 { // Ignore empty json like "{}\n" jsonData := jsonBuf.Bytes() jsonData = pretty.PrettyOptions(jsonData, jsonPrettyOpts) jsonData = pretty.Color(jsonData, jsonColorStyle) buf.Write(jsonData) } buf.WriteByte('\n') _, err := buf.WriteTo(pr.output) return err } func (pr *PrettyHandler) createSerializer(w io.Writer) slog.Handler { var jsonHandler slog.Handler = slog.NewJSONHandler(w, &slog.HandlerOptions{ Level: slog.LevelDebug, ReplaceAttr: pr.replaceAttr, }) if len(pr.withAttrs) > 0 { jsonHandler = jsonHandler.WithAttrs(pr.withAttrs) } if len(pr.withGroup) > 0 { for _, group := range pr.withGroup { jsonHandler = jsonHandler.WithGroup(group) } } return jsonHandler } func (pr *PrettyHandler) clone() *PrettyHandler { return &PrettyHandler{ opts: pr.opts, output: pr.output, replaceAttr: pr.replaceAttr, withAttrs: pr.withAttrs, withGroup: pr.withGroup, } } // WithAttrs implements slog.Handler interface. func (pr *PrettyHandler) WithAttrs(attrs []slog.Attr) slog.Handler { p := pr.clone() p.withAttrs = append(p.withAttrs, attrs...) return p } // WithGroup implements slog.Handler interface. func (pr *PrettyHandler) WithGroup(name string) slog.Handler { p := pr.clone() p.withGroup = append(p.withGroup, name) return p }