introduce xrayslog package

xrayslog/example_test.go

@@ -0,0 +1,40 @@
+//go:build go1.21
+// +build go1.21
+
+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
+}

xrayslog/xrayslog.go

@@ -0,0 +1,81 @@
+//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 (
+	"context"
+	"log/slog"
+
+	"github.com/shogo82148/aws-xray-yasdk-go/xray"
+)
+
+var _ slog.Handler = (*handler)(nil)
+
+type handler struct {
+	parent     slog.Handler
+	traceIDKey string
+	groups     []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 == "" && len(h.groups) == 0 {
+		// no trace ID and no groups. nothing to do.
+		return h.parent.Handle(ctx, record)
+	}
+
+	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))
+		}
+	}
+
+	if traceID != "" {
+		// add trace ID to the log record.
+		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 {
+	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.
+func NewHandler(parent slog.Handler, traceIDKey string) slog.Handler {
+	return &handler{
+		parent:     parent,
+		traceIDKey: traceIDKey,
+	}
+}

xrayslog/xrayslog_test.go

@@ -0,0 +1,135 @@
+//go:build go1.21
+// +build go1.21
+
+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_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)
+	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())
+	}
+}