Allow customfields to record any metadata type

`customfields` is an utility helper for attaching custom fields to gRPC
logs. Recently, it supports logging metadata whose keys and values are
string types only. This commit expand its ability to log data with any
type. Eventually, the custom fields are converted to `logrus.Fields`
type, which support any type already.

2 files changed, 48 insertions, 18 deletions

diff --git a/internal/log/customfields.go b/internal/log/customfields.go
index 1622810f2..38874e15c 100644
--- a/internal/log/customfields.go
+++ b/internal/log/customfields.go
@@ -14,7 +14,7 @@ type requestCustomFieldsKey struct{}
 // the object out with CustomFieldsFromContext.
 type CustomFields struct {
 	numericFields map[string]int
-	stringFields  map[string]string
+	anyFields     map[string]any
 	sync.Mutex
 }
 
@@ -45,11 +45,11 @@ func (fields *CustomFields) RecordMax(key string, value int) {
 }
 
 // RecordMetadata records a string metadata for the given key.
-func (fields *CustomFields) RecordMetadata(key string, value string) {
+func (fields *CustomFields) RecordMetadata(key string, value any) {
 	fields.Lock()
 	defer fields.Unlock()
 
-	fields.stringFields[key] = value
+	fields.anyFields[key] = value
 }
 
 // Fields returns all the fields as logrus.Fields
@@ -61,7 +61,7 @@ func (fields *CustomFields) Fields() logrus.Fields {
 	for k, v := range fields.numericFields {
 		f[k] = v
 	}
-	for k, v := range fields.stringFields {
+	for k, v := range fields.anyFields {
 		f[k] = v
 	}
 	return f
@@ -77,6 +77,6 @@ func CustomFieldsFromContext(ctx context.Context) *CustomFields {
 func InitContextCustomFields(ctx context.Context) context.Context {
 	return context.WithValue(ctx, requestCustomFieldsKey{}, &CustomFields{
 		numericFields: make(map[string]int),
-		stringFields:  make(map[string]string),
+		anyFields:     make(map[string]any),
 	})
 }
diff --git a/internal/log/customfields_test.go b/internal/log/customfields_test.go
index 07577b17c..b2d89b707 100644
--- a/internal/log/customfields_test.go
+++ b/internal/log/customfields_test.go
@@ -1,6 +1,7 @@
 package log_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/sirupsen/logrus"
@@ -73,17 +74,46 @@ func TestStatsFromContext_RecordMax(t *testing.T) {
 }
 
 func TestStatsFromContext_RecordMetadata(t *testing.T) {
-	ctx := testhelper.Context(t)
-
-	ctx = log.InitContextCustomFields(ctx)
-
-	customFields := log.CustomFieldsFromContext(ctx)
-
-	customFields.RecordMetadata("foo", "bar")
-	require.NotNil(t, customFields)
-	require.Equal(t, customFields.Fields(), logrus.Fields{"foo": "bar"})
-
-	customFields.RecordMetadata("foo", "baz") // override the existing value
-	require.NotNil(t, customFields)
-	require.Equal(t, customFields.Fields(), logrus.Fields{"foo": "baz"})
+	for _, tc := range []struct {
+		desc           string
+		setup          func(context.Context)
+		expectedFields logrus.Fields
+	}{
+		{
+			desc: "record a string metadata",
+			setup: func(ctx context.Context) {
+				customFields := log.CustomFieldsFromContext(ctx)
+				customFields.RecordMetadata("foo", "bar")
+			},
+			expectedFields: logrus.Fields{"foo": "bar"},
+		},
+		{
+			desc: "override metadata of the same key",
+			setup: func(ctx context.Context) {
+				customFields := log.CustomFieldsFromContext(ctx)
+				customFields.RecordMetadata("foo", "bar")
+				customFields.RecordMetadata("foo", "baz") // override the existing value
+			},
+			expectedFields: logrus.Fields{"foo": "baz"},
+		},
+		{
+			desc: "record metadata with different types",
+			setup: func(ctx context.Context) {
+				customFields := log.CustomFieldsFromContext(ctx)
+				customFields.RecordMetadata("hello", 1234)
+				customFields.RecordMetadata("hi", []int{1, 2, 3, 4})
+			},
+			expectedFields: logrus.Fields{
+				"hello": 1234,
+				"hi":    []int{1, 2, 3, 4},
+			},
+		},
+	} {
+		ctx := log.InitContextCustomFields(testhelper.Context(t))
+		tc.setup(ctx)
+
+		customFields := log.CustomFieldsFromContext(ctx)
+		require.NotNil(t, customFields)
+		require.Equal(t, customFields.Fields(), tc.expectedFields)
+	}
 }