colexecerror: add benchmarks for CatchVectorizedRuntimeError

Informs: cockroachdb#123235

Release note: None

Author: michae2

pkg/sql/colexecerror/BUILD.bazel

@@ -15,11 +15,34 @@ go_library(
 
 go_test(
     name = "colexecerror_test",
-    srcs = ["error_test.go"],
+    srcs = [
+        "error_test.go",
+        "main_test.go",
+    ],
     deps = [
         ":colexecerror",
+        "//pkg/base",
+        "//pkg/col/coldata",
+        "//pkg/col/coldataext",
+        "//pkg/security/securityassets",
+        "//pkg/security/securitytest",
+        "//pkg/server",
+        "//pkg/settings/cluster",
+        "//pkg/sql/colexec/colexectestutils",
+        "//pkg/sql/colmem",
+        "//pkg/sql/execinfra",
+        "//pkg/sql/pgwire/pgcode",
+        "//pkg/sql/pgwire/pgerror",
+        "//pkg/sql/sem/eval",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/skip",
+        "//pkg/testutils/sqlutils",
+        "//pkg/testutils/testcluster",
         "//pkg/util/leaktest",
         "//pkg/util/log",
+        "//pkg/util/mon",
+        "//pkg/util/randutil",
+        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/sql/colexecerror/error.go

@@ -100,12 +100,13 @@ func CatchVectorizedRuntimeError(operation func()) (retErr error) {
 // Multiple actual packages can have the same prefix as a single constant string
 // defined below, but all of such packages are allowed to be caught from.
 const (
-	colPackagesPrefix      = "github.com/cockroachdb/cockroach/pkg/col"
-	encodingPackagePrefix  = "github.com/cockroachdb/cockroach/pkg/util/encoding"
-	execinfraPackagePrefix = "github.com/cockroachdb/cockroach/pkg/sql/execinfra"
-	sqlColPackagesPrefix   = "github.com/cockroachdb/cockroach/pkg/sql/col"
-	sqlRowPackagesPrefix   = "github.com/cockroachdb/cockroach/pkg/sql/row"
-	sqlSemPackagesPrefix   = "github.com/cockroachdb/cockroach/pkg/sql/sem"
+	colPackagesPrefix        = "github.com/cockroachdb/cockroach/pkg/col"
+	encodingPackagePrefix    = "github.com/cockroachdb/cockroach/pkg/util/encoding"
+	execinfraPackagePrefix   = "github.com/cockroachdb/cockroach/pkg/sql/execinfra"
+	sqlColPackagesPrefix     = "github.com/cockroachdb/cockroach/pkg/sql/col"
+	sqlRowPackagesPrefix     = "github.com/cockroachdb/cockroach/pkg/sql/row"
+	sqlSemPackagesPrefix     = "github.com/cockroachdb/cockroach/pkg/sql/sem"
+	testSqlColPackagesPrefix = "pkg/sql/col"
 )
 
 // shouldCatchPanic checks whether the panic that was emitted from
@@ -135,7 +136,8 @@ func shouldCatchPanic(panicEmittedFrom string) bool {
 		strings.HasPrefix(panicEmittedFrom, execinfraPackagePrefix) ||
 		strings.HasPrefix(panicEmittedFrom, sqlColPackagesPrefix) ||
 		strings.HasPrefix(panicEmittedFrom, sqlRowPackagesPrefix) ||
-		strings.HasPrefix(panicEmittedFrom, sqlSemPackagesPrefix)
+		strings.HasPrefix(panicEmittedFrom, sqlSemPackagesPrefix) ||
+		strings.HasPrefix(panicEmittedFrom, testSqlColPackagesPrefix)
 }
 
 // StorageError is an error that was created by a component below the sql

pkg/sql/colexecerror/error_test.go

@@ -11,13 +11,22 @@
 package colexecerror_test
 
 import (
-	"errors"
+	"context"
+	gosql "database/sql"
+	"fmt"
+	"runtime"
 	"strings"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 )
 
@@ -74,3 +83,149 @@ func TestNonCatchablePanicIsNotCaught(t *testing.T) {
 		}))
 	})
 }
