Create way to add key-value pairs to new sinks (#4582)

Counterintuitively, a new sink added to an existing logger will not inherit the existing logger's key-value pairs. This has only caused a problem in practice with respect to a small number of known key-value pairs, so this PR creates a way to add specific key-value pairs to new sinks.

This supersedes #4579.

Author: rosecodym

pkg/log/log.go

@@ -136,18 +136,34 @@ func AddSentry(l logr.Logger, opts sentry.ClientOptions, tags map[string]string)
 	return AddSink(l, WithSentry(opts, tags))
 }
 
-// AddSink extends an existing logr.Logger with a new sink. It returns the new
-// logr.Logger, a cleanup function, and an error.
-func AddSink(l logr.Logger, sink logConfig) (logr.Logger, func() error, error) {
+// AddSink extends an existing logr.Logger with a new sink. It returns the new logr.Logger, a cleanup function, and an
+// error.
+//
+// The new sink will not inherit any of the existing logger's key-value pairs. Key-value pairs can be added to the new
+// sink specifically by passing them to this function.
+func AddSink(l logr.Logger, sink logConfig, keysAndValues ...any) (logr.Logger, func() error, error) {
 	if sink.err != nil {
 		return l, nil, sink.err
 	}
+
+	// New key-value pairs cannot be ergonomically added directly to cores. logr has code to do it, but that code is not
+	// exported. Rather than replicating it ourselves, we indirectly use it by creating a temporary logger for the new
+	// core, adding the key-value pairs to the temporary logger, and then extracting the temporary logger's modified
+	// core.
+	newSinkLogger := zapr.NewLogger(zap.New(sink.core))
+	newSinkLogger = newSinkLogger.WithValues(keysAndValues...)
+	newCoreLogger, err := getZapLogger(newSinkLogger)
+	if err != nil {
+		return l, nil, fmt.Errorf("error setting up new key-value pairs: %w", err)
+	}
+	newSinkCore := newCoreLogger.Core()
+
 	zapLogger, err := getZapLogger(l)
 	if err != nil {
 		return l, nil, errors.New("unsupported logr implementation")
 	}
 	zapLogger = zapLogger.WithOptions(zap.WrapCore(func(core zapcore.Core) zapcore.Core {
-		return zapcore.NewTee(core, sink.core)
+		return zapcore.NewTee(core, newSinkCore)
 	}))
 	return zapr.NewLogger(zapLogger), firstErrorFunc(zapLogger.Sync, sink.cleanup), nil
 }

pkg/log/log_test.go

@@ -163,6 +163,29 @@ func TestAddSink(t *testing.T) {
 	assert.Contains(t, buf2.String(), "line 2")
 }
 
+func TestAddSink_WithKeyValuePairs(t *testing.T) {
+	// Arrange: Create a logger with a key-value pair
+	var buf1 bytes.Buffer
+	logger, cleanup := New("service-name", WithConsoleSink(&buf1))
+	t.Cleanup(func() { _ = cleanup })
+	logger = logger.WithValues("sink 1 key", "sink 1 value")
+
+	// Arrange: Add a second sink with a new key-value pair
+	var buf2 bytes.Buffer
+	logger, flush, err := AddSink(logger, WithConsoleSink(&buf2), "sink 2 key", "sink 2 value")
+	require.NoError(t, err)
+
+	// Act
+	logger.Info("something")
+	require.NoError(t, flush())
+
+	// Assert: Confirm that each sink received only its own key-value pair
+	assert.Contains(t, buf1.String(), "sink 1")
+	assert.NotContains(t, buf1.String(), "sink 2")
+	assert.Contains(t, buf2.String(), "sink 2")
+	assert.NotContains(t, buf2.String(), "sink 1")
+}
+
 func TestStaticLevelSink(t *testing.T) {
 	var buf1, buf2 bytes.Buffer
 	l1 := zap.NewAtomicLevel()