sql: add cluster setting to disable RU estimation

This patch adds a cluster setting, `sql.tenant_ru_estimation.enabled`, which is used to determine whether tenants collect an RU estimate for queries run with `EXPLAIN ANALYZE`. This is an escape hatch so that the RU estimation logic can be more safely backported. Informs cockroachdb#74441 Release note: None
DrewKimball · Dec 8, 2022 · 0c7a184 · 0c7a184
1 parent 111e190
commit 0c7a184
Show file tree

Hide file tree

Showing 9 changed files with 36 additions and 12 deletions.
diff --git a/pkg/ccl/multitenantccl/tenantcostclient/tenant_side.go b/pkg/ccl/multitenantccl/tenantcostclient/tenant_side.go
@@ -789,10 +789,13 @@ func (c *tenantSideCostController) OnResponseWait(
 	}
 
 	// Record the number of RUs consumed by the IO request.
-	if sp := tracing.SpanFromContext(ctx); sp != nil && sp.RecordingType() != tracingpb.RecordingOff {
-		sp.RecordStructured(&roachpb.TenantConsumption{
-			RU: float64(totalRU),
-		})
+	if multitenant.TenantRUEstimateEnabled.Get(&c.settings.SV) {
+		if sp := tracing.SpanFromContext(ctx); sp != nil &&
+			sp.RecordingType() != tracingpb.RecordingOff {
+			sp.RecordStructured(&roachpb.TenantConsumption{
+				RU: float64(totalRU),
+			})
+		}
 	}
 
 	c.mu.Lock()

diff --git a/pkg/multitenant/BUILD.bazel b/pkg/multitenant/BUILD.bazel
@@ -16,6 +16,7 @@ go_library(
         "//pkg/kv",
         "//pkg/multitenant/tenantcostmodel",
         "//pkg/roachpb",
+        "//pkg/settings",
         "//pkg/sql/sqlliveness",
         "//pkg/sql/sqlutil",
         "//pkg/util/metric",

diff --git a/pkg/multitenant/cost_controller.go b/pkg/multitenant/cost_controller.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcostmodel"
+	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 )
@@ -131,3 +132,12 @@ type TenantSideExternalIORecorder interface {
 type exemptCtxValueType struct{}
 
 var exemptCtxValue interface{} = exemptCtxValueType{}
+
+// TenantRUEstimateEnabled determines whether EXPLAIN ANALYZE should return an
+// estimate for the number of RUs consumed by tenants.
+var TenantRUEstimateEnabled = settings.RegisterBoolSetting(
+	settings.TenantWritable,
+	"sql.tenant_ru_estimation.enabled",
+	"determines whether explain analyze should return an estimate for the query's RU consumption",
+	true,
+)
diff --git a/pkg/multitenant/multitenantcpu/cpu_usage.go b/pkg/multitenant/multitenantcpu/cpu_usage.go
@@ -50,7 +50,7 @@ func (h *CPUUsageHelper) StartCollection(
 
 // EndCollection should be called at the end of execution for a flow in order to
 // get the estimated number of RUs consumed due to CPU usage. It returns zero
-// for non-tenants.
+// for non-tenants. It is a no-op if StartCollection was never called.
 func (h *CPUUsageHelper) EndCollection(ctx context.Context) (ruFomCPU float64) {
 	if h.costController == nil || h.costController.GetCostConfig() == nil {
 		return 0

diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
@@ -22,6 +22,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/multitenant"
 	"github.com/cockroachdb/cockroach/pkg/multitenant/multitenantcpu"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
@@ -1038,9 +1039,11 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	ex.sessionTracing.TracePlanStart(ctx, stmt.AST.StatementTag())
 	ex.statsCollector.PhaseTimes().SetSessionPhaseTime(sessionphase.PlannerStartLogicalPlan, timeutil.Now())
 
-	if server := ex.server.cfg.DistSQLSrv; server != nil {
-		// Begin measuring CPU usage for tenants. This is a no-op for non-tenants.
-		ex.cpuStatsCollector.StartCollection(ctx, server.TenantCostController)
+	if multitenant.TenantRUEstimateEnabled.Get(ex.server.cfg.SV()) {
+		if server := ex.server.cfg.DistSQLSrv; server != nil {
+			// Begin measuring CPU usage for tenants. This is a no-op for non-tenants.
+			ex.cpuStatsCollector.StartCollection(ctx, server.TenantCostController)
+		}
 	}
 
 	// If adminAuditLogging is enabled, we want to check for HasAdminRole
@@ -1272,7 +1275,7 @@ func populateQueryLevelStatsAndRegions(
 	} else {
 		// If this query is being run by a tenant, record the RUs consumed by CPU
 		// usage and network egress to the client.
-		if cfg.DistSQLSrv != nil {
+		if multitenant.TenantRUEstimateEnabled.Get(cfg.SV()) && cfg.DistSQLSrv != nil {
 			if costController := cfg.DistSQLSrv.TenantCostController; costController != nil {
 				if costCfg := costController.GetCostConfig(); costCfg != nil {
 					networkEgressRUEstimate := costCfg.PGWireEgressCost(topLevelStats.networkEgressEstimate)

diff --git a/pkg/sql/distsql_running.go b/pkg/sql/distsql_running.go
@@ -23,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangecache"
+	"github.com/cockroachdb/cockroach/pkg/multitenant"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
@@ -631,7 +632,8 @@ func (dsp *DistSQLPlanner) Run(
 
 	recv.outputTypes = plan.GetResultTypes()
 	recv.contendedQueryMetric = dsp.distSQLSrv.Metrics.ContendedQueriesCount
-	if dsp.distSQLSrv.TenantCostController != nil && planCtx.planner != nil {
+	if multitenant.TenantRUEstimateEnabled.Get(&dsp.st.SV) &&
+		dsp.distSQLSrv.TenantCostController != nil && planCtx.planner != nil {
 		if instrumentation := planCtx.planner.curPlan.instrumentation; instrumentation != nil {
 			// Only collect the network egress estimate for a tenant that is running
 			// EXPLAIN ANALYZE, since the overhead is non-negligible.

diff --git a/pkg/sql/flowinfra/BUILD.bazel b/pkg/sql/flowinfra/BUILD.bazel
@@ -19,6 +19,7 @@ go_library(
     deps = [
         "//pkg/base",
         "//pkg/kv",
+        "//pkg/multitenant",
         "//pkg/roachpb",
         "//pkg/server/telemetry",
         "//pkg/settings",

diff --git a/pkg/sql/flowinfra/flow.go b/pkg/sql/flowinfra/flow.go
@@ -15,6 +15,7 @@ import (
 	"sync"
 
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/multitenant"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra/execopnode"
@@ -412,7 +413,8 @@ func (f *FlowBase) StartInternal(
 
 	f.status = flowRunning
 
-	if !f.Gateway && f.CollectStats {
+	if multitenant.TenantRUEstimateEnabled.Get(&f.Cfg.Settings.SV) &&
+		!f.Gateway && f.CollectStats {
 		// Remote flows begin collecting CPU usage here, and finish when the last
 		// outbox finishes. Gateway flows are handled by the connExecutor.
 		f.FlowCtx.TenantCPUMonitor.StartCollection(ctx, f.Cfg.TenantCostController)

diff --git a/pkg/sql/instrumentation.go b/pkg/sql/instrumentation.go
@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/multitenant"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings"
@@ -250,7 +251,8 @@ func (ih *instrumentationHelper) Setup(
 	ih.implicitTxn = implicitTxn
 	ih.codec = cfg.Codec
 	ih.origCtx = ctx
-	ih.isTenant = cfg.DistSQLSrv != nil && cfg.DistSQLSrv.TenantCostController != nil
+	ih.isTenant = multitenant.TenantRUEstimateEnabled.Get(cfg.SV()) && cfg.DistSQLSrv != nil &&
+		cfg.DistSQLSrv.TenantCostController != nil
 
 	switch ih.outputMode {
 	case explainAnalyzeDebugOutput: