Support tracing

Fixes 3

Makefile

@@ -1,7 +1,10 @@
 all: test
 
 test:
-	go test -v -shuffle=on -coverprofile=coverage.txt -count=1 ./...
+	go test -shuffle=on -coverprofile=coverage.txt -count=1 ./...
+
+bench:
+	go test -run=XXX -bench=. ./...
 
 vet:
 	go vet ./...

README.md

@@ -6,35 +6,81 @@ A lightweight [`log/slog.JSONHandler`](https://pkg.go.dev/log/slog#JSONHandler)
 wrapper that adapts the fields to the [Google Cloud Logging structured log
 format](https://cloud.google.com/logging/docs/structured-logging#structured_logging_special_fields).
 
-The handler merely reformats/renames the structured JSON log fields. It's
-still `JSONHandler` under the hood. It does NOT send logs to Cloud Logging
-directly (e.g. using the Cloud SDK).
-
 The intended use case is Cloud Run, but it should work in similar environments
 where logs are emitted to stdout/stderr and automatically picked up by Cloud
 Logging (e.g. App Engine, Cloud Functions, GKE).
 
+## Features
+
+- Lightweight. The handler merely reformats/renames the structured JSON log
+  fields. It's still [`log/slog.JSONHandler`](https://pkg.go.dev/log/slog#JSONHandler)
+  under the hood. It does NOT send logs to Cloud Logging directly (e.g. using
+  the Cloud SDK).
+
+- Tracing. A tracing middleware is provided to automatically extract tracing
+  information from `traceparent` or `X-Cloud-Trace-Context` HTTP request header,
+  and attaches it to the request context. The Handler automatically includes any
+  tracing information as log attributes.
+
+- Custom levels as supported by Google Cloud Logging, e.g. CRITICAL and NOTICE.
+
 ## Usage
 
 ```go
 import (
 	"log/slog"
 
+	"cloud.google.com/go/compute/metadata"
 	"github.com/ronny/clog"
 )
 
 func main() {
+	projectID, err := metadata.ProjectID()
+	if err != nil {
+		panic(err)
+	}
+
 	logger := slog.New(
 		clog.NewHandler(os.Stderr, clog.HandlerOptions{
 			Level: clog.LevelInfo,
+			GoogleProjectID: projectID,
 		}),
 	)
-	logger.Warn("flux capacitor is too warm", "tempCelsius", 42)
-	logger.Log(ctx, clog.LevelCritical, "flux capacitor is on fire")
+	slog.SetDefault(logger)
+
+	mux := http.NewServeMux()
+	mux.Handle("POST /warn", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+		// ⚠️ The log will have tracing attrs since we're using
+		// trace.Middleware below (assuming trace header is in the request):
+		// "logging.googleapis.com/trace"
+		// "logging.googleapis.com/spanId"
+		// "logging.googleapis.com/traceSampled"
+		slog.WarnContext(ctx, "flux capacitor is too warm",
+			"mycount", 42,
+		)
+	}))
+	mux.Handle("POST /critical", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := r.Context()
+		// ⚠️ Custom level CRITICAL
+		slog.Log(ctx, clog.LevelCritical, "flux capacitor is on fire")
+	}))
+
+	port := os.Getenv("PORT")
+	if port == "" {
+		port = "8080"
+	}
+	log.Printf("listening on port %s", port)
+
+	// ⚠️ `trace.Middleware` to make tracing information available in ctx in mux
+	// handlers.
+	if err := http.ListenAndServe(":"+port, trace.Middleware(mux)); err != nil {
+		log.Fatal(err)
+	}
 }
 ```
 
 ## Credits and acknowledgements
 
-Thank you to Remko Tronçon for doing all the hard work in
+Thank you to Remko Tronçon for doing most of the hard work in
 https://github.com/remko/cloudrun-slog which is the basis for this library.

handler.go

@@ -1,20 +1,67 @@
 package clog
 
 import (
+	"context"
+	"errors"
+	"fmt"
 	"io"
 	"log/slog"
+
+	"github.com/ronny/clog/trace"
 )
 
-type HandlerOptions = slog.HandlerOptions
+var ErrInvalidHandlerOptions = errors.New("invalid HandlerOptions")
 
-// NewHandler returns a [log/slog.JSONHandler] pre-configured for Google Cloud
-// Logging.
-//
-// Basically it replaces `opts.ReplaceAttr` with [ReplaceAttr].
-func NewHandler(w io.Writer, opts *HandlerOptions) *slog.JSONHandler {
-	if opts == nil {
-		opts = &HandlerOptions{}
-	}
+type HandlerOptions struct {
+	AddSource       bool
+	Level           slog.Leveler
+	ReplaceAttr     func(groups []string, a slog.Attr) slog.Attr
+	GoogleProjectID string
+}
+
+var _ slog.Handler = (*Handler)(nil)
+
+// Handler is a [log/slog.JSONHandler] preconfigured for Google Cloud Logging.
+type Handler struct {
+	opts    HandlerOptions
+	handler slog.Handler
+}
+
+func NewHandler(w io.Writer, opts HandlerOptions) (*Handler, error) {
 	opts.ReplaceAttr = ReplaceAttr
-	return slog.NewJSONHandler(w, opts)
+	if opts.GoogleProjectID == "" {
+		return nil, fmt.Errorf("%w: missing GoogleProjectID", ErrInvalidHandlerOptions)
+	}
+	return &Handler{
+		opts: opts,
+		handler: slog.NewJSONHandler(w, &slog.HandlerOptions{
+			AddSource:   opts.AddSource,
+			Level:       opts.Level,
+			ReplaceAttr: opts.ReplaceAttr,
+		}),
+	}, nil
+}
+
+// Handle implements [log/slog.Handler].
+func (h *Handler) Handle(ctx context.Context, record slog.Record) error {
+	return h.handler.Handle(ctx,
+		trace.NewRecord(ctx, record, h.opts.GoogleProjectID),
+	)
+}
+
+// Enabled implements [log/slog.Handler].
+func (h *Handler) Enabled(ctx context.Context, level slog.Level) bool {
+	return h.handler.Enabled(ctx, level)
+}
+
+// WithAttrs implements [log/slog.Handler].
+func (h *Handler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	h.handler = h.handler.WithAttrs(attrs)
+	return h
+}
+
+// WithGroup implements [log/slog.Handler].
+func (h *Handler) WithGroup(name string) slog.Handler {
+	h.handler = h.handler.WithGroup(name)
+	return h
 }

handler_test.go

@@ -3,7 +3,6 @@ package clog_test
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"log/slog"
 	"strings"
 	"testing"
@@ -17,10 +16,21 @@ func TestNewHandler(t *testing.T) {
 
 	buf := strings.Builder{}
 
-	handler := clog.NewHandler(&buf, &clog.HandlerOptions{
-		AddSource: true,
-		Level:     clog.LevelInfo,
+	_, err := clog.NewHandler(&buf, clog.HandlerOptions{})
+	if !assert.NotNil(t, err) {
+		return
+	}
+	assert.ErrorIs(t, err, clog.ErrInvalidHandlerOptions)
+	assert.ErrorContains(t, err, "missing GoogleProjectID")
+
+	handler, err := clog.NewHandler(&buf, clog.HandlerOptions{
+		AddSource:       true,
+		Level:           clog.LevelInfo,
+		GoogleProjectID: "my-project-id",
 	})
+	if !assert.Nil(t, err) {
+		return
+	}
 
 	logger := slog.New(handler)
 
@@ -37,7 +47,7 @@ func TestNewHandler(t *testing.T) {
 		clog.SpanIDKey, "banana",
 	)
 
-	fmt.Printf("%q\n", buf.String())
+	// fmt.Printf("%q\n", buf.String())
 
 	lines := strings.Split(buf.String(), "\n")
 	// the log lines are \n terminated, so the last line will always be empty since we split on \n
@@ -47,7 +57,7 @@ func TestNewHandler(t *testing.T) {
 	assert.Equal(t, 2, len(lines))
 
 	var warnEntry Entry
-	err := json.Unmarshal([]byte(lines[0]), &warnEntry)
+	err = json.Unmarshal([]byte(lines[0]), &warnEntry)
 	if err != nil {
 		t.Fatal(err)
 	}

keys.go

@@ -1,5 +1,7 @@
 package clog
 
+import "github.com/ronny/clog/trace"
+
 // Standard JSON log fields as per
 // https://cloud.google.com/logging/docs/structured-logging#structured_logging_special_fields
 const (
@@ -15,7 +17,7 @@ const (
 	LabelsKey         = "logging.googleapis.com/labels"
 	OperationKey      = "logging.googleapis.com/operation"
 	SourceLocationKey = "logging.googleapis.com/sourceLocation"
-	SpanIDKey         = "logging.googleapis.com/spanId"
-	TraceKey          = "logging.googleapis.com/trace"
-	TraceSampledKey   = "logging.googleapis.com/traceSampled"
+	TraceKey          = trace.TraceKey
+	SpanIDKey         = trace.SpanIDKey
+	TraceSampledKey   = trace.TraceSampledKey
 )

trace/cloud_trace.go

@@ -0,0 +1,87 @@
+package trace
+
+import (
+	"fmt"
+	"regexp"
+)
+
+var _ Trace = (*CloudTraceContext)(nil)
+
+// TRACE_ID/SPAN_ID();o=OPTIONS
+// Example: 70e0091f6f5d4643bb4eca9d81320c76/97123319527522;o=1
+var cloudTraceContextRegex = regexp.MustCompile(`^(?P<TraceID>[0-9a-fA-F]{32})/(?P<SpanID>[0-9]+);o=(?P<Options>.+)$`)
+
+// Deprecated: use [ParseW3CTraceParent] instead.
+func ParseCloudTraceContext(s string) (*CloudTraceContext, error) {
+	match := cloudTraceContextRegex.FindStringSubmatch(s)
+	result := make(map[string]string)
+	for i, name := range cloudTraceContextRegex.SubexpNames() {
+		if i != 0 && name != "" {
+			result[name] = match[i]
+		}
+	}
+	t := &CloudTraceContext{
+		TraceID: result["TraceID"],
+		SpanID:  result["SpanID"],
+		Options: result["Options"],
+	}
+	err := t.Validate()
+	if err != nil {
+		return nil, err
+	}
+	return t, nil
+}
+
+// Based on https://cloud.google.com/trace/docs/trace-context#http-requests.
+// See also https://cloud.google.com/run/docs/trace for Cloud Run specific info.
+//
+// Deprecated: use [W3CTraceParent] instead.
+type CloudTraceContext struct {
+	TraceID string
+	SpanID  string
+	Options string
+}
+
+func (t *CloudTraceContext) Validate() error {
+	if t == nil {
+		return fmt.Errorf("%w: nil CloudTraceContext", ErrInvalidTrace)
+	}
+	if len(t.TraceID) != 32 {
+		return fmt.Errorf("%w: invalid TraceID", ErrInvalidTrace)
+	}
+	if t.SpanID == "" {
+		return fmt.Errorf("%w: invalid SpanID", ErrInvalidTrace)
+	}
+	if t.Options == "" {
+		return fmt.Errorf("%w: invalid Options", ErrInvalidTrace)
+	}
+	return nil
+}
+
+func (t *CloudTraceContext) GetTraceID() string {
+	if t == nil {
+		return ""
+	}
+	return t.TraceID
+}
+
+func (t *CloudTraceContext) GetSpanID() string {
+	if t == nil {
+		return ""
+	}
+	return t.SpanID
+}
+
+func (t *CloudTraceContext) GetOptions() string {
+	if t == nil {
+		return ""
+	}
+	return t.Options
+}
+
+func (t *CloudTraceContext) Sampled() bool {
+	if t == nil {
+		return false
+	}
+	return t.Options == "1"
+}

trace/cloud_trace_test.go

@@ -0,0 +1,21 @@
+package trace_test
+
+import (
+	"testing"
+
+	"github.com/ronny/clog/trace"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCloudTraceContext(t *testing.T) {
+	validExample := "70e0091f6f5d4643bb4eca9d81320c76/97123319527522;o=1"
+
+	tr, err := trace.ParseCloudTraceContext(validExample)
+	if !assert.Nil(t, err) {
+		return
+	}
+
+	assert.Equal(t, "70e0091f6f5d4643bb4eca9d81320c76", tr.GetTraceID())
+	assert.Equal(t, "97123319527522", tr.GetSpanID())
+	assert.Equal(t, true, tr.Sampled())
+}

trace/context.go

@@ -0,0 +1,22 @@
+package trace
+
+import "context"
+
+type ctxKeyType int
+
+const ctxKey ctxKeyType = iota
+
+// NewContext returns a new context derived from ctx with trace attached.
+func NewContext(ctx context.Context, t Trace) context.Context {
+	return context.WithValue(ctx, ctxKey, t)
+}
+
+// FromContext extracts any [Trace] information from ctx and
+// returns it if found, otherwise returns nil.
+func FromContext(ctx context.Context) Trace {
+	t, ok := ctx.Value(ctxKey).(Trace)
+	if !ok {
+		return nil
+	}
+	return t
+}

trace/doc.go

@@ -0,0 +1,3 @@
+// Package trace contains utilities to extract and use Cloud Trace context in
+// logs.
+package trace

trace/errors.go

@@ -0,0 +1,8 @@
+package trace
+
+import "errors"
+
+var (
+	ErrUnparseable  = errors.New("unparseable")
+	ErrInvalidTrace = errors.New("invalid trace")
+)