Race condition in logger.SetConfig (#16)

* remove globalconfig variable to avoid data race

* Add WithRoutineLabel

* rename

* simplify logger a bit

* format

* doc

config/section.go

@@ -199,6 +199,10 @@ func (r *section) SetConfig(c Config) error {
 	r.lockObj.Lock()
 	defer r.lockObj.Unlock()
 
+	if reflect.TypeOf(c).Kind() != reflect.Ptr {
+		return fmt.Errorf("config must be a Pointer")
+	}
+
 	if !DeepEqual(r.config, c) {
 		r.config = c
 		r.isDirty.Store(true)

contextutils/context.go

@@ -4,22 +4,24 @@ package contextutils
 import (
 	"context"
 	"fmt"
+	"runtime/pprof"
 )
 
 type Key string
 
 const (
-	AppNameKey    Key = "app_name"
-	NamespaceKey  Key = "ns"
-	TaskTypeKey   Key = "tasktype"
-	ProjectKey    Key = "project"
-	DomainKey     Key = "domain"
-	WorkflowIDKey Key = "wf"
-	NodeIDKey     Key = "node"
-	TaskIDKey     Key = "task"
-	ExecIDKey     Key = "exec_id"
-	JobIDKey      Key = "job_id"
-	PhaseKey      Key = "phase"
+	AppNameKey      Key = "app_name"
+	NamespaceKey    Key = "ns"
+	TaskTypeKey     Key = "tasktype"
+	ProjectKey      Key = "project"
+	DomainKey       Key = "domain"
+	WorkflowIDKey   Key = "wf"
+	NodeIDKey       Key = "node"
+	TaskIDKey       Key = "task"
+	ExecIDKey       Key = "exec_id"
+	JobIDKey        Key = "job_id"
+	PhaseKey        Key = "phase"
+	RoutineLabelKey Key = "routine"
 )
 
 func (k Key) String() string {
@@ -35,6 +37,7 @@ var logKeys = []Key{
 	WorkflowIDKey,
 	TaskTypeKey,
 	PhaseKey,
+	RoutineLabelKey,
 }
 
 // Gets a new context with namespace set.
@@ -98,6 +101,14 @@ func WithTaskType(ctx context.Context, taskType string) context.Context {
 	return context.WithValue(ctx, TaskTypeKey, taskType)
 }
 
+// Gets a new context with Go Routine label key set and a label assigned to the context using pprof.Labels.
+// You can then call pprof.SetGoroutineLabels(ctx) to annotate the current go-routine and have that show up in
+// pprof analysis.
+func WithGoroutineLabel(ctx context.Context, routineLabel string) context.Context {
+	ctx = pprof.WithLabels(ctx, pprof.Labels(RoutineLabelKey.String(), routineLabel))
+	return context.WithValue(ctx, RoutineLabelKey, routineLabel)
+}
+
 func addFieldIfNotNil(ctx context.Context, m map[string]interface{}, fieldKey Key) {
 	val := ctx.Value(fieldKey)
 	if val != nil {
@@ -110,6 +121,7 @@ func addStringFieldWithDefaults(ctx context.Context, m map[string]string, fieldK
 	if val == nil {
 		val = ""
 	}
+
 	m[fieldKey.String()] = val.(string)
 }
 
@@ -120,6 +132,7 @@ func GetLogFields(ctx context.Context) map[string]interface{} {
 	for _, k := range logKeys {
 		addFieldIfNotNil(ctx, res, k)
 	}
+
 	return res
 }
 
@@ -128,6 +141,7 @@ func Value(ctx context.Context, key Key) string {
 	if val != nil {
 		return val.(string)
 	}
+
 	return ""
 }
 
@@ -136,5 +150,6 @@ func Values(ctx context.Context, keys ...Key) map[string]string {
 	for _, k := range keys {
 		addStringFieldWithDefaults(ctx, res, k)
 	}
+
 	return res
 }

contextutils/context_test.go

@@ -2,6 +2,7 @@ package contextutils
 
 import (
 	"context"
+	"runtime/pprof"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -111,3 +112,12 @@ func TestValues(t *testing.T) {
 	assert.Equal(t, "flyte", m[WorkflowIDKey.String()])
 	assert.Equal(t, "", m[ProjectKey.String()])
 }
+
+func TestWithGoroutineLabel(t *testing.T) {
+	ctx := context.Background()
+	ctx = WithGoroutineLabel(ctx, "my_routine_123")
+	pprof.SetGoroutineLabels(ctx)
+	m := Values(ctx, RoutineLabelKey)
+	assert.Equal(t, 1, len(m))
+	assert.Equal(t, "my_routine_123", m[RoutineLabelKey.String()])
+}

logger/config.go

@@ -21,12 +21,18 @@ const (
 	jsonDataKey string = "json"
 )
 
-var defaultConfig = &Config{
-	Formatter: FormatterConfig{
-		Type: FormatterJSON,
-	},
-	Level: InfoLevel,
-}
+var (
+	defaultConfig = &Config{
+		Formatter: FormatterConfig{
+			Type: FormatterJSON,
+		},
+		Level: InfoLevel,
+	}
+
+	configSection = config.MustRegisterSectionWithUpdates(configSectionKey, defaultConfig, func(ctx context.Context, newValue config.Config) {
+		onConfigUpdated(*newValue.(*Config))
+	})
+)
 
 // Global logger config.
 type Config struct {
@@ -47,13 +53,18 @@ type FormatterConfig struct {
 	Type FormatterType `json:"type" pflag:",Sets logging format type."`
 }
 
-var globalConfig = Config{}
-
 // Sets global logger config
-func SetConfig(cfg Config) {
-	globalConfig = cfg
+func SetConfig(cfg *Config) error {
+	if err := configSection.SetConfig(cfg); err != nil {
+		return err
+	}
 
-	onConfigUpdated(cfg)
+	onConfigUpdated(*cfg)
+	return nil
+}
+
+func GetConfig() *Config {
+	return configSection.GetConfig().(*Config)
 }
 
 // Level type.
@@ -78,9 +89,3 @@ const (
 	// DebugLevel level. Usually only enabled when debugging. Very verbose logging.
 	DebugLevel
 )
-
-func init() {
-	config.MustRegisterSectionWithUpdates(configSectionKey, defaultConfig, func(ctx context.Context, newValue config.Config) {
-		SetConfig(*newValue.(*Config))
-	})
-}

logger/config_test.go

@@ -1,10 +1,14 @@
 package logger
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
 
 func TestSetConfig(t *testing.T) {
 	type args struct {
-		cfg Config
+		cfg *Config
 	}
 	tests := []struct {
 		name string
@@ -14,7 +18,7 @@ func TestSetConfig(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			SetConfig(tt.args.cfg)
+			assert.NoError(t, SetConfig(tt.args.cfg))
 		})
 	}
 }