From 9e8aa43f891655988a57b1cf59a5ecd82395ce2a Mon Sep 17 00:00:00 2001 From: Russ Egan Date: Wed, 8 Jan 2025 17:30:27 -0600 Subject: [PATCH 1/2] Add support for adding headers Headers are optional attributes removed from the end of the line, and injected (values only) into the header of the line, after the level/source, and before the message. --- encoding.go | 43 ++++++++++---- handler.go | 91 ++++++++++++++++++++++++++---- handler_test.go | 146 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 256 insertions(+), 24 deletions(-) diff --git a/encoding.go b/encoding.go index 6088e65..130f0e8 100644 --- a/encoding.go +++ b/encoding.go @@ -86,8 +86,8 @@ func (e encoder) writeSource(buf *buffer, pc uintptr, cwd string) { buf.AppendString(frame.File) buf.AppendByte(':') buf.AppendInt(int64(frame.Line)) + buf.AppendByte(' ') }) - e.writeColoredString(buf, " > ", e.opts.Theme.AttrKey()) } func (e encoder) writeMessage(buf *buffer, level slog.Level, msg string) { @@ -98,6 +98,26 @@ func (e encoder) writeMessage(buf *buffer, level slog.Level, msg string) { } } +func (e encoder) writeHeaders(buf *buffer, headers []slog.Value) bool { + wrote := false + for _, a := range headers { + // todo: this skips empty values, omitting them entire from the header. + // alternately, I could print or something, so the number of + // headers in each log entry is always fixed... + if a.Kind() == slog.KindAny && a.Any() == nil { + continue + } + e.writeValue(buf, a, e.opts.Theme.Source()) + buf.AppendByte(' ') + wrote = true + } + return wrote +} + +func (e encoder) writeHeaderSeparator(buf *buffer) { + e.writeColoredString(buf, "> ", e.opts.Theme.AttrKey()) +} + func (e encoder) writeAttr(buf *buffer, a slog.Attr, group string) { // Elide empty Attrs. if a.Equal(slog.Attr{}) { @@ -123,38 +143,37 @@ func (e encoder) writeAttr(buf *buffer, a slog.Attr, group string) { buf.AppendString(a.Key) buf.AppendByte('=') }) - e.writeValue(buf, value) + e.writeValue(buf, value, e.opts.Theme.AttrValue()) } -func (e encoder) writeValue(buf *buffer, value slog.Value) { - attrValue := e.opts.Theme.AttrValue() +func (e encoder) writeValue(buf *buffer, value slog.Value, c ANSIMod) { switch value.Kind() { case slog.KindInt64: - e.writeColoredInt(buf, value.Int64(), attrValue) + e.writeColoredInt(buf, value.Int64(), c) case slog.KindBool: - e.writeColoredBool(buf, value.Bool(), attrValue) + e.writeColoredBool(buf, value.Bool(), c) case slog.KindFloat64: - e.writeColoredFloat(buf, value.Float64(), attrValue) + e.writeColoredFloat(buf, value.Float64(), c) case slog.KindTime: - e.writeColoredTime(buf, value.Time(), e.opts.TimeFormat, attrValue) + e.writeColoredTime(buf, value.Time(), e.opts.TimeFormat, c) case slog.KindUint64: - e.writeColoredUint(buf, value.Uint64(), attrValue) + e.writeColoredUint(buf, value.Uint64(), c) case slog.KindDuration: - e.writeColoredDuration(buf, value.Duration(), attrValue) + e.writeColoredDuration(buf, value.Duration(), c) case slog.KindAny: switch v := value.Any().(type) { case error: e.writeColoredString(buf, v.Error(), e.opts.Theme.AttrValueError()) return case fmt.Stringer: - e.writeColoredString(buf, v.String(), attrValue) + e.writeColoredString(buf, v.String(), c) return } fallthrough case slog.KindString: fallthrough default: - e.writeColoredString(buf, value.String(), attrValue) + e.writeColoredString(buf, value.String(), c) } } diff --git a/handler.go b/handler.go index 82ce3ea..58524c0 100644 --- a/handler.go +++ b/handler.go @@ -5,6 +5,7 @@ import ( "io" "log/slog" "os" + "slices" "strings" "sync" "time" @@ -38,6 +39,11 @@ type HandlerOptions struct { // Theme defines the colorized output using ANSI escape sequences Theme Theme + + // Headers are a list of attribute keys. These attributes will be removed from + // the trailing attr list, and the values will be inserted between + // the level/source and the message, in the configured order. + Headers []string } type Handler struct { @@ -45,6 +51,7 @@ type Handler struct { out io.Writer group string context buffer + headers []slog.Value enc *encoder } @@ -72,6 +79,7 @@ func NewHandler(out io.Writer, opts *HandlerOptions) *Handler { group: "", context: nil, enc: &encoder{opts: *opts}, + headers: make([]slog.Value, len(opts.Headers)), } } @@ -82,31 +90,66 @@ func (h *Handler) Enabled(_ context.Context, l slog.Level) bool { // Handle implements slog.Handler. func (h *Handler) Handle(_ context.Context, rec slog.Record) error { - buf := bufferPool.Get().(*buffer) + headerBuf, trailerBuf := bufferPool.Get().(*buffer), bufferPool.Get().(*buffer) + + h.enc.writeTimestamp(headerBuf, rec.Time) + h.enc.writeLevel(headerBuf, rec.Level) - h.enc.writeTimestamp(buf, rec.Time) - h.enc.writeLevel(buf, rec.Level) + var writeHeaderSeparator bool if h.opts.AddSource && rec.PC > 0 { - h.enc.writeSource(buf, rec.PC, cwd) + h.enc.writeSource(headerBuf, rec.PC, cwd) + writeHeaderSeparator = true } - h.enc.writeMessage(buf, rec.Level, rec.Message) - buf.copy(&h.context) + + h.enc.writeMessage(trailerBuf, rec.Level, rec.Message) + + trailerBuf.copy(&h.context) + + headers := h.headers + headersChanged := false rec.Attrs(func(a slog.Attr) bool { - h.enc.writeAttr(buf, a, h.group) + idx := slices.IndexFunc(h.opts.Headers, func(s string) bool { return s == a.Key }) + if idx >= 0 { + if !headersChanged { + headersChanged = true + // todo: this makes one allocation, but only if the headers weren't already + // satisfied by prior WithAttrs(). Could use a pool of *[]slog.Value, but + // I'm not sure it's worth it. + headers = make([]slog.Value, len(h.opts.Headers)) + copy(headers, h.headers) + } + headers[idx] = a.Value + } else { + h.enc.writeAttr(trailerBuf, a, h.group) + } return true }) - h.enc.NewLine(buf) - if _, err := buf.WriteTo(h.out); err != nil { - buf.Reset() - bufferPool.Put(buf) + h.enc.NewLine(trailerBuf) + + if len(headers) > 0 { + if h.enc.writeHeaders(headerBuf, headers) { + writeHeaderSeparator = true + } + } + + if writeHeaderSeparator { + h.enc.writeHeaderSeparator(headerBuf) + } + + if _, err := headerBuf.WriteTo(h.out); err != nil { + return err + } + if _, err := trailerBuf.WriteTo(h.out); err != nil { return err } - bufferPool.Put(buf) + bufferPool.Put(headerBuf) + bufferPool.Put(trailerBuf) return nil } // WithAttrs implements slog.Handler. func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler { + headers := h.extractHeaders(attrs) newCtx := h.context for _, a := range attrs { h.enc.writeAttr(&newCtx, a, h.group) @@ -118,6 +161,7 @@ func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler { group: h.group, context: newCtx, enc: h.enc, + headers: headers, } } @@ -133,5 +177,28 @@ func (h *Handler) WithGroup(name string) slog.Handler { group: name, context: h.context, enc: h.enc, + headers: h.headers, + } +} + +// extractHeaders scans the attributes for keys specified in HeaderKeys. +// If found, their values are saved in a new list. +// The original attribute list will be modified to remove the extracted attributes. +func (h *Handler) extractHeaders(attrs []slog.Attr) (headers []slog.Value) { + changed := false + headers = h.headers + for i, attr := range attrs { + idx := slices.IndexFunc(h.opts.Headers, func(s string) bool { return s == attr.Key }) + if idx >= 0 { + if !changed { + // make a copy of prefixes: + headers = make([]slog.Value, len(h.headers)) + copy(headers, h.headers) + } + headers[idx] = attr.Value + attrs[i] = slog.Attr{} // remove the prefix attribute + changed = true + } } + return } diff --git a/handler_test.go b/handler_test.go index eb09629..e96cd55 100644 --- a/handler_test.go +++ b/handler_test.go @@ -271,6 +271,152 @@ func TestHandler_Source(t *testing.T) { AssertEqual(t, fmt.Sprintf("%s INF foobar\n", now.Format(time.DateTime)), buf.String()) } +func TestHandler_Headers(t *testing.T) { + pc, file, line, _ := runtime.Caller(0) + cwd, _ := os.Getwd() + file, _ = filepath.Rel(cwd, file) + sourceField := fmt.Sprintf("%s:%d", file, line) + + tests := []struct { + name string + opts HandlerOptions + attrs []slog.Attr + withAttrs []slog.Attr + withGroups []string + want string + }{ + { + name: "no headers", + attrs: []slog.Attr{slog.String("foo", "bar")}, + want: "INF with headers foo=bar\n", + }, + { + name: "one header", + opts: HandlerOptions{Headers: []string{"foo"}}, + attrs: []slog.Attr{ + slog.String("foo", "bar"), + slog.String("bar", "baz"), + }, + want: "INF bar > with headers bar=baz\n", + }, + { + name: "two headers", + opts: HandlerOptions{Headers: []string{"foo", "bar"}}, + attrs: []slog.Attr{ + slog.String("foo", "bar"), + slog.String("bar", "baz"), + }, + want: "INF bar baz > with headers\n", + }, + { + name: "missing headers", + opts: HandlerOptions{Headers: []string{"foo", "bar"}}, + attrs: []slog.Attr{slog.String("bar", "baz"), slog.String("baz", "foo")}, + want: "INF baz > with headers baz=foo\n", + }, + { + name: "missing all headers", + opts: HandlerOptions{Headers: []string{"foo", "bar"}}, + want: "INF with headers\n", + }, + { + name: "header and source", + opts: HandlerOptions{Headers: []string{"foo"}, AddSource: true}, + attrs: []slog.Attr{ + slog.String("foo", "bar"), + slog.String("bar", "baz"), + }, + want: "INF " + sourceField + " bar > with headers bar=baz\n", + }, + { + name: "withattrs", + opts: HandlerOptions{Headers: []string{"foo"}}, + attrs: []slog.Attr{ + + slog.String("bar", "baz"), + }, + withAttrs: []slog.Attr{ + slog.String("foo", "bar"), + }, + want: "INF bar > with headers bar=baz\n", + }, + { + name: "withgroup", + opts: HandlerOptions{Headers: []string{"foo", "bar"}}, + attrs: []slog.Attr{ + slog.String("bar", "baz"), + slog.String("baz", "foo"), + }, + withGroups: []string{"group"}, + withAttrs: []slog.Attr{ + slog.String("foo", "bar"), + }, + want: "INF bar baz > with headers group.baz=foo\n", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + buf := bytes.Buffer{} + + opts := &test.opts + opts.NoColor = true + var h slog.Handler = NewHandler(&buf, &test.opts) + if test.withAttrs != nil { + h = h.WithAttrs(test.withAttrs) + } + for _, g := range test.withGroups { + h = h.WithGroup(g) + } + + rec := slog.NewRecord(time.Time{}, slog.LevelInfo, "with headers", pc) + + rec.AddAttrs(test.attrs...) + + AssertNoError(t, h.Handle(context.Background(), rec)) + AssertEqual(t, test.want, buf.String()) + }) + } + + t.Run("withAttrs state keeping", func(t *testing.T) { + // test to make sure the way that WithAttrs() copies the cached headers doesn't leak + // headers back to the parent handler or to subsequent Handle() calls (i.e. ensure that + // the headers slice is copied at the right times). + + buf := bytes.Buffer{} + h := NewHandler(&buf, &HandlerOptions{ + Headers: []string{"foo", "bar"}, + TimeFormat: "0", + NoColor: true, + }) + + assertLog := func(t *testing.T, handler slog.Handler, want string, attrs ...slog.Attr) { + buf.Reset() + rec := slog.NewRecord(time.Time{}, slog.LevelInfo, "with headers", pc) + + rec.AddAttrs(attrs...) + + AssertNoError(t, handler.Handle(context.Background(), rec)) + AssertEqual(t, want, buf.String()) + } + + assertLog(t, h, "INF bar > with headers\n", slog.String("foo", "bar")) + + h2 := h.WithAttrs([]slog.Attr{slog.String("foo", "baz")}) + assertLog(t, h2, "INF baz > with headers\n") + + h3 := h2.WithAttrs([]slog.Attr{slog.String("foo", "buz")}) + assertLog(t, h3, "INF buz > with headers\n") + // creating h3 should not have affected h2 + assertLog(t, h2, "INF baz > with headers\n") + + // overriding attrs shouldn't affect the handler + assertLog(t, h2, "INF biz > with headers\n", slog.String("foo", "biz")) + assertLog(t, h2, "INF baz > with headers\n") + + }) +} + func TestHandler_Err(t *testing.T) { w := writerFunc(func(b []byte) (int, error) { return 0, errors.New("nope") }) h := NewHandler(w, &HandlerOptions{NoColor: true}) From bb75924f2753f01db680d6e2cb21e62f6e7a467f Mon Sep 17 00:00:00 2001 From: Russ Egan Date: Wed, 8 Jan 2025 20:01:48 -0600 Subject: [PATCH 2/2] type --- handler.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/handler.go b/handler.go index 58524c0..6ab7f7f 100644 --- a/handler.go +++ b/handler.go @@ -181,7 +181,7 @@ func (h *Handler) WithGroup(name string) slog.Handler { } } -// extractHeaders scans the attributes for keys specified in HeaderKeys. +// extractHeaders scans the attributes for keys specified in Headers. // If found, their values are saved in a new list. // The original attribute list will be modified to remove the extracted attributes. func (h *Handler) extractHeaders(attrs []slog.Attr) (headers []slog.Value) {