Merge cockroachdb#108378

craig[bot] · sean- · craig[bot] · commit e3efe1130619 · 2023-08-10T18:42:37.000Z
108378: workload: cache DNS lookups by default r=sean- a=sean-

By default, DNS lookups used by jackc/pgx are now cached for 30s.  A negative duration will disable the caching behavior.

While performing load testing, at higher connection counts with lower connection rates, it was possible to overwhelm the default environment DNS resolvers.

Epic: none
Release note: None

Co-authored-by: Sean Chittenden &lt;sean@chittenden.org&gt;
diff --git a/DEPS.bzl b/DEPS.bzl
@@ -7666,6 +7666,16 @@ def go_deps():
             "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rs/cors/com_github_rs_cors-v1.8.0.zip",
         ],
     )
+    go_repository(
+        name = "com_github_rs_dnscache",
+        build_file_proto_mode = "disable_global",
+        importpath = "github.com/rs/dnscache",
+        sha256 = "11e1fa18f7a18eac97b54b3726363598577ac0df7a6ce806f4775088593c0047",
+        strip_prefix = "github.com/rs/dnscache@v0.0.0-20230804202142-fc85eb664529",
+        urls = [
+            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rs/dnscache/com_github_rs_dnscache-v0.0.0-20230804202142-fc85eb664529.zip",
+        ],
+    )
     go_repository(
         name = "com_github_rs_xid",
         build_file_proto_mode = "disable_global",
diff --git a/build/bazelutil/distdir_files.bzl b/build/bazelutil/distdir_files.bzl
@@ -901,6 +901,7 @@ DISTDIR_FILES = {
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rogpeppe/fastuuid/com_github_rogpeppe_fastuuid-v1.2.0.zip": "f9b8293f5e20270e26fb4214ca7afec864de92c73d03ff62b5ee29d1db4e72a1",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rogpeppe/go-internal/com_github_rogpeppe_go_internal-v1.9.0.zip": "7d777908b9c91a1685f2d709550e6b7478e14e9c4699dffd7f0a150e36dbc7e9",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rs/cors/com_github_rs_cors-v1.8.0.zip": "9aeb6b48d7ba5d34187b40adaed8280f0690e6d9b4fd6132eccbd62aa2c0efd9",
+    "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rs/dnscache/com_github_rs_dnscache-v0.0.0-20230804202142-fc85eb664529.zip": "11e1fa18f7a18eac97b54b3726363598577ac0df7a6ce806f4775088593c0047",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rs/xid/com_github_rs_xid-v1.3.0.zip": "809ed1d8845fe5d73f6973e9b7a33eefd786cc97b1aebe493243e420b7c89958",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/rs/zerolog/com_github_rs_zerolog-v1.15.0.zip": "8e98c48e7fd132aafbf129664e8fd65229d067d772bff4bd712a497b7a2f00c4",
     "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/russross/blackfriday/com_github_russross_blackfriday-v1.6.0.zip": "8dbd018a896577afef870d49adc9a7cbdcef54f6edd97dcbbe1b53e7cd6d66d4",
diff --git a/go.mod b/go.mod
@@ -205,6 +205,7 @@ require (
 	github.com/pseudomuto/protoc-gen-doc v1.3.2
 	github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475
 	github.com/robfig/cron/v3 v3.0.1
+	github.com/rs/dnscache v0.0.0-20230804202142-fc85eb664529
 	github.com/sasha-s/go-deadlock v0.3.1
 	github.com/shirou/gopsutil/v3 v3.21.12
 	github.com/slack-go/slack v0.9.5
diff --git a/go.sum b/go.sum
@@ -2062,6 +2062,8 @@ github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/f
 github.com/rs/cors v1.6.0/go.mod h1:gFx+x8UowdsKA9AchylcLynDq+nNFfI8FkUZdN/jGCU=
 github.com/rs/cors v1.7.0/go.mod h1:gFx+x8UowdsKA9AchylcLynDq+nNFfI8FkUZdN/jGCU=
 github.com/rs/cors v1.8.0/go.mod h1:EBwu+T5AvHOcXwvZIkQFjUN6s8Czyqw12GL/Y0tUyRM=
+github.com/rs/dnscache v0.0.0-20230804202142-fc85eb664529 h1:18kd+8ZUlt/ARXhljq+14TwAoKa61q6dX8jtwOf6DH8=
+github.com/rs/dnscache v0.0.0-20230804202142-fc85eb664529/go.mod h1:qe5TWALJ8/a1Lqznoc5BDHpYX/8HU60Hm2AwRmqzxqA=
 github.com/rs/xid v1.2.1/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ=
 github.com/rs/xid v1.3.0 h1:6NjYksEUlhurdVehpc7S7dk6DAmcKv8V9gG0FsVN2U4=
 github.com/rs/xid v1.3.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
diff --git a/pkg/ccl/testccl/workload/schemachange/schema_change_external_test.go b/pkg/ccl/testccl/workload/schemachange/schema_change_external_test.go
@@ -39,7 +39,9 @@ func TestWorkload(t *testing.T) {
 	scope := log.Scope(t)
 	defer scope.Close(t)
 	dir := scope.GetDirectory()
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel() // NOTE: Required to cleanup dnscache refresh Go routine
+
 	tc, _, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
 		t,
 		3, /* numServers */
diff --git a/pkg/workload/BUILD.bazel b/pkg/workload/BUILD.bazel
@@ -31,6 +31,7 @@ go_library(
         "@com_github_jackc_pgx_v5//pgconn",
         "@com_github_jackc_pgx_v5//pgxpool",
         "@com_github_lib_pq//:pq",
+        "@com_github_rs_dnscache//:dnscache",
         "@com_github_spf13_pflag//:pflag",
         "@org_golang_x_sync//errgroup",
     ],
diff --git a/pkg/workload/bank/bank_test.go b/pkg/workload/bank/bank_test.go
@@ -37,7 +37,9 @@ func TestBank(t *testing.T) {
 		{10, 100, 10}, // don't make more ranges than rows
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{UseDatabase: `test`})
 	defer s.Stopper().Stop(ctx)
 	sqlutils.MakeSQLRunner(db).Exec(t, `CREATE DATABASE test`)
diff --git a/pkg/workload/connection.go b/pkg/workload/connection.go
@@ -28,6 +28,7 @@ type ConnFlags struct {
 	Method      string // Method for issuing queries; see SQLRunner.
 
 	ConnHealthCheckPeriod time.Duration
+	DNSRefreshInterval    time.Duration
 	MaxConnIdleTime       time.Duration
 	MaxConnLifetime       time.Duration
 	MaxConnLifetimeJitter time.Duration
@@ -44,6 +45,7 @@ func NewConnFlags(genFlags *Flags) *ConnFlags {
 	c.IntVar(&c.Concurrency, `concurrency`, 2*runtime.GOMAXPROCS(0),
 		`Number of concurrent workers`)
 	c.StringVar(&c.Method, `method`, `cache_statement`, `SQL issue method (cache_statement, cache_describe, describe_exec, exec, simple_protocol)`)
+	c.DurationVar(&c.DNSRefreshInterval, `dns-refresh`, defaultDNSCacheRefresh, `Interval used to refresh cached DNS entries (<0 disables)`)
 	c.DurationVar(&c.ConnHealthCheckPeriod, `conn-healthcheck-period`, 30*time.Second, `Interval that health checks are run on connections`)
 	c.IntVar(&c.MinConns, `min-conns`, 0, `Minimum number of connections to attempt to keep in the pool`)
 	c.DurationVar(&c.MaxConnIdleTime, `max-conn-idle-time`, 150*time.Second, `Max time an idle connection will be kept around`)
@@ -58,6 +60,7 @@ func NewConnFlags(genFlags *Flags) *ConnFlags {
 		`concurrency`,
 		`conn-healthcheck-period`,
 		`db`,
+		`dns-refresh`,
 		`max-conn-idle-time`,
 		`max-conn-lifetime-jitter`,
 		`max-conn-lifetime`,
diff --git a/pkg/workload/csv_test.go b/pkg/workload/csv_test.go
@@ -76,7 +76,8 @@ func TestHandleCSV(t *testing.T) {
 }
 
 func BenchmarkWriteCSVRows(b *testing.B) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	var batches []coldata.Batch
 	for _, table := range tpcc.FromWarehouses(1).Tables() {
diff --git a/pkg/workload/insights/insights_test.go b/pkg/workload/insights/insights_test.go
@@ -38,7 +38,9 @@ func TestInsightsWorkload(t *testing.T) {
 		{10, 100, 10}, // don't make more ranges than rows
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{UseDatabase: `test`})
 	defer s.Stopper().Stop(ctx)
 	sqlutils.MakeSQLRunner(db).Exec(t, `CREATE DATABASE test`)
diff --git a/pkg/workload/pgx_helpers.go b/pkg/workload/pgx_helpers.go
@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/rs/dnscache"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -40,6 +41,9 @@ type MultiConnPool struct {
 	// NOTE(seanc@): method is on the hot-path, therefore make all reads to the
 	// query exec mode a dirty read.
 	method pgx.QueryExecMode
+
+	// Caching DNS resolver
+	resolver *dnscache.Resolver
 }
 
 // MultiConnPoolCfg encapsulates the knobs passed to NewMultiConnPool.
@@ -86,13 +90,17 @@ type MultiConnPoolCfg struct {
 	// max number of connections per pool.  A value less than 0 skips the
 	// connection warmup phase.
 	WarmupConns int
+
+	// Duration between refreshes of the DNS cache
+	DNSRefreshInterval time.Duration
 }
 
 // NewMultiConnPoolCfgFromFlags constructs a new MultiConnPoolCfg object based
 // on the connection flags.
 func NewMultiConnPoolCfgFromFlags(cf *ConnFlags) MultiConnPoolCfg {
 	return MultiConnPoolCfg{
 		ConnHealthCheckPeriod: cf.ConnHealthCheckPeriod,
+		DNSRefreshInterval:    cf.DNSRefreshInterval,
 		MaxConnIdleTime:       cf.MaxConnIdleTime,
 		MaxConnLifetime:       cf.MaxConnLifetime,
 		MaxConnLifetimeJitter: cf.MaxConnLifetimeJitter,
@@ -155,6 +163,36 @@ func NewMultiConnPool(
 	if cfg.MinConns > 0 {
 		minConns = cfg.MinConns
 	}
+	var dnsCacheRefreshInterval time.Duration
+	switch {
+	case cfg.DNSRefreshInterval > 0:
+		dnsCacheRefreshInterval = cfg.DNSRefreshInterval
+	case cfg.DNSRefreshInterval < 0:
+		// Disable DNS caching
+	default:
+		dnsCacheRefreshInterval = defaultDNSCacheRefresh
+	}
+	if dnsCacheRefreshInterval > 0 {
+		m.resolver = &dnscache.Resolver{}
+
+		// Refresh DNS cache
+		go func() {
+			ticker := time.NewTicker(dnsCacheRefreshInterval)
+			defer ticker.Stop()
+
+			for {
+				select {
+				case <-ticker.C:
+					m.resolver.RefreshWithOptions(dnscache.ResolverRefreshOptions{
+						ClearUnused:      true,
+						PersistOnFailure: true,
+					})
+				case <-ctx.Done():
+					return
+				}
+			}
+		}()
+	}
 
 	connsPerURL := distribute(cfg.MaxTotalConnections, len(urls))
 	maxConnsPerPool := cfg.MaxConnsPerPool
@@ -203,6 +241,9 @@ func NewMultiConnPool(
 			}
 
 			connCfg := poolCfg.ConnConfig
+			if m.resolver != nil {
+				connCfg.LookupFunc = m.resolver.LookupHost
+			}
 			connCfg.DefaultQueryExecMode = queryMode
 			p, err := pgxpool.NewWithConfig(ctx, poolCfg)
 			if err != nil {
diff --git a/pkg/workload/rand/rand_test.go b/pkg/workload/rand/rand_test.go
@@ -47,7 +47,9 @@ func TestRandRun(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	dbName := "rand_test"
 	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{UseDatabase: dbName})
 	defer s.Stopper().Stop(ctx)
diff --git a/pkg/workload/workload.go b/pkg/workload/workload.go
@@ -211,6 +211,9 @@ func Tuples(count int, fn func(int) []interface{}) BatchedTuples {
 const (
 	// timestampOutputFormat is used to output all timestamps.
 	timestampOutputFormat = "2006-01-02 15:04:05.999999-07:00"
+
+	// defaultDNSCacheRefresh is used to change the default --dns-refresh interval
+	defaultDNSCacheRefresh = 30 * time.Second
 )
 
 // TypedTuples returns a BatchedTuples where each batch has size 1. It's
diff --git a/pkg/workload/workloadsql/workloadsql_test.go b/pkg/workload/workloadsql/workloadsql_test.go
@@ -44,7 +44,9 @@ func TestSetup(t *testing.T) {
 		{10, 100, 4},
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{UseDatabase: `test`})
 	defer s.Stopper().Stop(ctx)
 	sqlutils.MakeSQLRunner(db).Exec(t, `CREATE DATABASE test`)

Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,8 @@ func TestHandleCSV(t *testing.T) {`
`76`	`76`	`}`
`77`	`77`
`78`	`78`	`func BenchmarkWriteCSVRows(b *testing.B) {`
`79`		`- ctx := context.Background()`
	`79`	`+ ctx, cancel := context.WithCancel(context.Background())`
	`80`	`+ defer cancel()`
`80`	`81`
`81`	`82`	`var batches []coldata.Batch`
`82`	`83`	`for _, table := range tpcc.FromWarehouses(1).Tables() {`