Merge #147554 #147720 #148209 #148237

147554: builtins: add more builtins to DistsqlBlocklist r=yuzefovich a=michae2

While adding a builtin for a test, I noticed that a neighboring builtin that accessed `evalctx.Planner` was not marked with `DistsqlBlocklist`. Then I noticed that many more functions accessing Planner were not marked with `DistsqlBlocklist`. I tried to ignore it, but I could not. Hence this yak-shave.

Epic: None

Release note: None

147720: bench/tpcc: use 3-node cluster, refactor, and create cleaner profiles r=mgartner a=mgartner

#### bench/tpcc: use 3-node cluster

BenchmarkTPCC now uses a 3-node cluster instead of a single-node
cluster. Starting-up a 3-node cluster takes several seconds, so the
benchmark now only does this once for all implementations (literal or
optimized) and mixes, rather than once per implemenation/mix
combination.

Also, we no longer generate a reusable store directory. Instead, we use
in-memory stores and initialize the TPC-C workload data each time the
benchmark is run. This provides a simpler user experience and the benchmark
runs just as quickly as reusing pre-initialized stores.

Release note: None

#### bench/tpcc: set cluster settings

Release note: None

#### bench/tpcc: move client code to same process

Moving the workload client to the same process vastly simplifies the
benchmark.

Release note: None

#### bench/tpcc: collect CPU and memory profile that exclude setup

CPU and heap profiles can be collected while running the benchmarks by
passing the "-profile" flag with a prefix for the profile file names.
The profile names will include the benchmark names, e.g.,
`foo_tpcc_literal_new_order.cpu` and `foo_tpcc_literal_new_order.mem`
will be created when running the benchmark:

    ./dev bench pkg/bench/tpcc -f BenchmarkTPCC/literal/new_order \
      --test-args '-test.benchtime=30s -profile=foo'

These profiles will omit CPU samples and allocations made during the setup
phase of the benchmark, unlike profiles collected with the `-test.cpuprofile`
and `-test.memprofile` flags.

Release note: None

#### tpcc: disable stdout in benchmark

The workload plumbing is cli-centric, but here we're using tpcc as a library.
While invoking it in this way doesn't trigger any workload-level logging
(behind the `--display-format` flag), tpcc has its own logging during data
loading. This commit makes it possible to override this logging, and does so in
BenchmarkTPCC.

In an ideal world, there would be configuration options for this sort of thing,
but this would be a large and painful refactor.

#### bench/tpcc: refactor and disable local client optimization

Release note: None

#### bench/tpcc: round-robin queries to all three nodes

Release note: None

#### bench/tpcc: use import data loader

Using the import data loader instead of the insert data loader speeds up
the setup time for the benchmark significantly.

Release note: None

#### bench/tpcc: run each mix on its own cluster

Each run of a benchmark is now run in its own cluster to make each
run independent of each other and avoid skewing results.

Release note: None

#### bench/tpcc: remove --profile flag

The `--profile` flag is no longer supported. Instead, CPU and heap
profiles will be collected if `-test.cpuprofile` or `-test.memprofile`
are provided. The benchmark will "hijack" the standard profiles and
create new profiles that omit CPU samples and allocations made during
the setup phase of the benchmark.

The profile names will include the prefix of the profile flags and the
benchmark names. For example, the profiles "foo_tpcc_literal_new_order.cpu"
and "foo_tpcc_literal_new_order.mem" will be created when running the
benchmark:

    ./dev bench pkg/bench/tpcc -f BenchmarkTPCC/literal/new_order \
      --test-args '-test.cpuprofile=foo.cpu -test.memprofile=foo.mem'

Note that the "foo.cpu" file will not be created because we must stop
the global CPU profiler in order to collect profiles that omit setup
samples. The "foo.mem" file will created and include all allocations
made during the entire duration of the benchmark.

Release note: None

#### bench/benchprof: add benchprof package

The benchprof package has been created to provide utilities for creating
CPU and memory profiles in benchmarks.

Release note: None

#### bench/benchprof: support mutex profiles

Release note: None

148209: sql/tests: speed-up sysbench micro-benchmarks r=mgartner a=mgartner

The sysbench micro-benchmarks have been sped up by avoiding repeatedly setting up a new cluster while Go's benchmarking harness ramps up `b.N`. Prior to this commit the benchmark command below took ~112 seconds and now takes ~36 seconds on my Mac M1 Pro.

    ./dev bench pkg/sql/tests \
      -f 'BenchmarkSysbench/SQL/1node_local/oltp_point_select' \
      --count 3

