Merge #123414 #123696

123414: sql: grant/revoke on a large number of objects can create a lot of jobs r=Dedej-Bergin a=Dedej-Bergin

Previously we would create multiple jobs for granting on multiple tables and types within one transaction.  This caused a performance slowdown, these code changes skip the making of multiple jobs and just execute the grants in a batch.

Fixes:#117643
Release note(performance improvement): multiple grants on tables and types within one transaction now run faster.

123696: sqlstats: fix sampling logic for first recording of a statement r=xinhaoz a=xinhaoz

By default we sample execution stats for 1% of statements,
(`sql.txn_stats.sample_rate`). We also turn on sampling if
this is the first time we're recording this fingerprint in the
application container. This 'sample on first encounter' logic did
not factor in the scenario where we may be unable to write new
fingerprints due to reaching max container capacities. Previously,
this meant that once the containers were full we'd begin to trace
every fingerprint that did not already get recorded.

This commit ensures that we also check the capacity when deciding
to sample a fingerprint we haven't recorded yet, skipping the sampling
if it is full.  Note that although the container may reach capacity
during the statement's execution, it's more important that we stop sampling
once we are sure the fingerprint cannot be written to the container.

Fixes: #123690

Release note: None

Co-authored-by: Bergin Dedej <bergin.dedej@cockroachlabs.com>
Co-authored-by: Xin Hao Zhang <xzhang@cockroachlabs.com>

pkg/bench/rttanalysis/testdata/benchmark_expectations

@@ -56,9 +56,9 @@ exp,benchmark
 5,GenerateObjects/generate_10_tables
 16,GenerateObjects/generate_10x10_schemas_and_tables_in_existing_db
 5,GenerateObjects/generate_50000_tables
-15,Grant/grant_all_on_1_table
-19,Grant/grant_all_on_2_tables
-23,Grant/grant_all_on_3_tables
+10,Grant/grant_all_on_1_table
+10,Grant/grant_all_on_2_tables
+10,Grant/grant_all_on_3_tables
 19,GrantRole/grant_1_role
 25,GrantRole/grant_2_roles
 6,Jobs/cancel_job
@@ -104,9 +104,9 @@ exp,benchmark
 4,ORMQueries/pg_type
 133,ORMQueries/prisma_column_descriptions
 3,ORMQueries/prisma_column_descriptions_updated
-15,Revoke/revoke_all_on_1_table
-19,Revoke/revoke_all_on_2_tables
-23,Revoke/revoke_all_on_3_tables
+10,Revoke/revoke_all_on_1_table
+10,Revoke/revoke_all_on_2_tables
+10,Revoke/revoke_all_on_3_tables
 17,RevokeRole/revoke_1_role
 21,RevokeRole/revoke_2_roles
 14,ShowGrants/grant_2_roles

pkg/sql/conn_executor.go

