From ce66b1dae4d1e9fc68d576c029463562e2324d63 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Tue, 10 Sep 2024 07:06:39 -0500
Subject: [PATCH] Fix data race with TraceLog.Config initialization

https://github.com/jackc/pgx/pull/2120
---
 go.mod                    |  2 +-
 go.sum                    |  4 +--
 tracelog/tracelog.go      | 19 +++++++++-----
 tracelog/tracelog_test.go | 53 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 0d952d691..20e948a3d 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/jackc/puddle/v2 v2.2.1
 	github.com/stretchr/testify v1.8.1
 	golang.org/x/crypto v0.17.0
+	golang.org/x/sync v0.8.0
 	golang.org/x/text v0.14.0
 )
 
@@ -15,7 +16,6 @@ require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/kr/pretty v0.3.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	golang.org/x/sync v0.1.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/go.sum b/go.sum
index 29fe452b2..c5bdd043d 100644
--- a/go.sum
+++ b/go.sum
@@ -31,8 +31,8 @@ github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKs
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k=
 golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
-golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
-golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
+golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
diff --git a/tracelog/tracelog.go b/tracelog/tracelog.go
index 212a25407..61015354a 100644
--- a/tracelog/tracelog.go
+++ b/tracelog/tracelog.go
@@ -6,6 +6,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"sync"
 	"time"
 	"unicode/utf8"
 
@@ -129,19 +130,25 @@ func DefaultTraceLogConfig() *TraceLogConfig {
 	}
 }
 
-// TraceLog implements pgx.QueryTracer, pgx.BatchTracer, pgx.ConnectTracer, and pgx.CopyFromTracer. All fields are
-// required.
+// TraceLog implements pgx.QueryTracer, pgx.BatchTracer, pgx.ConnectTracer, and pgx.CopyFromTracer. Logger and LogLevel
+// are required. Config will be automatically initialized on first use if nil.
 type TraceLog struct {
 	Logger   Logger
 	LogLevel LogLevel
-	Config   *TraceLogConfig
+
+	Config           *TraceLogConfig
+	ensureConfigOnce sync.Once
 }
 
 // ensureConfig initializes the Config field with default values if it is nil.
 func (tl *TraceLog) ensureConfig() {
-	if tl.Config == nil {
-		tl.Config = DefaultTraceLogConfig()
-	}
+	tl.ensureConfigOnce.Do(
+		func() {
+			if tl.Config == nil {
+				tl.Config = DefaultTraceLogConfig()
+			}
+		},
+	)
 }
 
 type ctxKey int
diff --git a/tracelog/tracelog_test.go b/tracelog/tracelog_test.go
index 6812a97ba..f1959b324 100644
--- a/tracelog/tracelog_test.go
+++ b/tracelog/tracelog_test.go
@@ -6,14 +6,17 @@ import (
 	"log"
 	"os"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jackc/pgx/v5/pgxtest"
 	"github.com/jackc/pgx/v5/tracelog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
 )
 
 var defaultConnTestRunner pgxtest.ConnTestRunner
@@ -35,18 +38,29 @@ type testLog struct {
 
 type testLogger struct {
 	logs []testLog
+
+	mux sync.Mutex
 }
 
 func (l *testLogger) Log(ctx context.Context, level tracelog.LogLevel, msg string, data map[string]any) {
+	l.mux.Lock()
+	defer l.mux.Unlock()
+
 	data["ctxdata"] = ctx.Value("ctxdata")
 	l.logs = append(l.logs, testLog{lvl: level, msg: msg, data: data})
 }
 
 func (l *testLogger) Clear() {
+	l.mux.Lock()
+	defer l.mux.Unlock()
+
 	l.logs = l.logs[0:0]
 }
 
 func (l *testLogger) FilterByMsg(msg string) (res []testLog) {
+	l.mux.Lock()
+	defer l.mux.Unlock()
+
 	for _, log := range l.logs {
 		if log.msg == msg {
 			res = append(res, log)
@@ -457,3 +471,42 @@ func TestLogPrepare(t *testing.T) {
 		require.Equal(t, err, logger.logs[0].data["err"])
 	})
 }
+
+// https://github.com/jackc/pgx/pull/2120
+func TestConcurrentUsage(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	logger := &testLogger{}
+	tracer := &tracelog.TraceLog{
+		Logger:   logger,
+		LogLevel: tracelog.LogLevelTrace,
+	}
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	config.ConnConfig.Tracer = tracer
+
+	for i := 0; i < 50; i++ {
+		func() {
+			pool, err := pgxpool.NewWithConfig(ctx, config)
+			require.NoError(t, err)
+
+			defer pool.Close()
+
+			eg := errgroup.Group{}
+
+			for i := 0; i < 5; i++ {
+				eg.Go(func() error {
+					_, err := pool.Exec(ctx, `select 1`)
+					return err
+				})
+			}
+
+			err = eg.Wait()
+			require.NoError(t, err)
+		}()
+	}
+}