+
+// BenchmarkCatchVectorizedRuntimeError measures the time for
+// CatchVectorizedRuntimeError to catch and process an error.
+func BenchmarkCatchVectorizedRuntimeError(b *testing.B) {
+	err := errors.New("oops")
+	storageErr := colexecerror.NewStorageError(err)
+	pgErr := pgerror.WithCandidateCode(err, pgcode.Warning)
+
+	cases := []struct {
+		name    string
+		thrower func()
+	}{
+		{
+			"noError",
+			func() {},
+		},
+		{
+			"expected",
+			func() {
+				colexecerror.ExpectedError(err)
+			},
+		},
+		{
+			"storage",
+			func() {
+				colexecerror.InternalError(storageErr)
+			},
+		},
+		{
+			"contextCanceled",
+			func() {
+				colexecerror.InternalError(context.Canceled)
+			},
+		},
+		{
+			"internalWithCode",
+			func() {
+				colexecerror.InternalError(pgErr)
+			},
+		},
+		{
+			"internal",
+			func() {
+				colexecerror.InternalError(err)
+			},
+		},
+		{
+			"runtime",
+			func() {
+				arr := []int{0, 1, 2}
+				_ = arr[3]
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		b.Run(tc.name, func(b *testing.B) {
+			b.RunParallel(func(pb *testing.PB) {
+				for pb.Next() {
+					_ = colexecerror.CatchVectorizedRuntimeError(tc.thrower)
+				}
+			})
+		})
+	}
+}
+
+func BenchmarkSQLCatchVectorizedRuntimeError(b *testing.B) {
+	defer leaktest.AfterTest(b)()
+	defer log.Scope(b).Close(b)
+
+	cases := []struct {
+		name    string
+		builtin string
+	}{
+		{
+			"noError",
+			"crdb_internal.void_func()",
+		},
+		{
+			"expectedWithCode",
+			"crdb_internal.force_error('01000', 'oops')",
+		},
+		{
+			"expectedAssertion",
+			"crdb_internal.force_assertion_error('oops')",
+		},
+		{
+			"internal",
+			"crdb_internal.force_vectorized_assertion_error('oops')",
+		},
+	}
+
+	sqlFmt := `SELECT %s,
+  0,
+  '',
+  0.0,
+  NULL,
+  '2000-01-01 00:00:00'::timestamptz,
+  b'00000000',
+  i + 0,
+  i * 1.5,
+  i / 100
+  FROM generate_series(0, 0) AS s(i)
+`
+
+	ctx := context.Background()
+	s := serverutils.StartServerOnly(b, base.TestServerArgs{SQLMemoryPoolSize: 10 << 30})
+	defer s.Stopper().Stop(ctx)
+
+	for _, parallelism := range []int{1, 20, 50} {
+		numConns := runtime.GOMAXPROCS(0) * parallelism
+		b.Run(fmt.Sprintf("conns=%d", numConns), func(b *testing.B) {
+			for _, tc := range cases {
+				stmt := fmt.Sprintf(sqlFmt, tc.builtin)
+				b.Run(tc.name, func(b *testing.B) {
+					// Create as many warm connections as we will need for the benchmark.
+					conns := make(chan *gosql.DB, numConns)
+					for range numConns {
+						conn := s.ApplicationLayer().SQLConn(b, serverutils.DBName(""))
+						// Make sure we're using local, vectorized execution.
+						sqlDB := sqlutils.MakeSQLRunner(conn)
+						sqlDB.Exec(b, "SET distsql = off")
+						sqlDB.Exec(b, "SET vectorize = on")
+						// Warm up the connection by executing the statement once. We should
+						// always go through the query plan cache after this.
+						_, _ = conn.Exec(stmt)
+						conns <- conn
+					}
+					b.SetParallelism(parallelism)
+					b.ResetTimer()
+					b.RunParallel(func(pb *testing.PB) {
+						var conn *gosql.DB
+						select {
+						case conn = <-conns:
+						default:
+							b.Fatal("not enough warm connections")
+						}
+						for pb.Next() {
+							_, _ = conn.Exec(stmt)
+						}
+					})
+				})
+			}
+		})
+	}
+}

pkg/sql/colexecerror/main_test.go

@@ -0,0 +1,90 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package colexecerror_test
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/col/coldata"
+	"github.com/cockroachdb/cockroach/pkg/col/coldataext"
+	"github.com/cockroachdb/cockroach/pkg/security/securityassets"
+	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexectestutils"
+	"github.com/cockroachdb/cockroach/pkg/sql/colexecerror"
+	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
+	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
+	"github.com/cockroachdb/cockroach/pkg/util/mon"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+//go:generate ../../util/leaktest/add-leaktest.sh *_test.go
+
+var (
+	// testAllocator is an Allocator with an unlimited budget for use in tests.
+	testAllocator     *colmem.Allocator
+	testColumnFactory coldata.ColumnFactory
+
+	// testMemMonitor and testMemAcc are a test monitor with an unlimited budget
+	// and a memory account bound to it for use in tests.
+	testMemMonitor *mon.BytesMonitor
+	testMemAcc     *mon.BoundAccount
+
+	// testDiskMonitor and testDiskmAcc are a test monitor with an unlimited budget
+	// and a disk account bound to it for use in tests.
+	testDiskMonitor *mon.BytesMonitor
+	testDiskAcc     *mon.BoundAccount
+)
+
+func TestMain(m *testing.M) {
+	securityassets.SetLoader(securitytest.EmbeddedAssets)
+	randutil.SeedForTests()
+	serverutils.InitTestServerFactory(server.TestServerFactory)
+	serverutils.InitTestClusterFactory(testcluster.TestClusterFactory)
+	os.Exit(func() int {
+		ctx := context.Background()
+		st := cluster.MakeTestingClusterSettings()
+		testMemMonitor = execinfra.NewTestMemMonitor(ctx, st)
+		defer testMemMonitor.Stop(ctx)
+		memAcc := testMemMonitor.MakeBoundAccount()
+		testMemAcc = &memAcc
+		evalCtx := eval.MakeTestingEvalContext(st)
+		testColumnFactory = coldataext.NewExtendedColumnFactory(&evalCtx)
+		testAllocator = colmem.NewAllocator(ctx, testMemAcc, testColumnFactory)
+		defer testMemAcc.Close(ctx)
+
+		testDiskMonitor = execinfra.NewTestDiskMonitor(ctx, cluster.MakeTestingClusterSettings())
+		defer testDiskMonitor.Stop(ctx)
+		diskAcc := testDiskMonitor.MakeBoundAccount()
+		testDiskAcc = &diskAcc
+		defer testDiskAcc.Close(ctx)
+
+		flag.Parse()
+		if !skip.UnderBench() {
+			// (If we're running benchmarks, don't set a random batch size.)
+			randomBatchSize := colexectestutils.GenerateBatchSize()
+			fmt.Printf("coldata.BatchSize() is set to %d\n", randomBatchSize)
+			if err := coldata.SetBatchSizeForTests(randomBatchSize); err != nil {
+				colexecerror.InternalError(err)
+			}
+		}
+		return m.Run()
+	}())
+}

pkg/sql/sem/builtins/builtins.go

@@ -5642,6 +5642,31 @@ SELECT
 		},
 	),
 
+	"crdb_internal.force_vectorized_assertion_error": makeBuiltin(
+		tree.FunctionProperties{
+			Category: builtinconstants.CategorySystemInfo,
+		},
+		tree.Overload{
+			Types:      tree.ParamTypes{{Name: "msg", Typ: types.String}},
+			ReturnType: tree.FixedReturnType(types.Int),
+			Fn: func(_ context.Context, _ *eval.Context, args tree.Datums) (tree.Datum, error) {
+				s, ok := tree.AsDString(args[0])
+				if !ok {
+					return nil, errors.Newf("expected string value, got %T", args[0])
+				}
+				msg := string(s)
+				err := errors.AssertionFailedf("%s", msg)
+				// Panic instead of returning the error. The vectorized panic-catcher
+				// will catch the panic and convert it into an internal error.
+				colexecerror.InternalError(err)
+				// This code is unreachable.
+				panic(err)
+			},
+			Info:       "This function is used only by CockroachDB's developers for testing purposes.",
+			Volatility: volatility.Volatile,
+		},
+	),
+
 	"crdb_internal.void_func": makeBuiltin(
 		tree.FunctionProperties{
 			Category: builtinconstants.CategorySystemInfo,

pkg/sql/sem/builtins/fixed_oids.go

@@ -2573,6 +2573,7 @@ var builtinOidsArray = []string{
 	2605: `merge_aggregated_stmt_metadata(arg1: jsonb) -> jsonb`,
 	2606: `crdb_internal.protect_mvcc_history(timestamp: decimal, expiration_window: interval, description: string) -> int`,
 	2607: `crdb_internal.extend_mvcc_history_protection(job_id: int) -> void`,
+	2608: `crdb_internal.force_vectorized_assertion_error(msg: string) -> int`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid