log: fix SetDefault for custom loggers (#31368)

Currently, even though it takes in a `Logger` interface,
`log.SetDefualt` enforces that the concrete type of the provided logger
is `*logger` because:
1. in `init` `root.Store` is called with a `*logger`
2. `atomic.Value` panics if the concrete type provided in `Store` is not
consistent across calls.
([ref](https://pkg.go.dev/sync/atomic#Value.Store))

> All calls to Store for a given Value must use values of the same
concrete type.

This PR changes to use `sync.RWMutex` and adds a test that panics on
`master`.

log/root.go

@@ -3,18 +3,24 @@ package log
 import (
 	"log/slog"
 	"os"
-	"sync/atomic"
+	"sync"
 )
 
-var root atomic.Value
+var (
+	rootLock sync.RWMutex
+	root     Logger
+)
 
 func init() {
-	root.Store(&logger{slog.New(DiscardHandler())})
+	root = &logger{slog.New(DiscardHandler())}
 }
 
 // SetDefault sets the default global logger
 func SetDefault(l Logger) {
-	root.Store(l)
+	rootLock.Lock()
+	defer rootLock.Unlock()
+
+	root = l
 	if lg, ok := l.(*logger); ok {
 		slog.SetDefault(lg.inner)
 	}
@@ -22,7 +28,10 @@ func SetDefault(l Logger) {
 
 // Root returns the root logger
 func Root() Logger {
-	return root.Load().(Logger)
+	rootLock.RLock()
+	defer rootLock.RUnlock()
+
+	return root
 }
 
 // The following functions bypass the exported logger methods (logger.Debug,

log/root_test.go

@@ -0,0 +1,19 @@
+package log
+
+import (
+	"testing"
+)
+
+// SetDefault should properly set the default logger when custom loggers are
+// provided.
+func TestSetDefaultCustomLogger(t *testing.T) {
+	type customLogger struct {
+		Logger // Implement the Logger interface
+	}
+
+	customLog := &customLogger{}
+	SetDefault(customLog)
+	if Root() != customLog {
+		t.Error("expected custom logger to be set as default")
+	}
+}