Make global logging reconfiguration race-safe

internal/logging/logging.go

@@ -54,10 +54,48 @@ var (
 	baseWriter    io.Writer = os.Stderr
 	baseComponent string
 	fileCloser    io.Closer
+	globalSink    = &dynamicWriter{writer: os.Stderr}
 
 	defaultTimeFmt = time.RFC3339
 )
 
+type dynamicWriter struct {
+	mu     sync.RWMutex
+	writer io.Writer
+}
+
+func (w *dynamicWriter) SetWriter(writer io.Writer) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if writer == nil {
+		writer = io.Discard
+	}
+	w.writer = writer
+}
+
+func (w *dynamicWriter) Write(p []byte) (int, error) {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+
+	if w.writer == nil {
+		return len(p), nil
+	}
+	return w.writer.Write(p)
+}
+
+type componentHook struct{}
+
+func (componentHook) Run(event *zerolog.Event, _ zerolog.Level, _ string) {
+	mu.RLock()
+	component := baseComponent
+	mu.RUnlock()
+
+	if component != "" {
+		event.Str("component", component)
+	}
+}
+
 var (
 	nowFn           = time.Now
 	isTerminalFn    = term.IsTerminal
@@ -85,11 +123,18 @@ var (
 
 func init() {
 	baseWriter = io.MultiWriter(baseWriter, GetBroadcaster())
-	baseLogger = zerolog.New(baseWriter).With().Timestamp().Logger()
+	globalSink.SetWriter(baseWriter)
+	zerolog.TimeFieldFormat = defaultTimeFmt
+	zerolog.ErrorStackMarshaler = pkgerrors.MarshalStack
+	baseLogger = newBaseLogger()
 	log.Logger = baseLogger
 }
 
-// Init configures zerolog globals and establishes the package baseline logger.
+func newBaseLogger() zerolog.Logger {
+	return zerolog.New(globalSink).Hook(componentHook{}).With().Timestamp().Logger()
+}
+
+// Init configures zerolog globals and refreshes the package logger output.
 func Init(cfg Config) zerolog.Logger {
 	mu.Lock()
 	defer mu.Unlock()
@@ -97,9 +142,7 @@ func Init(cfg Config) zerolog.Logger {
 	previousFileCloser := fileCloser
 	fileCloser = nil
 
-	zerolog.TimeFieldFormat = defaultTimeFmt
 	zerolog.SetGlobalLevel(parseLevel(cfg.Level))
-	zerolog.ErrorStackMarshaler = pkgerrors.MarshalStack
 
 	writer := selectWriter(cfg.Format)
 
@@ -117,15 +160,9 @@ func Init(cfg Config) zerolog.Logger {
 	}
 	component := strings.TrimSpace(cfg.Component)
 
-	contextBuilder := zerolog.New(writer).With().Timestamp()
-	if component != "" {
-		contextBuilder = contextBuilder.Str("component", component)
-	}
-
-	baseLogger = contextBuilder.Logger()
 	baseWriter = writer
 	baseComponent = component
-	log.Logger = baseLogger
+	globalSink.SetWriter(writer)
 
 	if previousFileCloser != nil {
 		if err := previousFileCloser.Close(); err != nil {
@@ -141,11 +178,15 @@ func Shutdown() {
 	mu.Lock()
 	defer mu.Unlock()
 
-	if fileCloser != nil {
-		if err := fileCloser.Close(); err != nil {
+	previousFileCloser := fileCloser
+	fileCloser = nil
+	baseWriter = io.MultiWriter(os.Stderr, GetBroadcaster())
+	globalSink.SetWriter(baseWriter)
+
+	if previousFileCloser != nil {
+		if err := previousFileCloser.Close(); err != nil {
 			fmt.Fprintf(os.Stderr, "logging: unable to close log file writer: %v\n", err)
 		}
-		fileCloser = nil
 	}
 
 	GetBroadcaster().Shutdown()

internal/logging/logging_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
+	"github.com/rs/zerolog/pkgerrors"
 	"golang.org/x/term"
 )
 
@@ -33,9 +34,11 @@ func resetLoggingState() {
 
 	baseWriter = os.Stderr
 	baseComponent = ""
-	baseLogger = zerolog.New(baseWriter).With().Timestamp().Logger()
+	globalSink.SetWriter(baseWriter)
+	baseLogger = newBaseLogger()
 	log.Logger = baseLogger
 	zerolog.TimeFieldFormat = defaultTimeFmt
+	zerolog.ErrorStackMarshaler = pkgerrors.MarshalStack
 	zerolog.SetGlobalLevel(zerolog.InfoLevel)
 	nowFn = time.Now
 	isTerminalFn = term.IsTerminal
@@ -242,6 +245,45 @@ func TestInitThreadSafety(t *testing.T) {
 	}
 }
 
+func TestInitThreadSafetyWithGlobalLogging(t *testing.T) {
+	t.Cleanup(resetLoggingState)
+
+	configs := []Config{
+		{Format: "json", Level: "debug", Component: "worker"},
+		{Format: "json", Level: "info", Component: "api"},
+		{Format: "json", Level: "warn", Component: "server"},
+	}
+
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+
+	for i := 0; i < 20; i++ {
+		wg.Add(1)
+		go func(worker int) {
+			defer wg.Done()
+			<-start
+			for j := 0; j < 20; j++ {
+				Init(configs[(worker+j)%len(configs)])
+				log.Info().Int("worker", worker).Int("iteration", j).Msg("concurrent init log")
+			}
+		}(i)
+	}
+
+	for i := 0; i < 20; i++ {
+		wg.Add(1)
+		go func(worker int) {
+			defer wg.Done()
+			<-start
+			for j := 0; j < 40; j++ {
+				log.Info().Int("worker", worker).Int("iteration", j).Msg("concurrent global log")
+			}
+		}(i)
+	}
+
+	close(start)
+	wg.Wait()
+}
+
 func TestIsLevelEnabled(t *testing.T) {
 	t.Cleanup(resetLoggingState)
 
@@ -1118,7 +1160,9 @@ func TestConcurrentLogging(t *testing.T) {
 	buf := &lockedBuffer{}
 	mu.Lock()
 	baseWriter = buf
-	baseLogger = zerolog.New(buf).With().Timestamp().Logger()
+	baseComponent = ""
+	globalSink.SetWriter(buf)
+	baseLogger = newBaseLogger()
 	log.Logger = baseLogger
 	mu.Unlock()