See the inline comments for more details.

Release note: None

148237: tree: avoid creating new DString in AdjustValueToType when possible r=yuzefovich a=yuzefovich

Previously, in `AdjustValueToType` we would always create a new `DString` object if the original datum wasn't modified in any way. This seems unnecessary, so we now track whether a modification might have occurred and allocate the new object only in that case, otherwise returning the original datum. We choose to not do the same for collated strings since there it wouldn't have been an equivalent change (we currently initialize a new CollationEnvironment).

Noticed this inefficiency when looking at CPU profile of a node running an IMPORT.

```
name                                old time/op    new time/op    delta
ImportFixture/tpcc/warehouses=1-24     586ms ± 9%     585ms ± 4%     ~     (p=0.853 n=10+10)

name                                old speed      new speed      delta
ImportFixture/tpcc/warehouses=1-24   148MB/s ±10%   148MB/s ± 5%     ~     (p=0.912 n=10+10)

name                                old alloc/op   new alloc/op   delta
ImportFixture/tpcc/warehouses=1-24    1.13GB ± 1%    1.10GB ± 1%   -2.76%  (p=0.000 n=10+10)

name                                old allocs/op  new allocs/op  delta
ImportFixture/tpcc/warehouses=1-24     6.08M ± 0%     4.12M ± 0%  -32.25%  (p=0.000 n=10+10)
```

Epic: None
Release note: None

Co-authored-by: Michael Erickson <[email protected]>
Co-authored-by: Marcus Gartner <[email protected]>
Co-authored-by: Tobias Grieger <[email protected]>
Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: 5 people

pkg/BUILD.bazel

@@ -858,6 +858,7 @@ GO_TARGETS = [
     "//pkg/base/serverident:serverident",
     "//pkg/base:base",
     "//pkg/base:base_test",
+    "//pkg/bench/benchprof:benchprof",
     "//pkg/bench/cmd/pgbenchsetup:pgbenchsetup",
     "//pkg/bench/cmd/pgbenchsetup:pgbenchsetup_lib",
     "//pkg/bench/hashbench:hashbench_test",

pkg/bench/benchprof/BUILD.bazel

@@ -0,0 +1,12 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "benchprof",
+    srcs = ["benchprof.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/bench/benchprof",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/testutils/sniffarg",
+        "@com_github_google_pprof//profile",
+    ],
+)

pkg/bench/benchprof/benchprof.go

