Bluemage/go/pkg/log/pretty.go

230 lines
5.5 KiB
Go

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 WriteLocker
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 WriteLocker, 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() {
pr.isFlushing.Store(true)
defer func() { pr.isFlushing.Store(false) }()
pr.wg.Wait()
for {
select {
// signal to handlers that flushing is done.
case pr.flushingChan <- struct{}{}:
default:
return
}
}
}
// Sync attemps to flush the buffer of the underlying writer.
func (pr *PrettyHandler) Close() error {
pr.Flush()
return nil
}
// 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
cwd, _ := os.Getwd()
f := strings.TrimPrefix(frame.File, cwd)
f = strings.TrimPrefix(f, string(os.PathSeparator))
levelColor.Fprint(buf, f)
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
}