diff --git a/pkg/bench/BUILD.bazel b/pkg/bench/BUILD.bazel index 874ac8491284..5feb15dac094 100644 --- a/pkg/bench/BUILD.bazel +++ b/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", diff --git a/pkg/bench/foreachdb.go b/pkg/bench/foreachdb.go index efcbf8c8a675..6364dcc58362 100644 --- a/pkg/bench/foreachdb.go +++ b/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)