From 180dac09794d8fe8196156318861a3764e1c651d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal.matczuk@saucelabs.com>
Date: Fri, 9 Feb 2024 13:39:33 +0100
Subject: [PATCH] martian: add TraceID aware logger wrapper

Move trace ID logging logic close to traceID.
Use logger stacking to add features to logs.

This is only temporary and traceID should be moved to the toplevel log package.
To make it work we would need to change the main Logger interface, see #544.
---
 internal/martian/trace.go | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/internal/martian/trace.go b/internal/martian/trace.go
index b0510844..f31b46fa 100644
--- a/internal/martian/trace.go
+++ b/internal/martian/trace.go
@@ -3,9 +3,12 @@
 package martian
 
 import (
+	"context"
 	"fmt"
 	"sync/atomic"
 	"time"
+
+	"github.com/saucelabs/forwarder/internal/martian/log"
 )
 
 // traceID is a unique identifier for a request.
@@ -32,3 +35,28 @@ func (t traceID) String() string {
 func (t traceID) Duration() time.Duration {
 	return time.Since(t.createdAt)
 }
+
+type TraceIDPrependingLogger struct {
+	log.Logger
+}
+
+func (l TraceIDPrependingLogger) Infof(ctx context.Context, format string, args ...any) {
+	l.Logger.Infof(ctx, l.format(ctx, format), args...)
+}
+
+func (l TraceIDPrependingLogger) Debugf(ctx context.Context, format string, args ...any) {
+	l.Logger.Debugf(ctx, l.format(ctx, format), args...)
+}
+
+func (l TraceIDPrependingLogger) Errorf(ctx context.Context, format string, args ...any) {
+	l.Logger.Errorf(ctx, l.format(ctx, format), args...)
+}
+
+var _ log.Logger = TraceIDPrependingLogger{}
+
+func (l TraceIDPrependingLogger) format(ctx context.Context, format string) string {
+	if id := TraceID(ctx); id != "" {
+		return fmt.Sprintf("[%s] %s", id, format)
+	}
+	return format
+}