From 12ac865f93eb53b03c1014db2df81acf9daa9481 Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Sun, 13 Aug 2023 01:08:43 +0900 Subject: [PATCH 1/7] introduce xrayslog package --- xrayslog/xrayslog.go | 59 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 xrayslog/xrayslog.go diff --git a/xrayslog/xrayslog.go b/xrayslog/xrayslog.go new file mode 100644 index 0000000..2f4a6df --- /dev/null +++ b/xrayslog/xrayslog.go @@ -0,0 +1,59 @@ +package xrayslog + +import ( + "context" + "log/slog" + + "github.com/shogo82148/aws-xray-yasdk-go/xray" +) + +var _ slog.Handler = (*handler)(nil) + +type handler struct { + parent slog.Handler + traceIDKey string +} + +// Enable implements slog.Handler interface. +func (h *handler) Enabled(ctx context.Context, level slog.Level) bool { + return h.parent.Enabled(ctx, level) +} + +// Handle implements slog.Handler interface. +func (h *handler) Handle(ctx context.Context, record slog.Record) error { + traceID := xray.ContextTraceID(ctx) + if traceID == "" { + // there is no trace ID in the context. + // don't add trace ID to the log record. + return h.parent.Handle(ctx, record) + } + + // add trace ID to the log record. + newRecord := record.Clone() + newRecord.AddAttrs(slog.String(h.traceIDKey, traceID)) + return h.parent.Handle(ctx, newRecord) +} + +// WithAttrs implements slog.Handler interface. +func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler { + return &handler{ + parent: h.parent.WithAttrs(attrs), + traceIDKey: h.traceIDKey, + } +} + +// WithGroup implements slog.Handler interface. +func (h *handler) WithGroup(name string) slog.Handler { + return &handler{ + parent: h.parent.WithGroup(name), + traceIDKey: h.traceIDKey, + } +} + +// NewHandler returns a slog.Handler that adds trace ID to the log record. +func NewHandler(parent slog.Handler, traceIDKey string) slog.Handler { + return &handler{ + parent: parent, + traceIDKey: traceIDKey, + } +} From 4ec25b67870f86e2f4dcc3fdaa56881da5e4048d Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Sun, 13 Aug 2023 11:32:21 +0900 Subject: [PATCH 2/7] add build tag go1.21 --- xrayslog/xrayslog.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xrayslog/xrayslog.go b/xrayslog/xrayslog.go index 2f4a6df..1ac0541 100644 --- a/xrayslog/xrayslog.go +++ b/xrayslog/xrayslog.go @@ -1,3 +1,6 @@ +//go:build go1.21 +// +build go1.21 + package xrayslog import ( From e7f8d0bf7635d2a88c0a418280d7a08011f0638d Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Fri, 1 Sep 2023 23:28:59 +0900 Subject: [PATCH 3/7] fix WithGroup --- xrayslog/xrayslog.go | 24 ++++++--- xrayslog/xrayslog_test.go | 110 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 6 deletions(-) create mode 100644 xrayslog/xrayslog_test.go diff --git a/xrayslog/xrayslog.go b/xrayslog/xrayslog.go index 1ac0541..e1d85f5 100644 --- a/xrayslog/xrayslog.go +++ b/xrayslog/xrayslog.go @@ -15,6 +15,7 @@ var _ slog.Handler = (*handler)(nil) type handler struct { parent slog.Handler traceIDKey string + groups []string } // Enable implements slog.Handler interface. @@ -27,12 +28,24 @@ func (h *handler) Handle(ctx context.Context, record slog.Record) error { traceID := xray.ContextTraceID(ctx) if traceID == "" { // there is no trace ID in the context. - // don't add trace ID to the log record. + // we don't need to add trace ID to the log record. return h.parent.Handle(ctx, record) } + newRecord := slog.NewRecord(record.Time, record.Level, record.Message, record.PC) + attrs := make([]any, 0, record.NumAttrs()) + record.Attrs(func(a slog.Attr) bool { + attrs = append(attrs, a) + return true + }) + for i := len(h.groups) - 1; i >= 0; i-- { + attrs = []any{slog.Group(h.groups[i], attrs...)} + } + for _, attr := range attrs { + newRecord.AddAttrs(attr.(slog.Attr)) + } + // add trace ID to the log record. - newRecord := record.Clone() newRecord.AddAttrs(slog.String(h.traceIDKey, traceID)) return h.parent.Handle(ctx, newRecord) } @@ -47,10 +60,9 @@ func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler { // WithGroup implements slog.Handler interface. func (h *handler) WithGroup(name string) slog.Handler { - return &handler{ - parent: h.parent.WithGroup(name), - traceIDKey: h.traceIDKey, - } + h2 := *h // shallow copy, but it is OK. + h2.groups = append(h2.groups, name) + return &h2 } // NewHandler returns a slog.Handler that adds trace ID to the log record. diff --git a/xrayslog/xrayslog_test.go b/xrayslog/xrayslog_test.go new file mode 100644 index 0000000..536535f --- /dev/null +++ b/xrayslog/xrayslog_test.go @@ -0,0 +1,110 @@ +package xrayslog + +import ( + "bytes" + "context" + "encoding/json" + "log/slog" + "testing" + + "github.com/shogo82148/aws-xray-yasdk-go/xray" +) + +func TestHandle_WithoutTraceID(t *testing.T) { + // build the logger + w := &bytes.Buffer{} + parent := slog.NewJSONHandler(w, nil) + h := NewHandler(parent, "trace_id") + logger := slog.New(h) + + ctx := context.Background() + + // test the logger + logger.InfoContext(ctx, "hello") + var v map[string]any + if err := json.Unmarshal(w.Bytes(), &v); err != nil { + t.Error(err) + } + if got, ok := v["trace_id"]; ok { + t.Errorf("trace_id should be empty, but got %s", got) + } +} + +func TestHandle_WithTraceID(t *testing.T) { + // build the logger + w := &bytes.Buffer{} + parent := slog.NewJSONHandler(w, nil) + h := NewHandler(parent, "trace_id") + logger := slog.New(h) + + // begin a new segment + ctx := context.Background() + ctx, segment := xray.BeginSegment(ctx, "my-segment") + defer segment.Close() + + // test the logger + logger.InfoContext(ctx, "hello") + var v map[string]any + if err := json.Unmarshal(w.Bytes(), &v); err != nil { + t.Error(err) + } + if v["trace_id"] != xray.ContextTraceID(ctx) { + t.Errorf("trace_id is not set: %s", w.String()) + } +} + +func TestWithAttrs(t *testing.T) { + // build the logger + w := &bytes.Buffer{} + parent := slog.NewJSONHandler(w, nil) + h := NewHandler(parent, "trace_id") + logger := slog.New(h.WithAttrs([]slog.Attr{ + slog.String("foo", "bar"), + })) + + // begin a new segment + ctx := context.Background() + ctx, segment := xray.BeginSegment(ctx, "my-segment") + defer segment.Close() + + // test the logger + logger.InfoContext(ctx, "hello") + var v map[string]any + if err := json.Unmarshal(w.Bytes(), &v); err != nil { + t.Error(err) + } + if v["trace_id"] != xray.ContextTraceID(ctx) { + t.Errorf("trace_id is not set: %s", w.String()) + } + if v["foo"] != "bar" { + t.Errorf("foo is not set: %s", w.String()) + } +} + +func TestWithGroup(t *testing.T) { + // build the logger + w := &bytes.Buffer{} + parent := slog.NewJSONHandler(w, nil) + h := NewHandler(parent, "trace_id") + logger := slog.New(h.WithGroup("my-group1").WithGroup("my-group2")) + + // begin a new segment + ctx := context.Background() + ctx, segment := xray.BeginSegment(ctx, "my-segment") + defer segment.Close() + + // test the logger + logger.InfoContext(ctx, "hello", slog.String("foo", "bar")) + var v map[string]any + if err := json.Unmarshal(w.Bytes(), &v); err != nil { + t.Error(err) + } + if v["trace_id"] != xray.ContextTraceID(ctx) { + t.Errorf("trace_id is not set: %s", w.String()) + } + group, _ := v["my-group1"].(map[string]any) + group, _ = group["my-group2"].(map[string]any) + if group == nil || group["foo"] != "bar" { + t.Errorf("foo is not set: %s", w.String()) + } +} From d0b28d9157d443a37d6f78b8c67da19c74b1186f Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Fri, 1 Sep 2023 23:39:19 +0900 Subject: [PATCH 4/7] fix WithGroup when there is no trace id --- xrayslog/xrayslog.go | 38 ++++++++++++++++++++++---------------- xrayslog/xrayslog_test.go | 24 +++++++++++++++++++++++- 2 files changed, 45 insertions(+), 17 deletions(-) diff --git a/xrayslog/xrayslog.go b/xrayslog/xrayslog.go index e1d85f5..9450280 100644 --- a/xrayslog/xrayslog.go +++ b/xrayslog/xrayslog.go @@ -26,27 +26,33 @@ func (h *handler) Enabled(ctx context.Context, level slog.Level) bool { // Handle implements slog.Handler interface. func (h *handler) Handle(ctx context.Context, record slog.Record) error { traceID := xray.ContextTraceID(ctx) - if traceID == "" { - // there is no trace ID in the context. - // we don't need to add trace ID to the log record. + if traceID == "" && len(h.groups) == 0 { + // no trace ID and no groups. nothing to do. return h.parent.Handle(ctx, record) } - newRecord := slog.NewRecord(record.Time, record.Level, record.Message, record.PC) - attrs := make([]any, 0, record.NumAttrs()) - record.Attrs(func(a slog.Attr) bool { - attrs = append(attrs, a) - return true - }) - for i := len(h.groups) - 1; i >= 0; i-- { - attrs = []any{slog.Group(h.groups[i], attrs...)} - } - for _, attr := range attrs { - newRecord.AddAttrs(attr.(slog.Attr)) + var newRecord slog.Record + if len(h.groups) == 0 { + newRecord = record.Clone() + } else { + newRecord = slog.NewRecord(record.Time, record.Level, record.Message, record.PC) + attrs := make([]any, 0, record.NumAttrs()) + record.Attrs(func(a slog.Attr) bool { + attrs = append(attrs, a) + return true + }) + for i := len(h.groups) - 1; i >= 0; i-- { + attrs = []any{slog.Group(h.groups[i], attrs...)} + } + for _, attr := range attrs { + newRecord.AddAttrs(attr.(slog.Attr)) + } } - // add trace ID to the log record. - newRecord.AddAttrs(slog.String(h.traceIDKey, traceID)) + if traceID != "" { + // add trace ID to the log record. + newRecord.AddAttrs(slog.String(h.traceIDKey, traceID)) + } return h.parent.Handle(ctx, newRecord) } diff --git a/xrayslog/xrayslog_test.go b/xrayslog/xrayslog_test.go index 536535f..bd6ea9b 100644 --- a/xrayslog/xrayslog_test.go +++ b/xrayslog/xrayslog_test.go @@ -81,7 +81,29 @@ func TestWithAttrs(t *testing.T) { } } -func TestWithGroup(t *testing.T) { +func TestWithGroup_WithoutTraceID(t *testing.T) { + // build the logger + w := &bytes.Buffer{} + parent := slog.NewJSONHandler(w, nil) + h := NewHandler(parent, "trace_id") + logger := slog.New(h.WithGroup("my-group1").WithGroup("my-group2")) + + ctx := context.Background() + + // test the logger + logger.InfoContext(ctx, "hello", slog.String("foo", "bar")) + var v map[string]any + if err := json.Unmarshal(w.Bytes(), &v); err != nil { + t.Error(err) + } + group, _ := v["my-group1"].(map[string]any) + group, _ = group["my-group2"].(map[string]any) + if group == nil || group["foo"] != "bar" { + t.Errorf("foo is not set: %s", w.String()) + } +} + +func TestWithGroup_WithTraceID(t *testing.T) { // build the logger w := &bytes.Buffer{} parent := slog.NewJSONHandler(w, nil) From 8ed220cf353bb02ddec7d0990a5271b8019c87fc Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Fri, 1 Sep 2023 23:43:09 +0900 Subject: [PATCH 5/7] add a build tag for Go 1.21 --- xrayslog/xrayslog_test.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xrayslog/xrayslog_test.go b/xrayslog/xrayslog_test.go index bd6ea9b..57412ac 100644 --- a/xrayslog/xrayslog_test.go +++ b/xrayslog/xrayslog_test.go @@ -1,3 +1,6 @@ +//go:build go1.21 +// +build go1.21 + package xrayslog import ( From fb7535efa06ffc453f3ec60a8719a4b75438fea3 Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Fri, 1 Sep 2023 23:51:19 +0900 Subject: [PATCH 6/7] add an example test --- xrayslog/example_test.go | 37 +++++++++++++++++++++++++++++++++++++ xrayslog/xrayslog.go | 1 + 2 files changed, 38 insertions(+) create mode 100644 xrayslog/example_test.go diff --git a/xrayslog/example_test.go b/xrayslog/example_test.go new file mode 100644 index 0000000..1f7a9ea --- /dev/null +++ b/xrayslog/example_test.go @@ -0,0 +1,37 @@ +package xrayslog_test + +import ( + "context" + "log/slog" + "os" + + "github.com/shogo82148/aws-xray-yasdk-go/xray" + "github.com/shogo82148/aws-xray-yasdk-go/xrayslog" +) + +func ExampleNewHandler() { + // it's for testing. + replace := func(groups []string, a slog.Attr) slog.Attr { + // Remove time. + if a.Key == slog.TimeKey && len(groups) == 0 { + return slog.Attr{} + } + return a + } + + // build the logger + parent := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ReplaceAttr: replace}) + h := xrayslog.NewHandler(parent, "trace_id") + logger := slog.New(h) + + // begin a new segment + ctx := context.Background() + ctx, segment := xray.BeginSegmentWithHeader(ctx, "my-segment", "Root=1-5e645f3e-1dfad076a177c5ccc5de12f5") + defer segment.Close() + + // output the log + logger.InfoContext(ctx, "hello") + + // Output: + // level=INFO msg=hello trace_id=1-5e645f3e-1dfad076a177c5ccc5de12f5 +} diff --git a/xrayslog/xrayslog.go b/xrayslog/xrayslog.go index 9450280..2eac01a 100644 --- a/xrayslog/xrayslog.go +++ b/xrayslog/xrayslog.go @@ -1,6 +1,7 @@ //go:build go1.21 // +build go1.21 +// Package xrayslog provides a [log/slog.Handler] that adds trace ID to the log record. package xrayslog import ( From a223744ecff98b11b4a57dd97ccd547cbdfb1e6b Mon Sep 17 00:00:00 2001 From: ICHINOSE Shogo Date: Fri, 1 Sep 2023 23:54:10 +0900 Subject: [PATCH 7/7] add the build tag for Go 1.21 --- xrayslog/example_test.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xrayslog/example_test.go b/xrayslog/example_test.go index 1f7a9ea..9a0613e 100644 --- a/xrayslog/example_test.go +++ b/xrayslog/example_test.go @@ -1,3 +1,6 @@ +//go:build go1.21 +// +build go1.21 + package xrayslog_test import (