@@ -1145,6 +1145,7 @@ func (s *Server) newConnExecutor(
 		applicationStats,
 		ex.server.insights.Writer(ex.sessionData().Internal),
 		ex.phaseTimes,
+		s.sqlStats.GetCounters(),
 		s.cfg.SQLStatsTestingKnobs,
 	)
 	ex.dataMutatorIterator.onApplicationNameChange = func(newName string) {

pkg/sql/conn_executor_exec.go

@@ -3255,17 +3255,17 @@ func (ex *connExecutor) recordTransactionStart(txnID uuid.UUID) {
 	ex.extraTxnState.transactionStatementsHash = util.MakeFNV64()
 	ex.extraTxnState.transactionStatementFingerprintIDs = nil
 	ex.extraTxnState.numRows = 0
-	ex.extraTxnState.shouldCollectTxnExecutionStats = false
 	ex.extraTxnState.accumulatedStats = execstats.QueryLevelStats{}
 	ex.extraTxnState.idleLatency = 0
 	ex.extraTxnState.rowsRead = 0
 	ex.extraTxnState.bytesRead = 0
 	ex.extraTxnState.rowsWritten = 0
 	ex.extraTxnState.rowsWrittenLogged = false
 	ex.extraTxnState.rowsReadLogged = false
-	if txnExecStatsSampleRate := collectTxnStatsSampleRate.Get(&ex.server.GetExecutorConfig().Settings.SV); txnExecStatsSampleRate > 0 {
-		ex.extraTxnState.shouldCollectTxnExecutionStats = txnExecStatsSampleRate > ex.rng.internal.Float64()
-	}
+
+	txnExecStatsSampleRate := collectTxnStatsSampleRate.Get(&ex.server.GetExecutorConfig().Settings.SV)
+	ex.extraTxnState.shouldCollectTxnExecutionStats = !ex.server.cfg.TestingKnobs.DisableProbabilisticSampling &&
+		txnExecStatsSampleRate > ex.rng.internal.Float64()
 
 	// Note ex.metrics is Server.Metrics for the connExecutor that serves the
 	// client connection, and is Server.InternalMetrics for internal executors.
@@ -3309,12 +3309,6 @@ func (ex *connExecutor) recordTransactionFinish(
 	txnTime := txnEnd.Sub(txnStart)
 	ex.totalActiveTimeStopWatch.Stop()
 
-	if ex.server.cfg.TestingKnobs.OnRecordTxnFinish != nil {
-		ex.server.cfg.TestingKnobs.OnRecordTxnFinish(
-			ex.executorType == executorTypeInternal, ex.phaseTimes, ex.planner.stmt.SQL,
-		)
-	}
-
 	// Note ex.metrics is Server.Metrics for the connExecutor that serves the
 	// client connection, and is Server.InternalMetrics for internal executors.
 	if contentionDuration := ex.extraTxnState.accumulatedStats.ContentionTime.Nanoseconds(); contentionDuration > 0 {
@@ -3370,6 +3364,12 @@ func (ex *connExecutor) recordTransactionFinish(
 		TxnErr:      txnErr,
 	}
 
+	if ex.server.cfg.TestingKnobs.OnRecordTxnFinish != nil {
+		ex.server.cfg.TestingKnobs.OnRecordTxnFinish(
+			ex.executorType == executorTypeInternal, ex.phaseTimes, ex.planner.stmt.SQL, recordedTxnStats,
+		)
+	}
+
 	ex.maybeRecordRetrySerializableContention(ev.txnID, transactionFingerprintID, txnErr)
 
 	if ex.extraTxnState.shouldLogToTelemetry {

pkg/sql/exec_util.go

@@ -1640,7 +1640,7 @@ type ExecutorTestingKnobs struct {
 
 	// OnRecordTxnFinish, if set, will be called as we record a transaction
 	// finishing.
-	OnRecordTxnFinish func(isInternal bool, phaseTimes *sessionphase.Times, stmt string)
+	OnRecordTxnFinish func(isInternal bool, phaseTimes *sessionphase.Times, stmt string, txnStats sqlstats.RecordedTxnStats)
 
 	// UseTransactionDescIDGenerator is used to force descriptor ID generation
 	// to use a transaction, and, in doing so, more deterministically allocate
@@ -1662,6 +1662,13 @@ type ExecutorTestingKnobs struct {
 	// ForceSQLLivenessSession will force the use of a sqlliveness session for
 	// transaction deadlines even in the system tenant.
 	ForceSQLLivenessSession bool
+
+	// DisableProbabilisticSampling, if set to true, will disable
+	// probabilistic transaction sampling. This is important for tests that
+	// want to deterministically test cases where we turn on transaction sampling
+	// due to some other condition. We can't set the probability to 0 since
+	// that would disable the feature entirely.
+	DisableProbabilisticSampling bool
 }
 
 // PGWireTestingKnobs contains knobs for the pgwire module.

pkg/sql/grant_revoke.go

@@ -391,17 +391,7 @@ func (n *changeDescriptorBackedPrivilegesNode) startExec(params runParams) error
 					DatabaseName:                   (*tree.Name)(&d.Name).String(),
 				})
 			}
-
 		case *tabledesc.Mutable:
-			// TODO (lucy): This should probably have a single consolidated job like
-			// DROP DATABASE.
-			if err := p.createOrUpdateSchemaChangeJob(
-				ctx, d,
-				fmt.Sprintf("updating privileges for table %d", d.ID),
-				descpb.InvalidMutationID,
-			); err != nil {
-				return err
-			}
 			if !d.Dropped() {
 				if err := p.writeSchemaChangeToBatch(ctx, d, b); err != nil {
 					return err
@@ -419,9 +409,11 @@ func (n *changeDescriptorBackedPrivilegesNode) startExec(params runParams) error
 				})
 			}
 		case *typedesc.Mutable:
-			err := p.writeTypeSchemaChange(ctx, d, fmt.Sprintf("updating privileges for type %d", d.ID))
-			if err != nil {
-				return err
+			if !d.Dropped() {
+				err := p.writeDescToBatch(ctx, d, b)
+				if err != nil {
+					return err
+				}
 			}
 			for _, grantee := range n.grantees {
 				privs := eventDetails // copy the granted/revoked privilege list.
@@ -435,12 +427,11 @@ func (n *changeDescriptorBackedPrivilegesNode) startExec(params runParams) error
 				})
 			}
 		case *schemadesc.Mutable:
-			if err := p.writeSchemaDescChange(
-				ctx,
-				d,
-				fmt.Sprintf("updating privileges for schema %d", d.ID),
-			); err != nil {
-				return err
+			if !d.Dropped() {
+				err := p.writeDescToBatch(ctx, d, b)
+				if err != nil {
+					return err
+				}
 			}
 			for _, grantee := range n.grantees {
 				privs := eventDetails // copy the granted/revoked privilege list.
@@ -454,8 +445,11 @@ func (n *changeDescriptorBackedPrivilegesNode) startExec(params runParams) error
 				})
 			}
 		case *funcdesc.Mutable:
