bench: exempt shared process tenants from rate limits

This exempts the tenant from the KV-side tenant limiter. We expect
this to be the default configuration of most shared-process tenants in
the near term. Thus, setting this during benchmark setups provides
more accurate benchmarks.

Epic: none

Release note: None

pkg/bench/BUILD.bazel

@@ -13,12 +13,15 @@ go_library(
     deps = [
         "//pkg/base",
         "//pkg/ccl",
+        "//pkg/multitenant/tenantcapabilities",
         "//pkg/roachpb",
         "//pkg/server",
+        "//pkg/testutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
+        "@com_github_cockroachdb_errors//:errors",
         "@com_github_go_sql_driver_mysql//:mysql",
         "@com_github_lib_pq//:pq",
         "@com_github_stretchr_testify//require",

pkg/bench/foreachdb.go

@@ -23,12 +23,15 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	_ "github.com/cockroachdb/cockroach/pkg/ccl"
+	"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
+	"github.com/cockroachdb/errors"
 	_ "github.com/go-sql-driver/mysql" // registers the MySQL driver to gosql
 	_ "github.com/lib/pq"              // registers the pg driver to gosql
 	"github.com/stretchr/testify/require"
@@ -78,6 +81,27 @@ func benchmarkSharedProcessTenantCockroach(b *testing.B, f BenchmarkFn) {
 	_, err = db.Exec(`ALTER TENANT ALL SET CLUSTER SETTING "spanconfig.tenant_limit" = 10000000`)
 	require.NoError(b, err)
 
+	// Exempt the tenant from rate limiting. We expect most
+	// shared-process tenants will run without rate limiting in
+	// the near term.
+	_, err = db.Exec(`ALTER TENANT benchtenant GRANT CAPABILITY exempt_from_rate_limiting`)
+	require.NoError(b, err)
+
+	var tenantID uint64
+	require.NoError(b, db.QueryRow(`SELECT id FROM [SHOW TENANT benchtenant]`).Scan(&tenantID))
+
+	err = testutils.SucceedsSoonError(func() error {
+		capabilities, found := s.(*server.TestServer).Server.TenantCapabilitiesReader().GetCapabilities(roachpb.MustMakeTenantID(tenantID))
+		if !found {
+			return errors.Newf("capabilities not yet ready")
+		}
+		if !capabilities.GetBool(tenantcapabilities.ExemptFromRateLimiting) {
+			return errors.Newf("capabilities not yet ready")
+		}
+		return nil
+	})
+	require.NoError(b, err)
+
 	_, err = tenantDB.Exec(`CREATE DATABASE bench`)
 	require.NoError(b, err)