github.com/oarkflow/log@v1.0.78/slog.go (about)

     1  //go:build go1.21
     2  // +build go1.21
     3  
     4  package log
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  	"log/slog"
    10  	"strconv"
    11  	"time"
    12  )
    13  
    14  func slogJSONAttrEval(e *Entry, a slog.Attr) *Entry {
    15  	if a.Equal(slog.Attr{}) {
    16  		return e
    17  	}
    18  	value := a.Value.Resolve()
    19  	switch value.Kind() {
    20  	case slog.KindBool:
    21  		return e.Bool(a.Key, value.Bool())
    22  	case slog.KindInt64:
    23  		return e.Int64(a.Key, value.Int64())
    24  	case slog.KindUint64:
    25  		return e.Uint64(a.Key, value.Uint64())
    26  	case slog.KindFloat64:
    27  		return e.Float64(a.Key, value.Float64())
    28  	case slog.KindString:
    29  		return e.Str(a.Key, value.String())
    30  	case slog.KindTime:
    31  		return e.TimeFormat(a.Key, time.RFC3339Nano, value.Time())
    32  	case slog.KindDuration:
    33  		return e.Int64(a.Key, int64(value.Duration()))
    34  	case slog.KindGroup:
    35  		if len(value.Group()) == 0 {
    36  			return e
    37  		}
    38  		if a.Key == "" {
    39  			for _, attr := range value.Group() {
    40  				e = slogJSONAttrEval(e, attr)
    41  			}
    42  			return e
    43  		}
    44  		e.buf = append(e.buf, ',', '"')
    45  		e.buf = append(e.buf, a.Key...)
    46  		e.buf = append(e.buf, '"', ':')
    47  		i := len(e.buf)
    48  		for _, attr := range value.Group() {
    49  			e = slogJSONAttrEval(e, attr)
    50  		}
    51  		e.buf[i] = '{'
    52  		e.buf = append(e.buf, '}')
    53  		return e
    54  	case slog.KindAny:
    55  		return e.Any(a.Key, value.Any())
    56  	default:
    57  		return e.Any(a.Key, value.Any())
    58  	}
    59  }
    60  
    61  type slogJSONHandler struct {
    62  	writer   io.Writer
    63  	options  *slog.HandlerOptions
    64  	fallback slog.Handler
    65  
    66  	entry Entry
    67  
    68  	grouping bool
    69  	groups   int
    70  }
    71  
    72  func (h *slogJSONHandler) Enabled(_ context.Context, level slog.Level) bool {
    73  	if h.options.Level != nil {
    74  		return h.options.Level.Level() <= level
    75  	}
    76  	return slog.LevelInfo <= level
    77  }
    78  
    79  func (h slogJSONHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
    80  	if h.options.ReplaceAttr != nil {
    81  		h.fallback = h.fallback.WithAttrs(attrs)
    82  		return &h
    83  	}
    84  	if len(attrs) == 0 {
    85  		return &h
    86  	}
    87  	i := len(h.entry.buf)
    88  	for _, attr := range attrs {
    89  		h.entry = *slogJSONAttrEval(&h.entry, attr)
    90  	}
    91  	if h.grouping {
    92  		h.entry.buf[i] = '{'
    93  	}
    94  	h.grouping = false
    95  	return &h
    96  }
    97  
    98  func (h slogJSONHandler) WithGroup(name string) slog.Handler {
    99  	if h.options.ReplaceAttr != nil {
   100  		h.fallback = h.fallback.WithGroup(name)
   101  		return &h
   102  	}
   103  	if name == "" {
   104  		return &h
   105  	}
   106  	if h.grouping {
   107  		h.entry.buf = append(h.entry.buf, '{')
   108  	} else {
   109  		h.entry.buf = append(h.entry.buf, ',')
   110  	}
   111  	h.entry.buf = append(h.entry.buf, '"')
   112  	h.entry.buf = append(h.entry.buf, name...)
   113  	h.entry.buf = append(h.entry.buf, '"', ':')
   114  	h.grouping = true
   115  	h.groups++
   116  	return &h
   117  }
   118  
   119  func (h *slogJSONHandler) Handle(ctx context.Context, r slog.Record) error {
   120  	if h.options.ReplaceAttr != nil {
   121  		return h.fallback.Handle(ctx, r)
   122  	}
   123  	return h.handle(ctx, r)
   124  }
   125  
   126  func (h *slogJSONHandler) handle(_ context.Context, r slog.Record) error {
   127  	e := epool.Get().(*Entry)
   128  	e.buf = e.buf[:0]
   129  
   130  	e.buf = append(e.buf, '{')
   131  
   132  	// time
   133  	if !r.Time.IsZero() {
   134  		e.buf = append(e.buf, '"')
   135  		e.buf = append(e.buf, slog.TimeKey...)
   136  		e.buf = append(e.buf, `":"`...)
   137  		e.buf = r.Time.AppendFormat(e.buf, time.RFC3339Nano)
   138  		e.buf = append(e.buf, `",`...)
   139  	}
   140  
   141  	// level
   142  	e.buf = append(e.buf, '"')
   143  	e.buf = append(e.buf, slog.LevelKey...)
   144  	switch r.Level {
   145  	case slog.LevelDebug:
   146  		e.buf = append(e.buf, `":"DEBUG"`...)
   147  	case slog.LevelInfo:
   148  		e.buf = append(e.buf, `":"INFO"`...)
   149  	case slog.LevelWarn:
   150  		e.buf = append(e.buf, `":"WARN"`...)
   151  	case slog.LevelError:
   152  		e.buf = append(e.buf, `":"ERROR"`...)
   153  	default:
   154  		e.buf = append(e.buf, `":"`...)
   155  		e.buf = append(e.buf, r.Level.String()...)
   156  		e.buf = append(e.buf, '"')
   157  	}
   158  
   159  	// source
   160  	if h.options.AddSource && r.PC != 0 {
   161  		name, file, line := pcNameFileLine(r.PC)
   162  		e.buf = append(e.buf, ',', '"')
   163  		e.buf = append(e.buf, slog.SourceKey...)
   164  		e.buf = append(e.buf, `":{"function":"`...)
   165  		e.buf = append(e.buf, name...)
   166  		e.buf = append(e.buf, `","file":"`...)
   167  		e.buf = append(e.buf, file...)
   168  		e.buf = append(e.buf, `","line":`...)
   169  		e.buf = strconv.AppendInt(e.buf, int64(line), 10)
   170  		e.buf = append(e.buf, '}')
   171  	}
   172  
   173  	// msg
   174  	e = e.Str(slog.MessageKey, r.Message)
   175  
   176  	// with
   177  	if b := h.entry.buf; len(b) != 0 {
   178  		e = e.Context(b)
   179  	}
   180  	i := len(e.buf)
   181  
   182  	// attrs
   183  	r.Attrs(func(attr slog.Attr) bool {
   184  		e = slogJSONAttrEval(e, attr)
   185  		return true
   186  	})
   187  
   188  	lastindex := func(buf []byte) int {
   189  		for i := len(buf) - 3; i >= 1; i-- {
   190  			if buf[i] == '"' && (buf[i-1] == ',' || buf[i-1] == '{') {
   191  				return i
   192  			}
   193  		}
   194  		return -1
   195  	}
   196  
   197  	// group attrs
   198  	if h.grouping {
   199  		if r.NumAttrs() > 0 {
   200  			e.buf[i] = '{'
   201  		} else if i = lastindex(e.buf); i > 0 {
   202  			e.buf = e.buf[:i-1]
   203  			h.groups--
   204  			for e.buf[len(e.buf)-1] == ':' {
   205  				if i = lastindex(e.buf); i > 0 {
   206  					e.buf = e.buf[:i-1]
   207  					h.groups--
   208  				}
   209  			}
   210  		} else {
   211  			e.buf = append(e.buf, '{')
   212  		}
   213  	}
   214  
   215  	// brackets closing
   216  	switch h.groups {
   217  	case 0:
   218  		e.buf = append(e.buf, '}', '\n')
   219  	case 1:
   220  		e.buf = append(e.buf, '}', '}', '\n')
   221  	case 2:
   222  		e.buf = append(e.buf, '}', '}', '}', '\n')
   223  	case 3:
   224  		e.buf = append(e.buf, '}', '}', '}', '}', '\n')
   225  	case 4:
   226  		e.buf = append(e.buf, '}', '}', '}', '}', '}', '\n')
   227  	default:
   228  		for i := 0; i <= h.groups; i++ {
   229  			e.buf = append(e.buf, '}')
   230  		}
   231  		e.buf = append(e.buf, '\n')
   232  	}
   233  
   234  	_, err := h.writer.Write(e.buf)
   235  
   236  	if cap(e.buf) <= bbcap {
   237  		epool.Put(e)
   238  	}
   239  
   240  	return err
   241  }
   242  
   243  func SlogNewJSONHandler(writer io.Writer, options *slog.HandlerOptions) slog.Handler {
   244  	h := &slogJSONHandler{
   245  		writer:   writer,
   246  		options:  options,
   247  		fallback: slog.NewJSONHandler(writer, options),
   248  		entry:    *NewContext(nil),
   249  	}
   250  	if h.options == nil {
   251  		h.options = new(slog.HandlerOptions)
   252  	}
   253  	return h
   254  }