@@ -0,0 +1,296 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package benchprof
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"regexp"
+	"runtime"
+	runtimepprof "runtime/pprof"
+	"strings"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/testutils/sniffarg"
+	"github.com/google/pprof/profile"
+)
+
+// StopFn is a function that stops profiling.
+type StopFn func(testing.TB)
+
+// Stop stops profiling.
+func (f StopFn) Stop(tb testing.TB) {
+	if f != nil {
+		f(tb)
+	}
+}
+
+// StartAllProfiles starts collection of CPU, heap, and mutex profiles, if the
+// "-test.cpuprofile", "-test.memprofile", or "-test.mutexprofile" flags are
+// set. See StartCPUProfile, StartMemProfile, and StartMutexProfile for more
+// details.
+//
+// Example usage:
+//
+//	func BenchmarkFoo(b *testing.B) {
+//	    // ..
+//	    b.Run("case", func(b *testing.B) {
+//	        defer benchprof.StartAllProfiles(b).Stop(b)
+//	        // Benchmark loop.
+//	    })
+//	}
+func StartAllProfiles(tb testing.TB) StopFn {
+	cpuStop := StartCPUProfile(tb)
+	memStop := StartMemProfile(tb)
+	mutexStop := StartMutexProfile(tb)
+	return func(b testing.TB) {
+		cpuStop(b)
+		memStop(b)
+		mutexStop(b)
+	}
+}
+
+// StartCPUProfile starts collection of a CPU profile if the "-test.cpuprofile"
+// flag has been set. The profile will omit CPU samples collected prior to
+// calling StartCPUProfile, and include all allocations made until the returned
+// StopFn is called.
+//
+// Example usage:
+//
+//	func BenchmarkFoo(b *testing.B) {
+//	    // ..
+//	    b.Run("case", func(b *testing.B) {
+//	        defer benchprof.StartCPUProfile(b).Stop(b)
+//	        // Benchmark loop.
+//	    })
+//	}
+//
+// The file name of the profile will include the prefix of the profile flags and
+// the benchmark names. For example, "foo_benchmark_thing_1.cpu" would be
+// created for a "BenchmarkThing/1" benchmark if the "-test.cpuprofile=foo.cpu"
+// flag is set.
+//
+// NB: The "foo.cpu" file will not be created because we must stop the global
+// CPU profiler in order to collect profiles that omit setup samples.
+func StartCPUProfile(tb testing.TB) StopFn {
+	var cpuProfFile string
+	if err := sniffarg.DoEnv("test.cpuprofile", &cpuProfFile); err != nil {
+		tb.Fatal(err)
+	}
+	if cpuProfFile == "" {
+		// Not CPU profile requested.
+		return nil
+	}
+
+	prefix := profilePrefix(cpuProfFile)
+	dir := outputDir(tb)
+	if dir != "" {
+		cpuProfFile = filepath.Join(dir, cpuProfFile)
+	}
+
+	// Hijack the harness's profile to make a clean profile.
+	// The flag is set, so likely a CPU profile started by the Go harness is
+	// running (unless -count is specified, but StopCPUProfile is idempotent).
+	runtimepprof.StopCPUProfile()
+
+	// Remove the harness's profile file. It would not be an accurate profile.
+	_ = os.Remove(cpuProfFile)
+
+	// Create a new profile file.
+	fileName := profileFileName(tb, dir, prefix, "cpu")
+	f, err := os.OpenFile(fileName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
+	if err != nil {
+		tb.Fatal(err)
+	}
+
+	// Start profiling.
+	if err := runtimepprof.StartCPUProfile(f); err != nil {
+		tb.Fatal(err)
+	}
+
+	return func(b testing.TB) {
+		runtimepprof.StopCPUProfile()
+		if err := f.Close(); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+// StartMemProfile starts collection of a heap profile if the "-test.memprofile"
+// flag has been set. The profile will omit memory allocations prior to calling
+// StartMemProfile, and include all allocations made until the returned StopFn
+// is called.
+//
+// Example usage:
+//
+//	func BenchmarkFoo(b *testing.B) {
+//	    // ..
+//	    b.Run("case", func(b *testing.B) {
+//	        defer benchprof.StartMemProfile(b).Stop(b)
+//	        // Benchmark loop.
+//	    })
+//	}
+//
+// The file name of the profile will include the prefix of the profile flags and
+// the benchmark names. For example, "foo_benchmark_thing_1.mem" would be
+// created for a "BenchmarkThing/1" benchmark if the "-test.memprofile=foo.mem"
+// flag is set.
+//
+// NB: The "foo.mem" file will still be created and include all allocations made
+// during the entire duration of the benchmark.
+func StartMemProfile(tb testing.TB) StopFn {
+	var memProfFile string
+	if err := sniffarg.DoEnv("test.memprofile", &memProfFile); err != nil {
+		tb.Fatal(err)
+	}
+	if memProfFile == "" {
+		// No heap profile requested.
+		return nil
+	}
+
+	prefix := profilePrefix(memProfFile)
+	dir := outputDir(tb)
+	if dir != "" {
+		memProfFile = filepath.Join(dir, memProfFile)
+	}
+
+	// Create a new profile file.
+	fileName := profileFileName(tb, dir, prefix, "mem")
+	diffAllocs := diffProfile(func() []byte {
+		p := runtimepprof.Lookup("allocs")
+		var buf bytes.Buffer
+		runtime.GC()
+		if err := p.WriteTo(&buf, 0); err != nil {
+			tb.Fatal(err)
+		}
+		return buf.Bytes()
+	})
+
+	return func(b testing.TB) {
+		if sl := diffAllocs(b); len(sl) > 0 {
+			if err := os.WriteFile(fileName, sl, 0644); err != nil {
+				b.Fatal(err)
+			}
+		}
+	}
+}
+
+// StartMutexProfile starts collection of a mutex profile if the
+// "-test.cpuprofile" flag has been set. The profile will only collect data
+// between calling StartMutexProfile and calling the returned StopFn.
+//
+// Example usage:
+//
+//	func BenchmarkFoo(b *testing.B) {
+//	    // ..
+//	    b.Run("case", func(b *testing.B) {
+//	        defer benchprof.StartMutexProfile(b).Stop(b)
+//	        // Benchmark loop.
+//	    })
+//	}
+//
+// The file name of the profile will include the prefix of the profile flags and
+// the benchmark names. For example, "foo_benchmark_thing_1.mutex" would be
+// created for a "BenchmarkThing/1" benchmark if the
+// "-test.mutexprofile=foo.mutex" flag is set.
+func StartMutexProfile(tb testing.TB) StopFn {
+	var mutexProfFile string
+	if err := sniffarg.DoEnv("test.mutexprofile", &mutexProfFile); err != nil {
+		tb.Fatal(err)
+	}
+	if mutexProfFile == "" {
+		// No mutex profile requested.
+		return nil
+	}
+
+	prefix := profilePrefix(mutexProfFile)
+	dir := outputDir(tb)
+	if dir != "" {
+		mutexProfFile = filepath.Join(dir, mutexProfFile)
+	}
+
+	// Create a new profile file.
+	fileName := profileFileName(tb, dir, prefix, "mutex")
+	diffAllocs := diffProfile(func() []byte {
+		p := runtimepprof.Lookup("mutex")
+		var buf bytes.Buffer
+		if err := p.WriteTo(&buf, 0); err != nil {
+			tb.Fatal(err)
+		}
+		return buf.Bytes()
+	})
+
+	return func(b testing.TB) {
+		if sl := diffAllocs(b); len(sl) > 0 {
+			if err := os.WriteFile(fileName, sl, 0644); err != nil {
+				b.Fatal(err)
+			}
+		}
+	}
+}
+
+func diffProfile(take func() []byte) func(testing.TB) []byte {
+	// The below is essentially cribbed from pprof.go in net/http/pprof.
+
+	baseBytes := take()
+	if baseBytes == nil {
+		return func(tb testing.TB) []byte { return nil }
+	}
+	return func(b testing.TB) []byte {
+		newBytes := take()
+		pBase, err := profile.ParseData(baseBytes)
+		if err != nil {
+			b.Fatal(err)
+		}
+		pNew, err := profile.ParseData(newBytes)
+		if err != nil {
+			b.Fatal(err)
+		}
+		pBase.Scale(-1)
+		pMerged, err := profile.Merge([]*profile.Profile{pBase, pNew})
+		if err != nil {
+			b.Fatal(err)
+		}
+		pMerged.TimeNanos = pNew.TimeNanos
+		pMerged.DurationNanos = pNew.TimeNanos - pBase.TimeNanos
+
+		buf := bytes.Buffer{}
+		if err := pMerged.Write(&buf); err != nil {
+			b.Fatal(err)
+		}
+		return buf.Bytes()
+	}
+}
+
+func outputDir(tb testing.TB) string {
+	var dir string
+	if err := sniffarg.DoEnv("test.outputdir", &dir); err != nil {
+		tb.Fatal(err)
+	}
+	return dir
+}
+
+func profilePrefix(profileArg string) string {
+	i := strings.Index(profileArg, ".")
+	if i == -1 {
+		return profileArg
+	}
+	return profileArg[:i]
+}
+
+func profileFileName(tb testing.TB, outputDir, prefix, suffix string) string {
+	saniRE := regexp.MustCompile(`\W+`)
+	testName := strings.TrimPrefix(tb.Name(), "Benchmark")
+	testName = strings.ToLower(testName)
+	testName = saniRE.ReplaceAllString(testName, "_")
+
+	fileName := prefix + "_" + testName + "." + suffix
+	if outputDir != "" {
+		fileName = filepath.Join(outputDir, fileName)
+	}
+	return fileName
+}

pkg/bench/tpcc/BUILD.bazel

@@ -4,35 +4,28 @@ go_test(
     name = "tpcc_test",
     srcs = [
         "bench_test.go",
-        "internal_test.go",
         "main_test.go",
-        "utils_test.go",
     ],
     data = glob(["testdata/**"]),
     deps = [
         "//pkg/base",
+        "//pkg/bench/benchprof",
+        "//pkg/ccl/workloadccl",
+        "//pkg/rpc/nodedialer",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",
-        "//pkg/storage",
-        "//pkg/testutils/pgurlutils",
+        "//pkg/settings/cluster",
+        "//pkg/sql/stats",
         "//pkg/testutils/serverutils",
-        "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
-        "//pkg/testutils/testfixtures",
-        "//pkg/util/envutil",
-        "//pkg/util/leaktest",
+        "//pkg/ts",
         "//pkg/util/log",
-        "//pkg/util/randutil",
         "//pkg/util/stop",
-        "//pkg/util/syncutil",
         "//pkg/workload",
         "//pkg/workload/histogram",
         "//pkg/workload/tpcc",
         "//pkg/workload/workloadsql",
-        "@com_github_cockroachdb_pebble//vfs",
-        "@com_github_jackc_pgx_v5//:pgx",
-        "@com_github_stretchr_testify//require",
     ],
 )