From f34053bdcc36a821840af4c83c915714fb83d105 Mon Sep 17 00:00:00 2001
From: Olivier Cano <ocano@scaleway.com>
Date: Fri, 20 Dec 2024 10:48:04 +0100
Subject: [PATCH] logging: add AddFields

Support the similar helpers from the v1 introduced in #91:
- [ctxzap.AddFields](https://github.com/grpc-ecosystem/go-grpc-middleware/blob/v1/logging/zap/ctxzap/context.go#L23-L32)
- [ctxlogrus.AddFields](https://github.com/grpc-ecosystem/go-grpc-middleware/blob/v1/logging/logrus/ctxlogrus/context.go#L22-L30)
- [ctxkit.AddFields](https://github.com/grpc-ecosystem/go-grpc-middleware/blob/v1/logging/kit/ctxkit/context.go#L22-L28)

Signed-off-by: Olivier Cano <ocano@scaleway.com>
---
 interceptors/logging/logging.go      | 25 ++++++++++++++++++++-----
 interceptors/logging/logging_test.go | 11 +++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/interceptors/logging/logging.go b/interceptors/logging/logging.go
index 510344fd7..c41317d35 100644
--- a/interceptors/logging/logging.go
+++ b/interceptors/logging/logging.go
@@ -22,6 +22,9 @@ var (
 )
 
 type fieldsCtxMarker struct{}
+type fieldsCtxValue struct {
+	fields Fields
+}
 
 var (
 	// fieldsCtxMarkerKey is the Context value marker that is used by logging middleware to read and write logging fields into context.
@@ -161,23 +164,23 @@ NextAddField:
 // If there are no fields in the context, it returns an empty Fields value.
 // Extracted fields are useful to construct your own logger that has fields from gRPC interceptors.
 func ExtractFields(ctx context.Context) Fields {
-	t, ok := ctx.Value(fieldsCtxMarkerKey).(Fields)
+	t, ok := ctx.Value(fieldsCtxMarkerKey).(*fieldsCtxValue)
 	if !ok {
 		return nil
 	}
-	n := make(Fields, len(t))
-	copy(n, t)
+	n := make(Fields, len(t.fields))
+	copy(n, t.fields)
 	return n
 }
 
-// InjectFields allows adding fields to any existing Fields that will be used by the logging interceptor or can be
+// InjectFields returns a new context with merged fields that will be used by the logging interceptor or can be
 // extracted further in ExtractFields.
 // For explicitness, in case of duplicates, the newest field occurrence wins. This allows nested components to update
 // popular fields like grpc.component (e.g. server invoking gRPC client).
 //
 // Don't overuse mutation of fields to avoid surprises.
 func InjectFields(ctx context.Context, f Fields) context.Context {
-	return context.WithValue(ctx, fieldsCtxMarkerKey, f.WithUnique(ExtractFields(ctx)))
+	return context.WithValue(ctx, fieldsCtxMarkerKey, &fieldsCtxValue{fields: f.WithUnique(ExtractFields(ctx))})
 }
 
 // InjectLogField is like InjectFields, just for one field.
@@ -185,6 +188,18 @@ func InjectLogField(ctx context.Context, key string, val any) context.Context {
 	return InjectFields(ctx, Fields{key, val})
 }
 
+// AddFields updates the fields already in the context that will be used by the logging interceptor or can be
+// extracted further in ExtractFields. For explicitness, in case of duplicates, the newest field occurrence wins.
+//
+// This function is not safe to call concurrently.
+func AddFields(ctx context.Context, f Fields) {
+	t, ok := ctx.Value(fieldsCtxMarkerKey).(*fieldsCtxValue)
+	if !ok {
+		return
+	}
+	t.fields = f.AppendUnique(t.fields)
+}
+
 // Logger requires Log method, similar to experimental slog, allowing logging interceptor to be interoperable. Official
 // adapters for popular loggers are in `provider/` directory (separate modules). It's totally ok to copy simple function
 // implementation over.
diff --git a/interceptors/logging/logging_test.go b/interceptors/logging/logging_test.go
index 150186f1f..f7ac27ccf 100644
--- a/interceptors/logging/logging_test.go
+++ b/interceptors/logging/logging_test.go
@@ -49,3 +49,14 @@ func TestFieldsDelete(t *testing.T) {
 	f.Delete("c")
 	require.Equal(t, Fields{}, f)
 }
+
+func TestAddFields(t *testing.T) {
+	c := InjectFields(context.Background(), Fields{"a", "2", "c", "3"})
+	f := ExtractFields(c)
+	require.Equal(t, Fields{"a", "2", "c", "3"}, f)
+	AddFields(c, Fields{"a", "1", "b", "2"})
+
+	// First context should have updated values.
+	f = ExtractFields(c)
+	require.Equal(t, Fields{"a", "1", "b", "2", "c", "3"}, f)
+}