Logging: Add WithContextualAttributes to pass log params based on the given context (#74428)

* suggestion to use `RegisterContextualLogProvider`

* add `pluginId`

* change to `WithContextualAttributes`

* move approach to instrumentation.go

* improve `WithContextualAttributes`

* unexport consts

* typo

* remove comment

* add `nil` check

pkg/infra/log/log.go

@@ -51,6 +51,14 @@ func init() {
 	}
 	logger := level.NewFilter(format(os.Stderr), level.AllowInfo())
 	root = newManager(logger)
+
+	RegisterContextualLogProvider(func(ctx context.Context) ([]any, bool) {
+		pFromCtx := ctx.Value(logParamsContextKey{})
+		if pFromCtx != nil {
+			return pFromCtx.([]any), true
+		}
+		return nil, false
+	})
 }
 
 // logManager manage loggers
@@ -270,6 +278,18 @@ func RegisterContextualLogProvider(mw ContextualLogProviderFunc) {
 	ctxLogProviders = append(ctxLogProviders, mw)
 }
 
+type logParamsContextKey struct{}
+
+// WithContextualAttributes adds contextual attributes to the logger based on the given context.
+// That allows loggers further down the chain to automatically log those attributes.
+func WithContextualAttributes(ctx context.Context, logParams []any) context.Context {
+	p := logParams
+	if ctx.Value(logParamsContextKey{}) != nil {
+		p = append(ctx.Value(logParamsContextKey{}).([]any), logParams...)
+	}
+	return context.WithValue(ctx, logParamsContextKey{}, p)
+}
+
 var logLevels = map[string]level.Option{
 	"trace":    level.AllowDebug(),
 	"debug":    level.AllowDebug(),

pkg/infra/log/log_test.go

@@ -267,6 +267,84 @@ func TestGetFilters(t *testing.T) {
 	})
 }
 
+func TestWithContextualAttributes_appendsContext(t *testing.T) {
+	t.Run("Logs arguments from context", func(t *testing.T) {
+		scenario := newLoggerScenario(t, false)
+
+		// logs `"k1", "v1"` with the first context
+		ctx := context.Background()
+		ctx = WithContextualAttributes(ctx, []any{"k1", "v1"})
+		ls := New("test").FromContext(ctx)
+
+		ls.Info("hello", "k2", "v2")
+
+		require.Len(t, scenario.loggedArgs, 1)
+		scenario.ValidateLineEquality(t, 0, []any{
+			"logger", "test",
+			"k1", "v1",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k2", "v2",
+		})
+	})
+	t.Run("Does not log arguments from different context", func(t *testing.T) {
+		scenario := newLoggerScenario(t, false)
+
+		// logs `"k1", "v1"` with the first context
+		ctx := context.Background()
+		ctx = WithContextualAttributes(ctx, []any{"k1", "v1"})
+		ls := New("test").FromContext(ctx)
+
+		ls.Info("hello", "k2", "v2")
+
+		require.Len(t, scenario.loggedArgs, 1)
+		scenario.ValidateLineEquality(t, 0, []any{
+			"logger", "test",
+			"k1", "v1",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k2", "v2",
+		})
+		// does not log `"k1", "v1"` with the new context
+		ctx = context.Background()
+		ls = New("test").FromContext(ctx)
+
+		ls.Info("hello", "k2", "v2")
+
+		require.Len(t, scenario.loggedArgs, 2)
+		scenario.ValidateLineEquality(t, 1, []any{
+			"logger", "test",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k2", "v2",
+		})
+	})
+	t.Run("Appends arguments set previously", func(t *testing.T) {
+		scenario := newLoggerScenario(t, false)
+
+		ctx := context.Background()
+		ctx = WithContextualAttributes(ctx, []any{"k1", "v1"})
+		ctx = WithContextualAttributes(ctx, []any{"k2", "v2"})
+		ls := New("test").FromContext(ctx)
+
+		ls.Info("hello", "k3", "v3")
+
+		require.Len(t, scenario.loggedArgs, 1)
+		scenario.ValidateLineEquality(t, 0, []any{
+			"logger", "test",
+			"k1", "v1",
+			"k2", "v2",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k3", "v3",
+		})
+	})
+}
+
 type scenarioContext struct {
 	loggedArgs [][]any
 	mockedTime time.Time
@@ -287,7 +365,12 @@ func (s *scenarioContext) ValidateLineEquality(t testing.TB, n int, expected []a
 	}
 }
 
-func newLoggerScenario(t testing.TB) *scenarioContext {
+func newLoggerScenario(t testing.TB, resetCtxLogProviders ...bool) *scenarioContext {
+	clearProviders := true
+	if len(resetCtxLogProviders) > 0 {
+		clearProviders = resetCtxLogProviders[0]
+	}
+
 	t.Helper()
 
 	scenario := &scenarioContext{
@@ -308,11 +391,13 @@ func newLoggerScenario(t testing.TB) *scenarioContext {
 		now = origNow
 	})
 
-	origContextHandlers := ctxLogProviders
-	ctxLogProviders = []ContextualLogProviderFunc{}
-	t.Cleanup(func() {
-		ctxLogProviders = origContextHandlers
-	})
+	if clearProviders {
+		origContextHandlers := ctxLogProviders
+		ctxLogProviders = []ContextualLogProviderFunc{}
+		t.Cleanup(func() {
+			ctxLogProviders = origContextHandlers
+		})
+	}
 
 	root = newManager(l)
 	return scenario