-			if err := p.writeFuncSchemaChange(ctx, d); err != nil {
-				return err
+			if !d.Dropped() {
+				err := p.writeDescToBatch(ctx, d, b)
+				if err != nil {
+					return err
+				}
 			}
 			for _, grantee := range n.grantees {
 				privs := eventDetails // copy the granted/revoked privilege list.

pkg/sql/grant_revoke_test.go

@@ -15,6 +15,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -239,3 +240,82 @@ func TestNoOpRevoke(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkGrantTables(b *testing.B) {
+	defer leaktest.AfterTest(b)()
+	defer log.Scope(b).Close(b)
+	ctx := context.Background()
+
+	for _, numTables := range []int{10, 100, 1000} {
+		b.Run(fmt.Sprintf("numTables=%d", numTables), func(b *testing.B) {
+			srv, sqlDB, _ := serverutils.StartServer(b, base.TestServerArgs{})
+			defer srv.Stopper().Stop(ctx)
+
+			sqlRun := sqlutils.MakeSQLRunner(sqlDB)
+			sqlRun.Exec(b, `CREATE DATABASE t;`)
+			sqlRun.Exec(b, `USE t;`)
+
+			sqlRun.Exec(b, `CREATE USER ROACH;`)
+
+			for i := 0; i < numTables; i++ {
+				sqlRun.Exec(b, fmt.Sprintf(`CREATE TABLE t.a%d (k INT PRIMARY KEY);`, i))
+			}
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				sqlRun.Exec(b, `GRANT ALL ON * TO ROACH;`)
+				sqlRun.Exec(b, `REVOKE ALL ON * FROM ROACH;`)
+			}
+		})
+	}
+}
+
+func BenchmarkGrantTypes(b *testing.B) {
+	defer leaktest.AfterTest(b)()
+	defer log.Scope(b).Close(b)
+	ctx := context.Background()
+
+	for _, numTypes := range []int{10, 100, 1000} {
+		b.Run(fmt.Sprintf("numTypes=%d", numTypes), func(b *testing.B) {
+			srv, sqlDB, _ := serverutils.StartServer(b, base.TestServerArgs{})
+			defer srv.Stopper().Stop(ctx)
+
+			sqlRun := sqlutils.MakeSQLRunner(sqlDB)
+			sqlRun.Exec(b, `CREATE DATABASE t;`)
+			sqlRun.Exec(b, `USE t;`)
+
+			sqlRun.Exec(b, `CREATE USER ROACH;`)
+
+			for i := 0; i < numTypes; i++ {
+				sqlRun.Exec(b, fmt.Sprintf(`CREATE TYPE a%d AS ENUM ('roach1', 'roach2', 'roach3');`, i))
+			}
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				txn := sqlRun.Begin(b)
+				for i := 0; i < numTypes; i++ {
+					_, err := txn.Exec(fmt.Sprintf(`GRANT ALL ON TYPE a%d TO ROACH;`, i))
+					if err != nil {
+						return
+					}
+				}
+				err := txn.Commit()
+				if err != nil {
+					return
+				}
+
+				txn = sqlRun.Begin(b)
+				for i := 0; i < numTypes; i++ {
+					_, err = txn.Exec(fmt.Sprintf(`REVOKE ALL ON TYPE a%d FROM ROACH;`, i))
+					if err != nil {
+						return
+					}
+				}
+				err = txn.Commit()
+				if err != nil {
+					return
+				}
+			}
+		})
+	}
+}

pkg/sql/instrumentation.go

@@ -470,20 +470,24 @@ func (ih *instrumentationHelper) Setup(
 		}
 	}
 
-	ih.collectExecStats = collectTxnExecStats
-
-	// Don't collect it if Stats Collection is disabled. If it is disabled the
-	// stats are not stored, so it always returns false for previouslySampled.
-	// TODO(117690): Unify StmtStatsEnable and TxnStatsEnable into a single cluster setting.
-	if !collectTxnExecStats && (!previouslySampled && sqlstats.StmtStatsEnable.Get(&cfg.Settings.SV)) {
-		// We don't collect the execution stats for statements in this txn, but
-		// this is the first time we see this statement ever, so we'll collect
-		// its execution stats anyway (unless the user disabled txn stats
-		// collection entirely).
-		statsCollectionDisabled := collectTxnStatsSampleRate.Get(&cfg.Settings.SV) == 0
-		ih.collectExecStats = !statsCollectionDisabled
+	shouldSampleFirstEncounter := func() bool {
+		// If this is the first time we see this statement in the current stats
+		// container, we'll collect its execution stats anyway (unless the user
+		// disabled txn or stmt stats collection entirely).
+		// TODO(117690): Unify StmtStatsEnable and TxnStatsEnable into a single cluster setting.
+		if collectTxnStatsSampleRate.Get(&cfg.Settings.SV) == 0 ||
+			!sqlstats.StmtStatsEnable.Get(&cfg.Settings.SV) {
+			return false
+		}
+
+		// We don't want to collect the stats if the stats container is full,
+		// since previouslySampled will always return false for statements
+		// not already in the container.
+		return !previouslySampled && !statsCollector.StatementsContainerFull()
 	}
 
+	ih.collectExecStats = collectTxnExecStats || shouldSampleFirstEncounter()
+
 	if !ih.collectBundle && ih.withStatementTrace == nil && ih.outputMode == unmodifiedOutput {
 		if ih.collectExecStats {
 			// If we need to collect stats, create a child span with structured