spanlogger: fix double logging of caller information

spanlogger/spanlogger.go

@@ -1,8 +1,13 @@
+// Provenance-includes-location: https://github.com/go-kit/log/blob/main/value.go
+// Provenance-includes-license: MIT
+// Provenance-includes-copyright: Go kit
+
 package spanlogger
 
 import (
 	"context"
 	"runtime"
+	"strconv"
 	"strings"
 
 	"go.uber.org/atomic" // Really just need sync/atomic but there is a lint rule preventing it.
@@ -163,9 +168,6 @@ func (s *SpanLogger) getLogger() log.Logger {
 		logger = log.With(logger, "trace_id", traceID)
 	}
 
-	// Replace the default valuer for the 'caller' attribute with one that gets the caller of the methods in this file.
-	logger = log.With(logger, "caller", spanLoggerAwareCaller())
-
 	// If the value has been set by another goroutine, fetch that other value and discard the one we made.
 	if !s.logger.CompareAndSwap(nil, &logger) {
 		pLogger := s.logger.Load()
@@ -188,46 +190,64 @@ func (s *SpanLogger) SetSpanAndLogTag(key string, value interface{}) {
 	s.logger.Store(&wrappedLogger)
 }
 
-// spanLoggerAwareCaller is like log.Caller, but ensures that the caller information is
-// that of the caller to SpanLogger, not SpanLogger itself.
-func spanLoggerAwareCaller() log.Valuer {
-	valuer := atomic.NewPointer[log.Valuer](nil)
-
+// Caller is like github.com/go-kit/log's Caller, but ensures that the caller information is
+// that of the caller to SpanLogger (if SpanLogger is being used), not SpanLogger itself.
+//
+// defaultStackDepth should be the number of stack frames to skip by default, as would be
+// passed to github.com/go-kit/log's Caller method.
+func Caller(defaultStackDepth int) log.Valuer {
 	return func() interface{} {
-		// If we've already determined the correct stack depth, use it.
-		existingValuer := valuer.Load()
-		if existingValuer != nil {
-			return (*existingValuer)()
-		}
-
-		// We haven't been called before, determine the correct stack depth to
-		// skip the configured logger's internals and the SpanLogger's internals too.
-		//
-		// Note that we can't do this in spanLoggerAwareCaller() directly because we
-		// need to do this when invoked by the configured logger - otherwise we cannot
-		// measure the stack depth of the logger's internals.
-
-		stackDepth := 3 // log.DefaultCaller uses a stack depth of 3, so start searching for the correct stack depth there.
+		stackDepth := defaultStackDepth + 1 // +1 to account for this method.
+		seenSpanLogger := false
+		pc := make([]uintptr, 1)
 
 		for {
-			_, file, _, ok := runtime.Caller(stackDepth)
+			function, file, line, ok := caller(stackDepth, pc)
 			if !ok {
 				// We've run out of possible stack frames. Give up.
-				valuer.Store(&unknownCaller)
-				return unknownCaller()
+				return "<unknown>"
 			}
 
-			if strings.HasSuffix(file, "spanlogger/spanlogger.go") {
-				stackValuer := log.Caller(stackDepth + 2) // Add one to skip the stack frame for the SpanLogger method, and another to skip the stack frame for the valuer which we'll invoke below.
-				valuer.Store(&stackValuer)
-				return stackValuer()
+			// If we're in a SpanLogger method, we need to continue searching.
+			//
+			// Matching on the exact function name like this does mean this will break if we rename or refactor SpanLogger, but
+			// the tests should catch this. In the worst case scenario, we'll log incorrect caller information, which isn't the
+			// end of the world.
+			if function == "github.com/grafana/dskit/spanlogger.(*SpanLogger).Log" || function == "github.com/grafana/dskit/spanlogger.(*SpanLogger).DebugLog" {
+				seenSpanLogger = true
+				stackDepth++
+				continue
 			}
 
-			stackDepth++
+			// We need to check for go-kit/log stack frames like this because using log.With, log.WithPrefix or log.WithSuffix
+			// (including the various level methods like level.Debug, level.Info etc.) to wrap a SpanLogger introduce an
+			// additional context.Log stack frame that calls into the SpanLogger. This is because the use of SpanLogger
+			// as the logger means the optimisation to avoid creating a new logger in
+			// https://github.com/go-kit/log/blob/c7bf81493e581feca11e11a7672b14be3591ca43/log.go#L141-L146 used by those methods
+			// can't be used, and so the SpanLogger is wrapped in a new logger.
+			if seenSpanLogger && function == "github.com/go-kit/log.(*context).Log" {
+				stackDepth++
+				continue
+			}
+
+			return formatCallerInfoForLog(file, line)
 		}
 	}
 }
 
-var unknownCaller log.Valuer = func() interface{} {
-	return "<unknown>"
+// caller is like runtime.Caller, but modified to allow reuse of the uintptr slice and return the function name.
+func caller(stackDepth int, pc []uintptr) (function string, file string, line int, ok bool) {
+	n := runtime.Callers(stackDepth+1, pc)
+	if n < 1 {
+		return "", "", 0, false
+	}
+
+	frame, _ := runtime.CallersFrames(pc).Next()
+	return frame.Function, frame.File, frame.Line, frame.PC != 0
+}
+
+// This is based on github.com/go-kit/log's Caller, but modified for use by Caller above.
+func formatCallerInfoForLog(file string, line int) string {
+	idx := strings.LastIndexByte(file, '/')
+	return file[idx+1:] + ":" + strconv.Itoa(line)
 }

spanlogger/spanlogger_test.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"path/filepath"
 	"runtime"
+	"slices"
 	"strings"
 	"testing"
 
@@ -45,9 +46,6 @@ func TestSpanLogger_CustomLogger(t *testing.T) {
 	}
 	resolver := fakeResolver{}
 
-	_, thisFile, thisLineNumber, ok := runtime.Caller(0)
-	require.True(t, ok)
-
 	span, ctx := New(context.Background(), logger, "test", resolver)
 	_ = span.Log("msg", "original spanlogger")
 
@@ -58,9 +56,9 @@ func TestSpanLogger_CustomLogger(t *testing.T) {
 	_ = span.Log("msg", "fallback spanlogger")
 
 	expect := [][]interface{}{
-		{"method", "test", "caller", toCallerInfo(thisFile, thisLineNumber+4), "msg", "original spanlogger"},
-		{"caller", toCallerInfo(thisFile, thisLineNumber+7), "msg", "restored spanlogger"},
-		{"caller", toCallerInfo(thisFile, thisLineNumber+10), "msg", "fallback spanlogger"},
+		{"method", "test", "msg", "original spanlogger"},
+		{"msg", "restored spanlogger"},
+		{"msg", "fallback spanlogger"},
 	}
 	require.Equal(t, expect, logged)
 }
@@ -88,9 +86,6 @@ func TestSpanLogger_SetSpanAndLogTag(t *testing.T) {
 		return nil
 	}
 
-	_, thisFile, thisLineNumber, ok := runtime.Caller(0)
-	require.True(t, ok)
-
 	spanLogger, _ := New(context.Background(), logger, "the_method", fakeResolver{})
 	require.NoError(t, spanLogger.Log("msg", "this is the first message"))
 
@@ -110,18 +105,15 @@ func TestSpanLogger_SetSpanAndLogTag(t *testing.T) {
 	expectedLogMessages := [][]interface{}{
 		{
 			"method", "the_method",
-			"caller", toCallerInfo(thisFile, thisLineNumber+4),
 			"msg", "this is the first message",
 		},
 		{
 			"method", "the_method",
-			"caller", toCallerInfo(thisFile, thisLineNumber+7),
 			"id", "123",
 			"msg", "this is the second message",
 		},
 		{
 			"method", "the_method",
-			"caller", toCallerInfo(thisFile, thisLineNumber+10),
 			"id", "123",
 			"more context", "abc",
 			"msg", "this is the third message",
@@ -206,7 +198,7 @@ func BenchmarkSpanLoggerWithRealLogger(b *testing.B) {
 		b.Run(name, func(b *testing.B) {
 			buf := bytes.NewBuffer(nil)
 			logger := dskit_log.NewGoKitWithWriter("logfmt", buf)
-			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", log.Caller(5))
+			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", Caller(5))
 
 			if debugEnabled {
 				logger = level.NewFilter(logger, level.AllowAll())
@@ -222,7 +214,7 @@ func BenchmarkSpanLoggerWithRealLogger(b *testing.B) {
 			resolver := fakeResolver{}
 			sl, _ := New(context.Background(), logger, "test", resolver, "bar")
 
-			b.Run("log", func(b *testing.B) {
+			b.Run("Log", func(b *testing.B) {
 				buf.Reset()
 				b.ResetTimer()
 
@@ -231,7 +223,7 @@ func BenchmarkSpanLoggerWithRealLogger(b *testing.B) {
 				}
 			})
 
-			b.Run("level.debug", func(b *testing.B) {
+			b.Run("level.Debug", func(b *testing.B) {
 				buf.Reset()
 				b.ResetTimer()
 
@@ -240,7 +232,7 @@ func BenchmarkSpanLoggerWithRealLogger(b *testing.B) {
 				}
 			})
 
-			b.Run("debuglog", func(b *testing.B) {
+			b.Run("DebugLog", func(b *testing.B) {
 				buf.Reset()
 				b.ResetTimer()
 
@@ -250,7 +242,31 @@ func BenchmarkSpanLoggerWithRealLogger(b *testing.B) {
 			})
 		})
 	}
+}
+
+func BenchmarkSpanLoggerAwareCaller(b *testing.B) {
+	runBenchmark := func(b *testing.B, caller log.Valuer) {
+		buf := bytes.NewBuffer(nil)
+		logger := dskit_log.NewGoKitWithWriter("logfmt", buf)
+		logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", caller)
 
+		b.ResetTimer()
+
+		for i := 0; i < b.N; i++ {
+			_ = logger.Log("msg", "foo", "more", "data")
+		}
+
+	}
+
+	const defaultStackDepth = 5
+
+	b.Run("with go-kit's Caller", func(b *testing.B) {
+		runBenchmark(b, log.Caller(defaultStackDepth))
+	})
+
+	b.Run("with dskit's spanlogger.Caller", func(b *testing.B) {
+		runBenchmark(b, Caller(defaultStackDepth))
+	})
 }
 
 // Logger which does nothing and implements the DebugEnabled interface used by SpanLogger.
@@ -267,24 +283,32 @@ type loggerWithDebugEnabled struct {
 
 func (l loggerWithDebugEnabled) DebugEnabled() bool { return l.debugEnabled }
 
-func TestSpanLogger_CallerInfo(t *testing.T) {
+func TestSpanLoggerAwareCaller(t *testing.T) {
 	testCases := map[string]func(w io.Writer) log.Logger{
 		// This is based on Mimir's default logging configuration: https://github.com/grafana/mimir/blob/50d1c27b4ad82b265ff5a865345bec2d726f64ef/pkg/util/log/log.go#L45-L46
 		"default logger": func(w io.Writer) log.Logger {
 			logger := dskit_log.NewGoKitWithWriter("logfmt", w)
-			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", log.Caller(5))
+			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", Caller(5))
 			logger = level.NewFilter(logger, level.AllowAll())
 			return logger
 		},
 
 		// This is based on Mimir's logging configuration with rate-limiting enabled: https://github.com/grafana/mimir/blob/50d1c27b4ad82b265ff5a865345bec2d726f64ef/pkg/util/log/log.go#L42-L43
 		"rate-limited logger": func(w io.Writer) log.Logger {
 			logger := dskit_log.NewGoKitWithWriter("logfmt", w)
-			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", log.Caller(6))
+			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", Caller(6))
 			logger = dskit_log.NewRateLimitedLogger(logger, 1000, 1000, nil)
 			logger = level.NewFilter(logger, level.AllowAll())
 			return logger
 		},
+
+		"default logger that has been wrapped with further information": func(w io.Writer) log.Logger {
+			logger := dskit_log.NewGoKitWithWriter("logfmt", w)
+			logger = log.With(logger, "ts", log.DefaultTimestampUTC, "caller", Caller(5))
+			logger = level.NewFilter(logger, level.AllowAll())
+			logger = log.With(logger, "user", "user-1")
+			return logger
+		},
 	}
 
 	resolver := fakeResolver{}
@@ -307,12 +331,15 @@ func TestSpanLogger_CallerInfo(t *testing.T) {
 		return buf, spanLogger, span.(*jaeger.Span)
 	}
 
-	requireSpanHasTwoLogLinesWithoutCaller := func(t *testing.T, span *jaeger.Span) {
+	requireSpanHasTwoLogLinesWithoutCaller := func(t *testing.T, span *jaeger.Span, extraFields ...otlog.Field) {
 		logs := span.Logs()
 		require.Len(t, logs, 2)
 
-		require.Equal(t, []otlog.Field{otlog.String("msg", "this is a test")}, logs[0].Fields)
-		require.Equal(t, []otlog.Field{otlog.String("msg", "this is another test")}, logs[1].Fields)
+		expectedFields := append(slices.Clone(extraFields), otlog.String("msg", "this is a test"))
+		require.Equal(t, expectedFields, logs[0].Fields)
+
+		expectedFields = append(slices.Clone(extraFields), otlog.String("msg", "this is another test"))
+		require.Equal(t, expectedFields, logs[1].Fields)
 	}
 
 	for name, loggerFactory := range testCases {
@@ -326,6 +353,7 @@ func TestSpanLogger_CallerInfo(t *testing.T) {
 
 				logged := logs.String()
 				require.Contains(t, logged, "caller="+toCallerInfo(thisFile, lineNumberTwoLinesBeforeFirstLogCall+2))
+				require.Equalf(t, 1, strings.Count(logged, "caller="), "expected to only have one caller field, but got: %v", logged)
 
 				logs.Reset()
 				_, _, lineNumberTwoLinesBeforeSecondLogCall, ok := runtime.Caller(0)
@@ -334,6 +362,7 @@ func TestSpanLogger_CallerInfo(t *testing.T) {
 
 				logged = logs.String()
 				require.Contains(t, logged, "caller="+toCallerInfo(thisFile, lineNumberTwoLinesBeforeSecondLogCall+2))
+				require.Equalf(t, 1, strings.Count(logged, "caller="), "expected to only have one caller field, but got: %v", logged)
 
 				requireSpanHasTwoLogLinesWithoutCaller(t, span)
 			})
@@ -346,6 +375,7 @@ func TestSpanLogger_CallerInfo(t *testing.T) {
 
 				logged := logs.String()
 				require.Contains(t, logged, "caller="+toCallerInfo(thisFile, lineNumberTwoLinesBeforeLogCall+2))
+				require.Equalf(t, 1, strings.Count(logged, "caller="), "expected to only have one caller field, but got: %v", logged)
 
 				logs.Reset()
 				_, _, lineNumberTwoLinesBeforeSecondLogCall, ok := runtime.Caller(0)
@@ -354,9 +384,33 @@ func TestSpanLogger_CallerInfo(t *testing.T) {
 
 				logged = logs.String()
 				require.Contains(t, logged, "caller="+toCallerInfo(thisFile, lineNumberTwoLinesBeforeSecondLogCall+2))
+				require.Equalf(t, 1, strings.Count(logged, "caller="), "expected to only have one caller field, but got: %v", logged)
 
 				requireSpanHasTwoLogLinesWithoutCaller(t, span)
 			})
+
+			t.Run("logging with SpanLogger wrapped in a level", func(t *testing.T) {
+				logs, spanLogger, span := setupTest(t, loggerFactory)
+
+				_, thisFile, lineNumberTwoLinesBeforeFirstLogCall, ok := runtime.Caller(0)
+				require.True(t, ok)
+				_ = level.Info(spanLogger).Log("msg", "this is a test")
+
+				logged := logs.String()
+				require.Contains(t, logged, "caller="+toCallerInfo(thisFile, lineNumberTwoLinesBeforeFirstLogCall+2))
+				require.Equalf(t, 1, strings.Count(logged, "caller="), "expected to only have one caller field, but got: %v", logged)
+
+				logs.Reset()
+				_, _, lineNumberTwoLinesBeforeSecondLogCall, ok := runtime.Caller(0)
+				require.True(t, ok)
+				_ = level.Info(spanLogger).Log("msg", "this is another test")
+
+				logged = logs.String()
+				require.Contains(t, logged, "caller="+toCallerInfo(thisFile, lineNumberTwoLinesBeforeSecondLogCall+2))
+				require.Equalf(t, 1, strings.Count(logged, "caller="), "expected to only have one caller field, but got: %v", logged)
+
+				requireSpanHasTwoLogLinesWithoutCaller(t, span, otlog.String("level", "info"))
+			})
 		})
 	}
 }